Feature Learning in Linear-Width Two-Layer Networks: Two vs. One Step of Gradient Descent

• URL: http://arxiv.org/abs/2605.17767v2
• Date: Thu, 21 May 2026 20:45:44 GMT
• Title: Feature Learning in Linear-Width Two-Layer Networks: Two vs. One Step of Gradient Descent
• Abstract summary: We study feature learning in two-layer neural networks within the linear-width regime.<n>By characterizing these early-phase evolutions, our work proposes a tractable framework for studying optimization and feature learning phenomenology.
• Score: 32.9638210129515
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: We study feature learning in two-layer neural networks within the linear-width regime, where the number of hidden neurons, sample size, and input dimension scale proportionally. While recent work has analyzed feature learning via a single step of gradient descent on the first layer weights in this regime, such one-step update schemes are fundamentally limited: the update to the weights is approximately rank-one, captures only a single direction, and requires the target function to have an information exponent of one. In this paper, we go beyond one-step updates to provide a full characterization of the features learned during the \textit{second step} of gradient descent with step-sizes $η_1\asymp N^{α_1}$ and $η_2 \asymp N^{α_2}$ for $α_1, α_2 \in [0,0.5)$, where $N$ is the number of hidden neurons. We derive a spectral characterization of the updated weights, demonstrating they behave as a spiked random matrix with multiple outliers, each corresponding to a learned direction. We show that the number of the outliers is determined by the parameters $α_1, α_2$ through $\lfloor \frac{α_2}{1/2 - α_1} \rfloor$. Furthermore, by analyzing the alignment between the learned directions and the target function, we identify a gap between training with independent versus reused batches. While independent batches restrict learning to directions with an information exponent of one, batch reuse enables the second update to capture directions even when the information exponent exceeds one, provided that $α_1, α_2$ are chosen properly. This shows that the benefits of batch reuse, previously observed in narrow-width regimes, persist in the linear-width limit as well. By characterizing these early-phase evolutions, our work proposes a tractable framework for studying optimization and feature learning phenomenology in modern overparameterized networks.

Related papers

• Mildly Overparameterized ReLU Networks on Orthogonal Data: Incremental Learning and Implicit Bias [11.187895893664484]
  We study the gradient flow dynamics of two-layer ReLU networks with training data.<n>We prove the limiting flow converges to a saddle-to-saddle jump process.<n>More broadly, our work provides the first rigorous proof of an incremental learning process for ReLU networks.
  arXiv  Detail & Related papers  (2026-05-26T14:39:20Z)
• Sharp feature-learning transitions and Bayes-optimal neural scaling laws in extensive-width networks [8.250374560598493]
  We study the information-theoretic limits of learning a one-hidden-layer teacher network with hierarchical features from noisy queries.<n>We show that a student trained with textscAdam near the effective width $k_c$ achieves these optimal scaling laws.
  arXiv  Detail & Related papers  (2026-05-11T11:39:03Z)
• Neural Networks Learn Generic Multi-Index Models Near Information-Theoretic Limit [66.20349460098275]
  We study the gradient descent learning of a general Gaussian Multi-index model $f(boldsymbolx)=g(boldsymbolUboldsymbolx)$ with hidden subspace $boldsymbolUin mathbbRrtimes d$.<n>We prove that under generic non-degenerate assumptions on the link function, a standard two-layer neural network trained via layer-wise gradient descent can agnostically learn the target with $o_d(1)$ test error.
  arXiv  Detail & Related papers  (2025-11-19T04:46:47Z)
• Beyond Softmax: A Natural Parameterization for Categorical Random Variables [61.709831225296305]
  We introduce the $textitcatnat$ function, a function composed of a sequence of hierarchical binary splits.<n>A rich set of experiments show that the proposed function improves the learning efficiency and yields models characterized by consistently higher test performance.
  arXiv  Detail & Related papers  (2025-09-29T12:55:50Z)
• Revisiting Gradient Descent: A Dual-Weight Method for Improved Learning [4.751362812627724]
  We introduce a novel framework for learning in neural networks by decomposing each neuron's weight vector into two distinct parts.<n>We show that this decomposition enhances generalization and resists overfitting, especially when training data are sparse or noisy.
  arXiv  Detail & Related papers  (2025-03-15T02:32:47Z)
• A Mean-Field Analysis of Neural Stochastic Gradient Descent-Ascent for Functional Minimax Optimization [90.87444114491116]
  This paper studies minimax optimization problems defined over infinite-dimensional function classes of overparametricized two-layer neural networks. We address (i) the convergence of the gradient descent-ascent algorithm and (ii) the representation learning of the neural networks. Results show that the feature representation induced by the neural networks is allowed to deviate from the initial one by the magnitude of $O(alpha-1)$, measured in terms of the Wasserstein distance.
  arXiv  Detail & Related papers  (2024-04-18T16:46:08Z)
• Asymptotics of feature learning in two-layer networks after one gradient-step [39.02152620420932]
  We investigate how two-layer neural networks learn from data, and improve over the kernel regime. We model the trained network by a spiked Random Features (sRF) model. We provide an exact description of the generalization error of the sRF in the high-dimensional limit.
  arXiv  Detail & Related papers  (2024-02-07T15:57:30Z)
• 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)
• Provable Multi-Task Representation Learning by Two-Layer ReLU Neural Networks [69.38572074372392]
  We present the first results proving that feature learning occurs during training with a nonlinear model on multiple tasks. Our key insight is that multi-task pretraining induces a pseudo-contrastive loss that favors representations that align points that typically have the same label across tasks.
  arXiv  Detail & Related papers  (2023-07-13T16:39:08Z)
• How Two-Layer Neural Networks Learn, One (Giant) Step at a Time [23.380148043514215]
  We investigate theoretically how the features of a two-layer neural network adapt to the structure of the target function.<n>We show that a batch size of $n = mathcalO(d)$ is indeed sufficient to learn multiple target directions satisfying a staircase property.
  arXiv  Detail & Related papers  (2023-05-29T17:43:44Z)
• Gradient Descent in Neural Networks as Sequential Learning in RKBS [63.011641517977644]
  We construct an exact power-series representation of the neural network in a finite neighborhood of the initial weights. We prove that, regardless of width, the training sequence produced by gradient descent can be exactly replicated by regularized sequential learning.
  arXiv  Detail & Related papers  (2023-02-01T03:18:07Z)
• 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)
• 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)

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.