Create partial sums for PQ codebook for use during diversity checks (#511)

* Create partial sums for PQ codebook for use during diversity checks of graph building

Signed-off-by: Jake Luciani <[email protected]>

Author: tjake

benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/IndexConstructionWithRandomSetBenchmark.java

@@ -41,7 +41,7 @@
 @BenchmarkMode(Mode.AverageTime)
 @OutputTimeUnit(TimeUnit.MILLISECONDS)
 @State(Scope.Thread)
-@Fork(1)
+@Fork(value = 1, jvmArgsAppend = {"--add-modules=jdk.incubator.vector", "--enable-preview", "-Djvector.experimental.enable_native_vectorization=false"})
 @Warmup(iterations = 2)
 @Measurement(iterations = 3)
 @Threads(1)
@@ -59,7 +59,7 @@ public class IndexConstructionWithRandomSetBenchmark {
     @Param({"0", "16"})
     private int numberOfPQSubspaces;
 
-    @Setup(Level.Invocation)
+    @Setup(Level.Trial)
     public void setup() throws IOException {
 
         final var baseVectors = new ArrayList<VectorFloat<?>>(numBaseVectors);

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

@@ -16,10 +16,18 @@
 
 package io.github.jbellis.jvector.quantization;
 
+import io.github.jbellis.jvector.graph.similarity.ScoreFunction;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import io.github.jbellis.jvector.vector.VectorUtil;
 import io.github.jbellis.jvector.vector.types.ByteSequence;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+
+import java.util.HashMap;
+import java.util.Map;
 
 public class ImmutablePQVectors extends PQVectors {
     private final int vectorCount;
+    private final Map<VectorSimilarityFunction, VectorFloat<?>> codebookPartialSumsMap;
 
     /**
      * Construct an immutable PQVectors instance with the given ProductQuantization and compressed data chunks.
@@ -33,6 +41,7 @@ public ImmutablePQVectors(ProductQuantization pq, ByteSequence<?>[] compressedDa
         this.compressedDataChunks = compressedDataChunks;
         this.vectorCount = vectorCount;
         this.vectorsPerChunk = vectorsPerChunk;
+        this.codebookPartialSumsMap = new HashMap<>();
     }
 
     @Override
@@ -44,4 +53,54 @@ protected int validChunkCount() {
     public int count() {
         return vectorCount;
     }
+
+    private synchronized VectorFloat<?> getOrCreateCodebookPartialSums(VectorSimilarityFunction vsf) {
+        return codebookPartialSumsMap.computeIfAbsent(vsf, pq::createCodebookPartialSums);
+    }
+
+    @Override
+    public ScoreFunction.ApproximateScoreFunction diversityFunctionFor(int node1, VectorSimilarityFunction similarityFunction) {
+        final int subspaceCount = pq.getSubspaceCount();
+        var node1Chunk = getChunk(node1);
+        var node1Offset = getOffsetInChunk(node1);
+        int clusterCount = pq.getClusterCount();
+
+        VectorFloat<?> codebookPartialSums = getOrCreateCodebookPartialSums(similarityFunction);
+
+        switch (similarityFunction) {
+            case DOT_PRODUCT:
+                return (node2) -> {
+                    var node2Chunk = getChunk(node2);
+                    var node2Offset = getOffsetInChunk(node2);
+                    // compute the euclidean distance between the query and the codebook centroids corresponding to the encoded points
+                    float sum = VectorUtil.assembleAndSumPQ(codebookPartialSums, subspaceCount, node1Chunk, node1Offset, node2Chunk, node2Offset, clusterCount);
+                    // scale to [0, 1]
+                    return (1 + sum) / 2;
+                };
+            case COSINE:
+                float norm1 = VectorUtil.assembleAndSumPQ(codebookPartialSums, subspaceCount, node1Chunk, node1Offset, node1Chunk, node1Offset, clusterCount);
+                return (node2) -> {
+                    var node2Chunk = getChunk(node2);
+                    var node2Offset = getOffsetInChunk(node2);
+                    // compute the dot product of the query and the codebook centroids corresponding to the encoded points
+                    float sum = VectorUtil.assembleAndSumPQ(codebookPartialSums, subspaceCount, node1Chunk, node1Offset, node2Chunk, node2Offset, clusterCount);
+                    float norm2 = VectorUtil.assembleAndSumPQ(codebookPartialSums, subspaceCount, node2Chunk, node2Offset, node2Chunk, node2Offset, clusterCount);
+                    float cosine = sum / (float) Math.sqrt(norm1 * norm2);
+                    // scale to [0, 1]
+                    return (1 + cosine) / 2;
+                };
+            case EUCLIDEAN:
+                return (node2) -> {
+                    var node2Chunk = getChunk(node2);
+                    var node2Offset = getOffsetInChunk(node2);
+                    // compute the euclidean distance between the query and the codebook centroids corresponding to the encoded points
+                    float sum = VectorUtil.assembleAndSumPQ(codebookPartialSums, subspaceCount, node1Chunk, node1Offset, node2Chunk, node2Offset, clusterCount);
+
+                    // scale to [0, 1]
+                    return 1 / (1 + sum);
+                };
+            default:
+                throw new IllegalArgumentException("Unsupported similarity function " + similarityFunction);
+        }
+    }
 }

jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java

@@ -22,6 +22,7 @@
 import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex;
 import io.github.jbellis.jvector.util.Accountable;
 import io.github.jbellis.jvector.util.PhysicalCoreExecutor;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
 import io.github.jbellis.jvector.vector.VectorUtil;
 import io.github.jbellis.jvector.vector.VectorizationProvider;
 import io.github.jbellis.jvector.vector.types.ByteSequence;
@@ -36,6 +37,7 @@
 import java.util.concurrent.ForkJoinPool;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -586,6 +588,35 @@ public void write(DataOutput out, int version) throws IOException
         }
     }
 
+    /**
+     * Creates a vector to hold partial sums for a single codebook.
+     * The partial sums are the dot products of each subvector centroid in the codebook with the other subvector centroids.
+     * Since the dot product is commutative, we only need to store the upper triangle of the matrix.
+     * There are M codebooks, and each codebook has k centroids, so the total number of partial sums is M * k * (k+1) / 2.
+     *
+     * @return a vector to hold partial sums for a single codebook
+     */
+    public VectorFloat<?> createCodebookPartialSums(VectorSimilarityFunction vectorSimilarityFunction) {
+        VectorFloat<?> partialSums = vectorTypeSupport.createFloatVector(getSubspaceCount() * getClusterCount() * (getClusterCount() + 1) / 2);
+        int index = 0;
+        for (int m = 0; m < M; m++) {
+            int size = subvectorSizesAndOffsets[m][0];
+            var codebook = codebooks[m];
+            for (int i = 0; i < clusterCount; i++) {
+                for (int j = i; j < clusterCount; j++) {
+
+                    float sum = vectorSimilarityFunction == VectorSimilarityFunction.EUCLIDEAN ?
+                            VectorUtil.squareL2Distance(codebook, i * size, codebook, j * size, size) :
+                            VectorUtil.dotProduct(codebook, i * size, codebook, j * size, size);
+
+                    partialSums.set(index++, sum);
+                }
+            }
+        }
+
+        return partialSums;
+    }
+
     @Override
     public int compressorSize() {
         int size = 0;
@@ -734,4 +765,8 @@ private static void checkClusterCount(int clusterCount) {
             LOG.warning("Using less than 256 PQ clusters will not reduce the memory footprint.");
         }
     }
+
+    public int getOriginalDimension() {
+        return originalDimension;
+    }
 }

jvector-base/src/main/java/io/github/jbellis/jvector/vector/DefaultVectorUtilSupport.java

@@ -308,6 +308,36 @@ public float assembleAndSum(VectorFloat<?> data, int dataBase, ByteSequence<?> b
     return sum;
   }
 
+  @Override
+  public float assembleAndSumPQ(
+          VectorFloat<?> codebookPartialSums,
+          int subspaceCount,                  // = M
+          ByteSequence<?> vector1Ordinals,
+          int vector1OrdinalOffset,
+          ByteSequence<?> vector2Ordinals,
+          int vector2OrdinalOffset,
+          int clusterCount
+  ) {
+      final int k          = clusterCount;
+      final int blockSize  = k * (k + 1) / 2;
+      float res = 0f;
+
+      for (int i = 0; i < subspaceCount; i++) {
+          int c1 = Byte.toUnsignedInt(vector1Ordinals.get(i + vector1OrdinalOffset));
+          int c2 = Byte.toUnsignedInt(vector2Ordinals.get(i + vector2OrdinalOffset));
+          int r  = Math.min(c1, c2);
+          int c  = Math.max(c1, c2);
+
+          int offsetRow  = r * k - (r * (r - 1) / 2);
+          int idxInBlock = offsetRow + (c - r);
+          int base       = i * blockSize;
+
+          res += codebookPartialSums.get(base + idxInBlock);
+      }
+
+      return res;
+  }
+
   @Override
   public int hammingDistance(long[] v1, long[] v2) {
     int hd = 0;

jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtil.java

@@ -170,6 +170,10 @@ public static float assembleAndSum(VectorFloat<?> data, int dataBase, ByteSequen
     return impl.assembleAndSum(data, dataBase, dataOffsets, dataOffsetsOffset, dataOffsetsLength);
   }
 
+  public static float assembleAndSumPQ(VectorFloat<?> data, int subspaceCount, ByteSequence<?> dataOffsets1, int dataOffsetsOffset1, ByteSequence<?> dataOffsets2, int dataOffsetsOffset2, int clusterCount) {
+    return impl.assembleAndSumPQ(data, subspaceCount, dataOffsets1, dataOffsetsOffset1, dataOffsets2, dataOffsetsOffset2, clusterCount);
+  }
+
   public static void bulkShuffleQuantizedSimilarity(ByteSequence<?> shuffles, int codebookCount, ByteSequence<?> quantizedPartials, float delta, float minDistance, VectorFloat<?> results, VectorSimilarityFunction vsf) {
     impl.bulkShuffleQuantizedSimilarity(shuffles, codebookCount, quantizedPartials, delta, minDistance, vsf, results);
   }
@@ -250,5 +254,4 @@ public static float nvqLoss(VectorFloat<?> vector, float growthRate, float midpo
   public static float nvqUniformLoss(VectorFloat<?> vector, float minValue, float maxValue, int nBits) {
     return impl.nvqUniformLoss(vector, minValue, maxValue, nBits);
   }
-
 }

jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtilSupport.java

@@ -24,6 +24,7 @@
 
 package io.github.jbellis.jvector.vector;
 
+import io.github.jbellis.jvector.quantization.ProductQuantization;
 import io.github.jbellis.jvector.vector.types.ByteSequence;
 import io.github.jbellis.jvector.vector.types.VectorFloat;
 
@@ -113,6 +114,22 @@ public interface VectorUtilSupport {
    */
   float assembleAndSum(VectorFloat<?> data, int baseIndex, ByteSequence<?> baseOffsets, int baseOffsetsOffset, int baseOffsetsLength);
 
+  /**
+   * Calculates the distance between 2 vectors, which were quantized using Product Quantization, using a precomputed table of partial results
+   *
+   * See {@link ProductQuantization#createCodebookPartialSums(VectorSimilarityFunction)}
+   *
+   * @param codebookPartialSums the vector of all PQ
+   * @param subspaceCount the number of PQ subspaces
+   * @param vector1Ordinals Specifies which centroid vector is used for each of node1's subvectors
+   * @param vector1OrdinalOffset the offset into the vector1Ordinals ByteSequence for node1 (in case vector1Ordinals is a chunk of many nodes)
+   * @param node2Ordinals Specifies which centroid vector is used for each of node2's subvectors
+   * @param node2OrdinalOffset the offset into the vector1Ordinals ByteSequence for node2 (in case vector1Ordinals is a chunk of many nodes)
+   * @param clusterCount the number of PQ clusters per subvector in the codebook
+   * @return the sum of the vectors
+   */
+  float assembleAndSumPQ(VectorFloat<?> codebookPartialSums, int subspaceCount, ByteSequence<?> vector1Ordinals, int vector1OrdinalOffset, ByteSequence<?> node2Ordinals, int node2OrdinalOffset, int clusterCount);
+
   int hammingDistance(long[] v1, long[] v2);
 
   // default implementation used here because Panama SIMD can't express necessary SIMD operations and degrades to scalar

jvector-native/src/main/java/io/github/jbellis/jvector/vector/NativeVectorUtilSupport.java

@@ -80,6 +80,20 @@ public float assembleAndSum(VectorFloat<?> data, int dataBase, ByteSequence<?> b
     }
 
 
+    @Override
+    public float assembleAndSumPQ(
+            VectorFloat<?> codebookPartialSums,
+            int subspaceCount,                  // = M
+            ByteSequence<?> vector1Ordinals,
+            int vector1OrdinalOffset,
+            ByteSequence<?> vector2Ordinals,
+            int vector2OrdinalOffset,
+            int clusterCount                    // = k
+    ) {
+        //Use the non-panama solution for now
+        return assembleAndSumPQ_128(codebookPartialSums, subspaceCount, vector1Ordinals, vector1OrdinalOffset, vector2Ordinals, vector2OrdinalOffset, clusterCount);
+    }
+
     @Override
     public void calculatePartialSums(VectorFloat<?> codebook, int codebookBase, int size, int clusterCount, VectorFloat<?> query, int queryOffset, VectorSimilarityFunction vsf, VectorFloat<?> partialSums) {
         switch (vsf) {

jvector-tests/src/test/java/io/github/jbellis/jvector/quantization/TestProductQuantization.java

@@ -34,6 +34,7 @@
 import java.nio.file.Files;
 import java.util.Arrays;
 import java.util.List;
+import java.util.concurrent.ForkJoinPool;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
@@ -406,5 +407,31 @@ public void testPQLayoutEdgeCases() {
         System.out.println("Test completed successfully");
     }
 
+    @Test
+    public void testPQCodebookSums() {
+        // Generate a PQ for random 2D vectors
+        var vectors = createRandomVectors(10000, 384);
+        var pq = ProductQuantization.compute(new ListRandomAccessVectorValues(vectors, 384), 48, 256, false);
+
+        MutablePQVectors pqm = new MutablePQVectors(pq);
+
+        // build the index vector-at-a-time (on disk)
+        for (int ordinal = 0; ordinal < vectors.size(); ordinal++)
+        {
+            VectorFloat<?> v = vectors.get(ordinal);
+            // compress the new vector and add it to the PQVectors
+            pqm.encodeAndSet(ordinal, v);
+        }
 
+        for (VectorSimilarityFunction vsf : VectorSimilarityFunction.values()) {
+            var sf = pqm.diversityFunctionFor(10, vsf);
+
+            ImmutablePQVectors pqi = ImmutablePQVectors.encodeAndBuild(pq, vectors.size(), new ListRandomAccessVectorValues(vectors, 384), ForkJoinPool.commonPool());
+            var sf2 = pqi.diversityFunctionFor(10, vsf);
+
+            for (int i = 0; i < vectors.size(); i++) {
+                assertEquals(vsf.name(), sf.similarityTo(i), sf2.similarityTo(i), 1e-6);
+            }
+        }
+    }
 }