Struct AdamW

pub struct AdamW<F: Float + ScalarOperand + Debug> { /* private fields */ }

AdamW optimizer for neural networks

Implements the AdamW optimizer from the paper: “Decoupled Weight Decay Regularization” by Loshchilov and Hutter (2017).

AdamW is a variant of Adam that correctly implements weight decay regularization, which helps improve generalization and training stability.

The key difference from Adam is that weight decay is applied directly to the weights rather than to the gradients.

Examples

use ndarray::Array1;
use scirs2_neural::optimizers::{AdamW, Optimizer};

// Create a simple AdamW optimizer with default parameters
let mut adamw = AdamW::<f64>::default_with_lr(0.001).unwrap();

// or with custom configuration
let mut adamw_custom = AdamW::new(0.001, 0.9, 0.999, 1e-8, 0.01);

Implementations

impl<F: Float + ScalarOperand + Debug> AdamW<F>

pub fn new( learning_rate: F, beta1: F, beta2: F, epsilon: F, weight_decay: F, ) -> Self

Creates a new AdamW optimizer with the given hyperparameters

Arguments

• learning_rate - The learning rate for parameter updates
• beta1 - Exponential decay rate for the first moment estimates (default: 0.9)
• beta2 - Exponential decay rate for the second moment estimates (default: 0.999)
• epsilon - Small constant for numerical stability (default: 1e-8)
• weight_decay - Weight decay factor (default: 0.01)

Examples found in repository

examples/advanced_optimizers_example.rs (line 124)

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("Advanced Optimizers Example");

    // Initialize random number generator
    let mut rng = SmallRng::seed_from_u64(42);

    // Create a synthetic binary classification dataset
    let num_samples = 1000;
    let num_features = 20;
    let num_classes = 2;

    println!(
        "Generating synthetic dataset with {} samples, {} features...",
        num_samples, num_features
    );

    // Generate random input features
    let mut x_data = Array2::<f32>::zeros((num_samples, num_features));
    for i in 0..num_samples {
        for j in 0..num_features {
            x_data[[i, j]] = rng.random_range(-1.0..1.0);
        }
    }

    // Create true weights and bias for data generation
    let mut true_weights = Array2::<f32>::zeros((num_features, 1));
    for i in 0..num_features {
        true_weights[[i, 0]] = rng.random_range(-1.0..1.0);
    }
    let true_bias = rng.random_range(-1.0..1.0);

    // Generate binary labels (0 or 1) based on linear model with logistic function
    let mut y_data = Array2::<f32>::zeros((num_samples, num_classes));
    for i in 0..num_samples {
        let mut logit = true_bias;
        for j in 0..num_features {
            logit += x_data[[i, j]] * true_weights[[j, 0]];
        }

        // Apply sigmoid to get probability
        let prob = 1.0 / (1.0 + (-logit).exp());

        // Convert to one-hot encoding
        if prob > 0.5 {
            y_data[[i, 1]] = 1.0; // Class 1
        } else {
            y_data[[i, 0]] = 1.0; // Class 0
        }
    }

    // Split into train and test sets (80% train, 20% test)
    let train_size = (num_samples as f32 * 0.8) as usize;
    let test_size = num_samples - train_size;

    let x_train = x_data.slice(ndarray::s![0..train_size, ..]).to_owned();
    let y_train = y_data.slice(ndarray::s![0..train_size, ..]).to_owned();
    let x_test = x_data.slice(ndarray::s![train_size.., ..]).to_owned();
    let y_test = y_data.slice(ndarray::s![train_size.., ..]).to_owned();

    println!("Training set: {} samples", train_size);
    println!("Test set: {} samples", test_size);

    // Create a simple neural network model
    let hidden_size = 64;
    let dropout_rate = 0.2;
    let seed_rng = SmallRng::seed_from_u64(42);

    // Shared function to create identical model architectures for fair comparison
    let create_model = || -> Result<Sequential<f32>, Box<dyn std::error::Error>> {
        let mut model = Sequential::new();

        // Input to hidden layer
        let dense1 = Dense::new(
            num_features,
            hidden_size,
            Some("relu"),
            &mut seed_rng.clone(),
        )?;
        model.add_layer(dense1);

        // Dropout for regularization
        let dropout = Dropout::new(dropout_rate, &mut seed_rng.clone())?;
        model.add_layer(dropout);

        // Hidden to output layer
        let dense2 = Dense::new(
            hidden_size,
            num_classes,
            Some("softmax"),
            &mut seed_rng.clone(),
        )?;
        model.add_layer(dense2);

        Ok(model)
    };

    // Create models for each optimizer
    let mut sgd_model = create_model()?;
    let mut adam_model = create_model()?;
    let mut adamw_model = create_model()?;
    let mut radam_model = create_model()?;
    let mut rmsprop_model = create_model()?;

    // Create the loss function
    let loss_fn = CrossEntropyLoss::new(1e-10);

    // Create optimizers
    let learning_rate = 0.001;
    let batch_size = 32;
    let epochs = 20;

    let mut sgd_optimizer = SGD::new_with_config(learning_rate, 0.9, 0.0);
    let mut adam_optimizer = Adam::new(learning_rate, 0.9, 0.999, 1e-8);
    let mut adamw_optimizer = AdamW::new(learning_rate, 0.9, 0.999, 1e-8, 0.01);
    let mut radam_optimizer = RAdam::new(learning_rate, 0.9, 0.999, 1e-8, 0.0);
    let mut rmsprop_optimizer = RMSprop::new_with_config(learning_rate, 0.9, 1e-8, 0.0);

    // Helper function to compute accuracy
    let compute_accuracy = |model: &Sequential<f32>, x: &Array2<f32>, y: &Array2<f32>| -> f32 {
        let predictions = model.forward(&x.clone().into_dyn()).unwrap();
        let mut correct = 0;

        for i in 0..x.shape()[0] {
            let mut max_idx = 0;
            let mut max_val = predictions[[i, 0]];

            for j in 1..num_classes {
                if predictions[[i, j]] > max_val {
                    max_val = predictions[[i, j]];
                    max_idx = j;
                }
            }

            let true_idx =
                y[[i, 0]] < y[[i, 1]] as usize as u8 as i8 as usize as isize as i32 as f32;
            if max_idx as i32 == true_idx as i32 {
                correct += 1;
            }
        }

        correct as f32 / x.shape()[0] as f32
    };

    // Helper function to train model
    let mut train_model =
        |model: &mut Sequential<f32>, optimizer: &mut dyn Optimizer<f32>, name: &str| -> Vec<f32> {
            println!("\nTraining with {} optimizer...", name);
            let start_time = Instant::now();

            let mut train_losses = Vec::new();
            let num_batches = train_size.div_ceil(batch_size);

            for epoch in 0..epochs {
                let mut epoch_loss = 0.0;

                // Create a permutation for shuffling the data
                let mut indices: Vec<usize> = (0..train_size).collect();
                indices.shuffle(&mut rng);

                for batch_idx in 0..num_batches {
                    let start = batch_idx * batch_size;
                    let end = (start + batch_size).min(train_size);
                    let batch_indices = &indices[start..end];

                    // Create batch data
                    let mut x_batch = Array2::<f32>::zeros((batch_indices.len(), num_features));
                    let mut y_batch = Array2::<f32>::zeros((batch_indices.len(), num_classes));

                    for (i, &idx) in batch_indices.iter().enumerate() {
                        for j in 0..num_features {
                            x_batch[[i, j]] = x_train[[idx, j]];
                        }
                        for j in 0..num_classes {
                            y_batch[[i, j]] = y_train[[idx, j]];
                        }
                    }

                    // Convert to dynamic dimension arrays
                    let x_batch_dyn = x_batch.into_dyn();
                    let y_batch_dyn = y_batch.into_dyn();

                    // Perform a training step
                    let batch_loss = model
                        .train_batch(&x_batch_dyn, &y_batch_dyn, &loss_fn, optimizer)
                        .unwrap();
                    epoch_loss += batch_loss;
                }

                epoch_loss /= num_batches as f32;
                train_losses.push(epoch_loss);

                // Calculate and print metrics every few epochs
                if epoch % 5 == 0 || epoch == epochs - 1 {
                    let train_accuracy = compute_accuracy(model, &x_train, &y_train);
                    let test_accuracy = compute_accuracy(model, &x_test, &y_test);

                    println!(
                        "Epoch {}/{}: loss = {:.6}, train_acc = {:.2}%, test_acc = {:.2}%",
                        epoch + 1,
                        epochs,
                        epoch_loss,
                        train_accuracy * 100.0,
                        test_accuracy * 100.0
                    );
                }
            }

            let elapsed = start_time.elapsed();
            println!("{} training completed in {:.2?}", name, elapsed);

            // Final evaluation
            let train_accuracy = compute_accuracy(model, &x_train, &y_train);
            let test_accuracy = compute_accuracy(model, &x_test, &y_test);
            println!("Final metrics for {}:", name);
            println!("  Train accuracy: {:.2}%", train_accuracy * 100.0);
            println!("  Test accuracy:  {:.2}%", test_accuracy * 100.0);

            train_losses
        };

    // Train models with different optimizers
    let sgd_losses = train_model(&mut sgd_model, &mut sgd_optimizer, "SGD");
    let adam_losses = train_model(&mut adam_model, &mut adam_optimizer, "Adam");
    let adamw_losses = train_model(&mut adamw_model, &mut adamw_optimizer, "AdamW");
    let radam_losses = train_model(&mut radam_model, &mut radam_optimizer, "RAdam");
    let rmsprop_losses = train_model(&mut rmsprop_model, &mut rmsprop_optimizer, "RMSprop");

    // Print comparison summary
    println!("\nOptimizer Comparison Summary:");
    println!("----------------------------");
    println!("Initial learning rate: {}", learning_rate);
    println!("Batch size: {}", batch_size);
    println!("Epochs: {}", epochs);
    println!();

    println!("Final Loss Values:");
    println!("  SGD:     {:.6}", sgd_losses.last().unwrap());
    println!("  Adam:    {:.6}", adam_losses.last().unwrap());
    println!("  AdamW:   {:.6}", adamw_losses.last().unwrap());
    println!("  RAdam:   {:.6}", radam_losses.last().unwrap());
    println!("  RMSprop: {:.6}", rmsprop_losses.last().unwrap());

    println!("\nLoss progression (first value, middle value, last value):");
    println!(
        "  SGD:     {:.6}, {:.6}, {:.6}",
        sgd_losses.first().unwrap(),
        sgd_losses[epochs / 2],
        sgd_losses.last().unwrap()
    );
    println!(
        "  Adam:    {:.6}, {:.6}, {:.6}",
        adam_losses.first().unwrap(),
        adam_losses[epochs / 2],
        adam_losses.last().unwrap()
    );
    println!(
        "  AdamW:   {:.6}, {:.6}, {:.6}",
        adamw_losses.first().unwrap(),
        adamw_losses[epochs / 2],
        adamw_losses.last().unwrap()
    );
    println!(
        "  RAdam:   {:.6}, {:.6}, {:.6}",
        radam_losses.first().unwrap(),
        radam_losses[epochs / 2],
        radam_losses.last().unwrap()
    );
    println!(
        "  RMSprop: {:.6}, {:.6}, {:.6}",
        rmsprop_losses.first().unwrap(),
        rmsprop_losses[epochs / 2],
        rmsprop_losses.last().unwrap()
    );

    println!("\nLoss improvement ratio (first loss / last loss):");
    println!(
        "  SGD:     {:.2}x",
        sgd_losses.first().unwrap() / sgd_losses.last().unwrap()
    );
    println!(
        "  Adam:    {:.2}x",
        adam_losses.first().unwrap() / adam_losses.last().unwrap()
    );
    println!(
        "  AdamW:   {:.2}x",
        adamw_losses.first().unwrap() / adamw_losses.last().unwrap()
    );
    println!(
        "  RAdam:   {:.2}x",
        radam_losses.first().unwrap() / radam_losses.last().unwrap()
    );
    println!(
        "  RMSprop: {:.2}x",
        rmsprop_losses.first().unwrap() / rmsprop_losses.last().unwrap()
    );

    println!("\nAdvanced optimizers demo completed successfully!");

    Ok(())
}

pub fn default_with_lr(learning_rate: F) -> Result<Self>

Creates a new AdamW optimizer with default hyperparameters

Arguments

• learning_rate - The learning rate for parameter updates

pub fn get_beta1(&self) -> F

Gets the beta1 parameter

pub fn set_beta1(&mut self, beta1: F) -> &mut Self

Sets the beta1 parameter

pub fn get_beta2(&self) -> F

Gets the beta2 parameter

pub fn set_beta2(&mut self, beta2: F) -> &mut Self

Sets the beta2 parameter

pub fn get_epsilon(&self) -> F

Gets the epsilon parameter

pub fn set_epsilon(&mut self, epsilon: F) -> &mut Self

Sets the epsilon parameter

pub fn get_weight_decay(&self) -> F

Gets the weight decay parameter

pub fn set_weight_decay(&mut self, weight_decay: F) -> &mut Self

Sets the weight decay parameter

pub fn reset(&mut self)

Resets the internal state of the optimizer

Trait Implementations

impl<F: Clone + Float + ScalarOperand + Debug> Clone for AdamW<F>

fn clone(&self) -> AdamW<F>

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl<F: Debug + Float + ScalarOperand + Debug> Debug for AdamW<F>

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

Formats the value using the given formatter.

impl<F: Float + ScalarOperand + Debug> Optimizer<F> for AdamW<F>

fn update( &mut self, params: &mut [Array<F, IxDyn>], grads: &[Array<F, IxDyn>], ) -> Result<()>

Update parameters based on gradients

fn get_learning_rate(&self) -> F

Get the current learning rate

fn set_learning_rate(&mut self, lr: F)

Set the learning rate

fn step_model(&mut self, model: &mut dyn ParamLayer<F>) -> Result<()>

Update model parameters using the optimizer (non-generic version for trait objects)

fn reset(&mut self)

Reset the optimizer’s internal state

fn name(&self) -> &'static str

Get the optimizer’s name