PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Qifan Yu, Xinyu Ma, Xingyan Bin, Di He

arXiv:2610.40362v1cs.LG

TL;DR

WN and WD can have markedly different loss dynamics under the same nominal LR schedule, motivating a better ELR scheduling principle. The paper derives AIR from directional-alignment analysis and reports lower loss than cosine alongside downstream and robustness gains. Its results identify ELR as the scheduling object and show the design extends beyond fixed-norm training.

  • Problem

    WN and WD exhibit markedly different training dynamics under the same nominal LR schedule, leaving the better ELR schedule as an open question.

    Source: §1 · pp. 1–2

  • Method

    The paper analyzes directional alignment under WN and adapts the derived inverse-square-root decay affinely to prescribed peak and terminal learning rates.

    Source: §1 · p. 2, §7 · p. 10

  • Results

    About 0.01 lower terminal training loss than cosine at matched cumulative LR, with AIR also improving downstream performance and hyperparameter robustness in MoE pretraining.

    Source: §1 · p. 3, §7 · p. 10

  • Takeaways & Limitations

    Transferring AIR’s ELR trajectory to WD reproduces WN training loss dynamics, showing the design extends beyond fixed-norm training.

    Source: §7 · p. 10

    The authors identify extension to more architectures, optimizers, and larger scales, and study of continued training, supervised fine-tuning, and reinforcement learning, as future work.

    Source: §7 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

Weight-normalized training can outperform weight decay under cosine decay yet underperform it under WSD, motivating an explanation of their differing loss dynamics. The paper identifies effective learning rate (ELR) as the key driver and proposes AIR, an affine inverse-square-root schedule for finite-horizon training.

1 Introduction

  • WN initially trails WD but eventually outperforms it under cosine, whereas it remains behind under WSD, despite using the same nominal LR schedule within each comparison.The paper attributes the contrast to differing ELR trajectories as parameter norms evolve.Source: §1 · pp. 1–2
  • Matching ELR trajectories largely collapses WN and WD training loss dynamics, identifying ELR rather than nominal LR alone as the key scheduling target.Matching uses norm replay or blockwise LR rescaling.Source: §1 · pp. 2–3
  • AIR adapts a theoretically derived inverse-square-root decay to prescribed peak and terminal LRs, while shared late-stage scheduling attenuates differences from earlier LR choices.The paper identifies late-stage ELR scheduling as a key source of AIR’s gains.Source: §1 · pp. 2–3
  • Approximately 0.01 lower terminal training loss than cosine at matched cumulative LR, with AIR also achieving the best downstream performance and hyperparameter robustness among the evaluated schedule families.The evaluation used 128 B-token MoE pretraining.Source: §1 · p. 3

2 Effective Learning Rate as the Scheduling Principle

The paper treats ELR as a more informative measure of an optimizer update’s effect than nominal LR alone. Experiments show that matching ELRs largely aligns WN and WD loss dynamics, while WN provides a clean setting for ELR schedule design.

2.1 Effective learning rate under WN and WD

  • ELR measures the optimizer update relative to a parameter block’s norm, capturing its effect on parameter direction more informatively than nominal LR alone.Under WN, the fixed norm makes ELR proportional to nominal LR; under WD, changing nominal LR also affects subsequent norms.Source: §2.1 · p. 3, §2.3 · p. 4

2.2 Aligning the training dynamics of WN and WD

  • ELR matching largely collapses WD and WN training loss dynamics; under cosine, norm replay leaves a gap below 2 × 10−4, while fixed-norm LR matching leaves several ×10−3.The comparison also finds close alignment under AIR and in reverse matching along recorded WD norms.Source: §2.2 · p. 4

2.3 WN as a setting for ELR schedule design

  • WN fixes each block’s norm, making ELR directly proportional to nominal LR and providing a clean setting for studying ELR scheduling.WD couples nominal LR changes with subsequent parameter norms.Source: §2.3 · p. 4

3 Designing a Better ELR Schedule under WN: AIR

The paper analyzes how learning-rate schedules shape expected directional alignment under WN and derives an unconstrained inverse-square-root optimum. It turns that shape into AIR, an affine schedule meeting prescribed peak and terminal learning rates.

3.1 Directional-alignment dynamics under WN

  • Under WN, the analysis models expected cosine alignment to a fixed idealized target direction, assuming a prespecified LR schedule and fixed-magnitude tangential updates.A larger LR increases new alignment linearly but strengthens retraction of existing alignment quadratically.Source: §3.1 · p. 4

3.2 Unconstrained optimal schedule: inverse-square-root decay

  • The time-shifted 1/√t schedule maximizes expected alignment at every time under the analyzed dynamics.The result assumes the unconstrained setting without a prescribed peak, terminal LR, or LR budget.Source: §3.2 · p. 5

3.3 AIR: A practical schedule design

  • AIR applies an affine transformation to the inverse-square-root curve to meet prescribed peak and terminal LRs over a finite training horizon.The resulting schedule decreases monotonically from peak LR p to terminal LR ηmin; when an LR budget is prescribed, p is calibrated from the discrete LR sum, including warmup.Source: §3.3 · p. 5

4 Empirical Evaluation

AIR is evaluated under matched schedule controls and achieves lower final loss than cosine. It also has the lowest LR sensitivity over the widest peak-LR range, indicating less dependence on narrowly tuned peaks.

4.2 Pretraining loss and downstream performance

  • AIR achieves lower final training loss than cosine with the same terminal and cumulative LRs.Figure 3 compares learning-rate schedules and training loss under WN.Source: Figure 3 · p. 6

4.1 Experimental setup

  • Schedules share linear warmup duration, terminal LR, and cumulative LR, with each peak calibrated to match the cumulative LR including warmup.The comparison uses each schedule at the end of training; the cosine baseline’s peak LR and warmup ratio are selected by hyperparameter search.Source: §4.1 · p. 6

4.3 Robustness to hyperparameter choices

  • AIR has the lowest LR sensitivity over the widest peak-LR range, so its low-loss region does not depend on a narrowly tuned peak.Figure 4 plots terminal training loss against peak LR; circles mark measured configurations, stars mark family bests, and dashed curves are family-wise convex fits.Source: §4.3 · p. 6, Figure 4 · p. 6

5 Late-Stage Scheduling Matters More

AIR has the lowest terminal losses and best downstream average among the compared decay shapes, while experiments and analysis identify late-stage scheduling as a major influence on final loss. A shared late schedule narrows gaps across early schedules, and effective late decay balances sustained annealing with gradual LR changes.

5 Late-Stage Scheduling Matters More

  • AIR achieves the lowest terminal losses and best downstream average among the matched-budget decay shapes at both model scales.Table 1 reports terminal losses and downstream performance for each shape’s best parameter selected during evaluation.Source: Table 1 · p. 7

5.1 Late-stage scheduling drives the final loss gap

  • Swapping the late stage closes about 90% of the AIR–Cycloid final-loss gap, while swapping early or middle stages leaves loss nearly unchanged.The schedules share peak, terminal, and cumulative LRs; their original final training losses differ by approximately 0.005.Source: §5.1 · pp. 7–8

5.2 Shared late scheduling reduces the final loss gap

  • A shared AIR late schedule lowers terminal training and evaluation loss for all six schedules, removing 42%–83% and 45%–94% of their respective gaps to AIR.Table 2 compares original schedules with versions retaining their first 70% and sharing the same late-stage schedule.Source: §5.2 · p. 8, Table 2 · p. 8

5.3 A trade-off in late-stage scheduling

  • Both abrupt decay at γ = 0.5 and early flattening at γ ∈ {2, 3} increase final loss; AIR and γ ∈ {0.75, 1} achieve the lowest losses.With junction and terminal LRs fixed, the results indicate a trade-off between sustained annealing and avoiding overly abrupt LR changes.Source: §5.3 · p. 8, Figure 6 · p. 9

5.4 A theoretical explanation of schedule dependence

  • A shared late stage attenuates the entering alignment difference without reversing its sign; higher LRs over a fixed late interval reduce its retention factor.The alignment bounds qualitatively explain the observed loss gaps, with one term retaining early-stage differences and another bounding late-stage contributions; the analysis ignores stochastic fluctuations around the mean.Source: §5.4 · p. 9

6 Related Work

Prior work studies parameter-scale control and learning-rate schedules, but scheduling theories focused on nominal LR do not account for evolving parameter norms. This work derives inverse-square-root decay by optimizing directional alignment under WN and adapts it to finite horizons as AIR.

6 Related Work

  • WN and related constraints control parameter scale; approaches vary in whether they separate magnitude from direction or fix weight and update norms.Manifold-constrained methods optimize tangent-space descent directions before retracting weights, while matrix-norm choices shape constraint geometry.Source: §6 · p. 9
  • Existing scheduling theories explain WSD, cooldown, and power decay but focus on nominal LR, which is not principled when parameter norms evolve.Concurrent work reports that matching ELR largely aligns WN and WD loss dynamics, but does not derive a decay law from an optimization objective.Source: §6 · pp. 9–10
  • The paper addresses this gap by optimizing directional alignment under WN to derive 1/√t decay, then applying an affine transformation for finite horizons to form AIR.Source: §6 · p. 10

7 Conclusion and Future Work

The conclusion identifies matched effective learning-rate (ELR) trajectories as the key to aligning WN and WD training dynamics, motivating ELR-targeted schedule design and AIR. The accompanying appendices develop the geometry, matching, and schedule-design analysis behind that approach.

7 Conclusion and Future Work

  • AIR reduces terminal loss by about 0.01 versus cosine at matched cumulative LR, improves downstream performance, and is robust across schedule sweeps.Sharing the late schedule substantially narrows final loss gaps between different early schedules; transferring AIR’s ELR trajectory to WD reproduces WN training loss dynamics.Source: §7 · p. 10

A Directional Geometry and ELR Matching

A.1 Exact normalized updates

  • ELR scales the optimizer update relative to the current parameter norm, while weight decay also contributes radial shrinkage.The exact normalized-update identities hold when the vector being normalized is nonzero.Source: §A.1 · p. 14

A.2 Radial and tangential components

  • The radial term is removed in the first-order directional change; ELR sets its scale for a given tangential update.Under WN, the expression uses λ = 0 and fixed norm R.Source: §A.2 · p. 14

A.3 Blockwise ELR transfer

  • Blockwise ELR transfer reverses coefficient matching, applying the rule separately to controlled blocks using recorded WD norms.A WN schedule supplies the target ELR for this matching.Source: §A.3 · p. 14

B WN Schedule Design and Calibration

  • The schedule-design appendix derives alignment dynamics and their unconstrained optimum, then establishes affine endpoint construction and discrete LR-sum calibration.The alignment derivation conditions on a fixed initial history and prescribed LR schedule; its continuous approximation assumes small relative updates, slow LR and alignment variation, and approximately constant update statistics.Source: B WN Schedule Design and Calibration · p. 15, §B.1 · p. 15

B.1 From the WN update to alignment dynamics

  • The alignment analysis averages over updates and optimizer-update randomness without requiring independence between the current parameters and update.The continuous approximation treats one optimizer step as the time unit.Source: §B.1 · p. 15

B.2 Derivation of the integral solution

The integral solution weights contributions introduced over time by the fraction retained through the endpoint, and the resulting state remains strictly positive.

B.2 Derivation of the integral solution

  • Each contribution introduced at time t is weighted by the fraction retained from t to the endpoint; the resulting solution is strictly positive throughout the interval.The positivity follows from a positive initial state, a positive retention factor, and a nonnegative integral contribution.Source: §B.2 · p. 15

B.3 Proof of the unconstrained optimal schedule

The proof establishes an unconstrained optimal schedule and its discrete counterpart, then describes the geometric and calibration conditions needed for finite-horizon use. An affine transformation preserves the inverse-square-root curve’s relative decay behavior while meeting prescribed endpoints.

B.3 Proof of the unconstrained optimal schedule

  • The unconstrained optimum requires m(t)η(t) = a/(2b), yielding a linear increase in m(t)^2 and a globally optimal terminal alignment.The proof notes that these squared states are squares of expected alignments, not expectations of squared alignments.Source: §B.3 · p. 16

B.4 Geometric range of the constant-coefficient assumptions

  • The continuous optimum’s interpretation as WN alignment is restricted to intervals where the geometric assumptions remain realizable.A constant positive drift cannot describe indefinitely growing physical alignment, and the continuous approximation also requires a small relative squared step.Source: §B.4 · p. 17

B.5 Affine endpoint construction

  • An affine transformation of the inverse-square-root curve uniquely sets the decay LR’s prescribed initial and final values while preserving relative decay-rate magnitudes.For p > ηmin, the resulting schedule decreases monotonically between its endpoints.Source: §B.5 · p. 17

B.6 Matching the discrete LR sum

  • Discrete LR-sum calibration determines a unique peak from a prescribed feasible sum when the peak coefficient is positive.Calibration uses the LR values actually executed, including the terminal endpoint only when its index is included.Source: §B.6 · pp. 17–18

C Experimental Settings and Evaluation Protocol

The appendix provides experimental-setting details for Sections 4 and 5.

C Experimental Settings and Evaluation Protocol

C.1 Training settings and comparison controls

The appendix specifies how schedule comparisons are controlled and analyzes how shared or differing LR stages affect expected directional alignment. Under a shared late schedule, an entering alignment difference is attenuated without reversing sign, ignoring stochastic fluctuations.

C.1 Training settings and comparison controls

  • Fixed-norm comparisons vary the complete optimizer-step LR sequence; matched-sum controls fix its sum, horizon, warmup, and terminal LR while allowing peak LR to differ.A shared late schedule means equal LR values at every shared optimizer step, not merely equal endpoints or integrated LR.Source: §C.1 · pp. 18–19

C.2 Compared decay shapes

  • Compared decay shapes use a common affine endpoint transformation, with peak calibration matching cosine’s cumulative LR, including warmup, while holding warmup duration and terminal LR fixed.The raw shapes are listed in Table 4; each schedule substitutes its own shape into the common transformation.Source: §C.2 · p. 19, Table 4 · p. 19

C.3 Loss summaries

  • Training loss uses an EMA of the last 100 finite logged values, while evaluation loss uses the last finite logged value without smoothing.These are terminal summaries of individual trajectories, not averages across independent training seeds.Source: §C.3 · p. 20

C.5 Hyperparameter sweeps and LR sensitivity

  • LR sensitivity is the mean terminal loss above a family’s best measured configuration per e-fold of peak LR, and is descriptive only over the measured range.Search ranges and configuration counts differ across families, so this normalization does not equalize tuning budgets.Source: §C.5 · p. 20

C.6 Stage swaps

  • Stage swaps follow a donor schedule only within the named interval; matching the original schedules’ cumulative LR does not ensure a replacement preserves that sum.Each replacement uses the recipient schedule outside the selected interval.Source: Table 1 · p. 22

C.7 Late decay exponents

  • Late-exponent interventions switch at about 0.79 of the decay stage and their cumulative LR differs from AIR by less than 1%.They use the same earlier schedule, junction LR, and configured terminal LR; cumulative LR is not exactly matched.Source: §C.7 · p. 22

D.2 Varying early-stage schedules

  • With a shared late schedule, the expected-alignment difference keeps its sign and cannot grow in magnitude, ignoring stochastic noise.This also holds for the exact discrete relation.Source: §D.2 · pp. 22–23

D.3 Varying late-stage schedules

  • When both stages vary, the exact ODE difference equation separates the propagated entering difference from differences caused by later LR choices.The schedule comparisons characterize how LR schedules promote expected directional alignment under common coefficients.Source: §D.3 · pp. 23–24

D.4 Slope bounds and endpoint information

The appendix bounds how LR-schedule differences affect expected alignment, under stated endpoint and regularity conditions. It also derives a sensitivity kernel identifying when and where an LR change most influences terminal alignment.

D.4 Slope bounds and endpoint information

  • The slope bound requires a common terminal LR; a jump at the switch means a shared early schedule alone does not impose that endpoint condition.For comparisons with exponent γ < 1, the terminal derivative is unbounded, so the finite-slope condition may fail and the general LR-difference bound applies instead.Source: §D.4 · p. 24

D.5 Affine power curves and the scope of the bounds

  • For the affine inverse-power family, β = 1/2 gives an inverse-square-root shape, and nearby exponents yield nearby terminal expected alignments from a common entering state.When the early stage also varies, the analysis incorporates the resulting entering-state difference.Source: §D.5 · p. 25

D.6 Infinitesimal sensitivity to the time of an LR change

  • The sensitivity kernel identifies when and where an LR perturbation most strongly influences terminal alignment.Later perturbations are retained more strongly through H(K, k), while a − 2bη(k)m(k) determines local alignment gain.Source: §D.6 · p. 25

Related papers

Loading 2610.40362v1…