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

2026-05-29

AI Summary

This study investigates how to assign optimal layer-wise learning rates in the early phase of training deep linear neural networks to minimize test loss. By deriving an exact closed-form solution after two steps of gradient descent, the authors characterize the initial training dynamics and propose a gradient-based inter-layer learning rate scaling strategy, constructing an analytically tractable surrogate loss function. Theoretical analysis reveals that employing unequal learning rates across layers at initialization enhances performance, whereas uniform learning rates become preferable in subsequent steps. This work presents the first precise theoretical characterization of the dynamics during the first two training iterations, and numerical experiments confirm that the proposed strategy significantly reduces test loss compared to baseline approaches.

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: https://github.com/TDCSZ327/Layer-Balancing.

Problem

• learning rate
• linear neural networks
• layer-wise optimization
• early training dynamics
• test loss

Innovation

• layer-wise learning rates
• exact dynamics
• linear neural networks
• optimal scaling
• early training regime