d803bfe2b1
git-subtree-dir: vendor/ruvector git-subtree-split: b64c21726f2bb37286d9ee36a7869fef60cc6900
133 lines
4.0 KiB
Rust
133 lines
4.0 KiB
Rust
//! 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);
|