feat: add metrics dot and cos

Signed-off-by: cutecutecat <junyuchen@tensorchord.ai>

Author: cutecutecat

crates/k_means/src/lib.rs

@@ -61,6 +61,26 @@ pub fn k_means_lookup<S: ScalarLike>(vector: &[S], centroids: &Vec2<S>) -> usize
     result.1
 }
 
+/// returns (centroid_dot_dis, (vector_l2_norm, centroids_l2_norm, index))
+pub fn k_means_lookup_by_dot<S: ScalarLike>(
+    vector: &[S],
+    centroids: &Vec2<S>,
+) -> (f32, f32, f32, usize) {
+    assert_ne!(centroids.shape_0(), 0);
+    let mut result = (f32::INFINITY, f32::INFINITY, f32::zero(), f32::zero(), 0);
+    for i in 0..centroids.shape_0() {
+        let dot = S::reduce_sum_of_xy(vector, &centroids[(i,)]);
+        let vector_square = S::reduce_sum_of_x2(vector);
+        let centroids_square = S::reduce_sum_of_x2(&centroids[(i,)]);
+
+        let l2_dis = vector_square + centroids_square - 2.0 * dot;
+        if l2_dis <= result.0 {
+            result = (l2_dis, -dot, vector_square, centroids_square, i);
+        }
+    }
+    (result.1, result.2, result.3, result.4)
+}
+
 pub fn k_means_lookup_many<S: ScalarLike>(vector: &[S], centroids: &Vec2<S>) -> Vec<(f32, usize)> {
     assert_ne!(centroids.shape_0(), 0);
     let mut seq = Vec::new();
@@ -71,6 +91,23 @@ pub fn k_means_lookup_many<S: ScalarLike>(vector: &[S], centroids: &Vec2<S>) ->
     seq
 }
 
+/// returns Vec of <l2_dis, (centroid_dot_dis, vector_l2_norm, centroids_l2_norm, index)>
+pub fn k_means_lookup_many_by_dot<S: ScalarLike>(
+    vector: &[S],
+    centroids: &Vec2<S>,
+) -> Vec<(f32, (f32, f32, f32, usize))> {
+    assert_ne!(centroids.shape_0(), 0);
+    let mut seq = Vec::new();
+    for i in 0..centroids.shape_0() {
+        let dot = S::reduce_sum_of_xy(vector, &centroids[(i,)]);
+        let vector_square = S::reduce_sum_of_x2(vector);
+        let centroids_square = S::reduce_sum_of_x2(&centroids[(i,)]);
+        let l2_dis = vector_square + centroids_square - 2.0 * dot;
+        seq.push((l2_dis, (-dot, vector_square, centroids_square, i)));
+    }
+    seq
+}
+
 fn spherical<S: ScalarLike>(vector: &mut [S]) {
     let l = S::reduce_sum_of_x2(vector).sqrt();
     S::vector_mul_scalar_inplace(vector, 1.0 / l);

crates/quantization/src/quantize.rs

@@ -152,9 +152,13 @@ mod mul_add {
 #[inline(always)]
 pub fn quantize<const N: u8>(lut: &[f32]) -> (f32, f32, Vec<u8>) {
     let (min, max) = f32::reduce_min_max_of_x(lut);
-    let k = 0.0f32.max((max - min) / (N as f32));
-    let b = min;
-    (k, b, mul_add::mul_add(lut, 1.0 / k, -b / k))
+    let delta = 0.0f32.max((max - min) / (N as f32));
+    let lower_bound = min;
+    (
+        delta,
+        lower_bound,
+        mul_add::mul_add(lut, 1.0 / delta, -lower_bound / delta),
+    )
 }
 
 #[inline(always)]

crates/rabitq/src/lib.rs

@@ -17,7 +17,7 @@ use common::json::Json;
 use common::mmap_array::MmapArray;
 use common::remap::RemappedCollection;
 use common::vec2::Vec2;
-use k_means::{k_means, k_means_lookup, k_means_lookup_many};
+use k_means::{k_means, k_means_lookup, k_means_lookup_by_dot, k_means_lookup_many_by_dot};
 use std::fs::create_dir;
 use std::path::Path;
 use stoppable_rayon as rayon;
@@ -69,14 +69,18 @@ impl<O: Op> Rabitq<O> {
     ) -> Box<dyn Iterator<Item = Element> + 'a> {
         let projected_query = O::proj(&self.projection, O::cast(vector));
         let lists = select(
-            k_means_lookup_many(&projected_query, &self.centroids),
+            k_means_lookup_many_by_dot(&projected_query, &self.centroids),
             opts.rabitq_nprobe as usize,
         );
         let mut heap = Vec::new();
-        for &(_, i) in lists.iter() {
-            let preprocessed = self
-                .quantization
-                .preprocess(&O::residual(&projected_query, &self.centroids[(i,)]));
+        for &(_, payload) in lists.iter() {
+            let (centroid_dot_dis, original_square, centroids_square, i) = payload;
+            let preprocessed = self.quantization.preprocess(
+                &O::residual(&projected_query, &self.centroids[(i,)]),
+                centroid_dot_dis,
+                original_square,
+                centroids_square,
+            );
             let start = self.offsets[i];
             let end = self.offsets[i + 1];
             self.quantization.push_batch(
@@ -150,8 +154,13 @@ fn from_nothing<O: Op>(
         collection.len(),
         |vector| {
             let vector = O::cast(collection.vector(vector));
-            let target = k_means_lookup(vector, &centroids);
-            O::proj(&projection, &O::residual(vector, &centroids[(target,)]))
+            let (centroid_dot_dis, original_square, _, target) =
+                k_means_lookup_by_dot(vector, &centroids);
+            (
+                O::proj(&projection, &O::residual(vector, &centroids[(target,)])),
+                centroid_dot_dis,
+                original_square,
+            )
         },
     );
     let projected_centroids = Vec2::from_vec(
@@ -194,7 +203,7 @@ fn open<O: Op>(path: impl AsRef<Path>) -> Rabitq<O> {
     }
 }
 
-fn select(mut lists: Vec<(f32, usize)>, n: usize) -> Vec<(f32, usize)> {
+fn select<T>(mut lists: Vec<(f32, T)>, n: usize) -> Vec<(f32, T)> {
     if lists.is_empty() || n == 0 {
         return Vec::new();
     }