Non-uniform vector quantization (#374)

Adds support for Non-uniform Vector Quantization (NVQ, pronounced as "new vec"). This new technique quantizes the values  in each vector with high accuracy by first applying a nonlinear transformation that is individually fit to each  vector. These nonlinearities are designed to be lightweight and have a negligible impact on distance computation  performance.

---------

Co-authored-by: Joel Knighton <[email protected]>
Co-authored-by: Ted Willke <[email protected]>

README.md

@@ -27,6 +27,7 @@ The graph is represented by an on-disk adjacency list per node, with additional
 
 The second pass can be performed with
 * Full resolution float32 vectors
+* NVQ, which uses a non-uniform technique to quantize vectors with high-accuracy
 
 [This two-pass design reduces memory usage and reduces latency while preserving accuracy](https://thenewstack.io/why-vector-size-matters/).  
 

UPGRADING.md

@@ -16,6 +16,12 @@
   `CompressedVectors` directly from `encodeAll()`.
 - `PQVectors::getProductQuantization` is removed; it duplicated `CompressedVectors::getCompressor` unnecessarily
 
+## New features
+- Support for Non-uniform Vector Quantization (NVQ, pronounced as "new vec"). This new technique quantizes the values
+  in each vector with high accuracy by first applying a nonlinear transformation that is individually fit to each
+  vector. These nonlinearities are designed to be lightweight and have a negligible impact on distance computation
+  performance.
+
 # Upgrading from 2.0.x to 3.0.x
 
 ## Critical API changes

jvector-base/src/main/java/io/github/jbellis/jvector/graph/GraphIndexBuilder.java

@@ -211,20 +211,20 @@ public static GraphIndexBuilder rescore(GraphIndexBuilder other, BuildScoreProvi
             var neighbors = other.graph.getNeighbors(i);
             var sf = newProvider.searchProviderFor(i).scoreFunction();
             var newNeighbors = new NodeArray(neighbors.size());
-            
+
             // Copy edges, compute new scores
             for (var it = neighbors.iterator(); it.hasNext(); ) {
                 int neighbor = it.nextInt();
                 // since we're using a different score provider, use insertSorted instead of addInOrder
                 newNeighbors.insertSorted(neighbor, sf.similarityTo(neighbor));
             }
-            
+
             newBuilder.graph.addNode(i, newNeighbors);
         }
 
         // Set the entry node
         newBuilder.graph.updateEntryNode(other.graph.entry());
-        
+
         return newBuilder;
     }
 

jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FeatureId.java

@@ -31,7 +31,8 @@
  */
 public enum FeatureId {
     INLINE_VECTORS(InlineVectors::load),
-    FUSED_ADC(FusedADC::load);
+    FUSED_ADC(FusedADC::load),
+    NVQ_VECTORS(NVQ::load);
 
     public static final Set<FeatureId> ALL = Collections.unmodifiableSet(EnumSet.allOf(FeatureId.class));
 

jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/FusedADC.java

@@ -19,9 +19,9 @@
 import io.github.jbellis.jvector.disk.RandomAccessReader;
 import io.github.jbellis.jvector.graph.GraphIndex;
 import io.github.jbellis.jvector.graph.similarity.ScoreFunction;
-import io.github.jbellis.jvector.pq.FusedADCPQDecoder;
-import io.github.jbellis.jvector.pq.PQVectors;
-import io.github.jbellis.jvector.pq.ProductQuantization;
+import io.github.jbellis.jvector.quantization.FusedADCPQDecoder;
+import io.github.jbellis.jvector.quantization.PQVectors;
+import io.github.jbellis.jvector.quantization.ProductQuantization;
 import io.github.jbellis.jvector.util.ExplicitThreadLocal;
 import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
 import io.github.jbellis.jvector.vector.VectorizationProvider;

jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/NVQ.java

@@ -0,0 +1,104 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.github.jbellis.jvector.graph.disk;
+
+import io.github.jbellis.jvector.disk.RandomAccessReader;
+import io.github.jbellis.jvector.graph.similarity.ScoreFunction;
+import io.github.jbellis.jvector.quantization.NVQScorer;
+import io.github.jbellis.jvector.quantization.NVQuantization;
+import io.github.jbellis.jvector.quantization.NVQuantization.QuantizedVector;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+
+import java.io.DataOutput;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+
+/**
+ * Implements the storage of NuVeQ vectors in an on-disk graph index.  These can be used for reranking.
+ */
+public class NVQ implements Feature {
+    private final NVQuantization nvq;
+    private final NVQScorer scorer;
+    private final ThreadLocal<QuantizedVector> reusableQuantizedVector;
+
+    public NVQ(NVQuantization nvq) {
+        this.nvq = nvq;
+        scorer = new NVQScorer(this.nvq);
+        reusableQuantizedVector = ThreadLocal.withInitial(() -> NVQuantization.QuantizedVector.createEmpty(nvq.subvectorSizesAndOffsets, nvq.bitsPerDimension));
+    }
+
+    @Override
+    public FeatureId id() {
+        return FeatureId.NVQ_VECTORS;
+    }
+
+    @Override
+    public int headerSize() {
+        return nvq.compressorSize();
+    }
+
+    @Override
+    public int inlineSize() { return nvq.compressedVectorSize();}
+
+    public int dimension() {
+        return nvq.globalMean.length();
+    }
+
+    static NVQ load(CommonHeader header, RandomAccessReader reader) {
+        try {
+            return new NVQ(NVQuantization.load(reader));
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
+    @Override
+    public void writeHeader(DataOutput out) throws IOException {
+        nvq.write(out, OnDiskGraphIndex.CURRENT_VERSION);
+    }
+
+    @Override
+    public void writeInline(DataOutput out, Feature.State state_) throws IOException {
+        var state = (NVQ.State) state_;
+        state.vector.write(out);
+    }
+
+    public static class State implements Feature.State {
+        public final QuantizedVector vector;
+
+        public State(QuantizedVector vector) {
+            this.vector = vector;
+        }
+    }
+
+    ScoreFunction.ExactScoreFunction rerankerFor(VectorFloat<?> queryVector,
+                                                 VectorSimilarityFunction vsf,
+                                                 FeatureSource source) {
+        var function = scorer.scoreFunctionFor(queryVector, vsf);
+
+        return node2 -> {
+            try {
+                var reader = source.inlineReaderForNode(node2, FeatureId.NVQ_VECTORS);
+                QuantizedVector.loadInto(reader, reusableQuantizedVector.get());
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+            return function.similarityTo(reusableQuantizedVector.get());
+        };
+    }
+}

jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndex.java

@@ -262,10 +262,13 @@ public void close() throws IOException {
 
         @Override
         public ScoreFunction.ExactScoreFunction rerankerFor(VectorFloat<?> queryVector, VectorSimilarityFunction vsf) {
-            if (!features.containsKey(FeatureId.INLINE_VECTORS)) {
-                throw new UnsupportedOperationException("No inline vectors in this graph");
+            if (features.containsKey(FeatureId.INLINE_VECTORS)) {
+                return RandomAccessVectorValues.super.rerankerFor(queryVector, vsf);
+            } else if (features.containsKey(FeatureId.NVQ_VECTORS)) {
+                return ((NVQ) features.get(FeatureId.NVQ_VECTORS)).rerankerFor(queryVector, vsf, this);
+            } else {
+                throw new UnsupportedOperationException("No reranker available for this graph");
             }
-            return RandomAccessVectorValues.super.rerankerFor(queryVector, vsf);
         }
 
         @Override

jvector-base/src/main/java/io/github/jbellis/jvector/graph/disk/OnDiskGraphIndexWriter.java

@@ -324,8 +324,10 @@ public OnDiskGraphIndexWriter build() throws IOException {
             int dimension;
             if (features.containsKey(FeatureId.INLINE_VECTORS)) {
                 dimension = ((InlineVectors) features.get(FeatureId.INLINE_VECTORS)).dimension();
+            } else if (features.containsKey(FeatureId.NVQ_VECTORS)) {
+                dimension = ((NVQ) features.get(FeatureId.NVQ_VECTORS)).dimension();
             } else {
-                throw new IllegalArgumentException("Inline vectors must be provided.");
+                throw new IllegalArgumentException("Inline or NVQ vectors must be provided.");
             }
 
             if (ordinalMapper == null) {

jvector-base/src/main/java/io/github/jbellis/jvector/graph/similarity/BuildScoreProvider.java

@@ -17,8 +17,8 @@
 package io.github.jbellis.jvector.graph.similarity;
 
 import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
-import io.github.jbellis.jvector.pq.BQVectors;
-import io.github.jbellis.jvector.pq.PQVectors;
+import io.github.jbellis.jvector.quantization.BQVectors;
+import io.github.jbellis.jvector.quantization.PQVectors;
 import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
 import io.github.jbellis.jvector.vector.VectorUtil;
 import io.github.jbellis.jvector.vector.VectorizationProvider;

jvector-base/src/main/java/io/github/jbellis/jvector/graph/similarity/CachingVectorValues.java

@@ -17,7 +17,7 @@
 package io.github.jbellis.jvector.graph.similarity;
 
 import io.github.jbellis.jvector.graph.RandomAccessVectorValues;
-import io.github.jbellis.jvector.pq.PQVectors;
+import io.github.jbellis.jvector.quantization.PQVectors;
 import io.github.jbellis.jvector.vector.types.VectorFloat;
 import org.agrona.collections.Int2ObjectHashMap;
 

jvector-base/src/main/java/io/github/jbellis/jvector/pq/BQVectors.java → jvector-base/src/main/java/io/github/jbellis/jvector/quantization/BQVectors.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package io.github.jbellis.jvector.pq;
+package io.github.jbellis.jvector.quantization;
 
 import io.github.jbellis.jvector.disk.RandomAccessReader;
 import io.github.jbellis.jvector.graph.similarity.ScoreFunction;

jvector-base/src/main/java/io/github/jbellis/jvector/pq/BinaryQuantization.java → jvector-base/src/main/java/io/github/jbellis/jvector/quantization/BinaryQuantization.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package io.github.jbellis.jvector.pq;
+package io.github.jbellis.jvector.quantization;
 
 import io.github.jbellis.jvector.disk.RandomAccessReader;
 import io.github.jbellis.jvector.graph.RandomAccessVectorValues;

jvector-base/src/main/java/io/github/jbellis/jvector/pq/CompressedVectors.java → jvector-base/src/main/java/io/github/jbellis/jvector/quantization/CompressedVectors.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package io.github.jbellis.jvector.pq;
+package io.github.jbellis.jvector.quantization;
 
 import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex;
 import io.github.jbellis.jvector.graph.similarity.ScoreFunction;

jvector-base/src/main/java/io/github/jbellis/jvector/pq/FusedADCPQDecoder.java → jvector-base/src/main/java/io/github/jbellis/jvector/quantization/FusedADCPQDecoder.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package io.github.jbellis.jvector.pq;
+package io.github.jbellis.jvector.quantization;
 
 import io.github.jbellis.jvector.graph.disk.FusedADC;
 import io.github.jbellis.jvector.graph.similarity.ScoreFunction;

jvector-base/src/main/java/io/github/jbellis/jvector/pq/ImmutableBQVectors.java → jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ImmutableBQVectors.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package io.github.jbellis.jvector.pq;
+package io.github.jbellis.jvector.quantization;
 
 public class ImmutableBQVectors extends BQVectors {
     public ImmutableBQVectors(BinaryQuantization bq, long[][] compressedVectors) {

jvector-base/src/main/java/io/github/jbellis/jvector/pq/ImmutablePQVectors.java → jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ImmutablePQVectors.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package io.github.jbellis.jvector.pq;
+package io.github.jbellis.jvector.quantization;
 
 import io.github.jbellis.jvector.vector.types.ByteSequence;
 

jvector-base/src/main/java/io/github/jbellis/jvector/pq/KMeansPlusPlusClusterer.java → jvector-base/src/main/java/io/github/jbellis/jvector/quantization/KMeansPlusPlusClusterer.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package io.github.jbellis.jvector.pq;
+package io.github.jbellis.jvector.quantization;
 
 import io.github.jbellis.jvector.vector.Matrix;
 import io.github.jbellis.jvector.vector.VectorUtil;

jvector-base/src/main/java/io/github/jbellis/jvector/pq/MutableBQVectors.java → jvector-base/src/main/java/io/github/jbellis/jvector/quantization/MutableBQVectors.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package io.github.jbellis.jvector.pq;
+package io.github.jbellis.jvector.quantization;
 
 public class MutableBQVectors extends BQVectors implements MutableCompressedVectors<long[]> {
     private static final int INITIAL_CAPACITY = 1024;

jvector-base/src/main/java/io/github/jbellis/jvector/pq/MutableCompressedVectors.java → jvector-base/src/main/java/io/github/jbellis/jvector/quantization/MutableCompressedVectors.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package io.github.jbellis.jvector.pq;
+package io.github.jbellis.jvector.quantization;
 
 public interface  MutableCompressedVectors<T> extends CompressedVectors {
     /**

jvector-base/src/main/java/io/github/jbellis/jvector/pq/MutablePQVectors.java → jvector-base/src/main/java/io/github/jbellis/jvector/quantization/MutablePQVectors.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package io.github.jbellis.jvector.pq;
+package io.github.jbellis.jvector.quantization;
 
 import io.github.jbellis.jvector.vector.VectorizationProvider;
 import io.github.jbellis.jvector.vector.types.ByteSequence;