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

AI summary

Overview

Research area: Machine learning optimization — specifically learning rate (LR) scheduling and training dynamics for large-scale pretraining, in the context of weight-normalized (WN) versus weight-decay (WD) training.

Technical level: Advanced. The paper reasons about effective learning rate trajectories, theoretical decay forms, and billion-token Mixture-of-Experts pretraining.

Scope (one sentence): The paper argues that the effective learning rate, not the nominal one, is the right object to schedule, and introduces a schedule (AIR) built on that principle and validated in large-scale pretraining.

What This Paper Is About

Weight-normalized training has recently been reported to beat conventional weight decay on pretraining loss, but it has not been clear why it helps or what practitioners should actually tune. This paper argues that the benefit comes from the effective learning rate trajectory that weight normalization implicitly produces, not from normalization itself. The goal is to redirect schedule design toward that effective quantity and to construct a schedule that targets it directly.

Key Contributions

1. Diagnosis of the source of WN's gains. The authors show that the reported pretraining improvements from weight-normalized training are driven primarily by the induced effective learning rate (ELR) trajectory rather than by normalization alone, and that matching ELR trajectories largely collapses the training loss dynamics of WN and WD.
2. A reframing of schedule design. They argue that schedules should be designed against ELR rather than the nominal learning rate alone, making ELR the principled scheduling object.
3. The AIR schedule. They propose AIR (Affine Inverse-square-Root decay), which pairs a theoretically derived (1/\sqrt{t}) decay with an affine adjustment so the schedule satisfies prescribed peak and terminal learning rates over a finite training horizon.
4. An account of where the gains come from, plus large-scale validation. They identify late-stage ELR scheduling as a key source of AIR's gains — because a shared late schedule attenuates the effect of earlier LR choices — and test the method in 128 B-token Mixture-of-Experts pretraining, including a transfer check showing the same ELR trajectory reproduces the corresponding dynamics under WD.

Main Findings

• ELR, not normalization, explains the WN advantage. The pretraining loss gains attributed to weight-normalized training are shown to arise primarily from the effective learning rate trajectory it induces.
• Matching ELR collapses the two regimes. When ELR trajectories are aligned, the training loss dynamics of WN and WD become largely the same — evidence that the nominal LR is the wrong variable to compare across methods.
• AIR outperforms cosine and other candidates. On 128 B-token Mixture-of-Experts pretraining, AIR leads on training loss, downstream performance, and hyperparameter robustness, reducing terminal training loss by approximately 0.01 relative to cosine at matched cumulative LR.
• Late-stage scheduling carries much of the benefit. A shared late schedule attenuates the effect of earlier LR choices, which the authors identify as a key reason AIR's gains appear.
• The effect transfers to weight decay. Applying the same ELR trajectory under WD reproduces the corresponding training dynamics, indicating the benefits are not specific to fixed-norm training.

Methodology in Plain English

The argument proceeds in two stages. First, the authors treat the effective learning rate — roughly, how large a step the optimizer actually takes relative to the size of the weights, rather than the number typed into the config — as the thing that governs training. They compare runs using weight normalization against runs using weight decay, align their effective learning rate trajectories, and check whether the loss curves converge on the same behavior. They largely do, which supports the claim that normalization mainly acts as an implicit LR schedule.

Second, they design a schedule directly for that effective quantity. The decay follows a (1/\sqrt{t}) form that the authors derive theoretically, and an affine (straight-line) adjustment is layered on top so the curve passes through a chosen peak LR at the start and a chosen terminal LR at the end of a fixed training budget. This is then evaluated in a large Mixture-of-Experts pretraining run and compared against cosine and other candidate schedules on training loss, downstream tasks, and sensitivity to hyperparameters. As a further check, the same effective trajectory is imposed on weight decay training to see whether the dynamics carry over.

The abstract does not report the derivation of the (1/\sqrt{t}) form in detail, the full experimental protocol, the downstream task set, or the specific hyperparameter ranges tested; those details would be in the full paper.

Why This Matters

Impact on research: The paper shifts the unit of analysis in schedule design from the nominal learning rate to the effective one. If ELR matching really does collapse WN and WD dynamics, then results reported for one training regime may be reproducible in the other simply by transferring the trajectory — which reframes a family of empirical comparisons as a single underlying phenomenon and gives schedule design a cleaner theoretical target.

Real-world applications (implied by the framing):

• Large language model pretraining, where LR schedule choice is one of the few high-leverage decisions made before a run costing substantial compute.
• Compute-budgeted training where the horizon is fixed in advance, since AIR is explicitly built to hit prescribed peak and terminal LRs over a finite budget.
• Reducing hyperparameter sweep cost: greater robustness to hyperparameters means fewer trial runs before committing to a full pretraining job.
• Training setups that rely on normalized weights for stability, where the nominal LR does not correspond to the actual step size and naive schedule tuning is therefore misleading.

Industry relevance: Pretraining runs are expensive and largely irreversible once launched, so a schedule that is both principled and less sensitive to hyperparameter choice has direct value. The reported roughly 0.01 terminal loss reduction at matched cumulative LR is the kind of margin that matters at scale, and the result that the same trajectory transfers to weight decay makes the recipe portable to existing training stacks.

Future Directions

• Scope of ELR equivalence. Testing whether ELR matching collapses dynamics as cleanly beyond the settings examined — other optimizers, architectures, model sizes, and training horizons — is left open by the abstract.
• Schedules beyond a fixed horizon. AIR is formulated for a prescribed peak and terminal LR over a finite budget; how it should be extended to continued training, restarts, or budgets that change mid-run is an open question.
• Tighter theory for the (1/\sqrt{t}) form. The decay is described as theoretically derived, but the abstract does not establish how general that derivation is or what assumptions it rests on.
• ELR as a general transfer tool. Since the trajectory reproduces dynamics across WN and WD, a natural next step is using ELR as a diagnostic for comparing, porting, or debugging schedules across training methods more broadly.

Target Audience

Optimization and training-dynamics researchers; engineers responsible for large-scale pretraining runs and LR schedule selection; and anyone working with weight normalization or weight decay who wants to understand why one outperforms the other. Readers should be comfortable with learning rate schedules and the mechanics of gradient-based training, as the paper assumes that background.

Authors’ abstract

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. Schedule design should therefore target ELR rather than the nominal learning rate (LR) alone. Building on this principle, we propose AIR (Affine Inverse-square-Root decay), which combines a theoretically derived $1/\sqrt{t}$ decay with an affine adjustment to satisfy prescribed peak and terminal LRs over a finite training horizon. We further identify late-stage ELR scheduling as a key source of its gains, because a shared late schedule attenuates the effect of earlier LR choices. Experiments on 128 B-token Mixture-of-Experts pretraining demonstrate that AIR outperforms cosine and other candidate schedules in training loss, downstream performance, and hyperparameter robustness, reducing terminal training loss by approximately 0.01 relative to cosine at matched cumulative LR. Applying the same ELR trajectory to WD reproduces the corresponding training dynamics, showing that the benefits extend beyond fixed-norm training and establishing ELR as the principled scheduling object.