[FEA] Support multivectors in all indexes
Nobody has claimed this yet.
- Dominant language
- Cuda
- Stars
- 854
- Forks
- 236
- Avg merge
- 3d 3h
- Merged PRs (30d)
- 62
Description
This feature is growing in popularity of LLMs where many "term" embeddings may come from the same "document" but we want to make sure the set of resulting "document" indices returned for a neighborhood query are unique for each query vector. At the moment, it seems this is being implemented at a higher layer in many systems as a de-duplication or filtering step. We should try and support such a feature as generally as possible. The ultimate goal being to apply a constraint, potentially in the k-selection step, if it can be done efficiently.
Thinking throught his a bit further, we have had an idea to support returning document ids as the k-closest neighbors, along with the corresponding distances. This method involves using a pre-filtering function and a global atomic, maintaining the following 2 arrays in global memory:
- An array of n_index_vectors, which maps each vector (value) to its doc id (key).
- An array of size n_documents, which maintains the ongoing sum of distances
The idea here is that we go through potential closest neighbors and each time the pre-filtering predicate function is invoked, we atomically add the distance to the corresponding doc id. This is a fairly naive approach, which leads to increased random writes and atomics, but I think we could find ways to reduce these with priors such as using a filtering threshold and ignoring distances beyond that threshold.
We should continue to find better ways to implement this, but this approach could yield reasonable results initially as something to build on.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
No files, tests, or entry points are named. Start by mapping the index implementations, neighborhood-query paths, filtering callbacks, and k-selection logic; done means all indexes can return unique document IDs for each query vector with corresponding distances, while preserving a viable efficiency path.
Written by the indexing model from the issue text.
Assessment
- Domain
- machine-learning, performance, search
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 20/100