Learning Rate Scheduling for Stable Training

The Stability Dilemma

Before we dissect the schedule, we must first understand the problem it solves. The choice of learning rate presents a classic dilemma:

• A large learning rate promises faster convergence, but if it's too large, especially at the beginning of training, it can cause the optimization process to become unstable and diverge. The loss can explode as the optimizer overshoots minima in the loss landscape.
• A small learning rate is safer and less likely to diverge, but it can lead to painfully slow convergence or cause the model to get trapped in suboptimal local minima.

This dilemma is particularly acute at the very start of training. A randomly initialized network is, by definition, in a chaotic state. The initial gradients can be extremely large and uninformative. Taking a large step based on these early gradients can permanently destabilize the training process. For Transformers, this problem is known to be exacerbated by the architecture itself.

The following video provides a concise explanation of why the original Transformer architecture, in particular, is sensitive to the learning rate at the start of training.

From a more formal perspective, this instability is related to the curvature of the loss landscape, which is mathematically described by the Hessian matrix (the matrix of second-order partial derivatives). The maximum stable learning rate for an optimizer like SGD is inversely proportional to the "sharpness" of the loss function, which is the largest eigenvalue of the Hessian, ${\lambda }^H$. Crudely, ${\eta }_c\approx 2/{\lambda }^H$. At the beginning of training, this sharpness can be very high, necessitating a small learning rate.

Phase 1: Linear Warmup

The solution to this initial instability is to not start with the target learning rate, but to build up to it gradually. This is the warmup phase.

Linear warmup is the most common approach. Here, the learning rate is increased linearly from a very small value (often zero) to its target maximum value, ${\eta }_{\text{trgt}}$, over a specified number of training steps, $T_{\text{wrm}}$. The formula for the learning rate ${\eta }_t$ at step $t\le T_{\text{wrm}}$ is:

Typically, ${\eta }_{\text{init}}$ is set to 0.

So, why does this work? There are two complementary, deep-seated reasons.

1. Stabilizing Adam's Estimates: As we saw in the last lesson, Adam relies on estimates of the first and second moments of the gradient ($m_t$ and $v_t$). At the very beginning of training, these estimates are based on very few data points and are initialized at zero, making them noisy and biased. A small learning rate during the warmup phase gives Adam's moment estimates time to accumulate more information and become more stable before the optimizer starts taking large steps.

2. Navigating the Loss Landscape: As a recent research paper explains, the primary benefit of warmup is that it allows the network to "self-stabilize" and find a region of the loss landscape that can tolerate a much higher learning rate.

In essence, warmup gently guides the model out of the initial, potentially very sharp and unstable, region of the parameter space. By the end of the warmup period, the model is in a "flatter" region where it can safely handle the much larger target learning rate, enabling rapid convergence for the rest of training.

Phase 2: Cosine Decay

Once the learning rate reaches its peak after the warmup phase, we don't want to keep it there indefinitely. As training progresses and the model gets closer to a good solution, it's beneficial to decrease the learning rate to allow for finer-grained adjustments and to help the optimizer settle into a good minimum.

While a simple linear or exponential decay would work, one of the most effective and popular strategies is cosine decay, also known as cosine annealing. In this scheme, the learning rate is decreased following the shape of a cosine curve.

The intuition is that the cosine function's shape is ideal for this process: it starts with a very gentle slope (decaying slowly), allowing the model to continue exploring the landscape with a relatively high learning rate. Then, the decay accelerates before finally slowing down again as the learning rate approaches its minimum value. This final, slow phase with a very small learning rate acts like a fine-tuning stage.

The following article provides an excellent, detailed mathematical breakdown of a schedule that combines linear warmup with cosine decay.

Let's quickly re-examine the core formula for the learning rate multiplier ${\alpha }_t$ during the cosine decay phase, for a step $t\ge T_{\text{wrm}}$:

Here:

• ${\tau }_w=\frac{t-T_{\text{wrm}}}{T_{\text{total}}-T_{\text{wrm}}}$ is the fraction of the decay period that has elapsed.
• ${\alpha }_f$ is the desired final learning rate multiplier (e.g., 0.1, meaning the final LR will be 10% of the max LR).

As the reading explains, when ${\tau }_w=0$ (the start of decay), the $\cos$ term is 1, and the expression simplifies to ${\alpha }_t=1$. When ${\tau }_w=1$ (the end of training), the $\cos$ term is -1, and the expression simplifies to ${\alpha }_t={\alpha }_f$. The cosine function ensures a smooth transition between these two points. The actual learning rate is then ${\eta }_t={\alpha }_t\times {\eta }_{\text{trgt}}$.

Conclusion

In this lesson, we've deconstructed the learning rate schedule, a critical component for stable and efficient LLM training. We've moved beyond simple heuristics to understand the deep reasons behind its characteristic "ramp up, decay down" shape.

Here are the key takeaways:

• The Stability Problem: A static learning rate is suboptimal. Starting with a high learning rate causes instability and divergence due to the high "sharpness" of the initial loss landscape. Starting with a low one leads to slow training.
• Linear Warmup: This phase gradually increases the learning rate, which serves two purposes: it allows Adam's internal moment estimates to stabilize, and it gently guides the model into flatter regions of the loss landscape that can tolerate a high learning rate.
• Cosine Decay: After the peak, this phase provides a principled and smooth decrease in the learning rate. Its shape is ideal for allowing continued exploration before gradually "fine-tuning" the model in a promising minimum.
• The Main Benefit: Ultimately, the entire purpose of a learning rate schedule is to allow the model to train successfully with a much higher peak learning rate (${\eta }_{\text{trgt}}$) than would otherwise be possible. This high learning rate is the engine of fast convergence.

We now have the core algorithm (Adam) and a strategy for managing its most important hyperparameter ($\alpha$). However, training LLMs requires processing data in massive batches, often too large to fit into the memory of a single GPU. Our next lesson will tackle this challenge by exploring gradient accumulation, a clever technique that simulates large batch sizes without requiring more hardware memory.