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

• URL: http://arxiv.org/abs/2606.00340v1
• Date: Fri, 29 May 2026 20:28:52 GMT
• Title: Balancing Learning Rates Across Layers: Exact Two-Step Dynamics and Optimal Scaling in Linear Neural Networks
• Abstract summary: We study optimal learning-rate selection in two-layer and three-layer linear neural networks trained to learn linear target functions.<n>We derive the exact closed-form expressions for the gradients and test loss after one and two steps of gradient descent.
• Score: 40.02433292363059
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Deep learning from the crowd Fundamentals of morphological galaxy classification [1.4375822229241138]
  We use Galaxy Zoo 1 as our experimental testbed and trained a convolutional neural network (CNN) for the automatic classification of galaxies' morphologies.<n>We analyze the impact of the following aspects on the classification accuracy and training efficiency.
  arXiv  Detail & Related papers  (2026-09-06T00:36:30Z)
• On Neural Scaling Laws for Weather Emulation through Continual Training [76.70310203875866]
  We study neural scaling in Scientific Machine Learning, focusing on models for weather forecasting.<n>To analyze scaling behavior in as simple a setting as possible, we adopt a minimal, scalable, general-purpose Swin Transformer architecture.<n>We show that models trained in this minimalist way follow predictable scaling trends and even outperform standard cosine learning rate schedules.
  arXiv  Detail & Related papers  (2026-03-26T17:37:25Z)
• Feature Learning Beyond the Edge of Stability [8.430481660019451]
  We propose a homogeneous multilayer perceptron parameterization with hidden layer width pattern and analyze its training dynamics under gradient descent.<n>We obtain formulas for the first three Taylor coefficients of the minibatch loss during training that illuminate the connection between sharpness and feature learning.
  arXiv  Detail & Related papers  (2025-02-18T18:23:33Z)
• A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networks [43.281323350357404]
  Feature learning is thought to be one of the fundamental reasons for the success of deep neural networks.<n>We show that with a learning rate that grows with the sample size, such training in fact introduces multiple rank-one components.
  arXiv  Detail & Related papers  (2023-10-11T20:55:02Z)
• Training a Two Layer ReLU Network Analytically [4.94950858749529]
  We will explore an algorithm for training two-layer neural networks with ReLU-like activation and the square loss. The method is faster than the gradient descent methods and has virtually no tuning parameters.
  arXiv  Detail & Related papers  (2023-04-06T09:57:52Z)
• Implicit Bias in Leaky ReLU Networks Trained on High-Dimensional Data [63.34506218832164]
  In this work, we investigate the implicit bias of gradient flow and gradient descent in two-layer fully-connected neural networks with ReLU activations. For gradient flow, we leverage recent work on the implicit bias for homogeneous neural networks to show that leakyally, gradient flow produces a neural network with rank at most two. For gradient descent, provided the random variance is small enough, we show that a single step of gradient descent suffices to drastically reduce the rank of the network, and that the rank remains small throughout training.
  arXiv  Detail & Related papers  (2022-10-13T15:09:54Z)
• Slimmable Networks for Contrastive Self-supervised Learning [69.9454691873866]
  Self-supervised learning makes significant progress in pre-training large models, but struggles with small models. We introduce another one-stage solution to obtain pre-trained small models without the need for extra teachers. A slimmable network consists of a full network and several weight-sharing sub-networks, which can be pre-trained once to obtain various networks.
  arXiv  Detail & Related papers  (2022-09-30T15:15:05Z)
• A Loss Curvature Perspective on Training Instability in Deep Learning [28.70491071044542]
  We study the evolution of the loss Hessian across many classification tasks in order to understand the effect curvature of the loss has on the training dynamics. Inspired by the conditioning perspective, we show that learning rate warmup can improve training stability just as much as batch normalization.
  arXiv  Detail & Related papers  (2021-10-08T20:25:48Z)
• Learning Neural Network Subspaces [74.44457651546728]
  Recent observations have advanced our understanding of the neural network optimization landscape. With a similar computational cost as training one model, we learn lines, curves, and simplexes of high-accuracy neural networks. With a similar computational cost as training one model, we learn lines, curves, and simplexes of high-accuracy neural networks.
  arXiv  Detail & Related papers  (2021-02-20T23:26:58Z)
• Path Sample-Analytic Gradient Estimators for Stochastic Binary Networks [78.76880041670904]
  In neural networks with binary activations and or binary weights the training by gradient descent is complicated. We propose a new method for this estimation problem combining sampling and analytic approximation steps. We experimentally show higher accuracy in gradient estimation and demonstrate a more stable and better performing training in deep convolutional models.
  arXiv  Detail & Related papers  (2020-06-04T21:51:21Z)
• The duality structure gradient descent algorithm: analysis and applications to neural networks [0.0]
  We propose an algorithm named duality structure gradient descent (DSGD) that is amenable to non-asymptotic performance analysis. We empirically demonstrate the behavior of DSGD in several neural network training scenarios.
  arXiv  Detail & Related papers  (2017-08-01T21:24:38Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.