ruvector-math 2.0.6

Advanced mathematics for next-gen vector search: Optimal Transport, Information Geometry, Product Manifolds
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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

//! Benchmarks for information geometry operations

use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use rand::prelude::*;
use rand_distr::StandardNormal;
use ruvector_math::information_geometry::{FisherInformation, KFACApproximation, NaturalGradient};

fn generate_gradients(n: usize, dim: usize, seed: u64) -> Vec<Vec<f64>> {
    let mut rng = StdRng::seed_from_u64(seed);
    (0..n)
        .map(|_| (0..dim).map(|_| rng.sample(StandardNormal)).collect())
        .collect()
}

fn bench_fisher_information(c: &mut Criterion) {
    let mut group = c.benchmark_group("fisher_information");

    for dim in [32, 64, 128, 256] {
        let samples = 100;
        let gradients = generate_gradients(samples, dim, 42);

        group.throughput(Throughput::Elements((samples * dim) as u64));

        // Diagonal FIM (fast)
        let fisher = FisherInformation::new();
        group.bench_with_input(
            BenchmarkId::new("diagonal_fim", dim),
            &gradients,
            |b, grads| {
                b.iter(|| fisher.diagonal_fim(black_box(grads)));
            },
        );

        // Full FIM (slower but more accurate)
        if dim <= 128 {
            group.bench_with_input(
                BenchmarkId::new("empirical_fim", dim),
                &gradients,
                |b, grads| {
                    b.iter(|| fisher.empirical_fim(black_box(grads)));
                },
            );
        }
    }

    group.finish();
}

fn bench_natural_gradient(c: &mut Criterion) {
    let mut group = c.benchmark_group("natural_gradient");

    for dim in [32, 64, 128] {
        let samples = 50;
        let gradients = generate_gradients(samples, dim, 42);
        let gradient = gradients[0].clone();

        group.throughput(Throughput::Elements(dim as u64));

        // Diagonal natural gradient (fast)
        let mut ng = NaturalGradient::new(0.01).with_diagonal(true);
        group.bench_with_input(
            BenchmarkId::new("diagonal_step", dim),
            &(&gradient, &gradients),
            |b, (g, gs)| {
                b.iter(|| {
                    ng.reset();
                    ng.step(black_box(g), Some(black_box(gs)))
                });
            },
        );
    }

    group.finish();
}

fn bench_kfac(c: &mut Criterion) {
    let mut group = c.benchmark_group("kfac");

    for (input_dim, output_dim) in [(32, 16), (64, 32), (128, 64)] {
        let batch_size = 32;
        let mut rng = StdRng::seed_from_u64(42);

        let activations: Vec<Vec<f64>> = (0..batch_size)
            .map(|_| (0..input_dim).map(|_| rng.sample(StandardNormal)).collect())
            .collect();

        let gradients: Vec<Vec<f64>> = (0..batch_size)
            .map(|_| {
                (0..output_dim)
                    .map(|_| rng.sample(StandardNormal))
                    .collect()
            })
            .collect();

        let weight_grad: Vec<Vec<f64>> = (0..output_dim)
            .map(|_| (0..input_dim).map(|_| rng.sample(StandardNormal)).collect())
            .collect();

        group.throughput(Throughput::Elements((input_dim * output_dim) as u64));

        // K-FAC update
        let mut kfac =
            ruvector_math::information_geometry::KFACApproximation::new(&[(input_dim, output_dim)]);

        group.bench_function(
            BenchmarkId::new("kfac_update", format!("{}x{}", input_dim, output_dim)),
            |b| {
                b.iter(|| kfac.update_layer(0, black_box(&activations), black_box(&gradients)));
            },
        );

        // K-FAC natural gradient
        kfac.update_layer(0, &activations, &gradients).unwrap();

        group.bench_function(
            BenchmarkId::new("kfac_nat_grad", format!("{}x{}", input_dim, output_dim)),
            |b| {
                b.iter(|| kfac.natural_gradient_layer(0, black_box(&weight_grad)));
            },
        );
    }

    group.finish();
}

criterion_group!(
    benches,
    bench_fisher_information,
    bench_natural_gradient,
    bench_kfac,
);
criterion_main!(benches);