Deep Learning Optimizer Implementations: SGD, Adam, and AdamW Explained

The harvard-edge/cs249r_book repository provides lightweight, pedagogical implementations of SGD, Adam, and AdamW optimizers that demonstrate the core algorithms driving modern neural network training.

Deep learning optimizer implementations form the backbone of neural network training, determining how model parameters update during backpropagation. The cs249r_book repository from Harvard Edge contains a minimal framework called tinytorch that implements three fundamental optimizers in pure Python. These implementations prioritize clarity over performance, making them ideal for understanding the mathematical mechanics behind stochastic gradient descent and adaptive moment estimation.

The Optimizer Base Class

All deep learning optimizer implementations in the repository inherit from a common abstract base class that standardizes the training interface.

Located in tinytorch/src/07_optimizers/07_optimizers.py (lines 238-300), the Optimizer class defines the shared API used by every concrete implementation:

• zero_grad(): Clears accumulated gradients from all parameters before the next backward pass
• step(): Abstract method that each subclass implements to perform the parameter update
• _extract_gradient(): Helper that normalizes gradient extraction from either Tensor objects or raw NumPy arrays

This abstraction allows the training loop to remain agnostic to the specific optimization algorithm while ensuring consistent gradient handling across different parameter types.

Stochastic Gradient Descent (SGD) Implementation

SGD represents the foundational deep learning optimizer implementation, extending basic gradient descent with optional momentum and weight decay regularization.

Core Algorithm

The implementation resides at lines 537-603 in tinytorch/src/07_optimizers/07_optimizers.py. The SGD class maintains three key hyperparameters:

• lr: The learning rate controlling step size
• momentum: Coefficient for exponential moving average of gradients (typically 0.9)
• weight_decay: L2 regularization coefficient applied to parameters

Velocity Buffers and Update Logic

SGD stores per-parameter velocity buffers (self.momentum_buffers) that accumulate past gradient directions. The update procedure follows this exact sequence:

1. Extract the current gradient using _extract_gradient
2. Apply weight decay: grad = grad + weight_decay * param
3. Update velocity: velocity = momentum * velocity + grad
4. Update parameters: param = param - learning_rate * velocity

This momentum term smooths optimization trajectories, helping escape shallow local minima and accelerating convergence in consistent gradient directions.

Adam Optimizer Implementation

Adam (Adaptive Moment Estimation) represents a more sophisticated deep learning optimizer implementation that maintains per-parameter adaptive learning rates based on gradient history.

Adaptive Moment Estimation Mechanics

Located at lines 834-904 in tinytorch/src/07_optimizers/07_optimizers.py, the Adam class implements adaptive optimization through:

• beta1: Exponential decay rate for first moment estimates (mean of gradients, default 0.9)
• beta2: Exponential decay rate for second moment estimates (uncentered variance, default 0.999)
• eps: Small constant for numerical stability (typically 1e-8)
• weight_decay: Optional L2 regularization

Bias Correction and Update Steps

Adam maintains two buffer dictionaries:

• m_buffers: First moment estimates (mean gradients)
• v_buffers: Second moment estimates (mean squared gradients)

The step() method implements the full Adam algorithm:

1. Increment global step_count for bias correction
2. For each parameter:
   • Extract gradient
   • Apply weight decay to gradient (if enabled)
   • Update first moment: m = beta1 * m + (1 - beta1) * grad
   • Update second moment: v = beta2 * v + (1 - beta2) * grad * grad
   • Compute bias-corrected estimates: m_hat = m / (1 - beta1^t), v_hat = v / (1 - beta2^t)
   • Update parameter: param -= lr * m_hat / (sqrt(v_hat) + eps)

This adaptive approach automatically adjusts learning rates per-parameter, making Adam particularly effective for sparse gradients and non-stationary objectives.

AdamW Optimizer Implementation

AdamW modifies the standard Adam deep learning optimizer implementation by decoupling weight decay from the gradient update, fixing a subtle but important regularization bug.

Decoupled Weight Decay

The implementation at lines 1222-1295 in tinytorch/src/07_optimizers/07_optimizers.py introduces the AdamW class. Unlike standard Adam, AdamW applies L2 regularization after the parameter update rather than adding it to the gradient.

This distinction matters because standard Adam's adaptive learning rates interact poorly with weight decay applied to gradients. When weight decay is added to gradients in Adam, the effective regularization strength varies with the historical gradient magnitude, weakening regularization for parameters with large gradients.

Update Procedure

AdamW follows the same moment estimation as Adam but modifies the final parameter update:

1. Perform standard Adam moment calculations (without weight decay in the gradient)
2. Compute the Adam update direction: update = lr * m_hat / (sqrt(v_hat) + eps)
3. Apply the update: param = param - update
4. Apply decoupled weight decay: param = param * (1 - lr * weight_decay)

This ensures that weight decay regularizes all parameters uniformly regardless of their gradient history, typically improving generalization in deep neural networks.

Practical Usage Examples

The following examples demonstrate how to instantiate and use these deep learning optimizer implementations with the tinytorch framework.

SGD with Momentum

from tinytorch.core.tensor import Tensor
from tinytorch.core.optimizers import SGD

# Initialize parameters

weights = Tensor([[0.5, -0.3], [0.8, 0.2]], requires_grad=True)
bias = Tensor([0.0, 0.0], requires_grad=True)

# Create optimizer with momentum and weight decay

optimizer = SGD([weights, bias], lr=0.01, momentum=0.9, weight_decay=0.001)

# Training loop

for epoch in range(100):
    # Forward pass (example: linear layer)

    output = weights @ input_data + bias
    loss = compute_loss(output, targets)
    
    # Backward pass

    loss.backward()
    
    # Parameter update

    optimizer.step()
    optimizer.zero_grad()

Adam Adaptive Optimization

from tinytorch.core.optimizers import Adam

# Adam works well with sparse gradients and varying scales

optimizer = Adam(
    model_parameters, 
    lr=0.001, 
    betas=(0.9, 0.999), 
    eps=1e-8,
    weight_decay=0.0
)

# Usage identical to SGD

loss.backward()
optimizer.step()
optimizer.zero_grad()

AdamW for Improved Regularization

from tinytorch.core.optimizers import AdamW

# AdamW applies decoupled weight decay

optimizer = AdamW(
    model_parameters,
    lr=0.001,
    betas=(0.9, 0.999),
    eps=1e-8,
    weight_decay=0.01  # Typically higher than Adam's weight_decay

)

# Training loop remains the same

Key Files in the Repository

Understanding the architecture of these deep learning optimizer implementations requires familiarity with the following source files:

File	Purpose	Location
07_optimizers.py	Contains the abstract Optimizer base class and concrete implementations of SGD, Adam, and AdamW	tinytorch/src/07_optimizers/07_optimizers.py
test_07_optimizers_progressive.py	Comprehensive test suite validating numerical correctness of optimizer updates	tinytorch/tests/07_optimizers/test_07_optimizers_progressive.py
tensor.py	Core tensor implementation supporting requires_grad and gradient storage	tinytorch/core/tensor.py
autograd.py	Automatic differentiation engine that populates Tensor.grad consumed by optimizers	tinytorch/core/autograd.py

These files collectively demonstrate how gradient computation, parameter storage, and optimization algorithms integrate within a minimal deep learning framework.

Summary

The cs249r_book repository provides educational deep learning optimizer implementations that clarify the algorithmic foundations of neural network training:

• SGD combines gradient descent with momentum smoothing and L2 regularization, using velocity buffers to stabilize updates across noisy mini-batches.
• Adam introduces adaptive learning rates through first and second moment estimation, automatically adjusting step sizes per-parameter based on historical gradient magnitudes.
• AdamW fixes Adam's regularization behavior by decoupling weight decay from the adaptive gradient updates, applying L2 penalty directly to parameters after the Adam step.
• All implementations inherit from a common Optimizer base class that standardizes gradient extraction and parameter management across the tinytorch framework.

Frequently Asked Questions

What is the difference between Adam and AdamW optimizers?

AdamW decouples weight decay from the gradient update, while Adam applies weight decay to the gradient itself. In the standard Adam implementation, weight decay gradients get scaled by the adaptive learning rate, causing parameters with large historical gradients to receive weaker regularization. AdamW applies the L2 penalty directly to the parameters after computing the Adam update, ensuring consistent regularization strength across all parameters regardless of their gradient history.

When should I use SGD instead of Adam for deep learning?

Use SGD with momentum when training large-scale vision models or when you have ample compute time for careful hyperparameter tuning. SGD often achieves better final generalization performance than adaptive methods like Adam, particularly in computer vision tasks with convolutional networks. However, SGD requires more careful learning rate scheduling and longer training times to converge. Adam works better for sparse gradients, NLP tasks, or when you need faster convergence with minimal hyperparameter tuning.

How does momentum improve the SGD optimizer?

Momentum accumulates exponentially decaying past gradients to smooth the optimization trajectory and accelerate convergence. The implementation stores a velocity buffer for each parameter that updates as v = μ * v + grad, where μ is the momentum coefficient (typically 0.9). This velocity term dampens oscillations in high-curvature directions while accelerating progress along consistent gradient directions. The result is faster convergence and reduced sensitivity to noisy mini-batch gradients compared to vanilla gradient descent.

What files contain the core optimizer logic in the cs249r_book repository?

The primary implementation resides in tinytorch/src/07_optimizers/07_optimizers.py, containing the Optimizer base class and concrete SGD, Adam, and AdamW implementations. Supporting infrastructure includes tinytorch/core/tensor.py for parameter storage and gradient tracking, and tinytorch/core/autograd.py for automatic differentiation. The test suite at tinytorch/tests/07_optimizers/test_07_optimizers_progressive.py validates numerical correctness against reference implementations.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}