Implementation of Random Rotation Matrix and Query Scaling Matrix with 1-bit Binary Quantization #2524

Vikasht34 · 2025-02-12T18:47:31Z

Description

This change enhances the One-Bit Scalar Quantization process in by refining the Below and Above Threshold Means computation and introducing Random Gaussian Rotation. These improvements optimize both quantization accuracy during indexing and query scaling during retrieval, ensuring better recall and efficiency in k-NN search.

Below and Above Threshold Means for Query Scaling
These means help in query-time scaling, enabling better similarity approximation.

During training, each vector dimension is analyzed, and a threshold (mean value) is computed.
Each dimension’s values are classified as:
- Below threshold: Values lower than the computed mean.
- Above threshold: Values higher than the computed mean.

Random Gaussian Rotation for Robustness
Applies a Gaussian orthonormal rotation matrix to input vectors before quantization.
Used only when the L2/L1 ratio exceeds 0.6, ensuring that rotation is applied selectively for improved data distribution.

Related Issues

Resolves #[Issue number to be closed when this PR is merged]

Check List

New functionality includes testing.
New functionality has been documented.
API changes companion pull request created.
Commits are signed per the DCO using --signoff.
Public documentation issue/PR created.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

…above and below thresholds for 1 bit Binary Quantization Signed-off-by: Vikasht34 <[email protected]>

Signed-off-by: Vikasht34 <[email protected]>

jmazanec15 · 2025-02-24T16:51:06Z

.../org/opensearch/knn/quantization/models/quantizationState/OneBitScalarQuantizationState.java

+     * Represents the mean of all values below the threshold for each dimension.
+     */
+    @Builder.Default
+    private float[] belowThresholdMeans = null;


Could you add examples for these, as was done for meanThresholds?

jmazanec15 · 2025-02-24T16:51:23Z

.../org/opensearch/knn/quantization/models/quantizationState/OneBitScalarQuantizationState.java

+    @Builder.Default
+    private float[] aboveThresholdMeans = null;
+    @Builder.Default
+    private double averageL2L1Ratio = 0.0;


Can you add a comment on this?

jmazanec15 · 2025-02-24T16:52:26Z

.../org/opensearch/knn/quantization/models/quantizationState/OneBitScalarQuantizationState.java

@@ -51,6 +76,20 @@ public void writeTo(StreamOutput out) throws IOException {
        out.writeVInt(Version.CURRENT.id); // Write the version
        quantizationParams.writeTo(out);
        out.writeFloatArray(meanThresholds);
+        out.writeOptionalArray(belowThresholdMeans != null ? new FloatArrayWrapper[] { new FloatArrayWrapper(belowThresholdMeans) } : null);


For BWC purposes, writeOptionalArray wont work in a mixed cluster state. We need to check the stream version.

jmazanec15 · 2025-02-24T16:53:12Z

.../org/opensearch/knn/quantization/models/quantizationState/OneBitScalarQuantizationState.java

+    @Builder.Default
+    private double averageL2L1Ratio = 0.0;
+    /**
+     * Rotation matrix used when L2/L1 ratio > 0.6


Can you add more detail? For instance, what are the properties of this matrix. Also, what will happen if L2/L1 ratio is < 0.6

jmazanec15 · 2025-02-24T16:54:09Z

.../org/opensearch/knn/quantization/models/quantizationState/OneBitScalarQuantizationState.java

@@ -63,6 +102,23 @@ public OneBitScalarQuantizationState(StreamInput in) throws IOException {
        int version = in.readVInt(); // Read the version
        this.quantizationParams = new ScalarQuantizationParams(in, version);
        this.meanThresholds = in.readFloatArray();
+        if (Version.fromId(version).onOrAfter(Version.V_3_0_0)) {
+            // Deserialize belowThresholdMeans using readOptionalArray
+            FloatArrayWrapper[] wrappedBelowThresholdMeans = in.readOptionalArray(FloatArrayWrapper::new, FloatArrayWrapper[]::new);


Why optional on this? Is there a case where its null?

Yes ..for Version less than 3.0.0

jmazanec15 · 2025-02-24T16:54:43Z

.../org/opensearch/knn/quantization/models/quantizationState/OneBitScalarQuantizationState.java

+            this.array = in.readFloatArray();
+        }
+
+        public float[] getArray() {


nit: lombok getter

jmazanec15 · 2025-02-24T16:58:38Z

src/main/java/org/opensearch/knn/quantization/quantizer/QuantizerHelper.java

+        float[] belowThresholdMeans = belowAboveMeans.getA();
+        float[] aboveThresholdMeans = belowAboveMeans.getB();
+
+        // Apply the same rotation to below and above threshold means if rotation was applied


In line 60, we apply the rotation on the meanThresholds. Then we calculate the belowAboveMeans with this rotated meanThresholds in line 65. So, shouldnt the returned thresholds already be rotated?

Instead, can we apply rotations to all thresholds at the end (including menathresholds)?

jmazanec15 · 2025-02-24T16:59:57Z

src/main/java/org/opensearch/knn/quantization/quantizer/QuantizerHelper.java

+     * @throws IOException              If an I/O error occurs while retrieving vector data.
+     * @throws IllegalArgumentException If any vector at the sampled indices is null.
+     */
+    private static Pair<float[], Double> calculateMeanAndL2L1Ratio(TrainingRequest<float[]> trainingRequest, int[] sampledIndices)


Instead of a pair, can we make this into some kind of object so that we can continue to extend with additional data?

jmazanec15 · 2025-02-24T17:02:01Z

src/main/java/org/opensearch/knn/quantization/quantizer/RandomGaussianRotation.java

+            }
+        }
+
+        // Step 2: Orthogonalize the matrix using the Gram-Schmidt process


Are there any libraries we can delegate this computation to?

jmazanec15 · 2025-02-24T17:04:36Z

.../org/opensearch/knn/quantization/models/quantizationState/OneBitScalarQuantizationState.java

@@ -63,6 +102,23 @@ public OneBitScalarQuantizationState(StreamInput in) throws IOException {
        int version = in.readVInt(); // Read the version


Is this going to be bwc? We need to check if the segment was created in a previous version as opposed to the node it is coming from.

jmazanec15 · 2025-02-26T15:11:20Z

@Vikasht34 For this, I think we probably can break up the random rotation and the computation of mean centers into 2 separate features. I think that they can be added independently.

Vikasht34 · 2025-02-26T16:28:13Z

@Vikasht34 For this, I think we probably can break up the random rotation and the computation of mean centers into 2 separate features. I think that they can be added independently.

IMO ..it should be part of whole ADC ...We have seen random rotation giving advantage of some data sets , It's beeter we get combo of ADC + Random Rotation.

Vikasht34 added 2 commits February 11, 2025 16:57

Implementation of Random Rotation Matrix with Query Scaling matix of …

72ce48f

…above and below thresholds for 1 bit Binary Quantization Signed-off-by: Vikasht34 <[email protected]>

Code Cleanup

68ac180

Signed-off-by: Vikasht34 <[email protected]>

Vikasht34 force-pushed the main branch from 13fcaba to 68ac180 Compare February 12, 2025 21:38

Vikasht34 added 2 commits February 12, 2025 13:39

Merge branch 'main' into main

01219ae

Signed-off-by: Vikasht34 <[email protected]>

Some Falky Test

b1afdef

Signed-off-by: Vikasht34 <[email protected]>

Vikasht34 force-pushed the main branch 3 times, most recently from f5e801f to dad6872 Compare February 13, 2025 22:42

Clean Code

4465832

Signed-off-by: Vikasht34 <[email protected]>

Vikasht34 force-pushed the main branch from dad6872 to 4465832 Compare February 13, 2025 22:55

Vikasht34 marked this pull request as ready for review February 14, 2025 00:12

Vikasht34 requested review from heemin32, navneet1v, VijayanB, vamshin, jmazanec15, naveentatikonda, junqiu-lei, martin-gaievski, ryanbogan, luyuncheng, shatejas and 0ctopus13prime as code owners February 14, 2025 00:12

Merge branch 'opensearch-project:main' into main

f1ff640

jmazanec15 reviewed Feb 24, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implementation of Random Rotation Matrix and Query Scaling Matrix with 1-bit Binary Quantization #2524

Implementation of Random Rotation Matrix and Query Scaling Matrix with 1-bit Binary Quantization #2524

Vikasht34 commented Feb 12, 2025

jmazanec15 Feb 24, 2025

jmazanec15 Feb 24, 2025

jmazanec15 Feb 24, 2025

jmazanec15 Feb 24, 2025

jmazanec15 Feb 24, 2025

Vikasht34 Feb 26, 2025

jmazanec15 Feb 24, 2025

jmazanec15 Feb 24, 2025

jmazanec15 Feb 24, 2025

jmazanec15 Feb 24, 2025

jmazanec15 Feb 24, 2025

jmazanec15 commented Feb 26, 2025

Vikasht34 commented Feb 26, 2025

		@@ -63,6 +102,23 @@ public OneBitScalarQuantizationState(StreamInput in) throws IOException {
		int version = in.readVInt(); // Read the version

Implementation of Random Rotation Matrix and Query Scaling Matrix with 1-bit Binary Quantization #2524

Are you sure you want to change the base?

Implementation of Random Rotation Matrix and Query Scaling Matrix with 1-bit Binary Quantization #2524

Conversation

Vikasht34 commented Feb 12, 2025

Description

Related Issues

Check List

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jmazanec15 commented Feb 26, 2025

Vikasht34 commented Feb 26, 2025