TinyTorch/modules/06_optimizers/optimizers_dev.ipynb

Module 06: Optimizers - Sophisticated Learning Algorithms

Welcome to Module 06! You'll build optimizers that enable neural networks to learn from gradients using sophisticated algorithms.

🔗 Prerequisites & Progress

You've Built: Tensor with gradients (Modules 01-05) You'll Build: SGD, Adam, and AdamW optimizers with sophisticated momentum and adaptive learning You'll Enable: Modern optimization algorithms that power state-of-the-art neural networks

Connection Map:

Gradients → Optimizers → Training
(Module 05)  (Module 06)  (Module 07)

Learning Objectives

By the end of this module, you will:

1. Implement SGD with momentum for stable gradient descent
2. Build Adam optimizer with adaptive learning rates
3. Create AdamW optimizer with decoupled weight decay
4. Understand memory and computational trade-offs in optimization algorithms

Let's get started!

📦 Where This Code Lives in the Final Package

Learning Side: You work in modules/06_optimizers/optimizers_dev.py
Building Side: Code exports to tinytorch.core.optimizers

# How to use this module:
from tinytorch.core.optimizers import SGD, Adam, AdamW

Why this matters:

• Learning: Complete optimization system for modern neural network training
• Production: Proper organization like PyTorch's torch.optim with all optimization algorithms together
• Consistency: All optimization logic and parameter updating in core.optimizers
• Integration: Works seamlessly with gradients from Module 05 for complete training capability

In [ ]:

#| default_exp core.optimizers
#| export

import numpy as np
from typing import List, Union, Optional, Dict, Any

# Import Tensor from Module 01 (now with gradient support from Module 05)
from tinytorch.core.tensor import Tensor

1. Introduction: What are Optimizers?

Optimizers are the engines that drive neural network learning. They take gradients computed from your loss function and use them to update model parameters toward better solutions. Think of optimization as navigating a complex landscape where you're trying to find the lowest valley (minimum loss).

The Optimization Challenge

Imagine you're hiking in dense fog, trying to reach the bottom of a valley. You can only feel the slope under your feet (the gradient), but you can't see where you're going. Different optimization strategies are like different hiking approaches:

Loss Landscape (2D visualization):
       🏔️
      /  \\
   🚶 /    \\
    /      \\
   /   🎯   \\  ← Global minimum (goal)
  /          \\
 🏔️          🏔️

Challenge: Navigate to 🎯 using only local slope information!

Our Optimizer Toolkit

SGD (Stochastic Gradient Descent)

• Strategy: Always step downhill
• Problem: Can get stuck oscillating in narrow valleys
• Solution: Add momentum to "coast" through oscillations

Adam (Adaptive Moment Estimation)

• Strategy: Adapt step size for each parameter individually
• Advantage: Different learning rates for different dimensions
• Key Insight: Some directions need big steps, others need small steps

AdamW (Adam with Weight Decay)

• Strategy: Adam + proper regularization
• Fix: Separates optimization from regularization
• Result: Better generalization and training stability

The Mathematics Behind Movement

At its core, optimization follows: θ_new = θ_old - α * direction

Where:

• θ = parameters (your position in the landscape)
• α = step size (learning rate)
• direction = where to step (gradient-based)

But sophisticated optimizers do much more than basic gradient descent!

2. Foundations: Mathematical Background

Understanding Momentum: The Physics of Optimization

Momentum in optimization works like momentum in physics. A ball rolling down a hill doesn't immediately change direction when it hits a small bump - it has momentum that carries it forward.

Without Momentum (SGD):           With Momentum:
     ↓                                ↘️
  ←  •  →  ← oscillation           →  •  → smooth path
     ↑                                ↙️

Narrow valley problem:            Momentum solution:
|\     /|                        |\     /|
| \ • / | ← ping-pong             | \ •→/ | ← smoother
|  \ /  |   motion                |  \ /  |   descent
|   ●   |                        |   ●   |

SGD with Momentum Formula:

velocity = β * previous_velocity + (1-β) * current_gradient
parameter = parameter - learning_rate * velocity

Where β ≈ 0.9 means "90% memory of previous direction"

Adam: Adaptive Learning for Each Parameter

Adam solves a key problem: different parameters need different learning rates. Imagine adjusting the focus and zoom on a camera - you need fine control for focus but coarse control for zoom.

Parameter Landscape (2 dimensions):

   param2
     ^
     |
   😞|    steep gradient
     |    (needs small steps)
     |
  ---+--●--→ param1
     |     \\
     |      \\ gentle gradient
     |       \\ (needs big steps)

Adam Solution: Automatic step size per parameter!

Adam's Two-Memory System:

1. First Moment (m): "Which direction am I usually going?"

   • m = β₁ * old_m + (1-β₁) * gradient
   • Like momentum, but for direction

2. Second Moment (v): "How big are my gradients usually?"

   • v = β₂ * old_v + (1-β₂) * gradient²
   • Tracks gradient magnitude

3. Adaptive Update:

   • step_size = m / √v
   • Big gradients → smaller steps
   • Small gradients → relatively bigger steps

AdamW: Fixing Weight Decay

Adam has a subtle bug in how it applies weight decay (regularization). AdamW fixes this:

Adam (incorrect):               AdamW (correct):
gradient += weight_decay * param    [compute gradient update]
update_param_with_gradient()        param -= learning_rate * gradient_update
                                   param *= (1 - weight_decay)  ← separate!

Why it matters:
- Adam: Weight decay affected by adaptive learning rates
- AdamW: Weight decay is consistent regardless of gradients

3. Implementation: Building Optimizers

Now we'll implement each optimizer step by step, following the pattern: understand the algorithm → implement it → test it immediately. Each optimizer builds on the foundation of the previous one.

Implementation Strategy

Optimizer Base Class
    ↓
SGD (foundation algorithm)
    ↓
SGD + Momentum (reduce oscillations)
    ↓
Adam (adaptive learning rates)
    ↓
AdamW (proper weight decay)

In [ ]:

#| export
class Optimizer:
    """
    Base class for all optimizers.

    This class defines the common interface that all optimizers must implement:
    - zero_grad(): Clear gradients from parameters
    - step(): Update parameters based on gradients
    """

    def __init__(self, params: List[Tensor]):
        """
        Initialize optimizer with parameters to optimize.

        TODO: Set up the parameter list for optimization

        APPROACH:
        1. Store parameters as a list for iteration
        2. Validate that all parameters require gradients
        3. Initialize step counter for algorithms that need it

        EXAMPLE:
        >>> linear = Linear(784, 128)
        >>> optimizer = SGD(linear.parameters(), lr=0.01)

        HINT: Check that each parameter has requires_grad=True
        """
        ### BEGIN SOLUTION
        # Validate and store parameters
        if not isinstance(params, list):
            params = list(params)

        # Check that parameters require gradients
        for i, param in enumerate(params):
            if not isinstance(param, Tensor):
                raise TypeError(f"Parameter {i} must be a Tensor, got {type(param)}")
            if not param.requires_grad:
                raise ValueError(f"Parameter {i} does not require gradients. Set requires_grad=True.")

        self.params = params
        self.step_count = 0  # For algorithms that need step counting
        ### END SOLUTION

    def zero_grad(self):
        """
        Clear gradients from all parameters.

        TODO: Reset all parameter gradients to None

        APPROACH:
        1. Iterate through all parameters
        2. Set each parameter's grad to None

        EXAMPLE:
        >>> optimizer.zero_grad()  # Clears all gradients
        >>> assert param.grad is None for param in optimizer.params

        WHY: Gradients accumulate by default, so we need to clear them between batches
        """
        ### BEGIN SOLUTION
        for param in self.params:
            param.grad = None
        ### END SOLUTION

    def step(self):
        """
        Update parameters based on gradients.

        This is abstract - each optimizer implements its own update rule.
        """
        raise NotImplementedError("Subclasses must implement step()")

🔬 Unit Test: Base Optimizer

This test validates our base Optimizer class works correctly. What we're testing: Parameter validation and zero_grad functionality Why it matters: Foundation for all specific optimizer implementations Expected: Proper parameter storage and gradient clearing

In [ ]:

def test_unit_optimizer_base():
    """🔬 Test base Optimizer functionality."""
    print("🔬 Unit Test: Base Optimizer...")

    # Create test parameters
    param1 = Tensor([1.0, 2.0], requires_grad=True)
    param2 = Tensor([[3.0, 4.0], [5.0, 6.0]], requires_grad=True)

    # Add some gradients
    param1.grad = Tensor([0.1, 0.2])
    param2.grad = Tensor([[0.3, 0.4], [0.5, 0.6]])

    # Create optimizer
    optimizer = Optimizer([param1, param2])

    # Test parameter storage
    assert len(optimizer.params) == 2
    assert optimizer.params[0] is param1
    assert optimizer.params[1] is param2
    assert optimizer.step_count == 0

    # Test zero_grad
    optimizer.zero_grad()
    assert param1.grad is None
    assert param2.grad is None

    # Test error handling
    try:
        bad_param = Tensor([1.0], requires_grad=False)
        Optimizer([bad_param])
        assert False, "Should have raised ValueError"
    except ValueError as e:
        assert "does not require gradients" in str(e)

    print("✅ Base Optimizer works correctly!")

if __name__ == "__main__":
    test_unit_optimizer_base()

SGD - Stochastic Gradient Descent

SGD is the foundation of neural network optimization. It implements the simple but powerful idea: "move in the direction opposite to the gradient."

Why SGD Works

Gradients point uphill (toward higher loss). To minimize loss, we go downhill:

Loss Surface (side view):

    Loss
     ^
     |
  📈 |     current position
     |    /
     |   • ← you are here
     |  / \
     | /   \ gradient points uphill
     |/     \
     ●-------\--→ parameters
      \        \
       \        ↘️ SGD steps downhill
        \        (opposite to gradient)
         \⭐ ← goal (minimum loss)

The Oscillation Problem

Pure SGD can get trapped oscillating in narrow valleys:

Narrow valley (top view):
  \     /
   \   /   ← steep sides
    \ /
  4← • →2  ← SGD bounces back and forth
    / \
   1   3   instead of going down the valley
  /     \
 ●       \
 goal     \

Momentum Solution

Momentum remembers the direction you were going and continues in that direction:

With momentum:
  \     /
   \   /
    \ /
     •  ← smooth path down the valley
    / ↓
   /   ↓
  ●    ↓  momentum carries us through oscillations
 goal

Implementation: SGD keeps a "velocity" buffer that accumulates momentum.

In [ ]:

#| export
class SGD(Optimizer):
    """
    Stochastic Gradient Descent with momentum.

    SGD is the foundational optimization algorithm that moves parameters
    in the direction opposite to gradients. With momentum, it remembers
    previous updates to reduce oscillations and accelerate convergence.
    """

    def __init__(self, params: List[Tensor], lr: float = 0.01, momentum: float = 0.0, weight_decay: float = 0.0):
        """
        Initialize SGD optimizer.

        TODO: Set up SGD with momentum and weight decay

        APPROACH:
        1. Call parent constructor to set up parameters
        2. Store learning rate, momentum, and weight decay
        3. Initialize momentum buffers for each parameter

        EXAMPLE:
        >>> optimizer = SGD(model.parameters(), lr=0.01, momentum=0.9)

        HINTS:
        - Momentum buffers should be initialized as None
        - They'll be created lazily on first step
        """
        ### BEGIN SOLUTION
        super().__init__(params)

        self.lr = lr
        self.momentum = momentum
        self.weight_decay = weight_decay

        # Initialize momentum buffers (created lazily)
        self.momentum_buffers = [None for _ in self.params]
        ### END SOLUTION

    def step(self):
        """
        Perform SGD update step with momentum.

        TODO: Implement SGD parameter update with momentum

        APPROACH:
        1. For each parameter with gradients:
           a. Apply weight decay if specified
           b. Update momentum buffer
           c. Update parameter using momentum

        FORMULA:
        - With weight decay: grad = grad + weight_decay * param
        - Momentum: v = momentum * v_prev + grad
        - Update: param = param - lr * v

        HINTS:
        - Skip parameters without gradients
        - Initialize momentum buffers on first use
        - Use in-place operations to save memory
        """
        ### BEGIN SOLUTION
        for i, param in enumerate(self.params):
            if param.grad is None:
                continue

            # Get gradient (param.grad is already a numpy array)
            grad = param.grad

            # Apply weight decay
            if self.weight_decay != 0:
                grad = grad + self.weight_decay * param.data

            # Update momentum buffer
            if self.momentum != 0:
                if self.momentum_buffers[i] is None:
                    # Initialize momentum buffer
                    self.momentum_buffers[i] = np.zeros_like(param.data)

                # Update momentum: v = momentum * v_prev + grad
                self.momentum_buffers[i] = self.momentum * self.momentum_buffers[i] + grad
                grad = self.momentum_buffers[i]

            # Update parameter: param = param - lr * grad
            param.data = param.data - self.lr * grad

        # Increment step counter
        self.step_count += 1
        ### END SOLUTION

🔬 Unit Test: SGD Optimizer

This test validates our SGD implementation works correctly. What we're testing: SGD updates with and without momentum Why it matters: Core optimization algorithm used in neural network training Expected: Correct parameter updates following SGD formulas

In [ ]:

def test_unit_sgd_optimizer():
    """🔬 Test SGD optimizer implementation."""
    print("🔬 Unit Test: SGD Optimizer...")

    # Test basic SGD without momentum
    param = Tensor([1.0, 2.0], requires_grad=True)
    param.grad = Tensor([0.1, 0.2])

    optimizer = SGD([param], lr=0.1)
    original_data = param.data.copy()

    optimizer.step()

    # Expected: param = param - lr * grad = [1.0, 2.0] - 0.1 * [0.1, 0.2] = [0.99, 1.98]
    expected = original_data - 0.1 * param.grad.data
    assert np.allclose(param.data, expected)
    assert optimizer.step_count == 1

    # Test SGD with momentum
    param2 = Tensor([1.0, 2.0], requires_grad=True)
    param2.grad = Tensor([0.1, 0.2])

    optimizer_momentum = SGD([param2], lr=0.1, momentum=0.9)

    # First step: v = 0.9 * 0 + [0.1, 0.2] = [0.1, 0.2]
    optimizer_momentum.step()
    expected_first = np.array([1.0, 2.0]) - 0.1 * np.array([0.1, 0.2])
    assert np.allclose(param2.data, expected_first)

    # Second step with same gradient
    param2.grad = Tensor([0.1, 0.2])
    optimizer_momentum.step()
    # v = 0.9 * [0.1, 0.2] + [0.1, 0.2] = [0.19, 0.38]
    expected_momentum = np.array([0.19, 0.38])
    expected_second = expected_first - 0.1 * expected_momentum
    assert np.allclose(param2.data, expected_second, rtol=1e-5)

    # Test weight decay
    param3 = Tensor([1.0, 2.0], requires_grad=True)
    param3.grad = Tensor([0.1, 0.2])

    optimizer_wd = SGD([param3], lr=0.1, weight_decay=0.01)
    optimizer_wd.step()

    # grad_with_decay = [0.1, 0.2] + 0.01 * [1.0, 2.0] = [0.11, 0.22]
    expected_wd = np.array([1.0, 2.0]) - 0.1 * np.array([0.11, 0.22])
    assert np.allclose(param3.data, expected_wd)

    print("✅ SGD optimizer works correctly!")

if __name__ == "__main__":
    test_unit_sgd_optimizer()

Adam - Adaptive Moment Estimation

Adam solves a fundamental problem with SGD: different parameters often need different learning rates. Think of tuning a complex system where some knobs need gentle adjustments and others need bold changes.

The Parameter Scaling Problem

Consider a neural network with both embedding weights and output weights:

Parameter Sensitivity Landscape:

  output_weight                 embedding_weight
       ↑                              ↑
       |                              |
    😱 |  steep cliff                  |  🐌 gentle slope
       |  (needs tiny steps)          |  (needs big steps)
       |                              |
    ━━━●━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━●━━━→

Same learning rate = disaster!
• Small LR: output weights learn fast, embeddings crawl
• Large LR: embeddings learn well, output weights explode

Adam's Adaptive Solution

Adam automatically adjusts learning rates by tracking two statistics:

1. MOMENTUM (first moment): "Which way am I usually going?"
   m = 0.9 * old_direction + 0.1 * current_gradient

   Visualization:
   old: →→→→
   new:     ↗️
   m:   →→→↗️  (weighted average)

2. SCALE (second moment): "How big are my steps usually?"
   v = 0.999 * old_scale + 0.001 * (current_gradient)²

   Big gradients → bigger v → smaller effective steps
   Small gradients → smaller v → bigger effective steps

3. ADAPTIVE UPDATE:
   step = momentum / √scale
   param = param - learning_rate * step

Bias Correction: The Cold Start Problem

Adam starts with m=0 and v=0, which creates a bias toward zero initially:

Without bias correction:    With bias correction:

Step 1: m = 0.9*0 + 0.1*g    Step 1: m̂ = m / (1-0.9¹) = m / 0.1
       = 0.1*g (too small!)           = g (correct!)

Step 2: m = 0.9*0.1*g + 0.1*g Step 2: m̂ = m / (1-0.9²) = m / 0.19
       = 0.19*g (still small)         ≈ g (better!)

Key Insight: Adam is like having an automatic transmission that adjusts gear ratios for each parameter individually.

In [ ]:

#| export
class Adam(Optimizer):
    """
    Adam optimizer with adaptive learning rates.

    Adam computes individual adaptive learning rates for different parameters
    from estimates of first and second moments of the gradients.
    This makes it effective for problems with sparse gradients or noisy data.
    """

    def __init__(self, params: List[Tensor], lr: float = 0.001, betas: tuple = (0.9, 0.999), eps: float = 1e-8, weight_decay: float = 0.0):
        """
        Initialize Adam optimizer.

        TODO: Set up Adam with adaptive learning rates

        APPROACH:
        1. Call parent constructor
        2. Store hyperparameters (lr, betas, eps, weight_decay)
        3. Initialize first and second moment buffers

        PARAMETERS:
        - lr: Learning rate (default: 0.001)
        - betas: Coefficients for computing running averages (default: (0.9, 0.999))
        - eps: Small constant for numerical stability (default: 1e-8)
        - weight_decay: L2 penalty coefficient (default: 0.0)

        EXAMPLE:
        >>> optimizer = Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999))
        """
        ### BEGIN SOLUTION
        super().__init__(params)

        self.lr = lr
        self.beta1, self.beta2 = betas
        self.eps = eps
        self.weight_decay = weight_decay

        # Initialize moment buffers (created lazily)
        self.m_buffers = [None for _ in self.params]  # First moment (mean)
        self.v_buffers = [None for _ in self.params]  # Second moment (variance)
        ### END SOLUTION

    def step(self):
        """
        Perform Adam update step.

        TODO: Implement Adam parameter update with adaptive learning rates

        APPROACH:
        1. For each parameter with gradients:
           a. Apply weight decay if specified
           b. Update first moment estimate (momentum of gradient)
           c. Update second moment estimate (momentum of squared gradient)
           d. Compute bias-corrected moments
           e. Update parameter using adaptive learning rate

        FORMULAS:
        - m_t = β₁ * m_{t-1} + (1-β₁) * g_t
        - v_t = β₂ * v_{t-1} + (1-β₂) * g_t²
        - m̂_t = m_t / (1-β₁^t)
        - v̂_t = v_t / (1-β₂^t)
        - θ_t = θ_{t-1} - lr * m̂_t / (√v̂_t + ε)

        HINTS:
        - Initialize buffers as zeros on first use
        - Use step_count for bias correction
        - Square gradients element-wise for second moment
        """
        ### BEGIN SOLUTION
        # Increment step counter first (needed for bias correction)
        self.step_count += 1

        for i, param in enumerate(self.params):
            if param.grad is None:
                continue

            # Get gradient (param.grad is already a numpy array)
            grad = param.grad

            # Apply weight decay
            if self.weight_decay != 0:
                grad = grad + self.weight_decay * param.data

            # Initialize buffers if needed
            if self.m_buffers[i] is None:
                self.m_buffers[i] = np.zeros_like(param.data)
                self.v_buffers[i] = np.zeros_like(param.data)

            # Update biased first moment estimate
            self.m_buffers[i] = self.beta1 * self.m_buffers[i] + (1 - self.beta1) * grad

            # Update biased second moment estimate
            self.v_buffers[i] = self.beta2 * self.v_buffers[i] + (1 - self.beta2) * (grad ** 2)

            # Compute bias correction
            bias_correction1 = 1 - self.beta1 ** self.step_count
            bias_correction2 = 1 - self.beta2 ** self.step_count

            # Compute bias-corrected moments
            m_hat = self.m_buffers[i] / bias_correction1
            v_hat = self.v_buffers[i] / bias_correction2

            # Update parameter
            param.data = param.data - self.lr * m_hat / (np.sqrt(v_hat) + self.eps)
        ### END SOLUTION

🔬 Unit Test: Adam Optimizer

This test validates our Adam implementation works correctly. What we're testing: Adam updates with adaptive learning rates and bias correction Why it matters: Most popular optimizer for modern neural networks Expected: Correct parameter updates following Adam formulas

In [ ]:

def test_unit_adam_optimizer():
    """🔬 Test Adam optimizer implementation."""
    print("🔬 Unit Test: Adam Optimizer...")

    # Test basic Adam functionality
    param = Tensor([1.0, 2.0], requires_grad=True)
    param.grad = Tensor([0.1, 0.2])

    optimizer = Adam([param], lr=0.01, betas=(0.9, 0.999), eps=1e-8)
    original_data = param.data.copy()

    # First step
    optimizer.step()

    # Manually compute expected values
    grad = np.array([0.1, 0.2])

    # First moment: m = 0.9 * 0 + 0.1 * grad = 0.1 * grad
    m = 0.1 * grad

    # Second moment: v = 0.999 * 0 + 0.001 * grad^2 = 0.001 * grad^2
    v = 0.001 * (grad ** 2)

    # Bias correction
    bias_correction1 = 1 - 0.9 ** 1  # = 0.1
    bias_correction2 = 1 - 0.999 ** 1  # = 0.001

    m_hat = m / bias_correction1  # = grad
    v_hat = v / bias_correction2  # = grad^2

    # Update
    expected = original_data - 0.01 * m_hat / (np.sqrt(v_hat) + 1e-8)

    assert np.allclose(param.data, expected, rtol=1e-6)
    assert optimizer.step_count == 1

    # Test second step to verify moment accumulation
    param.grad = Tensor([0.1, 0.2])
    optimizer.step()

    # Should have updated moments
    assert optimizer.m_buffers[0] is not None
    assert optimizer.v_buffers[0] is not None
    assert optimizer.step_count == 2

    # Test with weight decay
    param2 = Tensor([1.0, 2.0], requires_grad=True)
    param2.grad = Tensor([0.1, 0.2])

    optimizer_wd = Adam([param2], lr=0.01, weight_decay=0.01)
    optimizer_wd.step()

    # Weight decay should modify the effective gradient
    # grad_with_decay = [0.1, 0.2] + 0.01 * [1.0, 2.0] = [0.11, 0.22]
    # The exact computation is complex, but we can verify parameter changed
    assert not np.array_equal(param2.data, np.array([1.0, 2.0]))

    print("✅ Adam optimizer works correctly!")

if __name__ == "__main__":
    test_unit_adam_optimizer()

AdamW - Adam with Decoupled Weight Decay

AdamW fixes a subtle but important bug in Adam's weight decay implementation. The bug affects how regularization interacts with adaptive learning rates.

The Adam Weight Decay Bug

In standard Adam, weight decay is added to gradients before the adaptive scaling:

Adam's approach (problematic):
1. gradient = computed_gradient + weight_decay * parameter
2. m = β₁ * m + (1-β₁) * gradient
3. v = β₂ * v + (1-β₂) * gradient²
4. step = m / √v
5. parameter = parameter - learning_rate * step

Problem: Weight decay gets "adapted" by the learning rate scaling!

Why This Matters

Weight decay should be a consistent regularization force, but Adam makes it inconsistent:

Parameter Update Comparison:

Large gradients → small adaptive LR → weak weight decay effect
Small gradients → large adaptive LR → strong weight decay effect

This is backwards! We want consistent regularization.

AdamW's Fix: Decoupled Weight Decay

AdamW separates gradient-based updates from weight decay:

AdamW's approach (correct):
1. m = β₁ * m + (1-β₁) * pure_gradient  ← NO weight decay here
2. v = β₂ * v + (1-β₂) * pure_gradient²
3. step = m / √v
4. parameter = parameter - learning_rate * step        ← gradient update
5. parameter = parameter * (1 - weight_decay_rate)    ← separate decay

Result: Consistent regularization independent of gradient magnitudes!

Visual Comparison

Adam weight decay:               AdamW weight decay:

gradient ──┐                    gradient ──→ adaptive ──→ param
           ├─→ adaptive ──→ param                  update
weight ────┘   scaling
decay
                                weight ─────────→ param
                                decay           shrinkage

Coupled (inconsistent)          Decoupled (consistent)

Key Insight: AdamW treats optimization and regularization as separate, independent processes, leading to better training dynamics and generalization.

In [ ]:

#| export
class AdamW(Optimizer):
    """
    AdamW optimizer with decoupled weight decay.

    AdamW fixes a bug in Adam's weight decay implementation by decoupling
    weight decay from the gradient-based update. This leads to better
    regularization and is the preferred version for most applications.
    """

    def __init__(self, params: List[Tensor], lr: float = 0.001, betas: tuple = (0.9, 0.999), eps: float = 1e-8, weight_decay: float = 0.01):
        """
        Initialize AdamW optimizer.

        TODO: Set up AdamW with decoupled weight decay

        APPROACH:
        1. Call parent constructor
        2. Store hyperparameters (note higher default weight_decay)
        3. Initialize moment buffers like Adam

        KEY DIFFERENCE from Adam:
        - Weight decay is applied directly to parameters, not added to gradients
        - This provides better regularization behavior

        EXAMPLE:
        >>> optimizer = AdamW(model.parameters(), lr=0.001, weight_decay=0.01)
        """
        ### BEGIN SOLUTION
        super().__init__(params)

        self.lr = lr
        self.beta1, self.beta2 = betas
        self.eps = eps
        self.weight_decay = weight_decay

        # Initialize moment buffers (same as Adam)
        self.m_buffers = [None for _ in self.params]
        self.v_buffers = [None for _ in self.params]
        ### END SOLUTION

    def step(self):
        """
        Perform AdamW update step with decoupled weight decay.

        TODO: Implement AdamW parameter update

        APPROACH:
        1. For each parameter with gradients:
           a. Update moments using gradients (NOT modified by weight decay)
           b. Compute bias-corrected moments
           c. Apply gradient-based update
           d. Apply weight decay directly to parameters

        KEY DIFFERENCE from Adam:
        - Weight decay: θ_t = θ_t - lr * weight_decay * θ_t (applied after gradient update)
        - NOT: grad = grad + weight_decay * param (Adam's incorrect approach)

        FORMULAS:
        - Same moment updates as Adam (using unmodified gradients)
        - Gradient update: θ_t = θ_{t-1} - lr * m̂_t / (√v̂_t + ε)
        - Weight decay: θ_t = θ_t * (1 - lr * weight_decay)

        HINT: Apply weight decay after gradient update for proper decoupling
        """
        ### BEGIN SOLUTION
        # Increment step counter first
        self.step_count += 1

        for i, param in enumerate(self.params):
            if param.grad is None:
                continue

            # Get gradient (NOT modified by weight decay) - param.grad is already a numpy array
            grad = param.grad

            # Initialize buffers if needed
            if self.m_buffers[i] is None:
                self.m_buffers[i] = np.zeros_like(param.data)
                self.v_buffers[i] = np.zeros_like(param.data)

            # Update moments using pure gradients
            self.m_buffers[i] = self.beta1 * self.m_buffers[i] + (1 - self.beta1) * grad
            self.v_buffers[i] = self.beta2 * self.v_buffers[i] + (1 - self.beta2) * (grad ** 2)

            # Compute bias correction
            bias_correction1 = 1 - self.beta1 ** self.step_count
            bias_correction2 = 1 - self.beta2 ** self.step_count

            # Compute bias-corrected moments
            m_hat = self.m_buffers[i] / bias_correction1
            v_hat = self.v_buffers[i] / bias_correction2

            # Apply gradient-based update
            param.data = param.data - self.lr * m_hat / (np.sqrt(v_hat) + self.eps)

            # Apply decoupled weight decay
            if self.weight_decay != 0:
                param.data = param.data * (1 - self.lr * self.weight_decay)
        ### END SOLUTION

🔬 Unit Test: AdamW Optimizer

This test validates our AdamW implementation with decoupled weight decay. What we're testing: AdamW updates with proper weight decay decoupling Why it matters: State-of-the-art optimizer for transformer models Expected: Correct separation of gradient updates and weight decay

In [ ]:

def test_unit_adamw_optimizer():
    """🔬 Test AdamW optimizer implementation."""
    print("🔬 Unit Test: AdamW Optimizer...")

    # Test AdamW vs Adam difference in weight decay
    # Create identical parameters for comparison
    param_adam = Tensor([1.0, 2.0], requires_grad=True)
    param_adamw = Tensor([1.0, 2.0], requires_grad=True)

    param_adam.grad = Tensor([0.1, 0.2])
    param_adamw.grad = Tensor([0.1, 0.2])

    # Create optimizers with same settings
    adam = Adam([param_adam], lr=0.01, weight_decay=0.01)
    adamw = AdamW([param_adamw], lr=0.01, weight_decay=0.01)

    # Take one step
    adam.step()
    adamw.step()

    # Results should be different due to weight decay implementation
    assert not np.allclose(param_adam.data, param_adamw.data, rtol=1e-6)

    # Test AdamW basic functionality
    param = Tensor([1.0, 2.0], requires_grad=True)
    param.grad = Tensor([0.1, 0.2])

    optimizer = AdamW([param], lr=0.01, weight_decay=0.01)
    original_data = param.data.copy()

    optimizer.step()

    # Parameter should have changed
    assert not np.array_equal(param.data, original_data)
    assert optimizer.step_count == 1

    # Test that moment buffers are created
    assert optimizer.m_buffers[0] is not None
    assert optimizer.v_buffers[0] is not None

    # Test zero weight decay behaves like Adam
    param1 = Tensor([1.0, 2.0], requires_grad=True)
    param2 = Tensor([1.0, 2.0], requires_grad=True)

    param1.grad = Tensor([0.1, 0.2])
    param2.grad = Tensor([0.1, 0.2])

    adam_no_wd = Adam([param1], lr=0.01, weight_decay=0.0)
    adamw_no_wd = AdamW([param2], lr=0.01, weight_decay=0.0)

    adam_no_wd.step()
    adamw_no_wd.step()

    # Should be very similar (within numerical precision)
    assert np.allclose(param1.data, param2.data, rtol=1e-10)

    print("✅ AdamW optimizer works correctly!")

if __name__ == "__main__":
    test_unit_adamw_optimizer()

4. Integration: Bringing It Together

Now let's see how our optimizers perform in realistic scenarios. We'll compare their behavior on the same optimization problem to understand their different characteristics.

Optimizer Behavior Comparison

Each optimizer takes a different approach to the same problem:

Optimization Problem: Find minimum of f(x) = x²

SGD approach:        Adam approach:        AdamW approach:
  ↓                    ↓                     ↓
 x ──→ minimize       x ──→ minimize       x ──→ minimize
  ↑                    ↑                     ↑
fixed LR           adaptive LR          adaptive LR + decay

5. Systems Analysis: Optimizer Performance and Memory

Different optimizers have very different resource requirements. Understanding these trade-offs is crucial for production ML systems.

Memory Usage Patterns

Optimizer Memory Requirements (per parameter):

SGD:           Adam/AdamW:
┌────────┐     ┌────────┐
│ param  │     │ param  │
├────────┤     ├────────┤
│momentum│     │   m    │ ← first moment
└────────┘     ├────────┤
               │   v    │ ← second moment
               └────────┘

2× memory       3× memory

Computational Complexity

Per-step Operations:

SGD:                     Adam:
• 1 multiplication       • 3 multiplications
• 1 addition            • 4 additions
• 1 subtraction         • 1 subtraction
                        • 1 square root
                        • 1 division

O(n) simple ops         O(n) complex ops

In [ ]:

def analyze_optimizer_memory_usage():
    """📊 Analyze memory usage of different optimizers."""
    print("📊 Analyzing Optimizer Memory Usage...")

    # Create test parameters of different sizes
    param_sizes = [1000, 10000, 100000]  # 1K, 10K, 100K parameters

    print("Optimizer Memory Analysis (per parameter tensor):")
    print("=" * 60)
    print(f"{'Size':<10} {'SGD':<10} {'Adam':<10} {'AdamW':<10} {'Ratio':<10}")
    print("-" * 60)

    for size in param_sizes:
        # Create parameter
        param = Tensor(np.random.randn(size), requires_grad=True)
        param.grad = Tensor(np.random.randn(size))

        # SGD memory (parameter + momentum buffer)
        sgd = SGD([param], momentum=0.9)
        sgd.step()  # Initialize buffers
        sgd_memory = size * 2  # param + momentum buffer

        # Adam memory (parameter + 2 moment buffers)
        param_adam = Tensor(np.random.randn(size), requires_grad=True)
        param_adam.grad = Tensor(np.random.randn(size))
        adam = Adam([param_adam])
        adam.step()  # Initialize buffers
        adam_memory = size * 3  # param + m_buffer + v_buffer

        # AdamW memory (same as Adam)
        adamw_memory = adam_memory

        # Memory ratio (Adam/SGD)
        ratio = adam_memory / sgd_memory

        print(f"{size:<10} {sgd_memory:<10} {adam_memory:<10} {adamw_memory:<10} {ratio:.1f}x")

    print("\n💡 Key Insights:")
    print("- SGD: 2× parameter memory (momentum buffer)")
    print("- Adam/AdamW: 3× parameter memory (two moment buffers)")
    print("- Memory scales linearly with model size")
    print("- Trade-off: More memory for better convergence")

In [ ]:

def analyze_optimizer_convergence_behavior():
    """📊 Analyze convergence behavior of different optimizers."""
    print("📊 Analyzing Optimizer Convergence Behavior...")

    # Simulate optimization of a quadratic function: f(x) = 0.5 * x^2
    # Optimal solution: x* = 0, gradient = x

    def quadratic_loss(x):
        """Simple quadratic function for optimization testing."""
        return 0.5 * (x ** 2).sum()

    def compute_gradient(x):
        """Gradient of quadratic function: df/dx = x."""
        return x.copy()

    # Starting point
    x_start = np.array([5.0, -3.0, 2.0])  # Far from optimum [0, 0, 0]

    # Test different optimizers
    optimizers_to_test = [
        ("SGD", SGD, {"lr": 0.1}),
        ("SGD+Momentum", SGD, {"lr": 0.1, "momentum": 0.9}),
        ("Adam", Adam, {"lr": 0.1}),
        ("AdamW", AdamW, {"lr": 0.1, "weight_decay": 0.01})
    ]

    print("Convergence Analysis (quadratic function f(x) = 0.5 * x²):")
    print("=" * 70)
    print(f"{'Optimizer':<15} {'Step 0':<12} {'Step 5':<12} {'Step 10':<12} {'Final Loss':<12}")
    print("-" * 70)

    for name, optimizer_class, kwargs in optimizers_to_test:
        # Reset parameter
        param = Tensor(x_start.copy(), requires_grad=True)
        optimizer = optimizer_class([param], **kwargs)

        losses = []

        # Run optimization for 10 steps
        for step in range(11):
            # Compute loss and gradient
            loss = quadratic_loss(param.data)
            param.grad = Tensor(compute_gradient(param.data))

            losses.append(loss)

            # Update parameters
            if step < 10:  # Don't update after last evaluation
                optimizer.step()
                optimizer.zero_grad()

        # Format results
        step0 = f"{losses[0]:.6f}"
        step5 = f"{losses[5]:.6f}"
        step10 = f"{losses[10]:.6f}"
        final = f"{losses[10]:.6f}"

        print(f"{name:<15} {step0:<12} {step5:<12} {step10:<12} {final:<12}")

    print("\n💡 Key Insights:")
    print("- SGD: Steady progress but can be slow")
    print("- SGD+Momentum: Faster convergence, less oscillation")
    print("- Adam: Adaptive rates help with different parameter scales")
    print("- AdamW: Similar to Adam with regularization effects")

"""

🧪 Module Integration Test

Final validation that everything works together correctly. """

def import_previous_module(module_name: str, component_name: str): import sys import os sys.path.append(os.path.join(os.path.dirname(file), '..', module_name)) module = import(f"{module_name.split('_')[1]}_dev") return getattr(module, component_name)

In [ ]:

def test_module():
    """
    Comprehensive test of entire module functionality.

    This final test runs before module summary to ensure:
    - All unit tests pass
    - Functions work together correctly
    - Module is ready for integration with TinyTorch
    """
    print("🧪 RUNNING MODULE INTEGRATION TEST")
    print("=" * 50)

    # Run all unit tests
    print("Running unit tests...")
    test_unit_optimizer_base()
    test_unit_sgd_optimizer()
    test_unit_adam_optimizer()
    test_unit_adamw_optimizer()

    print("\nRunning integration scenarios...")

    # Test realistic neural network optimization scenario
    print("🔬 Integration Test: Multi-layer Network Optimization...")

    # Import components from previous modules using standardized helper
    Tensor = import_previous_module('01_tensor', 'Tensor')
    Linear = import_previous_module('03_layers', 'Linear')
    ReLU = import_previous_module('02_activations', 'ReLU')
    MSELoss = import_previous_module('04_losses', 'MSELoss')

    # Create parameters for a 2-layer network
    # Layer 1: 3 inputs -> 4 hidden
    W1 = Tensor(np.random.randn(3, 4) * 0.1, requires_grad=True)
    b1 = Tensor(np.zeros(4), requires_grad=True)

    # Layer 2: 4 hidden -> 2 outputs
    W2 = Tensor(np.random.randn(4, 2) * 0.1, requires_grad=True)
    b2 = Tensor(np.zeros(2), requires_grad=True)

    params = [W1, b1, W2, b2]

    # Add realistic gradients
    W1.grad = Tensor(np.random.randn(3, 4) * 0.01)
    b1.grad = Tensor(np.random.randn(4) * 0.01)
    W2.grad = Tensor(np.random.randn(4, 2) * 0.01)
    b2.grad = Tensor(np.random.randn(2) * 0.01)

    # Test all optimizers on same network
    optimizers = [
        SGD(params, lr=0.01, momentum=0.9),
        Adam([p for p in params], lr=0.001),  # Fresh param list for Adam
        AdamW([p for p in params], lr=0.001, weight_decay=0.01)  # Fresh param list for AdamW
    ]

    # Save original parameter values
    original_params = [p.data.copy() for p in params]

    # Test SGD
    optimizers[0].step()
    sgd_params = [p.data.copy() for p in params]

    # Restore parameters and test Adam
    for i, p in enumerate(params):
        p.data = original_params[i].copy()
        # Re-add gradients since they may have been modified
        if i == 0:
            p.grad = Tensor(np.random.randn(3, 4) * 0.01)
        elif i == 1:
            p.grad = Tensor(np.random.randn(4) * 0.01)
        elif i == 2:
            p.grad = Tensor(np.random.randn(4, 2) * 0.01)
        else:
            p.grad = Tensor(np.random.randn(2) * 0.01)

    # Update parameter references for Adam
    optimizers[1].params = params
    optimizers[1].step()
    adam_params = [p.data.copy() for p in params]

    # Restore parameters and test AdamW
    for i, p in enumerate(params):
        p.data = original_params[i].copy()
        # Re-add gradients
        if i == 0:
            p.grad = Tensor(np.random.randn(3, 4) * 0.01)
        elif i == 1:
            p.grad = Tensor(np.random.randn(4) * 0.01)
        elif i == 2:
            p.grad = Tensor(np.random.randn(4, 2) * 0.01)
        else:
            p.grad = Tensor(np.random.randn(2) * 0.01)

    # Update parameter references for AdamW
    optimizers[2].params = params
    optimizers[2].step()
    adamw_params = [p.data.copy() for p in params]

    # Verify parameters changed differently for each optimizer
    for i in range(len(params)):
        # Parameters should be different from original
        assert not np.array_equal(sgd_params[i], original_params[i])
        assert not np.array_equal(adam_params[i], original_params[i])
        assert not np.array_equal(adamw_params[i], original_params[i])

        # Different optimizers should produce different results
        assert not np.allclose(sgd_params[i], adam_params[i], rtol=1e-6)

    print("✅ Multi-layer network optimization works!")

    # Test optimizer state management
    print("🔬 Integration Test: Optimizer State Management...")

    param = Tensor([1.0, 2.0], requires_grad=True)
    param.grad = Tensor([0.1, 0.2])

    optimizer = Adam([param], lr=0.001)

    # First step should initialize buffers
    optimizer.step()
    assert optimizer.m_buffers[0] is not None
    assert optimizer.v_buffers[0] is not None
    assert optimizer.step_count == 1

    # Zero grad should clear gradients but preserve optimizer state
    optimizer.zero_grad()
    assert param.grad is None
    assert optimizer.m_buffers[0] is not None  # State preserved
    assert optimizer.step_count == 1  # Step count preserved

    print("✅ Optimizer state management works!")

    print("\n" + "=" * 50)
    print("🎉 ALL TESTS PASSED! Module ready for export.")
    print("Run: tito module complete 06_optimizers")

In [ ]:

# Run comprehensive module test
if __name__ == "__main__":
    test_module()

🎯 MODULE SUMMARY: Optimizers

Congratulations! You've built sophisticated optimization algorithms that power modern neural network training!

Key Accomplishments

• Built SGD optimizer with momentum for stable gradient descent and oscillation reduction
• Implemented Adam optimizer with adaptive learning rates and bias correction for different parameter scales
• Created AdamW optimizer with decoupled weight decay for proper regularization
• Analyzed memory trade-offs: SGD (2×), Adam/AdamW (3× parameter memory)
• All tests pass ✅ (validated by test_module())

Ready for Next Steps

Your optimizer implementations enable sophisticated neural network training! With gradients from Module 05 and optimizers from Module 06, you're ready to build complete training loops.

Export with: tito module complete 06_optimizers

Next: Module 07 will add training loops, learning rate scheduling, and checkpointing for complete end-to-end neural network training!