diskann_quantization/minmax/
quantizer.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use super::vectors::{DataMutRef, FullQueryMut, MinMaxCompensation, MinMaxIP, MinMaxL2Squared};
7use core::f32;
8
9use crate::{
10    AsFunctor, CompressInto,
11    algorithms::Transform,
12    alloc::{GlobalAllocator, ScopedAllocator},
13    bits::{Representation, Unsigned},
14    minmax::{MinMaxCosine, MinMaxCosineNormalized, vectors::FullQueryMeta},
15    num::Positive,
16    scalar::{InputContainsNaN, bit_scale},
17};
18
19/// Recall that from the module-level documentation, MinMaxQuantizer, quantizes X
20/// into `n` bit vectors as follows  -
21/// ```math
22/// X' = round((X - s) * (2^n - 1) / c).clamp(0, 2^n - 1))
23/// ```
24/// where `s` is a shift value and `c` is a scaling parameter computed from the range of values.
25///
26/// For most bit widths (>1), given a positive scaling parameter `grid_scale : f32`,
27/// these are computed as:
28/// ```math
29/// - m = (max_i X[i] + min_i X[i]) / 2.0
30/// - w = max_i X[i] - min_i X[i]
31///
32/// - s = m - w * grid_scale
33/// - c = 2 * w * grid_scale
34///
35/// where `grid_scale` is an input to the quantizer.
36/// ```
37/// For 1-bit quantization, to avoid outliers, `s` and `c` are derived differently:
38/// - Values are first split into two groups: those below and above the mean.
39/// - `s` is the average of values below the mean.
40/// - `c` is the difference between the average of values above the mean and `s`.
41///
42/// See [`MinMaxCompensation`] for notation.
43/// We have then that
44/// ```math
45/// X = X' * (c / (2^n - 1)) + s
46///          --------------    -
47///                 |          |
48///                ax          bx
49/// ```
50pub struct MinMaxQuantizer {
51    /// Support for different strategies of pre-transforming vectors before applying compression.
52    /// See [`Transform`] for more details on supported types. The input dimension of vectors
53    /// to the quantizer is derived from `transform.input_dim()`.
54    transform: Transform<GlobalAllocator>,
55
56    /// Scaling parameter used to scale the range (min, max) in order to avoid outliers.
57    /// The input must be a positive value. In general, any value between [0.8, 1] does well.
58    grid_scale: Positive<f32>,
59}
60
61impl MinMaxQuantizer {
62    /// Instantiates a new quantizer with specific transform.
63    pub fn new(transform: Transform<GlobalAllocator>, grid_scale: Positive<f32>) -> Self {
64        Self {
65            transform,
66            grid_scale,
67        }
68    }
69
70    /// Input dimension of vectors to quantizer.
71    pub fn dim(&self) -> usize {
72        self.transform.input_dim()
73    }
74
75    /// Output dimension of vectors after applying transform.
76    ///
77    /// Output storage vectors should use this dimension instead of `self.dim()` because
78    /// in general, the output dim **may** be different from the input dimension.
79    pub fn output_dim(&self) -> usize {
80        self.transform.output_dim()
81    }
82
83    /// Outputs the minimum and maximum value of the range of values
84    /// for an input vector `vec`. The function cases based on the
85    /// intended number of bits `NBITS` per dimension.
86    ///
87    /// * `1-bit` - In order to avoid outlier values, the range
88    ///   is defined by taking the values larger and smaller than
89    ///   the numeric mean, and then taking the respective means of
90    ///   each of these sets as the `max` and `min`.
91    ///
92    /// * `N-bits` - Computes the `min` and `max` of the vector values.
93    ///
94    /// # Returns
95    ///
96    /// * `(m - w * g, m + w * g)` - the lower and upper end of the range, where,
97    ///   `m = (max + min) / 2.0`, `w = (max - min) / 2.0`, and `g = self.grid_scale`.
98    fn get_range<const NBITS: usize>(&self, vec: &[f32]) -> (f32, f32) {
99        let (min, max) = match NBITS {
100            1 => {
101                let (mut min, mut min_count) = (0.0f32, 0.0f32);
102                let (mut max, mut max_count) = (0.0f32, 0.0f32);
103
104                let mean = vec.iter().sum::<f32>() / (vec.len() as f32);
105
106                vec.iter().for_each(|x| {
107                    let m = f32::from((*x < mean) as u8);
108                    min += m * x;
109                    min_count += m;
110                    max += (1.0 - m) * x;
111                    max_count += 1.0 - m;
112                });
113
114                ((min / min_count).min(mean), (max / max_count).max(mean))
115            }
116            _ => {
117                vec // Using `f32::NAN` since [`core::f32::min`] and `max` output the other value if one of them is NAN .
118                    .iter()
119                    .fold((f32::NAN, f32::NAN), |(cmin, cmax), &e| {
120                        (cmin.min(e), cmax.max(e))
121                    })
122            }
123        };
124
125        let width = (max - min) / 2.0;
126        let mid = min + width;
127
128        (
129            mid - width * self.grid_scale.into_inner(),
130            mid + width * self.grid_scale.into_inner(),
131        )
132    }
133
134    fn compress<const NBITS: usize, T>(
135        &self,
136        from: &[T],
137        mut into: DataMutRef<'_, NBITS>,
138    ) -> Result<L2Loss, InputContainsNaN>
139    where
140        T: Copy + Into<f32>,
141        Unsigned: Representation<NBITS>,
142    {
143        let mut into_vec = into.vector_mut();
144
145        assert_eq!(from.len(), self.dim());
146        assert_eq!(self.output_dim(), into_vec.len());
147
148        let domain = Unsigned::domain_const::<NBITS>();
149        let domain_min = *domain.start() as f32;
150        let domain_max = *domain.end() as f32;
151
152        let mut vec = vec![f32::default(); self.output_dim()];
153
154        // We know vec.len() == self.output_dim() and `from.len() == self.dim`
155        #[allow(clippy::unwrap_used)]
156        self.transform
157            .transform_into(
158                &mut vec,
159                &from.iter().map(|&x| x.into()).collect::<Vec<f32>>(),
160                ScopedAllocator::global(),
161            )
162            .unwrap();
163
164        let (min, max) = self.get_range::<NBITS>(&vec);
165
166        let inverse_scale = (max - min).max(1e-8) / bit_scale::<NBITS>(); // To avoid NaN. This is ONLY possible if the vector is all the same value.
167        let mut norm_squared: f32 = 0.0;
168        let mut code_sum: f32 = 0.0;
169        let mut loss: f32 = 0.0;
170
171        let mut nan_check = false;
172
173        vec.iter().enumerate().for_each(|(i, &v)| {
174            nan_check |= v.is_nan();
175
176            let code = ((v - min) / inverse_scale)
177                .clamp(domain_min, domain_max)
178                .round();
179
180            let v_r = (code * inverse_scale) + min; // reconstructed value for `v`.
181            norm_squared += v_r * v_r;
182            code_sum += code;
183            loss += (v_r - v).powi(2);
184
185            //SAFETY: we checked that the lengths of `from` and `into_vec` are the same.
186            unsafe {
187                into_vec.set_unchecked(i, code as u8);
188            }
189        });
190
191        let meta = MinMaxCompensation {
192            dim: self.output_dim() as u32,
193            b: min,
194            a: inverse_scale,
195            n: inverse_scale * code_sum,
196            norm_squared,
197        };
198
199        into.set_meta(meta);
200
201        if nan_check {
202            Err(InputContainsNaN)
203        } else {
204            Ok(match Positive::new(loss) {
205                Ok(p) => L2Loss::Positive(p),
206                Err(_) => L2Loss::Zero,
207            })
208        }
209    }
210}
211
212/////////////////
213// Compression //
214/////////////////
215
216/// A struct defining euclidean loss from quantization.
217///
218/// For an input vector `x` and its representation `x'`,
219/// this is supposed to store `||x - x'||^2`.
220#[derive(Clone, Copy, Debug)]
221pub enum L2Loss {
222    Zero,
223    Positive(Positive<f32>),
224}
225
226impl L2Loss {
227    /// Euclidean loss as a `f32` value
228    pub fn as_f32(&self) -> f32 {
229        match self {
230            L2Loss::Zero => 0.0,
231            L2Loss::Positive(p) => p.into_inner(),
232        }
233    }
234}
235
236impl<const NBITS: usize, T> CompressInto<&[T], DataMutRef<'_, NBITS>> for MinMaxQuantizer
237where
238    T: Copy + Into<f32>,
239    Unsigned: Representation<NBITS>,
240{
241    type Error = InputContainsNaN;
242
243    type Output = L2Loss;
244
245    /// Compress the input vector `from` into a mut ref of Data `to`.
246    ///
247    /// This method computes and stores the compensation coefficients required for computing
248    /// distances correctly.
249    ///
250    /// # Error
251    ///
252    /// Returns an error if the input contains `NaN`.
253    ///
254    /// # Panics
255    ///
256    /// Panics if:
257    /// * `from.len() != self.dim()`: Vector to be compressed must have the same
258    ///   dimensionality as the quantizer.
259    /// * `to.vector().len() != self.output_dim()`: Compressed vector must have the same dimensionality
260    ///   as the quantizer.
261    fn compress_into(&self, from: &[T], to: DataMutRef<'_, NBITS>) -> Result<L2Loss, Self::Error> {
262        self.compress::<NBITS, T>(from, to)
263    }
264}
265
266impl<'a, T> CompressInto<&[T], FullQueryMut<'a>> for MinMaxQuantizer
267where
268    T: Copy + Into<f32>,
269{
270    type Error = InputContainsNaN;
271
272    type Output = ();
273
274    /// Compress the input vector `from` into a [`FullQueryMut`] `to`.
275    ///
276    /// This method simply applies the transformation to the input without
277    /// any compression.
278    ///
279    /// # Error
280    ///
281    /// Returns an error if the input contains `NaN`.
282    ///
283    /// # Panics
284    ///
285    /// Panics if:
286    /// * `from.len() != self.dim()`: Vector to be compressed must have the same
287    ///   dimensionality as the quantizer.
288    /// * `to.len() != self.output_dim()`: Compressed vector must have the same dimensionality
289    ///   as the quantizer.
290    fn compress_into(&self, from: &[T], mut to: FullQueryMut<'a>) -> Result<(), Self::Error> {
291        assert_eq!(from.len(), self.dim());
292        assert_eq!(self.output_dim(), to.len());
293
294        // Transform the input vector and return error if it contains NaN
295        let from: Vec<f32> = from.iter().map(|&x| x.into()).collect();
296        if from.iter().any(|x| x.is_nan()) {
297            return Err(InputContainsNaN);
298        }
299
300        // We know vec.len() == self.output_dim() and `from.len() == self.dim`
301        #[allow(clippy::unwrap_used)]
302        self.transform
303            .transform_into(to.vector_mut(), &from, ScopedAllocator::global())
304            .unwrap();
305
306        let norm_squared = to.vector().iter().map(|x| *x * *x).sum::<f32>();
307        let sum = to.vector().iter().sum::<f32>();
308
309        *to.meta_mut() = FullQueryMeta { norm_squared, sum };
310
311        Ok(())
312    }
313}
314
315///////////////////////
316// Distance Functors //
317///////////////////////
318
319macro_rules! impl_functor {
320    ($dist:ident) => {
321        impl AsFunctor<$dist> for MinMaxQuantizer {
322            // no need to do any work here.
323            fn as_functor(&self) -> $dist {
324                $dist
325            }
326        }
327    };
328}
329
330impl_functor!(MinMaxIP);
331impl_functor!(MinMaxL2Squared);
332impl_functor!(MinMaxCosine);
333impl_functor!(MinMaxCosineNormalized);
334
335///////////
336// Tests //
337///////////
338#[cfg(test)]
339mod minmax_quantizer_tests {
340    use std::num::NonZeroUsize;
341
342    use diskann_utils::{Reborrow, ReborrowMut};
343    use diskann_vector::{PureDistanceFunction, distance::SquaredL2};
344    use rand::{
345        SeedableRng,
346        distr::{Distribution, Uniform},
347        rngs::StdRng,
348    };
349
350    use super::*;
351    use crate::{
352        algorithms::transforms::NullTransform,
353        alloc::GlobalAllocator,
354        minmax::vectors::{Data, DataRef, FullQuery, FullQueryMut},
355    };
356
357    fn reconstruct_minmax<const NBITS: usize>(v: DataRef<'_, NBITS>) -> Vec<f32>
358    where
359        Unsigned: Representation<NBITS>,
360    {
361        (0..v.len())
362            .map(|i| {
363                let m = v.meta();
364                v.vector().get(i).unwrap() as f32 * m.a + m.b
365            })
366            .collect()
367    }
368
369    fn test_quantizer_encoding_random<const NBITS: usize>(
370        dim: usize,
371        rng: &mut StdRng,
372        relative_err: f32,
373        scale: f32,
374    ) where
375        Unsigned: Representation<NBITS>,
376        MinMaxQuantizer: for<'a, 'b> CompressInto<&'a [f32], DataMutRef<'b, NBITS>, Output = L2Loss>
377            + for<'a, 'b> CompressInto<&'a [f32], FullQueryMut<'b>, Output = ()>,
378    {
379        let distribution = Uniform::new_inclusive::<f32, f32>(-1.0, 1.0).unwrap();
380
381        let quantizer = MinMaxQuantizer::new(
382            Transform::Null(NullTransform::new(NonZeroUsize::new(dim).unwrap())),
383            Positive::new(scale).unwrap(),
384        );
385
386        assert_eq!(quantizer.dim(), dim);
387
388        let vector: Vec<f32> = distribution.sample_iter(rng).take(dim).collect();
389
390        let mut encoded = Data::new_boxed(dim);
391        let loss = quantizer
392            .compress_into(&*vector, encoded.reborrow_mut())
393            .unwrap();
394
395        let reconstructed = reconstruct_minmax::<NBITS>(encoded.reborrow());
396        assert_eq!(reconstructed.len(), dim);
397
398        let reconstruction_error: f32 = SquaredL2::evaluate(&*vector, &*reconstructed);
399        let norm = vector.iter().map(|x| x * x).sum::<f32>();
400        assert!(
401            (reconstruction_error / norm) <= relative_err,
402            "Expected vector : {:?} to be reconstructed within error {} but instead got : {:?}, with error {} for dim : {}",
403            &vector,
404            relative_err,
405            &reconstructed,
406            reconstruction_error / norm,
407            dim,
408        );
409
410        assert!((loss.as_f32() - reconstruction_error) <= 1e-4);
411
412        let expected_code_sum = (0..dim)
413            .map(|i| encoded.vector().get(i).unwrap() as f32)
414            .sum::<f32>();
415        let code_sum = encoded.reborrow().meta().n / encoded.reborrow().meta().a;
416        assert!(
417            (code_sum - expected_code_sum).abs() <= 2e-5 * (dim as f32),
418            "Encoded vector with dim : {dim} is {:?}, got error : {} for vector : {:?}",
419            encoded.reborrow(),
420            (code_sum - expected_code_sum).abs(),
421            &vector,
422        );
423        let recon_norm_sq = reconstructed.iter().map(|x| x * x).sum::<f32>();
424        assert!((encoded.reborrow().meta().norm_squared - recon_norm_sq).abs() <= 1e-3);
425
426        // FullQuery
427        let mut f = FullQuery::new_in(dim, GlobalAllocator).unwrap();
428        quantizer
429            .compress_into(vector.as_slice(), f.reborrow_mut())
430            .unwrap();
431
432        f.vector()
433            .iter()
434            .enumerate()
435            .zip(vector.iter())
436            .for_each(|((i, x), y)| {
437                assert!(
438                    (*x - *y).abs() < 1e-10,
439                    "Full Query did not compress dimension {i} with value {} correctly, got {} instead.",
440                    *y,
441                    *x,
442                )
443            });
444
445        assert!(
446            (f.meta().norm_squared - norm).abs() < 1e-10,
447            "Full Query norm in meta should be {norm} but instead got {}",
448            f.meta().norm_squared
449        );
450
451        let sum = vector.iter().sum::<f32>();
452        assert!(
453            (f.meta().sum - sum) < 1e-10,
454            "Full Query norm in meta should be {sum} but instead got {}",
455            f.meta().sum
456        );
457    }
458
459    cfg_if::cfg_if! {
460        if #[cfg(miri)] {
461            // The max dim does not need to be as high for `CompensatedVectors` because they
462            // defer their distance function implementation to `BitSlice`, which is more
463            // heavily tested.
464            const TRIALS: usize = 2;
465        } else {
466            const TRIALS: usize = 10;
467        }
468    }
469
470    macro_rules! test_minmax_quantizer_encoding {
471        ($name:ident, $dim:literal, $nbits:literal, $seed:literal, $err:expr) => {
472            #[test]
473            fn $name() {
474                let mut rng = StdRng::seed_from_u64($seed);
475                let scales = [1.0, 1.1, 0.9];
476                for (s, e) in scales.iter().zip($err) {
477                    for d in 10..$dim {
478                        for _ in 0..TRIALS {
479                            test_quantizer_encoding_random::<$nbits>(d, &mut rng, e, *s);
480                        }
481                    }
482                }
483            }
484        };
485    }
486    test_minmax_quantizer_encoding!(
487        test_minmax_encoding_1bit,
488        100,
489        1,
490        0xa32d5658097a1c35,
491        vec![0.5, 0.5, 0.5]
492    );
493    test_minmax_quantizer_encoding!(
494        test_minmax_encoding_2bit,
495        100,
496        2,
497        0xf60c0c8d1aadc126,
498        vec![0.5, 0.5, 0.5]
499    );
500    test_minmax_quantizer_encoding!(
501        test_minmax_encoding_4bit,
502        100,
503        4,
504        0x09fa14c42a9d7d98,
505        vec![1.0e-2, 1.0e-2, 3.0e-2]
506    );
507    test_minmax_quantizer_encoding!(
508        test_minmax_encoding_8bit,
509        100,
510        8,
511        0xaedf3d2a223b7b77,
512        vec![2.0e-3, 2.0e-3, 7.0e-3]
513    );
514
515    macro_rules! expand_to_bitrates {
516        ($name:ident, $func:ident) => {
517            #[test]
518            fn $name() {
519                $func::<1>();
520                $func::<2>();
521                $func::<4>();
522                $func::<8>();
523            }
524        };
525    }
526
527    /// Tests the edge case where min == max but both are non-zero.
528    fn test_all_same_value_vector<const NBITS: usize>()
529    where
530        Unsigned: Representation<NBITS>,
531        MinMaxQuantizer:
532            for<'a, 'b> CompressInto<&'a [f32], DataMutRef<'b, NBITS>, Output = L2Loss>,
533    {
534        let dim = 30;
535        let quantizer = MinMaxQuantizer::new(
536            Transform::Null(NullTransform::new(NonZeroUsize::new(dim).unwrap())),
537            Positive::new(1.0).unwrap(),
538        );
539        let constant_value = 42.5f32;
540        let vector = vec![constant_value; dim];
541
542        let mut encoded = Data::new_boxed(dim);
543        let result = quantizer.compress_into(&vector, encoded.reborrow_mut());
544
545        assert!(
546            result.is_ok(),
547            "Constant-value vector should compress successfully"
548        );
549
550        assert!(result.unwrap().as_f32().abs() <= 1e-6);
551
552        // Reconstruction should yield the original constant value (approximately)
553        let reconstructed = reconstruct_minmax(encoded.reborrow());
554        for &val in &reconstructed {
555            assert!(
556                (val - constant_value).abs() < 1e-3,
557                "Reconstructed value {} should be close to original {}. Compressed vector is {:?}",
558                val,
559                constant_value,
560                encoded.meta(),
561            );
562        }
563    }
564
565    /// This tests boundary conditions in the quantization logic.
566    fn test_two_distinct_values<const NBITS: usize>()
567    where
568        Unsigned: Representation<NBITS>,
569        MinMaxQuantizer:
570            for<'a, 'b> CompressInto<&'a [f32], DataMutRef<'b, NBITS>, Output = L2Loss>,
571    {
572        let dim = 20;
573        let quantizer = MinMaxQuantizer::new(
574            Transform::Null(NullTransform::new(NonZeroUsize::new(dim).unwrap())),
575            Positive::new(1.0).unwrap(),
576        );
577
578        let val1 = -10.0f32;
579        let val2 = 15.0f32;
580        let mut vector = vec![val1; dim];
581        // Make half the vector the second value
582        for i in vector.iter_mut().skip(dim) {
583            *i = val2;
584        }
585
586        let mut encoded = Data::new_boxed(dim);
587        let result = quantizer.compress_into(&vector, encoded.reborrow_mut());
588
589        assert!(
590            result.is_ok(),
591            "Two-value vector should compress successfully"
592        );
593
594        assert!(result.unwrap().as_f32().abs() <= 1e-6);
595
596        // Verify that only two distinct codes are used
597        let mut codes_used = std::collections::HashSet::new();
598        for i in 0..dim {
599            codes_used.insert(encoded.vector().get(i).unwrap());
600        }
601
602        // For most bit widths, we should see exactly 2 codes (min and max of domain)
603        if NBITS > 1 {
604            assert!(
605                codes_used.len() <= 2,
606                "Should use at most 2 distinct codes for 2-value input, but used: {:?}",
607                codes_used
608            );
609        }
610
611        // Verify reconstruction maintains the two-value structure approximately
612        let reconstructed = reconstruct_minmax(encoded.reborrow());
613        for ((i, val), v) in reconstructed.into_iter().enumerate().zip(&vector) {
614            // Round to nearest 0.1 to account for quantization error
615            assert!(
616                (val - v).abs() < 1e-4,
617                "Reconstructed value in dim : {i} is {val}, when it should be {v}."
618            );
619        }
620    }
621
622    /// Verifies that NaN values in the input cause the expected error but
623    /// dimension in meta is correctly set.
624    fn test_nan_input_error<const NBITS: usize>()
625    where
626        Unsigned: Representation<NBITS>,
627        MinMaxQuantizer:
628            for<'a, 'b> CompressInto<&'a [f32], DataMutRef<'b, NBITS>, Output = L2Loss>,
629    {
630        let dim = 100;
631        let quantizer = MinMaxQuantizer::new(
632            Transform::Null(NullTransform::new(NonZeroUsize::new(dim).unwrap())),
633            Positive::new(1.0).unwrap(),
634        );
635
636        // Test vector with NaN in the middle.
637        let mut vector_nan = vec![1.0f32; dim];
638        vector_nan[33] = f32::NAN;
639        let mut encoded = Data::new_boxed(dim);
640        let result = quantizer.compress_into(&vector_nan, encoded.reborrow_mut());
641        assert!(result.is_err(), "Vector with NaN should cause an error");
642
643        let meta = encoded.meta();
644        assert_eq!(meta.dim as usize, dim);
645    }
646
647    expand_to_bitrates!(all_same_values_vector, test_all_same_value_vector);
648    expand_to_bitrates!(two_distinct_values, test_two_distinct_values);
649    expand_to_bitrates!(nan_input_error, test_nan_input_error);
650
651    /// Verifies that providing a vector with wrong dimensionality causes a panic.
652    #[test]
653    #[should_panic(expected = "assertion `left == right` failed\n  left: 15\n right: 10")]
654    fn test_dimension_mismatch_panic()
655    where
656        Unsigned: Representation<8>,
657        MinMaxQuantizer: for<'a, 'b> CompressInto<&'a [f32], DataMutRef<'b, 8>, Output = L2Loss>,
658    {
659        let expected_dim = 10;
660        let quantizer = MinMaxQuantizer::new(
661            Transform::Null(NullTransform::new(NonZeroUsize::new(expected_dim).unwrap())),
662            Positive::new(1.0).unwrap(),
663        );
664
665        // Provide vector with wrong dimension
666        let wrong_vector = vec![1.0f32; expected_dim + 5]; // Too many dimensions
667        let mut encoded = Data::new_boxed(expected_dim);
668
669        // This should panic due to assertion in compress_into
670        let _ = quantizer.compress_into(&wrong_vector, encoded.reborrow_mut());
671    }
672}
diskann_quantization/minmax/quantizer.rs

diskann_quantization/minmax/
quantizer.rs
