PaperlayerView exact revision on arXiv

Source-linked AI summary

Measuring the local shape of learning-rate transfer curves

Satoki Ishikawa, Rio Yokota, Ryo Karakida

arXiv:2610.21814v1cs.LG

TL;DR

The local shape of learning-rate transfer curves is hard to characterize because trained performance depends on learning dynamics, architecture, and implicit biases, while hypergradient estimation is challenging. The paper proposes complementary theoretical and approximate expressions for the LR gradient, which support bounds, normalization-layer analysis, and layer-wise sensitivity profiles.

  • Problem

    LR transfer-curve shape varies with model size, architecture, and other hyperparameters, while its theoretical characterization and accurate hypergradient evaluation remain difficult.

    Source: §1 · pp. 1–2

  • Method

    The paper derives an EoS-based LR-gradient formula and an empirically tested fixed-point approximation that treats LR sensitivity as a hypergradient.

    Source: §1 · p. 2

  • Results

    The complementary expressions explain normalization-layer effects and bound the LR gradient; FP-LR profiles identify strong sensitivities in ResNet’s final layer and GPT’s embedding and output layers.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The LR gradient offers a framework for analyzing curve shape, including layer-wise contributions and the effects of normalization and weight decay.

    Source: §1 · p. 2

    The characterization is restricted to evaluating first derivatives, and the EoS formula assumes twice-differentiable loss and differentiable inverse sharpness.

    Source: §3.1 · p. 4, §6 · p. 10

Abstract

from arXiv · show

Understanding the learning-rate (LR) dependence of trained model performance is essential for better optimization of neural networks, but the appropriate size of the LR depends on learning dynamics, architecture, and their hidden biases, and thus our understanding of this complex dependence remains limited. The behavior of the loss curve as a function of the LR has recently been observed in extensive experiments on LR transfer across different settings, such as model sizes, and it is becoming more important to explore methodologies for quantitatively characterizing the shape of this LR curve. In this work, we provide two relational equations that give a good empirical explanation of the gradient of the LR transfer curve, i.e., the LR gradient. One is a theoretical formula that holds at the edge of stability, and the other is an approximate expression obtained by naively treating this gradient as a hyper-parameter gradient. Using these two relational equations in a complementary manner, we show a lower bound on the gradient of the curve and a zero gradient under specific conditions for normalization layers in theory. The approximate expression also applies to a wide range of experimental settings and allows us to decompose the LR gradient into layer-wise components and quantify the strong contribution of specific layers to the curve shape. Thus, this serves as a foundation for better understanding the internal mechanisms of LR dependence.

1 INTRODUCTION

Learning-rate transfer curves vary with model size, architecture, and other hyperparameters, yet their local shape remains difficult to characterize theoretically. This work analyzes their slope and proposes two complementary expressions for the LR gradient.

Motivation

  • Characterizing the curve is difficult because performance depends on learning dynamics, architecture, and implicit biases, while large learning rates produce nonlinear training behavior.Large learning rates commonly extend beyond the classical descent-lemma regime, where edge-of-stability behavior appears.Source: §1 · pp. 1–2
  • Direct hypergradient computation is challenging and can cause gradient explosion or short-horizon bias in deep architectures.The curve’s geometry as a function of learning rate naturally calls for a hypergradient.Source: §1 · p. 2

Contributions

  • Two relational expressions explain the LR gradient: an EoS-based theoretical formula and a fixed-point approximation that works empirically under slow hypergradient dynamics.The fixed-point expression is intended for a broader range of settings, while the EoS formula assumes edge-of-stability conditions.Source: §1 · p. 2
  • Both expressions address normalization-layer effects and bound the LR gradient; EoS-LR further characterizes a flat curve under a specific weight-decay and normalization combination.FP-LR also supports layer-wise decomposition under a block-diagonal Hessian approximation.Source: §1 · p. 2
  • FP-LR profiles identify strong LR sensitivities in ResNet’s final layer and GPT’s embedding and output layers.In ResNet under µP, changing the output multiplier changes accuracy and LR transfer across widths, whereas changing the input multiplier has little effect.Source: §1 · p. 2

Contribution

  • The paper presents its analysis as a first step toward understanding the universal shape of LR curves for tuning and transfer.Source: §1 · p. 2

2 PRELIMINARIES

Prior work studies LR transfer and the dynamics that shape its curves, but theoretical characterization remains limited. This section frames the curve as final loss versus learning rate and motivates gradient-based descriptions beyond finite differences.

2.1 Related Work

  • LR transfer research grew from µP studies of width scaling, while later evidence shows transfer also depends on weight-feature alignment, weight decay, loss, and normalization.µP’s appropriate LR scale early in training does not necessarily hold as training progresses.Source: §2.1 · p. 2
  • EoS relates learning rate to loss curvature, and prior evidence found width-independent maximum Hessian eigenvalues under EoS, supporting better LR transfer.The paper presents its theoretical evaluation of how EoS governs the LR gradient as a first result of its kind.Source: §2.1 · p. 3
  • Standard hypergradients can suffer gradient explosion and short-horizon bias; here they measure LR sensitivity of trained models rather than optimize hyperparameters.The authors report iterative hypergradient instability, but a heuristic fixed-point expression approximates the LR gradient accurately in their experiments.Source: §2.1 · p. 3

2.2 The shape of LR transfer curves

  • A transfer curve plots final loss against LR; its finite-difference local shape depends on grid resolution and requires experiments across diverse settings.The trained parameters are denoted θ*, LR by η, and grid resolution by Δ; the paper also notes its methods can use validation loss when validation samples are available.Source: §2.2 · pp. 3–4
  • A derivative-based curve description requires the hypergradient ∂θ*/∂η, whose evaluation in modern neural networks is nontrivial.Section 3 evaluates it implicitly, while Section 4 gives an empirically validated approximation for explicit evaluation.Source: §2.2 · p. 4

3 EOS-CONSTRAINED LR GRADIENT

The section derives an LR-gradient relation by treating the edge-of-stability condition η = F(θ) as a constraint on loss minimization. It shows how the choice of F affects whether that relation captures LR sensitivity, including at small batch sizes.

3.1 Basic formula

  • The derivation uses η = F(θ) as an EoS constraint and applies Lagrange multipliers to the constrained loss minimum.For GD, F is the inverse of sharpness; the formulation can incorporate algorithm-dependent expressions for F.Source: §3 · p. 4, §3.1 · p. 4
  • When EoS holds, substituting trained parameters into the derived relation accurately predicts finite-difference LR gradients.The expression's denominator involves ∇λmax, a third-derivative quantity whose stable numerical computation is generally unexplored; the authors estimate the directional derivative by finite differences.Source: §3.1 · p. 5
  • At ideal minima, ∇L⊤∇λmax is positive when the LR curve decreases with η, and its sign need not match that along training trajectories.The authors report a positive inner product at averaged trained parameters.Source: §3.1 · p. 5
  • At small batch sizes, Fbatch matches finite-difference LR gradients where Fnaive fails.As batch size decreases, λmax deviates from 2/η, motivating the Fbatch-based relation in this regime.Source: §3.1 · p. 5, Figure 2 · p. 5

4 FIXED-POINT LR GRADIENT

This section introduces a time-averaged fixed-point approximation to the LR hypergradient that does not require an explicit EoS constraint. It agrees with finite-difference estimates for SGD and signGD, and selected-layer sweeps reveal differing LR sensitivities across architectures.

4 FIXED-POINT LR GRADIENT

  • FP-LR gradients agree with finite-difference estimates for both SGD and signGD, with agreement improving as the averaging window grows.The figure compares FP-LR with finite differences across the two optimizers and varies the averaging window.Source: Figure 3 · p. 6, §4 · p. 6
  • Naive recursive hypergradients tend to diverge, while the naive fixed-point approximation is unstable or differs from finite-difference estimates.The fixed-point version assumes ∂θ/∂η has approximately reached a stationary point.Source: §4 · p. 6
  • The time-averaged fixed-point approximation is computationally stable and approximates the LR gradient without an explicit EoS constraint.Its time averages span [t − τ, t]; under common finite training schedules, the averaged gradient need not vanish near the end of training.Source: §4 · pp. 5–6
  • Time averaging may remove noisy short-timescale dynamics: under EoS, gradients can flip sign period-two while parameters move toward flatter regions over longer timescales.The authors present this as one interpretation of the improved agreement with larger averaging windows.Source: §4 · pp. 6–7
  • Selected-layer sweeps show ResNet sensitivity concentrated in the final layer, while GPT's MLP input layer is highly sensitive to its LR.In ResNet, layers other than the final layer and BatchNorm layers are scale-invariant; in GPT, QK layers are immediately followed by QK normalization.Source: Figure 4 · p. 7

5 QUANTITATIVE UNDERSTANDING OF LR TRANSFER CURVES

The section uses EoS-LR and FP-LR analyses to characterize LR-curve gradients, including scale-invariant cases, empirical bounds, and layerwise sensitivity. Experiments identify the output-layer scale as an important factor in LR transfer across widths.

5.1 Analysis of Scale-invariant models

  • For scale-invariant models, EoS-LR gives a zero LR gradient when F is positively homogeneous, and the gradient is also zero for signGD and simplified Adam when a = 1.More generally, the section reports a zero gradient for scale-invariant layers and says the gradient is generally nonzero if either L2 regularization or scale invariance is absent.Source: §5.1 · pp. 7–8
  • With an L2 constraint ∥θ∥² ≤ Cη², scale-invariant layers have a zero LR gradient when a = 1; otherwise, the expression depends on a and the constraint.The analysis connects the parameter-norm constraint to weight decay with ηλ held constant.Source: §5.1 · p. 8
  • In ResNet, layers followed by normalization are scale-invariant and their loss is not sensitive to their learning rates; in GPT, the MLP input layer is highly sensitive because normalization does not immediately follow it.The stated scale-invariant ResNet exceptions are the final layer and BatchNorm layers; GPT’s QK layers are immediately followed by QK normalization.Source: §5.1 · pp. 7–8

5.2 Bounds on LR gradients

  • Across CNNs, ResNets, and GPT, LR-gradient amplitudes scale approximately as η−3/2 over a broad LR range, between η−1 and η−2.Under EoS, the bound |dL/dη| = Ω(η−3/2) follows from ∥∇λmax∥ scaling as η−1 and ∥∇L∥² = C/η; FP-LR bounds can scale as 1/η or 1/η² but do not themselves yield η−3/2.Source: §5.2 · p. 8, Figure 5 · p. 9

5.3 Layerwise Profiling

  • ResNet’s classifier has a larger layerwise FP-LR gradient than its other layers, while GPT shows large contributions from embedding, output, and MLP input layers.These layerwise profiles quantify each layer’s contribution to LR sensitivity; the decomposition uses a layerwise block-diagonal approximation to the time-averaged Hessian.Source: §5.3 · p. 9, Figure 6 · p. 9
  • With both multipliers set to 1, Base Width=1 is more accurate than Base Width=32; changing the output multiplier substantially affects accuracy and LR transfer, unlike changing the input multiplier.At Base Width=32, reducing the output multiplier to 1/9 brings accuracy to roughly the Base Width=1 default level, while increasing it to 3 lowers accuracy.Source: §5.3 · pp. 9–10, Figure 7 · p. 10

6 LIMITATIONS AND DISCUSSIONS

The discussion frames the results as a foundation for studying stable LR tuning, while identifying open theoretical and methodological questions.

6 LIMITATIONS AND DISCUSSIONS

  • The EoS-LR characterization can incorporate newly established EoS relations, but EoS for practical preconditioned optimizers with moving averages or damping remains unresolved.The authors identify resolving those relations as a step toward characterizing LR curves for such optimizers.Source: §6 · p. 10
  • The work evaluates only first derivatives of the LR curve; higher-order derivatives may help explain curve shape but introduce additional terms that are unlikely to be as simple.Even the current expressions require careful numerical computation.Source: §6 · p. 10
  • Extending model-size analysis beyond the scale-invariant case requires understanding the universality of the largest-eigenvalue derivative, involving third derivatives of the loss.The authors note that theory for large-size limits of these higher-order terms remains under development.Source: §6 · p. 10

A DERIVATIONS OF LR GRADIENTS

The appendix derives LR-gradient results using scale invariance, constrained minimization, and a fixed-point formulation.

A.1 Scale invariance

  • Under scale invariance and positive homogeneity of F, the constrained minimization yields dL/dη = 0.The proof uses the Lagrange-multiplier conditions and the homogeneity relation for F.Source: §A.1 · p. 14

A.2 Scale invariance with weight decay

  • With scale-invariant layers and ηλ held constant, bounded gradients imply a parameter-norm bound proportional to η² for sufficiently large training steps.This bound motivates imposing ∥θ∥² ≤ Cη² in the constrained analysis; the resulting zero-gradient condition is stated for a = 1.Source: §A.2 · pp. 14–15

A.3 FP-LR with scale invariance

  • For GD, the appendix interprets the FP-LR expression as a fixed point of the sensitivity recurrence under stationary sensitivity and time-averaged independence assumptions.It replaces the averaged Hessian by the Hessian at averaged parameters to reduce curvature-computation cost.Source: §A.3 · p. 16

B IMPLEMENTATION

The implementation appendix describes fixed-point approximations for several optimizers, numerical evaluation of EoS-LR, and a K-FAC approximation for layerwise profiles.

B.1.1 SGD

  • For SGD, FP-LR is obtained by treating the averaged parameter sensitivity as stationary, factorizing averaged Hessian–sensitivity terms, and solving for the sensitivity.The implementation approximates the averaged Hessian with the Hessian at averaged parameters to reduce curvature-computation cost.Source: §B.1.1 · p. 17

B.1.2 signGD

  • For signGD, the FP-LR procedure time-averages the sensitivity recurrence and approximates the optimizer derivative using coordinatewise regression slopes.Under approximately Gaussian gradient fluctuations, the regularized population regression slope approaches the expected derivative as ρ → 0; the exact signGD inverse formula cannot be obtained by setting ε = 0.Source: §B.1.2 · p. 18

B.1.3 AdamW

  • For AdamW, FP-LR uses a diagonal regression approximation for the update sensitivity and heuristic independence and curvature replacements.The method uses forward-mode direct hypergradient computation and conjugate gradients for inverse-curvature–vector products where applicable.Source: §B.1.3 · p. 19

B.2 EoS-LR gradient

  • EoS-LR’s gradient is evaluated by approximating F’s directional derivative with a centered finite difference, evaluating Eq. (6)’s intermediate expression without computing third derivatives.Replacing Fnaive with this expression scales the EoS-LR gradient by (1 + 2β)/(1 + β).Source: §B.2 · p. 19

B.3 Layerwise FP-LR Gradient with K-FAC

  • Layerwise FP-LR profiles use empirical K-FAC curvature factors for linear layers, excluding normalization parameters, and evaluate inverse products through factor eigendecompositions.The curvature approximation is formed from layer inputs and backpropagated loss derivatives, with quantities averaged over the final training epoch.Source: §B.3 · p. 20

C.1 EoS-LR Estimates and Gradient Alignment

The section tests EoS-LR estimates against finite differences and examines how sharpness and gradient alignment relate to their agreement. Longer full-batch training improves agreement in the tested MLP and CNN settings.

C.1 EoS-LR Estimates and Gradient Alignment

  • EoS-LR and time-averaged FP-LR agree with finite differences over much of the decreasing-loss LR range for cross-entropy CNNs.EoS-LR also broadly captures the magnitude and LR dependence for ResNet on CIFAR-10; remaining discrepancies near zero at larger LRs vary across seeds.Source: §C.1 · p. 20, Figure 8 · p. 20, Figure 9 · p. 21
  • The largest Hessian eigenvalue is above the EoS threshold 2/η at smaller LRs and below it at larger LRs, consistent with constraint directions associated with decreasing and increasing curve branches.The inequality formulation gives a nonzero LR gradient only when the constraint is active; an inactive constraint has multiplier zero.Source: §C.1 · pp. 20–21, Figure 9 · p. 21
  • At averaged parameters, sharpness-gradient alignment supports the EoS-LR sign, while alignment at individual training steps fluctuates more and can have the opposite sign.For the MLP and CNN, alignment toward −1 at larger LRs becomes more evident with training; the CNN shows it across many LRs after 8192 iterations.Source: §C.1 · pp. 21–22, Figure 10 · p. 21
  • The sharpness-gradient norm scales with an LR exponent close to −1 across three widths, supporting the empirical scaling used in the EoS-LR bound.At exponentially averaged parameters, batch sharpness remains below 2/η while the largest Hessian eigenvalue stays close to that threshold; unaveraged batch sharpness is closer to it.Source: §C.1 · p. 21, Figure 11 · p. 22, Figure 12 · p. 22
  • Longer full-batch training improves EoS-LR agreement with finite differences for the tested MLP and CNN.For the MLP, agreement extends toward lower LRs; for the CNN, differences at 8192 iterations are reduced by 32768 iterations.Source: §C.1 · p. 22

C.2 FP-LR Estimates across Batch Sizes and Architectures

The section evaluates time-averaged FP-LR across batch sizes, architectures, and optimizers, finding improved agreement in several settings but persistent discrepancies in others. It also tests gradient-direction alignment assumptions used to motivate the approximation.

C.2 FP-LR Estimates across Batch Sizes and Architectures

  • Averaged gradients stay closely aligned with augmentation and batch size 2,048, but alignment falls toward zero as LR increases without augmentation at batch size 128.These results do not test changes in gradient magnitude or the independence and derivative assumptions noted in the passage.Source: §C.2 · p. 22, Figure 14 · p. 24, Figure 15
  • Time-averaged FP-LR captures the overall decrease in LR-gradient magnitude across tested CNN batch sizes, but exact-Hessian estimates do not consistently improve agreement over Gauss–Newton estimates.The models use the same epoch count, so smaller batches receive more parameter updates.Source: §C.2 · p. 23, Figure 16 · p. 25
  • Time-averaged FP-LR agrees more closely with finite differences than the naive fixed-point estimate on Fashion-MNIST and improves on that estimate for ResNet8 over much of the LR range.Removing batch normalization increases LR sensitivity in ResNet8.Source: §C.2 · p. 23, Figure 17 · p. 25, Figure 18 · p. 25
  • With a fixed update budget, time-averaged Gauss–Newton FP-LR closely follows finite differences at batch size 4,096 but often overestimates magnitude at batch size 1,024 and below.Time averaging improves agreement but does not remove discrepancies at smaller batch sizes.Source: §C.2 · p. 23, Figure 19 · p. 26

C.3 FP-LR Estimates for signGD and AdamW

  • For ResNet8 trained with signGD, time averaging reduces the naive fixed-point estimate’s overestimation, but the estimate still differs from finite differences, particularly without batch normalization.For the AdamW CNN, the time-averaged Adam approximation largely follows finite differences, with poorer agreement near the nearly flat loss minimum.Source: §C.3 · p. 26, Figure 20 · p. 26, Figure 21 · p. 27
  • For ResNet8, restricting the AdamW LR-gradient estimate to the output layer improves agreement with finite differences over much of the LR range.This motivates layer-wise analysis but does not establish that other layers are generally negligible.Source: §C.3 · p. 26, Figure 22 · p. 27

C.4 Normalization, Weight Decay, and LR Transfer

The section tests how normalization, weight decay, and LR assignment affect LR sensitivity and the LR range associated with low loss or high accuracy. Effects differ across the tested model configurations.

C.4 Normalization, Weight Decay, and LR Transfer

  • QK normalization broadens the low-validation-loss LR range for GPT models trained with a common LR across layers.Across most tested configurations, adding weight decay to the common-LR setting also reduces local LR sensitivity; sensitivity does not vary consistently with model width.Source: §C.4 · p. 27, Figure 23 · p. 27, Figure 24 · p. 28
  • In ResNet, combining batch normalization with LR-independent weight decay broadens the high-test-accuracy LR range compared with using either alone.At width 32, the accuracy-derivative magnitude is also smaller over much of the LR range; this sensitivity measure is an accuracy derivative, unlike the loss derivatives in the other comparisons.Source: §C.4 · pp. 27–28, Figure 25 · p. 28
  • Under 1/d attention scaling, Pre-LN has lower local LR sensitivity than Post-LN in the GPT experiment, while QK-normalization differences are smaller and validation loss is higher with QK normalization.The sensitivity difference between Pre-LN and Post-LN exceeds the difference between models with and without QK normalization in these runs.Source: §C.4 · p. 28, Figure 26 · p. 29
  • With QK normalization, rescaling only QK parameters leaves loss nearly constant, whereas rescaling all parameters changes it, so QK normalization alone does not make the whole model scale-invariant.The QK scaling curves are nearly flat for both Pre-LN and Post-LN despite their different LR sensitivities.Source: §C.4 · p. 28, Figure 26 · p. 29
  • A two-layer Transformer has a broader low-training-loss LR range under Pre-LN than under Post-LN.The compared model widths are 128–4,096.Source: §C.4 · pp. 28–29, Figure 27 · p. 29

D.1 CNN for EoS-LR Experiments

The appendices specify the architectures, datasets, optimizers, and training schedules used across the CNN, ResNet, and GPT experiments. The setups differ by experiment, including their loss functions, augmentation, and LR schedules.

D.1 CNN for EoS-LR Experiments

  • The EoS-LR CNN has three convolutional layers followed by a fully connected hidden layer and a 10-output linear layer, without batch normalization.Experiments use a fixed subset of 8,192 CIFAR-10 training images, squared loss, and constant-LR SGD without momentum or weight decay, for up to 200 epochs.Source: §D.1 · pp. 29–30

D.2 CNN and ResNet8 for FP-LR Experiments

  • The FP-LR experiments use cross-entropy on CIFAR-10 for CNNs and ResNet8, except for the Fashion-MNIST experiment.The CNN variants use SGD and standard CIFAR augmentation; ResNet8 uses stronger augmentation and cosine LR decay.Source: §D.2 · p. 30

D.3 ResNet Experiments with AirBench

  • The ResNet experiments use an AirBench CIFAR-10 implementation with residual blocks, batch normalization, and a fixed whitening convolution.Training uses Nesterov SGD, cross-entropy with label smoothing, data augmentation, and an exponential moving average of weights; Figure 7 evaluates test accuracy.Source: §D.3 · p. 30

D.4 GPT Training

  • The GPT experiments use a 12-block Transformer trained on FineWeb with next-token cross-entropy, comparing Pre-LN and Post-LN with and without QK normalization.They use AdamW and µP scaling; reference training uses 8,192-token sequences, 10,000 iterations including 1,000 warmup iterations, and a separate validation split.Source: §D.4 · p. 31

Related papers

Loading 2610.21814v1…