Struct SamOptimizer 

pub struct SamOptimizer<O: Optimizer> { /* private fields */ }

SAM optimizer (Sharpness Aware Minimization).

SAM seeks parameters that lie in neighborhoods having uniformly low loss, improving model generalization. It requires two forward-backward passes per step: one to compute the adversarial perturbation, and one to compute the actual gradient.

Reference: Foret et al. “Sharpness-Aware Minimization for Efficiently Improving Generalization” (ICLR 2021)

Note: This is a wrapper optimizer. SAM requires special handling in the training loop to perform two gradient computations per step. The typical usage is:

1. Compute gradients at current parameters
2. Compute adversarial perturbation
3. Compute gradients at perturbed parameters
4. Update with the perturbed gradients

Implementations

impl<O: Optimizer> SamOptimizer<O>

pub fn new(base_optimizer: O, rho: f64) -> TrainResult<Self>

Create a new SAM optimizer.

Arguments

• base_optimizer - The base optimizer to use (SGD, Adam, etc.)
• rho - Perturbation radius (typically 0.05)

pub fn first_step( &mut self, parameters: &mut HashMap<String, Array<f64, Ix2>>, gradients: &HashMap<String, Array<f64, Ix2>>, ) -> TrainResult<()>

Compute adversarial perturbations.

This should be called with the first set of gradients to compute the perturbation direction.

pub fn second_step( &mut self, parameters: &mut HashMap<String, Array<f64, Ix2>>, gradients: &HashMap<String, Array<f64, Ix2>>, ) -> TrainResult<()>

Perform the actual optimization step.

This should be called with the second set of gradients (computed at the perturbed parameters). It will remove the perturbations and update the parameters using the base optimizer.

Trait Implementations

impl<O: Debug + Optimizer> Debug for SamOptimizer<O>

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

Formats the value using the given formatter.

impl<O: Optimizer> Optimizer for SamOptimizer<O>

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

Update parameters with computed gradients.

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 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.