Deep Learning Optimizers: Deriving Momentum, RMSProp, and AdamW mathematically

Scenario context

Optimizers determine how we update neural network weights during training. While Stochastic Gradient Descent (SGD) updates parameters along the negative gradient direction, modern models use adaptive learning rates and momentum. This guide derives the mathematical models behind Momentum, RMSProp, Adam, and AdamW, and implements them from scratch in Python using raw NumPy.

When training deep neural networks, th…

This walkthrough focuses on Concept: The Need for Adaptive Learning Rates.

Before making a design choice, clarify:

• When training deep neural networks, the objective of our optimization algorithm is to navigate a complex, high-dimensional loss landscape to find a global minimum.
• The simplest optimization algorithm, Stochastic Gradient Descent (SGD) , updates the parameters $W$ by taking a step in the direction of the negative gradient of the loss: $$ W {t+1} = W t - \alpha \nabla L(W t) $$ wher…
• However, SGD exhibits major vulnerabilities in complex loss landscapes: Oscillations in Canyons: If the loss landscape is steep in one direction and gentle in another (forming a ravine or canyon), SGD will oscillate wil…

Concepts in play: deep-learning-foundations, caching, jvm-internals-and-garbage-collection

Stage 1 — Mechanics: Gradient Descent Step Modifications

• Modern optimizers evolve by adding historical tracking loops to the standard SGD update: 1.
• Momentum: Accelerates gradient descent by accumulating a moving average of past gradients.
• This acts like a heavy ball rolling down a hill, carrying momentum that helps it skip past local minima and dampen oscillations.
• RMSProp (Root Mean Square Propagation): Dampens oscillations by dividing the learning rate for a parameter by a running average of the magnitudes of its recent gradients.

Observable outcome: RMSProp (Root Mean Square Propagation): Dampens oscillations by dividing the learning rate for a parameter by a running average of the magnitudes of its recent gradients.

Stage 2 — Flow: Optimization Path Transitions

• The flowchart below visualizes the mathematical progression and logical updates introduced by each optimizer stage: The table below summarizes the key hyperparameters used to configure these mathematical update paths:

flowchart TD
    SGD[Stochastic Gradient Descent
Standard step: W - alpha * grad] -->|Add Velocity/History| Mom[Momentum
Accumulate past step directions]
    Mom -->|Combine with Squared Gradients| Adam[Adam
First and second moments tracking]
    SGD -->|Divide by Squared Gradients| RMS[RMSProp
Normalize step size by gradient magnitude]
    RMS -->|Combine with Velocity| Adam
    Adam -->|Decouple Weight Decay| AdamW[AdamW
Subtract weight decay directly from weights]

Observable outcome: The flowchart below visualizes the mathematical progression and logical updates introduced by each optimizer stage: The table below summarizes the key hyperparameters used to configure these mathematical update paths:

Stage 3 — Deep Dive: Vectorized updates and Weight Decay

• Let us derive the exact equations and analyze the performance constraints of vectorized optimizers.
• Optimizer State Update Internals During backpropagation, we calculate the gradient matrix $g t = \nabla L(W t)$.
• To run updates, the optimizer must maintain historical matrices of the same shape as $W t$ in executor memory.
• For Momentum, we maintain a velocity matrix $v t$.

Observable outcome: For Momentum, we maintain a velocity matrix $v t$.

Stage 4 — Advanced Concepts: Decoupling Weight Decay in AdamW

• L2 regularization (commonly called weight decay) adds a penalty term $\frac{1}{2} \lambda \ W\ ^2$ to the loss function to prevent overfitting.
• In standard gradient descent, the gradient of the regularized loss is: $$ \nabla L {reg}(W t) = g t + \lambda W t $$ In Adam, if we feed this regularized gradient into the moment updates: $$ m t = \beta 1 m {t-1} + (1 -…
• If a parameter has historically large gradients, its $v t$ will be large, which divides and suppresses the weight decay penalty .
• Conversely, parameters with small gradients will experience stronger weight decay.

Observable outcome: Conversely, parameters with small gradients will experience stronger weight decay.

Stage 5 — Applications: From Neural Networks to NLP Transformers

• Large Language Model Training: Almost all modern Transformers (including GPT and Llama) are trained using AdamW because it provides stable convergence for deep architectures.
• Computer Vision (ResNet): Often trained using SGD with Momentum, as it can achieve slightly higher generalization accuracy when tuned correctly.
• Generative Adversarial Networks (GANs): Use Adam to stabilize the competitive training loops between generator and discriminator networks.

Observable outcome: Generative Adversarial Networks (GANs): Use Adam to stabilize the competitive training loops between generator and discriminator networks.

End-to-end outcome

After working through the stages, you should be able to explain:

• SGD: Baseline optimizer; updates weights along the negative gradient vector.
• Momentum: Dampens oscillations by adding a running average of past gradients ($v t$).

Alternatives considered

Decision Guide: SGD vs. Adam vs. AdamW

Decision Guide: SGD vs. Adam vs. AdamW

Practical Implementation: NumPy Optimizers Comparison

Here is a complete, runnable Python script implementing SGD, Momentum, RMSProp, and AdamW from scratch.

Failure and degraded modes

Trade-offs and Failure Modes

• Detection: Hyperparameter Sensitivity: Adam and AdamW introduce several hyperparameters ($\beta 1$, $\beta 2$, $\epsilon$, $\lambda$).
• User or operator impact: Hyperparameter Sensitivity: Adam and AdamW introduce several hyperparameters ($\beta 1$, $\beta 2$, $\epsilon$, $\lambda$).

Lessons Learned: Common Optimization Pitfalls

• Detection: Confusing L2 Regularization with Decoupled Weight Decay: Standard deep learning libraries (like PyTorch) historically named their L2 regularization parameter weight decay in the Adam optimizer class.
• User or operator impact: Confusing L2 Regularization with Decoupled Weight Decay: Standard deep learning libraries (like PyTorch) historically named their L2 regularization parameter weight decay in the Adam optimizer class.