🤖 machine learning

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

This paper derives exact closed-form expressions for the early training dynamics of linear neural networks to demonstrate that minimizing test loss requires unequal layer-wise learning rates in the first step, followed by equal rates in subsequent steps.

Original authors: Tianyu Pang, Vignesh Kothapalli, Shenyang Deng, Haohui Wang, Dawei Zhou, Yaoqing Yang

Published 2026-06-02

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Tianyu Pang, Vignesh Kothapalli, Shenyang Deng, Haohui Wang, Dawei Zhou, Yaoqing Yang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are training a team of workers to solve a puzzle. In a deep neural network, these "workers" are the different layers of the network. Usually, we tell every worker to learn at the exact same speed (using the same learning rate). But this paper asks: Is that actually the best way to start?

The authors studied simple, linear versions of these networks (like a straight line of workers) to see exactly what happens in the very first few steps of training. They discovered that the "one-size-fits-all" approach is actually a mistake at the beginning, but becomes the best strategy a moment later.

Here is the breakdown of their findings using everyday analogies:

1. The Setup: The Relay Race

Think of a two-layer network as a relay race with two runners:

Runner 1 (The First Layer): Their job is to catch the baton (the raw data) and figure out how to run with it.
Runner 2 (The Second Layer): Their job is to take what Runner 1 did and pass the final message to the finish line (the prediction).

Usually, coaches (the algorithms) tell both runners to sprint at the same speed. The paper shows that in the very first second of the race, this is a bad idea.

2. The "First Step" Surprise: Asymmetry is King

The authors found that right at the start, the two runners have different jobs that require different speeds.

The Metaphor: Imagine Runner 1 is trying to untangle a knot, while Runner 2 is just holding the baton steady. If you make them both run at full speed, Runner 2 might get confused or run out of sync because they aren't ready to sprint yet.
The Finding: To get the best result after just one step, you need unequal learning rates. You need to speed up the first runner (who is doing the heavy lifting of understanding the data) and slow down the second runner (who is just relaying the signal). If you force them to run at the same speed, the team performs worse.

3. The "Second Step" Shift: Balance Takes Over

Now, imagine the race continues to the second step.

The Metaphor: After that first split second, Runner 1 has untangled the knot and is now running smoothly. Runner 2 has caught up and is ready to sprint. Now, they are working as a synchronized team. If one is faster than the other, they start pulling against each other.
The Finding: After two steps, the math shows that the best strategy flips. Now, the team performs best when both runners have the exact same speed. The layers have "coordinated," and balancing their learning rates minimizes errors.

4. The "Goldilocks" Zone: How Big Should the Steps Be?

The paper also looked at how "big" these learning steps should be.

The Metaphor: If you tell the runners to take steps that are too huge, they will trip and fall (the math breaks down). If the steps are too tiny, they will never finish the race.
The Finding: There is a specific "sweet spot" for the size of the learning rate.
- For a two-layer network, the steps can be quite large (scaling with the width of the network).
- For a three-layer network (adding a third runner), the steps must be smaller. The extra layer makes the system more sensitive, so you have to be more careful not to overshoot.

5. The Big Picture: A Dynamic Strategy

The most important takeaway is that learning rates shouldn't be static.

Early Training: You need asymmetry. Treat the layers differently because they are doing different things.
Later Training: You need symmetry. Once the layers have aligned, treat them equally to keep them in sync.

Summary

The paper proves that in the very beginning of training a neural network, treating all layers the same is actually the wrong move. You need to give the "input" layer a different learning rate than the "output" layer to get the best start. However, just a few steps later, the network naturally wants to balance out, and equal learning rates become the optimal choice.

It's like teaching a new employee: at first, you might need to give them very specific, different instructions than your senior staff. But once they get the hang of it, they should all follow the same standard operating procedures to work efficiently together.

Technical Summary: Balancing Learning Rates Across Layers in Linear Neural Networks

Problem Statement

The paper investigates the optimal selection of layer-wise learning rates in deep neural networks, specifically focusing on two-layer and three-layer linear networks trained to learn linear target functions. While practical deep learning often employs separate or adaptive learning rates across layers to accelerate convergence, the theoretical understanding of how the relative scale of these rates affects early training dynamics and generalization remains unclear. Previous analyses have largely relied on continuous-time gradient flow, asymptotic regimes, or kernel-based approximations, which fail to capture the discrete, finite-step dynamics where early layer-wise interactions and the coupling between signal-aligned components and residual terms are critical. The authors aim to provide a precise, finite-step characterization of how layer-wise learning rates shape early generalization.

Methodology

The authors develop a theoretical framework based on exact closed-form expressions for gradients and test loss after one and two steps of gradient descent (GD). The core of their approach involves:

Gradient Decomposition: They decompose the exact gradients into dominant, label-aligned components ( $A_t$ ) and smaller residual terms ( $B_t$ ) arising from weight self-interactions. Under random orthogonal initialization and sufficiently wide networks, they prove that the residual terms are negligible in operator norm compared to the signal-aligned terms for learning rates below specific critical thresholds.
Finite-Step Analysis: Instead of relying on asymptotic limits, the authors derive exact expressions for the test loss after one and two GD updates. This allows them to isolate the leading-order contributions of learning rates to the loss landscape.
Scaling Regimes: They analyze the behavior of the network under different learning rate scalings relative to the network width $h$ . They identify critical thresholds (e.g., $O(h\sqrt{h})$ for two-layer networks and $O(h)$ for three-layer networks) beyond which the gradient dynamics and test loss behavior qualitatively change, connecting their findings to maximal-update and mean-field scaling regimes.
Optimization of Allocation: Under a constrained learning rate budget ( $\eta_1 + \eta_2 = \text{constant}$ ), they analyze the test loss landscape to determine whether symmetric ( $\eta_1 = \eta_2$ ) or asymmetric learning rate allocations are optimal at different stages of training.

Key Contributions

Exact Two-Step Dynamics: The paper derives exact closed-form expressions for the test loss of two-layer and three-layer linear networks after one and two GD steps. This provides a tractable surrogate for analyzing early training dynamics that was previously unavailable.
Gradient Approximation Validity: The authors rigorously bound the residual terms in operator norm, establishing conditions under which the approximate gradients (signal-aligned components only) accurately capture the test loss dynamics.
Phase Transition in Learning Rate Symmetry: A central finding is a distinct early-training regime characterized by a transition from asymmetric to balanced learning rates:
- One-Step Regime: Symmetric learning rates are suboptimal. The test loss is minimized by asymmetric allocations because the first and second layers play distinct roles (representation formation vs. readout) that are best served by unequal learning rates initially.
- Two-Step Regime: In sufficiently wide networks, symmetric learning rates become locally optimal after two steps. This indicates that as training progresses, layers become sufficiently coordinated, making balanced updates advantageous for minimizing test loss.
Depth-Dependent Scaling: The analysis extends to three-layer networks, revealing that the admissible learning rate regime is reduced to $O(h)$ (compared to $O(h\sqrt{h})$ for two layers) due to the additional output layer and initialization schemes. The two-step test loss in three-layer networks exhibits stronger dependence on higher-order products of learning rates, highlighting deeper cross-layer interactions.

Results

Theoretical Validation: The derived closed-form expressions for test loss closely match empirical test losses measured in numerical experiments for both one-step and two-step updates.
Optimal Allocation:
- For two-layer networks, if the total learning rate scale is within $O(h\sqrt{h})$ , the point $\eta_1 = \eta_2$ is not a local minimum after one step but becomes a local minimum after two steps (for sufficiently large width and specific scaling ranges).
- For three-layer networks, a similar transition occurs, but the optimal symmetric regime is established under a stricter learning rate scale of $O(h)$ .
Generalization: The authors demonstrate that these findings hold under Gaussian initialization and in the presence of label noise. Furthermore, experiments on deeper linear networks (4-layer, 8-layer) and a three-layer nonlinear network (with ReLU activation) suggest that the transition from asymmetry to balance generalizes beyond the specific linear settings analyzed theoretically.

Significance and Claims

The paper claims to provide a principled foundation for understanding how layer-wise learning rates shape early generalization in linear neural networks. By moving beyond asymptotic or continuous-time analyses, the authors offer explicit formulas linking learning rate allocation directly to test loss.

The significance of this work lies in:

Theoretical Insight: It elucidates that optimal layer-wise learning rates are dynamic rather than static. Early-stage training benefits from asymmetry to leverage layer-specific signal propagation, while later stages promote balance to enable coordinated alignment.
Practical Guidance: The results suggest a strategy for designing learning rate schedulers: starting with asymmetric rates and transitioning to balanced rates as cross-layer interactions develop. The authors propose a simple scheduler based on layer norm balancing that captures this transition and achieves lower training and test loss compared to uniform baselines in their experiments.
Foundational Framework: The work connects explicit gradient norms, test loss formulas, and width scaling to a theoretical understanding of when and why learning-rate balancing is beneficial, offering a framework that could guide the design of learning rate schedules in more complex architectures, even though the current analysis is restricted to linear networks.

The authors remain modest, noting that their analysis is restricted to linear networks with orthogonal (or Gaussian) initialization and focuses on the first few steps of gradient descent. They acknowledge that extending these results to nonlinear activations, stochastic optimization, and adaptive methods remains a challenging direction for future work.