ByteDance Seed proposes AIR learning rate schedule: final loss 0.01 lower than with cosine

Researchers at ByteDance Seed (Qifan Yu, Xinyu Ma, Xingyan Bin, Di He) have proposed AIR, a learning rate schedule. In pretraining a Mixture-of-Experts (MoE) model on 128 billion tokens, it cut the final training loss by about 0.01 compared with a cosine decay (cosine). The comparison used the same total sum of LR over the whole run. Here, the learning rate (LR) is the step size an optimizer uses to update a model's weights. It is usually reduced during training, most often along a cosine curve.

Why weight normalization came up at all

To control the size of the weights, practitioners usually use weight decay (WD), which pulls the weights slightly toward zero at every step. Weight normalization (WN), which trains with normalized weights, has recently been considered as an alternative. Several papers have reported a noticeable loss improvement in pretraining.

The authors checked where that improvement comes from and arrived at a different explanation. Under WN, the nominal LR translates into a different real update size than under WD, which produces its own effective learning rate (ELR) trajectory. The researchers adjusted the ELR of both methods so that the trajectories matched. The WN and WD loss curves then nearly merged. So the gain came from the shape of the ELR, not from the normalization itself.

The practical takeaway is that a schedule should be designed for the ELR, not only for the nominal LR.

How AIR works

The name stands for Affine Inverse-square-Root decay. The schedule has two parts.

• Base. The step falls in proportion to 1/√t, where t is the training step number. The authors derived this form theoretically. • Correction. A linear (affine) transformation. Training lasts a finite time, and a pure 1/√t decay does not give the right values at the start and the end. The correction lets users set the peak and final LR in advance.

As with cosine, the user specifies two endpoints: the peak and the final value.

The authors also make an observation about the late phase. If runs share the same tail of the ELR schedule, that tail weakens the influence of earlier LR choices. The authors consider late-stage ELR scheduling one of the main sources of the gain.

What the experiment showed

The test was a pretraining run of an MoE model on 128 billion tokens. In an MoE, only a subset of expert blocks processes each token. AIR was compared with cosine and other candidate schedules. The results:

• AIR has lower training loss. Its final loss is about 0.01 lower than cosine's, with the same total sum of LR over the whole run. • AIR is also better on downstream task quality. • The schedule is more robust to hyperparameter choice.

In a separate experiment, the same ELR trajectory was applied to ordinary weight decay, and the training dynamics were reproduced. The authors conclude that the benefit is not limited to training with a fixed weight norm, and they name the ELR as the object worth scheduling.

What this means in practice

A 0.01 difference in loss is small, but it came from changing just one step-size schedule. The work offers three practical pointers:

• check how your weight scheme and optimizer convert the nominal LR into an effective one; • try the AIR schedule, where you set only the peak and final LR; • pay more attention to the end of training, because a shared schedule tail weakens the influence of earlier settings.

The abstract describes experiments only on MoE pretraining with 128 billion tokens. It does not report results for other model sizes or data volumes.

The paper is published on arXiv under number 2610.40362 and was featured in the Hugging Face Daily Papers selection.