Optimizers for Transformers (Adam, AdamW)

Optimizers play a fundamental role in managing the complex loss surfaces associated with training deep neural networks. While standard Stochastic Gradient Descent (SGD) established initial methods, its fixed learning rate often proves inadequate for the scale and complexity of Transformer models. Adaptive learning rate algorithms, which dynamically adjust the step size for each parameter, have become essential tools. Among these, Adam and its refinement, AdamW, are the most widely adopted optimizers for training Transformers.

Adam: Adaptive Moment Estimation

Adam, short for Adaptive Moment Estimation, integrates ideas from both momentum methods and RMSprop. It computes individual adaptive learning rates for different parameters based on estimates of the first and second moments of the gradients.

1. First Moment (Momentum): Adam maintains an exponentially decaying average of past gradients. This acts like momentum, helping the optimizer accelerate through shallow ravines in the loss surface and dampen oscillations in directions of high curvature. The first moment estimate $m_t$ at timestep $t$ is calculated using the gradient $g_t$: $m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t$ The hyperparameter $\beta_1$ controls the exponential decay rate, typically set to a value like 0.9.

2. Second Moment (Variance Adaptation): Adam also maintains an exponentially decaying average of past squared gradients. This serves as an estimate of the variance (or more accurately, the uncentered second moment) of the gradients for each parameter. This estimate is used to scale the learning rate inversely; parameters with larger or noisier gradient histories receive smaller updates, while those with smaller or sparser gradients receive larger updates. The second moment estimate $v_t$ is calculated as: $v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2$ The hyperparameter $\beta_2$ controls this decay rate, often set to 0.999.

3. Bias Correction: Because the moment estimates $m_t$ and $v_t$ are initialized at zero, they are biased toward zero during the early stages of training. Adam incorporates bias correction terms to counteract this: $\hat{m}_t = \frac{m_t}{1 - \beta_1^t}$ $\hat{v}_t = \frac{v_t}{1 - \beta_2^t}$ These corrected estimates, $\hat{m}_t$ and $\hat{v}_t$, provide better estimates of the true moments, especially early on.

4. Parameter Update: Finally, the parameters $w$ are updated using the bias-corrected moment estimates. The update rule scales the base learning rate $\alpha$ by the ratio of the corrected first moment to the square root of the corrected second moment: $w_{t+1} = w_t - \alpha \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$ The small constant $\epsilon$ (e.g., $10^{-8}$) is added to the denominator for numerical stability, preventing division by zero.

Adam often demonstrates rapid initial convergence and performance across various deep learning tasks. However, a subtle issue arises when combining Adam with standard L2 regularization (weight decay). L2 regularization aims to prevent overfitting by adding a penalty term $\frac{\lambda}{2} ||w||^2$ to the loss function. This results in an additional term $-\lambda w_t$ being added to the gradient $g_t$. In the conventional Adam implementation, this gradient component $g_t = \nabla \mathcal{L}(w_t) + \lambda w_t$ is used to compute both $m_t$ and $v_t$. This couples the weight decay effect with the adaptive learning rate mechanism. Specifically, the effective weight decay applied during the update ($\alpha \frac{\lambda w_t}{\sqrt{\hat{v}_t} + \epsilon}$) gets scaled by the adaptive denominator $\sqrt{\hat{v}_t}$. This means weights associated with large historical gradient magnitudes (large $\hat{v}_t$) experience diminished effective decay, potentially hindering the regularization effect.

AdamW: Adam with Decoupled Weight Decay

AdamW was introduced to resolve the suboptimal interaction between weight decay and the adaptive learning rates in the original Adam algorithm. The core idea proposed by Loshchilov & Hutter (2019) is simple yet effective: decouple the weight decay update from the gradient-based update performed by Adam.

Instead of incorporating the L2 penalty into the gradient $g_t$ used for moment estimation, AdamW calculates the moment estimates and the main adaptive update step using only the gradient derived from the primary loss function, $\nabla \mathcal{L}(w_t)$. The weight decay is then applied separately and directly to the weights after the adaptive step.

The process can be summarized as:

1. Compute the gradient $g_t = \nabla \mathcal{L}(w_t)$ based only on the task loss.
2. Update the first and second moment estimates ($m_t$, $v_t$) and compute their bias-corrected versions ($\hat{m}_t$, $\hat{v}_t$) using $g_t$, as in standard Adam.
3. Calculate the main Adam update step (without the learning rate factored in yet): $\Delta w'_t = \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$
4. Update the weights by applying both the scaled adaptive step and the decoupled weight decay. Using $\eta_t$ as the scheduled learning rate at timestep $t$: $w_{t+1} = w_t - \eta_t (\Delta w'_t + \lambda w_t)$ Alternatively, and closer to the original paper's formulation, the decay can be applied multiplicatively before the update: $w_{t+1} = (1 - \eta_t \lambda) w_t - \eta_t \Delta w'_t$. The main insight is that the decay term $\lambda w_t$ is no longer scaled by $1 / (\sqrt{\hat{v}_t} + \epsilon)$.

By decoupling, AdamW ensures that weight decay applies more consistently across all parameters, proportional only to the weight magnitude $w_t$ and the learning rate $\eta_t$, irrespective of the gradient history captured in $\hat{v}_t$. This modification often leads to improved model generalization and better final performance compared to standard Adam with L2 regularization, particularly for complex models like Transformers where effective regularization is highly beneficial. Consequently, AdamW has become the de facto standard optimizer for training modern Transformers.

Hyperparameters and Approaches

Effectively using Adam or AdamW requires setting several hyperparameters:

• Learning Rate ($\alpha$ or $\eta_t$): This remains a highly sensitive hyperparameter. For Transformers, it's almost never constant; instead, a learning rate schedule involving an initial "warmup" phase followed by a decay phase is critical (as discussed in the subsequent section). Typical peak learning rates fall within the $10^{-5}$ to $10^{-3}$ range, influenced by factors like model size, batch size, and the specific schedule used.
• $\beta_1$: The exponential decay rate for the first moment estimate. The default value of 0.9 works well in most cases.
• $\beta_2$: The exponential decay rate for the second moment estimate. While the default of 0.999 is common, some studies suggest slightly lower values (e.g., 0.98, 0.99) might occasionally offer benefits.
• $\epsilon$: A small value to ensure numerical stability. Defaults like $10^{-8}$ or $10^{-6}$ are usually sufficient.
• Weight Decay ($\lambda$): Applicable directly in AdamW, or as an L2 penalty coefficient if using standard Adam. This is an important regularization hyperparameter requiring tuning through experimentation. Common starting points are around 0.01 or 0.1, but the optimal value is dataset and model-dependent. Setting $\lambda=0$ effectively disables weight decay.

In summary, while Adam represented a major step forward in optimization, AdamW provides a more refined approach to incorporating weight decay, leading to demonstrable improvements in training large language models. It is generally the recommended choice for Transformer architectures, always used in conjunction with carefully tuned learning rate schedules and regularization parameters.