A learning-rate schedule built around update size
Your LR Schedule May Not Be Principled: Effective Learning Rate Scheduling in Weight-Normalized Training · arXiv:2610.40362
In plain words
The paper asks why two ways of training a model can behave differently even when they use the same learning-rate schedule. One method keeps parameter sizes fixed; the other lets them change.
The authors argue that the update size relative to each parameter block matters more than the scheduled rate alone.
They propose AIR, a schedule shaped by an inverse-square-root curve. In tests on two mixture-of-experts model sizes, AIR had about 0.01 lower final training loss than cosine decay under matched schedule conditions. Those figures summarize individual runs, not averages across independent runs.
The research
Why can weight-normalized training finish ahead of weight decay under one learning-rate schedule, yet behind it under another? The paper investigates whether the difference lies in the update size relative to each parameter block, rather than in the nominal rate alone.
The authors introduce AIR, a finite-horizon learning-rate schedule based on effective update size, and test it against other schedules in mixture-of-experts pretraining.
The authors define effective learning rate (ELR) as the nominal rate divided by a parameter block’s norm. Because weight-normalized training fixes block norms and weight decay does not, the same nominal schedule can produce different relative update trajectories.
Matching those trajectories largely aligns the methods’ training-loss dynamics, supporting ELR as a scheduling target.
The authors derive AIR from an idealized alignment model, then compare schedules in mixture-of-experts pretraining with matched warmup, ending rate, cumulative rate, and horizon. AIR’s terminal training loss is about 0.01 below cosine’s at both tested scales; each value summarizes one run’s ending trajectory. AIR also has the best reported downstream average among the compared shapes.
The empirical results cover the reported mixture-of-experts pretraining comparisons, and the terminal loss summaries come from individual training trajectories rather than averages across independent runs.
The alignment theory uses idealized assumptions, and broader architectures, optimizers, scales, and training stages remain untested.
What the evidence shows
- A schedule-dependent contrast
The paper proposes AIR, a learning-rate schedule for finite-horizon pretraining. It begins with a contrast: weight-normalized training (WN) ends ahead of weight decay (WD) under cosine decay, but behind it under warmup-stable-decay.
Since both methods can use the same nominal schedule, the authors ask whether their different loss trajectories reflect the size of their updates relative to their parameter blocks.
- A different update measure
The authors define effective learning rate (ELR) as nominal learning rate divided by a parameter block’s norm. WN fixes that norm, while WD’s norm can change during training.
So matching nominal rates does not necessarily match relative update sizes. This gives the authors a concrete test: if ELR helps explain the contrast, making the two methods’ ELR trajectories more alike should also make their loss trajectories more alike.
- Testing the proposed explanation
In a cosine comparison, the authors match WN to WD by replaying WD’s recorded parameter-block norms, which makes the WN loss trajectory differ from WD by less than 2x 10−4 loss units. Adjusting WN’s blockwise learning rates instead leaves a residual gap of several x10−3.
Matching WN’s ELR in WD also brings its loss dynamics close to the WN baseline under cosine and AIR. These results support ELR as a useful scheduling target.
- Building and comparing AIR
To design a schedule for WN, the authors derive an inverse-square-root learning-rate decay from an idealized model of directional alignment. AIR, or Affine Inverse-square-root decay, adjusts that shape to meet chosen peak and ending rates over a finite training horizon.
In matched pretraining comparisons on two mixture-of-experts model sizes, AIR’s terminal training loss was about 0.01 lower than cosine’s, and AIR had the best reported downstream average among the compared schedule shapes.
- Examining the late stage
The authors also test where schedule differences matter. In one matched AIR–Cycloid pair, swapping the late stage closed about 90% of the pair’s final training-loss gap, while swapping early or middle stages had little effect.
Across six tested schedules, using a shared AIR late segment narrowed their training and evaluation loss gaps to AIR. These interventions point to late-stage choices as an important source of the tested differences.
Deep
- Original paper authors
- Qifan Yu, Xinyu Ma, Xingyan Bin, Di He
- Submitted to arXiv
- License
- CC BY 4.0