Add MUVERA processors for multi-vector search #3164
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,221 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.knn.processor.muvera; | ||
|
|
||
| import java.util.ArrayList; | ||
| import java.util.List; | ||
| import java.util.Random; | ||
|
|
||
| /** | ||
| * MUVERA (Multi-Vector Retrieval via Fixed Dimensional Encodings) | ||
| * | ||
| * Converts variable-length multi-vector embeddings (e.g. ColBERT/ColPali token embeddings) | ||
| * into fixed-size single vectors that approximate MaxSim scoring via dot product. | ||
| * | ||
| * The output FDE dimension = rReps * 2^kSim * dimProj. | ||
| * | ||
| * Document processing: normalizes cluster centers by count, fills empty clusters via Hamming nearest neighbor. | ||
| * Query processing: raw sum (no normalization), no empty cluster filling. | ||
| * | ||
| * The random seed must be identical between index time and query time to produce compatible encodings. | ||
| */ | ||
| public class MuveraEncoder { | ||
| private final int dim; | ||
| private final int kSim; | ||
| private final int dimProj; | ||
| private final int rReps; | ||
| private final int numPartitions; | ||
| private final double[][] simhashVectors; | ||
| private final double[][][] dimReductionProjections; | ||
|
|
||
| /** | ||
| * Maximum allowed FDE dimension. Matches the k-NN engine max dimension limit (16,000) | ||
| * to ensure the FDE vector can be indexed in any supported engine. | ||
| */ | ||
| static final int MAX_FDE_DIMENSION = 16_000; | ||
|
Contributor
There was a problem hiding this comment. is there chances people want to configure the max dimension? this can be a configuration setting in the processor ?
Author
There was a problem hiding this comment. @mingshl The 16,000 cap mirrors the k-NN engine's max vector dimension, going above it would fail at indexing regardless. Users already control the actual FDE size through the MUVERA parameters (dim, k_sim, dim_proj, r_reps). Adding a configurable max would let users set a tighter bound, but it feels like an edge case since the engine enforces the real ceiling. Happy to add it as an optional max_fde_dimension processor config if you think it's worth it though. |
||
|
|
||
| public MuveraEncoder(int dim, int kSim, int dimProj, int rReps, long seed) { | ||
| if (dim <= 0) { | ||
| throw new IllegalArgumentException("dim must be positive, got: " + dim); | ||
| } | ||
| if (kSim < 0) { | ||
| throw new IllegalArgumentException("k_sim must be non-negative, got: " + kSim); | ||
| } | ||
| if (dimProj <= 0) { | ||
| throw new IllegalArgumentException("dim_proj must be positive, got: " + dimProj); | ||
| } | ||
| if (rReps <= 0) { | ||
| throw new IllegalArgumentException("r_reps must be positive, got: " + rReps); | ||
| } | ||
|
|
||
| // Validate that the resulting FDE dimension does not exceed the k-NN engine limit. | ||
| // FDE dimension = rReps * 2^kSim * dimProj | ||
| long fdeDimension = (long) rReps * (1L << kSim) * dimProj; | ||
| if (fdeDimension > MAX_FDE_DIMENSION) { | ||
| throw new IllegalArgumentException( | ||
| "MUVERA parameters produce an FDE dimension of [" | ||
| + fdeDimension | ||
| + "] which exceeds the maximum allowed dimension of [" | ||
| + MAX_FDE_DIMENSION | ||
| + "]. Reduce r_reps, k_sim, or dim_proj. (r_reps=" | ||
| + rReps | ||
| + " * 2^k_sim=" | ||
| + (1L << kSim) | ||
| + " * dim_proj=" | ||
| + dimProj | ||
| + ")" | ||
| ); | ||
| } | ||
|
|
||
| this.dim = dim; | ||
| this.kSim = kSim; | ||
| this.dimProj = dimProj; | ||
| this.rReps = rReps; | ||
| this.numPartitions = 1 << kSim; | ||
|
praveenMprasad marked this conversation as resolved.
|
||
|
|
||
| Random rng = new Random(seed); | ||
|
|
||
| // SimHash hyperplanes: rReps sets of kSim hyperplanes, each dim-dimensional | ||
| simhashVectors = new double[rReps][kSim * dim]; | ||
| for (int r = 0; r < rReps; r++) { | ||
| for (int i = 0; i < kSim * dim; i++) { | ||
| simhashVectors[r][i] = rng.nextGaussian(); | ||
| } | ||
| } | ||
|
|
||
| // Random projection matrices with entries in {-1, +1} | ||
| dimReductionProjections = new double[rReps][dim][dimProj]; | ||
| for (int r = 0; r < rReps; r++) { | ||
| for (int i = 0; i < dim; i++) { | ||
| for (int j = 0; j < dimProj; j++) { | ||
| dimReductionProjections[r][i][j] = rng.nextBoolean() ? 1.0 : -1.0; | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Returns the output FDE dimension: rReps * 2^kSim * dimProj. | ||
| */ | ||
| public int getEmbeddingSize() { | ||
| return rReps * numPartitions * dimProj; | ||
| } | ||
|
|
||
| private int getClusterId(double[] vector, int repIndex) { | ||
| int clusterId = 0; | ||
| for (int k = 0; k < kSim; k++) { | ||
| double dot = 0; | ||
| int offset = k * dim; | ||
| for (int d = 0; d < this.dim; d++) { | ||
| dot += vector[d] * simhashVectors[repIndex][offset + d]; | ||
| } | ||
| if (dot > 0) { | ||
| clusterId |= (1 << k); | ||
| } | ||
| } | ||
| return clusterId; | ||
| } | ||
|
|
||
| private int hammingDistance(int a, int b) { | ||
| return Integer.bitCount(a ^ b); | ||
| } | ||
|
|
||
| /** | ||
| * Core FDE encoding. | ||
| * | ||
| * @param vectors multi-vector input, shape [numVectors][dim] | ||
| * @param fillEmpty if true, fill empty clusters from Hamming-nearest non-empty cluster (document mode) | ||
| * @param normalize if true, normalize cluster centers by count (document mode) | ||
| * @return FDE vector of length getEmbeddingSize() | ||
| */ | ||
| public float[] process(double[][] vectors, boolean fillEmpty, boolean normalize) { | ||
| int nVectors = vectors.length; | ||
| float[] output = new float[getEmbeddingSize()]; | ||
| int outOffset = 0; | ||
| double scale = 1.0 / Math.sqrt(dimProj); | ||
|
|
||
| for (int r = 0; r < rReps; r++) { | ||
| double[][] centers = new double[numPartitions][dim]; | ||
|
Member
There was a problem hiding this comment. This seems too heavy. Can we optimize this by reusing the arrays somehow? Allocation for every loop can get expensive.
Author
There was a problem hiding this comment. Pre-allocated centers, counts, and clusterVecIndices arrays outside the loop and reset them per repetition. No more per-loop allocation. |
||
| int[] counts = new int[numPartitions]; | ||
| List<List<Integer>> clusterVecIndices = new ArrayList<>(numPartitions); | ||
| for (int i = 0; i < numPartitions; i++) { | ||
| clusterVecIndices.add(new ArrayList<>()); | ||
| } | ||
|
|
||
| // Assign vectors to clusters and accumulate | ||
| for (int v = 0; v < nVectors; v++) { | ||
| int cid = getClusterId(vectors[v], r); | ||
| counts[cid]++; | ||
| clusterVecIndices.get(cid).add(v); | ||
| for (int d = 0; d < dim; d++) { | ||
| centers[cid][d] += vectors[v][d]; | ||
| } | ||
| } | ||
|
|
||
| // Normalize by count (document mode) | ||
| if (normalize) { | ||
| for (int c = 0; c < numPartitions; c++) { | ||
| if (counts[c] > 0) { | ||
| for (int d = 0; d < dim; d++) { | ||
| centers[c][d] /= counts[c]; | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| // Fill empty clusters from Hamming-nearest non-empty cluster (document mode) | ||
| if (fillEmpty) { | ||
| for (int c = 0; c < numPartitions; c++) { | ||
| if (counts[c] == 0) { | ||
| int nearest = -1; | ||
| int minDist = Integer.MAX_VALUE; | ||
| for (int other = 0; other < numPartitions; other++) { | ||
| if (counts[other] > 0) { | ||
| int dist = hammingDistance(c, other); | ||
| if (dist < minDist) { | ||
| minDist = dist; | ||
| nearest = other; | ||
| } | ||
| } | ||
| } | ||
| if (nearest >= 0) { | ||
| int vecIdx = clusterVecIndices.get(nearest).get(0); | ||
| System.arraycopy(vectors[vecIdx], 0, centers[c], 0, dim); | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| // Random projection: centers[c] (dim) x projection (dim x dimProj) -> projected (dimProj) | ||
| for (int c = 0; c < numPartitions; c++) { | ||
| for (int j = 0; j < dimProj; j++) { | ||
| double val = 0; | ||
| for (int d = 0; d < dim; d++) { | ||
| val += centers[c][d] * dimReductionProjections[r][d][j]; | ||
| } | ||
| output[outOffset++] = (float) (scale * val); | ||
| } | ||
| } | ||
| } | ||
| return output; | ||
| } | ||
|
|
||
| /** | ||
| * Encode document multi-vectors into a single FDE vector. | ||
| * Normalizes cluster centers and fills empty clusters. | ||
| */ | ||
| public float[] processDocument(double[][] vectors) { | ||
| return process(vectors, true, true); | ||
| } | ||
|
|
||
| /** | ||
| * Encode query multi-vectors into a single FDE vector. | ||
| * No normalization, no empty cluster filling. | ||
| */ | ||
| public float[] processQuery(double[][] vectors) { | ||
| return process(vectors, false, false); | ||
| } | ||
| } | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Any particular reason to use double instead of float? The return type signatures are all floats.
The operations for similarity match can benefit with SIMD/Panama calculations and the memory consumption can also reduce by half.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Good point on the memory savings. The internal arrays use double for numerical precision during accumulation — with ~1000 vectors per document, the cluster center sums can get large and float accumulation could introduce meaningful rounding errors. The output is cast to float at the end since that's what the knn_vector field stores.
That said, for the simhash and projection matrices (which are generated once at init), float would be fine. I'll convert those to float in a follow-up to reduce the encoder's memory footprint. The cluster center accumulation should stay double to avoid precision loss during summation.
Open to converting everything to float if you think the precision tradeoff is acceptable — happy to benchmark the quality impact.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think this makes sense. Let's reduce the footprint as much as we can
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Converted simhash and projection matrices to float. Cluster center accumulation stays double to avoid precision loss during summation. See latest commit.