Balancing Learning Rates Across Layers: Exact Two-Step Dynamics and Optimal Scaling in Linear Neural Networks
Summary
This paper derives exact closed-form expressions for gradients and test loss after one and two steps of gradient descent in two-layer and three-layer linear neural networks, characterizing optimal learning rate selection and revealing a distinct early-training regime where unequal layer-wise learning rates are initially optimal.
View Cached Full Text
Cached at: 06/02/26, 03:41 PM
# Balancing Learning Rates Across Layers: Exact Two-Step Dynamics and Optimal Scaling in Linear Neural Networks Source: [https://arxiv.org/abs/2606.00340](https://arxiv.org/abs/2606.00340) [View PDF](https://arxiv.org/pdf/2606.00340) > Abstract: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:[this https URL](https://github.com/TDCSZ327/Layer-Balancing)\. ## Submission history From: Tianyu Pang \[[view email](https://arxiv.org/show-email/83449265/2606.00340)\] **\[v1\]**Fri, 29 May 2026 20:28:52 UTC \(7,485 KB\)
Similar Articles
Optimal Learning Rate Scaling Depends on Data in Deep Scalar Linear Networks
This paper demonstrates that optimal learning rate scaling in deep scalar linear networks is inherently data-dependent, contradicting prior data-agnostic scaling rules. It shows that with data-dependent scaling, convergence becomes depth-independent, including at infinite depth.
Gradient Descent with Large Step Size Restores Symmetry in Deep Linear Networks with Multi-Pathway
This paper shows that discrete Gradient Descent with large step sizes restores symmetry in multi-pathway Deep Linear Networks, countering the symmetry-breaking predicted by Gradient Flow, and leads to signal re-balancing across pathways. The authors theoretically prove that balanced solutions are flatter (less sharp) than sparse ones, and large learning rates drive the network toward stable, balanced configurations.
Spectral Asymptotics of Neural Network Loss Landscapes: An Exact Decomposition of the Curvature Exponent
This paper presents an exact decomposition of the curvature exponent α in neural network loss landscapes, explaining why it varies across layer types. It introduces the spectral alignment decomposition and derives a spectral transfer identity linking curvature, gradient rank decay, and Hessian exponents, validated across architectures and datasets.
Can Neural Networks Achieve Optimal Computational-statistical Tradeoff? An Analysis on Single-Index Model
This paper demonstrates that two-layer neural networks trained with gradient-based methods can achieve the optimal computational-statistical tradeoff for learning Gaussian single-index models, matching the SQ lower bound up to polylogarithmic factors for all generative exponents and extending to sparse settings with a novel weight perturbation technique.
Scaling Laws, Carefully (25 minute read)
A comprehensive overview of scaling laws in deep learning, tracing their theoretical roots and empirical findings, and explaining how loss decreases predictably with model size, data, and compute.