379 lines
11 KiB
Rust
379 lines
11 KiB
Rust
//! Learning Mechanism Performance Benchmarks
|
|
//!
|
|
//! Benchmarks for MicroLoRA, SONA, and adaptive learning.
|
|
//! Focus on parameter efficiency and training speed.
|
|
//!
|
|
//! Run with: cargo bench --bench learning_performance
|
|
|
|
use criterion::{criterion_group, criterion_main, Criterion, BenchmarkId, Throughput};
|
|
|
|
/// Generate random f32 vector
|
|
fn random_vector(dim: usize, seed: u64) -> Vec<f32> {
|
|
(0..dim)
|
|
.map(|i| {
|
|
let x = ((seed.wrapping_mul(i as u64 + 1).wrapping_mul(0x5DEECE66D)) % 1000) as f32;
|
|
(x / 500.0) - 1.0
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn bench_micro_lora(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("micro_lora");
|
|
|
|
// Test different ranks and dimensions
|
|
for (dim, rank) in [(64, 4), (128, 8), (256, 16), (512, 32)].iter() {
|
|
let input = random_vector(*dim, 42);
|
|
let gradients = random_vector(*dim, 123);
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("forward", format!("d{}_r{}", dim, rank)),
|
|
&(&input, dim, rank),
|
|
|b, (inp, d, r)| {
|
|
// TODO: When MicroLoRA is implemented:
|
|
// let lora = MicroLoRA::new(*d, *r);
|
|
|
|
b.iter(|| {
|
|
// Placeholder: simulate LoRA forward pass
|
|
// output = input + B @ A @ input
|
|
|
|
// A: dim x rank projection
|
|
let mut projected = vec![0.0f32; **r];
|
|
for i in 0..**r {
|
|
for j in 0..**d {
|
|
projected[i] += inp[j] * 0.01;
|
|
}
|
|
}
|
|
|
|
// B: rank x dim projection
|
|
let mut output = vec![0.0f32; **d];
|
|
for i in 0..**d {
|
|
for j in 0..**r {
|
|
output[i] += projected[j] * 0.01;
|
|
}
|
|
}
|
|
|
|
// Add residual
|
|
for (o, i) in output.iter_mut().zip(inp.iter()) {
|
|
*o += i;
|
|
}
|
|
|
|
output
|
|
});
|
|
},
|
|
);
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("backward", format!("d{}_r{}", dim, rank)),
|
|
&(&gradients, dim, rank),
|
|
|b, (grad, d, r)| {
|
|
b.iter(|| {
|
|
// Placeholder: simulate LoRA backward pass
|
|
// Compute gradients for A and B
|
|
|
|
let grad_a = vec![vec![0.01f32; **r]; **d];
|
|
let grad_b = vec![vec![0.01f32; **d]; **r];
|
|
|
|
(grad_a, grad_b)
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_sona_adaptation(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("sona");
|
|
|
|
let input_dim = 64;
|
|
let hidden_dim = 128;
|
|
let output_dim = 32;
|
|
|
|
let input = random_vector(input_dim, 42);
|
|
let target = random_vector(output_dim, 123);
|
|
|
|
group.bench_function("forward", |b| {
|
|
// TODO: When SONA is implemented:
|
|
// let sona = SONA::new(input_dim, hidden_dim, output_dim);
|
|
|
|
b.iter(|| {
|
|
// Placeholder: simulate SONA forward
|
|
let mut hidden = vec![0.0f32; hidden_dim];
|
|
for i in 0..hidden_dim {
|
|
for j in 0..input_dim {
|
|
hidden[i] += input[j] * 0.01;
|
|
}
|
|
hidden[i] = hidden[i].tanh();
|
|
}
|
|
|
|
let mut output = vec![0.0f32; output_dim];
|
|
for i in 0..output_dim {
|
|
for j in 0..hidden_dim {
|
|
output[i] += hidden[j] * 0.01;
|
|
}
|
|
}
|
|
output
|
|
});
|
|
});
|
|
|
|
group.bench_function("adapt_architecture", |b| {
|
|
b.iter(|| {
|
|
// Placeholder: simulate architecture adaptation
|
|
// Analyze activation statistics and prune/grow neurons
|
|
|
|
let neuron_activities = random_vector(hidden_dim, 456);
|
|
let threshold = 0.1;
|
|
|
|
let active_count: usize = neuron_activities.iter()
|
|
.filter(|&a| a.abs() > threshold)
|
|
.count();
|
|
|
|
active_count
|
|
});
|
|
});
|
|
|
|
group.bench_function("prune_neurons", |b| {
|
|
b.iter(|| {
|
|
// Placeholder: simulate neuron pruning
|
|
let neuron_importance = random_vector(hidden_dim, 789);
|
|
let prune_threshold = 0.05;
|
|
|
|
let kept_indices: Vec<usize> = neuron_importance.iter()
|
|
.enumerate()
|
|
.filter(|(_, &imp)| imp.abs() > prune_threshold)
|
|
.map(|(i, _)| i)
|
|
.collect();
|
|
|
|
kept_indices
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_online_learning(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("online_learning");
|
|
|
|
let dim = 64;
|
|
let num_samples = 100;
|
|
|
|
// Generate training samples
|
|
let samples: Vec<(Vec<f32>, Vec<f32>)> = (0..num_samples)
|
|
.map(|i| {
|
|
let input = random_vector(dim, i as u64);
|
|
let target = random_vector(dim, (i + 1000) as u64);
|
|
(input, target)
|
|
})
|
|
.collect();
|
|
|
|
group.throughput(Throughput::Elements(num_samples as u64));
|
|
|
|
group.bench_function("single_sample_update", |b| {
|
|
b.iter(|| {
|
|
// Placeholder: simulate single-sample SGD update
|
|
let (input, target) = &samples[0];
|
|
let learning_rate = 0.01;
|
|
|
|
// Forward
|
|
let mut output = vec![0.0f32; dim];
|
|
for i in 0..dim {
|
|
for j in 0..dim {
|
|
output[i] += input[j] * 0.01;
|
|
}
|
|
}
|
|
|
|
// Compute gradients
|
|
let mut gradients = vec![0.0f32; dim];
|
|
for i in 0..dim {
|
|
gradients[i] = 2.0 * (output[i] - target[i]);
|
|
}
|
|
|
|
// Update (simulated)
|
|
let weight_update: f32 = gradients.iter().map(|g| g * learning_rate).sum();
|
|
|
|
weight_update
|
|
});
|
|
});
|
|
|
|
group.bench_function("batch_update_100", |b| {
|
|
b.iter(|| {
|
|
// Placeholder: simulate batch update
|
|
let mut total_gradient = vec![0.0f32; dim];
|
|
|
|
for (input, target) in &samples {
|
|
// Forward
|
|
let mut output = vec![0.0f32; dim];
|
|
for i in 0..dim {
|
|
for j in 0..dim {
|
|
output[i] += input[j] * 0.01;
|
|
}
|
|
}
|
|
|
|
// Accumulate gradients
|
|
for i in 0..dim {
|
|
total_gradient[i] += 2.0 * (output[i] - target[i]) / (num_samples as f32);
|
|
}
|
|
}
|
|
|
|
total_gradient
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_experience_replay(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("experience_replay");
|
|
|
|
let dim = 64;
|
|
let buffer_size = 10000;
|
|
let batch_size = 32;
|
|
|
|
// Pre-fill buffer
|
|
let buffer: Vec<(Vec<f32>, Vec<f32>)> = (0..buffer_size)
|
|
.map(|i| {
|
|
(random_vector(dim, i as u64), random_vector(dim, (i + 10000) as u64))
|
|
})
|
|
.collect();
|
|
|
|
group.bench_function("uniform_sampling", |b| {
|
|
b.iter(|| {
|
|
// Uniform random sampling
|
|
let mut batch = Vec::with_capacity(batch_size);
|
|
for i in 0..batch_size {
|
|
let idx = (i * 313 + 7) % buffer_size; // Pseudo-random
|
|
batch.push(&buffer[idx]);
|
|
}
|
|
batch
|
|
});
|
|
});
|
|
|
|
group.bench_function("prioritized_sampling", |b| {
|
|
// Priorities (simulated)
|
|
let priorities: Vec<f32> = (0..buffer_size)
|
|
.map(|i| 1.0 + (i as f32 / buffer_size as f32))
|
|
.collect();
|
|
let total_priority: f32 = priorities.iter().sum();
|
|
|
|
b.iter(|| {
|
|
// Prioritized sampling (simplified)
|
|
let mut batch = Vec::with_capacity(batch_size);
|
|
let segment = total_priority / batch_size as f32;
|
|
|
|
for i in 0..batch_size {
|
|
let target = (i as f32 + 0.5) * segment;
|
|
let mut cumsum = 0.0;
|
|
for (idx, &p) in priorities.iter().enumerate() {
|
|
cumsum += p;
|
|
if cumsum >= target {
|
|
batch.push(&buffer[idx]);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
batch
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_meta_learning(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("meta_learning");
|
|
|
|
let dim = 32;
|
|
let num_tasks = 5;
|
|
let shots_per_task = 5;
|
|
|
|
// Generate task data
|
|
let tasks: Vec<Vec<(Vec<f32>, Vec<f32>)>> = (0..num_tasks)
|
|
.map(|t| {
|
|
(0..shots_per_task)
|
|
.map(|s| {
|
|
let input = random_vector(dim, (t * 100 + s) as u64);
|
|
let target = random_vector(dim, (t * 100 + s + 50) as u64);
|
|
(input, target)
|
|
})
|
|
.collect()
|
|
})
|
|
.collect();
|
|
|
|
group.bench_function("inner_loop_adaptation", |b| {
|
|
b.iter(|| {
|
|
// MAML-style inner loop (simplified)
|
|
let task = &tasks[0];
|
|
let inner_lr = 0.01;
|
|
let inner_steps = 5;
|
|
|
|
let mut adapted_weights = vec![0.0f32; dim * dim];
|
|
|
|
for _ in 0..inner_steps {
|
|
let mut gradient = vec![0.0f32; dim * dim];
|
|
|
|
for (input, target) in task {
|
|
// Forward
|
|
let mut output = vec![0.0f32; dim];
|
|
for i in 0..dim {
|
|
for j in 0..dim {
|
|
output[i] += input[j] * adapted_weights[i * dim + j];
|
|
}
|
|
}
|
|
|
|
// Backward
|
|
for i in 0..dim {
|
|
for j in 0..dim {
|
|
gradient[i * dim + j] += 2.0 * (output[i] - target[i]) * input[j];
|
|
}
|
|
}
|
|
}
|
|
|
|
// Update
|
|
for (w, g) in adapted_weights.iter_mut().zip(gradient.iter()) {
|
|
*w -= inner_lr * g / (shots_per_task as f32);
|
|
}
|
|
}
|
|
|
|
adapted_weights
|
|
});
|
|
});
|
|
|
|
group.bench_function("outer_loop_update", |b| {
|
|
b.iter(|| {
|
|
// Meta-gradient computation (simplified)
|
|
let outer_lr = 0.001;
|
|
let mut meta_gradient = vec![0.0f32; dim * dim];
|
|
|
|
for task in &tasks {
|
|
// Simulate adapted performance gradient
|
|
for (input, target) in task {
|
|
for i in 0..dim {
|
|
for j in 0..dim {
|
|
meta_gradient[i * dim + j] += input[j] * target[i] * 0.001;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Scale by learning rate
|
|
for g in &mut meta_gradient {
|
|
*g *= outer_lr / (num_tasks as f32);
|
|
}
|
|
|
|
meta_gradient
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
criterion_group!(
|
|
benches,
|
|
bench_micro_lora,
|
|
bench_sona_adaptation,
|
|
bench_online_learning,
|
|
bench_experience_replay,
|
|
bench_meta_learning
|
|
);
|
|
|
|
criterion_main!(benches);
|