vsa-optim-rs

Deterministic training optimization using Vector Symbolic Architecture (VSA), ternary quantization, and closed-form gradient prediction.

A pure Rust implementation enabling efficient large model fine-tuning on consumer hardware through mathematically principled gradient compression and prediction.

Key Properties

• Deterministic: Identical inputs produce identical outputs—no stochastic variance in predictions
• Closed-form: Weighted least squares with Cramer's rule—no iterative optimization
• Memory-efficient: ~90% gradient storage reduction via VSA compression
• Compute-efficient: ~80% backward pass reduction via gradient prediction

Installation

[dependencies]
vsa-optim-rs = "0.1"

Quick Start

Deterministic Phase Training (Recommended)

The DeterministicPhaseTrainer orchestrates training through mathematically rigorous phases with guaranteed reproducibility:

use vsa_optim_rs::{DeterministicPhaseTrainer, DeterministicPhaseConfig, DeterministicPhase};
use candle_core::Device;
use std::collections::HashMap;

// Define parameter shapes
let shapes = vec![
    ("layer1.weight".into(), vec![768, 768]),
    ("layer2.weight".into(), vec![768, 3072]),
];

// Configure deterministic training
let config = DeterministicPhaseConfig {
    warmup_steps: 10,       // Initial gradient collection
    full_steps: 5,          // Full computation per cycle
    predict_steps: 20,      // Predicted steps per cycle
    correct_every: 5,       // Correction frequency
    adaptive_phases: true,  // Auto-adjust on loss increase
    ..Default::default()
};

let mut trainer = DeterministicPhaseTrainer::new(&shapes, config, &Device::Cpu)?;

// Training loop
for step in 0..100 {
    let info = trainer.begin_step()?;
    
    match info.phase {
        DeterministicPhase::Warmup | DeterministicPhase::Full | DeterministicPhase::Correct => {
            // Compute gradients via backpropagation
            let gradients = compute_gradients(&model, &batch);
            trainer.record_full_gradients(&gradients)?;
        }
        DeterministicPhase::Predict => {
            // Use deterministically predicted gradients (no backward pass)
            let gradients = trainer.get_predicted_gradients()?;
            apply_gradients(&mut model, &gradients);
        }
    }
    
    trainer.end_step(loss)?;
}

let stats = trainer.get_stats();
println!("Speedup: {:.2}x ({} full, {} predicted)", 
    stats.speedup, stats.full_steps, stats.predicted_steps);

VSA Gradient Compression

Compress gradients using hyperdimensional computing with bind/bundle/unbind operations:

use vsa_optim_rs::{VSAGradientCompressor, VSAConfig};

let config = VSAConfig::builder()
    .dimension(8192)          // Hypervector dimension
    .compression_ratio(0.1)   // 10x compression target
    .seed(42)                 // Reproducible projections
    .build();

let param_shapes = vec![
    ("weight".into(), vec![1024, 1024]),
];

let mut compressor = VSAGradientCompressor::new(&param_shapes, config, &device)?;

// Compress gradients
let compressed = compressor.compress(&gradients)?;
println!("Compression: {:.1}x", compressed.stats.compression_ratio);

// Decompress when needed
let restored = compressor.decompress(&compressed)?;

Ternary Gradient Accumulation

Memory-efficient accumulation using balanced ternary {-1, 0, +1}:

use vsa_optim_rs::{TernaryGradientAccumulator, TernaryConfig};

let config = TernaryConfig::builder()
    .accumulation_steps(8)
    .use_stochastic_rounding(true)  // Unbiased quantization
    .build();

let mut accumulator = TernaryGradientAccumulator::new(&param_shapes, config, &device)?;

for micro_batch in micro_batches {
    let gradients = compute_gradients(&model, &micro_batch);
    accumulator.accumulate(&gradients)?;  // ~93% memory savings
}

// Retrieve accumulated gradients for optimizer step
let accumulated = accumulator.get_accumulated()?;
optimizer.step(&accumulated)?;
accumulator.reset()?;

Architecture

Deterministic Gradient Prediction

The core innovation: predict gradients using weighted least squares model fitting with a closed-form solution (no iterative optimization):

Gradient Model: g(t) = baseline + velocity × t + residual

Where:
  - baseline: Weighted mean of historical gradients
  - velocity: Gradient change rate (fitted via normal equations)
  - residual: Exponentially-averaged prediction error for drift correction

Warmup Phase: Collect initial gradient samples to establish prediction baseline.

Prediction Fitting: Solve normal equations using Cramer's rule:

[Σw    Σwt  ] [b]   [Σwg   ]
[Σwt   Σwt²] [v] = [Σwtg  ]

Residual Tracking: Maintain exponentially-decayed average of prediction errors to correct systematic drift without stochastic noise.

Training Phase Cycle

┌─────────────────────────────────────────────────────────────────┐
│                                                                 │
│  WARMUP ──► FULL ──► PREDICT ──► CORRECT ──► FULL ──► ...      │
│  (N steps)  (M)      (P steps)    (periodic)  (M)               │
│                         │              │                        │
│                         └──────────────┘                        │
│                         (correction cycle)                      │
│                                                                 │
└─────────────────────────────────────────────────────────────────┘

VSA Compression Pipeline

Gradients ──► Project to HD ──► Bind with keys ──► Bundle (majority) ──► Compressed
                                                                              │
Decompressed ◄── Inverse Project ◄── Unbind with keys ◄────────────────────┘

Operations leverage the quasi-orthogonality of random vectors in high dimensions (Johnson-Lindenstrauss lemma) for information-preserving compression.

Performance Characteristics

Speedup Analysis

With default configuration (warmup=10, full=5, predict=20, correct_every=5):

100 steps = 10 warmup + (5 full + 20 predict) × cycles
         = 10 warmup + ~25 full + ~65 predict
         ≈ 35 backward passes instead of 100
         = 2.9x theoretical speedup

Actual speedup depends on backward pass cost relative to forward pass.

Configuration Reference

DeterministicPhaseConfig

DeterministicPhaseConfig {
    warmup_steps: 10,              // Steps to collect initial gradients
    full_steps: 5,                 // Full gradient steps per cycle
    predict_steps: 20,             // Predicted steps per cycle
    correct_every: 5,              // Correction frequency during predict
    adaptive_phases: true,         // Auto-adjust on loss increase
    loss_increase_threshold: 0.1,  // Threshold to trigger adaptation
    history_window: 8,             // Gradients to keep for model fitting
    prediction_horizon: 1,         // Steps ahead to predict
    history_decay: 0.95,           // Exponential decay for weighting
    residual_threshold: 0.1,       // When to apply residual correction
}

VSAConfig

VSAConfig::builder()
    .dimension(8192)          // HD space dimension (↑ = better reconstruction)
    .compression_ratio(0.1)   // Target compression factor
    .seed(42)                 // RNG seed for reproducibility
    .build()

TernaryConfig

TernaryConfig::builder()
    .accumulation_steps(8)         // Micro-batches per optimizer step
    .use_stochastic_rounding(true) // Unbiased quantization to {-1, 0, +1}
    .build()

Requirements

• Rust: 1.92+ (2021 edition)
• Dependencies: candle-core 0.9+, trit-vsa 0.1+

Integration with axolotl-rs

For YAML-driven LLM fine-tuning with automatic VSA acceleration:

use axolotl_rs::{VSAAccelerator, VSAAcceleratorConfig};

let config = VSAAcceleratorConfig::default();  // Or ::conservative(), ::aggressive()
let mut accel = VSAAccelerator::new(&trainable_params, config, &device)?;

for batch in dataloader {
    let info = accel.begin_step()?;
    
    if info.needs_backward {
        loss.backward();
        accel.record_gradients(&trainable_params)?;
    } else {
        let grads = accel.get_predicted_gradients()?;
        // Apply predicted gradients
    }
    
    accel.end_step(loss_value)?;
}

println!("{}", accel.get_stats());  // "VSA: 100 steps (35 full, 65 predicted), 2.86x speedup"

Sister Crates

License

MIT License. See LICENSE-MIT for details.

References

• Kanerva, P. (2009). Hyperdimensional Computing: An Introduction to Computing in Distributed Representation with High-Dimensional Random Vectors.
• Rahimi, A. et al. (2016). High-Dimensional Computing as a Nanoscalable Paradigm.
• Johnson, W. & Lindenstrauss, J. (1984). Extensions of Lipschitz mappings into a Hilbert space.
• Ma, S. et al. (2024). The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits.

"Simplicity is the ultimate sophistication." — Leonardo da Vinci