When Both Layers Learn: Training Dynamics of Representing Linear Models via ReLU Networks

• URL: http://arxiv.org/abs/2606.04476v1
• Date: Wed, 03 Jun 2026 05:44:30 GMT
• Title: When Both Layers Learn: Training Dynamics of Representing Linear Models via ReLU Networks
• Abstract summary: We study the gradient descent dynamics for jointly training both layers of a one-hidden-layer ReLU network to fit a linear target function.<n>Our analysis tracks the trajectory through three phases: an alignment phase in which hidden weights progressively align with the planted direction while the output weights maintain the correct sign pattern.<n>We establish novel uniform concentration results that hold along the entire trajectory, and are essential for obtaining order-wise optimal sample complexity.
• Score: 22.81761236732655
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: In this paper, we study the gradient descent dynamics for jointly training both layers of a one-hidden-layer ReLU network to fit a linear target function. Concretely, we consider a realizable setting where inputs are drawn i.i.d. from a Gaussian distribution and labels follow a planted linear model. This stylized framework captures salient features of end-to-end training in inverse problems and certain auto-encoder models. Despite its apparent simplicity, the dynamics remain poorly understood, in part because the loss landscape contains multiple non-strict saddle points, making it unclear why gradient descent from random initialization reliably escapes bad stationary regions. We provide a detailed characterization of the optimization landscape and prove that gradient descent from a moderately small random initialization-simultaneously training both layers-converges to a global minimizer at a linear rate with order-wise optimal sample complexity. Our analysis tracks the trajectory through three phases: an alignment phase in which hidden weights progressively align with the planted direction while the output weights maintain the correct sign pattern; a growth phase in which the norms of both layers increase while preserving alignment; and a local refinement phase in which the aligned neurons rapidly converge to the planted direction, yielding fast local convergence. To rigorously show that GD avoids non-strict saddles, we develop trajectory-level control arguments for the end-to-end dynamics. In addition, we establish novel uniform concentration results that hold along the entire trajectory, and are essential for obtaining order-wise optimal sample complexity. We corroborate our theory with extensive experiments across a range of configurations.

Related papers

• Across the Loss Landscape with Progressive Growth [51.366966420881646]
  We show how grow-and-optimize strategies to train toward flatter regions by viewing growth as progressive relaxation.<n>We empirically validate these predictions in controlled toy basins and in a realistic ResNet/CI-100 setting.<n>Although progressive subspace growth reliably produces flatter solutions, curvature reductions do not universally translate into improved test performance, highlighting subtleties in the flatness-generalization connection.
  arXiv  Detail & Related papers  (2026-08-25T13:50:48Z)
• Need We Teach Foundation Models What is a Generative Image? Gradient-Free Generative Artifact Detection via Analytic Spectral Adaptation [3.881211374324378]
  Adapting foundation models to detect generative artifacts via gradient-based updates compromises their intrinsic representations.<n>We propose a gradient-free methodology reframing detection from binary classification to an out-of-distribution (OOD) anomaly measurement problem.
  arXiv  Detail & Related papers  (2026-06-03T13:04:23Z)
• Enhancing Gradient Inversion Attacks in Federated Learning via Hierarchical Feature Optimization [56.95448807869383]
  Federated Learning (FL) has emerged as a compelling paradigm for privacy-preserving distributed machine learning.<n>Recent studies find that the gradients exchanged in the FL system are also vulnerable to privacy leakage.<n>We propose textbfGradient textbfInversion over textbfFeature textbfDomains (GIFD)
  arXiv  Detail & Related papers  (2026-04-01T14:32:15Z)
• Parallel Diffusion Solver via Residual Dirichlet Policy Optimization [88.7827307535107]
  Diffusion models (DMs) have achieved state-of-the-art generative performance but suffer from high sampling latency due to their sequential denoising nature.<n>Existing solver-based acceleration methods often face significant image quality degradation under a low-dimensional budget.<n>We propose the Ensemble Parallel Direction solver (dubbed as EPD-EPr), a novel ODE solver that mitigates these errors by incorporating multiple gradient parallel evaluations in each step.
  arXiv  Detail & Related papers  (2025-12-28T05:48:55Z)
• Description of the Training Process of Neural Networks via Ergodic Theorem : Ghost nodes [3.637162892228131]
  We present a unified framework for understanding and accelerating deep neural networks via training gradient descent (SGD)<n>We introduce a practical diagnostic, the running estimate of the largest Lyapunov exponent, which distinguishes genuine convergence toward stablers.<n>We propose a ghost category extension for standard classifiers that adds auxiliary ghost output nodes so the model gains extra descent directions.
  arXiv  Detail & Related papers  (2025-07-01T17:54:35Z)
• Gradient Descent Converges Linearly to Flatter Minima than Gradient Flow in Shallow Linear Networks [0.0]
  We study the gradient descent dynamics of a depth-2 linear neural network with a single input and output.<n>We show that GD converges at an explicit linear rate to a global minimum of the training loss, even with a large stepsize.
  arXiv  Detail & Related papers  (2025-01-15T20:43:36Z)
• Hallmarks of Optimization Trajectories in Neural Networks: Directional Exploration and Redundancy [75.15685966213832]
  We analyze the rich directional structure of optimization trajectories represented by their pointwise parameters. We show that training only scalar batchnorm parameters some while into training matches the performance of training the entire network.
  arXiv  Detail & Related papers  (2024-03-12T07:32:47Z)
• Training Dynamics of Multi-Head Softmax Attention for In-Context Learning: Emergence, Convergence, and Optimality [54.20763128054692]
  We study the dynamics of gradient flow for training a multi-head softmax attention model for in-context learning of multi-task linear regression. We prove that an interesting "task allocation" phenomenon emerges during the gradient flow dynamics.
  arXiv  Detail & Related papers  (2024-02-29T18:43:52Z)
• Convergence of mean-field Langevin dynamics: Time and space discretization, stochastic gradient, and variance reduction [49.66486092259376]
  The mean-field Langevin dynamics (MFLD) is a nonlinear generalization of the Langevin dynamics that incorporates a distribution-dependent drift. Recent works have shown that MFLD globally minimizes an entropy-regularized convex functional in the space of measures. We provide a framework to prove a uniform-in-time propagation of chaos for MFLD that takes into account the errors due to finite-particle approximation, time-discretization, and gradient approximation.
  arXiv  Detail & Related papers  (2023-06-12T16:28:11Z)
• 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)
• On regularization of gradient descent, layer imbalance and flat minima [9.08659783613403]
  We analyze the training dynamics for deep linear networks using a new metric - imbalance - which defines the flatness of a solution. We demonstrate that different regularization methods, such as weight decay or noise data augmentation, behave in a similar way.
  arXiv  Detail & Related papers  (2020-07-18T00:09:14Z)

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.