Learning Rate Warmup: Linear Warmup and Large Batch Training

Article links

Make inline references clickable

Training a large neural network from scratch is a delicate process. The parameters start randomly initialized, the gradients are enormous, and the optimizer has no accumulated history to guide it. If you immediately start training with your target learning rate, those first few updates can be catastrophically large, sending weights into regions of the loss landscape that are hard to recover from. Learning rate warmup is the solution: instead of jumping straight to your intended learning rate, you start very small and gradually increase it over a defined number of steps, giving the model time to settle into a stable training regime before committing to large gradient steps.

Warmup is a practical heuristic with firm theoretical grounding, and it is standard practice in virtually every large-scale language model training run. GPT, BERT, T5, LLaMA, and nearly every transformer trained since 2018 uses some form of learning rate warmup. Understanding why warmup works, how to implement it, and how to tune it is essential knowledge for anyone training neural networks at scale.

Think of warmup as the equivalent of warming up a car engine before driving it hard. The engine works best at operating temperature, not straight from a cold start. Similarly, the optimizer works best when its internal statistics have had a chance to reflect the actual training dynamics of the model, not the chaotic, high-noise conditions of step one. The metaphor extends further: just as a cold engine can handle gentle acceleration but stalls under aggressive throttle, the optimizer at step one can handle small updates safely but collapses under large ones.

To understand warmup, you first need to understand what goes wrong without it. The instability of early training is not a single problem but a combination of three overlapping issues that each make the others worse.

At the start of training, all model parameters are drawn from a random distribution. For a transformer with billions of parameters, this means the function the model computes is essentially random noise. The loss is high, the gradients are large, and the curvature of the loss landscape is steep and jagged. There is no structure yet, no learned representations, and no reliable gradient signal.

The problem is that adaptive optimizers like Adam, which you encountered in the Adam Optimizer chapter, estimate gradient statistics from past updates. Adam maintains two moving averages:

• : the first moment, tracking the exponentially weighted mean of recent gradients
• : the second moment, tracking the exponentially weighted mean of squared gradients (a proxy for variance)

At step , both and are zero, having just been initialized. Adam applies bias correction to account for this, but the estimates are still unreliable because they are based on a single gradient observation. No single sample tells you much about a distribution.

To understand the magnitude of this problem, consider how Adam computes the effective parameter update. Given the corrected moment estimates, the update rule is:

where:

• is the parameter update applied at step
• is the learning rate you specify
• is the bias-corrected first moment estimate (mean of recent gradients)
• is the bias-corrected second moment estimate (mean of squared gradients)
• is a small constant (typically ) added for numerical stability

The bias correction factor deserves careful attention. With typical , this correction factor at step evaluates to:

This means the corrected second moment is 1000 times the raw second moment at step 1. Since the second moment estimate was initialized at zero and has only received a single noisy gradient sample, this correction amplifies even a tiny initial gradient variance estimate to a very large value. The result is that the effective step Adam takes at step 1 is not the step you intended when you chose your learning rate. It behaves unpredictably.

This is a fundamental property of bias correction in exponential moving averages: the correction is necessary to get unbiased estimates, but it comes at the cost of high-variance estimates in the earliest steps. The statistics simply have not had enough data yet to be trusted.

Early in training, the gradient direction is particularly unreliable. The model has not yet learned any meaningful representations, so gradients are computed over a random loss landscape. Two batches drawn at step 1 might produce gradient directions that point nearly opposite to each other, whereas by step 10,000, gradients have accumulated structure and roughly agree on which direction to move.

This unreliability of gradient direction is the fundamental bottleneck that warmup is designed to address. The key insight is that gradient variance decreases as the model begins to learn. Once the model has identified some structure in the data (even rudimentary patterns like common word frequencies or basic syntactic relationships), the gradients computed on different batches start to point in more consistent directions. In the first few hundred steps, however, the optimizer has little evidence about which direction will reduce the loss consistently.

Taking large steps when the gradient direction is uncertain is doubly harmful. First, the large step itself may land in a bad region. Second, the Adam second-moment estimate records this large, noisy gradient as a signal that the true gradient magnitude is high, which will affect subsequent updates. A bad early step can poison the optimizer's statistics for thousands of subsequent steps, slowing down training long after the immediate instability has passed.

Transformers are particularly vulnerable to early instability because of the scaled dot-product attention mechanism. As covered in the Scaled Dot-Product Attention chapter, attention computes softmax over scaled dot products of queries and keys. When weights are randomly initialized, the dot products can be arbitrarily large, pushing the softmax into saturation, where all attention mass concentrates on a single token. This creates near-zero gradients for all other tokens, essentially making the attention layer blind to most of its input.

If the learning rate is large in those first steps, the optimizer reinforces these degenerate attention patterns before the model has had a chance to develop meaningful attention distributions. Recovery from softmax saturation is slow and painful: the gradients for non-attended tokens are near-zero, so the optimizer receives no signal pushing it to attend to those tokens, and the model gets stuck in a degenerate solution.

This sensitivity is one reason why transformer training is more warmup-dependent than simpler architectures like MLPs. The softmax function introduces a hard nonlinearity that can create these attention collapse failures, and they are much easier to prevent (via warmup) than to recover from after the fact.

Warmup addresses all three problems simultaneously. By starting with a very small learning rate:

• Adam's accumulated statistics ( and ) have time to build up reliable estimates before large steps are taken
• Gradient directions have time to stabilize as the model begins learning basic representations
• Attention weights can develop reasonable distributions before the optimizer locks them in

The core intuition is that you want the optimizer to be in a reliable, well-characterized state before making large, consequential weight updates. The first few hundred warmup steps are almost like a reconnaissance phase: the model learns the rough terrain of the loss landscape at a safe, slow pace before committing to the large steps that characterize the main training run.

There is also a subtler benefit that is often overlooked. During warmup, the model learns from every batch at a learning rate that is small enough to prevent overcommitting to any single batch's gradient. This gives the optimizer a better picture of the average gradient direction across many batches before it starts making large updates based on that direction. It is analogous to averaging multiple noisy measurements before acting on them.

The concept of starting with a small learning rate and increasing it was practiced informally for years before it received its modern formulation. Early practitioners of neural network training in the 1990s and 2000s would sometimes manually reduce their learning rate if training diverged, or start with conservative values and adjust over time. But these practices were ad hoc and not systematically understood.

The modern warmup protocol entered widespread use through a confluence of three developments. First, the emergence of very deep networks around 2015 made training stability a first-class concern: networks with 100+ layers were dramatically more sensitive to early gradient dynamics than shallow networks. Second, the introduction of adaptive optimizers like Adam in 2014 created a new class of optimizer-specific instability that was not present with SGD (where the moment statistics issue does not arise in the same form). Third, the explosive growth of batch sizes driven by multi-GPU training created the large-batch regime where warmup became essential rather than optional.

The "Attention Is All You Need" paper (Vaswani et al., 2017) brought warmup to mass attention by baking it directly into the transformer training recipe. Their inverse-square-root schedule was the first widely replicated warmup-plus-decay formula, and every major language model since has used some variant of the pattern. The "Accurate, Large Minibatch SGD" paper (Goyal et al., 2017) provided the theoretical justification for why warmup is necessary at large batch sizes, linking it formally to the linear scaling rule. Together, these two papers established warmup as a non-negotiable ingredient of large-scale neural network training.

To build deeper intuition for warmup, it helps to think about the loss landscape directly. This high-dimensional function maps every possible set of model weights to a loss value. Training moves the weights from a random starting point toward a region of low loss, typically a flat basin that generalizes well to unseen data.

At the start of training, the model's parameters are in a region of the loss landscape that has never been visited before. Random initialization places you somewhere on a roughly random surface. The local gradient at that starting point tells you the direction of steepest descent, but in a high-dimensional, non-convex landscape, following the steepest local gradient blindly is not a reliable strategy. Steepest descent may lead you into a narrow valley with sharp walls (a so-called sharp minimum), or it may take you to a saddle point where the gradient appears to be zero in some directions but the curvature is positive, making it a trap.

Large learning rates at the beginning of training are dangerous precisely because the curvature information at the starting point is unreliable. The loss landscape near a random initialization is not like the landscape near a well-trained point. Near a random initialization, many loss directions are steep and many others are nearly flat, creating a jagged, highly anisotropic surface. An optimizer that takes large steps in this terrain is likely to skip over useful regions and land in areas of high loss.

Research on the connection between learning rate and generalization has revealed an important insight: large learning rates, when applied during the main training phase (after warmup), tend to find flatter minima than small learning rates. Flat minima generalize better because they are robust to small perturbations in the weight space. Sharp minima, by contrast, have a narrow basin of attraction, and even small changes in the weights can push the model into high-loss territory.

This is one of the reasons why warmup is designed to be followed by a relatively high peak learning rate, not a conservative one. The warmup phase gets you to a region where the gradient signal is reliable and the optimizer can make informed decisions. Then the high peak learning rate allows the optimizer to explore broadly, avoiding sharp minima and finding flatter, more generalizable solutions.

The warmup-then-decay pattern divides optimization into distinct phases. During warmup, small steps move the parameters into a region where gradient estimates are more reliable. At peak learning rate, larger steps search for a good basin. During cosine decay, progressively smaller steps refine the parameters within that basin. This sequence matches the empirical observation that models trained with warmup-plus-decay consistently outperform models trained with a fixed learning rate throughout.

There is a deeper mathematical reason why warmup improves optimization on the loss landscape. Adam's second-moment estimate is an approximation to the diagonal of the gradient covariance matrix. This covariance matrix describes how much the gradient varies in each direction in parameter space. Parameters with high gradient variance are uncertain, and Adam compensates by taking smaller steps along those dimensions. Parameters with low gradient variance are reliable, and Adam can take larger steps.

At initialization, the gradient covariance is essentially unknown. The first few gradient samples are dominated by the noise of random initialization, not by the true curvature structure of the loss landscape. As training proceeds through the warmup phase, the optimizer accumulates gradient samples and builds up a picture of the curvature. By the time warmup ends and the learning rate reaches its peak, the optimizer has a reasonably accurate model of which directions in parameter space are risky (high curvature) and which are safe (low curvature). This is exactly the kind of terrain information you want before making large steps.

The warmup duration, viewed through this lens, is the minimum number of gradient samples needed to estimate the curvature structure reliably enough to take large steps safely. This connects back to the earlier derivation: the threshold of is the number of samples after which the exponential moving average of squared gradients has a coefficient of variation below 10%, meaning it is a reasonably accurate estimate of the true gradient variance in each direction.

The most common warmup strategy is linear warmup, where the learning rate increases from zero (or a small initial value) to the target learning rate in a straight line over steps.

At training step , the learning rate during warmup is computed as a simple linear fraction of the peak learning rate:

where:

• is the learning rate at step
• is the target (peak) learning rate, the value you have tuned for stable training
• is the total number of warmup steps (the duration of the warmup phase)
• ranges from to

At step , the learning rate is , which is nearly zero for typical warmup durations. At step , the fraction equals 1 and the learning rate reaches its peak value .

After warmup completes (at ), the learning rate transitions to a decay schedule. If using cosine decay (covered in the next chapter), the combined schedule becomes:

where:

• is the minimum learning rate at the end of training (often )
• is the total number of training steps across the entire training run

The cosine term evaluates to at the start of the decay phase (progress ), giving , and evaluates to at the end (progress ), giving . This creates a smooth, S-shaped transition from peak to minimum.

Why Linear and Not Exponential?Link Copied

How long should warmup last? This is one of the most commonly asked hyperparameter questions in large-scale training, and the answer depends on several factors.

For most transformer training runs, warmup should last between 0.5% and 2% of total training steps. Common specific recommendations from landmark training runs:

• BERT used 10,000 warmup steps out of approximately 1,000,000 total steps (1%)
• GPT-2 used 2,000 warmup steps
• T5 used 10,000 warmup steps
• LLaMA and its successors typically use 2,000 warmup steps

For smaller models or shorter training runs, you can use fewer warmup steps. A model trained for 10,000 steps might use only 100-500 warmup steps. For very large models (billions of parameters) trained for trillions of tokens, warmup up to 2,000-5,000 steps is typical because the peak learning rate is often higher and the consequences of early instability are more severe.

One important observation is that the absolute number of warmup steps matters more than the percentage. A model trained for 200,000 steps and one trained for 50,000 steps may both use 2,000 warmup steps, because the constraint on warmup duration comes from the optimizer's statistics stabilization time, not from the total training budget.

The minimum warmup duration is bounded by how long it takes for Adam's second-moment estimate to become reliable. For a parameter with gradient variance , the bias-corrected second-moment estimate after steps has variance:

where:

• is Adam's bias-corrected second-moment estimate at step
• is the true variance of the gradient for a given parameter
• is Adam's second-moment decay rate (typically 0.999)

This estimate becomes reliable, meaning its coefficient of variation drops below 10%, approximately when the number of steps satisfies:

With , that threshold is steps. This matches the empirical observation that 1,000 to 2,000 warmup steps work well for most runs. The formula tells us that the warmup duration is essentially the "memory span" of Adam's second moment estimator: you need to collect enough samples for the exponential moving average to represent the recent gradient distribution accurately.

There is a direct connection to the parameter: using a smaller (such as 0.99 or 0.95) allows Adam's second-moment estimate to adapt faster, because the exponential moving average forgets old samples more quickly and new samples have more weight. This means you can use shorter warmup with a smaller . Some recent work on adaptive optimizers explicitly reduces for the first few thousand steps for this reason. But for the standard configuration, you should plan for at least 1,000 warmup steps.

The consequences of getting warmup duration wrong in either direction are asymmetric, which makes this a safety-critical hyperparameter.

If warmup is too short (say, 100 steps when training for 100,000 steps with a large learning rate), the learning rate spikes before Adam has reliable statistics, potentially causing:

• Loss spikes where training destabilizes temporarily and the loss jumps to a much higher value before recovering
• Slow recovery as the optimizer backs away from bad parameter regions
• In severe cases, training divergence that requires a full restart from an earlier checkpoint

If warmup is too long (say, 20% of total training), you lose training efficiency by spending too many steps at low learning rates. The model converges more slowly, and you effectively reduce the average learning rate across training. The cost is not catastrophic but measurable: a training run with 20% warmup may require 10-15% more total compute to reach the same loss as one with 2% warmup.

The sweet spot is the range where warmup is long enough for Adam's statistics to stabilize but short enough that the model spends the majority of its budget training at or near the peak learning rate. In practice, erring slightly on the side of longer warmup is safer than shorter warmup: the efficiency cost of slightly too-long warmup is small, while the cost of slightly too-short warmup can be a diverged training run.

The relationship between warmup and batch size is one of the most practically important aspects of this topic. When you increase batch size to speed up training, you must also increase the learning rate, and warmup becomes significantly more important.

As covered in more detail in the Large Batch Training chapter, the linear scaling rule states that when you multiply batch size by a factor , you should multiply the learning rate by the same factor to maintain equivalent training dynamics:

where:

• is the batch size scaling factor (new batch size divided by the baseline)
• is the learning rate that was tuned and validated for the reference batch size
• is the learning rate to use at the scaled batch size

The intuition behind this rule is that a larger batch provides a lower-variance gradient estimate, because you are averaging over more samples. This makes each gradient step more informative and more representative of the true gradient direction. You can therefore take a proportionally larger step without the same risk of moving in a noisy direction. The learning rate compensates for the reduced noise by taking a larger step, maintaining roughly the same effective update per token seen.

The linear scaling rule is an approximation that holds well for moderate batch size increases but begins to break down at very large batch sizes. For most practical purposes (batch sizes up to ~8,192), it is the standard starting point.

The linear scaling rule interacts with warmup in a critical way. When you scale up the learning rate proportionally with batch size, the peak learning rate becomes much larger. A 16x batch size increase means a 16x larger peak learning rate, which amplifies all the instability problems of early training:

• Adam's unreliable early statistics lead to even larger initial updates when multiplied by the higher learning rate
• Gradient noise affects updates by a larger magnitude, so a bad step in the wrong direction has worse consequences
• Attention saturation happens faster if a large step locks in degenerate patterns before meaningful representations form

This is why the scaling rule must be paired with a proportionally longer warmup. If you double the batch size (and thus the learning rate), you should also roughly double the warmup duration. The general recommendation from Goyal et al. (2017) is to scale warmup duration linearly with the learning rate scaling factor:

where:

• is the warmup duration to use at the new (larger) batch size
• is the warmup duration tuned for the baseline batch size
• is the same batch size scaling factor from the learning rate rule above

For example, if your baseline configuration uses a batch size of 256 with a learning rate of and 1,000 warmup steps, and you scale to a batch size of 4,096 (a 16x increase), you should use:

• Learning rate:
• Warmup steps:

Without the warmup increase, the 1.6e-3 learning rate at step 1 would produce enormous updates on an optimizer with unreliable statistics, almost certainly causing training instability.

The Goyal et al. 2017 paper "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour" made warmup standard practice for large-batch training. The paper's contribution was twofold: it formulated the linear scaling rule precisely, and it demonstrated empirically that gradual warmup was necessary to make the rule work in practice.

The key experiment showed that training ResNet-50 on 256 GPUs with a batch size of 8,192 failed without warmup. The loss would spike and the model would fail to converge. With 5 epochs of gradual linear warmup, the model trained successfully and matched the accuracy of a model trained with a 32x smaller batch size. This result established warmup as a prerequisite for large-batch training, not a minor tuning convenience.

For language models, subsequent transformer training papers extended these findings. Without warmup, the initial destabilization from high learning rates can cause loss spikes that occasionally grow so large that the model never recovers, especially in deeper architectures where the gradient signal must propagate through many layers.

For very large batch sizes (beyond 8,192), the linear scaling rule begins to break down. The fundamental reason is that gradient noise reduction does not scale linearly with batch size: doubling the batch size reduces the gradient standard deviation by a factor of , not by a factor of 2. The noise reduction follows the law of large numbers, which is a square-root relationship.

At extreme batch sizes, the signal-to-noise ratio plateaus because the dominant source of gradient variance shifts from sampling noise (which scales as ) to structural noise from the data distribution itself. Continuing to increase the learning rate linearly once this plateau is reached leads to diminishing or negative returns.

An alternative scaling rule uses square-root rather than linear scaling:

This is more conservative at large batch sizes and reduces the peak learning rate, which in turn reduces the required warmup duration. Both scaling rules are in active use, and the choice between them depends on the specific model architecture, training setup, and the batch size regime you are operating in. A practical heuristic is to use linear scaling up to (batch sizes below ~8,000 if your baseline is 256) and to consider square-root scaling for larger scale factors.

Linear warmup dominates modern practice, but it is worth knowing the other variants that practitioners use and why they exist. Each variant makes a different tradeoff between how fast the learning rate grows in the early steps and how quickly the optimizer transitions to peak-rate training.

Some practitioners use a two-phase approach: start at a fixed small learning rate for a few hundred steps, then ramp linearly to the peak. The logic is that the very first steps are so noisy that even a small linear growth feels premature. By holding constant at, say, for the first 100 steps, you give the model a chance to take a handful of consistent small steps before beginning the ramp. In practice, this variant produces results nearly identical to linear warmup, and the extra hyperparameter (the length of the flat phase) is rarely worth the tuning cost.

Polynomial warmup generalizes linear warmup by allowing a power in the growth function:

With you recover linear warmup. With (such as or ), the schedule is convex, growing slowly at first and then accelerating. This creates a gentler start than linear warmup at the cost of spending less time at intermediate learning rates. With (such as , the square-root case), the schedule is concave, growing quickly early and then plateauing. This gives the model more time at moderate learning rates during warmup.

For most applications, (linear) is the right choice. The gains from tuning are small compared to the cost of searching that hyperparameter space. But polynomial warmup appears in some frameworks (such as the Hugging Face transformers library's get_polynomial_decay_schedule_with_warmup) because it generalizes both the linear and square-root cases.

A less common variant uses the cosine function itself for warmup, analogous to how cosine is used for decay:

This produces an S-shaped growth curve that starts very slowly, accelerates through the middle of warmup, and then slows down as it approaches the peak. The motivation is that this shape mirrors the cosine decay phase, creating a symmetric schedule where both warmup and decay transition smoothly. Some practitioners find this aesthetically pleasing and claim a slight improvement in the earliest training stability, but empirical evidence for this is weak.

A recent and practically important variant is the Warmup-Stable-Decay schedule, which divides training into three phases rather than two. After the warmup ramp, the learning rate stays constant at the peak value (the stable phase) for the majority of training, then decays sharply in the final phase. This pattern emerged from work on continued pretraining and incremental training, where practitioners discovered that a long flat phase followed by a sharp cooldown often outperforms continuous cosine decay. The WSD schedule is increasingly used in modern LLM training because it is more compatible with training runs that are not committed to a fixed total step count from the beginning.

The case for warmup changes depending on which optimizer you use. Warmup was originally motivated largely by the properties of Adam, but the same principles apply, with different degrees of urgency, across the optimizer landscape.

For Adam and AdamW, which are the overwhelmingly dominant choice for language model training, warmup is essentially mandatory. The bias correction problem described earlier applies directly, and the second-moment stabilization argument gives you a concrete lower bound for warmup duration. The standard configuration requires at least 1,000 steps of warmup, and most runs use 2,000-5,000 steps.

The AdamW variant, which corrects the weight decay implementation (decoupling it from the gradient update), does not change the warmup requirements relative to Adam. The moment estimation dynamics are identical. AdamW is uniformly preferred over Adam for large-scale pretraining because of better regularization behavior, but your warmup settings should be the same.

For SGD with momentum, warmup is less theoretically necessary but still empirically beneficial. SGD does not have second-moment estimation, so the bias correction issue does not apply. However, SGD still accumulates a momentum buffer that can be problematic at high learning rates in early training. The momentum buffer starts at zero and can oscillate when the gradients are noisy and the learning rate is large, producing oscillatory parameter trajectories that slow convergence. Linear warmup reduces this risk.

The Goyal et al. 2017 paper used SGD for its ImageNet training experiments, establishing that warmup is beneficial for SGD as well, particularly at large batch sizes. For language models, SGD is rarely used (Adam dominates), but the principle transfers: warmup is a safe default for any optimizer.

Newer optimizers like Adafactor and Sophia are designed with large-scale training stability in mind, and they often incorporate warmup-like behavior internally. Adafactor, which factorizes the second-moment estimate to reduce memory usage, has built-in scale-invariant updates that partially compensate for early estimation noise. Some Adafactor implementations use a relative learning rate that grows as , which is effectively an automatic warmup. If you use Adafactor with its default schedule, you may not need to add explicit warmup. But when using a fixed learning rate with Adafactor, the same warmup logic applies.

Modern large language model development rarely follows a single training run from scratch to deployment. Instead, it involves multiple training stages: pretraining on a large general corpus, continued pretraining on domain-specific data, supervised fine-tuning, and reinforcement learning from human feedback. Each stage transition creates a new set of challenges for learning rate scheduling, and warmup plays a role at every stage.

When you transition from one training stage to the next, the model's weights are in a good state (they were learned during the previous stage), but the optimizer's accumulated statistics may not match the new training distribution. The Adam moment vectors and reflect the gradient distribution of the previous stage. If the new stage uses a different dataset, a different loss function, or a different mix of tasks, the pre-accumulated moment estimates can be misleading.

The conservative approach is to reset the optimizer state entirely at each stage transition and apply a fresh warmup from a low learning rate. This ensures the optimizer re-adapts to the new gradient distribution before taking large steps. The cost is that you lose the momentum information from the previous stage, which means the optimizer may converge slightly more slowly at the start of each new stage.

A less conservative approach is to keep the optimizer state but apply a short warmup of a few hundred steps. This treats the residual moment estimates as prior information that can be updated quickly by the new stage's gradients. In practice, this approach works well when the stage transition is not too dramatic (for example, continuing pretraining on a different text domain) but may fail if the distribution shift is large (for example, switching from masked language modeling to reinforcement learning from human feedback).

One approach that has been studied for multi-stage training is the cosine annealing with warm restarts (SGDR) strategy. In this approach, training alternates between a warmup ramp and a cosine decay, cycling repeatedly. Each cycle starts fresh from the peak learning rate, performs a short warmup, then decays back to a minimum. The cycle length is often increased geometrically (each restart cycle is longer than the previous).

The motivation is that each warmup-decay cycle gives the model a chance to explore a new region of the loss landscape (during the high learning rate phase) and then refine its position within that region (during the decay phase). The repeated restarts act like a sequence of increasingly refined searches. In practice, this strategy works well for continued pretraining and fine-tuning scenarios where you want the model to adapt to new data without forgetting the structure learned in previous cycles.

When fine-tuning a pretrained language model on a downstream task, the warmup requirements are quite different from pretraining. The model starts from a very good initial point (the pretrained weights), and the goal is to adapt the weights gently to the task without disrupting the broad representations learned during pretraining.

For fine-tuning, shorter warmup periods are generally appropriate. A common recommendation is 6-10% of total fine-tuning steps for warmup, which is higher as a fraction of total steps than in pretraining, because the total number of fine-tuning steps is much smaller (often 1,000-10,000 rather than 100,000-1,000,000). In absolute terms, fine-tuning warmup is often only a few hundred steps.

The peak learning rate for fine-tuning should also be much smaller than the pretraining peak: typically to rather than to . This conservative learning rate reduces the risk of catastrophic forgetting (covered in the Catastrophic Forgetting chapter), where fine-tuning on the task overwrites the general representations acquired during pretraining.

Instruction tuning, which trains a model to follow natural language instructions by exposing it to instruction-following demonstrations, has become a standard step in the development of conversational AI systems. Instruction tuning is a form of supervised fine-tuning, but it typically uses a much more diverse set of tasks than traditional fine-tuning, including summarization, question answering, code generation, and mathematical reasoning.

For instruction tuning, the warmup configuration is influenced by the number of gradient steps and the peak learning rate. Instruction tuning runs are often very short (a few thousand to tens of thousands of steps), so warmup of 100-500 steps is typical. The learning rate is kept conservative to prevent the instruction tuning from washing out the pretraining representations. The combination of short warmup, low peak learning rate, and gradual decay has been validated across many instruction-tuning setups including the original InstructGPT experiments and subsequent work.

To make warmup concrete, let's trace through a specific training configuration and see exactly what the learning rate schedule looks like.

Consider training a medium-size transformer (300M parameters) with these settings:

• Total training steps: 100,000
• Warmup steps: 2,000
• Peak learning rate:
• Minimum learning rate:
• Decay schedule: cosine decay after warmup

At step , the warmup formula gives:

This is nearly zero. Adam's statistics are unreliable at this point, and this tiny learning rate ensures the first update is tiny and harmless regardless of the gradient magnitude.

At step , halfway through warmup:

Adam has now seen 1,000 gradient observations and has reasonable second-moment estimates. The learning rate is half the peak, which is a meaningful training rate while still being conservative enough for the optimizer's current state.

At step , warmup complete:

The learning rate reaches its peak. Adam's second-moment estimates have stabilized (we crossed the threshold), and the model has begun to develop basic representations. Cosine decay begins from this point.

At step , end of training, the cosine decay progress is , giving:

The learning rate ends at the minimum, one-tenth of the peak. This final low learning rate allows the model to fine-tune its parameters without overshooting the local minimum it has found during the main training phase.

Now consider what happens when you scale the same 300M parameter model to a larger compute cluster, increasing batch size from 256 to 4,096 (a 16x increase). Applying the linear scaling rule:

New peak learning rate:

New warmup steps:

The first update at step 1 with the new configuration is:

Notice that the first-step learning rate is identical to the original configuration. This is not a coincidence. When both the peak learning rate and the warmup duration scale by the same factor , their ratio stays constant, and the per-step learning rate at any fractional position within warmup is preserved. This is the mathematical reason the two configurations are training-equivalent during warmup: the gradient step sizes at each warmup fraction are the same, and the model sees the same effective learning rate trajectory relative to the total warmup budget.

This equivalence breaks at the level of absolute step count: the large-batch run completes the same fractional training milestone in far fewer wall-clock steps (because each step processes 16x more data), but the per-token compute and gradient statistics at each checkpoint are comparable. This is why the linear scaling rule with proportional warmup maintains training quality: it preserves the statistical structure of the training trajectory across batch size changes.

To make the optimizer-statistics argument tangible, let's trace exactly what happens to Adam's second-moment estimate for a single parameter across the first 2,000 steps of warmup.

Suppose a particular embedding parameter has a true gradient standard deviation of . The true gradient variance is . We want Adam's second-moment estimate to converge to this value before taking large steps.

With , the bias-corrected second moment at step converges toward the true variance as:

But the expected value alone does not tell the whole story. The variance of around its expected value is the issue. At step (very early in warmup), the estimate is based on only 10 noisy gradient observations, heavily weighted toward the initial zero. The relative error in the estimate can easily exceed 50%, meaning the optimizer might behave as if the true gradient variance is or instead of . A 2x error in the second-moment estimate translates directly to a error in the step size, since Adam's update scales as .

By step , the exponential moving average has accumulated enough samples that the estimate is reliable. The initialization influence has decayed to , meaning only 37% of the second-moment estimate's weight is still on the zero initialization, while 63% comes from actual gradient observations. At this point, the step size variation due to estimation error is under 10%, which is acceptable for taking large updates safely.

This is precisely why warmup duration and the Adam parameter are tightly linked. If you reduce to 0.99, the initialization influence at step 100 is already , matching what takes 1,000 steps with . A lower effectively speeds up the statistics stabilization, allowing shorter warmup. This relationship is exploited in some modern training setups that use to enable faster adaptation during early training and continued training stages.

Let's implement warmup schedules from scratch and visualize how they behave. We'll build a general-purpose schedule that combines linear warmup with cosine decay.

PyTorch IntegrationLink Copied

For practitioners working with the Hugging Face transformers library, built-in schedulers provide these schedules without needing to implement them from scratch:

Let's implement the linear scaling rule with proportional warmup scaling:

The important parameters for configuring learning rate warmup are:

• warmup_steps: Number of steps to ramp from 0 to peak learning rate. Typical values: 1,000-5,000 for most transformer training runs.
• peak_lr: The target learning rate reached after warmup. Must be tuned based on model size, batch size, and architecture.
• min_lr: The floor for cosine decay at end of training. Usually 10% of peak_lr (min_lr = 0.1 * peak_lr).
• beta2: Adam's second-moment decay rate. Affects how quickly statistics stabilize; lower values (0.95-0.99) allow shorter warmup.
• batch_size: Larger batches require proportionally higher learning rates and longer warmup durations.

When training fails or degrades, understanding whether the problem is warmup-related is an important diagnostic skill. Warmup failures leave characteristic signatures in the loss curve and optimizer statistics that you can learn to recognize.

The most common symptom of insufficient warmup is a loss spike in the first few hundred steps. The loss begins decreasing normally for a handful of steps (while the learning rate is very small), then suddenly jumps to a value higher than the initial loss. If warmup is too short, this spike happens at the point where the learning rate crosses the instability threshold. After the spike, training may recover and continue to converge (if the spike was mild), or the loss may plateau at a high value, or diverge entirely.

You can distinguish a warmup-related spike from other causes by checking the timing. A warmup spike will occur at approximately or shortly after, when the learning rate first reaches values high enough to cause instability. A spike at a later point in training (say, at step 50,000 in a 100,000 step run) is more likely due to a bad batch, a numerical overflow in the model, or gradient clipping interaction.

A useful diagnostic is to monitor the gradient norm (the L2 norm of all gradient tensors concatenated) throughout training. Under healthy training, the gradient norm should start relatively high, decline during the warmup phase as the model begins to learn structure, and then stabilize. If the gradient norm spikes sharply during early training, this indicates that warmup was not sufficient to prevent destabilizing updates.

Many training frameworks log gradient norms automatically, and plotting the gradient norm against the learning rate schedule gives you a direct view of the interaction between warmup dynamics and gradient behavior.

A more targeted diagnostic is to track the second-moment estimate for a representative set of parameters. The second-moment estimate should start near zero, grow rapidly in the first few hundred steps, and then stabilize as it converges to its long-run estimate. If you observe growing erratically or oscillating in early training, this is a sign that the learning rate was too high before the second-moment estimate had stabilized, confirming that warmup was too short.

In practice, you can log optimizer.state[param]['exp_avg_sq'] for a few parameters in PyTorch to inspect the second-moment estimates directly. Watching these values across the warmup phase gives you direct evidence of whether the optimizer statistics have stabilized.

If a training run has already failed due to insufficient warmup, the options depend on how far training has progressed. If you catch the failure early (within the first 1-2% of training), the safest option is to restart from scratch with a longer warmup. If the model has a checkpoint from before the spike, you can restart from that checkpoint with a lower peak learning rate or longer warmup.

If the failure occurs after significant training, a common practice is to load the checkpoint from before the spike and continue training with a short warmup of a few hundred steps. This "mini-warmup" lets the optimizer re-stabilize after the disruption. This approach works because the model's weights are much better initialized at this point than they were at step 0, so the optimizer only needs a short period to rebuild reliable moment estimates.

Warmup is simple and powerful, but it comes with practical caveats that are important to understand before applying it to your own training runs.

One frustrating property of warmup is that the optimal warmup duration is not universal. It depends on the model size, the learning rate, the batch size, the optimizer configuration, and even the initialization scheme. A warmup duration that works well for one configuration may be too short or too long for another.

In practice, this means that when you change any major hyperparameter, including the learning rate, batch size, model architecture, or optimizer settings, you should revisit the warmup duration. The rule of thumb of 1-2% of total steps is a reasonable starting point, but significant changes in training setup may require re-tuning. This adds to the overall hyperparameter tuning burden of large-scale training. For very expensive training runs (thousands of GPU-hours), practitioners often run shorter "proxy" training runs to validate the warmup configuration before committing to the full run.

Warmup stabilizes the early training dynamics, but it does not eliminate the need for careful weight initialization. As discussed in the Weight Initialization chapter, poor initialization (such as weights that are far too large or too small) can still cause problems even with warmup. Warmup and initialization work together: initialization sets the starting point, and warmup controls how aggressively you move from it.

Models with residual connections (like transformers) are particularly sensitive. The Pre-LayerNorm vs. Post-LayerNorm placement affects gradient magnitudes at initialization, and some architectures (like GPT-3's scaled initialization) explicitly reduce the variance of residual stream projections to make training stable at scale, reducing the reliance on warmup. Warmup is a compensatory mechanism, not a replacement for proper initialization design.

Most large-scale training runs use gradient clipping alongside warmup. These two techniques interact: gradient clipping prevents individual large updates from destabilizing training, while warmup prevents the systematic pattern of large updates in early training. They complement each other, but you should not rely on gradient clipping as a substitute for warmup.

Gradient clipping operates reactively: it truncates a gradient after it has been computed but before it is applied, so it can prevent any single catastrophically large step. Warmup operates proactively: it scales the learning rate down so that even a large gradient leads to a small update. The proactive approach is more predictable and does not distort the gradient direction (clipping does). Using both together provides defense in depth.

An important modern extension of warmup is the warmup-then-cooldown approach used when continuing training from a checkpoint. If you load a pretrained model and continue training (for domain adaptation, instruction tuning, or extending the context window), you should apply a short warmup even though the model is already trained.

This is because the optimizer state (Adam's and estimates) may not transfer perfectly when training resumes after a pause or when using a different dataset. The accumulated statistics reflect the gradient distribution of the previous training dataset, which may differ significantly from the new dataset. A brief warmup period lets the optimizer re-stabilize its statistics to reflect the new training distribution.

For continued pretraining with a cosine schedule, it is common to restart the cosine decay from the checkpoint's peak learning rate with a warmup of a few hundred steps, rather than simply continuing the previous schedule. This pattern is increasingly used in the era of fine-tuning and instruction tuning of large language models, where multiple training phases are common.

Before warmup became standard practice around 2015-2018, training very large neural networks reliably was significantly harder. Loss spikes were common, training runs frequently diverged, and practitioners often needed to carefully hand-tune the learning rate schedule with conservative early learning rates that were different from their peak rates. The tooling and best practices for handling early training instability were ad hoc.

Warmup, combined with adaptive optimizers and gradient clipping, made it possible to train transformers with billions of parameters reliably across many different configurations. The combination of these techniques reduced the failure rate of large training runs dramatically and lowered the amount of engineering effort required to avoid instability. This reliability improvement is what enabled the scale-up of language models from millions to billions to trillions of parameters: when training a model costs millions of dollars of compute, you cannot afford frequent divergences due to training instability.

Warmup is not a silver bullet, but it is one of the foundational engineering decisions that makes modern large language model training tractable. It is a small piece of the puzzle that, when combined with the other techniques covered in Part XXXII: Training Optimization, enables the reliable, reproducible training of the largest language models in existence.

Learning rate warmup starts training with a very small learning rate and gradually increases it to the target value over a defined number of steps. The core motivation is threefold: it gives Adam's second-moment estimates time to stabilize, it allows gradients to find more reliable directions, and it prevents transformer attention weights from saturating before meaningful representations form.

Linear warmup is the dominant strategy, computing during the warmup phase before transitioning to a decay schedule. The warmup duration is typically 0.5-2% of total training steps, with a theoretical minimum of approximately steps based on Adam's statistics stabilization. For the standard , this means at least 1,000 warmup steps.

When scaling to larger batch sizes, both the learning rate and warmup duration should scale proportionally with the batch size ratio. This linear scaling rule is essential for preserving training stability when increasing parallelism. The square-root scaling rule provides a more conservative alternative for very large batch sizes where gradient variance reduction plateaus.

Beyond linear warmup, variants like polynomial warmup, cosine warmup, and the Warmup-Stable-Decay schedule offer different tradeoffs. For most applications, linear warmup is the right default choice. When warmup fails, characteristic loss spikes and gradient norm patterns allow you to diagnose the problem and adjust either the warmup duration or the peak learning rate.

The next chapter covers Learning Rate Decay, which describes what happens after warmup ends: how cosine schedules, step decay, and inverse square-root decay shape the learning rate throughout the bulk of training.