Skip to content

Add faster version of KL-Sum using numpy - #200

Open
mamei16 wants to merge 3 commits into
miso-belica:mainfrom
mamei16:main
Open

mamei16 wants to merge 3 commits into
miso-belica:mainfrom
mamei16:main

Conversation

@mamei16

@mamei16 mamei16 commented Dec 2, 2023

Copy link
Copy Markdown

Hi, thanks for creating this project!
I was looking for a summarization method for a project I'm working on and really liked KL-Sum. However, it was too slow for my specific case. Therefore, I rewrote it using vectorized operations. It is not a perfect replacement for the original, since the results can slightly differ. However, all original tests still pass. Feel free to reject this PR if you'd rather keep the original only.

Cheers

This comment was marked as outdated.

miso-belica added a commit that referenced this pull request Aug 14, 2026
Every content word of a sentence is also a content word of the whole
document, so the document's vocabulary can index a count array for each
sentence and for the summary. Merging two of them and computing the
divergence then becomes NumPy arithmetic instead of a Python loop over
a dictionary, which is where all the time was going.

The dictionary of a joint frequency contained only the words of the
sentence and of the summary, so a word that neither has, which now sits
in the array as a zero, has to be left out of the divergence just like a
word that the document does not have. PR #200 does that with a 42.0
sentinel and a mask; a zero frequency says the same thing without one,
because the word counts made it here consistently normalized and
filtered.

Divergences agree with the Python ones to 6.1e-16 over 60 random
documents, and the 13 picks that changed were all between candidates
within 1.1e-16 of each other.

400 sentence benchmark: 4.0s to 0.96s, which matches PR #200's 0.90s.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GhszyWASs6ccjLXvq7Kz5d
miso-belica added a commit that referenced this pull request Aug 14, 2026
Each pass called into NumPy about six times per candidate, on arrays as
short as the document's vocabulary, so the per call overhead was most of
the cost. Stack the sentence counts into one row per sentence and let a
pass work on the whole table at once, keeping the words on the last axis
so the divergence sums to one value per row.

Two things had to give up their scalar shape to make that work. The
empty word list case cannot be an early return any more, so the divisor
is held away from zero instead, which leaves the frequencies zero
because the counts of an empty list are zero anyway. And the divergence
cannot drop the words that neither side has by indexing with a mask,
because each row keeps its own words, so both sides get a frequency of
one there and contribute log(1) instead.

Summing over those extra zeros makes no difference to a sum, but NumPy
splits a longer sum differently: values move by at most 2.2e-16, and the
24 picks that changed over 60 random documents were all between
candidates that far apart. Summing a whole table at once is bit for bit
what summing its rows one by one gives.

Sentences are no longer removed from the parallel lists as they are
picked, since the table is indexed by the position in the document; a
list of the remaining positions decides what a pass scores.

400 sentence benchmark: 0.96s to 0.49s, against PR #200's 0.90s.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GhszyWASs6ccjLXvq7Kz5d
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants