Escape dynamics and implicit bias of one-pass SGD in overparameterized quadratic networks
| Authors | Dario Bocchi et al. |
| Year | 2026 |
| Field | Statistics / ML |
| arXiv | 2604.03068 |
| Download | |
| Categories | stat.ML |
Abstract
We analyze the one-pass stochastic gradient descent dynamics of a two-layer neural network with quadratic activations in a teacher--student framework. In the high-dimensional regime, where the input dimension and the number of samples diverge at fixed ratio , and for finite hidden widths of the student and teacher, respectively, we study the low-dimensional ordinary differential equations that govern the evolution of the student--teacher and student--student overlap matrices. We show that overparameterization (p>p^*) only modestly accelerates escape from a plateau of poor generalization by modifying the prefactor of the exponential decay of the loss. We then examine how unconstrained weight norms introduce a continuous rotational symmetry that results in a nontrivial manifold of zero-loss solutions for p>1. From this manifold the dynamics consistently selects the closest solution to the random initialization, as enforced by a conserved quantity in the ODEs governing the evolution of the overlaps. Finally, a Hessian analysis of the population-loss landscape confirms that the plateau and the solution manifold correspond to saddles with at least one negative eigenvalue and to marginal minima in the population-loss geometry, respectively.
Engineering Breakdown
Plain English
This paper analyzes how two-layer neural networks with quadratic activations learn in a teacher-student setup when trained with one-pass stochastic gradient descent in high-dimensional settings. The key finding is that overparameterization (when the student network is wider than the teacher) provides only modest speedups in escaping poor generalization plateaus—it changes the exponential decay rate's prefactor but not the fundamental dynamics. The authors derive low-dimensional ODEs governing the evolution of overlap matrices between student and teacher networks across the input dimension and sample count ratio (α = M/N) space, and they identify continuous rotational symmetries introduced by unconstrained weight norms as a critical structural property affecting learning dynamics.
Core Technical Contribution
The paper's core novelty is a rigorous theoretical characterization of one-pass SGD dynamics for quadratic neural networks in the limit where both input dimension N and sample count M diverge while maintaining fixed ratio α = M/N, with finite but arbitrary hidden layer widths (p for student, p* for teacher). Unlike prior work that either studies infinite-width limits or requires symmetric initialization, this analysis applies to realistic finite-width networks and reveals that overparameterization's benefit is quantitatively weaker than commonly assumed—it modulates only the prefactor of exponential convergence, not the timescale itself. The identification of continuous rotational symmetry as a structural property that emerges from unconstrained weight norms is novel and provides mechanistic insight into why certain plateau-escape dynamics occur. This bridges the gap between infinite-width mean-field theories and practical finite-width learning, with explicit dependence on the ratio α = M/N.
How It Works
The framework begins with a student network trying to learn a teacher network's function using quadratic activations. Both student and teacher have two layers with finite widths p and p* respectively, and receive N-dimensional input samples. At each step, the student updates weights via one-pass SGD on a batch of M samples, where M/N = α is held constant as both dimensions grow large. The authors derive closed-form ODEs that track two key quantities: the student-teacher overlap matrix (measuring alignment between student and teacher weight directions) and the student-student overlap matrix (measuring internal student weight structure). These ODEs capture how the student network gradually aligns with the teacher through learning, and crucially reveal that when p > p* (overparameterization), the overlap matrices decay toward correct values exponentially, but the escape from poor generalization plateaus is governed by exponential timescales with prefactors that only modestly improve with extra width. The continuous rotational symmetry arises because the loss depends on weights only through their norms and pairwise products, creating redundancy in the parameter space that affects convergence dynamics.
Production Impact
For engineers training neural networks at scale, this paper quantifies a counterintuitive but important insight: simply making your network wider than necessary provides diminishing returns in training speed, contrary to the intuition from infinite-width limits. When training with finite sample budgets and fixed computational capacity, this suggests that careful initialization and learning rate tuning may be more valuable than blindly overparameterizing. The identification of continuous rotational symmetries suggests that weight normalization, layer normalization, or parametrization schemes that remove these symmetries could meaningfully accelerate training—this has direct relevance to architecture design choices. The explicit dependence on α = M/N implies that for fixed model size, adjusting the batch size and sample count ratio can substantially affect convergence, providing concrete guidance for hyperparameter tuning. However, the restriction to quadratic activations and one-pass (single-epoch) SGD limits direct applicability to ReLU or attention-based models trained for multiple epochs, so results should be validated empirically on your specific architecture before redesigning training pipelines.
Limitations and When Not to Use This
This analysis is restricted to quadratic activations, which are rarely used in modern deep learning—ReLU, GELU, and other piecewise-linear or smooth activations have fundamentally different dynamics that are not captured here. The one-pass SGD assumption (single epoch) is unrealistic for practical training, which typically involves many passes over the data; extending to multi-epoch training requires analyzing how the overlap matrices evolve under repeated sampling, a substantially harder problem. The finite-width results assume the ratio α = M/N is fixed and finite, but in modern practice, α may change during training (e.g., due to data augmentation, iterative data selection, or curriculum learning), and these scenarios are not addressed. The framework assumes a well-specified teacher-student setup with matching architectures up to width; generalization to mismatched architectures, different activation functions between student and teacher, or the presence of label noise is unclear. Finally, the paper does not address how these dynamics scale when stacking more than two layers or when incorporating modern architectural components like batch normalization, skip connections, or attention mechanisms, limiting transferability to real production systems.
Research Context
This work builds on the growing literature of neural network learning theory in the high-dimensional regime, extending prior results that typically either study infinite-width limits (Jacot et al.'s Neural Tangent Kernel regime) or make restrictive symmetry assumptions. The paper advances the teacher-student framework, a fundamental model in statistical learning theory used to isolate core learning phenomena, by making it more realistic through finite widths and removing symmetry constraints. It contributes to the line of research initiated by recent work on quadratic networks and one-pass SGD in mean-field settings, but with explicit handling of finite widths and fixed sample-to-dimension ratios characteristic of modern high-dimensional statistics. The identification of rotational symmetry's role in plateau dynamics opens new research directions into how parametrization choices and explicit symmetry-breaking mechanisms (e.g., normalization schemes) can accelerate learning, potentially leading to improved training algorithms and theoretical understanding of modern architectural choices.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::