[WIP] Multi-Vector support for HNSW search #13525
Conversation
pending: - merge fields - scorer changes - default scorer etc. - reader changes
|
As mentioned earlier, here is my rough plan for splitting this change into smaller PRs. Some of these steps could be merged if the impl. warrants it:
|
|
The more I think about it, the less I feel like the knn codec is the best choice for this feature (assuming that this issue is focused on late interaction models).
Using the knn codec to handle multi-vectors seems limiting, especially since it treats multi-vectors as a single unit for scoring. This works well for late interaction models, where we’re dealing with a collection of embeddings, but it’s restrictive if we want to index each vector separately. It could be helpful to explore other options instead of relying on the knn codec alone. Along those lines, I created a quick draft of a What do you think of this approach? It feels like we could skip the full knn framework if our main goal is just to score a bag of embeddings. This would keep things simpler and allow us to focus specifically on max similarity scoring without the added weight of the full knn codec. My main worry is that adding multi-vectors to the knn codec as a late interaction model might add complexity later. It’s really two different approaches, and it seems valuable to keep the option for indexing each vector separately. We could expose this flexibility through the aggregation function, but that might complicate things across all codecs, as they’d need to handle both aggregate and independent cases efficiently. |
|
One use-case for multi-vectors is indexing product aspects as separate embeddings for e-commerce search. At Amazon Product Search (where I work), we'd like to experiment with separate embeddings to represent product attributes, user product opinions, and product images. Such e-commerce use-cases would have a limited set of embeddings, but leverage similarity computations across all of them. I see your point about scaling challenges with very high cardinality multi-vectors like token level ColBERT embeddings. Keeping them in a However, I do think there is space for both solutions. It's not obvious to me how the knn codec gets polluted for future complexity. We would still support single vectors as is. My mental model is: if you want to use multi-vectors in nearest neighbor search (hnsw or newer algos later), index them in the knn field. Otherwise, index them separately as doc-values used only for re-ranking top results. |
This seems like just more than one knn field, or the nested field support. But, I understand the desire to add a multi-vector support to the flat codecs. I am honestly torn around whats the best path forward for the majority of users in Lucene. |
|
I tried to find some blogs and benchmarks on other library implementations. Astra Db, Vespa, faiss and nmslib, all seem to support multi-vectors in some form. From what I can tell, Astra Db and Vespa have ColBERT style multi-vector support in ANN [1] [2]. Benchmarks indicate ColBERT outperforms other techniques in quality, but full ColBERT on ANN has higher latency [3]. For large scale applications, users seem to overquery on ANN with single vector representations, and rerank them with ColBERT token vectors [4]. However, there's also ongoing work/research on reducing the no. of embeddings in ColBERT, like PLAID which replaces a bunch of vectors with their centroids [5]. ...
I hear you! And I don't want to add complexity only because we have some body of work in this PR. Thanks for raising the concern Jim, it led me to some interesting reading. ... My current thinking is that this is a rapidly evolving field, and it's early to lean one way or another. Adding this support unlocks experimentation. We might add different, scalable, ANN algos going forward, and our flat storage format should work with most of them. Meanwhile, there's research on different ways to run late interaction with multiple but fewer vectors. This change will help users experiment with what works at their scale, for their cost/performance/quality requirements. I'm happy to change my perspective, and would like to hear more opinions. One reason to not add this would be if it makes the single vector setup hard to evolve. I'd like to understand if (and how) this is happening, and think on how we can address those concerns. 1: https://docs.datastax.com/en/ragstack/examples/colbert.html |
Amen! This ends up being so domain specific. Multi-embeddings become key when you deal with domain voids in the LLMs used to create the embeddings. That's most big corpuses. So at least being able to experiment would get you far more feedback. I would be ok with writing some tests if that helps. |
I believe we should carefully consider the approach to adding multi-vector support through an aggregate function. From the outset, we assume that multi-vectors should be scored together, which is an important principle. Moreover, the default aggregate function proposed in the PR relies on brute force, which is not practical for any indexing setup. My concern is that this proposal doesn’t truly add support for independent multi-vectors. Instead, it introduces a block of vectors that must be scored together, which feels like a workaround rather than a comprehensive solution. This approach doesn’t address the key challenges of implementing true multi-vector support in the codec. The root issue is that the current KNN codec assumes the number of vectors is bounded by a single integer, a limitation that needs to be addressed first. Removing this constraint is a complex task but essential for properly supporting multi-vectors. Once that foundation is in place, adding support for setups like ColBERT should become relatively straightforward. Finally, while the max-sim function proposed in this PR may work as a ranking function, it isn’t suitable for indexing any documents. A true solution should allow for independent multi-vectors to be queried and scored flexibly without these constraints. |
That's a valid concern. I've been thinking about a more comprehensive multi-vector solution. Sharing some raw thoughts below, would love to get feedback. We support a default aggregation value of Once this is in place, we can add support for "dependent" multi-vector values like ColBERT. They'll take an aggregation function. Each graph node will represent all vectors for a document and use aggregated similarity (like in this PR). This will let us experiment with full ANN on ColBERT style multi-vectors. |
|
...contd. from above – thoughts on supporting independent multi-vectors specified via The Our codec today has single unique sequentially increasing vector ordinal per doc, which we can store and fetch with the DirectMonotonicWriter. For multi-vectors, we need to handle multiple nodeIds mapping to a single document. I'm thinking of using "ordinals" and "sub-ordinals" to identify each vector value. 'Ordinal' is incremented when docId changes. 'Sub-ordinals' start at 0 for each new doc and are incremented for subsequent vector values in the doc. A nodeId in the graph, is a "long" with ordinals and sub-ordinals packed into MSB and LSB bits separately. For flat storage, we can continue to use the technique in this PR; i.e. have one DirectMonotonicWriter object for docIds indexed by "ordinals", and another that stores start offsets for each docId, again indexed by ordinals. The sub-ordinal bits help us seek to exact vector values from this metadata. int ordToDoc(long nodeId) {
// get int ordinal from most-significant 32 bits
// get docId for the ordinal from DirectMonotonicWriter
}
float[] vectorValue(int nodeId) {
// get int ordinal from most-significant 32 bits
// get "startOffset" for ordinal
// get subOrdinal from least-signifant 32 bits
// read vector value from startOffset + (subOrdinal * dimension * byteSize)
}
float[] getAllVectorValues(int nodeId) {
// get int ordinal from most-significant 32 bits
// get "startOffset" for ordinal
// get "endOffset" from offset value for ordinal + 1
// return values from [startOffset, endOffset)
}With this setup, we won't need parent-block join queries for multiple vector values. And we can use I'm skeptical if this'll give a visible performance boost. It should at least be similar to the block-join setup we have today, but hopefully more convenient to use. And it sets us up for "dependent" multi-vector values like ColBERT. We'll need to code this up to iron out any wrinkles. I can work on a draft PR if the idea makes sense. Note that this still doesn't allow >2B vector values. While the "long" nodeId can support it, our ANN impl. returns arrays containing all nodeIds is various places. I don't think java can support >2B array length. But we can address this limitation separately, perhaps with a different ANN algo for such high cardinality graphs. |
Your proposal to implement (sidenote: if you are doing max/average, you can do that during index time though, right?) I'm currently conducting A/B tests on three methods to retrieve and rank documents with multiple vectors:
The third approach is particularly promising for domain-specific applications, where standard aggregation methods may not suffice. For instance, embedding tags could be linked to user access controls, unlocking certain vectors at query time, or to specific n-grams, activating them based on query content. Incorporating a mechanism to override the default aggregation method would facilitate experimentation with these strategies. |
|
Thank you for sharing these use-cases @krickert !
Honestly, I think the existing parent-block join can achieve most use-cases for independent multi-vectors (the passage vector use case). But the approach above might make it easier to use? We also need it for dependent multi-vectors like ColBERT, though it's a separate question on whether ANN is even viable for ColBERT (v/s only for reranking). I'd like to know what issues or limitations do people face with the existing parent-child support for multiple vector values, so we can address them here. |
Not sure. But it is frustrating for me: we only calculate K chunks and not N documents. I want to return N documents all the time, and keep running K until N is reached. Since it runs K on the chunks, I'd rather it return all thee chunks that it can until it reaches N amount of documents. Then we can return the chunks that match which can be used by highlighting.
Indexing the child docs requires making more docs. We just care about the resulting embedding, so why not treat it like a tensor instead of an entire document? It's frustrating to always make a child doc for multiple vectors when I can just do a keyword-value style instead. Also, there's def some limitations with how you can use it with scoring and the query ends up looking like a mess. If we can simplify the query syntax that would help a lot. If you can get a unit test going for your PR, I'd be glad to expand on it and play with it a bit. |
|
This PR has not had activity in the past 2 weeks, labeling it as stale. If the PR is waiting for review, notify the [email protected] list. Thank you for your contribution! |
|
I pivoted to an approach that handles independent multi-vectors within flat storage, instead of requiring index time parent-block joins. Have raised a draft PR here – #14173 |
|
This PR has not had activity in the past 2 weeks, labeling it as stale. If the PR is waiting for review, notify the [email protected] list. Thank you for your contribution! |
|
I believe this proposal to add a new field for multi-vector support is facing significant challenges primarily because we aim to support HNSW-based search on it. However, if our goal were limited to enabling MAXSIM scoring between two multi-vectors using only exact search (i.e., without KNN indexing for faster retrieval), many of these challenges could be avoided. This exact search with MAXSIM scoring approach would still be valuable for re-scoring documents after an initial KNN query retrieves a small candidate set. If we want to represent multi-vectors as an array of vectors within a single document—rather than relying on a parent-child document structure—it might make sense to integrate this into the existing field type. However, such an approach should clearly demonstrate a performance advantage to justify the change. |
|
On a second thought, if user just diable indexing, implementing multi vector in the same field might be same as having a separate field without indexing capability. |
one of the challenges, and perhaps the main concern with this PR, has been practical scaling challenges of working with multi-vectors on the entire corpus. This would apply to exact search as well. A viable middle ground would be to over-collect documents via regular single-vector knn search in the first pass, and later rerank them using late-interaction multi-vectors in the second pass. This idea was shared by jimczi@ earlier in this thread. I'm working on a PR that combines his approach with a FunctionScore query for reranking. Should be able to raise it for review this week. |
|
This PR has not had activity in the past 2 weeks, labeling it as stale. If the PR is waiting for review, notify the [email protected] list. Thank you for your contribution! |
Adds support for multi-valued vectors to Lucene.
In addition to max-similarity aggregations like parent-block joins, this change supports ColBERT style distance functions that compute interaction across all query and document vector values. Documents can have a variable number of vector values, but to support distance function computations, we require all values to have the same dimension.
Addresses #12313 .
.
Approach
We define a new "Tensor" field that comprises multiple vector values, and a new
TensorSimilarityFunctionto compute distance across multiple vectors (uses SumMax() currently). Node ordinal is assigned to the tensor value, giving us one ordinal per document. All vector values of a tensor field are processed together during writing, reading and scoring. They are passed around as a packedfloat[]orbyte[]array with all vector values concatenated. Consumers (like theTensorSimilarityFunction) slice this array by dimension to get individual vector values.Tensors are stored using a new FlatVectorStorage that supports writing/reading variable length values per field (allowing us to have a different number of vectors per tensor). We reuse the existing HNSW readers and writers. Each graph node is a tensor and maps to a single document. I also added a new codec tensor format, to allow both tensors and vectors to coexist. I'm not yet sure how to integrate with the quantization changes (separate later date change) and didn't want to force everything into a single format. Tensors continue to work with
KnnVectorWriter/Readerand extend theFlatVectorWriter/Readerclasses.Finally, I named the field and format "Tensors" though technically these are only rank-2 tensors. The thought was that we might extend this field and format if we ever went for higher rank tensors support. I'm open to renaming based on community feedback.
.
Major Changes
The PR has a lot of files which is not practical to review. Here are the files with key changes. If we align on the approach, I'm happy to reraise separate split PRs with different changes.
Lucene99FlatTensorsWriterfor writing in the new flat tensor format - lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99FlatTensorsWriter.javaLucene99FlatTensorsReaderfor reading the flat tensor format - lucene/core/src/java/org/apache/lucene/codecs/lucene99/Lucene99FlatTensorsReader.javaHnswTensorFormatthat uses FlatTensorFormat to initialize the flat storage readers/writers underlying HNSW reader/writer..
Open Questions
vectorEncodingandvectorDimensionattributes in FieldInfo instead of a separate tensor encoding and dimension