AIR schedule targets effective learning rate, not nominal rate
Original title: Your LR Schedule May Not Be Principled: Effective Learning Rate Scheduling in Weight-Normalized Training
On 128 billion-token MoE pretraining, AIR lowers terminal training loss by about 0.01 versus cosine at matched cumulative learning rate.
The paper argues that reported pretraining gains from weight-normalized training come mainly from the induced effective learning rate trajectory rather than normalization alone. Matching those trajectories largely aligns the loss dynamics of weight normalization and weight decay. AIR combines a 1/t decay with an affine adjustment so peak and terminal rates are met over a finite horizon, with late-stage scheduling identified as the main source of gains.
- Matching effective-learning-rate paths largely collapses weight-normalization and weight-decay loss dynamics.
- The same trajectory applied to weight decay reproduces the corresponding training dynamics.
- AIR beats cosine and other candidate schedules on loss, downstream performance, and hyperparameter robustness.
- A shared late schedule attenuates the effect of earlier learning-rate choices.
From the original article
1 min read at huggingface.co
Weight-normalized (WN) training is emerging as a compelling alternative to conventional weight decay (WD), with recent studies reporting substantial pretraining loss gains. We show that these gains are driven primarily by the induced effective learning rate (ELR) trajectory rather than by normalization alone. Matching ELR trajectories largely collapses the training loss dynamics of WN and WD.
By Qifan Yu, Xinyu Ma, Xingyan Bin · huggingface.co
A short excerpt; the article and its images belong to the publisher. Summary, key facts and assessment above are AI-written from the article and can be wrong.