Struct AdamW 

pub struct AdamW { /* private fields */ }

AdamW optimizer

AdamW decouples weight decay from the gradient-based update, making it more effective than L2 regularization. Instead of adding weight decay to the gradient, it applies weight decay directly to the parameters.

Standard Adam with L2: θ_t = θ_{t-1} - lr * (m_t / (√v_t + ε) + λ * θ_{t-1}) AdamW: θ_t = (1 - lr * λ) * θ_{t-1} - lr * m_t / (√v_t + ε)

Implementations

impl AdamW

pub fn new( lr: f32, beta1: f32, beta2: f32, epsilon: f32, weight_decay: f32, ) -> Self

Create a new AdamW optimizer

pub fn default_params(lr: f32) -> Self

Create AdamW with default parameters (weight_decay = 0.01)

pub fn step_count(&self) -> u64

Get optimizer step counter.

pub fn set_step_count(&mut self, t: u64)

Set optimizer step counter (for checkpoint resume).

pub fn first_moments(&self) -> &[Option<Array1<f32>>]

Get first moment buffers (m) as f32 slices.

pub fn second_moments(&self) -> &[Option<Array1<f32>>]

Get second moment buffers (v) as f32 slices.

pub fn set_first_moment(&mut self, idx: usize, data: Array1<f32>)

Set first moment buffer at index.

pub fn set_second_moment(&mut self, idx: usize, data: Array1<f32>)

Set second moment buffer at index.

pub fn beta1(&self) -> f32

Get beta1 hyperparameter.

pub fn beta2(&self) -> f32

Get beta2 hyperparameter.

pub fn weight_decay(&self) -> f32

Get weight decay hyperparameter.

Trait Implementations

impl Optimizer for AdamW

fn step(&mut self, params: &mut [Tensor])

Perform a single optimization step

fn step_refs(&mut self, params: &mut [&mut Tensor])

Perform optimization step on referenced parameters

fn lr(&self) -> f32

Get learning rate

fn set_lr(&mut self, lr: f32)

Set learning rate

fn zero_grad(&mut self, params: &mut [Tensor])

Zero out all gradients

fn zero_grad_refs(&mut self, params: &mut [&mut Tensor])

Zero gradients on referenced parameters