Non-uniform vector quantization

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,107 @@
+/*
+ * 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;
+
+    public NVQ(NVQuantization nvq) {
+        this.nvq = nvq;
+        scorer = new NVQScorer(this.nvq);
+    }
+
+    @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 new ScoreFunction.ExactScoreFunction() {
+            private final QuantizedVector scratch = NVQuantization.QuantizedVector.createEmpty(nvq.subvectorSizesAndOffsets, nvq.bitsPerDimension);
+
+            @Override
+            public float similarityTo(int node2) {
+                try {
+                    var reader = source.inlineReaderForNode(node2, FeatureId.NVQ_VECTORS);
+                    QuantizedVector.loadInto(reader, scratch);
+                } catch (IOException e) {
+                    throw new RuntimeException(e);
+                }
+                return function.similarityTo(scratch);
+            }
+        };
+    }
+}

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

@@ -262,10 +262,16 @@ 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,6 +324,8 @@ 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.");
             }

jvector-base/src/main/java/io/github/jbellis/jvector/graph/similarity/BuildScoreProvider.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.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/optimization/LossFunction.java

@@ -0,0 +1,137 @@
+/*
+ * 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.optimization;
+
+import java.util.Arrays;
+import java.util.stream.IntStream;
+
+/**
+ * Class that models a loss function that maps a multidimensional vector onto a real number.
+ */
+public abstract class LossFunction {
+    // The number of dimensions
+    final private int nDims;
+
+    // The box constraints that define the feasible set.
+    private float[] minBounds;
+    private float[] maxBounds;
+
+    /**
+     * Constructs a LossFunction acting on vectors of the specified number of dimensions.
+     * @param nDims the number of dimensions
+     */
+    public LossFunction(int nDims) {
+        if (nDims <= 0) {
+            throw new IllegalArgumentException("The standard deviation initSigma must be positive");
+        }
+        minBounds = new float[nDims];
+        maxBounds = new float[nDims];
+        for (int d = 0; d < nDims; d++) {
+            minBounds[d] = Float.NEGATIVE_INFINITY;
+            maxBounds[d] = Float.POSITIVE_INFINITY;
+        }
+
+        this.nDims = nDims;
+    }
+
+    /**
+     * Computes the loss function. It assumes that input is within the feasible set
+     * @param x the input vector
+     * @return the loss
+     */
+    public abstract float compute(float[] x);
+
+    /**
+     * Computes the loss function and projects the input in-place onto the feasible set
+     * @param x the input vector
+     * @return the loss
+     */
+    public float projectCompute(float[] x) {
+        project(x);
+        return compute(x);
+    }
+
+    /**
+     * Sets the minimum values of the box constraints.
+     * @param bounds the specified minimum bound
+     */
+    public void setMinBounds(float[] bounds) {
+        if (nDims != bounds.length) {
+            throw new IllegalArgumentException("The length of bounds should match the number of dimensions");
+        }
+        minBounds = bounds;
+    }
+
+    /**
+     * Gets the minimum values of the box constraints.
+     * @return the minimum bound
+     */
+    public float[] getMinBounds() {
+        return minBounds;
+    }
+
+    /**
+     * Sets the maximum values of the box constraints.
+     * @param bounds the specified maximum bound
+     */
+    public void setMaxBounds(float[] bounds) {
+        if (nDims != bounds.length) {
+            throw new IllegalArgumentException("The length of bounds should match the number of dimensions");
+        }
+        maxBounds = bounds;
+    }
+
+    /**
+     * Gets the maximum values of the box constraints.
+     * @return the maximum bound
+     */
+    public float[] getMaxBounds() {
+        return maxBounds;
+    }
+
+    /**
+     * Projects the input onto the feasible set. If in-place, the input array is modified;
+     * otherwise, a copy is created and then projected.
+     * @param x the input vector
+     * @param inPlace If true, the input array is modified; otherwise, a copy is created and then projected.
+     * @return the projected vector
+     */
+    public float[] project(float[] x, boolean inPlace) {
+        float[] copy;
+        if (inPlace) {
+            copy = x;
+        }
+        else {
+            copy = Arrays.copyOf(x, x.length);
+        }
+        for (int d = 0; d < nDims; d++) {
+            copy[d] = Math.min(Math.max(x[d], minBounds[d]), maxBounds[d]);
+        }
+        return copy;
+    }
+
+    /**
+     * Projects the input in-place onto the feasible set.
+     * @param x the input vector
+     */
+    public void project(float[] x) {
+        project(x, true);
+    }
+
+    public boolean minimumGoalAchieved(float lossValue) {
+        return false;
+    }
+}