diskann_quantization/minmax/
mod.rs

1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6//! # MinMax Quantization
7//!
8//! MinMax quantization provides memory-efficient vector compression by converting
9//! floating-point values to small n-bit integers on a per-vector basis.
10//!
11//! ## Core Concept
12//!
13//! Each vector is independently quantized using the formula:
14//! ```math
15//! X' = round((X - s) * (2^n - 1) / c).clamp(0, 2^n - 1)
16//! ```
17//! where `s` is a shift value and `c` is a scaling parameter computed from the
18//! range of values.
19//!
20//! For most bit widths (>1), given a positive scaling parameter `grid_scale : f32`,
21//! these are computed as:
22//! ```math
23//! - m = (max_i X[i] + min_i X[i]) / 2.0
24//! - w = max_i X[i] - min_i X[i]
25//!
26//! - s = m - w * grid_scale
27//! - c = 2 * w * grid_scale
28//! ```
29//! For 1-bit quantization, to avoid outliers, `s` and `c` are derived differently:
30//!   i) Values are first split into two groups: those below and above the mean.
31//!  ii) `s` is the average of values below the mean.
32//! iii) `c` is the difference between the average of values above the mean and `s`.
33//!
34//! This encoding is similar to scalar quantization, but, since both 's' and 'c'
35//! are computed on a per-vector basis, this allows this quantization mechanism
36//! to be applied in a **streaming setting**; making it qualitatively different
37//! than scalar quantization.
38//!
39//! ## Module Components
40//!
41//! - [`MinMaxQuantizer`]: Handles vector encoding and decoding
42//! - [`Data`]: Stores quantized vectors with compensation parameters
43//! - Distance functions:
44//!   - [`MinMaxIP`]: Inner product distance for quantized vectors.
45//!   - [`MinMaxL2Squared`]: L2 (Euclidean) distance for quantized vectors.
46//!   - [`MinMaxCosine`]: Cosine similarity for quantized vectors.
47//!   - [`MinMaxCosineNormalized`]: Cosine similarity for quantized vectors assuming the
48//!     original full-precision vectors were normalized.
49//!
50//! To reconstruct the original vector, the inverse operation is applied:
51//! ```math
52//! X = X' * c / (2^n - 1) + s
53//! ```
54//!
55//! ## Multi-vector Support
56//!
57//! [`MinMaxMeta`] and [`MinMaxKernel`] support storing and computing distances between
58//! multi-vector representations that use MinMax quantization.
59//!
60//! ```rust
61//! use std::num::NonZeroUsize;
62//! use diskann_quantization::{
63//!     algorithms::{transforms::NullTransform, Transform},
64//!     minmax::{MinMaxMeta, MinMaxQuantizer},
65//!     multi_vector::{
66//!         distance::{Chamfer, MaxSim, QueryMatRef},
67//!         Defaulted, Mat, MatRef, Standard,
68//!     },
69//!     num::Positive,
70//!     CompressInto,
71//! };
72//! use diskann_utils::{Reborrow, ReborrowMut};
73//! use diskann_vector::{DistanceFunctionMut, PureDistanceFunction};
74//!
75//! const NBITS: usize = 8;
76//! let dim = 4;
77//! let num_query_vectors = 2;
78//! let num_doc_vectors = 3;
79//!
80//! // Create a MinMax quantizer (using NullTransform for simplicity)
81//! let quantizer = MinMaxQuantizer::new(
82//!     Transform::Null(NullTransform::new(NonZeroUsize::new(dim).unwrap())),
83//!     Positive::new(1.0).unwrap(),
84//! );
85//!
86//! // Full-precision query multi-vector (2 vectors × 4 dimensions)
87//! let query_data: Vec<f32> = vec![
88//!     1.0, 0.0, 0.0, 0.0,  // query vector 0
89//!     0.0, 1.0, 0.0, 0.0,  // query vector 1
90//! ];
91//! let query_input = MatRef::new(
92//!     Standard::new(num_query_vectors, dim).unwrap(), &query_data
93//! ).unwrap();
94//!
95//! // Full-precision document multi-vector (3 vectors × 4 dimensions)
96//! let doc_data: Vec<f32> = vec![
97//!     0.5, 0.5, 0.0, 0.0,  // doc vector 0
98//!     1.0, 0.0, 0.0, 0.0,  // doc vector 1
99//!     0.0, 0.0, 1.0, 0.0,  // doc vector 2
100//! ];
101//! let doc_input = MatRef::new(
102//!     Standard::new(num_doc_vectors, dim).unwrap(), &doc_data
103//! ).unwrap();
104//!
105//! // Create owned matrices for quantized output using Mat::new
106//! let mut query_out: Mat<MinMaxMeta<NBITS>> =
107//!     Mat::new(MinMaxMeta::new(num_query_vectors, dim), Defaulted).unwrap();
108//! let mut doc_out: Mat<MinMaxMeta<NBITS>> =
109//!     Mat::new(MinMaxMeta::new(num_doc_vectors, dim), Defaulted).unwrap();
110//!
111//! // Quantize both multi-vectors
112//! quantizer.compress_into(query_input, query_out.reborrow_mut()).unwrap();
113//! quantizer.compress_into(doc_input, doc_out.reborrow_mut()).unwrap();
114//!
115//! // Get immutable views via reborrow for distance computation
116//! let query_mv = query_out.reborrow();
117//! let doc_mv = doc_out.reborrow();
118//!
119//! // Compute MaxSim: per-query-vector max similarities
120//! let mut scores = vec![0.0f32; num_query_vectors];
121//! MaxSim::new(&mut scores).evaluate(query_mv.into(), doc_mv);
122//! // scores[i] = min over all doc vectors of distance(query[i], doc[j])
123//!
124//! // Compute Chamfer distance (sum of MaxSim scores)
125//! let chamfer = Chamfer::evaluate(query_mv.into(), doc_mv);
126//! ```
127mod multi;
128mod quantizer;
129mod recompress;
130mod vectors;
131
132/////////////
133// Exports //
134/////////////
135
136pub use multi::{MinMaxKernel, MinMaxMeta};
137pub use quantizer::{L2Loss, MinMaxQuantizer};
138pub use recompress::{RecompressError, Recompressor};
139pub use vectors::{
140    Data, DataMutRef, DataRef, DecompressError, FullQuery, FullQueryMeta, FullQueryMut,
141    FullQueryRef, MetaParseError, MinMaxCompensation, MinMaxCosine, MinMaxCosineNormalized,
142    MinMaxIP, MinMaxL2Squared,
143};
diskann_quantization/minmax/mod.rs

diskann_quantization/minmax/
mod.rs
