change the centroids initialization

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

vortex-tensor/src/encodings/turboquant/centroids.rs

@@ -3,11 +3,23 @@
 
 //! Max-Lloyd centroid computation for TurboQuant scalar quantizers.
 //!
-//! Pre-computes optimal scalar quantizer centroids for the marginal distribution of coordinates
-//! after random rotation of a unit-norm vector. In high dimensions, each coordinate of a randomly
-//! rotated unit vector follows a distribution proportional to `(1 - x^2)^((d-3)/2)` on `[-1, 1]`,
-//! which converges to `N(0, 1/d)`. The Max-Lloyd algorithm finds optimal quantization centroids
-//! that minimize MSE for this distribution.
+//! Pre-computes and caches optimal scalar quantizer centroids for the marginal distribution of
+//! coordinates after a random orthogonal transform of a unit-norm vector.
+//!
+//! In high dimensions, each coordinate of a randomly transformed unit vector follows a
+//! distribution proportional to `(1 - x^2)^((d-3)/2)` on `[-1, 1]`, which converges to
+//! `N(0, 1/d)`.
+//!
+//! The Max-Lloyd algorithm finds optimal quantization centroids that minimize MSE for this
+//! distribution.
+//!
+//! Centroids are not stored in TurboQuant arrays. They are deterministically derived from
+//! `(padded_dim, bit_width)` and cached process-locally.
+//!
+//! The centroid model follows the random orthogonal transform marginal used by the TurboQuant
+//! paper. This encoder applies a SORF-style structured transform instead of a dense random Gaussian
+//! or orthogonal matrix, so paper-level error bounds should not be treated as verified for this
+//! implementation without separate empirical validation.
 
 use std::sync::LazyLock;
 
@@ -34,9 +46,9 @@ static CENTROID_CACHE: LazyLock<DashMap<(u32, u8), Buffer<f32>>> = LazyLock::new
 /// Get or compute cached centroids for the given dimension and bit width.
 ///
 /// Returns `2^bit_width` centroids sorted in ascending order, representing optimal scalar
-/// quantization levels for the coordinate distribution after random rotation in
+/// quantization levels for the coordinate distribution after a random orthogonal transform in
 /// `dimension`-dimensional space.
-pub fn compute_or_get_centroids(dimension: u32, bit_width: u8) -> VortexResult<Buffer<f32>> {
+pub(crate) fn compute_or_get_centroids(dimension: u32, bit_width: u8) -> VortexResult<Buffer<f32>> {
     vortex_ensure!(
         (1..=MAX_BIT_WIDTH).contains(&bit_width),
         "TurboQuant bit_width must be 1-{}, got {bit_width}",
@@ -86,22 +98,35 @@ impl HalfIntExponent {
 
 /// Compute optimal centroids via the Max-Lloyd (Lloyd-Max) algorithm.
 ///
-/// Operates on the marginal distribution of a single coordinate of a randomly rotated unit vector
-/// in d dimensions.
+/// Operates on the marginal distribution of a single coordinate of a randomly transformed unit
+/// vector in d dimensions.
 ///
 /// The probability distribution function is:
 ///   `f(x) = C_d * (1 - x^2)^((d-3)/2)` on `[-1, 1]`
 /// where `C_d` is the normalizing constant.
+///
+/// Centroids are seeded uniformly on `[±sqrt(bit_width) * sigma]` (where `sigma` is the standard
+/// deviation of the normal distribution that hyphershere dimension values take, and specifically
+/// `sigma = 1/sqrt( dimension)`) rather than across the full `[-1, 1]`, which strands most of the
+/// centroids in the near-zero-mass tails.
+///
+/// Note that the `sqrt(bit_width)` is mostly empirically derived, we do not have a theoretical
+/// basis for choosing this other than the fact that it seems to produce good results.
 fn max_lloyd_centroids(dimension: u32, bit_width: u8) -> Buffer<f32> {
     debug_assert!((1..=MAX_BIT_WIDTH).contains(&bit_width));
     let num_centroids = 1usize << bit_width;
 
     // For the marginal distribution on [-1, 1], we use the exponent (d-3)/2.
     let exponent = HalfIntExponent::from_numerator(dimension as i32 - 3);
 
-    // Initialize centroids uniformly on [-1, 1].
+    // The coordinate marginal concentrates around 0 with this standard deviation.
+    let sigma = 1.0 / f64::from(dimension).sqrt();
+    let init_half = (f64::from(bit_width).sqrt() * sigma).min(1.0);
+
+    // Initialize centroids uniformly on [-init_half, init_half], where the mass lives, so no cell
+    // starts in a zero-mass region and freezes.
     let mut centroids: Vec<f64> = (0..num_centroids)
-        .map(|idx| -1.0 + (2.0 * (idx as f64) + 1.0) / (num_centroids as f64))
+        .map(|idx| -init_half + (2.0 * (idx as f64) + 1.0) * init_half / (num_centroids as f64))
         .collect();
 
     let mut boundaries: Vec<f64> = vec![0.0; num_centroids + 1];
@@ -193,7 +218,7 @@ fn pdf_unnormalized(x_val: f64, exponent: HalfIntExponent) -> f64 {
 /// For `k` centroids, returns `k-1` boundaries. A value below `boundaries[0]` maps to centroid 0, a
 /// value in `[boundaries[i-1], boundaries[i])` maps to centroid `i`, and a
 /// value `>= boundaries[k-2]` maps to centroid `k-1`.
-pub fn compute_centroid_boundaries(centroids: &[f32]) -> Vec<f32> {
+pub(crate) fn compute_centroid_boundaries(centroids: &[f32]) -> Vec<f32> {
     centroids.windows(2).map(|w| (w[0] + w[1]) * 0.5).collect()
 }
 
@@ -203,7 +228,7 @@ pub fn compute_centroid_boundaries(centroids: &[f32]) -> Vec<f32> {
 /// centroids. Uses binary search on the midpoints, avoiding distance comparisons
 /// in the inner loop.
 #[inline]
-pub fn find_nearest_centroid(value: f32, boundaries: &[f32]) -> u8 {
+pub(crate) fn find_nearest_centroid(value: f32, boundaries: &[f32]) -> u8 {
     debug_assert!(
         boundaries.windows(2).all(|w| w[0] <= w[1]),
         "boundaries must be sorted"

vortex-turboquant/src/centroids.rs

@@ -106,16 +106,29 @@ impl HalfIntExponent {
 /// The probability distribution function is:
 ///   `f(x) = C_d * (1 - x^2)^((d-3)/2)` on `[-1, 1]`
 /// where `C_d` is the normalizing constant.
+///
+/// Centroids are seeded uniformly on `[±sqrt(bit_width) * sigma]` (where `sigma` is the standard
+/// deviation of the normal distribution that hyphershere dimension values take, and specifically
+/// `sigma = 1/sqrt( dimension)`) rather than across the full `[-1, 1]`, which strands most of the
+/// centroids in the near-zero-mass tails.
+///
+/// Note that the `sqrt(bit_width)` is mostly empirically derived, we do not have a theoretical
+/// basis for choosing this other than the fact that it seems to produce good results.
 fn max_lloyd_centroids(dimension: u32, bit_width: u8) -> Buffer<f32> {
     debug_assert!((1..=MAX_BIT_WIDTH).contains(&bit_width));
     let num_centroids = 1usize << bit_width;
 
     // For the marginal distribution on [-1, 1], we use the exponent (d-3)/2.
     let exponent = HalfIntExponent::from_numerator(dimension as i32 - 3);
 
-    // Initialize centroids uniformly on [-1, 1].
+    // The coordinate marginal concentrates around 0 with this standard deviation.
+    let sigma = 1.0 / f64::from(dimension).sqrt();
+    let init_half = (f64::from(bit_width).sqrt() * sigma).min(1.0);
+
+    // Initialize centroids uniformly on [-init_half, init_half], where the mass lives, so no cell
+    // starts in a zero-mass region and freezes.
     let mut centroids: Vec<f64> = (0..num_centroids)
-        .map(|idx| -1.0 + (2.0 * (idx as f64) + 1.0) / (num_centroids as f64))
+        .map(|idx| -init_half + (2.0 * (idx as f64) + 1.0) * init_half / (num_centroids as f64))
         .collect();
 
     let mut boundaries: Vec<f64> = vec![0.0; num_centroids + 1];