feat: vendor midstream and sublinear-time-solver libraries

Add ruvnet/midstream (AIMDS real-time inference) and
ruvnet/sublinear-time-solver (sublinear optimization algorithms)
as vendored dependencies under vendor/.

Co-Authored-By: claude-flow <ruv@ruv.net>
This commit is contained in:
ruv
2026-03-02 23:32:45 -05:00
parent 14902e6b4e
commit e91bb8a1d5
1600 changed files with 1852646 additions and 0 deletions
@@ -0,0 +1,198 @@
use ndarray::{Array2, Array1, Axis};
use std::f32;
/// Self-attention mechanism for temporal sequences
pub struct TemporalAttention {
d_model: usize,
n_heads: usize,
d_k: usize,
// Query, Key, Value projections for each head
w_q: Vec<Array2<f32>>,
w_k: Vec<Array2<f32>>,
w_v: Vec<Array2<f32>>,
w_o: Array2<f32>,
// Positional encoding
pos_encoding: Array2<f32>,
}
impl TemporalAttention {
pub fn new(d_model: usize, n_heads: usize, max_seq_len: usize) -> Self {
assert_eq!(d_model % n_heads, 0, "d_model must be divisible by n_heads");
let d_k = d_model / n_heads;
use rand::{thread_rng, Rng};
let mut rng = thread_rng();
let scale = (1.0 / d_k as f32).sqrt();
// Initialize projection matrices for each head
let mut w_q = Vec::new();
let mut w_k = Vec::new();
let mut w_v = Vec::new();
for _ in 0..n_heads {
w_q.push(Array2::from_shape_fn((d_k, d_model), |_|
rng.gen::<f32>() * scale - scale/2.0));
w_k.push(Array2::from_shape_fn((d_k, d_model), |_|
rng.gen::<f32>() * scale - scale/2.0));
w_v.push(Array2::from_shape_fn((d_k, d_model), |_|
rng.gen::<f32>() * scale - scale/2.0));
}
let w_o = Array2::from_shape_fn((d_model, d_model), |_|
rng.gen::<f32>() * scale - scale/2.0);
// Create sinusoidal positional encoding
let pos_encoding = Self::create_positional_encoding(max_seq_len, d_model);
Self {
d_model,
n_heads,
d_k,
w_q,
w_k,
w_v,
w_o,
pos_encoding,
}
}
fn create_positional_encoding(max_len: usize, d_model: usize) -> Array2<f32> {
let mut encoding = Array2::zeros((max_len, d_model));
for pos in 0..max_len {
for i in 0..d_model/2 {
let angle = pos as f32 / (10000.0_f32.powf(2.0 * i as f32 / d_model as f32));
encoding[[pos, 2*i]] = angle.sin();
encoding[[pos, 2*i + 1]] = angle.cos();
}
}
encoding
}
/// Scaled dot-product attention
fn attention(&self, q: &Array2<f32>, k: &Array2<f32>, v: &Array2<f32>) -> Array2<f32> {
let d_k_sqrt = (self.d_k as f32).sqrt();
// Compute attention scores: Q @ K^T / sqrt(d_k)
let scores = q.dot(&k.t()) / d_k_sqrt;
// Apply softmax
let exp_scores = scores.mapv(|x| x.exp());
let sum_exp = exp_scores.sum_axis(Axis(1));
let attention_weights = &exp_scores / &sum_exp.insert_axis(Axis(1));
// Apply attention to values
attention_weights.dot(v)
}
/// Multi-head attention forward pass
pub fn forward(&self, x: &Array2<f32>) -> Array2<f32> {
let seq_len = x.nrows();
let batch_d = x.ncols();
// Add positional encoding
let pos_slice = self.pos_encoding.slice(s![..seq_len, ..batch_d]);
let x_pos = x + &pos_slice;
let mut head_outputs = Vec::new();
// Process each attention head
for h in 0..self.n_heads {
let q = x_pos.dot(&self.w_q[h].t());
let k = x_pos.dot(&self.w_k[h].t());
let v = x_pos.dot(&self.w_v[h].t());
let head_out = self.attention(&q, &k, &v);
head_outputs.push(head_out);
}
// Concatenate heads
let mut concat = Array2::zeros((seq_len, self.d_model));
for (h, head_out) in head_outputs.iter().enumerate() {
let start = h * self.d_k;
let end = start + self.d_k;
concat.slice_mut(s![.., start..end]).assign(head_out);
}
// Final linear projection
concat.dot(&self.w_o.t())
}
/// Extract temporal features with attention
pub fn extract_features(&self, sequence: &[Vec<f32>]) -> Vec<f32> {
let seq_len = sequence.len();
let feat_dim = sequence[0].len();
// Convert to ndarray
let mut x = Array2::zeros((seq_len, feat_dim));
for (i, features) in sequence.iter().enumerate() {
for (j, &val) in features.iter().enumerate() {
x[[i, j]] = val;
}
}
// Apply attention
let attended = self.forward(&x);
// Global average pooling over time
attended.mean_axis(Axis(0))
.unwrap()
.to_vec()
}
}
/// Causal attention for autoregressive prediction
pub struct CausalAttention {
attention: TemporalAttention,
mask: Array2<bool>,
}
impl CausalAttention {
pub fn new(d_model: usize, n_heads: usize, max_seq_len: usize) -> Self {
let attention = TemporalAttention::new(d_model, n_heads, max_seq_len);
// Create causal mask (lower triangular)
let mut mask = Array2::from_elem((max_seq_len, max_seq_len), false);
for i in 0..max_seq_len {
for j in 0..=i {
mask[[i, j]] = true;
}
}
Self { attention, mask }
}
/// Apply causal masking to attention scores
pub fn forward_causal(&self, x: &Array2<f32>) -> Array2<f32> {
let seq_len = x.nrows();
// Apply standard attention
let output = self.attention.forward(x);
// Apply causal mask (in practice, this would be done inside attention computation)
let mask_slice = self.mask.slice(s![..seq_len, ..seq_len]);
// Masked output
output.masked_fill(&mask_slice.mapv(|b| !b), 0.0)
}
}
// Simplified implementation for masked_fill
trait MaskedFill {
fn masked_fill(&self, mask: &Array2<bool>, value: f32) -> Self;
}
impl MaskedFill for Array2<f32> {
fn masked_fill(&self, mask: &Array2<bool>, value: f32) -> Self {
let mut result = self.clone();
for ((i, j), &m) in mask.indexed_iter() {
if !m {
result[[i, j]] = value;
}
}
result
}
}
use ndarray::s;
@@ -0,0 +1,17 @@
use crate::data::{Sample, to_class};
/// Naive baseline: next = last_in_window; class from that
pub struct Baseline;
impl Baseline {
pub fn predict_reg(samples: &[Sample], window: usize) -> Vec<f32> {
samples.iter().map(|s| s.x[window-1]).collect()
}
pub fn predict_cls(samples: &[Sample], window: usize) -> Vec<usize> {
samples.iter().map(|s| {
let yhat = s.x[window-1];
to_class(yhat)
}).collect()
}
}
@@ -0,0 +1,62 @@
use rand::{rngs::StdRng, SeedableRng, Rng};
use rand_distr::{Normal, Distribution};
#[derive(Clone)]
pub struct Sample {
pub x: Vec<f32>, // windowed features
pub y: f32, // next value or class
}
pub struct Dataset {
pub train: Vec<Sample>,
pub val: Vec<Sample>,
pub test: Vec<Sample>,
}
/// Synthetic temporal process with regime shifts and delays.
/// Tasks:
/// 1) next_value regression
/// 2) future_bucket classification (coarse future state)
pub fn make_synthetic(window: usize, n: usize, seed: u64) -> Dataset {
let mut rng = StdRng::seed_from_u64(seed);
let noise = Normal::new(0.0, 0.3).unwrap();
// latent regime that flips with low prob
let mut regime = 0.0f32;
let mut series: Vec<f32> = Vec::with_capacity(n + window + 10);
let mut last = 0.0f32;
for t in 0..(n + window + 10) {
if rng.gen::<f32>() < 0.02 { regime = if regime == 0.0 { 1.0 } else { 0.0 }; }
let drift = if regime == 0.0 { 0.02 } else { -0.015 };
let val = 0.8 * last + drift + noise.sample(&mut rng) as f32;
series.push(val);
last = val;
// occasional delayed impulse
if t % 37 == 0 { last += 0.9; }
}
let mut rows = Vec::new();
for i in 0..n {
let w = &series[i..i+window];
let y = series[i+window];
let mut x = w.to_vec();
// add simple time features
x.push(((i as f32) % 24.0) / 24.0);
x.push(regime);
rows.push(Sample { x, y });
}
let split1 = (0.7 * rows.len() as f32) as usize;
let split2 = (0.85 * rows.len() as f32) as usize;
Dataset {
train: rows[..split1].to_vec(),
val: rows[split1..split2].to_vec(),
test: rows[split2..].to_vec(),
}
}
pub fn to_class(y: f32) -> usize {
// 3-bucket classification for coarse future state
if y < -0.25 { 0 } else if y > 0.25 { 2 } else { 1 }
}
@@ -0,0 +1,300 @@
use crate::mlp::Mlp;
use crate::mlp_optimized::OptimizedMlp;
use crate::mlp_ultra::UltraMlp;
use crate::mlp_classifier::ClassifierMlp;
use rayon::prelude::*;
use std::sync::Arc;
/// Ensemble model combining multiple predictors with weighted voting
pub struct EnsembleModel {
models: Vec<ModelType>,
weights: Vec<f32>,
use_adaptive_weights: bool,
}
enum ModelType {
Simple(Mlp),
Optimized(OptimizedMlp),
Ultra(UltraMlp),
Classifier(ClassifierMlp),
}
impl EnsembleModel {
pub fn new() -> Self {
Self {
models: Vec::new(),
weights: Vec::new(),
use_adaptive_weights: true,
}
}
pub fn add_model_simple(&mut self, input: usize, hidden: usize, output: usize) {
self.models.push(ModelType::Simple(Mlp::new(input, hidden, output)));
self.weights.push(1.0);
}
pub fn add_model_optimized(&mut self, input: usize, hidden: usize, output: usize) {
self.models.push(ModelType::Optimized(OptimizedMlp::new(input, hidden, output)));
self.weights.push(1.0);
}
pub fn add_model_ultra(&mut self, input: usize, hidden: usize, output: usize) {
self.models.push(ModelType::Ultra(UltraMlp::new(input, hidden, output)));
self.weights.push(1.0);
}
pub fn add_model_classifier(&mut self, input: usize, output: usize) {
self.models.push(ModelType::Classifier(ClassifierMlp::new(input, output)));
self.weights.push(1.0);
}
/// Train all models in parallel with different random initializations
pub fn train_ensemble(&mut self, x: &Vec<Vec<f32>>, y: &Vec<f32>,
epochs: usize, lr: f32, val_x: &Vec<Vec<f32>>, val_y: &Vec<usize>) {
use rand::seq::SliceRandom;
use rand::thread_rng;
// Train each model with different data shuffling for diversity
let model_count = self.models.len();
// Parallel training with thread-safe access
let x_arc = Arc::new(x.clone());
let y_arc = Arc::new(y.clone());
// Sequential training (models contain mutable state)
for (i, model) in self.models.iter_mut().enumerate() {
let mut indices: Vec<usize> = (0..x.len()).collect();
let mut rng = thread_rng();
use rand::Rng;
indices.shuffle(&mut rng);
// Bootstrap sampling for diversity
let bootstrap_size = (x.len() as f32 * 0.8) as usize;
let bootstrap_indices: Vec<usize> = (0..bootstrap_size)
.map(|_| indices[rng.gen::<usize>() % indices.len()])
.collect();
let x_bootstrap: Vec<Vec<f32>> = bootstrap_indices.iter()
.map(|&i| x[i].clone())
.collect();
let y_bootstrap: Vec<f32> = bootstrap_indices.iter()
.map(|&i| y[i])
.collect();
// Train based on model type
match model {
ModelType::Simple(ref mut m) => {
m.train_regression(&x_bootstrap, &y_bootstrap, epochs, lr);
}
ModelType::Optimized(ref mut m) => {
m.train_batch(&x_bootstrap, &y_bootstrap, epochs, lr, 32);
}
ModelType::Ultra(ref mut m) => {
m.train_batch_parallel(&x_bootstrap, &y_bootstrap, epochs, lr, 32);
}
ModelType::Classifier(ref mut m) => {
m.train_classification(&x_bootstrap, &y_bootstrap, epochs, 32);
}
}
println!("Trained model {}/{}", i + 1, model_count);
}
// Update weights based on validation performance
if self.use_adaptive_weights && !val_x.is_empty() {
self.update_weights(val_x, val_y);
}
}
/// Update ensemble weights based on validation accuracy
fn update_weights(&mut self, val_x: &Vec<Vec<f32>>, val_y: &Vec<usize>) {
let accuracies: Vec<f32> = self.models.iter_mut().map(|model| {
let predictions = match model {
ModelType::Simple(ref m) => m.predict_cls3(val_x),
ModelType::Optimized(ref m) => m.predict_cls3(val_x),
ModelType::Ultra(ref mut m) => m.predict_cls3_parallel(val_x),
ModelType::Classifier(ref mut m) => m.predict_cls3(val_x),
};
let correct = predictions.iter().zip(val_y.iter())
.filter(|(p, y)| p == y)
.count();
correct as f32 / val_y.len() as f32
}).collect();
// Convert accuracies to weights (squared for emphasis)
let total_acc: f32 = accuracies.iter().map(|&a| a * a).sum();
if total_acc > 0.0 {
self.weights = accuracies.iter()
.map(|&a| (a * a) / total_acc)
.collect();
}
println!("Ensemble weights: {:?}", self.weights);
}
/// Predict using weighted voting
pub fn predict_ensemble(&mut self, x: &[Vec<f32>]) -> Vec<usize> {
let predictions: Vec<Vec<usize>> = self.models.iter_mut().map(|model| {
match model {
ModelType::Simple(ref m) => m.predict_cls3(x),
ModelType::Optimized(ref m) => m.predict_cls3(x),
ModelType::Ultra(ref mut m) => m.predict_cls3_parallel(x),
ModelType::Classifier(ref mut m) => m.predict_cls3(x),
}
}).collect();
// Weighted voting for each sample
(0..x.len()).map(|i| {
let mut votes = vec![0.0; 3]; // 3 classes
for (model_idx, model_preds) in predictions.iter().enumerate() {
let pred = model_preds[i];
if pred < 3 {
votes[pred] += self.weights[model_idx];
}
}
// Return class with highest weighted votes
votes.iter()
.enumerate()
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
.map(|(idx, _)| idx)
.unwrap_or(1) // Default to middle class
}).collect()
}
/// Fast ensemble prediction with majority voting (no weights)
pub fn predict_fast(&mut self, x: &[Vec<f32>]) -> Vec<usize> {
let predictions: Vec<Vec<usize>> = self.models.par_iter_mut().map(|model| {
match model {
ModelType::Simple(ref m) => m.predict_cls3(x),
ModelType::Optimized(ref m) => m.predict_cls3(x),
ModelType::Ultra(ref mut m) => m.predict_cls3_parallel(x),
ModelType::Classifier(ref mut m) => m.predict_cls3(x),
}
}).collect();
// Simple majority voting
(0..x.len()).into_par_iter().map(|i| {
let mut votes = [0u32; 3];
for model_preds in &predictions {
if model_preds[i] < 3 {
votes[model_preds[i]] += 1;
}
}
votes.iter()
.enumerate()
.max_by_key(|&(_, &v)| v)
.map(|(idx, _)| idx)
.unwrap_or(1)
}).collect()
}
}
/// Boosted ensemble using AdaBoost-style weighting
pub struct BoostedEnsemble {
weak_learners: Vec<Mlp>,
alphas: Vec<f32>,
input_dim: usize,
hidden_dim: usize,
}
impl BoostedEnsemble {
pub fn new(input: usize, hidden: usize) -> Self {
Self {
weak_learners: Vec::new(),
alphas: Vec::new(),
input_dim: input,
hidden_dim: hidden,
}
}
pub fn train_boosted(&mut self, x: &Vec<Vec<f32>>, y: &Vec<usize>,
n_estimators: usize, epochs: usize) {
let n_samples = x.len();
let mut weights = vec![1.0 / n_samples as f32; n_samples];
for t in 0..n_estimators {
// Train weak learner on weighted data
let mut learner = Mlp::new(self.input_dim, self.hidden_dim, 3);
// Convert classes to continuous for regression
let y_cont: Vec<f32> = y.iter().map(|&c| c as f32 - 1.0).collect();
// Weight samples by resampling
let mut weighted_x = Vec::new();
let mut weighted_y = Vec::new();
use rand::thread_rng;
use rand::distributions::{Distribution, WeightedIndex};
let mut rng = thread_rng();
let dist = WeightedIndex::new(&weights).unwrap();
for _ in 0..n_samples {
let idx = dist.sample(&mut rng);
weighted_x.push(x[idx].clone());
weighted_y.push(y_cont[idx]);
}
learner.train_regression(&weighted_x, &weighted_y, epochs, 0.01);
// Calculate error
let predictions = learner.predict_cls3(x);
let mut error = 0.0;
for i in 0..n_samples {
if predictions[i] != y[i] {
error += weights[i];
}
}
// Avoid division by zero
if error >= 0.5 {
break; // Stop if not better than random
}
// Calculate alpha
let alpha = 0.5 * ((1.0 - error) / error.max(1e-10)).ln();
// Update weights
for i in 0..n_samples {
if predictions[i] != y[i] {
weights[i] *= (alpha.exp());
} else {
weights[i] *= ((-alpha).exp());
}
}
// Normalize weights
let sum: f32 = weights.iter().sum();
weights.iter_mut().for_each(|w| *w /= sum);
self.weak_learners.push(learner);
self.alphas.push(alpha);
println!("Boosting round {}/{}, error: {:.4}", t + 1, n_estimators, error);
}
}
pub fn predict_boosted(&self, x: &[Vec<f32>]) -> Vec<usize> {
let n_classes = 3;
x.iter().map(|xi| {
let mut class_scores = vec![0.0; n_classes];
for (learner, &alpha) in self.weak_learners.iter().zip(&self.alphas) {
let pred = learner.predict_cls3(&[xi.clone()])[0];
if pred < n_classes {
class_scores[pred] += alpha;
}
}
class_scores.iter()
.enumerate()
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
.map(|(idx, _)| idx)
.unwrap_or(1)
}).collect()
}
}
@@ -0,0 +1,181 @@
use ndarray::{Array1, Array2};
use rand::{thread_rng, Rng, distributions::Uniform};
use std::f32::consts::PI;
/// Random Fourier Features for kernel approximation
/// Maps input to high-dimensional space where linear methods work well
pub struct FourierFeatures {
// Random projection parameters
omega: Array2<f32>, // Random frequencies (D x d)
b: Array1<f32>, // Random phase shifts
input_dim: usize,
feature_dim: usize,
sigma: f32, // RBF kernel bandwidth
// Linear model on top
weights: Array1<f32>,
bias: f32,
}
impl FourierFeatures {
pub fn new(input_dim: usize, feature_dim: usize, sigma: f32) -> Self {
let mut rng = thread_rng();
let normal = Uniform::new(0.0, 2.0 * PI);
// Sample random frequencies from Gaussian (for RBF kernel)
let scale = 1.0 / sigma;
let omega = Array2::from_shape_fn((feature_dim, input_dim), |_|
rng.gen::<f32>() * scale);
// Random phase shifts
let b = Array1::from_shape_fn(feature_dim, |_| rng.sample(normal));
Self {
omega,
b,
input_dim,
feature_dim,
sigma,
weights: Array1::zeros(feature_dim),
bias: 0.0,
}
}
/// Transform input using random Fourier features
pub fn transform(&self, x: &[f32]) -> Array1<f32> {
let x_arr = Array1::from_vec(x.to_vec());
let projections = self.omega.dot(&x_arr) + &self.b;
// Apply cosine to get Fourier features
let scale = (2.0 / self.feature_dim as f32).sqrt();
projections.mapv(|p| (p.cos() * scale))
}
/// Batch transform
pub fn transform_batch(&self, x: &[Vec<f32>]) -> Array2<f32> {
let n_samples = x.len();
let mut features = Array2::zeros((n_samples, self.feature_dim));
for (i, xi) in x.iter().enumerate() {
let feat = self.transform(xi);
features.row_mut(i).assign(&feat);
}
features
}
/// Train using closed-form ridge regression
pub fn train(&mut self, x: &[Vec<f32>], y: &[f32], lambda: f32) {
let features = self.transform_batch(x);
let n = features.nrows();
// Closed-form solution: w = (X^T X + λI)^{-1} X^T y
let xtx = features.t().dot(&features);
let xty = features.t().dot(&Array1::from_vec(y.to_vec()));
// Add regularization
let mut reg_xtx = xtx + Array2::<f32>::eye(self.feature_dim) * lambda * n as f32;
// Solve (simplified - production would use LAPACK)
self.weights = self.solve_regularized(&reg_xtx, &xty);
// Compute bias
let predictions = features.dot(&self.weights);
let mean_pred = predictions.mean().unwrap();
let mean_y = y.iter().sum::<f32>() / y.len() as f32;
self.bias = mean_y - mean_pred;
}
fn solve_regularized(&self, a: &Array2<f32>, b: &Array1<f32>) -> Array1<f32> {
// Simplified solver using gradient descent
let mut x = Array1::zeros(self.feature_dim);
let lr = 0.01;
for _ in 0..100 {
let grad = a.dot(&x) - b;
x = x - &grad * lr;
}
x
}
pub fn predict(&self, x: &[Vec<f32>]) -> Vec<f32> {
x.iter().map(|xi| {
let features = self.transform(xi);
self.weights.dot(&features) + self.bias
}).collect()
}
pub fn predict_class(&self, x: &[Vec<f32>]) -> Vec<usize> {
self.predict(x).iter().map(|&y| {
if y < -0.25 { 0 }
else if y > 0.25 { 2 }
else { 1 }
}).collect()
}
}
/// Adaptive Fourier Features with frequency learning
pub struct AdaptiveFourierFeatures {
base: FourierFeatures,
frequency_lr: f32,
adapt_frequencies: bool,
}
impl AdaptiveFourierFeatures {
pub fn new(input_dim: usize, feature_dim: usize, sigma: f32) -> Self {
Self {
base: FourierFeatures::new(input_dim, feature_dim, sigma),
frequency_lr: 0.001,
adapt_frequencies: true,
}
}
/// Train with frequency adaptation
pub fn train_adaptive(&mut self, x: &[Vec<f32>], y: &[f32], epochs: usize) {
for epoch in 0..epochs {
// Standard training
self.base.train(x, y, 0.01);
if self.adapt_frequencies && epoch % 10 == 0 {
// Adapt frequencies based on gradient
self.adapt_frequencies_step(x, y);
}
}
}
fn adapt_frequencies_step(&mut self, x: &[Vec<f32>], y: &[f32]) {
// Compute gradient w.r.t frequencies
let mut grad_omega = Array2::zeros((self.base.feature_dim, self.base.input_dim));
for (xi, &yi) in x.iter().zip(y.iter()) {
let x_arr = Array1::from_vec(xi.clone());
let projections = self.base.omega.dot(&x_arr) + &self.base.b;
let features = projections.mapv(|p| p.cos());
let pred = self.base.weights.dot(&features) + self.base.bias;
let error = pred - yi;
// Gradient through cosine
for j in 0..self.base.feature_dim {
let grad_cos = -projections[j].sin();
let grad_j = error * self.base.weights[j] * grad_cos;
for k in 0..self.base.input_dim {
grad_omega[[j, k]] += grad_j * xi[k];
}
}
}
// Update frequencies
self.base.omega = &self.base.omega - &grad_omega * self.frequency_lr;
}
pub fn predict(&self, x: &[Vec<f32>]) -> Vec<f32> {
self.base.predict(x)
}
pub fn predict_class(&self, x: &[Vec<f32>]) -> Vec<usize> {
self.base.predict_class(x)
}
}
@@ -0,0 +1,23 @@
// lib.rs - Expose modules for library use
pub mod data;
pub mod metrics;
pub mod baseline;
pub mod mlp;
pub mod mlp_optimized;
pub mod mlp_ultra;
pub mod mlp_classifier;
pub mod ensemble;
pub mod attention;
pub mod reservoir;
pub mod fourier;
pub mod sparse;
pub mod mlp_avx512;
pub mod quantization;
pub mod mlp_quantized;
#[cfg(feature = "ruv-fann")]
pub mod ruv_fann_impl;
#[cfg(not(feature = "ruv-fann"))]
pub mod ruv_fann_adapter;
@@ -0,0 +1,298 @@
mod data;
mod metrics;
mod baseline;
mod mlp;
mod mlp_optimized;
mod mlp_ultra;
mod mlp_classifier;
mod ensemble;
mod attention;
mod reservoir;
mod fourier;
mod sparse;
mod mlp_avx512;
mod quantization;
mod mlp_quantized;
mod ruv_fann_adapter;
mod ruv_fann_impl;
use clap::{Parser, ValueEnum};
use data::{make_synthetic, to_class};
use metrics::{mse, acc};
use baseline::Baseline;
use mlp::Mlp;
use mlp_optimized::OptimizedMlp;
use mlp_ultra::UltraMlp;
use mlp_classifier::ClassifierMlp;
use ensemble::{EnsembleModel, BoostedEnsemble};
use reservoir::{ReservoirComputer, QuantumReservoir};
use fourier::{FourierFeatures, AdaptiveFourierFeatures};
use sparse::{SparseNetwork, LotteryTicketNetwork};
use mlp_avx512::DynamicAvx512Mlp;
use mlp_quantized::QuantizedMlpBackend;
#[cfg(feature = "ruv-fann")]
use ruv_fann_impl::RuvFannModel;
#[cfg(not(feature = "ruv-fann"))]
use ruv_fann_adapter::ruv_fann_backend::RuvFannModel;
#[derive(Copy, Clone, ValueEnum)]
enum Backend { Mlp, MlpOpt, MlpUltra, MlpAvx512, MlpQuantized, MlpClassifier, Ensemble, Boosted, Reservoir, QuantumReservoir, Fourier, AdaptiveFourier, Sparse, LotteryTicket, RuvFann, Baseline }
#[derive(Parser)]
struct Args {
#[arg(long, default_value_t=32)]
window: usize,
#[arg(long, default_value_t=5000)]
n: usize,
#[arg(long, default_value_t=42)]
seed: u64,
#[arg(long, value_enum, default_value_t=Backend::Mlp)]
backend: Backend,
#[arg(long, default_value_t=64)]
hidden: usize,
#[arg(long, default_value_t=8)]
epochs: usize,
#[arg(long, default_value_t=0.01)]
lr: f32,
#[arg(long, default_value_t=false)]
classify: bool
}
fn main() {
let args = Args::parse();
let ds = make_synthetic(args.window, args.n, args.seed);
// prepare tensors
let to_xy = |v: &Vec<data::Sample>| {
let x: Vec<Vec<f32>> = v.iter().map(|s| s.x.clone()).collect();
let y_reg: Vec<f32> = v.iter().map(|s| s.y).collect();
let y_cls: Vec<usize> = v.iter().map(|s| to_class(s.y)).collect();
(x, y_reg, y_cls)
};
let (xtr, ytr, ytrc) = to_xy(&ds.train);
let (xva, yva, yvac) = to_xy(&ds.val);
let (xte, yte, ytec) = to_xy(&ds.test);
match args.backend {
Backend::Baseline => {
let yhat = Baseline::predict_reg(&ds.test, args.window);
let yhatc = Baseline::predict_cls(&ds.test, args.window);
println!("baseline_mse_test={:.6}", mse(&yte, &yhat));
println!("baseline_acc_test={:.4}", acc(&ytec, &yhatc));
}
Backend::Mlp => {
let mut model = if args.classify { Mlp::new(xtr[0].len(), args.hidden, 3) }
else { Mlp::new(xtr[0].len(), args.hidden, 1) };
if args.classify {
// train via regression to continuous y, then map to buckets
model.train_regression(&xtr, &ytr, args.epochs, args.lr);
let yhat = model.predict_cls3(&xte);
println!("mlp_acc_test={:.4}", acc(&ytec, &yhat));
} else {
model.train_regression(&xtr, &ytr, args.epochs, args.lr);
let yhat = model.predict_reg(&xte);
println!("mlp_mse_val={:.6}", mse(&yva, &model.predict_reg(&xva)));
println!("mlp_mse_test={:.6}", mse(&yte, &yhat));
}
}
Backend::MlpOpt => {
let mut model = if args.classify { OptimizedMlp::new(xtr[0].len(), args.hidden, 3) }
else { OptimizedMlp::new(xtr[0].len(), args.hidden, 1) };
if args.classify {
model.train_batch(&xtr, &ytr, args.epochs, args.lr, 32);
let yhat = model.predict_cls3(&xte);
println!("mlp_opt_acc_test={:.4}", acc(&ytec, &yhat));
} else {
model.train_batch(&xtr, &ytr, args.epochs, args.lr, 32);
let yhat = model.predict_reg(&xte);
println!("mlp_opt_mse_val={:.6}", mse(&yva, &model.predict_reg(&xva)));
println!("mlp_opt_mse_test={:.6}", mse(&yte, &yhat));
}
}
Backend::MlpUltra => {
let mut model = if args.classify { UltraMlp::new(xtr[0].len(), args.hidden, 3) }
else { UltraMlp::new(xtr[0].len(), args.hidden, 1) };
if args.classify {
model.train_batch_parallel(&xtr, &ytr, args.epochs, args.lr, 32);
let yhat = model.predict_cls3_parallel(&xte);
println!("mlp_ultra_acc_test={:.4}", acc(&ytec, &yhat));
} else {
model.train_batch_parallel(&xtr, &ytr, args.epochs, args.lr, 32);
let yhat = model.predict_parallel(&xte);
println!("mlp_ultra_mse_val={:.6}", mse(&yva, &model.predict_parallel(&xva)));
println!("mlp_ultra_mse_test={:.6}", mse(&yte, &yhat));
}
}
Backend::MlpAvx512 => {
let model = DynamicAvx512Mlp::new(xtr[0].len(), args.hidden, if args.classify { 3 } else { 1 });
if args.classify {
// Note: AVX512 model doesn't have train method yet, using predict only
let yhat = model.predict_class(&xte);
println!("mlp_avx512_acc_test={:.4} (inference only)", acc(&ytec, &yhat));
} else {
let yhat = model.predict(&xte);
println!("mlp_avx512_mse_test={:.6} (inference only)", mse(&yte, &yhat));
}
}
Backend::MlpQuantized => {
let mut model = QuantizedMlpBackend::new(xtr[0].len(), args.hidden, if args.classify { 3 } else { 1 });
// Train in FP32
model.train(&xtr, &ytr, args.epochs, args.lr);
// Benchmark INT8 vs FP32
println!("\n=== INT8 Quantization Performance ===");
model.benchmark_inference(&xte[..100.min(xte.len())], 100);
if args.classify {
let yhat = model.predict_class(&xte);
println!("\nmlp_quantized_acc_test={:.4}", acc(&ytec, &yhat));
} else {
// Compare FP32 vs INT8 accuracy
let yhat_fp32 = model.predict_fp32(&xte);
let yhat_int8 = model.predict(&xte);
println!("\nmlp_quantized_mse_test_fp32={:.6}", mse(&yte, &yhat_fp32));
println!("mlp_quantized_mse_test_int8={:.6}", mse(&yte, &yhat_int8));
// Calculate accuracy loss
let mse_fp32 = mse(&yte, &yhat_fp32);
let mse_int8 = mse(&yte, &yhat_int8);
let accuracy_loss = ((mse_int8 - mse_fp32).abs() / mse_fp32) * 100.0;
println!("Quantization accuracy loss: {:.2}%", accuracy_loss);
}
let (orig, quant, ratio) = model.get_compression_stats();
println!("\nModel compression: {} -> {} bytes ({:.2}x)", orig, quant, ratio);
}
Backend::MlpClassifier => {
let mut model = ClassifierMlp::new(xtr[0].len(), 3);
if args.classify {
model.train_classification(&xtr, &ytr, args.epochs, 32);
let yhat = model.predict_cls3(&xte);
let yhat_val = model.predict_cls3(&xva);
println!("mlp_classifier_acc_val={:.4}", acc(&yvac, &yhat_val));
println!("mlp_classifier_acc_test={:.4}", acc(&ytec, &yhat));
} else {
println!("MlpClassifier is for classification only, use --classify flag");
}
}
Backend::Ensemble => {
if args.classify {
let mut ensemble = EnsembleModel::new();
let input_dim = xtr[0].len();
// Add diverse models
ensemble.add_model_simple(input_dim, args.hidden, 3);
ensemble.add_model_optimized(input_dim, args.hidden * 2, 3);
ensemble.add_model_ultra(input_dim, args.hidden, 3);
ensemble.add_model_classifier(input_dim, 3);
ensemble.train_ensemble(&xtr, &ytr, args.epochs, args.lr, &xva, &yvac);
let yhat = ensemble.predict_ensemble(&xte);
println!("ensemble_acc_test={:.4}", acc(&ytec, &yhat));
} else {
println!("Ensemble is for classification only, use --classify flag");
}
}
Backend::Boosted => {
if args.classify {
let mut boosted = BoostedEnsemble::new(xtr[0].len(), args.hidden);
boosted.train_boosted(&xtr, &ytrc, 10, args.epochs / 2);
let yhat = boosted.predict_boosted(&xte);
println!("boosted_acc_test={:.4}", acc(&ytec, &yhat));
} else {
println!("Boosted is for classification only, use --classify flag");
}
}
Backend::Reservoir => {
let mut model = ReservoirComputer::new(xtr[0].len(), 100, 1);
if args.classify {
model.train_ridge(&xtr, &ytr, 0.001);
let yhat = model.predict_class(&xte);
println!("reservoir_acc_test={:.4}", acc(&ytec, &yhat));
} else {
model.train_ridge(&xtr, &ytr, 0.001);
let yhat = model.predict(&xte);
println!("reservoir_mse_test={:.6}", mse(&yte, &yhat));
}
}
Backend::QuantumReservoir => {
let mut model = QuantumReservoir::new(xtr[0].len(), 100, 1);
if args.classify {
model.train(&xtr, &ytr);
let yhat = model.predict_quantum(&xte);
println!("quantum_reservoir_acc_test={:.4}", acc(&ytec, &yhat));
} else {
model.train(&xtr, &ytr);
// For regression, use classical prediction
let yhat = model.classical_reservoir.predict(&xte);
println!("quantum_reservoir_mse_test={:.6}", mse(&yte, &yhat));
}
}
Backend::Fourier => {
let mut model = FourierFeatures::new(xtr[0].len(), 500, 1.0);
if args.classify {
model.train(&xtr, &ytr, 0.001);
let yhat = model.predict_class(&xte);
println!("fourier_acc_test={:.4}", acc(&ytec, &yhat));
} else {
model.train(&xtr, &ytr, 0.001);
let yhat = model.predict(&xte);
println!("fourier_mse_test={:.6}", mse(&yte, &yhat));
}
}
Backend::AdaptiveFourier => {
let mut model = AdaptiveFourierFeatures::new(xtr[0].len(), 500, 1.0);
if args.classify {
model.train_adaptive(&xtr, &ytr, 50);
let yhat = model.predict_class(&xte);
println!("adaptive_fourier_acc_test={:.4}", acc(&ytec, &yhat));
} else {
model.train_adaptive(&xtr, &ytr, 50);
let yhat = model.predict(&xte);
println!("adaptive_fourier_mse_test={:.6}", mse(&yte, &yhat));
}
}
Backend::Sparse => {
let mut model = SparseNetwork::new(xtr[0].len(), args.hidden, if args.classify { 3 } else { 1 }, 0.1);
if args.classify {
model.train(&xtr, &ytr, args.epochs * 10, 0.01);
let yhat = model.predict_class(&xte);
let (active, pruned, sparsity) = model.get_sparsity_stats();
println!("sparse_acc_test={:.4} (active={}, pruned={}, sparsity={:.1}%)",
acc(&ytec, &yhat), active, pruned, sparsity * 100.0);
} else {
model.train(&xtr, &ytr, args.epochs * 10, 0.01);
let yhat = model.predict(&xte);
let (active, pruned, sparsity) = model.get_sparsity_stats();
println!("sparse_mse_test={:.6} (active={}, pruned={}, sparsity={:.1}%)",
mse(&yte, &yhat), active, pruned, sparsity * 100.0);
}
}
Backend::LotteryTicket => {
let mut model = LotteryTicketNetwork::new(xtr[0].len(), args.hidden, if args.classify { 3 } else { 1 });
if args.classify {
model.find_winning_ticket(&xtr, &ytr, 0.2, 5);
let yhat = model.predict_class(&xte);
println!("lottery_ticket_acc_test={:.4}", acc(&ytec, &yhat));
} else {
model.find_winning_ticket(&xtr, &ytr, 0.2, 5);
let yhat = model.predict(&xte);
println!("lottery_ticket_mse_test={:.6}", mse(&yte, &yhat));
}
}
Backend::RuvFann => {
let mut model = if args.classify { RuvFannModel::new(xtr[0].len(), args.hidden, 3) }
else { RuvFannModel::new(xtr[0].len(), args.hidden, 1) };
model.train_regression(&xtr, &ytr, args.epochs, args.lr);
if args.classify {
let yhat = model.predict_cls3(&xte);
println!("ruv_fann_acc_test={:.4}", acc(&ytec, &yhat));
} else {
let yhat = model.predict_reg(&xte);
println!("ruv_fann_mse_test={:.6}", mse(&yte, &yhat));
}
}
}
}
@@ -0,0 +1,8 @@
pub fn mse(y: &[f32], yhat: &[f32]) -> f32 {
y.iter().zip(yhat).map(|(a,b)| (a-b)*(a-b)).sum::<f32>() / y.len() as f32
}
pub fn acc(y: &[usize], yhat: &[usize]) -> f32 {
let c = y.iter().zip(yhat).filter(|(a,b)| a==b).count();
c as f32 / y.len() as f32
}
@@ -0,0 +1,108 @@
use ndarray::{Array2, Array1};
use rand::{thread_rng, Rng};
pub struct Mlp {
w1: Array2<f32>,
b1: Array1<f32>,
w2: Array2<f32>,
b2: Array1<f32>,
output_dim: usize,
}
fn relu(x: &mut Array1<f32>) {
x.iter_mut().for_each(|v| if *v < 0.0 { *v = 0.0 });
}
impl Mlp {
pub fn new(input: usize, hidden: usize, output: usize) -> Self {
let mut rng = thread_rng();
let scale = (2.0 / input as f32).sqrt();
let w1 = Array2::from_shape_fn((hidden, input), |_| rng.gen::<f32>() * scale - scale/2.0);
let b1 = Array1::zeros(hidden);
let w2 = Array2::from_shape_fn((output, hidden), |_| rng.gen::<f32>() * scale - scale/2.0);
let b2 = Array1::zeros(output);
Self { w1, b1, w2, b2, output_dim: output }
}
fn forward(&self, x: &Array1<f32>) -> Array1<f32> {
let mut h = self.w1.dot(x) + &self.b1;
relu(&mut h);
self.w2.dot(&h) + &self.b2
}
pub fn train_regression(&mut self, x: &Vec<Vec<f32>>, y: &Vec<f32>, epochs: usize, lr: f32) {
// Simplified SGD without full backprop for now
for _ in 0..epochs {
for (xi, yi) in x.iter().zip(y) {
let x_arr = Array1::from_vec(xi.clone());
let output = self.forward(&x_arr);
// Only train if we have 1D output for regression
if self.output_dim == 1 {
let err = output[0] - yi;
// Numerical gradient approximation for simplicity
let _eps = 0.0001;
// Update w2 and b2
let mut h = self.w1.dot(&x_arr) + &self.b1;
relu(&mut h);
for i in 0..self.w2.nrows() {
for j in 0..self.w2.ncols() {
self.w2[[i, j]] -= lr * err * h[j] * 0.1;
}
self.b2[i] -= lr * err * 0.1;
}
// Update w1 and b1 with smaller learning rate
for i in 0..self.w1.nrows() {
for j in 0..self.w1.ncols() {
self.w1[[i, j]] -= lr * err * x_arr[j] * 0.01;
}
self.b1[i] -= lr * err * 0.01;
}
}
}
}
}
pub fn predict_reg(&self, x: &[Vec<f32>]) -> Vec<f32> {
x.iter().map(|xi| {
let out = self.forward(&Array1::from_vec(xi.clone()));
if self.output_dim == 1 {
out[0]
} else {
// For multi-output, return first element as regression
out[0]
}
}).collect()
}
pub fn predict_cls3(&self, x: &[Vec<f32>]) -> Vec<usize> {
x.iter().map(|xi| {
let out = self.forward(&Array1::from_vec(xi.clone()));
if self.output_dim >= 3 {
// Find argmax for 3-class classification
let mut best = 0;
let mut best_val = out[0];
for i in 1..3 {
if out[i] > best_val {
best_val = out[i];
best = i;
}
}
best
} else {
// Fall back to threshold-based classification
let val = out[0];
if val < -0.25 { 0 }
else if val > 0.25 { 2 }
else { 1 }
}
}).collect()
}
}
@@ -0,0 +1,341 @@
use std::arch::x86_64::*;
use std::alloc::{alloc, dealloc, Layout};
use std::ptr;
/// Ultra-optimized CPU-only MLP with AVX-512 for lowest possible latency
/// Key optimizations:
/// - AVX-512: Process 16 floats at once (vs 8 with AVX2)
/// - Cache-aligned memory: Prevent false sharing
/// - Prefetching: Hide memory latency
/// - Loop unrolling: Reduce branch overhead
/// - Compile-time dimensions: Enable better optimization
pub struct UltraLowLatencyMlp<const INPUT: usize, const HIDDEN: usize, const OUTPUT: usize> {
// Cache-aligned weight storage (64-byte aligned for AVX-512)
w1: *mut f32, // HIDDEN x INPUT (row-major for sequential access)
b1: *mut f32, // HIDDEN
w2: *mut f32, // OUTPUT x HIDDEN
b2: *mut f32, // OUTPUT
// Pre-allocated buffers for zero-copy operation
hidden_buf: *mut f32, // HIDDEN (aligned)
}
impl<const I: usize, const H: usize, const O: usize> UltraLowLatencyMlp<I, H, O> {
pub fn new() -> Self {
unsafe {
// Allocate cache-aligned memory (64 bytes for AVX-512)
let align = 64;
let w1 = Self::alloc_aligned(H * I, align);
let b1 = Self::alloc_aligned(H, align);
let w2 = Self::alloc_aligned(O * H, align);
let b2 = Self::alloc_aligned(O, align);
let hidden_buf = Self::alloc_aligned(H, align);
// Initialize with small random weights
Self::init_weights(w1, H * I);
Self::init_weights(w2, O * H);
ptr::write_bytes(b1, 0, H);
ptr::write_bytes(b2, 0, O);
Self { w1, b1, w2, b2, hidden_buf }
}
}
#[inline(always)]
unsafe fn alloc_aligned(size: usize, align: usize) -> *mut f32 {
let layout = Layout::from_size_align(size * 4, align).unwrap();
alloc(layout) as *mut f32
}
unsafe fn init_weights(ptr: *mut f32, size: usize) {
let scale = (2.0 / size as f32).sqrt();
for i in 0..size {
*ptr.add(i) = (rand::random::<f32>() - 0.5) * scale;
}
}
/// Ultra-fast forward pass with AVX-512
/// Latency target: <100ns for typical sizes
#[target_feature(enable = "avx512f")]
#[inline]
pub unsafe fn forward_avx512(&self, input: &[f32; I], output: &mut [f32; O]) {
// Layer 1: Input -> Hidden with AVX-512 (16 floats at once)
self.matmul_avx512(input.as_ptr(), self.w1, self.b1, self.hidden_buf, H, I);
// ReLU activation (vectorized)
self.relu_avx512(self.hidden_buf, H);
// Layer 2: Hidden -> Output
self.matmul_avx512(self.hidden_buf, self.w2, self.b2, output.as_mut_ptr(), O, H);
}
/// AVX-512 matrix multiply with prefetching
#[target_feature(enable = "avx512f")]
unsafe fn matmul_avx512(&self, x: *const f32, w: *const f32, b: *const f32,
out: *mut f32, rows: usize, cols: usize) {
const SIMD_WIDTH: usize = 16; // AVX-512 processes 16 floats
// Process each output neuron
for i in 0..rows {
let row_offset = i * cols;
// Prefetch next cache line
_mm_prefetch(w.add(row_offset + 64) as *const i8, _MM_HINT_T0);
// Initialize accumulator with bias
let mut sum = _mm512_set1_ps(*b.add(i));
// Process 16 elements at a time
let chunks = cols / SIMD_WIDTH;
let mut j = 0;
// Unroll by 4 for better pipelining
while j + 3 < chunks {
let w0 = _mm512_load_ps(w.add(row_offset + j * SIMD_WIDTH));
let x0 = _mm512_load_ps(x.add(j * SIMD_WIDTH));
sum = _mm512_fmadd_ps(w0, x0, sum);
let w1 = _mm512_load_ps(w.add(row_offset + (j + 1) * SIMD_WIDTH));
let x1 = _mm512_load_ps(x.add((j + 1) * SIMD_WIDTH));
sum = _mm512_fmadd_ps(w1, x1, sum);
let w2 = _mm512_load_ps(w.add(row_offset + (j + 2) * SIMD_WIDTH));
let x2 = _mm512_load_ps(x.add((j + 2) * SIMD_WIDTH));
sum = _mm512_fmadd_ps(w2, x2, sum);
let w3 = _mm512_load_ps(w.add(row_offset + (j + 3) * SIMD_WIDTH));
let x3 = _mm512_load_ps(x.add((j + 3) * SIMD_WIDTH));
sum = _mm512_fmadd_ps(w3, x3, sum);
j += 4;
}
// Handle remaining chunks
while j < chunks {
let wv = _mm512_load_ps(w.add(row_offset + j * SIMD_WIDTH));
let xv = _mm512_load_ps(x.add(j * SIMD_WIDTH));
sum = _mm512_fmadd_ps(wv, xv, sum);
j += 1;
}
// Horizontal sum (reduce 16 -> 1)
let result = _mm512_reduce_add_ps(sum);
// Handle remaining elements (non-SIMD)
let mut scalar_sum = result;
for k in (chunks * SIMD_WIDTH)..cols {
scalar_sum += *w.add(row_offset + k) * *x.add(k);
}
*out.add(i) = scalar_sum;
}
}
/// Vectorized ReLU with AVX-512
#[target_feature(enable = "avx512f")]
unsafe fn relu_avx512(&self, data: *mut f32, size: usize) {
let zero = _mm512_setzero_ps();
let chunks = size / 16;
for i in 0..chunks {
let val = _mm512_load_ps(data.add(i * 16));
let relu = _mm512_max_ps(val, zero);
_mm512_store_ps(data.add(i * 16), relu);
}
// Handle remainder
for i in (chunks * 16)..size {
let val = *data.add(i);
*data.add(i) = val.max(0.0);
}
}
/// Fallback for non-AVX512 systems (still optimized)
pub fn forward_fallback(&self, input: &[f32; I], output: &mut [f32; O]) {
unsafe {
// Use AVX2 if available, otherwise scalar
#[cfg(target_feature = "avx2")]
{
self.forward_avx512(input, output);
}
#[cfg(not(target_feature = "avx2"))]
{
// Scalar fallback implementation
unsafe {
for i in 0..H {
let mut sum = *self.b1.add(i);
for j in 0..I {
sum += input[j] * *self.w1.add(i * I + j);
}
*self.hidden_buf.add(i) = sum.max(0.0);
}
for i in 0..O {
let mut sum = *self.b2.add(i);
for j in 0..H {
sum += *self.hidden_buf.add(j) * *self.w2.add(i * H + j);
}
output[i] = sum;
}
}
}
}
}
/// Train with low-latency SGD (no allocations in hot path)
pub fn train_fast(&mut self, x: &[f32; I], y: f32, lr: f32) {
unsafe {
let mut output = [0.0; O];
// Forward pass
self.forward_avx512(x, &mut output);
// Compute error
let error = output[0] - y;
// Backward pass (simplified, no allocation)
// Update output weights
for i in 0..H {
let grad = error * (*self.hidden_buf.add(i));
*self.w2.add(i) -= lr * grad;
}
*self.b2 -= lr * error;
// Backprop to hidden
for i in 0..H {
if *self.hidden_buf.add(i) > 0.0 { // ReLU gradient
let hidden_error = error * (*self.w2.add(i));
// Update hidden weights
for j in 0..I {
*self.w1.add(i * I + j) -= lr * hidden_error * x[j];
}
*self.b1.add(i) -= lr * hidden_error;
}
}
}
}
/// Batch prediction with minimal latency
#[inline]
pub fn predict_batch(&self, inputs: &[[f32; I]], outputs: &mut [[f32; O]]) {
// Sequential processing for thread safety
unsafe {
for (x, y) in inputs.iter().zip(outputs.iter_mut()) {
self.forward_avx512(x, y);
}
}
}
}
impl<const I: usize, const H: usize, const O: usize> Drop for UltraLowLatencyMlp<I, H, O> {
fn drop(&mut self) {
unsafe {
let align = 64;
dealloc(self.w1 as *mut u8, Layout::from_size_align(H * I * 4, align).unwrap());
dealloc(self.b1 as *mut u8, Layout::from_size_align(H * 4, align).unwrap());
dealloc(self.w2 as *mut u8, Layout::from_size_align(O * H * 4, align).unwrap());
dealloc(self.b2 as *mut u8, Layout::from_size_align(O * 4, align).unwrap());
dealloc(self.hidden_buf as *mut u8, Layout::from_size_align(H * 4, align).unwrap());
}
}
}
/// Wrapper for dynamic dimensions
pub struct DynamicAvx512Mlp {
weights_flat: Vec<f32>,
dims: (usize, usize, usize),
}
impl DynamicAvx512Mlp {
pub fn new(input: usize, hidden: usize, output: usize) -> Self {
let total_params = (input * hidden) + hidden + (hidden * output) + output;
let mut weights_flat = Vec::with_capacity(total_params);
// Initialize
let scale = (2.0 / input as f32).sqrt();
for _ in 0..total_params {
weights_flat.push((rand::random::<f32>() - 0.5) * scale);
}
Self { weights_flat, dims: (input, hidden, output) }
}
/// Predict with dynamic dimensions (still uses SIMD where possible)
pub fn predict(&self, x: &[Vec<f32>]) -> Vec<f32> {
let (input_dim, hidden_dim, _) = self.dims;
x.iter().map(|xi| {
let mut hidden = vec![0.0f32; hidden_dim];
// Layer 1: Use AVX2 if available
#[cfg(target_feature = "avx2")]
unsafe {
self.matmul_avx2_dynamic(&xi, &self.weights_flat[0..input_dim * hidden_dim],
&mut hidden, hidden_dim, input_dim);
}
#[cfg(not(target_feature = "avx2"))]
{
// Scalar fallback
for i in 0..hidden_dim {
let mut sum = self.weights_flat[input_dim * hidden_dim + i]; // bias
for j in 0..input_dim {
sum += xi[j] * self.weights_flat[i * input_dim + j];
}
hidden[i] = sum.max(0.0);
}
}
// Layer 2 (simplified for single output)
let w2_start = input_dim * hidden_dim + hidden_dim;
let mut output = self.weights_flat[w2_start + hidden_dim]; // bias
for i in 0..hidden_dim {
output += hidden[i] * self.weights_flat[w2_start + i];
}
output
}).collect()
}
#[cfg(target_feature = "avx2")]
#[target_feature(enable = "avx2")]
unsafe fn matmul_avx2_dynamic(&self, x: &[f32], w: &[f32], out: &mut [f32],
rows: usize, cols: usize) {
const SIMD_WIDTH: usize = 8;
for i in 0..rows {
let row_offset = i * cols;
let mut sum = _mm256_setzero_ps();
// Process 8 elements at a time
let chunks = cols / SIMD_WIDTH;
for j in 0..chunks {
let idx = row_offset + j * SIMD_WIDTH;
let wv = _mm256_loadu_ps(&w[idx]);
let xv = _mm256_loadu_ps(&x[j * SIMD_WIDTH]);
sum = _mm256_fmadd_ps(wv, xv, sum);
}
// Horizontal sum
let sum_array: [f32; 8] = std::mem::transmute(sum);
let mut result: f32 = sum_array.iter().sum();
// Handle remainder
for j in (chunks * SIMD_WIDTH)..cols {
result += w[row_offset + j] * x[j];
}
out[i] = result.max(0.0); // ReLU
}
}
pub fn predict_class(&self, x: &[Vec<f32>]) -> Vec<usize> {
self.predict(x).iter().map(|&y| {
if y < -0.25 { 0 }
else if y > 0.25 { 2 }
else { 1 }
}).collect()
}
}
@@ -0,0 +1,325 @@
use ndarray::{Array2, Array1, Axis};
use rand::{thread_rng, Rng};
/// Specialized classifier with proper softmax, dropout, and batch norm
pub struct ClassifierMlp {
// Weights
w1: Array2<f32>,
b1: Array1<f32>,
w2: Array2<f32>,
b2: Array1<f32>,
w3: Array2<f32>, // Extra layer for better classification
b3: Array1<f32>,
// Batch normalization parameters
bn1_gamma: Array1<f32>,
bn1_beta: Array1<f32>,
bn1_running_mean: Array1<f32>,
bn1_running_var: Array1<f32>,
bn2_gamma: Array1<f32>,
bn2_beta: Array1<f32>,
bn2_running_mean: Array1<f32>,
bn2_running_var: Array1<f32>,
// Adam optimizer state
mw1: Array2<f32>, mb1: Array1<f32>,
mw2: Array2<f32>, mb2: Array1<f32>,
mw3: Array2<f32>, mb3: Array1<f32>,
sw1: Array2<f32>, sb1: Array1<f32>,
sw2: Array2<f32>, sb2: Array1<f32>,
sw3: Array2<f32>, sb3: Array1<f32>,
t: f32,
// Architecture
input_dim: usize,
hidden1_dim: usize,
hidden2_dim: usize,
output_dim: usize,
// Training settings
dropout_rate: f32,
is_training: bool,
lr_schedule: LRSchedule,
}
pub enum LRSchedule {
Constant(f32),
Cosine { initial: f32, min: f32, period: usize },
StepDecay { initial: f32, decay: f32, step_size: usize },
}
impl ClassifierMlp {
pub fn new(input: usize, output: usize) -> Self {
let mut rng = thread_rng();
// Deeper network for better classification
let hidden1 = 128;
let hidden2 = 64;
// Xavier initialization
let scale1 = (2.0 / input as f32).sqrt();
let scale2 = (2.0 / hidden1 as f32).sqrt();
let scale3 = (2.0 / hidden2 as f32).sqrt();
let w1 = Array2::from_shape_fn((hidden1, input), |_|
rng.gen::<f32>() * scale1 - scale1/2.0);
let b1 = Array1::zeros(hidden1);
let w2 = Array2::from_shape_fn((hidden2, hidden1), |_|
rng.gen::<f32>() * scale2 - scale2/2.0);
let b2 = Array1::zeros(hidden2);
let w3 = Array2::from_shape_fn((output, hidden2), |_|
rng.gen::<f32>() * scale3 - scale3/2.0);
let b3 = Array1::zeros(output);
// Batch norm parameters
let bn1_gamma = Array1::ones(hidden1);
let bn1_beta = Array1::zeros(hidden1);
let bn1_running_mean = Array1::zeros(hidden1);
let bn1_running_var = Array1::ones(hidden1);
let bn2_gamma = Array1::ones(hidden2);
let bn2_beta = Array1::zeros(hidden2);
let bn2_running_mean = Array1::zeros(hidden2);
let bn2_running_var = Array1::ones(hidden2);
// Adam states
let mw1 = Array2::zeros((hidden1, input));
let mb1 = Array1::zeros(hidden1);
let mw2 = Array2::zeros((hidden2, hidden1));
let mb2 = Array1::zeros(hidden2);
let mw3 = Array2::zeros((output, hidden2));
let mb3 = Array1::zeros(output);
let sw1 = Array2::zeros((hidden1, input));
let sb1 = Array1::zeros(hidden1);
let sw2 = Array2::zeros((hidden2, hidden1));
let sb2 = Array1::zeros(hidden2);
let sw3 = Array2::zeros((output, hidden2));
let sb3 = Array1::zeros(output);
Self {
w1, b1, w2, b2, w3, b3,
bn1_gamma, bn1_beta, bn1_running_mean, bn1_running_var,
bn2_gamma, bn2_beta, bn2_running_mean, bn2_running_var,
mw1, mb1, mw2, mb2, mw3, mb3,
sw1, sb1, sw2, sb2, sw3, sb3,
t: 0.0,
input_dim: input,
hidden1_dim: hidden1,
hidden2_dim: hidden2,
output_dim: output,
dropout_rate: 0.3,
is_training: true,
lr_schedule: LRSchedule::Cosine {
initial: 0.001,
min: 0.00001,
period: 1000
},
}
}
fn batch_norm(&self, x: &Array1<f32>, gamma: &Array1<f32>, beta: &Array1<f32>,
running_mean: &Array1<f32>, running_var: &Array1<f32>) -> Array1<f32> {
if self.is_training {
let mean = x.mean().unwrap();
let var = x.var(0.0);
let x_norm = (x - mean) / (var + 1e-5).sqrt();
gamma * &x_norm + beta
} else {
let x_norm = (x - running_mean) / (running_var + 1e-5).mapv(f32::sqrt);
gamma * &x_norm + beta
}
}
fn dropout(&self, x: &Array1<f32>) -> Array1<f32> {
if self.is_training && self.dropout_rate > 0.0 {
let mut rng = thread_rng();
x.mapv(|v| if rng.gen::<f32>() > self.dropout_rate {
v / (1.0 - self.dropout_rate)
} else {
0.0
})
} else {
x.clone()
}
}
fn leaky_relu(x: &Array1<f32>) -> Array1<f32> {
x.mapv(|v| if v > 0.0 { v } else { 0.01 * v })
}
pub fn forward(&self, x: &Array1<f32>) -> (Array1<f32>, Array1<f32>, Array1<f32>) {
// Layer 1: Linear + BatchNorm + LeakyReLU + Dropout
let z1 = self.w1.dot(x) + &self.b1;
let bn1 = self.batch_norm(&z1, &self.bn1_gamma, &self.bn1_beta,
&self.bn1_running_mean, &self.bn1_running_var);
let h1 = Self::leaky_relu(&bn1);
let h1_drop = self.dropout(&h1);
// Layer 2: Linear + BatchNorm + LeakyReLU + Dropout
let z2 = self.w2.dot(&h1_drop) + &self.b2;
let bn2 = self.batch_norm(&z2, &self.bn2_gamma, &self.bn2_beta,
&self.bn2_running_mean, &self.bn2_running_var);
let h2 = Self::leaky_relu(&bn2);
let h2_drop = self.dropout(&h2);
// Output layer: Linear (no activation, will apply softmax in loss)
let logits = self.w3.dot(&h2_drop) + &self.b3;
(logits, h1_drop, h2_drop)
}
fn softmax(logits: &Array1<f32>) -> Array1<f32> {
let max = logits.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
let exp_vals = logits.mapv(|v| (v - max).exp());
let sum = exp_vals.sum();
exp_vals / sum
}
fn get_lr(&self, epoch: usize) -> f32 {
match &self.lr_schedule {
LRSchedule::Constant(lr) => *lr,
LRSchedule::Cosine { initial, min, period } => {
let progress = (epoch % period) as f32 / *period as f32;
min + (initial - min) * (1.0 + (std::f32::consts::PI * progress).cos()) / 2.0
}
LRSchedule::StepDecay { initial, decay, step_size } => {
initial * decay.powi((epoch / step_size) as i32)
}
}
}
pub fn backward(&mut self, x: &Array1<f32>, y_class: usize, epoch: usize) {
let (logits, h1, h2) = self.forward(x);
let probs = Self::softmax(&logits);
// Cross-entropy gradient
let mut grad_logits = probs;
if y_class < self.output_dim {
grad_logits[y_class] -= 1.0;
}
// Gradient w.r.t W3, b3
let grad_w3 = grad_logits.clone().insert_axis(Axis(1)) * h2.clone().insert_axis(Axis(0));
let grad_b3 = grad_logits.clone();
// Backprop to h2
let grad_h2 = self.w3.t().dot(&grad_logits);
let grad_h2_relu = grad_h2 * h2.mapv(|v| if v > 0.0 { 1.0 } else { 0.01 });
// Gradient w.r.t W2, b2
let grad_w2 = grad_h2_relu.clone().insert_axis(Axis(1)) * h1.clone().insert_axis(Axis(0));
let grad_b2 = grad_h2_relu.clone();
// Backprop to h1
let grad_h1 = self.w2.t().dot(&grad_h2_relu);
let grad_h1_relu = grad_h1 * h1.mapv(|v| if v > 0.0 { 1.0 } else { 0.01 });
// Gradient w.r.t W1, b1
let grad_w1 = grad_h1_relu.clone().insert_axis(Axis(1)) * x.clone().insert_axis(Axis(0));
let grad_b1 = grad_h1_relu;
// Adam update with scheduled learning rate
let lr = self.get_lr(epoch);
self.adam_update(
grad_w1, grad_b1,
grad_w2.into_shape(self.w2.dim()).unwrap(), grad_b2,
grad_w3.into_shape(self.w3.dim()).unwrap(), grad_b3,
lr
);
}
fn adam_update(&mut self, grad_w1: Array2<f32>, grad_b1: Array1<f32>,
grad_w2: Array2<f32>, grad_b2: Array1<f32>,
grad_w3: Array2<f32>, grad_b3: Array1<f32>, lr: f32) {
let beta1 = 0.9;
let beta2 = 0.999;
let epsilon = 1e-8;
self.t += 1.0;
// Update first moments
self.mw1 = &self.mw1 * beta1 + &grad_w1 * (1.0 - beta1);
self.mb1 = &self.mb1 * beta1 + &grad_b1 * (1.0 - beta1);
self.mw2 = &self.mw2 * beta1 + &grad_w2 * (1.0 - beta1);
self.mb2 = &self.mb2 * beta1 + &grad_b2 * (1.0 - beta1);
self.mw3 = &self.mw3 * beta1 + &grad_w3 * (1.0 - beta1);
self.mb3 = &self.mb3 * beta1 + &grad_b3 * (1.0 - beta1);
// Update second moments
self.sw1 = &self.sw1 * beta2 + grad_w1.mapv(|x| x * x) * (1.0 - beta2);
self.sb1 = &self.sb1 * beta2 + grad_b1.mapv(|x| x * x) * (1.0 - beta2);
self.sw2 = &self.sw2 * beta2 + grad_w2.mapv(|x| x * x) * (1.0 - beta2);
self.sb2 = &self.sb2 * beta2 + grad_b2.mapv(|x| x * x) * (1.0 - beta2);
self.sw3 = &self.sw3 * beta2 + grad_w3.mapv(|x| x * x) * (1.0 - beta2);
self.sb3 = &self.sb3 * beta2 + grad_b3.mapv(|x| x * x) * (1.0 - beta2);
// Bias correction
let bias1 = 1.0 - beta1.powf(self.t);
let bias2 = 1.0 - beta2.powf(self.t);
// Update weights
self.w1 = &self.w1 - lr * &self.mw1 / bias1 / ((&self.sw1 / bias2).mapv(f32::sqrt) + epsilon);
self.b1 = &self.b1 - lr * &self.mb1 / bias1 / ((&self.sb1 / bias2).mapv(f32::sqrt) + epsilon);
self.w2 = &self.w2 - lr * &self.mw2 / bias1 / ((&self.sw2 / bias2).mapv(f32::sqrt) + epsilon);
self.b2 = &self.b2 - lr * &self.mb2 / bias1 / ((&self.sb2 / bias2).mapv(f32::sqrt) + epsilon);
self.w3 = &self.w3 - lr * &self.mw3 / bias1 / ((&self.sw3 / bias2).mapv(f32::sqrt) + epsilon);
self.b3 = &self.b3 - lr * &self.mb3 / bias1 / ((&self.sb3 / bias2).mapv(f32::sqrt) + epsilon);
}
pub fn train_classification(&mut self, x: &Vec<Vec<f32>>, y: &Vec<f32>,
epochs: usize, batch_size: usize) {
self.is_training = true;
for epoch in 0..epochs {
let mut indices: Vec<usize> = (0..x.len()).collect();
let mut rng = thread_rng();
use rand::seq::SliceRandom;
indices.shuffle(&mut rng);
for batch_start in (0..x.len()).step_by(batch_size) {
let batch_end = (batch_start + batch_size).min(x.len());
for i in batch_start..batch_end {
let idx = indices[i];
let x_arr = Array1::from_vec(x[idx].clone());
// Convert continuous y to class
let y_class = if y[idx] < -0.25 { 0 }
else if y[idx] > 0.25 { 2 }
else { 1 };
self.backward(&x_arr, y_class, epoch);
}
// Update batch norm running stats
self.bn1_running_mean = &self.bn1_running_mean * 0.9 + &self.bn1_beta * 0.1;
self.bn1_running_var = &self.bn1_running_var * 0.9 + &self.bn1_gamma * 0.1;
self.bn2_running_mean = &self.bn2_running_mean * 0.9 + &self.bn2_beta * 0.1;
self.bn2_running_var = &self.bn2_running_var * 0.9 + &self.bn2_gamma * 0.1;
}
}
}
pub fn predict_cls3(&mut self, x: &[Vec<f32>]) -> Vec<usize> {
self.is_training = false;
x.iter().map(|xi| {
let x_arr = Array1::from_vec(xi.clone());
let (logits, _, _) = self.forward(&x_arr);
let probs = Self::softmax(&logits);
// Return argmax
let mut best = 0;
for i in 1..self.output_dim.min(3) {
if probs[i] > probs[best] {
best = i;
}
}
best
}).collect()
}
}
@@ -0,0 +1,246 @@
use ndarray::{Array2, Array1, Axis};
use rand::{thread_rng, Rng};
use rayon::prelude::*;
pub struct OptimizedMlp {
w1: Array2<f32>,
b1: Array1<f32>,
w2: Array2<f32>,
b2: Array1<f32>,
// Momentum terms
vw1: Array2<f32>,
vb1: Array1<f32>,
vw2: Array2<f32>,
vb2: Array1<f32>,
// Adam optimizer state
mw1: Array2<f32>,
mb1: Array1<f32>,
mw2: Array2<f32>,
mb2: Array1<f32>,
sw1: Array2<f32>,
sb1: Array1<f32>,
sw2: Array2<f32>,
sb2: Array1<f32>,
t: f32,
output_dim: usize,
use_adam: bool,
}
impl OptimizedMlp {
pub fn new(input: usize, hidden: usize, output: usize) -> Self {
let mut rng = thread_rng();
// He initialization for ReLU
let scale1 = (2.0 / input as f32).sqrt();
let scale2 = (2.0 / hidden as f32).sqrt();
let w1 = Array2::from_shape_fn((hidden, input), |_|
rng.gen::<f32>() * scale1 - scale1/2.0);
let b1 = Array1::zeros(hidden);
let w2 = Array2::from_shape_fn((output, hidden), |_|
rng.gen::<f32>() * scale2 - scale2/2.0);
let b2 = Array1::zeros(output);
// Initialize momentum and Adam states
let vw1 = Array2::zeros((hidden, input));
let vb1 = Array1::zeros(hidden);
let vw2 = Array2::zeros((output, hidden));
let vb2 = Array1::zeros(output);
let mw1 = Array2::zeros((hidden, input));
let mb1 = Array1::zeros(hidden);
let mw2 = Array2::zeros((output, hidden));
let mb2 = Array1::zeros(output);
let sw1 = Array2::zeros((hidden, input));
let sb1 = Array1::zeros(hidden);
let sw2 = Array2::zeros((output, hidden));
let sb2 = Array1::zeros(output);
Self {
w1, b1, w2, b2,
vw1, vb1, vw2, vb2,
mw1, mb1, mw2, mb2,
sw1, sb1, sw2, sb2,
t: 0.0,
output_dim: output,
use_adam: true,
}
}
pub fn forward(&self, x: &Array1<f32>) -> (Array1<f32>, Array1<f32>) {
let z1 = self.w1.dot(x) + &self.b1;
let h = z1.mapv(|v| v.max(0.0)); // ReLU
let output = self.w2.dot(&h) + &self.b2;
(output, h)
}
pub fn backward(&mut self, x: &Array1<f32>, y_true: f32, lr: f32) {
let (output, h) = self.forward(x);
// Compute loss gradient
let grad_output = if self.output_dim == 1 {
Array1::from_elem(1, output[0] - y_true)
} else {
// Softmax + cross-entropy gradient for classification
let exp_out = output.mapv(|v| v.exp());
let sum_exp = exp_out.sum();
let softmax = &exp_out / sum_exp;
// Create one-hot target
let class = if y_true < -0.25 { 0 }
else if y_true > 0.25 { 2 }
else { 1 };
let mut target = Array1::zeros(self.output_dim);
if class < self.output_dim {
target[class] = 1.0;
}
softmax - target
};
// Gradient w.r.t w2 and b2
let grad_w2 = grad_output.clone().insert_axis(Axis(1)) * h.clone().insert_axis(Axis(0));
let grad_b2 = grad_output.clone();
// Backprop through hidden layer
let grad_h = self.w2.t().dot(&grad_output);
// Gradient of ReLU: 1 if input > 0, else 0
let z1 = self.w1.dot(x) + &self.b1;
let grad_z1 = grad_h * z1.mapv(|v| if v > 0.0 { 1.0 } else { 0.0 });
// Gradient w.r.t w1 and b1
let grad_w1 = grad_z1.clone().insert_axis(Axis(1)) * x.clone().insert_axis(Axis(0));
let grad_b1 = grad_z1;
// Update weights using Adam or momentum
if self.use_adam {
self.adam_update(grad_w1, grad_b1, grad_w2.into_shape(self.w2.dim()).unwrap(), grad_b2, lr);
} else {
self.momentum_update(grad_w1, grad_b1, grad_w2.into_shape(self.w2.dim()).unwrap(), grad_b2, lr);
}
}
fn adam_update(&mut self, grad_w1: Array2<f32>, grad_b1: Array1<f32>,
grad_w2: Array2<f32>, grad_b2: Array1<f32>, lr: f32) {
let beta1 = 0.9;
let beta2 = 0.999;
let epsilon = 1e-8;
self.t += 1.0;
// Update biased first moment estimate
self.mw1 = &self.mw1 * beta1 + &grad_w1 * (1.0 - beta1);
self.mb1 = &self.mb1 * beta1 + &grad_b1 * (1.0 - beta1);
self.mw2 = &self.mw2 * beta1 + &grad_w2 * (1.0 - beta1);
self.mb2 = &self.mb2 * beta1 + &grad_b2 * (1.0 - beta1);
// Update biased second raw moment estimate
self.sw1 = &self.sw1 * beta2 + grad_w1.mapv(|x| x * x) * (1.0 - beta2);
self.sb1 = &self.sb1 * beta2 + grad_b1.mapv(|x| x * x) * (1.0 - beta2);
self.sw2 = &self.sw2 * beta2 + grad_w2.mapv(|x| x * x) * (1.0 - beta2);
self.sb2 = &self.sb2 * beta2 + grad_b2.mapv(|x| x * x) * (1.0 - beta2);
// Compute bias-corrected moments
let bias_correction1 = 1.0 - beta1.powf(self.t);
let bias_correction2 = 1.0 - beta2.powf(self.t);
// Update weights
self.w1 = &self.w1 - lr * &self.mw1 / bias_correction1 / ((&self.sw1 / bias_correction2).mapv(f32::sqrt) + epsilon);
self.b1 = &self.b1 - lr * &self.mb1 / bias_correction1 / ((&self.sb1 / bias_correction2).mapv(f32::sqrt) + epsilon);
self.w2 = &self.w2 - lr * &self.mw2 / bias_correction1 / ((&self.sw2 / bias_correction2).mapv(f32::sqrt) + epsilon);
self.b2 = &self.b2 - lr * &self.mb2 / bias_correction1 / ((&self.sb2 / bias_correction2).mapv(f32::sqrt) + epsilon);
}
fn momentum_update(&mut self, grad_w1: Array2<f32>, grad_b1: Array1<f32>,
grad_w2: Array2<f32>, grad_b2: Array1<f32>, lr: f32) {
let momentum = 0.9;
// Update velocity
self.vw1 = &self.vw1 * momentum - &grad_w1 * lr;
self.vb1 = &self.vb1 * momentum - &grad_b1 * lr;
self.vw2 = &self.vw2 * momentum - &grad_w2 * lr;
self.vb2 = &self.vb2 * momentum - &grad_b2 * lr;
// Update weights
self.w1 = &self.w1 + &self.vw1;
self.b1 = &self.b1 + &self.vb1;
self.w2 = &self.w2 + &self.vw2;
self.b2 = &self.b2 + &self.vb2;
}
pub fn train_regression(&mut self, x: &Vec<Vec<f32>>, y: &Vec<f32>, epochs: usize, lr: f32) {
for _ in 0..epochs {
// Shuffle indices for SGD
let mut indices: Vec<usize> = (0..x.len()).collect();
let mut rng = thread_rng();
use rand::seq::SliceRandom;
indices.shuffle(&mut rng);
for &i in &indices {
let x_arr = Array1::from_vec(x[i].clone());
self.backward(&x_arr, y[i], lr);
}
}
}
pub fn train_batch(&mut self, x: &Vec<Vec<f32>>, y: &Vec<f32>, epochs: usize,
lr: f32, batch_size: usize) {
for _ in 0..epochs {
let mut indices: Vec<usize> = (0..x.len()).collect();
let mut rng = thread_rng();
use rand::seq::SliceRandom;
indices.shuffle(&mut rng);
for batch_start in (0..x.len()).step_by(batch_size) {
let batch_end = (batch_start + batch_size).min(x.len());
// Accumulate gradients for batch
for i in batch_start..batch_end {
let idx = indices[i];
let x_arr = Array1::from_vec(x[idx].clone());
self.backward(&x_arr, y[idx], lr / batch_size as f32);
}
}
}
}
pub fn predict_reg(&self, x: &[Vec<f32>]) -> Vec<f32> {
x.par_iter().map(|xi| {
let (out, _) = self.forward(&Array1::from_vec(xi.clone()));
if self.output_dim == 1 {
out[0]
} else {
out[0]
}
}).collect()
}
pub fn predict_cls3(&self, x: &[Vec<f32>]) -> Vec<usize> {
x.par_iter().map(|xi| {
let (out, _) = self.forward(&Array1::from_vec(xi.clone()));
if self.output_dim >= 3 {
// Softmax + argmax for classification
let exp_out = out.mapv(|v| v.exp());
let sum_exp = exp_out.sum();
let probs = exp_out / sum_exp;
let mut best = 0;
let mut best_val = probs[0];
for i in 1..3.min(probs.len()) {
if probs[i] > best_val {
best_val = probs[i];
best = i;
}
}
best
} else {
let val = out[0];
if val < -0.25 { 0 }
else if val > 0.25 { 2 }
else { 1 }
}
}).collect()
}
}
@@ -0,0 +1,267 @@
use crate::quantization::{QuantizedMlp, QuantizedWeights};
use rand::Rng;
/// Quantized MLP wrapper for temporal-compare
/// Provides 4x model size reduction with minimal accuracy loss
pub struct QuantizedMlpBackend {
quantized: Option<QuantizedMlp>,
// Training happens in FP32, quantize after
weights1: Vec<f32>,
bias1: Vec<f32>,
weights2: Vec<f32>,
bias2: Vec<f32>,
input_dim: usize,
hidden_dim: usize,
output_dim: usize,
// Track compression stats
original_size: usize,
quantized_size: usize,
}
impl QuantizedMlpBackend {
pub fn new(input_dim: usize, hidden_dim: usize, output_dim: usize) -> Self {
let mut rng = rand::thread_rng();
// Xavier initialization
let scale1 = (2.0 / input_dim as f32).sqrt();
let scale2 = (2.0 / hidden_dim as f32).sqrt();
let weights1: Vec<f32> = (0..hidden_dim * input_dim)
.map(|_| rng.gen_range(-scale1..scale1))
.collect();
let weights2: Vec<f32> = (0..output_dim * hidden_dim)
.map(|_| rng.gen_range(-scale2..scale2))
.collect();
let original_size = (weights1.len() + weights2.len() + hidden_dim + output_dim) * 4;
Self {
quantized: None,
weights1,
bias1: vec![0.0; hidden_dim],
weights2,
bias2: vec![0.0; output_dim],
input_dim,
hidden_dim,
output_dim,
original_size,
quantized_size: 0,
}
}
/// Train in FP32 for best accuracy
pub fn train(&mut self, x: &[Vec<f32>], y: &[f32], epochs: usize, lr: f32) {
for epoch in 0..epochs {
let mut total_loss = 0.0;
for (xi, &yi) in x.iter().zip(y.iter()) {
// Forward pass (FP32)
let mut hidden = vec![0.0f32; self.hidden_dim];
// Layer 1
for i in 0..self.hidden_dim {
let mut sum = self.bias1[i];
for j in 0..self.input_dim {
sum += self.weights1[i * self.input_dim + j] * xi[j];
}
hidden[i] = sum.max(0.0); // ReLU
}
// Layer 2
let mut output = self.bias2[0];
for i in 0..self.hidden_dim {
output += self.weights2[i] * hidden[i];
}
// Loss (MSE)
let error = output - yi;
total_loss += error * error;
// Backward pass
// Output layer gradients
for i in 0..self.hidden_dim {
self.weights2[i] -= lr * error * hidden[i];
}
self.bias2[0] -= lr * error;
// Hidden layer gradients
for i in 0..self.hidden_dim {
if hidden[i] > 0.0 {
let grad = error * self.weights2[i];
for j in 0..self.input_dim {
self.weights1[i * self.input_dim + j] -= lr * grad * xi[j];
}
self.bias1[i] -= lr * grad;
}
}
}
if epoch % 10 == 0 {
println!("Epoch {}: Loss = {:.6}", epoch, total_loss / x.len() as f32);
}
}
// Quantize after training
self.quantize();
}
/// Quantize the trained FP32 model to INT8
pub fn quantize(&mut self) {
let qmlp = QuantizedMlp::from_float_mlp(
&self.weights1,
&self.bias1,
&self.weights2,
&self.bias2,
self.input_dim,
self.hidden_dim,
self.output_dim
);
self.quantized_size = qmlp.model_size();
self.quantized = Some(qmlp);
}
/// Predict using quantized weights (fast)
pub fn predict(&self, x: &[Vec<f32>]) -> Vec<f32> {
match &self.quantized {
Some(qmlp) => {
x.iter().map(|xi| {
let mut output = vec![0.0f32; self.output_dim];
#[cfg(target_arch = "x86_64")]
{
qmlp.forward_avx2(xi, &mut output);
}
#[cfg(not(target_arch = "x86_64"))]
{
qmlp.forward(xi, &mut output);
}
output[0]
}).collect()
}
None => {
// Fallback to FP32 if not quantized
self.predict_fp32(x)
}
}
}
/// Predict using FP32 weights (for comparison)
pub fn predict_fp32(&self, x: &[Vec<f32>]) -> Vec<f32> {
x.iter().map(|xi| {
let mut hidden = vec![0.0f32; self.hidden_dim];
// Layer 1
for i in 0..self.hidden_dim {
let mut sum = self.bias1[i];
for j in 0..self.input_dim {
sum += self.weights1[i * self.input_dim + j] * xi[j];
}
hidden[i] = sum.max(0.0);
}
// Layer 2
let mut output = self.bias2[0];
for i in 0..self.hidden_dim {
output += self.weights2[i] * hidden[i];
}
output
}).collect()
}
/// Classification prediction
pub fn predict_class(&self, x: &[Vec<f32>]) -> Vec<usize> {
self.predict(x).iter().map(|&y| {
if y < -0.25 { 0 }
else if y > 0.25 { 2 }
else { 1 }
}).collect()
}
/// Get compression statistics
pub fn get_compression_stats(&self) -> (usize, usize, f32) {
let ratio = if self.quantized_size > 0 {
self.original_size as f32 / self.quantized_size as f32
} else {
1.0
};
(self.original_size, self.quantized_size, ratio)
}
/// Compare FP32 vs INT8 performance
pub fn benchmark_inference(&self, x: &[Vec<f32>], iterations: usize) {
use std::time::Instant;
// Warm up
let _ = self.predict_fp32(&x[..1.min(x.len())]);
let _ = self.predict(&x[..1.min(x.len())]);
// Benchmark FP32
let start = Instant::now();
for _ in 0..iterations {
let _ = self.predict_fp32(x);
}
let fp32_time = start.elapsed();
// Benchmark INT8
let start = Instant::now();
for _ in 0..iterations {
let _ = self.predict(x);
}
let int8_time = start.elapsed();
let speedup = fp32_time.as_secs_f32() / int8_time.as_secs_f32();
println!("\n=== Quantization Benchmark ===");
println!("FP32 time: {:.3}s", fp32_time.as_secs_f32());
println!("INT8 time: {:.3}s", int8_time.as_secs_f32());
println!("Speedup: {:.2}x", speedup);
let (orig, quant, ratio) = self.get_compression_stats();
println!("\n=== Model Size ===");
println!("Original: {} bytes", orig);
println!("Quantized: {} bytes", quant);
println!("Compression: {:.2}x", ratio);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_quantized_mlp() {
let mut model = QuantizedMlpBackend::new(32, 64, 1);
// Create dummy data
let x: Vec<Vec<f32>> = (0..100)
.map(|_| (0..32).map(|_| rand::random()).collect())
.collect();
let y: Vec<f32> = (0..100).map(|_| rand::random()).collect();
// Train and quantize
model.train(&x, &y, 10, 0.01);
// Check predictions work
let pred_fp32 = model.predict_fp32(&x[..10]);
let pred_int8 = model.predict(&x[..10]);
// Should be similar but not identical
for (p32, p8) in pred_fp32.iter().zip(&pred_int8) {
let diff = (p32 - p8).abs();
assert!(diff < 0.1, "Quantization error too large: {}", diff);
}
// Check compression
let (_, _, ratio) = model.get_compression_stats();
assert!(ratio > 3.0, "Compression ratio too low: {}", ratio);
}
}
@@ -0,0 +1,314 @@
use std::arch::x86_64::*;
use ndarray::Axis;
use rand::{thread_rng, Rng};
use rayon::prelude::*;
use std::sync::Arc;
/// Ultra-optimized MLP with SIMD, cache blocking, and parallel processing
pub struct UltraMlp {
// Weights stored in row-major for cache efficiency
w1_flat: Vec<f32>,
b1: Vec<f32>,
w2_flat: Vec<f32>,
b2: Vec<f32>,
// Dimensions
input_dim: usize,
hidden_dim: usize,
output_dim: usize,
// Momentum buffers (flat)
vw1: Vec<f32>,
vb1: Vec<f32>,
vw2: Vec<f32>,
vb2: Vec<f32>,
// Pre-allocated buffers for forward pass
hidden_buffer: Vec<f32>,
output_buffer: Vec<f32>,
}
impl UltraMlp {
pub fn new(input: usize, hidden: usize, output: usize) -> Self {
let mut rng = thread_rng();
// He initialization
let scale1 = (2.0 / input as f32).sqrt();
let scale2 = (2.0 / hidden as f32).sqrt();
// Initialize weights flat for SIMD
let w1_flat: Vec<f32> = (0..hidden*input)
.map(|_| rng.gen::<f32>() * scale1 - scale1/2.0)
.collect();
let w2_flat: Vec<f32> = (0..output*hidden)
.map(|_| rng.gen::<f32>() * scale2 - scale2/2.0)
.collect();
Self {
w1_flat,
b1: vec![0.0; hidden],
w2_flat,
b2: vec![0.0; output],
input_dim: input,
hidden_dim: hidden,
output_dim: output,
vw1: vec![0.0; hidden * input],
vb1: vec![0.0; hidden],
vw2: vec![0.0; output * hidden],
vb2: vec![0.0; output],
hidden_buffer: vec![0.0; hidden],
output_buffer: vec![0.0; output],
}
}
/// SIMD-accelerated matrix-vector multiplication
#[target_feature(enable = "avx2")]
#[inline]
unsafe fn simd_matmul(weights: &[f32], input: &[f32], output: &mut [f32],
rows: usize, cols: usize) {
const SIMD_WIDTH: usize = 8;
for i in 0..rows {
let row_offset = i * cols;
let mut sum = _mm256_setzero_ps();
// Process 8 elements at a time with AVX2
let chunks = cols / SIMD_WIDTH;
for j in 0..chunks {
let idx = j * SIMD_WIDTH;
let w = _mm256_loadu_ps(&weights[row_offset + idx]);
let x = _mm256_loadu_ps(&input[idx]);
sum = _mm256_fmadd_ps(w, x, sum);
}
// Horizontal sum
let sum_array = std::mem::transmute::<__m256, [f32; 8]>(sum);
let mut result = sum_array.iter().sum::<f32>();
// Handle remaining elements
for j in (chunks * SIMD_WIDTH)..cols {
result += weights[row_offset + j] * input[j];
}
output[i] = result;
}
}
/// Vectorized ReLU activation
#[target_feature(enable = "avx2")]
#[inline]
unsafe fn simd_relu(data: &mut [f32]) {
const SIMD_WIDTH: usize = 8;
let zero = _mm256_setzero_ps();
let chunks = data.len() / SIMD_WIDTH;
for i in 0..chunks {
let idx = i * SIMD_WIDTH;
let val = _mm256_loadu_ps(&data[idx]);
let relu = _mm256_max_ps(val, zero);
_mm256_storeu_ps(&mut data[idx], relu);
}
// Handle remaining
for i in (chunks * SIMD_WIDTH)..data.len() {
data[i] = data[i].max(0.0);
}
}
/// Ultra-fast forward pass
pub fn forward_fast(&mut self, x: &[f32]) -> &[f32] {
unsafe {
// Layer 1: input -> hidden
Self::simd_matmul(&self.w1_flat, x, &mut self.hidden_buffer,
self.hidden_dim, self.input_dim);
// Add bias with SIMD
for i in 0..self.hidden_dim {
self.hidden_buffer[i] += self.b1[i];
}
// ReLU activation
Self::simd_relu(&mut self.hidden_buffer);
// Layer 2: hidden -> output
Self::simd_matmul(&self.w2_flat, &self.hidden_buffer, &mut self.output_buffer,
self.output_dim, self.hidden_dim);
// Add bias
for i in 0..self.output_dim {
self.output_buffer[i] += self.b2[i];
}
}
&self.output_buffer
}
/// Optimized backpropagation with momentum
pub fn backward_fast(&mut self, x: &[f32], y_true: f32, lr: f32) {
// Clone output to avoid borrow issues
let output_copy = self.forward_fast(x).to_vec();
let momentum = 0.9;
// Output gradient
let grad_out = if self.output_dim == 1 {
vec![output_copy[0] - y_true]
} else {
// Softmax gradient for classification
let max = output_copy.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
let exp_sum = output_copy.iter().map(|&v| (v - max).exp()).sum::<f32>();
let mut grad = output_copy.iter().map(|&v| (v - max).exp() / exp_sum).collect::<Vec<_>>();
let class = if y_true < -0.25 { 0 } else if y_true > 0.25 { 2 } else { 1 };
if class < grad.len() {
grad[class] -= 1.0;
}
grad
};
// Gradient w.r.t W2 and b2 (vectorized)
for i in 0..self.output_dim {
let grad_i = grad_out[i];
// Update bias with momentum
self.vb2[i] = momentum * self.vb2[i] - lr * grad_i;
self.b2[i] += self.vb2[i];
// Update weights with momentum (vectorized)
let w2_offset = i * self.hidden_dim;
for j in 0..self.hidden_dim {
let idx = w2_offset + j;
let grad_w = grad_i * self.hidden_buffer[j];
self.vw2[idx] = momentum * self.vw2[idx] - lr * grad_w;
self.w2_flat[idx] += self.vw2[idx];
}
}
// Gradient backprop to hidden layer
let mut grad_hidden = vec![0.0; self.hidden_dim];
for i in 0..self.hidden_dim {
for j in 0..self.output_dim {
grad_hidden[i] += self.w2_flat[j * self.hidden_dim + i] * grad_out[j];
}
// ReLU gradient
if self.hidden_buffer[i] <= 0.0 {
grad_hidden[i] = 0.0;
}
}
// Update W1 and b1
for i in 0..self.hidden_dim {
let grad_i = grad_hidden[i];
// Update bias
self.vb1[i] = momentum * self.vb1[i] - lr * grad_i;
self.b1[i] += self.vb1[i];
// Update weights
let w1_offset = i * self.input_dim;
for j in 0..self.input_dim {
let idx = w1_offset + j;
let grad_w = grad_i * x[j];
self.vw1[idx] = momentum * self.vw1[idx] - lr * grad_w;
self.w1_flat[idx] += self.vw1[idx];
}
}
}
/// Parallel batch training with cache-friendly access
pub fn train_batch_parallel(&mut self, x: &Vec<Vec<f32>>, y: &Vec<f32>,
epochs: usize, lr: f32, batch_size: usize) {
use rand::seq::SliceRandom;
for _ in 0..epochs {
let mut indices: Vec<usize> = (0..x.len()).collect();
let mut rng = thread_rng();
indices.shuffle(&mut rng);
// Process in mini-batches
for batch_start in (0..x.len()).step_by(batch_size) {
let batch_end = (batch_start + batch_size).min(x.len());
// Sequential within batch for weight updates
for i in batch_start..batch_end {
let idx = indices[i];
self.backward_fast(&x[idx], y[idx], lr / batch_size as f32);
}
}
}
}
/// Parallel prediction for multiple samples
pub fn predict_parallel(&mut self, x: &[Vec<f32>]) -> Vec<f32> {
// Create thread-local copies for parallel execution
let w1 = Arc::new(self.w1_flat.clone());
let b1 = Arc::new(self.b1.clone());
let w2 = Arc::new(self.w2_flat.clone());
let b2 = Arc::new(self.b2.clone());
let hidden_dim = self.hidden_dim;
let input_dim = self.input_dim;
let output_dim = self.output_dim;
x.par_iter().map(|xi| {
let mut hidden = vec![0.0; hidden_dim];
let mut output = vec![0.0; output_dim];
unsafe {
// Forward pass with local buffers
Self::simd_matmul(&w1, xi, &mut hidden, hidden_dim, input_dim);
for i in 0..hidden_dim {
hidden[i] += b1[i];
}
Self::simd_relu(&mut hidden);
Self::simd_matmul(&w2, &hidden, &mut output, output_dim, hidden_dim);
for i in 0..output_dim {
output[i] += b2[i];
}
}
if output_dim == 1 { output[0] } else { output[0] }
}).collect()
}
pub fn predict_cls3_parallel(&mut self, x: &[Vec<f32>]) -> Vec<usize> {
let w1 = Arc::new(self.w1_flat.clone());
let b1 = Arc::new(self.b1.clone());
let w2 = Arc::new(self.w2_flat.clone());
let b2 = Arc::new(self.b2.clone());
let hidden_dim = self.hidden_dim;
let input_dim = self.input_dim;
let output_dim = self.output_dim;
x.par_iter().map(|xi| {
let mut hidden = vec![0.0; hidden_dim];
let mut output = vec![0.0; output_dim];
unsafe {
Self::simd_matmul(&w1, xi, &mut hidden, hidden_dim, input_dim);
for i in 0..hidden_dim {
hidden[i] += b1[i];
}
Self::simd_relu(&mut hidden);
Self::simd_matmul(&w2, &hidden, &mut output, output_dim, hidden_dim);
for i in 0..output_dim {
output[i] += b2[i];
}
}
if output_dim >= 3 {
// Argmax
let mut best = 0;
for i in 1..3.min(output.len()) {
if output[i] > output[best] {
best = i;
}
}
best
} else {
let val = output[0];
if val < -0.25 { 0 } else if val > 0.25 { 2 } else { 1 }
}
}).collect()
}
}
@@ -0,0 +1,323 @@
use std::ops::{Add, Mul};
/// INT8 Quantization for 4x size reduction and 2x faster inference
/// Symmetric quantization: maps [-max, max] to [-127, 127]
#[derive(Clone, Debug)]
pub struct Int8Quantizer {
pub scale: f32,
pub zero_point: i8,
}
impl Int8Quantizer {
/// Create quantizer from float weights
pub fn from_weights(weights: &[f32]) -> Self {
// Find min/max for dynamic range
let (min, max) = weights.iter().fold(
(f32::INFINITY, f32::NEG_INFINITY),
|(min, max), &w| (min.min(w), max.max(w))
);
// Symmetric quantization for better accuracy
let abs_max = min.abs().max(max.abs());
let scale = abs_max / 127.0;
Self {
scale,
zero_point: 0, // Symmetric around zero
}
}
/// Quantize float to INT8
#[inline]
pub fn quantize(&self, value: f32) -> i8 {
let scaled = value / self.scale;
scaled.round().clamp(-127.0, 127.0) as i8
}
/// Dequantize INT8 back to float
#[inline]
pub fn dequantize(&self, value: i8) -> f32 {
value as f32 * self.scale
}
/// Quantize entire array
pub fn quantize_array(&self, values: &[f32]) -> Vec<i8> {
values.iter().map(|&v| self.quantize(v)).collect()
}
/// Dequantize entire array
pub fn dequantize_array(&self, values: &[i8]) -> Vec<f32> {
values.iter().map(|&v| self.dequantize(v)).collect()
}
}
/// Quantized weight storage for neural networks
pub struct QuantizedWeights {
pub weights: Vec<i8>,
pub quantizer: Int8Quantizer,
pub shape: (usize, usize),
}
impl QuantizedWeights {
pub fn from_float_matrix(weights: &[f32], rows: usize, cols: usize) -> Self {
let quantizer = Int8Quantizer::from_weights(weights);
let quantized = quantizer.quantize_array(weights);
Self {
weights: quantized,
quantizer,
shape: (rows, cols),
}
}
/// Get dequantized weight at position
#[inline]
pub fn get(&self, row: usize, col: usize) -> f32 {
let idx = row * self.shape.1 + col;
self.quantizer.dequantize(self.weights[idx])
}
/// Quantized matrix multiply with on-the-fly dequantization
/// Still faster than FP32 due to better cache utilization
pub fn matmul_quantized(&self, input: &[f32], output: &mut [f32]) {
let (rows, cols) = self.shape;
for i in 0..rows {
let mut sum = 0.0f32;
let row_offset = i * cols;
// Process in chunks for better cache performance
for j in 0..cols {
let w_int8 = self.weights[row_offset + j];
// Delay dequantization to minimize float operations
sum += (w_int8 as f32) * input[j];
}
// Apply scale once at the end
output[i] = sum * self.quantizer.scale;
}
}
/// SIMD-accelerated quantized matmul (AVX2)
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
pub unsafe fn matmul_quantized_avx2(&self, input: &[f32], output: &mut [f32]) {
use std::arch::x86_64::*;
let (rows, cols) = self.shape;
let scale = _mm256_set1_ps(self.quantizer.scale);
for i in 0..rows {
let mut sum = _mm256_setzero_ps();
let row_offset = i * cols;
// Process 8 elements at once
let chunks = cols / 8;
for j in 0..chunks {
let idx = j * 8;
// Load 8 INT8 weights and convert to float
let w_ptr = self.weights.as_ptr().add(row_offset + idx);
let w_i8 = _mm_loadl_epi64(w_ptr as *const __m128i);
let w_i32 = _mm256_cvtepi8_epi32(w_i8);
let w_f32 = _mm256_cvtepi32_ps(w_i32);
// Load 8 input floats
let x = _mm256_loadu_ps(&input[idx]);
// Multiply and accumulate
sum = _mm256_fmadd_ps(w_f32, x, sum);
}
// Horizontal sum
let sum_array: [f32; 8] = std::mem::transmute(sum);
let mut result = sum_array.iter().sum::<f32>();
// Handle remainder
for j in (chunks * 8)..cols {
result += (self.weights[row_offset + j] as f32) * input[j];
}
// Apply scale
output[i] = result * self.quantizer.scale;
}
}
/// Get memory size in bytes
pub fn memory_size(&self) -> usize {
self.weights.len() // INT8 = 1 byte each
}
}
/// Quantized MLP with INT8 weights
pub struct QuantizedMlp {
pub w1: QuantizedWeights,
pub b1: Vec<f32>, // Keep biases as FP32 (small overhead)
pub w2: QuantizedWeights,
pub b2: Vec<f32>,
hidden_dim: usize,
}
impl QuantizedMlp {
/// Create from existing float MLP
pub fn from_float_mlp(w1: &[f32], b1: &[f32], w2: &[f32], b2: &[f32],
input_dim: usize, hidden_dim: usize, output_dim: usize) -> Self {
Self {
w1: QuantizedWeights::from_float_matrix(w1, hidden_dim, input_dim),
b1: b1.to_vec(),
w2: QuantizedWeights::from_float_matrix(w2, output_dim, hidden_dim),
b2: b2.to_vec(),
hidden_dim,
}
}
/// Forward pass with INT8 weights
pub fn forward(&self, input: &[f32], output: &mut [f32]) {
let mut hidden = vec![0.0f32; self.hidden_dim];
// Layer 1: Input -> Hidden
self.w1.matmul_quantized(input, &mut hidden);
// Add bias and ReLU
for (h, &b) in hidden.iter_mut().zip(&self.b1) {
*h = (*h + b).max(0.0);
}
// Layer 2: Hidden -> Output
self.w2.matmul_quantized(&hidden, output);
// Add bias
for (o, &b) in output.iter_mut().zip(&self.b2) {
*o += b;
}
}
/// Forward pass with AVX2 acceleration
#[cfg(target_arch = "x86_64")]
pub fn forward_avx2(&self, input: &[f32], output: &mut [f32]) {
unsafe {
let mut hidden = vec![0.0f32; self.hidden_dim];
// Layer 1 with SIMD
self.w1.matmul_quantized_avx2(input, &mut hidden);
// Vectorized bias + ReLU
use std::arch::x86_64::*;
let zero = _mm256_setzero_ps();
for i in (0..self.hidden_dim).step_by(8) {
if i + 8 <= self.hidden_dim {
let h = _mm256_loadu_ps(&hidden[i]);
let b = _mm256_loadu_ps(&self.b1[i]);
let sum = _mm256_add_ps(h, b);
let relu = _mm256_max_ps(sum, zero);
_mm256_storeu_ps(&mut hidden[i], relu);
} else {
// Handle remainder
for j in i..self.hidden_dim {
hidden[j] = (hidden[j] + self.b1[j]).max(0.0);
}
}
}
// Layer 2 with SIMD
self.w2.matmul_quantized_avx2(&hidden, output);
// Add output bias
for (o, &b) in output.iter_mut().zip(&self.b2) {
*o += b;
}
}
}
/// Get total model size in bytes
pub fn model_size(&self) -> usize {
self.w1.memory_size() +
self.w2.memory_size() +
(self.b1.len() + self.b2.len()) * 4 // FP32 biases
}
/// Compression ratio vs FP32
pub fn compression_ratio(&self, original_params: usize) -> f32 {
let original_bytes = original_params * 4; // FP32
let quantized_bytes = self.model_size();
original_bytes as f32 / quantized_bytes as f32
}
}
/// Quantization-aware training for better INT8 accuracy
pub struct QuantizationAwareTraining {
pub fake_quantize: bool,
pub num_bits: u8,
}
impl QuantizationAwareTraining {
pub fn new() -> Self {
Self {
fake_quantize: true,
num_bits: 8,
}
}
/// Fake quantization during training to simulate INT8 effects
pub fn fake_quantize_weights(&self, weights: &mut [f32]) {
if !self.fake_quantize {
return;
}
let quantizer = Int8Quantizer::from_weights(weights);
for w in weights.iter_mut() {
// Quantize and immediately dequantize to simulate INT8 effects
let quantized = quantizer.quantize(*w);
*w = quantizer.dequantize(quantized);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_quantization_round_trip() {
let weights = vec![-1.5, -0.5, 0.0, 0.5, 1.5];
let quantizer = Int8Quantizer::from_weights(&weights);
for &w in &weights {
let q = quantizer.quantize(w);
let dq = quantizer.dequantize(q);
// Should be close but not exact due to quantization
assert!((w - dq).abs() < 0.02);
}
}
#[test]
fn test_compression_ratio() {
let input_dim = 32;
let hidden_dim = 64;
let output_dim = 3;
let total_params = (input_dim * hidden_dim) + hidden_dim +
(hidden_dim * output_dim) + output_dim;
let w1 = vec![0.1; input_dim * hidden_dim];
let b1 = vec![0.0; hidden_dim];
let w2 = vec![0.1; hidden_dim * output_dim];
let b2 = vec![0.0; output_dim];
let qmlp = QuantizedMlp::from_float_mlp(
&w1, &b1, &w2, &b2,
input_dim, hidden_dim, output_dim
);
let ratio = qmlp.compression_ratio(total_params);
// Should achieve ~3.5-4x compression (weights are INT8, biases stay FP32)
assert!(ratio > 3.0 && ratio < 4.5);
println!("Compression ratio: {:.2}x", ratio);
println!("Original size: {} bytes", total_params * 4);
println!("Quantized size: {} bytes", qmlp.model_size());
}
}
@@ -0,0 +1,311 @@
use ndarray::{Array2, Array1, ArrayView1};
use rand::{thread_rng, Rng, distributions::Uniform};
use std::f32::consts::PI;
/// Echo State Network (Reservoir Computing) - Ultra-efficient for time series
/// Key insight: Random reservoir + simple readout often beats complex architectures
pub struct ReservoirComputer {
// Reservoir parameters
reservoir_size: usize,
spectral_radius: f32,
sparsity: f32,
leak_rate: f32,
// Reservoir matrices
w_in: Array2<f32>, // Input weights (random, fixed)
w_res: Array2<f32>, // Reservoir weights (random, sparse, fixed)
w_out: Array2<f32>, // Output weights (trained)
// State
state: Array1<f32>,
// Memory optimization: ring buffer for states
state_history: Vec<Array1<f32>>,
history_index: usize,
max_history: usize,
}
impl ReservoirComputer {
pub fn new(input_dim: usize, reservoir_size: usize, output_dim: usize) -> Self {
let mut rng = thread_rng();
let uniform = Uniform::new(-1.0, 1.0);
// Input weights: random projection
let w_in = Array2::from_shape_fn((reservoir_size, input_dim), |_|
rng.sample(uniform));
// Reservoir weights: sparse random matrix
let sparsity = 0.9; // 90% zeros for efficiency
let mut w_res = Array2::zeros((reservoir_size, reservoir_size));
// Create sparse reservoir with specific spectral radius
for i in 0..reservoir_size {
for j in 0..reservoir_size {
if rng.gen::<f32>() > sparsity {
w_res[[i, j]] = rng.sample(uniform);
}
}
}
// Scale to desired spectral radius (0.9 for edge of chaos)
let spectral_radius = 0.9;
w_res = Self::scale_spectral_radius(w_res, spectral_radius);
// Output weights will be learned
let w_out = Array2::zeros((output_dim, reservoir_size + input_dim));
Self {
reservoir_size,
spectral_radius,
sparsity,
leak_rate: 0.3,
w_in,
w_res,
w_out,
state: Array1::zeros(reservoir_size),
state_history: Vec::new(),
history_index: 0,
max_history: 1000,
}
}
/// Scale matrix to desired spectral radius
fn scale_spectral_radius(mut matrix: Array2<f32>, target_radius: f32) -> Array2<f32> {
// Approximate largest eigenvalue with power iteration
let n = matrix.nrows();
let mut v = Array1::from_shape_fn(n, |_| rand::random::<f32>());
for _ in 0..20 {
v = matrix.dot(&v);
let norm = v.dot(&v).sqrt();
if norm > 0.0 {
v = v / norm;
}
}
let eigenvalue = v.dot(&matrix.dot(&v)) / v.dot(&v);
if eigenvalue.abs() > 0.0 {
matrix = matrix * (target_radius / eigenvalue.abs());
}
matrix
}
/// Update reservoir state with new input
pub fn update_state(&mut self, input: &Array1<f32>) -> Array1<f32> {
// Reservoir dynamics: state = (1-α)*state + α*tanh(W_in*input + W_res*state)
let input_contribution = self.w_in.dot(input);
let recurrent_contribution = self.w_res.dot(&self.state);
let new_state = &input_contribution + &recurrent_contribution;
let activated = new_state.mapv(|x| x.tanh());
// Leaky integration
self.state = &self.state * (1.0 - self.leak_rate) + &activated * self.leak_rate;
// Store in history (ring buffer)
if self.state_history.len() < self.max_history {
self.state_history.push(self.state.clone());
} else {
self.state_history[self.history_index] = self.state.clone();
}
self.history_index = (self.history_index + 1) % self.max_history;
self.state.clone()
}
/// Collect states for training (washout period to remove initial transients)
pub fn collect_states(&mut self, inputs: &[Vec<f32>], washout: usize)
-> (Array2<f32>, Vec<Array1<f32>>) {
let n_samples = inputs.len();
let mut all_states = Vec::new();
// Reset reservoir
self.state = Array1::zeros(self.reservoir_size);
for (i, input) in inputs.iter().enumerate() {
let input_arr = Array1::from_vec(input.clone());
let state = self.update_state(&input_arr);
if i >= washout {
// Concatenate state with input (for direct connections)
let mut extended = Vec::from(state.as_slice().unwrap());
extended.extend_from_slice(input);
all_states.push(Array1::from_vec(extended));
}
}
// Convert to matrix for ridge regression
let n_train = all_states.len();
let state_dim = self.reservoir_size + inputs[0].len();
let mut state_matrix = Array2::zeros((n_train, state_dim));
for (i, state) in all_states.iter().enumerate() {
state_matrix.row_mut(i).assign(state);
}
(state_matrix, all_states)
}
/// Train output weights using ridge regression
pub fn train_ridge(&mut self, x: &Vec<Vec<f32>>, y: &Vec<f32>,
regularization: f32) {
let washout = 100.min(x.len() / 10);
let (states, _) = self.collect_states(x, washout);
// Prepare targets (skip washout)
let targets = &y[washout..];
// Ridge regression: W_out = Y * X^T * (X * X^T + λI)^-1
let lambda = regularization;
let n = states.nrows();
let d = states.ncols();
// X * X^T
let gram = states.t().dot(&states);
// Add regularization
let mut reg_gram = gram + Array2::<f32>::eye(d) * lambda;
// Solve using pseudo-inverse (in practice, use better solver)
let y_mat = Array2::from_shape_vec((1, targets.len()), targets.to_vec())
.expect("Shape mismatch");
// Simple solution (production would use LAPACK)
// w_out should be (output_dim x extended_dim) = (1 x (reservoir_size + input_dim))
// y_mat is (1 x n), states is (n x d), so compute Y * pinv(X^T)
let pinv = Self::simple_pinv(&reg_gram); // (d x d)
let temp = states.dot(&pinv); // (n x d)
self.w_out = y_mat.dot(&temp); // (1 x d)
}
/// Simple pseudo-inverse (production code would use LAPACK)
fn simple_pinv(matrix: &Array2<f32>) -> Array2<f32> {
// Simplified: just add strong regularization for stability
let n = matrix.nrows();
let reg = Array2::<f32>::eye(n) * 0.01;
let stabilized = matrix + &reg;
// Return stabilized inverse approximation
// In production, use proper SVD or QR decomposition
stabilized.mapv(|x| 1.0 / (x + 0.001))
}
/// Predict using trained reservoir
pub fn predict(&mut self, x: &[Vec<f32>]) -> Vec<f32> {
// Reset state for prediction
self.state = Array1::zeros(self.reservoir_size);
let mut predictions = Vec::new();
for input in x {
let input_arr = Array1::from_vec(input.clone());
let state = self.update_state(&input_arr);
// Extended state (reservoir + input)
let mut extended = Vec::from(state.as_slice().unwrap());
extended.extend_from_slice(input);
let extended_arr = Array1::from_vec(extended);
// Linear readout
let output = self.w_out.dot(&extended_arr);
predictions.push(output[0]);
}
predictions
}
/// Classify using reservoir
pub fn predict_class(&mut self, x: &[Vec<f32>]) -> Vec<usize> {
let outputs = self.predict(x);
outputs.iter().map(|&val| {
if val < -0.25 { 0 }
else if val > 0.25 { 2 }
else { 1 }
}).collect()
}
}
/// Quantum Reservoir Computing - Exploit quantum superposition
pub struct QuantumReservoir {
pub classical_reservoir: ReservoirComputer,
quantum_layer_size: usize,
phase_matrix: Array2<f32>,
entanglement_strength: f32,
}
impl QuantumReservoir {
pub fn new(input_dim: usize, reservoir_size: usize, output_dim: usize) -> Self {
let classical_reservoir = ReservoirComputer::new(input_dim, reservoir_size, output_dim);
let quantum_layer_size = 16; // Small quantum layer
// Random phase matrix for quantum interference
let mut rng = thread_rng();
let phase_matrix = Array2::from_shape_fn((quantum_layer_size, quantum_layer_size),
|_| rng.gen::<f32>() * 2.0 * PI);
Self {
classical_reservoir,
quantum_layer_size,
phase_matrix,
entanglement_strength: 0.5,
}
}
/// Simulate quantum interference patterns
fn quantum_transform(&self, state: &Array1<f32>) -> Array1<f32> {
let n = self.quantum_layer_size.min(state.len());
let mut quantum_state = Array1::zeros(n);
// Take first n elements
for i in 0..n {
quantum_state[i] = state[i % state.len()];
}
// Apply phase rotations (simulated quantum gates)
let mut result = Array1::<f32>::zeros(n);
for i in 0..n {
for j in 0..n {
let phase = self.phase_matrix[[i, j]];
let amplitude = quantum_state[j] * phase.cos()
+ quantum_state[(j + 1) % n] * phase.sin();
result[i] += amplitude * self.entanglement_strength;
}
}
// Normalize (quantum measurement)
let norm = result.dot(&result).sqrt();
if norm > 0.0 {
result = result / norm;
}
result
}
pub fn train(&mut self, x: &Vec<Vec<f32>>, y: &Vec<f32>) {
// First pass through classical reservoir
self.classical_reservoir.train_ridge(x, y, 0.001);
// Enhance with quantum layer (in practice, would train quantum parameters)
}
pub fn predict_quantum(&mut self, x: &[Vec<f32>]) -> Vec<usize> {
// Classical prediction
let classical_pred = self.classical_reservoir.predict_class(x);
// Quantum enhancement (simplified)
classical_pred.iter().enumerate().map(|(i, &pred)| {
// Add quantum fluctuation based on position
let quantum_factor = (i as f32 * 0.1).sin();
if quantum_factor > 0.3 && pred < 2 {
pred + 1
} else if quantum_factor < -0.3 && pred > 0 {
pred - 1
} else {
pred
}
}).collect()
}
}
@@ -0,0 +1,31 @@
#[cfg(feature = "ruv-fann")]
pub mod ruv_fann_backend {
// Implement this against ruv-fann's API.
// Expected surface:
// - new(input, hidden, output) -> Self
// - train_regression(x, y, epochs, lr)
// - predict_reg(x) -> Vec<f32>
// - predict_cls3(x) -> Vec<usize>
// Keep identical signatures to mlp.rs so the runner can swap backends.
// Example skeleton:
pub struct RuvFannModel { /* fields from ruv-fann */ }
impl RuvFannModel {
pub fn new(_input: usize, _hidden: usize, _output: usize) -> Self { Self { } }
pub fn train_regression(&mut self, _x: &Vec<Vec<f32>>, _y: &Vec<f32>, _epochs: usize, _lr: f32) { /* ... */ }
pub fn predict_reg(&self, _x: &[Vec<f32>]) -> Vec<f32> { vec![] }
pub fn predict_cls3(&self, _x: &[Vec<f32>]) -> Vec<usize> { vec![] }
}
}
#[cfg(not(feature = "ruv-fann"))]
pub mod ruv_fann_backend {
// No-op stub so the crate compiles without the feature.
pub struct RuvFannModel;
impl RuvFannModel {
pub fn new(_: usize, _: usize, _: usize) -> Self { Self }
pub fn train_regression(&mut self, _: &Vec<Vec<f32>>, _: &Vec<f32>, _: usize, _: f32) {}
pub fn predict_reg(&self, _: &[Vec<f32>]) -> Vec<f32> { Vec::new() }
pub fn predict_cls3(&self, _: &[Vec<f32>]) -> Vec<usize> { Vec::new() }
}
}
@@ -0,0 +1,109 @@
#[cfg(feature = "ruv-fann")]
use ruv_fann_dep::{Network, Activation};
#[cfg(feature = "ruv-fann")]
pub struct RuvFannModel {
network: Network,
input_dim: usize,
output_dim: usize,
}
#[cfg(feature = "ruv-fann")]
impl RuvFannModel {
pub fn new(input: usize, hidden: usize, output: usize) -> Self {
// Create a 3-layer network: input -> hidden -> output
let network = Network::new(&[input, hidden, output])
.with_activation(Activation::Sigmoid)
.with_learning_rate(0.7);
Self {
network,
input_dim: input,
output_dim: output,
}
}
pub fn train_regression(&mut self, x: &Vec<Vec<f32>>, y: &Vec<f32>, epochs: usize, lr: f32) {
// Get mutable access to network
let network = Arc::get_mut(&mut self.network)
.expect("Cannot get mutable reference to network");
network.set_learning_rate(lr);
// Prepare training data in FANN format
let mut inputs: Vec<Vec<f32>> = Vec::new();
let mut outputs: Vec<Vec<f32>> = Vec::new();
for (xi, yi) in x.iter().zip(y) {
inputs.push(xi.clone());
if self.output_dim == 1 {
outputs.push(vec![*yi]);
} else {
// For classification, create one-hot encoding
let class = if *yi < -0.25 { 0 }
else if *yi > 0.25 { 2 }
else { 1 };
let mut one_hot = vec![0.0; self.output_dim];
if class < self.output_dim {
one_hot[class] = 1.0;
}
outputs.push(one_hot);
}
}
// Train using batch training
for _ in 0..epochs {
for (input, output) in inputs.iter().zip(&outputs) {
network.train(input, output);
}
}
}
pub fn predict_reg(&self, x: &[Vec<f32>]) -> Vec<f32> {
x.iter().map(|xi| {
let output = self.network.run(xi).expect("Failed to run network");
if self.output_dim == 1 {
output[0]
} else {
// Return first output for regression
output[0]
}
}).collect()
}
pub fn predict_cls3(&self, x: &[Vec<f32>]) -> Vec<usize> {
x.iter().map(|xi| {
let output = self.network.run(xi).expect("Failed to run network");
if self.output_dim >= 3 {
// Find argmax for classification
let mut best = 0;
let mut best_val = output[0];
for i in 1..3.min(output.len()) {
if output[i] > best_val {
best_val = output[i];
best = i;
}
}
best
} else {
// Threshold-based classification for single output
let val = output[0];
if val < -0.25 { 0 }
else if val > 0.25 { 2 }
else { 1 }
}
}).collect()
}
}
#[cfg(not(feature = "ruv-fann"))]
pub struct RuvFannModel;
#[cfg(not(feature = "ruv-fann"))]
impl RuvFannModel {
pub fn new(_: usize, _: usize, _: usize) -> Self { Self }
pub fn train_regression(&mut self, _: &Vec<Vec<f32>>, _: &Vec<f32>, _: usize, _: f32) {}
pub fn predict_reg(&self, _: &[Vec<f32>]) -> Vec<f32> { Vec::new() }
pub fn predict_cls3(&self, _: &[Vec<f32>]) -> Vec<usize> { Vec::new() }
}
@@ -0,0 +1,397 @@
use ndarray::{Array1, Array2};
use rand::{thread_rng, Rng};
use std::collections::HashMap;
/// Sparse Neural Network with dynamic sparsity
/// Only keeps top-k% connections active
pub struct SparseNetwork {
// Sparse weight storage
w1_indices: Vec<(usize, usize)>, // Active connections
w1_values: Vec<f32>, // Weight values
w2_indices: Vec<(usize, usize)>,
w2_values: Vec<f32>,
// Biases (dense)
b1: Array1<f32>,
b2: Array1<f32>,
// Dimensions
input_dim: usize,
hidden_dim: usize,
output_dim: usize,
sparsity: f32, // Fraction of weights to keep
// Hidden activations
hidden: Array1<f32>,
// Pruning statistics
pruning_threshold: f32,
pruned_count: usize,
}
impl SparseNetwork {
pub fn new(input: usize, hidden: usize, output: usize, sparsity: f32) -> Self {
let mut rng = thread_rng();
// Initialize with random sparse connections
let n_connections1 = ((input * hidden) as f32 * sparsity) as usize;
let n_connections2 = ((hidden * output) as f32 * sparsity) as usize;
let mut w1_indices = Vec::new();
let mut w1_values = Vec::new();
// Random sparse initialization for layer 1
let scale1 = (2.0 / input as f32).sqrt();
let mut used1 = std::collections::HashSet::new();
while w1_indices.len() < n_connections1 {
let i = rng.gen_range(0..hidden);
let j = rng.gen_range(0..input);
if used1.insert((i, j)) {
w1_indices.push((i, j));
w1_values.push(rng.gen::<f32>() * scale1 - scale1/2.0);
}
}
// Layer 2
let mut w2_indices = Vec::new();
let mut w2_values = Vec::new();
let scale2 = (2.0 / hidden as f32).sqrt();
let mut used2 = std::collections::HashSet::new();
while w2_indices.len() < n_connections2 {
let i = rng.gen_range(0..output);
let j = rng.gen_range(0..hidden);
if used2.insert((i, j)) {
w2_indices.push((i, j));
w2_values.push(rng.gen::<f32>() * scale2 - scale2/2.0);
}
}
Self {
w1_indices,
w1_values,
w2_indices,
w2_values,
b1: Array1::zeros(hidden),
b2: Array1::zeros(output),
input_dim: input,
hidden_dim: hidden,
output_dim: output,
sparsity,
hidden: Array1::zeros(hidden),
pruning_threshold: 0.01,
pruned_count: 0,
}
}
/// Sparse forward pass
pub fn forward(&mut self, x: &[f32]) -> Vec<f32> {
// Reset hidden
self.hidden.fill(0.0);
// Sparse matrix-vector multiplication for layer 1
for ((&(i, j), &val)) in self.w1_indices.iter().zip(&self.w1_values) {
if j < x.len() {
self.hidden[i] += x[j] * val;
}
}
// Add bias and apply ReLU
self.hidden = &self.hidden + &self.b1;
self.hidden.mapv_inplace(|x| x.max(0.0));
// Layer 2
let mut output = vec![0.0; self.output_dim];
for ((&(i, j), &val)) in self.w2_indices.iter().zip(&self.w2_values) {
output[i] += self.hidden[j] * val;
}
// Add bias
for i in 0..self.output_dim {
output[i] += self.b2[i];
}
output
}
/// Dynamic pruning - remove small weights
pub fn prune_weights(&mut self, threshold: f32) {
// Prune layer 1
let mut new_indices1 = Vec::new();
let mut new_values1 = Vec::new();
for (idx, val) in self.w1_indices.iter().zip(&self.w1_values) {
if val.abs() > threshold {
new_indices1.push(*idx);
new_values1.push(*val);
}
}
let pruned1 = self.w1_indices.len() - new_indices1.len();
self.w1_indices = new_indices1;
self.w1_values = new_values1;
// Prune layer 2
let mut new_indices2 = Vec::new();
let mut new_values2 = Vec::new();
for (idx, val) in self.w2_indices.iter().zip(&self.w2_values) {
if val.abs() > threshold {
new_indices2.push(*idx);
new_values2.push(*val);
}
}
let pruned2 = self.w2_indices.len() - new_indices2.len();
self.w2_indices = new_indices2;
self.w2_values = new_values2;
self.pruned_count += pruned1 + pruned2;
self.pruning_threshold = threshold;
}
/// Regrow connections (lottery ticket hypothesis)
pub fn regrow_connections(&mut self, n_regrow: usize) {
let mut rng = thread_rng();
// Regrow layer 1
let mut used1: std::collections::HashSet<_> = self.w1_indices.iter().cloned().collect();
let scale1 = (2.0 / self.input_dim as f32).sqrt();
for _ in 0..n_regrow/2 {
let i = rng.gen_range(0..self.hidden_dim);
let j = rng.gen_range(0..self.input_dim);
if used1.insert((i, j)) {
self.w1_indices.push((i, j));
self.w1_values.push(rng.gen::<f32>() * scale1 - scale1/2.0);
}
}
// Regrow layer 2
let mut used2: std::collections::HashSet<_> = self.w2_indices.iter().cloned().collect();
let scale2 = (2.0 / self.hidden_dim as f32).sqrt();
for _ in 0..n_regrow/2 {
let i = rng.gen_range(0..self.output_dim);
let j = rng.gen_range(0..self.hidden_dim);
if used2.insert((i, j)) {
self.w2_indices.push((i, j));
self.w2_values.push(rng.gen::<f32>() * scale2 - scale2/2.0);
}
}
}
/// Train with sparse gradient updates
pub fn train(&mut self, x: &[Vec<f32>], y: &[f32], epochs: usize, lr: f32) {
for epoch in 0..epochs {
let mut total_loss = 0.0;
for (xi, &yi) in x.iter().zip(y.iter()) {
let output = self.forward(xi);
let pred = if self.output_dim == 1 { output[0] } else { output[0] };
let error = pred - yi;
total_loss += error * error;
// Sparse backpropagation
self.sparse_backward(xi, error, lr);
}
// Dynamic sparsity - prune and regrow periodically
if epoch > 0 && epoch % 10 == 0 {
let old_count = self.w1_indices.len() + self.w2_indices.len();
self.prune_weights(self.pruning_threshold);
let pruned = old_count - (self.w1_indices.len() + self.w2_indices.len());
if pruned > 0 {
self.regrow_connections(pruned);
}
}
if epoch % 100 == 0 {
println!("Sparse epoch {}: loss={:.6}, active_weights={}",
epoch, total_loss / x.len() as f32,
self.w1_indices.len() + self.w2_indices.len());
}
}
}
fn sparse_backward(&mut self, x: &[f32], error: f32, lr: f32) {
// Output gradient
let grad_out = error;
// Update layer 2 weights (sparse)
let mut grad_hidden = Array1::zeros(self.hidden_dim);
for i in 0..self.w2_indices.len() {
let (out_idx, hid_idx) = self.w2_indices[i];
if out_idx == 0 { // For single output
// Update weight
self.w2_values[i] -= lr * grad_out * self.hidden[hid_idx];
// Accumulate gradient for hidden layer
grad_hidden[hid_idx] += self.w2_values[i] * grad_out;
}
}
// Update bias
self.b2[0] -= lr * grad_out;
// Apply ReLU gradient
grad_hidden.mapv_inplace(|g| if self.hidden[0] > 0.0 { g } else { 0.0 });
// Update layer 1 weights (sparse)
for i in 0..self.w1_indices.len() {
let (hid_idx, in_idx) = self.w1_indices[i];
if in_idx < x.len() {
self.w1_values[i] -= lr * grad_hidden[hid_idx] * x[in_idx];
}
}
// Update hidden bias
self.b1 = &self.b1 - &grad_hidden * lr;
}
pub fn predict(&mut self, x: &[Vec<f32>]) -> Vec<f32> {
x.iter().map(|xi| {
let output = self.forward(xi);
if self.output_dim == 1 { output[0] } else { output[0] }
}).collect()
}
pub fn predict_class(&mut self, x: &[Vec<f32>]) -> Vec<usize> {
self.predict(x).iter().map(|&y| {
if y < -0.25 { 0 }
else if y > 0.25 { 2 }
else { 1 }
}).collect()
}
pub fn get_sparsity_stats(&self) -> (usize, usize, f32) {
let active = self.w1_indices.len() + self.w2_indices.len();
let total = self.input_dim * self.hidden_dim + self.hidden_dim * self.output_dim;
let sparsity = active as f32 / total as f32;
(active, self.pruned_count, sparsity)
}
}
/// Lottery Ticket Network - finds winning sparse subnetworks
pub struct LotteryTicketNetwork {
base_network: SparseNetwork,
initial_weights1: HashMap<(usize, usize), f32>,
initial_weights2: HashMap<(usize, usize), f32>,
winning_mask1: HashMap<(usize, usize), bool>,
winning_mask2: HashMap<(usize, usize), bool>,
iteration: usize,
}
impl LotteryTicketNetwork {
pub fn new(input: usize, hidden: usize, output: usize) -> Self {
let base = SparseNetwork::new(input, hidden, output, 1.0); // Start dense
// Store initial weights
let mut initial_weights1 = HashMap::new();
for (&idx, &val) in base.w1_indices.iter().zip(&base.w1_values) {
initial_weights1.insert(idx, val);
}
let mut initial_weights2 = HashMap::new();
for (&idx, &val) in base.w2_indices.iter().zip(&base.w2_values) {
initial_weights2.insert(idx, val);
}
Self {
base_network: base,
initial_weights1,
initial_weights2,
winning_mask1: HashMap::new(),
winning_mask2: HashMap::new(),
iteration: 0,
}
}
/// Iterative magnitude pruning to find winning tickets
pub fn find_winning_ticket(&mut self, x: &[Vec<f32>], y: &[f32],
prune_rate: f32, iterations: usize) {
for iter in 0..iterations {
println!("Lottery iteration {}/{}", iter + 1, iterations);
// Reset to initial weights with current mask
self.reset_to_initial();
// Train
self.base_network.train(x, y, 100, 0.01);
// Prune based on magnitude
self.magnitude_prune(prune_rate);
self.iteration += 1;
}
}
fn reset_to_initial(&mut self) {
// Reset to initial weights, keeping only winning connections
let mut new_indices1 = Vec::new();
let mut new_values1 = Vec::new();
for (&idx, &init_val) in &self.initial_weights1 {
if self.iteration == 0 || self.winning_mask1.get(&idx) == Some(&true) {
new_indices1.push(idx);
new_values1.push(init_val);
}
}
self.base_network.w1_indices = new_indices1;
self.base_network.w1_values = new_values1;
// Layer 2
let mut new_indices2 = Vec::new();
let mut new_values2 = Vec::new();
for (&idx, &init_val) in &self.initial_weights2 {
if self.iteration == 0 || self.winning_mask2.get(&idx) == Some(&true) {
new_indices2.push(idx);
new_values2.push(init_val);
}
}
self.base_network.w2_indices = new_indices2;
self.base_network.w2_values = new_values2;
}
fn magnitude_prune(&mut self, prune_rate: f32) {
// Collect all weight magnitudes
let mut magnitudes: Vec<f32> = self.base_network.w1_values.iter()
.chain(&self.base_network.w2_values)
.map(|v| v.abs())
.collect();
magnitudes.sort_by(|a, b| a.partial_cmp(b).unwrap());
let cutoff_idx = (magnitudes.len() as f32 * prune_rate) as usize;
let threshold = if cutoff_idx < magnitudes.len() {
magnitudes[cutoff_idx]
} else {
0.0
};
// Update masks
self.winning_mask1.clear();
for (&idx, &val) in self.base_network.w1_indices.iter()
.zip(&self.base_network.w1_values) {
self.winning_mask1.insert(idx, val.abs() > threshold);
}
self.winning_mask2.clear();
for (&idx, &val) in self.base_network.w2_indices.iter()
.zip(&self.base_network.w2_values) {
self.winning_mask2.insert(idx, val.abs() > threshold);
}
}
pub fn predict(&mut self, x: &[Vec<f32>]) -> Vec<f32> {
self.base_network.predict(x)
}
pub fn predict_class(&mut self, x: &[Vec<f32>]) -> Vec<usize> {
self.base_network.predict_class(x)
}
}