Measuring the local shape of learning-rate transfer curves

Satoki Ishikawa; Rio Yokota; Ryo Karakida

cs.LGOct 15, 2026

Measuring the local shape of learning-rate transfer curves

Authors: Satoki Ishikawa, Rio Yokota, Ryo Karakida

Organizations: Institute of Science Tokyo · AIST

Abstract

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.

Figures & tables

Figure 1: Learning-rate transfer curves and their local shape at the edge of stability. (a) LR transfer curves show how the final training loss depends on the learning rate

\eta

across network widths. (b) We focus on the EoS regime, where the largest Hessian eigenvalue satisfies

\lambda_{\max}(H)\approx 2/\eta

. (c) The EoS-LR gradient of

F_{\mathrm{naive}}

closely tracks the finite difference estimate, thereby capturing the local shape of the LR transfer curve (shown for width 256).

Figure 2: EoS-LR gradients based on

F_{batch}

capture LR sensitivity at small batch sizes. (a) As the batch size decreases,

\lambda_{max}

deviates from

2/\eta

. (b,c) At small batch sizes, the LR gradient based on

F_{batch}

closely matches the finite-difference estimate.

Figure 3: FP-LR gradients agree with finite-difference estimates. (a,b) FP-LR captures LR sensitivity for both SGD and signGD. (c) Increasing the averaging window improves agreement between the FP-LR gradient and the finite-difference estimate.

Figure 4: Validation loss versus selected-layer LR. We sweep the LR of each indicated parameter group while fixing the others: (a,b) ResNet on CIFAR-10; (c) GPT on FineWeb.

Figure 5: LR gradients of the evaluation loss exhibit

\eta^{-3/2}

scaling. Log–log plots show the absolute LR gradients for CNN, ResNet, and GPT, estimated by finite differences from LR sweeps. We use the test loss for CNN and the validation loss for ResNet and GPT; colors indicate network width. Reference lines show the slopes

\eta^{-2}

\eta^{-3/2}

, and

\eta^{-1}

suggested by the bounds.

Figure 6: Layerwise FP-LR gradient. (a,b) ResNet on CIFAR-10 with two different LRs. (c) GPT on FineWeb. (d) Blocks

9

–

11

of (c), shown on a linear scale.

Figure 7: Accuracy and LR transfer under different input and output multipliers. Test accuracy for ResNet on CIFAR-10 under

\mu

P, comparing Base Width=

1

and

32

. (a) Both multipliers are

1

. (b,c) Only the output multiplier changes to

3

1/9

, respectively. (d,e) Only the input multiplier changes to

3

1/3

, respectively.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: LR-gradient estimates for cross-entropy training. CNNs on CIFAR-10 with widths

128

256

, and

512

. (a) At width

256

, we compare finite differences, EoS-LR based on

F_{\mathrm{naive}}

and

F_{\mathrm{batch}}

, and time-averaged FP-LR using Gauss–Newton curvature. (b–d) The three estimators across widths. The plots show both the sign and magnitude of the LR gradient.

Figure 9: EoS-LR estimates for ResNet on CIFAR-10. Left: training loss versus LR. Middle: EoS-LR estimates and finite-difference LR gradients. Right: the largest Hessian eigenvalue and the EoS boundary

2(1+\beta)/\eta(1+2\beta)

. Each curve corresponds to one of three seeds.

Figure 10: Parameter averaging reduces fluctuations in sharpness-gradient alignment. ResNet on CIFAR-10 at width

64

, trained for

5{,}000

epochs with LRs

0.1

and

1

. We plot

\cos(\nabla\mathcal L,\nabla\lambda_{\max})

at (a) averaged parameters and (b) current parameters. Curves and shaded regions show seed means and standard errors.

Figure 11: LR scaling of the sharpness-gradient norm. CNNs on CIFAR-10 at widths

128

256

, and

512

. We estimate

\|\nabla_{\theta}\lambda_{\max}\|

from finite differences of

F_{\mathrm{naive}}

. Reference lines show exponents

-3/2

-1

, and

-1/2

; the observed scaling is close to

-1

. Points and error bars show seed means and standard errors.

Figure 12: Batch sharpness falls below the EoS threshold at averaged parameters. Sharpness versus LR on Fashion-MNIST using an MLP. The first two panels show results after

200

epochs, and the last two after

2000

epochs. Within each pair, the left panel uses the final parameters and the right panel uses exponentially averaged parameters (EWA). We compare the largest Hessian eigenvalue and batch sharpness with

2/\eta

(dashed).

Figure 13: LR-gradient agreement with longer full-batch GD training. Finite differences (gray) and EoS-LR (teal) for the MLP (top) and CNN (bottom) trained for increasing numbers of iterations

T

512

examples from two classes of Fashion-MNIST. The averaging window remains

512

updates.

Figure 14: Gradient alignment with longer training. Cosine similarity

\cos(\nabla\mathcal{L},\nabla\lambda_{\max})

at parameters averaged over the last

512

updates, for the same MLP (top) and CNN (bottom) as in Figure 13 .

Figure 15: Alignment of gradients at successive averaged parameters. A width-

256

CNN trained on CIFAR-10 with mean-squared loss for

200

epochs. We compare batch sizes

128

and

2{,}048

, with and without augmentation. We plot the cosine similarity of full-data gradients evaluated at successive epoch-averaged parameters near the end of training. Error bars show standard errors across seeds.

Figure 16: FP-LR estimates across batch sizes. A width-

32

CNN is trained on CIFAR-10 for

200

epochs, with batch sizes from

64

4{,}096

. Each panel compares the magnitude of the test-loss LR gradient from finite differences with truncated forward-mode, naive fixed-point, and time-averaged FP-LR estimates. “Hessian” distinguishes exact-Hessian estimates from those using Gauss–Newton curvature.

Figure 17: LR-gradient estimates on Fashion-MNIST. Width-

32

CNNs with and without batch normalization are trained with batch size

1{,}024

. Left: test loss versus LR. Middle and right: LR-gradient magnitudes for the models with and without batch normalization, respectively.

Figure 18: Batch normalization reduces ResNet8 LR sensitivity. ResNet8 is trained on CIFAR-10 for

100

epochs with SGD, width

8

, and batch size

256

. Left: test-loss curves with BN and weight decay, without BN, and without weight decay. Remaining panels: corresponding LR-gradient magnitudes, using the exact Hessian or the Gauss–Newton approximation.

Figure 19: MLP LR-gradient estimates with a fixed update budget. A width-

1{,}024

MLP with ELU activations is trained on CIFAR-10 for

20{,}000

updates. We compare batch sizes

64

256

1{,}024

, and

4{,}096

. Four panels show the magnitude of the test-loss LR gradient; the upper-right panel shows test loss.

Figure 20: Time averaging reduces overestimation by the naive signGD estimator. ResNet8 on CIFAR-10 with width

8

and batch size

256

, with and without batch normalization. Left: test loss. Middle and right: LR-gradient magnitudes, comparing finite differences, time-averaged signGD, naive signGD, and the time-averaged GD approximation.

Figure 21: Time-averaged estimates for CNN training with AdamW. A width-

32

CNN with batch normalization is trained on CIFAR-10 with batch size

1{,}024

. Left: test loss. Right: LR-gradient magnitudes from finite differences and the GD, signGD, and Adam approximations. The naive Adam estimate is included for comparison.

Figure 22: Output-layer estimates more closely follow finite differences for ResNet8. ResNet8 with batch normalization is trained on CIFAR-10 using AdamW, width

8

, and batch size

256

. Left: test loss. Right: LR-gradient magnitudes, including estimates restricted to the output layer. The vertical line marks the LR minimizing test loss.

Figure 23: QK normalization broadens the low-loss LR range with a common LR across layers. GPT models trained on FineWeb with AdamW and weight decay. Panels compare Pre-LN and Post-LN, with and without QK normalization; colors indicate model width.

Figure 24: Layerwise LRs and weight decay reduce local LR sensitivity in most tested GPT configurations. Columns compare Pre-LN and Post-LN, with and without QK normalization. The three curves use layer-wise LRs without weight decay, a common LR without weight decay, or a common LR with weight decay. At each model width, sensitivity is the mean of the five smallest nonzero absolute finite-difference estimates of the validation-loss LR gradient across the LR sweep.

Figure 25: Batch normalization and weight decay broaden the high-accuracy LR range. ResNet on CIFAR-10 trained with AdamW and LR-independent weight decay. (a–c) Test accuracy without BN, without weight decay, and with both, across widths

32

–

512

. Curves are smoothed in log-LR space; faint points show seed means, shaded bands show standard errors, and dashed lines mark the width-

32

optimum. (d) At width

32

, the absolute finite-difference derivative of smoothed accuracy, measured in percentage points per unit LR.

Figure 26: Pre-LN reduces local LR sensitivity under

1/d

attention scaling. GPT models are trained on FineWeb with AdamW. Top: validation loss versus LR for model widths

256

–

4{,}096

, comparing Pre-LN and Post-LN with and without QK normalization. Bottom left: local LR sensitivity, measured by the mean of the five smallest nonzero LR-gradient magnitudes estimated by finite differences. Bottom middle and right: loss after multiplying all parameters or only QK parameters by

c

, at model width

256

and LR

0.5

Figure 27: Pre-LN gives a broader low-loss LR range in a two-layer Transformer. We use the implementation of Yang et al. (2021) with

1/d

attention scaling. Panels compare Pre-LN and Post-LN, with and without QK normalization. Curves show training loss for model widths

128

–

4{,}096

Neural Network Optimization

Neural network optimization studies algorithms and training dynamics that help models learn efficiently, reliably, and generalize well.

See all 321 papers

Fast Learning Rate Transfer in Shallow Linear Networks at Growing Training Horizons
Sep 28, 2026Mana Sakai, Masaaki ImaizumiHyperparameter Transfer Neural Network Training Dynamics
The University of Tokyo · RIKEN Center for Advanced Intelligence Project · Kyoto University
Balancing Learning Rates Across Layers: Exact Two-Step Dynamics and Optimal Scaling in Linear Neural Networks
May 29, 2026Tianyu Pang, Vignesh Kothapalli, Shenyang Deng +3Deep Learning Optimization Neural Network Training Dynamics
Dartmouth College · Stanford University · Virginia Tech
On the Nonlinearity of Learning Rate Scaling for LLM Training
Jun 28, 2026Zaiwen Yang, Huaqing Zhang, Jing Xu +1Language Model Scaling Laws Hyperparameter Transfer
Tsinghua University · Tsinghua University, Shanghai Qi Zhi Institute

Sep 28, 2026J/K move · Enter open · S save

Fast Learning Rate Transfer in Shallow Linear Networks at Growing Training Horizons

Mana Sakai, Masaaki Imaizumi

The University of Tokyo · RIKEN Center for Advanced Intelligence Project · Kyoto University

Hyperparameter Transfer Neural Network Training Dynamics

Open paper AskPDF arXiv

Hyperparameter transfer across model width can substantially reduce the cost of tuning large neural networks, but its behavior when the training horizon grows with width is not fully understood. Building on the framework of fast hyperparameter transfer (Ghosh et al., 2026), which formalizes when transfer is effective, we investigate conditions that ensure fast transfer in the growing-horizon regime. Specifically, we study learning-rate transfer in a shallow linear network with a single trainable hidden matrix, trained by full-batch gradient descent. Under additional spectral assumptions, our main results are threefold. (i) We prove fast learning-rate transfer as

n,T\to\infty

whenever

T=o(\sqrt{n})

. (ii) We characterize the transfer rates through the finite-width perturbation scale, the first-order sensitivities of the loss and its learning-rate derivative to finite-width perturbations, and the local loss curvature. (iii) We derive limiting distributions for the optimal learning rate and optimized loss, governed by fluctuations associated with the extreme eigenvalues of the data Gram matrix. These results clarify how spectral structure and local loss sensitivities govern learning-rate transfer at growing horizons.

Abstract

Neural Network Optimization

Explore similar work

Fast Learning Rate Transfer in Shallow Linear Networks at Growing Training Horizons

Balancing Learning Rates Across Layers: Exact Two-Step Dynamics and Optimal Scaling in Linear Neural Networks

On the Nonlinearity of Learning Rate Scaling for LLM Training