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 η across network widths. (b) We focus on the EoS regime, where the largest Hessian eigenvalue satisfies λmax(H)≈2/η . (c) The EoS-LR gradient of Fnaive 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 Fbatch capture LR sensitivity at small batch sizes. (a) As the batch size decreases, λmax deviates from 2/η . (b,c) At small batch sizes, the LR gradient based on Fbatch 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 η−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 η−2 , η−3/2 , and η−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 μ P, comparing Base Width= 1 and 32 . (a) Both multipliers are 1 . (b,c) Only the output multiplier changes to 3 or 1/9 , respectively. (d,e) Only the input multiplier changes to 3 or 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 Fnaive and Fbatch , 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+β)/η(1+2β) . 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(∇L,∇λ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 ∥∇θλmax∥ from finite differences of Fnaive . 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/η (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 on 512 examples from two classes of Fashion-MNIST. The averaging window remains 512 updates.
Figure 14: Gradient alignment with longer training. Cosine similarity cos(∇L,∇λ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 to 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 .
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→∞ whenever T=o(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.
Mana Sakai, Masaaki Imaizumi
The University of Tokyo · RIKEN Center for Advanced Intelligence Project · Kyoto University
We study optimal learning-rate selection in two-layer and three-layer linear neural networks trained to learn linear target functions. In particular, we derive the exact closed-form expressions for the gradients and test loss after one and two steps of gradient descent, enabling a precise characterization of early training dynamics. We characterize how learning rates should scale under the gradient approximation in the first two steps, and prove that performing updates with this approximation yields a tractable surrogate loss with a tight, small approximation error. This formulation enables the theoretical analysis of layer-wise learning rates and reveals a distinct early-training regime: test loss can be minimized by unequal learning rates at the initial step, while equal learning rates become optimal in subsequent steps. Our numerical experiments validate the theory and demonstrate the importance of balancing layer-wise learning rates early during training. The code is available at: https://github.com/TDCSZ327/Layer-Balancing.
Tianyu Pang, Vignesh Kothapalli, Shenyang Deng +3
Dartmouth College · Stanford University · Virginia Tech
Learning-rate transfer can reduce the cost of training large language models: instead of sweeping learning rates at target scale, practitioners extrapolate from smaller runs. Existing approaches often assume that the optimal learning rate follows a log-linear scaling law in data scale and model size. We carefully examine and evaluate this scaling law. In our empirical study of GPT-2--style models from 22M to 707M parameters trained on 5B to 100B tokens, the optimal learning rate develops upward curvature at larger scales, leading to inaccurate extrapolation. We find that this curvature largely disappears when learning rates are replaced by effective learning rate (the step size in normalized weight space), and when data D extrapolation is used instead of model size N extrapolation. Next, we explain nonlinearity in scaling: weight-norm converges to equilibrium slower when optimal learning is small, requiring a larger step size to reduce the transient phase. Experiments with AdamH, which directly controls the effective learning rate, further support this explanation.
Zaiwen Yang, Huaqing Zhang, Jing Xu +1
Tsinghua University · Tsinghua University, Shanghai Qi Zhi Institute