Optimal Learning Rate Scaling Depends on Data in Deep Scalar Linear Networks

Yedi Zhang; Peter E. Latham; Leena Chennuru Vankadara; Andrew Saxe

cs.LGJul 8, 2026

Optimal Learning Rate Scaling Depends on Data in Deep Scalar Linear Networks

Authors: Yedi Zhang, Peter E. Latham, Leena Chennuru Vankadara, Andrew Saxe

Organizations: Gatsby Computational Neuroscience Unit, University College London · Sainsbury Wellcome Centre, University College London

Abstract

In this short note we consider the gradient descent dynamics of deep scalar linear networks, $f(x) = \prod_{l=1}^L w_l x$ , which enjoy exact time-course solutions for any integer depth. We show that even in this minimal model, the optimal depth-wise learning rate scaling depends on data, whereas data-agnostic scaling rules fail to transfer across depths. Under the data-dependent optimal scaling, the learning dynamics is independent of data and weakly dependent on depth, resulting in a constant linear convergence rate across all depths including infinity. We further show similar data-dependent effects in deep scalar linear networks with residual connections.

Explore similar work

Balancing Learning Rates Across Layers: Exact Two-Step Dynamics and Optimal Scaling in Linear Neural Networks
May 29, 2026Tianyu Pang, Vignesh Kothapalli, Shenyang Deng +3Learning Rates Two-Layer Neural Networks
The Map Behind the Flow: Finite-Step Gradient Descent as a Dynamical System
Jul 6, 2026Thomas HofmannGradient Descent Wasserstein Gradient Flows
On the Nonlinearity of Learning Rate Scaling for LLM Training
Jun 28, 2026Zaiwen Yang, Huaqing Zhang, Jing Xu +1Learning Rates Large Language Model Training

May 29, 2026 · cs.LGJ/K move · Enter open · S save

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

Tianyu Pang, Vignesh Kothapalli, Shenyang Deng, Haohui Wang+2

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.

Learning Rates Two-Layer Neural Networks Gradient Descent Rate Scaling

Open paper Ask about it PDF

Abstract

Explore similar work

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

The Map Behind the Flow: Finite-Step Gradient Descent as a Dynamical System

On the Nonlinearity of Learning Rate Scaling for LLM Training