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

arXiv: https://arxiv.org/abs/2606.00340 · HTML: https://mlsanity.com/paper/2606.00340/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 49 references that are in the map.
- First posted on arXiv 29 May 2026.

## Builds on (most cited first, 20 of 20)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [On the Variance of the Adaptive Learning Rate and Beyond](https://mlsanity.com/paper/1908.03265/) (2019)
- [Exact solutions to the nonlinear dynamics of learning in deep linear neural networks](https://mlsanity.com/paper/1312.6120/) (2013)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [Large Batch Training of Convolutional Networks](https://mlsanity.com/paper/1708.03888/) (2017)
- [Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer](https://mlsanity.com/paper/2203.03466/) (2022)
- [A Convergence Analysis of Gradient Descent for Deep Linear Neural Networks](https://mlsanity.com/paper/1810.02281/) (2018)
- [Outlier Weighed Layerwise Sparsity (OWL): A Missing Secret Sauce for Pruning LLMs to High Sparsity](https://mlsanity.com/paper/2310.05175/) (2023)
- [Adam-mini: Use Fewer Learning Rates To Gain More](https://mlsanity.com/paper/2406.16793/) (2024)
- [Layer-adaptive sparsity for the Magnitude-based Pruning](https://mlsanity.com/paper/2010.07611/) (2020)
- [Implicit Regularization of Discrete Gradient Dynamics in Linear Neural Networks](https://mlsanity.com/paper/1904.13262/) (2019)
- [Tensor Programs VI: Feature Learning in Infinite-Depth Neural Networks](https://mlsanity.com/paper/2310.02244/) (2023)
- [Channel Pruning via Automatic Structure Search](https://mlsanity.com/paper/2001.08565/) (2020)
- [Universal Language Model Fine-tuning for Text Classification](https://mlsanity.com/paper/1801.06146/) (2018)
- [Large Learning Rate Tames Homogeneity: Convergence and Balancing Effect](https://mlsanity.com/paper/2110.03677/) (2021)
- [The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training](https://mlsanity.com/paper/2502.19002/) (2025)
- [AlphaLoRA: Assigning LoRA Experts Based on Layer Training Quality](https://mlsanity.com/paper/2410.10054/) (2024)
- [AlphaDecay: Module-wise Weight Decay for Heavy-Tailed Balancing in LLMs](https://mlsanity.com/paper/2506.14562/) (2025)
- [Layer-wise Importance Matters: Less Memory for Better Performance in Parameter-efficient Fine-tuning of Large Language Models](https://mlsanity.com/paper/2410.11772/) (2024)

## Built on by (most cited first, 1 of 1)

- [A Full Adam Theorem for Spectral Heavy-Tail Onset](https://mlsanity.com/paper/2609.12996/) (2026)

## Abstract (arXiv)

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.