Struct SAM 

pub struct SAM<A, O, D>
where
    A: Float + ScalarOperand + Debug,
    O: Optimizer<A, D> + Clone,
    D: Dimension,
{ /* private fields */ }

Sharpness-Aware Minimization (SAM) optimizer

SAM is an optimization technique that seeks parameters that lie in neighborhoods having uniformly low loss values, which improves generalization. It achieves this by performing a two-step update process:

1. Compute and take a step in the direction of the “sharpness” gradient (perturbed parameters)
2. Compute the gradient at these perturbed parameters and use it to update the original parameters

This implementation wraps around a base optimizer and modifies its behavior to implement the SAM algorithm.

Parameters

• inner_optimizer - The optimizer to use for the parameter updates
• rho - The neighborhood size for perturbation (default: 0.05)
• epsilon - Small constant for numerical stability (default: 1e-12)
• adaptive - Whether to use adaptive perturbation size (SAM-A) (default: false)

Example

use scirs2_core::ndarray::Array1;
use optirs_core::optimizers::{SAM, SGD};
use optirs_core::Optimizer;

// Create a base optimizer
let sgd = SGD::new(0.1);

// Wrap it with SAM
let mut optimizer = SAM::new(sgd);

// First step to compute perturbed parameters and store perturbed gradients
let params = Array1::zeros(10);
let gradients = Array1::ones(10);
let (perturbed_params_) = optimizer.first_step(&params, &gradients).unwrap();

// Second step to update original parameters using gradients at perturbed parameters
// Normally, you would compute new gradients at perturbed_params
let new_gradients = Array1::ones(10) * 0.5; // Example new gradients
let updated_params = optimizer.second_step(&params, &new_gradients).unwrap();

Implementations

impl<A, O, D> SAM<A, O, D>

where A: Float + ScalarOperand + Debug, O: Optimizer<A, D> + Clone, D: Dimension,

pub fn new(inner_optimizer: O) -> Self

Creates a new SAM optimizer with the given inner optimizer and default settings

pub fn with_config(inner_optimizer: O, rho: A, adaptive: bool) -> Self

Creates a new SAM optimizer with the given inner optimizer and configuration

pub fn with_rho(self, rho: A) -> Self

Set the rho parameter (neighborhood size)

pub fn with_epsilon(self, epsilon: A) -> Self

Set the epsilon parameter (numerical stability)

pub fn with_adaptive(self, adaptive: bool) -> Self

Set whether to use adaptive perturbation size (SAM-A)

pub fn inner_optimizer(&self) -> &O

Get the inner optimizer

pub fn inner_optimizer_mut(&mut self) -> &mut O

Get a mutable reference to the inner optimizer

pub fn rho(&self) -> A

Get the rho parameter

pub fn epsilon(&self) -> A

Get the epsilon parameter

pub fn is_adaptive(&self) -> bool

Check if using adaptive perturbation size

pub fn first_step( &mut self, params: &Array<A, D>, gradients: &Array<A, D>, ) -> Result<(Array<A, D>, A)>

First step of SAM: compute perturbed parameters by moving in the direction of the gradient

Arguments

• params - Current parameters
• gradients - Gradients of the loss with respect to the parameters

Returns

Tuple containing (perturbed_parameters, perturbation_size)

pub fn second_step( &mut self, params: &Array<A, D>, gradients: &Array<A, D>, ) -> Result<Array<A, D>>

Second step of SAM: update the original parameters using gradients at the perturbed point

Arguments

• _params - Original parameters (used for validation)
• gradients - Gradients of the loss with respect to the perturbed parameters

Returns

Updated parameters after applying the “sharpness-aware” update

pub fn reset(&mut self)

Reset the internal state

Trait Implementations

impl<A, O, D> Clone for SAM<A, O, D>

where A: Float + ScalarOperand + Debug, O: Optimizer<A, D> + Clone, D: Dimension,

fn clone(&self) -> Self

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl<A, O, D> Debug for SAM<A, O, D>

where A: Float + ScalarOperand + Debug, O: Optimizer<A, D> + Clone + Debug, D: Dimension,

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

Formats the value using the given formatter.

impl<A, O, D> Optimizer<A, D> for SAM<A, O, D>

where A: Float + ScalarOperand + Debug + Send + Sync, O: Optimizer<A, D> + Clone + Send + Sync, D: Dimension,

fn step( &mut self, params: &Array<A, D>, gradients: &Array<A, D>, ) -> Result<Array<A, D>>

Updates parameters using the given gradients

fn set_learning_rate(&mut self, learning_rate: A)

Sets a new learning rate

fn get_learning_rate(&self) -> A

Gets the current learning rate

fn step_list( &mut self, params_list: &[&Array<A, D>], gradients_list: &[&Array<A, D>], ) -> Result<Vec<Array<A, D>>>

Updates multiple parameter arrays at once