AIR schedule targets effective learning rate in weight-normalized training
An arXiv preprint proposes AIR (Affine Inverse-square-Root decay), a learning-rate schedule that targets the effective learning rate rather than the nominal learning rate alone. AIR combines a theoretically derived inverse-square-root decay with an affine adjustment to meet prescribed peak and terminal learning rates. The authors report that, in 128B-token Mixture-of-Experts pretraining, AIR outperformed cosine and other candidate schedules in training loss, downstream performance, and hyperparameter robustness, with terminal training loss approximately 0.01 lower than cosine at matched cumulative learning rate.
The paper reports that matching effective-learning-rate trajectories largely aligns the training-loss dynamics of weight-normalized and weight-decay training, supporting the effective learning rate as a more principled scheduling target.
Computer Science > Machine Learning
arXiv:2610.40362 (cs)
[Submitted on 29 Oct 2026]
Title:Your LR Schedule May Not Be Principled: Effective Learning Rate Scheduling in Weight-Normalized Training
Authors:Qifan Yu, Xinyu Ma, Xingyan Bin, Di He
View a PDF of the paper titled Your LR Schedule May Not Be Principled: Effective Learning Rate Scheduling in Weight-Normalized Training, by Qifan Yu and 3 other authors