Étapes d’échauffement

Warmup steps are a crucial technique in training large-scale neural networks, particularly for models like Transformers and large language models (LLMs). They involve starting the training process with a very small learning rate and gradually increasing it to the initial (target) learning rate over a specified number of steps. This approach is widely used to stabilize training and improve final model performance.

Here’s a detailed breakdown of warmup steps, why they are used, and how they work:

1. The Core Concept

In standard training, the learning rate is often set to a constant value or follows a decay schedule from the start. Warmup steps introduce a gradual increase in the learning rate at the beginning of training. Instead of jumping straight to a high learning rate, the optimizer starts with a tiny value (often close to zero) and linearly or non-linearly ramps it up to the desired initial learning rate.

2. Why Are Warmup Steps Needed?

The primary reasons for using warmup steps are related to the stability and health of the training process, especially in the early iterations:

• Avoiding Destructive Updates: At the start of training, the model's weights are randomly initialized. The gradients computed in these early steps are often large and noisy. If you apply a high learning rate immediately, the model can take massive, erratic steps in the parameter space, potentially leading to a phenomenon called "loss explosion" (where the loss becomes NaN or infinity) or pushing the model into a poor region of the loss landscape from which it cannot recover.
• Stabilizing Batch Normalization (for CNNs): In convolutional networks, warmup gives the Batch Normalization layers time to estimate the correct running statistics (mean and variance) of the activations. Early on, these statistics are inaccurate, and a high learning rate can cause instability.
• Adaptive Optimizer Bias Correction: Optimizers like Adam and RMSprop maintain an exponentially moving average of past gradients and squared gradients. At the beginning, these estimates are biased towards zero. While these optimizers have built-in bias correction, warmup provides an additional safety net, allowing the optimizer to "settle" its internal state before applying large updates.
• The "Early Phase" of Training: Research (e.g., the paper "On the Variance of the Adaptive Learning Rate and Beyond") suggests that the early phase of training is critical for finding a good "basin" in the loss landscape. Warmup helps the model navigate this initial phase more cautiously, leading to better generalization in the long run.

3. How Warmup Steps Work (Common Schedules)

Warmup is typically implemented as a schedule that is a function of the current training step (t). The two most common types are:

• Linear Warmup: The learning rate increases linearly from 0 (or a very small value) to the target learning rate (lr_max).
  • Formula: lr_t = lr_max * (t / warmup_steps)
  • This is the most common and simplest approach.

• Exponential Warmup: The learning rate increases exponentially from a small value to the target.
  • Formula: lr_t = lr_min * (lr_max / lr_min)^(t / warmup_steps)
  • This can be useful for very deep models where a slower start is beneficial.

4. The Typical Training Pipeline

Warmup is almost always combined with a subsequent decay schedule. The full learning rate schedule often looks like this:

1. Warmup Phase (e.g., first 1% to 10% of total steps): Learning rate increases from ~0 to lr_max.
2. Main Training Phase: Learning rate holds at lr_max for a period or immediately starts decaying.
3. Decay Phase: Learning rate decreases (e.g., cosine decay, linear decay, or step decay) to a near-zero value to fine-tune the model.

5. Practical Considerations

• How many steps? A common heuristic is to use 1% to 10% of the total training steps. For example, if you train for 100,000 steps, a warmup of 1,000 to 10,000 steps is typical. For very large models (like GPT-3), warmup steps can be in the thousands.
• Batch Size Interaction: Warmup is especially critical when using very large batch sizes. Larger batches provide more accurate gradient estimates, but they can also lead to sharper minima. Warmup helps to mitigate the instability associated with large-batch training.
• Implementation: Most deep learning frameworks (PyTorch, TensorFlow, Hugging Face Transformers) have built-in support for warmup schedules. For example, in Hugging Face's transformers library, you can easily set num_warmup_steps in the scheduler configuration.

Summary

In essence, warmup steps are a form of "safe start" for training. They prevent the model from making wild, destructive jumps in the parameter space during the fragile initial phase. By gradually introducing the learning rate, warmup allows the model to find a stable trajectory, leading to more robust convergence and better final performance. It is a standard, almost mandatory, component in the training recipe for modern large-scale models.