Knuth–Plass chunker to fix my obscurely broken semantic search

I found out recently that all my users were getting semantic search results that were slightly off--just sometimes. And the only reason I found out was an unrelated change that made that existing bug a lot more obvious.

I've been building this plugin that adds semantic search to Obsidian.MD--and the core algorithm I used evolved over the years becoming less and less naive;

• V1 'Full Naive': Take the note's first ~200 words that can fit in the embedding model I was using and use that as the semantic representative for the note.
• V2 'Mean Naive': Chunk all the notes in ~200 word windows with a bit of overlap between chunks--then let the average of all chunk's embedding be the representative.
• V3 'Mean Naive + Atomic Sentences': Same as 'Mean Naive', but I 'cut' the chunks at the sentence level--to prevent cutting a word or sentence in half.

Then I started looking into including this new functionality I've been experimenting with; using centroid chunks as pseudo-descriptors for notes, which required every chunk to be kept. So I refactored the core algorithm once again to V4;

• V4: 'Chunk representatives': V3, but rather than averaging every chunk, they're all saved and treated as a representative of the note.

It was at this point that a previously hidden artifact got surfaced in my testing: I was seeing some random notes show up in completely unrelated semantic searches. And to be honest, I didn't immediately scour my codebase for the issue--I took note of it in my mind and moved on. It wasn't until I decided to give my core functions a proper human review in the middle of the refactor that I saw the problem; the chunker was way too naive--it was taking ~200 word chunks out of a note until it ran out of words. Which leaves the possibility that the last chunk consists of just a couple words.

Do I think I was wrong for not jumping on the odd artifacts in the search results? I don't think so. The reality in designing an embedding-based system is that the results are always kind of vague. You see that one note is 57% similar to another--I read both notes and think 'yeah, I guess'.

The tricky thing with these systems is that it isn't blatant when they're off. It won't show an error or completely false result, it'll just be slightly off, making bugs harder to spot.

This artifact of the last embedding of the note being potential garbage wasn't a huge issue in V3, since all the chunk's embeddings were averaged. So, in the worst two-chunk case scenario the garbage chunk's impact would be halved by the healthy full chunk. But the V4 update made every chunk a canonical semantic representation of the entire note--leading to the random unrelated notes showing up in some search results. Garbage embeddings that were obscured before are now significantly overthrowing the search quality.

Proper chunking with the Knuth–Plass algorithm

This leaves a pretty rudimentary algorithm problem; how to chunk the text as equally as possible--while not cutting sentences and allow for an overlap? Which, conveniently for me, has already been solved somewhere else; sentence breaking using the Knuth-Plass algorithm. That algorithm ensures paragraphs can be displayed in neat equal length lines, or at least as much as it can without cutting a word. It does this by efficiently calculating the 'badness' of every possible configuration and taking the least bad option.

Greedy/naive fills each window to capacity and starts a new one, so whatever's left at the end is a chunk regardless of size. Knuth–Plass scores the chunking as a whole; a chunking configuration containing a single sentence end chunk scores far worse than one where six chunks are each just slightly under capacity.

This handles the odd ends; even dispersal is prioritized. All I have to do is change the input from words to whole sentences to get a proper chunking algorithm. Overlap can simply be added after the fact by budgeting some room for it in the window sizes given to the algorithm.

Conclusion

Chunking in semantic search (often a cornerstone in RAG systems) is easy to get wrong, since the output is so vague. It's difficult to tell when it's off. And in my case it wasn't only difficult to tell, it also happened pretty rarely. It wasn't until I made a change that the mistake that was always there became a lot more apparent--and I still didn't even pay attention to it until I realized the mistake when I reviewed the code.

In my case, the bug only affected about 500 people by giving them a less optimal search result when they're browsing through their notes--not the end of the world. Though when I'm developing RAG systems for clinical settings, one of those odd artifacts could just mean an LLM is looking at a document it has no business considering.