Concepts / Cosine Annealing for Learning Rate Schedules

Cosine Annealing for Learning Rate Schedules

Linear warmup gradually increases the learning rate from initial_lr to peak_lr during the first warmup_steps batches of training.

  • Programming

A Learning Rate That Moves

The learning rate controls how large a change the optimizer makes to a model's weights on each update. In many simple training loops it stays at one value for the entire run. The training function in this article does something different: it recalculates the learning rate on every step of the loop, and it uses one rule for the early steps and another rule for the rest of training. The two rules are called the warmup phase and the cosine annealing phase.

What do you think happens?

The warmup branch raises the learning rate from a starting value toward a peak value over a fixed number of steps. From one step to the next, does the learning rate go up by the same amount each time, or by a different amount each time?

  • The same amount on every step

  • A larger amount on every step

  • A smaller amount on every step

  • A random amount on each step

Reveal answer
Answer:

The same amount on every step

The word linear describes the shape of the change: the learning rate rises by the same fixed increment on every warmup step, so plotted against steps it forms a straight line. That increment is computed once, before the training loop begins.

The rest of this article follows the training function step by step. You will compute the fixed warmup increment, trace the learning rate through the warmup ramp, see the exact condition that hands control to cosine annealing, and finish by learning how to find the line responsible for a learning rate value that looks wrong.

Recalculating the Rate on Every Step

The training function takes initial_lr and min_lr as arguments, and both have default values in the function signature. It does not receive a peak learning rate as an argument. Instead, it reads the peak from the optimizer, using the learning rate stored in the optimizer's first parameter group. The function therefore assumes that the optimizer's starting learning rate is the peak learning rate.

The function also calculates total_training_steps as the length of the training loader multiplied by the number of epochs. The warmup increment is computed once, before the epoch loop starts. The learning rate itself is recalculated inside the innermost loop, once per batch, because the step counter changes once per batch. This means the learning rate is a value that is set again and again as training proceeds, not a single number chosen at the start.

The Warmup Increment

During warmup, the learning rate rises by a fixed amount on every step. That amount, called lr_increment in the training function, is computed from three values. In the notation below, P stands for peak_lr, I stands for initial_lr, and W stands for warmup_steps.

Formula

Here Delta is the per-step increment. The total rise from the initial value to the peak is P minus I, and dividing that total rise by the number of warmup steps spreads it evenly across those steps. Because the increment is the same on every step, the total rise is exactly the increment multiplied by the number of warmup steps.

Once the increment is known, the learning rate on any warmup step follows directly from the step counter. The source writes the rule as lr equal to initial_lr plus global_step times lr_increment, where global_step is the step counter.

Formula, step by step
  1. 1
    lr_t is the learning rate at warmup step t, where t is the value of global_step.
  2. 2
    I is the starting value, initial_lr, which is the learning rate at step 0.
  3. 3
    t times Delta is the number of steps taken so far multiplied by the fixed increment.
  4. 4
    Adding the two terms gives the learning rate for that step, rising in a straight line as t grows.
python
Output (expected)
0 0.01
1 0.02
2 0.03
3 0.04
Notes on the code
  • The increment is derived from the three values: (0.05 - 0.01) / 4 = 0.01.
  • The increment is assigned once, outside the loop, mirroring the source's description.
  • Each pass computes initial_lr plus global_step times lr_increment, the warmup rule from the source.
  • This loop is a simplified illustration of the warmup branch and omits the rest of the training function.

Tracing the Warmup Ramp

Worked example

Tracing a four-step warmup

initial_lr is 0.01, peak_lr is 0.05, and warmup_steps is 4. Find the increment and the learning rate at each global_step, and identify the step where the cosine phase begins.

  1. 1
    Compute the increment

    The total rise is 0.05 minus 0.01, which equals 0.04. Dividing by warmup_steps of 4 gives an increment of 0.01 per step.

  2. 2
    Evaluate steps 0 through 3

    Each of these steps is less than warmup_steps, so the linear rule applies. At global_step 0 the rate is 0.01 plus 0 times 0.01, which is 0.01. At step 1 it is 0.02, at step 2 it is 0.03, and at step 3 it is 0.04.

  3. 3
    Check step 4

    At global_step 4 the condition global_step less than warmup_steps is false, because 4 is not less than 4. The training function leaves the warmup branch here and uses the cosine branch.

  4. 4
    Confirm the handoff value

    The source states that the cosine schedule begins from the peak learning rate, so the rate at step 4 starts at 0.05. The linear rule would also give 0.01 plus 4 times 0.01, which is 0.05, so the two phases meet at the peak.

Result

The learning rate takes the values 0.01, 0.02, 0.03, and 0.04 on the four warmup steps, then the cosine phase begins at 0.05. The rise is a straight line with a constant step size of 0.01.

The diagram below lays out the same ramp as a timeline. Each box is one value of global_step, and the label on each arrow is the fixed increment that carries the rate forward. The last box is where the cosine phase takes over.

global_step 0lr 0.01global_step 1lr 0.02global_step 2lr 0.03global_step 3lr 0.04global_step 4lr 0.05 peak+0.01+0.01+0.01+0.01
How does the learning rate change step by step during warmup, and where does it reach peak_lr?

Switching to Cosine Annealing

Warmup and cosine annealing are connected through the same step counter and the same peak_lr. Warmup climbs toward the peak, and cosine annealing starts from that peak and decreases toward the minimum. The handoff is decided on every batch by a single comparison, which the training function performs before it chooses a formula.

Linear warmupglobal_step < warmup_stepsCosine annealingglobal_step not less thanwarmup_stepsglobal_step reaches warmup_steps
At which global_step does the training loop stop using the warmup branch and start using the cosine branch, and what changes at that boundary?

Tracing a Wrong Learning Rate to Its Line

When a learning rate looks wrong during training, the fastest path to the cause is to ask three questions in order. Each question points to one part of the training function, so the divergence can be located rather than guessed. The flowchart below shows where control passes on each batch before the learning rate is set.

global_stepstep counter for this batchglobal_step <warmup_stepswhich branch applieslr = initial_lr +global_step *lr_incrementwarmup branchlearning rate forthis stepprogress from stepsafter warmupcosine branchcomparetruefalsewarmup valuecosine value
Which line in the loop sets the learning rate on each step, and how does control pass between the warmup and cosine branches before that line runs?
  • First, which branch applied at that step? Was global_step below warmup_steps?
  • Second, for a warmup step, what were the values of initial_lr, peak_lr, and warmup_steps that fed into lr_increment?
  • Third, for a cosine step, what were total_training_steps and warmup_steps that defined progress?
  • Common mistake

    Assuming the learning rate stays fixed for the whole run

    Why it fails

    The training function recalculates the learning rate on every step of the loop, so a single value read at the start does not describe the rate later in training.

    Do this instead

    Identify the phase at each step first, then trace the value that phase's formula produces.

  • Common mistake

    Recomputing the warmup increment inside the batch loop

    Why it fails

    The increment is computed once, before the epoch loop starts. Only the learning rate itself is recalculated per batch.

    Do this instead

    Keep the increment assignment outside the loop and reuse it on every warmup step.

  • Common mistake

    Expecting the peak learning rate to come from an argument

    What it looks like

    A user passes a peak value to the function but the optimizer was created with a different learning rate.

    Why it fails

    The function does not take a peak argument. It reads the peak from the optimizer's first parameter group.

    Do this instead

    Check the optimizer's starting learning rate, since that value is what the schedule treats as the peak.

  • Common mistake

    Treating global_step equal to warmup_steps as still in warmup

    Why it fails

    The warmup branch requires global_step to be strictly less than warmup_steps, so the boundary step already uses the cosine rule.

    Do this instead

    Use the strict less-than comparison when reasoning about which step uses which phase.

Practice Tracing the Schedule

EASY

Suppose initial_lr is 0.02, peak_lr is 0.10, and warmup_steps is 4. Compute lr_increment, then give the learning rate at global_step 0, 2, and 3. State which phase applies at global_step 4.

Hints
  • The total rise is peak_lr minus initial_lr, and the increment divides that rise by warmup_steps.
  • Use lr equal to initial_lr plus global_step times lr_increment for every step below warmup_steps.
  • The comparison is global_step less than warmup_steps, so check whether 4 is strictly less than 4.
MEDIUM

A training run shows a learning rate of 0.10 at global_step 4 and a learning rate of 0.02 at global_step 0, but the rate at global_step 2 is 0.04 instead of 0.06. List the lines you would check, in order, to locate the cause.

Hints
  • Start by confirming which branch applied at global_step 2.
  • Then check the values that fed into lr_increment, including the peak read from the optimizer.

Key Takeaways

  • Linear warmup raises the learning rate from initial_lr toward peak_lr over the first warmup_steps batches, by a fixed increment of (peak_lr - initial_lr) divided by warmup_steps.
  • During warmup the learning rate on each step is initial_lr plus global_step times lr_increment, and the increment is computed once before the epoch loop begins.
  • The training function switches to cosine annealing once global_step is no longer less than warmup_steps, and that cosine phase starts from the peak and measures progress against total_training_steps minus warmup_steps.
  • When a learning rate looks wrong, identify the branch at that step, then trace the warmup or cosine inputs back to the line that sets the rate on each batch.

More concepts starting with C