Struct ScheduleFreeAdamW 

pub struct ScheduleFreeAdamW { /* private fields */ }

Schedule-free AdamW optimizer.

Maintains both training sequence (x) and evaluation sequence (y). During training, gradients are computed w.r.t. x_t. During evaluation, use y_t for better generalization.

Algorithm

# Training step:
m_t = β₁ * m_{t-1} + (1 - β₁) * g_t
v_t = β₂ * v_{t-1} + (1 - β₂) * g_t²
x_t = x_{t-1} - lr * m_t / (√v_t + ε) - lr * λ * x_{t-1}  # weight decay

# Evaluation sequence (exponential moving average):
y_t = (1 - γ) * x_t + γ * y_{t-1}

# At test time, use y_t instead of x_t

Example

use tensorlogic_train::{ScheduleFreeAdamW, ScheduleFreeConfig};
use scirs2_core::ndarray::Array2;
use std::collections::HashMap;

let config = ScheduleFreeConfig::default()
    .with_lr(0.001)
    .with_warmup_steps(1000);

let mut optimizer = ScheduleFreeAdamW::new(config);

// Training mode - use training parameters
optimizer.set_training_mode(true);

// ... compute gradients ...

// Evaluation mode - switch to averaged parameters
optimizer.set_training_mode(false);
let eval_params = optimizer.get_eval_parameters();

Implementations

impl ScheduleFreeAdamW

pub fn new(config: ScheduleFreeConfig) -> Self

Create a new Schedule-Free AdamW optimizer.

pub fn set_training_mode(&mut self, training: bool)

Set training mode.

When training_mode = true, use train_params (x_t) for gradients. When training_mode = false, use eval_params (y_t) for evaluation.

pub fn is_training(&self) -> bool

Get current training mode.

pub fn get_eval_parameters(&self) -> &HashMap<String, Array2<f64>>

Get evaluation parameters (for inference).

pub fn get_train_parameters(&self) -> &HashMap<String, Array2<f64>>

Get training parameters (for gradient computation).

Trait Implementations

impl Clone for ScheduleFreeAdamW

fn clone(&self) -> ScheduleFreeAdamW

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for ScheduleFreeAdamW

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Optimizer for ScheduleFreeAdamW

fn zero_grad(&mut self)

Zero all gradients.

fn get_lr(&self) -> f64

Get current learning rate.

fn set_lr(&mut self, lr: f64)

Set learning rate.

fn step( &mut self, parameters: &mut HashMap<String, Array2<f64>>, gradients: &HashMap<String, Array2<f64>>, ) -> TrainResult<()>

Update parameters with computed gradients.

fn state_dict(&self) -> HashMap<String, Vec<f64>>

Get optimizer state for checkpointing.

fn load_state_dict(&mut self, _state: HashMap<String, Vec<f64>>)

Load optimizer state from checkpoint.