Your LR Schedule May Not Be Principled: Effective Learning Rate Scheduling in Weight-Normalized Training

Your LR Schedule May Not Be Principled. Effective Learning Rate Scheduling in Weight-Normalized Training. By Qifan Yu and three others, from ByteDance Seed and Peking University. Published on the twenty-ninth of October.

Weight-normalized training is emerging as an alternative to conventional weight decay, and recent studies report substantial pretraining loss gains from it. The claim here is about where those gains come from. Not, primarily, from normalization itself, but from the effective learning rate trajectory that normalization induces. If that is right, a schedule should be designed around the effective learning rate, not around the nominal learning rate alone.

The first line of evidence is a matching test. When the effective learning rate trajectories of the two regimes are matched, their training loss dynamics largely collapse onto each other. Building on that principle, the authors propose AIR, Affine Inverse-square-Root decay. It combines a theoretically derived decay with an affine adjustment, so that the schedule hits a prescribed peak and a prescribed terminal learning rate over a finite training horizon. They also identify late-stage scheduling of the effective learning rate as a key source of its gains, because a shared late schedule attenuates the effect of earlier learning rate choices.

In Mixture-of-Experts pretraining on one hundred twenty-eight billion tokens, AIR outperforms cosine and other candidate schedules in training loss, downstream performance and hyperparameter robustness. Terminal training loss drops by about zero point zero one relative to cosine, at matched cumulative learning rate. And applying the same trajectory to weight decay reproduces the corresponding training dynamics.

That last result sets the reach of the claim. The benefits are not confined to fixed-norm training, since weight decay follows the same trajectory to the same dynamics. The experiments reported are Mixture-of-Experts pretraining at that one token scale, and the loss improvement over cosine is small in absolute terms, about one hundredth. The matching result is stated as a large collapse of the dynamics, not an exact one.

The upshot is a change of target. The quantity a schedule should be shaping is the effective learning rate, and once it is, the apparent advantage of one regime over the other mostly reduces to which trajectory each one happened to follow.