In our last lesson, we navigated the complex data engineering pipeline required to transform the raw, chaotic web into a high-quality, curated dataset for pre-training. Now that we have trillions of clean tokens, we turn to the core mechanical question of learning: how does the model actually adjust its billions of parameters to get better at predicting the next token? This is the job of the optimizer.
This lesson delves into the Adam optimizer, the algorithm that has become the de facto standard for training large language models. We will dissect its mathematical formulation to understand how it cleverly combines two powerful ideas—momentum and adaptive learning rates. Our goal is not just to learn the mechanics of Adam, but to understand precisely why its design is so uniquely effective for the heterogeneous architecture of Transformers.
We'll start from the familiar ground of gradient descent and see how adding estimates of the gradient's first and second moments gives rise to Adam's robust performance.
From Vanilla Gradient Descent to Intelligent Updates
You're already familiar with the basic principle of Stochastic Gradient Descent (SGD). At each step, we update the model's weights by taking a small step in the direction opposite to the gradient of the loss function :
While effective, SGD has limitations. It uses a single, fixed learning rate for all parameters, which can lead to slow convergence if the loss landscape is shaped like a long, narrow ravine. Furthermore, it can oscillate or get stuck. To improve upon this, two key concepts were introduced:
- Momentum: Instead of only using the current gradient, we can accumulate a "velocity" based on an exponentially weighted average of past gradients. This helps the optimizer build speed in consistent directions and dampens oscillations. This moving average is an estimate of the first moment (the mean) of the gradient.
- Adaptive Learning Rates: Different parameters might need different step sizes. For instance, parameters for rare features might need larger updates than parameters for common features. Algorithms like AdaGrad and RMSProp achieve this by scaling each parameter's update based on the historical magnitude of its gradients. This is achieved by tracking the second moment (the uncentered variance) of the gradient.
Adam, which stands for Adaptive Moment Estimation, elegantly combines both of these principles.
The Building Blocks: Momentum and RMSProp
To fully appreciate Adam, let's quickly review the two core ideas it merges. The following video provides excellent visual intuitions for both momentum and RMSProp.
Who's Adam and What's He Optimizing? | Deep Dive into Optimizers for Machine Learning!
This video from Sourish Kundu, "Who's Adam and What's He Optimizing?", clearly explains the concepts that form the basis of Adam.
First, watch the segment on Momentum from the start of the section to the mathematical explanation. Focus on the physical analogy of a ball rolling down a hill, gaining speed. Next, watch the explanation of RMSProp from its introduction to the discussion on AdaGrad. The key insight here is how RMSProp adapts the learning rate separately for each parameter, dampening updates for parameters with large gradients and amplifying them for those with small gradients.
The Adam Algorithm: A Step-by-Step Breakdown
Adam takes the "velocity" term from momentum and the per-parameter scaling from RMSProp and combines them into a single, powerful update rule. It maintains two moving averages for each parameter:
- First Moment Estimate (): The mean of the gradients (the "momentum" part).
- Second Moment Estimate (): The uncentered variance of the gradients (the "adaptive" part).
Let's walk through the algorithm. The following video by Andrew Ng provides a clear, formal derivation of the steps.
Adam Optimization Algorithm (C2W2L08)
This video, part of the DeepLearning.AI specialization, gives a concise and mathematically precise explanation of the Adam algorithm.
Watch the video from the beginning of the implementation details. As you watch, follow the steps below, which correspond to the video's explanation: Update Moment Estimates (Momentum and RMSprop updates): At each iteration t, with gradient g_t, we update the moving averages m_t and v_t using hyperparameters \beta_1 (typically ~0.9) and \beta_2 (typically ~0.999): m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2 Apply Bias Correction (Bias correction): Since m_0 and v_0 are initialized to zero, the estimates are biased towards zero, especially during the initial steps. Adam corrects for this bias: \hat{m}_t = \frac{m_t}{1 - \beta_1^t} \hat{v}_t = \frac{v_t}{1 - \beta_2^t} Perform the Weight Update (Weight update): Finally, the weights are updated using the bias-corrected estimates. Notice how \hat{m}_t determines the direction of the update, while \sqrt{\hat{v}_t} normalizes its magnitude on a per-parameter basis. w_{t+1} = w_t - \alpha \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} Here, \alpha is the master learning rate, and \epsilon is a small constant (e.g., 10^{-8}) for numerical stability. For a textual reference of these equations, you can consult the following article.
For a textual reference of these equations, you can consult the following article.
The Math behind Adam Optimizer
This article from Towards Data Science provides a clear, step-by-step written description of the Adam algorithm's mathematics.
Focus on the section "The Mathematics Behind Adam", which lists the six steps of the algorithm, including initialization, moment updates, bias correction, and the final parameter update.
The empirical performance of Adam is remarkable. As shown in the chart below, when training a neural network on the MNIST dataset, Adam (purple line) converges significantly faster and to a lower training cost than other popular optimizers.
Convergence of various optimization algorithms on an MNIST classification task. The y-axis (training cost) is on a log scale. Adam demonstrates faster convergence compared to SGD with Nesterov momentum, AdaGrad, AdaDelta, and RMSProp.
Why is Adam So Well-Suited for Transformers?
We've seen how Adam works. Now for the crucial question: why is it the optimizer of choice for virtually all large language models? The answer lies in the unique architecture of the Transformer.
Unlike a Convolutional Neural Network (CNN), which is a largely homogeneous stack of identical layers, a Transformer is highly heterogeneous. It's composed of functionally distinct parameter blocks:
- Token and positional embedding layers
- Query, Key, and Value projection matrices in each attention head
- Output projection matrices in each attention head
- The two large linear layers of the Position-wise Feed-Forward Network
- Layer normalization parameters
- The final output head mapping to the vocabulary
Each of these blocks performs a different role, and their parameters can have vastly different gradient magnitudes and dynamics during training. From a mathematical standpoint, this means the curvature of the loss landscape varies dramatically across different parameter blocks.
This is a scenario where standard SGD, with its single learning rate, would struggle. A learning rate optimized for the embedding layer might be catastrophically large for the attention projections, or vice-versa. This is where Adam's per-parameter adaptivity becomes essential. The term in the denominator automatically normalizes the update for each parameter, effectively giving each one its own adaptive learning rate.
A recent research paper provides a compelling explanation for this phenomenon through the lens of the Hessian matrix—the matrix of second partial derivatives that describes the local curvature of a function.
Why Transformers Need Adam: A Hessian Perspective
This paper from NeurIPS 2024 offers a rigorous explanation for Adam's success in training Transformers. Given your background, the Hessian-based analysis will provide the deep theoretical insight you're looking for.
First, read the Abstract to get the paper's main thesis. Next, focus on the section "Why SGD underperforms...". The key concept introduced here is "block heterogeneity". Pay attention to how the authors define this as the large differences in the Hessian spectra across parameter blocks in Transformers, a feature not present in more homogeneous architectures like CNNs.
The core finding of the paper is that Transformers exhibit significant block heterogeneity. The Hessian spectrum (the set of eigenvalues of the Hessian, which characterize its curvature) is vastly different for the embedding block versus the attention block versus the MLP block. SGD is ill-equipped to handle this because its single learning rate cannot simultaneously be optimal for all blocks.
Adam, by normalizing each parameter's update by an estimate of its historical gradient variance, effectively counteracts this heterogeneity. It automatically "dampens" the updates for parameters in high-curvature regions (large gradients) and "accelerates" updates for parameters in low-curvature regions (small gradients), allowing the entire model to train in a much more stable and efficient manner.
Conclusion
In this lesson, we dissected the Adam optimizer and explored its profound suitability for training Transformers. You now have a deep, theory-grounded understanding of this critical component of the LLM training stack.
The key takeaways are:
- Adam combines two powerful ideas: It uses an estimate of the first moment of the gradients (mean) to incorporate momentum and an estimate of the second moment (uncentered variance) to create adaptive, per-parameter learning rates.
- Bias correction is crucial: Early in training, the moment estimates are corrected to counteract their initialization bias towards zero.
- The update rule is a normalized step: The final update uses the mean of the gradients for direction and the square root of the variance of the gradients for normalization, effectively balancing the step size for every single parameter.
- Transformers are heterogeneous: The diverse components of a Transformer create a complex loss landscape with varying curvature across different parameter blocks.
- Adam tames heterogeneity: Adam's adaptive, per-parameter nature is precisely what is needed to navigate this complex landscape, making it far more effective than an optimizer with a single, global learning rate like SGD.
While Adam automatically adapts the learning rate for each parameter, the overall magnitude of the updates is still governed by the master learning rate, . Choosing the right value for and, more importantly, changing it over the course of a multi-trillion-token training run, is critical for success. In our next lesson, we will explore learning rate scheduling, examining common strategies like linear warmup and cosine decay that guide the optimizer toward a good solution.