Add hadamard rotation to vector fields - #16092
Conversation
A lightweight wrapping FlatVectorsFormat that applies a randomized Hadamard rotation to vectors before handing them to a delegate format (e.g. Lucene104ScalarQuantizedVectorsFormat), and rotates query vectors at search time. Because the rotation is orthogonal, dot product, cosine similarity, and Euclidean distance are all preserved, so the delegate's similarity math is unchanged. The rotation redistributes variance across dimensions, which makes OSQ's assumption of Gaussian components hold on datasets whose raw components are skewed or uniform (image pixels, histograms, non-transformer embeddings). Motivation and approach come from the discussion in Apache Lucene PR apache#15903 (TurboQuant) and Elastic's April 2026 blog on BBQ preconditioning, which measured 41-74% recall improvements on GIST / SIFT / Fashion-MNIST at ~2-4% query overhead. Implementation: - HadamardRotation: immutable, thread-safe, O(d log d) via Fast Walsh-Hadamard Transform with random sign flips and a Fisher-Yates permutation. Supports non-power-of-2 dimensions through a block decomposition (e.g. 768 = 512 + 256). Provides both forward and inverse rotations. - RotationPreconditionedVectorsFormat: public FlatVectorsFormat with a no-arg constructor (required for SPI) that defaults to wrapping Lucene104ScalarQuantizedVectorsFormat, plus constructors that take a custom delegate and seed. - RotationPreconditionedVectorsWriter: intercepts addValue to rotate each vector before forwarding to the delegate's field writer. Byte vectors pass through unchanged. - RotationPreconditionedVectorsReader: rotates float query vectors before scoring, inverse-rotates stored vectors in getFloatVectorValues for rescore/CheckIndex callers, and exposes the raw rotated values via getMergeInstance() so that the delegate's merge runs entirely in rotated space (preserving byte-copy merge where the delegate supports it). - Global deterministic rotation per (dim, seed): the same rotation across all segments enables byte-copy merges in the underlying format. SPI wiring: - META-INF/services/org.apache.lucene.codecs.KnnVectorsFormat registers the new format alongside FaissKnnVectorsFormat. - module-info.java exports the package and lists the format under the KnnVectorsFormat 'provides' clause. Tests: - TestHadamardRotation: 11 unit tests covering orthogonality (L2 norm preservation, dot product preservation, Euclidean distance preservation), determinism (same seed -> same rotation), non-identity (different seeds differ), block decomposition correctness, and the spreading property on concentrated inputs. - TestRotationPreconditionedVectorsFormat: extends BaseKnnVectorsFormatTestCase and runs the full suite of KNN-format correctness tests (merge, sort, delete, multi-field, mismatched fields, random exceptions, etc.). Eight tests in the base suite assert bit-exact round-trip equality of indexed vectors; those are overridden with explanatory comments because rotate+inverse_rotate introduces ~1e-7 floating-point drift. Search correctness is unaffected because the rotation is orthogonal. All sandbox tests pass (336 tests, 64 pre-existing skips).
Refactors the rotation preconditioner out of a sandbox FlatVectorsFormat wrapper and into the existing Lucene104ScalarQuantizedVectorsFormat (OSQ) as an opt-in feature controlled by a rotationSeed constructor argument. Rotation is now a first-class capability of OSQ: when a non-zero seed is supplied, every incoming vector is Hadamard-rotated before centroid computation and quantization, and every query is rotated the same way at search time. Because the rotation is orthogonal, all similarity functions (dot product, cosine, Euclidean) are preserved, but per-coordinate distributions become much more Gaussian — which makes OSQ's initialization assumption hold on datasets with skewed or uniform components (image pixels, histograms, non-transformer embeddings). Motivation comes from Apache Lucene PR apache#15903 (TurboQuant) discussion and Elastic's April 2026 blog on BBQ preconditioning, which measured 41-74% recall improvements on GIST / SIFT / Fashion-MNIST at ~2-4% query overhead. Changes: - Move HadamardRotation + its test from lucene/sandbox to lucene/core/src/java/org/apache/lucene/util/quantization/ so it lives next to the existing OptimizedScalarQuantizer. - Lucene104ScalarQuantizedVectorsFormat: add a rotationSeed constructor parameter (default ROTATION_DISABLED = 0 preserves existing behaviour). Bump the on-disk format to VERSION_PRECONDITIONED (1). Old segments (version 0) are still readable; their seed is implicitly 0. - Lucene104HnswScalarQuantizedVectorsFormat: add a matching ctor overload so the HNSW wrapper can enable preconditioning. - Lucene104ScalarQuantizedVectorsWriter: constructor takes the seed; FieldWriter.addValue rotates the incoming vector up front so all downstream OSQ math (centroid accumulation, raw storage, quantization) runs in the rotated basis. writeMeta persists the seed. - Lucene104ScalarQuantizedVectorsReader: FieldEntry now carries rotationSeed; readField reads it when the version supports it. getRandomVectorScorer(String, float[]) rotates the query before scoring. getFloatVectorValues wraps the raw delegate with an InverseRotatedFloatVectorValues so external callers (rerank, CheckIndex, FieldExistsQuery, etc.) see the original vectors they indexed. getMergeInstance() returns a lightweight MergeReader that skips the inverse rotation — the downstream merge then operates entirely in rotated space, preserving consistency across segments. - Remove the sandbox/rotation package and its tests; revert the sandbox module-info and SPI service registration. - Update OSQ and HNSW toString() tests to include rotationSeed. Add TestLucene104ScalarQuantizedVectorsFormatPreconditioning covering end-to-end search with rotation enabled, round-tripping vectors through rotate+inverseRotate via getFloatVectorValues, seed=0 equivalence to the default format, and toString observability. All existing OSQ flat/HNSW/backward-compat tests continue to pass. The 4 new preconditioning tests and the 11 HadamardRotation math tests pass.
Replaces the previous attempt that modified Lucene104 in place. Since
Lucene104 is a shipped codec with a frozen on-disk format, any layout
change belongs in a new codec family.
This commit:
- Restores Lucene104ScalarQuantizedVectorsFormat (and the matching
HNSW wrapper / writer / reader / tests) to their exact pre-patch
state. Anyone with a Lucene104 index can still read it byte-for-byte
the same as before.
- Introduces Lucene105ScalarQuantizedVectorsFormat + the HNSW wrapper
as a new codec family (package
org.apache.lucene.codecs.lucene105). The codec-name headers and
internal NAME strings all use 'Lucene105' so the new layout can be
distinguished at read time. File extensions (.veq, .vemq) are the
same because the codec-name header in each file is what
disambiguates.
- Adds rotation preconditioning natively to Lucene105 as an opt-in
feature controlled by a rotationSeed constructor argument:
* Default / sentinel value ROTATION_DISABLED (0) keeps the format
layout shape matching Lucene104 aside from one extra long per
field in metadata.
* A non-zero seed enables Hadamard rotation at index and query
time. The rotation is orthogonal so dot product / cosine /
Euclidean distances are preserved end to end; what changes is
the per-coordinate distribution of the stored vectors, which
becomes much more Gaussian. This helps OSQ initialization on
datasets with skewed / uniform components (image pixels,
histograms, non-transformer embeddings).
* The seed is persisted in per-field metadata. Reader rotates
queries in getRandomVectorScorer, inverse-rotates stored values
in getFloatVectorValues (so external rerank / CheckIndex /
FieldExistsQuery callers see the original vectors), and exposes
an unrotated view via getMergeInstance so merges stay in the
rotated basis end to end.
- Clones the scorer (Lucene105ScalarQuantizedVectorScorer) and the two
Off-heap value classes (OffHeapScalarQuantizedVectorValues,
OffHeapScalarQuantizedFloatVectorValues) into the new package so the
Lucene104 package-private members don't have to be made public for
the Lucene105 codec to use them. HadamardRotation lives once, in
lucene/core/src/java/org/apache/lucene/util/quantization/, because
it's a utility rather than a codec.
- Registers Lucene105ScalarQuantizedVectorsFormat and
Lucene105HnswScalarQuantizedVectorsFormat via SPI (META-INF
services file and the module-info 'provides' clause), and exports
the new package.
- Adds TestLucene105ScalarQuantizedVectorsFormatPreconditioning with
four targeted tests covering end-to-end preconditioned search,
vector round-trip through rotate+inverseRotate via
getFloatVectorValues, seed=0 equivalence to the default format, and
toString observability.
All existing Lucene104 OSQ tests, Lucene105 preconditioning tests,
HadamardRotation math tests, backward-compat Lucene99 OSQ tests, and
sandbox tests pass.
Usage:
// Pick the old codec for backward compat.
new Lucene104ScalarQuantizedVectorsFormat();
// Pick the new codec with no rotation (default).
new Lucene105ScalarQuantizedVectorsFormat();
// Pick the new codec with rotation preconditioning enabled.
new Lucene105ScalarQuantizedVectorsFormat(
ScalarEncoding.UNSIGNED_BYTE, 0x5eedCafeBabeL);
…ring merge
During force_merge, the HNSW graph builder needs to compare documents
against each other. For 1-bit and 2-bit (asymmetric) encodings, this
requires building a temporary 4-bit "query" representation of each
document by reading back its float vector and re-quantizing it against
the segment centroid.
The bug: getRandomVectorScorerSupplierForMerge() called
getFloatVectorValues(), which inverse-rotates the stored vectors back
to original space (designed for external callers). These un-rotated
vectors were then quantized against the centroid, which lives in
rotated space (computed from rotated vectors during indexing). The
centering step (vector[i] - centroid[i]) mixed original-space vectors
with a rotated-space centroid, producing meaningless 4-bit
representations. The HNSW graph built from these scores was
essentially random, dropping recall from 0.695 to 0.050 (1-bit) and
0.816 to 0.055 (2-bit).
Only 1-bit and 2-bit are affected. 4-bit, 7-bit, and 8-bit use
symmetric scoring which reads already-quantized bytes directly — no
float vectors involved, no rotation mismatch possible.
The fix: use rawVectorsReader.getFloatVectorValues() to read the
stored rotated vectors directly, matching the rotated centroid.
Indices built with the buggy code have corrupted HNSW graphs for
1-bit and 2-bit segments and need reindexing or re-merging.
Benchmark results (Cohere v3 1024d, 500K docs, DOT_PRODUCT):
bits baseline before-fix after-fix
1 0.695 0.050 0.729
2 0.816 0.055 0.841
4 0.918 0.937 0.937
7 0.970 0.982 0.982
8 0.976 0.985 0.985
This reverts commit 305839b.
shubhamvishu
added
the
skip-changelog
|
I also ran the same luceneutil with 4K dimensional vectors and I see even higher impact to recall(~6-7% improvement) overall net-net with slight slowness in indexing-rate(~5%) due to rotation overhead. With internal Amazon 4K vectors embeddingsSetup:
|
| FieldInfo info = fieldInfos.fieldInfo(field); | ||
| return info != null | ||
| && "true" | ||
| .equals(info.getAttribute(Lucene104ScalarQuantizedVectorsFormat.ROTATION_ENABLED_KEY)); |
There was a problem hiding this comment.
I haven't ever seen attributes used to enable codec features like this, at least it doesn't appear to be common practice in the core codecs. More typically this would tick VERSION_CURRENT and something (probably just the rotation seed?) would be encoded as part of field metadata. I don't have a good sense as to why maybe @mikemccand or @benwtrent has a firmer and better reasoned opinion.
There was a problem hiding this comment.
Yeah I also reasonate here. The initial version (1st commit) tried to retain it in the sandox module -> then moved it into the main core Codec but it was writing a random seed to the metadata of the field and I had to bump the codec to 105. I honestly didn't wanted to bump the Codec for this given usecase so I switched to use a constant seed and assuming rotation is enabled for all vector fields by default(which obviously simplified all of this but takes away the capability from user to configure it on per field basis). So eventually I biased towards sharing the rotationEnabled flag (configured per vector field) to query time via FieldInfos and avoiding Codec bump since this way we were not breaking the backward compatibility and also how simple it was in nature. I'm open to ideas whichever we would want to choose or if there is a better way to share this info(rotation enabled/disabled) hopefully avoiding the codec bump.
| if (isRotationEnabled(field) && target != null) { | ||
| HadamardRotation rotation = rotationFor(field, fi.dimension); | ||
| float[] rotated = new float[target.length]; | ||
| rotation.rotate(target, rotated); |
There was a problem hiding this comment.
IIUC this rotation operation is probably fairly expensive -- at least as expensive as quantization but possibly even more expensive. In Lucene's segment structure this operation will be repeated for every segment searched. In your tests you probably ran with everything merged down to a single segment but I'm interested in what this costs in a more typical multi-segment setup. A microbenchmark for the rotation or a full luceneutil run would be helpful here.
Depending on the cost we might want to figure out a way to reuse this computation across segments somehow which probably requires upstream integration with the knn query classes.
There was a problem hiding this comment.
rotate is a O(d log d) operation here. I ran the luceneutil without forceMerge too but I don't have the Cohere results handy anymore(I can pull those up again). Though I have the results with 4K dim embeddings in multi segment index. Sharing those below for you reference(will share the Cohere one also soon). I didn't seem to regress the latency as such. We can do more JMH benchmarking or whatever could give us more confidence but so far it appears to be a cheap cost(Idk if the quantization overhead was significant in past).
There was a problem hiding this comment.
We cannot do a rotation per segment, that is a non starter.
This would significantly impact latency in the most common scenario, which is many segments.
Instead I suggest queries need to have an "additional phase" that looks to see if any of the KnnFormats can apply a "globalPrecondition" step or something and then apply it once for the query.
There was a problem hiding this comment.
To see the significance of the performance impact, you will need many vectors spread over many segments (which is very common, I mean 10s of millions spread over 30-50+ segments).
There was a problem hiding this comment.
@benwtrent I see, yeah I think moving in upstream and doing it once makes sense to me. Maybe just moving it to KnnQuery as Trevor mentioned below?
To see the significance of the performance impact, you will need many vectors spread over many segments
That could be the case yes. Sharing below the cohere v3 run I put with multi segment and it seems to have not much impact but larger segments of 30M might move the needle.
Result : Cohere V3 without forceMerge
Baseline :
Results:
NOTE: nDoc = 500000 for all runs; skipping column
NOTE: searchType = KNN for all runs; skipping column
NOTE: topK = 100 for all runs; skipping column
NOTE: fanout = 100 for all runs; skipping column
NOTE: resultSimilarity = N/A for all runs; skipping column
NOTE: decay = N/A for all runs; skipping column
NOTE: resultCount = 100.000 for all runs; skipping column
NOTE: maxConn = 64 for all runs; skipping column
NOTE: beamWidth = 250 for all runs; skipping column
NOTE: force_merge(s) = 0.00 for all runs; skipping column
NOTE: filterStrategy = null for all runs; skipping column
NOTE: filterSelectivity = N/A for all runs; skipping column
NOTE: overSample = 1.000 for all runs; skipping column
NOTE: bp-reorder = false for all runs; skipping column
NOTE: indexType = HNSW for all runs; skipping column
NOTE: rerank = no for all runs; skipping column
recall latency(ms) netCPU avgCpuCount quantized visited index(s) index_docs/s num_segments index_size(MB) vec_disk(MB) vec_RAM(MB)
0.695 1.916 5.285 2.758 1 bits 54581 94.09 5314.12 5 2092.52 2020.836 67.711
0.817 1.887 6.052 3.208 2 bits 50314 95.73 5223.08 5 2147.02 2081.871 128.746
0.921 2.107 7.275 3.453 4 bits 50877 97.97 5103.60 6 2266.48 2204.895 251.770
0.976 3.300 12.843 3.892 7 bits 61241 101.92 4906.00 9 2506.59 2449.036 495.911
0.983 3.280 12.810 3.905 8 bits 61080 99.65 5017.51 9 2506.62 2449.036 495.911
Candidate:
Results:
NOTE: nDoc = 500000 for all runs; skipping column
NOTE: searchType = KNN for all runs; skipping column
NOTE: topK = 100 for all runs; skipping column
NOTE: fanout = 100 for all runs; skipping column
NOTE: resultSimilarity = N/A for all runs; skipping column
NOTE: decay = N/A for all runs; skipping column
NOTE: resultCount = 100.000 for all runs; skipping column
NOTE: maxConn = 64 for all runs; skipping column
NOTE: beamWidth = 250 for all runs; skipping column
NOTE: force_merge(s) = 0.00 for all runs; skipping column
NOTE: filterStrategy = null for all runs; skipping column
NOTE: filterSelectivity = N/A for all runs; skipping column
NOTE: overSample = 1.000 for all runs; skipping column
NOTE: bp-reorder = false for all runs; skipping column
NOTE: indexType = HNSW for all runs; skipping column
NOTE: rerank = no for all runs; skipping column
recall latency(ms) netCPU avgCpuCount quantized visited index(s) index_docs/s num_segments index_size(MB) vec_disk(MB) vec_RAM(MB)
0.729 1.845 5.056 2.740 1 bits 51755 96.87 5161.40 5 2090.45 2020.836 67.711
0.843 1.870 5.904 3.157 2 bits 49123 98.38 5082.18 5 2146.54 2081.871 128.746
0.940 2.089 7.143 3.419 4 bits 50584 95.99 5209.04 6 2266.39 2204.895 251.770
0.988 3.255 12.872 3.955 7 bits 60998 100.92 4954.52 9 2506.59 2449.036 495.911
0.993 3.244 12.887 3.973 8 bits 60955 100.85 4958.01 9 2506.60 2449.036 495.911
There was a problem hiding this comment.
I ran the luceneutil without forceMerge too but I don't have the Cohere results handy anymore(I can pull those up again).
Note that knnPerfTest.py in luceneutil now enables autologger so all output from every run should be in your $LOGS_DIR. I added autologger for exactly this situation!! I got so annoyed losing track of my runs...
There was a problem hiding this comment.
rotateis aO(d log d)operation here.
Wait, this is not so costly? It's like log(d) dot products for each segment, when we do many dot products per vector per (large-ish) segment. It's fixed cost per segment, and amortizes to zero as the index gets bigger.
To see the significance of the performance impact, you will need many vectors spread over many segments (which is very common, I mean 10s of millions spread over 30-50+ segments).
Actually, this should be the best case scenario -- the larger the index/segment, the more dot products you'll need to crawl its graph, reducing this small added fixed per-segment cost of log(d) dot products?
I don't really understand the fuss. The added cost (amortizes to zero) seems reasonable if we really do see such recall lift.
@shubhamvishu you could create a tiny tool in luceneutil to just do a one-time rotation of .vec file? Then we all could quickly evaluate impact on our own vector corpora, even without this PR, for starters.
But anyway +1 to make a single global rotation matrix for each vector field, not because I'm worried about the performance impact of per-segment query vector rotation, but for a different reason: because I think this would then unlock data blind quantization as long as vectors are L2 unit-sphere normalized. Then (once your vectors are truly isotropic), optimal quantization is solely a function of int D (dimension) and not the actual vectors, which would be amazing:
- Faster indexing since we wouldn't need the extra pass through the vectors to get histogram of values across all dimensions / centroid for the segment (to tune the quantizer, which is not data blind today)
- Merging could then bulk-copy vectors in quantized form, since their quantized encoding will not change
- We could fully drop full precision vectors (!!) since they are no longer needed for tuning quantizer. Of course user should still have option to keep them, e.g. maybe they want to re-rank.
But let's keep this PR focused on pre-conditioning! Data blind quantizer, options to keep full precision vectors or not, can come later...
There was a problem hiding this comment.
I addressed this piece and other comments(of having single rotation matrix per dim). We now rotate the incoming query only once in the query rewrite(KnnFloatVectorQuery#rewrite) itself (no more per segment overhead) and same with rotation matrices.
you could create a tiny tool in luceneutil to just do a one-time rotation of .vec file?
@mikemccand similar to having rotation in luceneutil, now in the latest revision of this PR the rotation has a new specific knn vector format(takes another knn vector format as input) which completely decouples rotation feature.
Luceneutil with Amazon 4K vectors embeddings (forceMerge=False)NOTE : Run 1 and 2 are on separate 4K embedding dataset(500K) so sharing both Run 1Baseline :Candidate:Run 2Baseline :Candidate:cc - @mccullocht |
| private final FieldInfos fieldInfos; | ||
|
|
||
| /** Lazily built Hadamard rotations, keyed by field name. */ | ||
| private final Map<String, HadamardRotation> rotations = new ConcurrentHashMap<>(); |
There was a problem hiding this comment.
I don't know about this. This isn't cheap. We have "fieldnum_segmentssize(HadamardRotationObject)",
Seems to me the rotation matrix should just be stored off heap :/ Or the rotation is by dimension, not by field name (idk why we need a new random matrix for this for two fields that have the same dimension).
There was a problem hiding this comment.
idk why we need a new random matrix for this for two fields that have the same dimension
It was because the random seed is taking the field name into picture but I agree to your point we could reuse the same matrix across vectors of same dimension actually (likely there is no/much benefit of having random seeds other than avoiding the possibility of choosing a bad seed for all vector fields but this simplification overshadows that possibility without completely discarding). I'll try to stick to a single rotation matrix for a dimension only. Do you think there is enough value of moving it off heap after having 1 rotation or it'll be an overkill?
There was a problem hiding this comment.
I made it a single global matrix (per dimension) and not maintaining it per segment now. Thanks!
|
Its a nice idea. I think we should strive to have a general "precondition vectors" interface. I am sure on the idea of having it integrated via field infos...I need to do some thinking here. Two big issues besides the API that are bothering me:
I don't know of another API in Lucene that has lazy state cached that is global over segments...this would be a fairly new thing here. Maybe we can "hack it" and add something to the KnnFormat reader interface, e.g. "globalPreconditioning" or something that queries can iterate and utilize...ugh, but then queries don't have the very nice API of just "search", and need to do this other step of "precondition"... we don't want to make things super complex for ALL other vector queries that Lucene has or that others wrote :( this is a tough one. |
|
Maybe this information should appear as part of FieldInfo? That would solve the "global" aspect of configuration as it would be uniform across all segments, and rotation could be handled in IndexChain and KnnQuery. |
|
@benwtrent @mccullocht Currently we are creating the rotation seed from the field name and caching that for each segment reader so this would be calculated once and reused for a field but I agree with Ben likely there could be benefits of keeping it off heap (or) I'm thinking we could even drop the field from the seed so that its only driven by the dimension (like 1 seed per unique dimensions in the vectors fields indexed?). That way we don't need to have this per segment (just one global object for a specific dimension)? Thoughts?
Right, I like the approach to do the rotation upfront into the KnnQuery using the FieldInfos setting. That way it would be global + less intrusive. |
| dimension, Lucene104ScalarQuantizedVectorsFormat.rotationSeed(f))); | ||
| } | ||
|
|
||
| private boolean isRotationEnabled(String field) { |
There was a problem hiding this comment.
Instead of adding rotation code to a specific format, can we create a new vector format that simply rotates the incoming vector in all functions, and passes it on to a delegate?
This would cleanly separate rotation functionality + make it re-usable across any arbitrary vector format.
There was a problem hiding this comment.
This is a really cool approach @kaivalnp! Then any KnnVectorsFormat out in the wild could quickly/easily insert preconditioning. Maybe look at existing delegating formats (PerFieldXXX)?
There was a problem hiding this comment.
+1 I too love the idea @kaivalnp. I like the part of any knn vector format should be able to integrate rotation without extra wiring or duplication. I added a RotationAwareKnnVectorsFormat which pretty much does this now (rotated the vectors when indexing and provides rotated/unrotated views of it at read time accordingly). I confirmed with a luceneutil the recall improvement still holds confirming it works as expected(similar to earlier approach).
InverseRotatedFloatVectorValues did not override scorer() or rescorer(), causing UnsupportedOperationException when the exact-search fallback or rescore path is triggered with rotation enabled. Since the query target arrives pre-rotated (from KnnFloatVectorQuery.rewrite()) and the delegate holds rotated vectors, both methods delegate directly to the underlying rotated-space FloatVectorValues for correct scoring.
… wrapper Adds a KnnVectorsFormat wrapper that rotates vectors at index time and provides getMergeInstance() to keep merge in rotated space. Removes rotation logic from Lucene104 codec internals. Query rotation is handled once globally in KnnFloatVectorQuery.rewrite() via FieldInfo attribute.
| && fi.getVectorDimension() == target.length | ||
| && "true".equals(fi.getAttribute(RotationAwareKnnVectorsFormat.ROTATION_ENABLED_KEY))) { | ||
| float[] rotated = new float[target.length]; | ||
| HadamardRotation.forDimension(fi.getVectorDimension()).rotate(target, rotated); |
There was a problem hiding this comment.
I'm not a fan of this pattern -- the responsibility of processing the query vector should lie with the format.
This leads to a possibility of vector queries producing garbage results if they do not consume this attribute correctly (for example, FloatVectorSimilarityQuery was missed in this PR).
Perhaps one way to catch such cases is to randomly wrap the KNN format being tested with a rotating one in BaseKnnVectorsFormatTestCase (or if we can accept slower tests: explicitly test everything with and without rotation).
That said, readers / writers in KnnVectorsFormat are segment-specific and not suitable to process rotation here, perhaps we need a global pre-processing step in the format itself?
I haven't thought too deeply about what the API would look like though.
There was a problem hiding this comment.
perhaps we need a global pre-processing step in the format itself?
This is difficult, every segment could have a different format version. There isn't really an idea of a "global format for all the segments".
I think it is possible to simply provide hints that preprocessing has occurred. For example, first iterate all segments, preprocessing the query and this provides some metadata on the query indicating that its been appropriately pre-processed. Then when its passed during the query, that metadata can be checked for. If its there, huzzah! Processed! Otherwise, go the slow path and preprocess.
I am not 100% sure this type of API will work either :/
There was a problem hiding this comment.
@kaivalnp @benwtrent I agree with the concerns here and the challenge of such complicated API. I'm thinking if for starter we should keep the rotation per segment and keep this change restricted to rotation like a recall vs latency tradeoff. We could follow with more ideas on how we could clawback the latency. One idea would be to try and see if SIMD could help if some parts of rotation logic or we could think of the sort of API as mentioned which could do this. But doing this in this PR I feel would make this more complicated(or errorprone since we don't have a clear vision for such an API yet). Another good thing of doing it per segment would be that way the whole rotation logic would be retained in the KNN vector format and not spilling into Query implementations etc so we won't be unintentionally breaking any current use cases but providing user an way to improve recall at the cost of some latency overhead for starters and later optimize (keeping this change focussed on capability addition).
I also ran JMH benchmarks(below) on Graviton 3 to see how many dot products(byte vectors) a single rotation operation is equivalent to in some scenarios (it would change with machine vs machine or different quantization etc scenarios). When compared to byte vectors dot product (using Panama based vectorization), a single rotation operation is roughly equivalent to ~50-60 dot products max and when using optimized native kernels it is roughly equivalent to ~200-350 dot products max. This extra rotation overhead might be (or maybe not) within the range we could swallow but looking for suggestions if this overhead seems larger than the potential overhead we initially thought or less than that? (at least for some specific use cases [dim, segment size] the overhead might be less though we could easily think of a situation where the overhead could be large like with 1000 visited nodes(dot products) visited per segment 300 could be <= 20-30% overhead roughly but with large segments it becomes less dominant and even with Panama when compared to other faster ways).
Graviton 3
Benchmark (size) Mode Cnt Score Error Units
HadamardRotationBenchmark.inverseRotate 128 thrpt 15 1.493 ± 0.024 ops/us
HadamardRotationBenchmark.inverseRotate 256 thrpt 15 0.708 ± 0.013 ops/us
HadamardRotationBenchmark.inverseRotate 702 thrpt 15 0.257 ± 0.002 ops/us
HadamardRotationBenchmark.inverseRotate 1024 thrpt 15 0.160 ± 0.001 ops/us
HadamardRotationBenchmark.inverseRotate 4096 thrpt 15 0.035 ± 0.001 ops/us
HadamardRotationBenchmark.rotate 128 thrpt 15 1.470 ± 0.004 ops/us
HadamardRotationBenchmark.rotate 256 thrpt 15 0.692 ± 0.010 ops/us
HadamardRotationBenchmark.rotate 702 thrpt 15 0.246 ± 0.001 ops/us
HadamardRotationBenchmark.rotate 1024 thrpt 15 0.155 ± 0.001 ops/us
HadamardRotationBenchmark.rotate 4096 thrpt 15 0.034 ± 0.001 ops/us
VectorUtilBenchmark.binaryDotProductScalar 128 thrpt 15 17.796 ± 0.139 ops/us
VectorUtilBenchmark.binaryDotProductScalar 256 thrpt 15 9.014 ± 0.032 ops/us
VectorUtilBenchmark.binaryDotProductScalar 702 thrpt 15 3.362 ± 0.017 ops/us
VectorUtilBenchmark.binaryDotProductScalar 1024 thrpt 15 2.285 ± 0.021 ops/us
VectorUtilBenchmark.binaryDotProductScalar 4096 thrpt 15 0.566 ± 0.004 ops/us
VectorUtilBenchmark.binaryDotProductVector 128 thrpt 15 60.224 ± 0.010 ops/us
VectorUtilBenchmark.binaryDotProductVector 256 thrpt 15 31.283 ± 0.004 ops/us
VectorUtilBenchmark.binaryDotProductVector 702 thrpt 15 11.609 ± 0.005 ops/us
VectorUtilBenchmark.binaryDotProductVector 1024 thrpt 15 8.039 ± 0.001 ops/us
VectorUtilBenchmark.binaryDotProductVector 4096 thrpt 15 2.024 ± 0.001 ops/us
VectorUtilBenchmark.dot8sNative 128 thrpt 15 135.797 ± 2.069 ops/us
VectorUtilBenchmark.dot8sNative 256 thrpt 15 90.647 ± 3.913 ops/us
VectorUtilBenchmark.dot8sNative 702 thrpt 15 49.021 ± 3.444 ops/us
VectorUtilBenchmark.dot8sNative 1024 thrpt 15 42.404 ± 0.713 ops/us
VectorUtilBenchmark.dot8sNative 4096 thrpt 15 12.222 ± 0.227 ops/us
| Dim | rotate (ops/μs) | binaryDotProductScalar (ops/μs) | binaryDotProductVector (ops/μs) | dot8sNative (ops/μs) |
|---|---|---|---|---|
| 128 | 1.47 | 17.8 (12x) | 60.2 (41x) | 135.8 (92x) |
| 256 | 0.69 | 9.0 (13x) | 31.3 (45x) | 90.6 (131x) |
| 702 | 0.25 | 3.4 (14x) | 11.6 (47x) | 49.0 (199x) |
| 1024 | 0.155 | 2.3 (15x) | 8.0 (52x) | 42.4 (274x) |
| 4096 | 0.034 | 0.57 (17x) | 2.0 (60x) | 12.2 (359x) |
There was a problem hiding this comment.
I'm thinking if for starter we should keep the rotation per segment and keep this change restricted to rotation like a recall vs latency tradeoff.
+1 for PnP ("progress not perfection"). New features should rarely be blocked because of performance, when they are opt-in (not the default) -- that can always be iteratively improved. Still, it's always good to debate approaches / APIs. We should anyways mark the new APIs @lucene.experimental so we are free to iterate them even across Lucene releases.
We could follow with more ideas on how we could clawback the latency.
I'm still confused -- why is latency a problem? It amortizes to 0% as the index gets bigger and bigger? (Because it'd be a fixed per-segment cost, like the cost of seeking the terms dict to locate postings for that segment). Also, "matrix times vector" is such an insanely well optimized function these days... surely we could (eventually) SIMD-ify it? If I read the JMH results right, it's ~6.5 microseconds per rotate for 1K vector, say times 30 segments, 200 microseconds (0.2 msec). That's fine for starters -- it's opt-in -- optimize later!
Another good thing of doing it per segment would be that way the whole rotation logic would be retained in the KNN vector format and not spilling into Query implementations etc so we won't be unintentionally breaking any current use cases but providing user an way to improve recall at the cost of some latency overhead for starters and later optimize (keeping this change focussed on capability addition).
+1 for that simplicity.
Perhaps one way to catch such cases is to randomly wrap the KNN format being tested with a rotating one in BaseKnnVectorsFormatTestCase (or if we can accept slower tests: explicitly test everything with and without rotation).
Big +1 for that -- random testing is powerful, given enough time all bugs are shallow! (The monkey Shakespeare typewriter thing, remixed with Linus' Law heh)
There was a problem hiding this comment.
It amortizes to 0% as the index gets bigger and bigger?
It does amortize, but the constant overhead is significant for single bit quantized, which I expect to be even faster than native int8 (which is almost 400x better throughput than the rotation?!?).
I agree, we should aim for opt in and simplicity. I am simply saying we shouldn't incur significant constant overheads when they are avoidable. I am fine with iterating on the opt-in API as it stands now. I have just been pushing a little bit to see if we can land on an API that doesn't have a significantly adverse segments with less than 1M vectors.
There was a problem hiding this comment.
I pushed a new revision moving the complete rotation specific logic into the new vector format and keeping it all simple and does not change anything with existing query impl or code paths. Enabling rotation simply means wrapping the any knn vector format with the new RotationAwareKnnVectorsFormat(<some knn format>) and added randomized tests testing it. Also confirmed the earlier rotation benchmark numbers remain the same (so functionally all same)
There was a problem hiding this comment.
@shubhamvishu thank you a bunch for iterating on the API and the discussion here :)
Without this, the SPI no-arg constructor of RotationAwareKnnVectorsFormat hardcoded a fixed delegate (Lucene104HnswScalarQuantizedVectorsFormat), which would silently mis-read indices written with any other delegate. Now: write time persists delegate.getName() as a per-field attribute (DELEGATE_FORMAT_KEY); read time resolves the delegate per attribute via KnnVectorsFormat.forName(...). The no-arg constructor is read-only and throws on writes, removing the hardcoded delegate dependency.
mikemccand
left a comment
mikemccand
left a comment
There was a problem hiding this comment.
Thank you @shubhamvishu -- it'd be awesome if we can really pull this off, globally, not per-segment, but I'm nervous, I think there are subtle merge bugs now (I added specific comments), and it's going against the Lucene grain.
I also asked Claude (the new Fable 5, "free" for two weeks if you have a paid Claude account already, then pay-to-play) and it looks like a great review too on quick glance: https://claude.ai/share/bc62ac58-9b15-454e-b52c-7b658babb062
| /** Wraps the given delegate format with rotation preconditioning. */ | ||
| public RotationAwareKnnVectorsFormat(KnnVectorsFormat delegate) { | ||
| super("RotationAwareKnnVectorsFormat"); | ||
| this.delegate = delegate; |
There was a problem hiding this comment.
Maybe this.delegate = Objects.requireNonNull(delegate)?
| */ | ||
| private final KnnVectorsFormat delegate; | ||
|
|
||
| /** No-arg constructor for SPI registration. Read-only — writing requires an explicit delegate. */ |
There was a problem hiding this comment.
Maybe add link to the next constructor? "For indexing, use {@link ...}"?
In fact, Lucene users should never directly use this constructor (since SPI creates all instances at read-time)? Maybe strengthen it to:
/** No-arg constructor for read-only mode, only for SPI registration. Do not use this for indexing. Use {@link...} instead.
| for (FieldInfo fi : state.fieldInfos) { | ||
| String name = fi.getAttribute(DELEGATE_FORMAT_KEY); | ||
| if (name != null) { | ||
| actualDelegate = |
There was a problem hiding this comment.
Hmm I think this won't work if I have two fields, one with this cool rotating wrapper around vector format A, and then another rotating wrapper around vector format B? Because we are globally going through all fields in this loop, breaking on the first delegate field we see. Maybe see how PerFieldXXXFormat does it? They have the same challenge... they may persist their own mapping file, or maybe they have a simple way to use FieldInfo.
| public class RotationAwareKnnVectorsFormat extends KnnVectorsFormat { | ||
|
|
||
| /** FieldInfo attribute key signaling rotation is enabled. */ | ||
| public static final String ROTATION_ENABLED_KEY = "Lucene104SQVecRotation"; |
There was a problem hiding this comment.
Maybe RotationAwareKnnVectorsFormat_Enabled or so? (So people seeing this FieldInfo attr will know it comes from this wrapper format?
There was a problem hiding this comment.
Actually maybe RotationAwareKnnVectors* -> RotatingKnnVectors* everywhere?
| * delegate at read time, so an index written with a non-default delegate is still readable when | ||
| * the SPI no-arg constructor is used. | ||
| */ | ||
| public static final String DELEGATE_FORMAT_KEY = "RotationAwareDelegateFormat"; |
There was a problem hiding this comment.
Maybe RotationAwareKnnVectorsFormat_DelegateFormat or so?
| } | ||
|
|
||
| /** Field-level writer that rotates each float vector before forwarding. */ | ||
| private static final class RotatingFieldWriter extends KnnFieldVectorsWriter<float[]> { |
There was a problem hiding this comment.
I asked Claude (Fable 5 -- free for two weeks, and then pay-to-play (by usage)!!) to review -- it had some great items: https://claude.ai/share/bc62ac58-9b15-454e-b52c-7b658babb062 some of which are overlapping my comments, some not!
| if (values == null) { | ||
| return null; | ||
| } | ||
| if (isRotationEnabled(field)) { |
There was a problem hiding this comment.
Actually, I'm confused, why wouldn't it be enabled? Shouldn't it always be enabled? I.e., only use this rotating wrapper format for fields you want to rotate?
Is it because of the byte[] case? Can we just throw exc if user tries to rotate those, for now? Add // TODO saying we may someday add rotation for that case too (no reason it shouldn't work?)?.
| public KnnVectorsReader getMergeInstance() throws IOException { | ||
| // Merge operates in rotated space — return the delegate directly so vectors are not | ||
| // inverse-rotated. The merge writer will see already-rotated vectors and write them as-is. | ||
| return delegateReader.getMergeInstance(); |
There was a problem hiding this comment.
Same danger as other merging place (hmm, why are there two?) -- we might end up mixing rotated / non-rotated vectors together?
| } | ||
|
|
||
| private float[] maybeRotateTarget(String field, float[] target) { | ||
| if (target != null && isRotationEnabled(field)) { |
There was a problem hiding this comment.
Hmm I don't like this null check / leniency -- caller should not pass null, that's weird.
Can we take maybe out of the name? Only call this method if you want to rotate.
|
|
||
| private boolean isRotationEnabled(String field) { | ||
| FieldInfo info = fieldInfos.fieldInfo(field); | ||
| return info != null && "true".equals(info.getAttribute(ROTATION_ENABLED_KEY)); |
There was a problem hiding this comment.
Throw exception if info is null?
|
This PR has not had activity in the past 2 weeks, labeling it as stale. If the PR is waiting for review, notify the dev@lucene.apache.org list. Thank you for your contribution! |
5 participants
Description
I worked with CC(Claude Code; did a great job in all phases from initial impl to testing) to have this PR which adds the Hadamard rotation(Fast Walsh Hadamard Transform) to vector fields(default false; configurable codec param; no codec bump required) inspired from the @xande 's TurboQuant PR (who works with me on Amazon Product Search) but a stripped down version just adding rotation to vectors in isolation. This address the 2nd item
Implement random rotation of vectors and queries.from Data-blind scalar quantization issue @mccullocht is working on.I'm opening this to gather community feedback, as it shows promising recall improvements. I'd like to see whether we want to incorporate this into Lucene, reuse some of these ideas, or discard the approach if there are concerns.
The shows upto ~5-7% recall improvement in luceneutil benchmarks with Cohere V3 and Amazon's internal 4K dim vector embeddings. Current approach rotates the incoming float vectors at insertion (so we index the vectors in rotated space in .vec file) and rest of the flow continues as is. It stores whether to do rotation for a vector field or not info in the
FieldInfos. At query time, it checks if the field has rotation enabled and rotates the query it true.TL;DR : Randomized orthogonal rotation (sign flips + permutation + FWHT) that Gaussianizes vector dimensions distributions to favor the scalar quantization(OSQ) accuracy while preserving distances.
Here's an ASCII diagram Claude generated explaining the Hadamard rotation steps :
Cohere V3
Baseline (
main) :Candidate (
main+ rotation i.e. this PR) :Recall