Introduce Projected Eigen distance function for multi-vector re-ranking (#1203)

<!--
Thanks for contributing a pull request! Please ensure you have taken a
look at
the contribution guidelines:
https://github.com/microsoft/DiskANN/blob/main/CONTRIBUTING.md
-->
- [x] Does this PR have a descriptive title that could go in our release
notes?
- [ ] Does this PR add any new dependencies?
- [ ] Does this PR modify any existing APIs?
- [ ] Is the change to the API backwards compatible?
- [ ] Should this result in any changes to our documentation, either
updating existing docs or adding new ones?

#### Reference Issues/PRs
#1201 
<!--
Example: Fixes #1234. See also #3456.
Please use keywords (e.g., Fixes) to create link to the issues or pull
requests
you resolved, so that they will automatically be closed when your pull
request
is merged. See
https://github.com/blog/1506-closing-issues-via-pull-requests
-->

#### What does this implement/fix? Briefly explain your changes.
Introduce a new distance function for multi-vector re-ranking:
`ProjectedEigen`.
Given a query multivector `Q = [q_1, q_2, .. q_K]` and doc multivector
`D = [d_1, d_2, ... d_N]`, `ProjectedEigen` distance between the two
(asymmetric) is defined as $$\sum_{i=1}^{i=K} \sum_{j=1}^{j=N} -
(IP(q_i, d_j)^2)$$.

This PR introduces preliminary support for this distance function inside
the `diskann-quantization` crate (similar to `MaxSim`).
#### Any other comments?

Author: suhasjs

diskann-quantization/src/multi_vector/distance/fallback.rs

@@ -9,6 +9,7 @@ use diskann_vector::distance::InnerProduct;
 use diskann_vector::{DistanceFunctionMut, PureDistanceFunction};
 
 use super::max_sim::{Chamfer, MaxSim};
+use super::projected_eigen::ProjectedEigen;
 use crate::multi_vector::{MatRef, MaxSimError, Repr, Standard};
 
 /////////////////
@@ -100,6 +101,50 @@ impl FallbackKernel {
             f(i, min_dist);
         }
     }
+
+    /// Core kernel for computing per-query-vector projected-eigen scores.
+    ///
+    /// For each `query` vector, sums the negated squared inner product
+    /// against every document vector, then calls `f(index, score)` with the
+    /// result. If there are no vectors in the `doc`, the kernel returns
+    /// immediately.
+    ///
+    /// The callback can be used to aggregate scores as needed - as is the
+    /// case with [`ProjectedEigen`].
+    ///
+    /// # Arguments
+    ///
+    /// * `query` - The query multi-vector (wrapped as [`QueryMatRef`])
+    /// * `doc` - The document multi-vector
+    /// * `f` - Callback invoked with `(query_index, score)` for each query vector
+    #[inline]
+    pub(crate) fn projected_eigen_kernel<F, T: Copy>(
+        query: QueryMatRef<'_, Standard<T>>,
+        doc: MatRef<'_, Standard<T>>,
+        mut f: F,
+    ) where
+        F: FnMut(usize, f32),
+        InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>,
+    {
+        // Early exit if no doc vectors - callback should never be invoked
+        if doc.num_vectors() == 0 {
+            return;
+        }
+
+        for (i, q_vec) in query.rows().enumerate() {
+            let mut sum = 0.0f32;
+
+            for d_vec in doc.rows() {
+                // `InnerProduct::evaluate` returns the negated inner product;
+                // squaring discards the sign, so negate the squared value to
+                // obtain `-IP(q, d)²`.
+                let ip = InnerProduct::evaluate(q_vec, d_vec);
+                sum += -(ip * ip);
+            }
+
+            f(i, sum);
+        }
+    }
 }
 
 ////////////
@@ -159,6 +204,27 @@ where
     }
 }
 
+/////////////////////
+// ProjectedEigen //
+/////////////////////
+
+impl<T: Copy> PureDistanceFunction<QueryMatRef<'_, Standard<T>>, MatRef<'_, Standard<T>>, f32>
+    for ProjectedEigen
+where
+    InnerProduct: for<'a, 'b> PureDistanceFunction<&'a [T], &'b [T], f32>,
+{
+    #[inline(always)]
+    fn evaluate(query: QueryMatRef<'_, Standard<T>>, doc: MatRef<'_, Standard<T>>) -> f32 {
+        let mut sum = 0.0f32;
+
+        FallbackKernel::projected_eigen_kernel(query, doc, |_i, score| {
+            sum += score;
+        });
+
+        sum
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -185,6 +251,17 @@ mod tests {
             .fold(f32::MAX, f32::min)
     }
 
+    /// Naive implementation of projected-eigen for a single query vector
+    /// against all doc vectors: `\sum_{j} -IP(q, d_{j})^2`.
+    fn naive_projected_eigen_single(query_vec: &[f32], doc: &MatRef<'_, Standard<f32>>) -> f32 {
+        doc.rows()
+            .map(|d_vec| {
+                let ip: f32 = query_vec.iter().zip(d_vec.iter()).map(|(a, b)| a * b).sum();
+                -(ip * ip)
+            })
+            .sum()
+    }
+
     /// Generate deterministic test data.
     fn make_test_data(len: usize, ceil: usize, shift: usize) -> Vec<f32> {
         (0..len).map(|v| ((v + shift) % ceil) as f32).collect()
@@ -286,6 +363,22 @@ mod tests {
                     nd,
                     dim
                 );
+
+                // Test ProjectedEigen
+                let projected = ProjectedEigen::evaluate(query, doc);
+                let expected_projected: f32 = query
+                    .rows()
+                    .map(|q_vec| naive_projected_eigen_single(q_vec, &doc))
+                    .sum();
+
+                assert!(
+                    (projected - expected_projected).abs()
+                        < 1e-6 * expected_projected.abs().max(1.0),
+                    "ProjectedEigen mismatch for ({},{},{})",
+                    nq,
+                    nd,
+                    dim
+                );
             }
         }
 
@@ -303,5 +396,15 @@ mod tests {
 
             assert_eq!(result, 0.0);
         }
+
+        #[test]
+        fn projected_eigen_with_zero_docs_returns_zero() {
+            let query = make_query(&[1.0, 0.0, 0.0, 1.0], 2, 2);
+            let doc = make_doc(&[], 0, 2);
+
+            // No document vectors means no pairs contribute, so the sum is 0.
+            let result = ProjectedEigen::evaluate(query, doc);
+            assert_eq!(result, 0.0);
+        }
     }
 }

diskann-quantization/src/multi_vector/distance/mod.rs

@@ -47,9 +47,11 @@ mod isa;
 mod kernel;
 mod kernels;
 mod max_sim;
+mod projected_eigen;
 
 pub use factory::{MaxSimElement, build_max_sim};
 pub use fallback::QueryMatRef;
 pub use isa::{MaxSimIsa, NotSupported};
 pub use kernel::{BoxErase, Erase, MaxSimKernel};
 pub use max_sim::{Chamfer, MaxSim, MaxSimError};
+pub use projected_eigen::ProjectedEigen;

diskann-quantization/src/multi_vector/distance/projected_eigen.rs

@@ -0,0 +1,27 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+//! Projected-eigen distance type for multi-vector representations.
+
+/////////////////////
+// ProjectedEigen //
+/////////////////////
+
+/// Projected-eigen distance for multi-vector similarity.
+///
+/// Computes the negated sum of squared inner products over *all*
+/// query/document vector pairs:
+///
+/// ```text
+/// ProjectedEigen(Q, D) = \sum_{i} \sum_{j} -IP(q_i, d_j)²
+/// ```
+///
+/// Unlike [`Chamfer`](super::Chamfer), which keeps only the best-matching
+/// document vector per query vector, this accumulates a contribution from
+/// every pair, so the score reflects the full query–document interaction. As
+/// with the other multi-vector distances, lower is better.
+///
+/// Implements [`PureDistanceFunction`](diskann_vector::PureDistanceFunction)
+/// for matrix view types.
+#[derive(Debug, Clone, Copy)]
+pub struct ProjectedEigen;

diskann-quantization/src/multi_vector/mod.rs

@@ -58,7 +58,7 @@ pub(crate) mod matrix;
 pub use block_transposed::{BlockTransposed, BlockTransposedMut, BlockTransposedRef};
 pub use distance::{
     BoxErase, Chamfer, Erase, MaxSim, MaxSimElement, MaxSimError, MaxSimIsa, MaxSimKernel,
-    NotSupported, QueryMatRef, build_max_sim,
+    NotSupported, ProjectedEigen, QueryMatRef, build_max_sim,
 };
 pub use matrix::{
     Defaulted, LayoutError, Mat, MatMut, MatRef, NewCloned, NewMut, NewOwned, NewRef, Overflow,