Struct GradientPredictor 

pub struct GradientPredictor { /* private fields */ }

Predict future gradients from history.

Gradient prediction reduces compute by ~80% (4 predicted steps per 1 computed step) while maintaining convergence quality through periodic correction cycles.

The predictor maintains a history of recent gradients and uses a momentum-based extrapolation to predict future gradients. Corrections are computed as the difference between predicted and actual gradients.

Example

use vsa_optim_rs::prediction::GradientPredictor;
use vsa_optim_rs::PredictionConfig;

let shapes = vec![
    ("layer1.weight".to_string(), vec![64, 128]),
];
let mut predictor = GradientPredictor::new(&shapes, PredictionConfig::default(), &Device::Cpu)?;

// Training loop
for step in 0..total_steps {
    if predictor.should_compute_full() {
        // loss.backward() - compute full gradients
        predictor.record_gradient(&gradients)?;
        predictor.apply_correction(&mut gradients);
    } else {
        let predicted = predictor.predict_gradient()?;
        // Use predicted gradients for optimizer step
    }
}

Implementations

impl GradientPredictor

pub fn new( param_shapes: &[(String, Vec<usize>)], config: PredictionConfig, device: &Device, ) -> Result<GradientPredictor, OptimError>

Create a new gradient predictor.

Arguments

• param_shapes - List of (name, shape) tuples for parameters
• config - Prediction configuration
• device - Device for tensor storage

Errors

Returns error if initialization fails.

pub fn should_compute_full(&self) -> bool

Check if full gradient computation is needed.

Full computation is needed:

1. At the start (insufficient history)
2. After prediction_steps predicted steps (correction cycle)
3. When prediction quality degrades below threshold

pub fn record_gradient( &mut self, gradients: &HashMap<String, Tensor>, ) -> Result<(), OptimError>

Record current gradients to history.

Called after full gradient computation to update history.

Arguments

• gradients - Map of parameter names to gradient tensors

Errors

Returns error if tensor cloning fails.

pub fn predict_gradient( &mut self, ) -> Result<HashMap<String, Tensor>, OptimError>

Predict gradients based on history.

Uses momentum-based extrapolation from gradient history:

g_pred = g[-1] + momentum * (g[-1] - g[-2])

Returns

Dictionary mapping parameter names to predicted gradients.

Errors

Returns error if tensor operations fail.

pub fn compute_correction( &mut self, actual_gradients: &HashMap<String, Tensor>, ) -> Result<HashMap<String, Tensor>, OptimError>

Compute correction between predicted and actual gradients.

The correction term captures the prediction error and is accumulated to apply a “catch-up” adjustment.

Arguments

• actual_gradients - The actual computed gradients

Returns

Dictionary of correction terms.

Errors

Returns error if tensor operations fail.

pub fn apply_correction( &mut self, gradients: &mut HashMap<String, Tensor>, ) -> Result<(), OptimError>

Apply accumulated corrections to gradients.

After computing full gradients, adds the accumulated correction to account for prediction errors from previous predicted steps.

Arguments

• gradients - Mutable map of gradients to modify in-place

Errors

Returns error if tensor operations fail.

pub fn get_stats(&self) -> PredictorStats

Get prediction statistics.

pub const fn total_steps(&self) -> usize

Get total steps.

pub fn reset(&mut self)

Reset predictor state.