Fixing quantization interval initialization for optimized sq (#14374)

* Fixing quantization interval initialization for optimized sq

* adding changes along with original binary quantization change

* adjusting test

Author: benwtrent

lucene/CHANGES.txt

@@ -86,6 +86,11 @@ New Features
   These queries allow for the vector search entry points to be initialized via a `seed` query. This follows
   the research provided via https://arxiv.org/abs/2307.16779. (Sean MacAvaney, Ben Trent).
 
+* GITHUB#14078, GITHUB#14374: Adds two new binary quantized vector codecs `Lucene102HnswBinaryQuantizedVectorsFormat`
+  and `Lucene102BinaryQuantizedVectorsFormat`. These new formats allow a ~32x from raw float32 vectors while
+  achieving high recall (with oversampling and reranking vectors) for typical modern knn vector search tasks.
+  (Mayya Sharipova, Ben Trent, Tom Veasey, John Wagster, Chris Hegarty)
+
 * GITHUB#13974,GITHUB#14276: Introducing DocValuesMultiRangeQuery.SortedSetStabbingBuilder into sandbox.
   (Mikhail Khludnev)
 

lucene/core/src/java/org/apache/lucene/util/quantization/OptimizedScalarQuantizer.java

@@ -149,9 +149,9 @@ public QuantizationResult[] multiScalarQuantize(
       // Linearly scale the interval to the standard deviation of the vector, ensuring we are within
       // the min/max bounds
       intervalScratch[0] =
-          (float) clamp((MINIMUM_MSE_GRID[bits[i] - 1][0] + vecMean) * vecStd, min, max);
+          (float) clamp(MINIMUM_MSE_GRID[bits[i] - 1][0] * vecStd + vecMean, min, max);
       intervalScratch[1] =
-          (float) clamp((MINIMUM_MSE_GRID[bits[i] - 1][1] + vecMean) * vecStd, min, max);
+          (float) clamp(MINIMUM_MSE_GRID[bits[i] - 1][1] * vecStd + vecMean, min, max);
       optimizeIntervals(intervalScratch, vector, norm2, points);
       float nSteps = ((1 << bits[i]) - 1);
       float a = intervalScratch[0];
@@ -214,10 +214,8 @@ public QuantizationResult scalarQuantize(
     double vecStd = Math.sqrt(vecVar);
     // Linearly scale the interval to the standard deviation of the vector, ensuring we are within
     // the min/max bounds
-    intervalScratch[0] =
-        (float) clamp((MINIMUM_MSE_GRID[bits - 1][0] + vecMean) * vecStd, min, max);
-    intervalScratch[1] =
-        (float) clamp((MINIMUM_MSE_GRID[bits - 1][1] + vecMean) * vecStd, min, max);
+    intervalScratch[0] = (float) clamp(MINIMUM_MSE_GRID[bits - 1][0] * vecStd + vecMean, min, max);
+    intervalScratch[1] = (float) clamp(MINIMUM_MSE_GRID[bits - 1][1] * vecStd + vecMean, min, max);
     optimizeIntervals(intervalScratch, vector, norm2, points);
     float nSteps = ((1 << bits) - 1);
     // Now we have the optimized intervals, quantize the vector

lucene/core/src/test/org/apache/lucene/util/quantization/TestOptimizedScalarQuantizer.java

@@ -27,6 +27,64 @@
 public class TestOptimizedScalarQuantizer extends LuceneTestCase {
   static final byte[] ALL_BITS = new byte[] {1, 2, 3, 4, 5, 6, 7, 8};
 
+  static float[] deQuantize(byte[] quantized, byte bits, float[] interval, float[] centroid) {
+    float[] dequantized = new float[quantized.length];
+    float a = interval[0];
+    float b = interval[1];
+    int nSteps = (1 << bits) - 1;
+    double step = (b - a) / nSteps;
+    for (int h = 0; h < quantized.length; h++) {
+      double xi = (double) (quantized[h] & 0xFF) * step + a;
+      dequantized[h] = (float) (xi + centroid[h]);
+    }
+    return dequantized;
+  }
+
+  public void testQuantizationQuality() {
+    int dims = 16;
+    int numVectors = 32;
+    float[][] vectors = new float[numVectors][];
+    float[] centroid = new float[dims];
+    for (int i = 0; i < numVectors; ++i) {
+      vectors[i] = new float[dims];
+      for (int j = 0; j < dims; ++j) {
+        vectors[i][j] = randomFloat();
+        centroid[j] += vectors[i][j];
+      }
+    }
+    for (int j = 0; j < dims; ++j) {
+      centroid[j] /= numVectors;
+    }
+    // similarity doesn't matter for this test
+    OptimizedScalarQuantizer osq =
+        new OptimizedScalarQuantizer(VectorSimilarityFunction.DOT_PRODUCT);
+    float[] scratch = new float[dims];
+    for (byte bit : ALL_BITS) {
+      float eps = (1f / (float) (1 << (bit)));
+      byte[] destination = new byte[dims];
+      for (int i = 0; i < numVectors; ++i) {
+        System.arraycopy(vectors[i], 0, scratch, 0, dims);
+        OptimizedScalarQuantizer.QuantizationResult result =
+            osq.scalarQuantize(scratch, destination, bit, centroid);
+        assertValidResults(result);
+        assertValidQuantizedRange(destination, bit);
+
+        float[] dequantized =
+            deQuantize(
+                destination,
+                bit,
+                new float[] {result.lowerInterval(), result.upperInterval()},
+                centroid);
+        float mae = 0;
+        for (int k = 0; k < dims; ++k) {
+          mae += Math.abs(dequantized[k] - vectors[i][k]);
+        }
+        mae /= dims;
+        assertTrue("bits: " + bit + " mae: " + mae + " > eps: " + eps, mae <= eps);
+      }
+    }
+  }
+
   public void testAbusiveEdgeCases() {
     // large zero array
     for (VectorSimilarityFunction vectorSimilarityFunction : VectorSimilarityFunction.values()) {