turboquant 0.1.1

Implementation of Google's TurboQuant algorithm for vector quantization
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

use turboquant::turboquant_mse::TurboQuantMSE;
use turboquant::turboquant_prod::TurboQuantProd;
use turboquant::utils::{inner_product, normalize};

fn main() {
    println!("=== TurboQuant Demo ===\n");

    let dim = 128;
    let seed = 42u64;

    // Generate a random unit vector
    use rand::SeedableRng;
    use rand_distr::{Distribution, Normal};
    let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
    let normal = Normal::new(0.0, 1.0).unwrap();
    let raw: Vec<f64> = (0..dim).map(|_| normal.sample(&mut rng)).collect();
    let x = normalize(&raw).unwrap();

    let query_raw: Vec<f64> = (0..dim).map(|_| normal.sample(&mut rng)).collect();
    let query = normalize(&query_raw).unwrap();

    let true_ip = inner_product(&x, &query);

    println!("Vector dimension: {}", dim);
    println!("True inner product: {:.6}\n", true_ip);

    println!("--- TurboQuantMSE ---");
    for bits in [2u8, 3, 4] {
        let tq = TurboQuantMSE::new(dim, bits, seed).unwrap();
        let q = tq.quantize(&x).unwrap();
        let mse = tq.actual_mse(&x).unwrap();
        let bound = tq.distortion_bound();
        println!(
            "  {}-bit: MSE={:.6} (bound={:.6}), ratio={:.1}x",
            bits,
            mse,
            bound,
            q.compression_ratio()
        );
    }

    println!("\n--- TurboQuantProd ---");
    for bits in [3u8, 4] {
        let tq = TurboQuantProd::new(dim, bits, seed).unwrap();
        let q = tq.quantize(&x).unwrap();
        let est_ip = tq.estimate_inner_product(&q, &query).unwrap();
        let err = (true_ip - est_ip).abs();
        println!(
            "  {}-bit: estimated IP={:.6}, error={:.6}, ratio={:.1}x",
            bits,
            est_ip,
            err,
            q.compression_ratio()
        );
    }
}