Learning Rate Adjustment in Training Loops

A 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. Warmup is the first of these phases and is the focus of this article.

As you read the training function, keep one question in mind: at a given step of training, which rule decides the learning rate, and why? The rest of this article follows that question from the warmup formula, to the point where the rule changes, to the line of code that sets the value the optimizer uses.

Linear warmup gradually increases the learning rate from initial_lr to peak_lr during the first warmup_steps batches of training. The word linear describes the shape of the change: the learning rate rises by the same amount on every step, so plotted against steps it forms a straight line. The source describes this increment as a fixed quantity computed once, before the training loop begins.

During warmup, the learning rate at each step is the initial value plus the number of steps taken so far, multiplied by the increment. The source writes this as lr = initial_lr + global_step * lr_increment, where global_step is the step counter.

In these formulas, I stands for initial_lr, P stands for peak_lr, W stands for warmup_steps, s stands for global_step, and Delta stands for lr_increment. The first formula is the increment, computed once. The second formula gives the learning rate on each warmup step.

The diagram uses the example values initial_lr = 0.0001, peak_lr = 0.001, and warmup_steps = 5. Each step adds the same amount, so the points lie on a straight line from the starting value to the peak. Each step label shows the learning rate at that step.

The training function takes initial_lr and min_lr as arguments, with 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 source notes that this assumes the optimizer's starting learning rate is used as 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_steps value is the number that controls how long the ramp lasts. The code sample above is an illustration written for this article to show where each value comes from; the names match the source, but the default values and the warmup_steps line are example choices.

The learning rate itself is recalculated inside the innermost loop, once per batch, because the step counter changes once per batch. On each pass, the code first checks the condition global_step < warmup_steps. When that condition is true, the code uses the linear formula. When it is false, the code uses the cosine formula.

The flowchart shows the decision that runs on every batch. The check is made against the current value of global_step, so the chosen branch can change from one batch to the next without any other change to the code.

Once global_step reaches warmup_steps, the training function stops using the linear rule and starts using cosine annealing. Cosine annealing is a technique used to adjust the learning rate after the warmup phase. The two phases share 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 variable progress measures how far training has moved through the post-warmup period. The source defines it as the number of steps completed after warmup, divided by the number of steps remaining after warmup. The number of steps remaining is the total number of training steps minus warmup_steps. Because progress is measured from the end of warmup, the cosine schedule begins from the peak learning rate and moves toward min_lr as training continues.

Most unexpected learning rate values come from a small set of misreadings. The items below are the ones most likely to send you looking in the wrong place.

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 line of the training function, so the divergence can be located rather than guessed.