Add ivfaster: a segment-lifecycle-aware IVF vector codec for the sandbox - #16567
Add ivfaster: a segment-lifecycle-aware IVF vector codec for the sandbox#16567RKSPD wants to merge 2 commits into
Conversation
An IVF vector format built around two quantization tiers and an exhaustive
router. Documents are clustered into nlist cells by Lloyd iteration, and a
query selects a few cells and scores only their documents.
* Coarse tier (nitrox2): a symmetric 2-bit thermometer code whose summed
per-dimension level distance equals popcount(q ^ d) over the whole code,
so scoring a cell is one XOR and popcount over a contiguous byte string.
* Fine tier (int8): one byte per dimension, ranking the coarse survivors.
* The Reaper: after a centroid update, only documents that could have
changed cell are re-routed, on a per-pair movement bound that provably
contains every document whose assignment changes.
* Clustering runs to convergence, stopping once a pass changes at most
0.5% of assignments, with lloydIters as a backstop for a field that does
not converge.
* Spill: near-boundary documents are written into additional cells chosen
by the SOAR objective, so a query reaches them at a lower nprobe.
* A centroid graph serves search-time cell selection, built from the final
centroids. Index-time routing is always an exhaustive scan and never
consults the graph.
Centroids are unit norm for every similarity, so argmin -dot and
argmin ||v - c||^2 coincide, both Lloyd steps minimize the same objective,
and every distance in the codec reduces to one dot product.
Tests cover the coarse and fine codes, the routing cascade, clustering and
the Reaper's sufficiency claim, spill completeness, merge, and end-to-end
recall through the real codec.
|
exciting! the numbers you shared look compelling. Can you address the failing checks? It looks like a reference to the new module crept in to another module somehow? And there are some warnings that need to be cleaned up. But beyond that, this is a lot of new code to review. I wonder if you can help reviewers by (1) providing some overview of the algorithm and the new classes. If you included that in the PR javadocs or module.java, that's ideal, but please reference from the PR description if so. Also (2) is there any code here that's optional for a first implementation? For example I see you included multiple quantization options. Can we split this up and look at a stripped down version without quantization support? I'm afraid this may languish if we can't simplify it. |
Hi Mike! Thank you for commenting on my PR. I agree this is really huge and hard to review. I'm working on breaking the codec down into reviewable pieces, but just wanted to get the results out so the community can get a feel for the results and the implementation scope. Unlike other codecs, IVFaster necessarily uses a cascading coarse/fine quantizer in both the index and search phases, which makes it difficult to think about the IVFaster algorithm without these pieces. The performance characteristics are largely based on the new quantizer design. I'll get more comprehensive documentation and architecture diagrams soon. Once again thank you. |
|
Disclaimer: Rikhil and I work together and started a conversation about this earlier today, but I wanted to bring it here so others have an opportunity to join the discussion. A property that really excites me about IVF is that docIDs can be traversed in monotonically increasing order within each cluster since it's fundamentally just a postings list of docs (and it also follows that a disjunction of clusters can produce a union of ordered docIDs, exactly the same as something like I'm wondering if we could restructure the approach here to allow for this doc-at-a-time scoring. I'm going to look at the code in more detail soon (so I may have an overly-naive mental model of what's going on), but it seems to me like we could create a new Query/Collector implementations that do the following:
What I'd love to achieve here is a way to run filters doc-at-a-time before doing expensive full-precision scoring. We fundamentally can't do this with HNSW since it can't visit docs in docID order (same problem we have with the points index, just much more expensive because of vector scoring). I'm positive I'm missing some important bits of IVF and your proposed implementation, but I wanted to float this structure to see what you think. |
3 participants
IVFaster: a segment-lifecycle-aware IVF vector codec (sandbox)
It's no secret that vector search sits awkwardly on Lucene's segment model. Segments are immutable and merge constantly, and both properties are cheap for an inverted index and expensive for an ANN structure. There have been attempts to improve this behavior, namely
IncrementalHnswGraphMerger, which adopts the largest incoming graph as a base and inserts the remaining documents instead of rebuilding a new graph from the document contents of both segments. However, HNSW construction itself is a very expensive operation: every inserted document is a beam search with a random access per hop, and insertion mutates the base graph's neighbor lists. Furthermore, with incremental merges, a base graph with more than 40% deletions is declined, because its connectivity has degraded, and the merge falls back to a full rebuild. All of these problems are made worse with quantization. Since standard quantization schemes create segment-dependent compressed vectors, the codec needs to keep full-precision vectors on disk just to requantize on merge. Related discussion: #15612.IVFaster takes the position that a vector index should look like a posting list, and is built around the Lucene segment architecture rather than adapted to it afterwards.
Read literally, that is what the IVF family already is: a cell id is a key, and the vectors routed to it are the values under that key. In IVFaster, a single document also can live in multiple cells. Document vectors near a cell boundary are listed under several keys so that a query probing any one of them finds the document, with
spillBitssetting how many extra cells a document may occupy.The structure a merge settles is
nlist-sized, not N-sized. The largest incoming segment donates its centroids, the other segments' documents are routed into them by a tiled scan over centroid codes, and a few Lloyd iterations settle the means. The only thing rebuilt from scratch is the centroid graph, overnlistnodes rather than N: in the benchmarks below, 8000 centroids against 1M documents.Quantization is global. The grid is derived in closed form from
dimafter a shared Hadamard rotation, with no trained or per-segment statistic, so a document's code is the same in every segment. Merges copy codes, and the codec never has to retain float32 vectors just to requantize later.Two tiers: a 2-bit coarse code quantizer (Nitrox2) narrows each probed cell to a shortlist with one XOR+popcount pass over a contiguous byte string, and an int8 fine tier ranks the survivors. Cells are chosen by greedy descent over a navigable small-world graph on the centroids. Each node's payload in the graph is the centroid's 2-bit code interleaved with its neighbor list, so a hop stays L3-resident and descent cost is set by
efrather than bynlist.The codec is capable of disk-based search. Slots are grouped by cell, so the byte range of every probed cell is known before the scan begins, and the reader hints all
nproberuns throughIndexInput#prefetchup front, so cold faults overlap in place of one synchronous fault per cell. A graph descent has no equivalent, because the next hop's address is unknown until the current node is scored, which makes its misses serial by construction. The hint is adaptive and collapses to a counter increment once pages are resident, so it is on by default.Note: an async or batched I/O path is deliberately absent, because it only benefits when nearly everything is cold and carries heavy overhead when warm. Every number below is page-cache warm, and disk-resident performance is coming soon.
Results
1M Cohere Embed multilingual-v3 Wikipedia-en, dim 1024, unit norm, dot product, 1000 held-out queries, recall@100 against exact NN, force-merged to one segment. Graviton3, Corretto 25, index on standard gp2 EBS (250MiB/s). Latency is warm and single-threaded through
luceneutil. IVFaster dialed withnprobeandnlist=8000, HNSW with search fanout. Every number below was measured in one session on one machine.spillBits=3spillBits=2The gap widens with recall, because HNSW buys recall by visiting more nodes at a memory latency each while IVFaster scans more cells sequentially.
spillBits=3spillBits=2Build is storage-bound on EBS rather than CPU-bound: about half the IVFaster build wall clock is
fsyncon the data file with the cores idle, so these figures track bytes written more than clustering work.Index size is the column IVFaster loses at
spillBits=3. A slot costs 1344 B and both the fine records and the coarse planes replicate per slot, so atspillMargin=1.40the index holds almost exactly four slots per document.spillBits=2trades that back: 25% smaller, 16% faster to build, and roughly 1.5x thenprobeto match a given recall, which costs 9 to 15% more latency in the 0.94-and-above bands.IVFaster config:
nlist=8000,spillMargin=1.40,soarLambda=1.0,lloydIters=10,bruteN=700,nprobeMargin=0.75,verifyMultiplier=2, Nitrox2 coarse + int8 fine. All codec defaults exceptnlist, whose default is 1000;nlistis corpus-dependent and not auto-scaled, and we see best results near 100-150 docs per cell. HNSW baseline:Lucene104HnswScalarQuantizedVectorsFormat, 7-bit scalar quantization,M=16,beamWidth=100, dialed with search fanout.Caveats: each point is a single run and latency noise is about 5%.
Scope
Hi all. This is a sandbox codec, marked
@lucene.experimentalwith no back-compat guarantee. It is a pretty big PR with 37 files underlucene/sandbox. It is self-contained with no other modifications to Lucene.I would really appreciate any feedback from the community. Thank you.