🤖 machine learning

Implicit Bias of SGD in Multivariate ReLU Networks: Effective Width Collapse

This paper demonstrates that noisy stochastic gradient descent training of wide two-layer ReLU networks for multivariate regression induces an implicit bias toward a unique, effectively finite-width predictor where neurons align along a bounded number of directions determined by the combinatorial geometry of the training data, despite the network's infinite overparameterization.

Original authors: Shuang Liang, Tom Jacobs, Guido Montúfar

Published 2026-07-07

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Shuang Liang, Tom Jacobs, Guido Montúfar

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a massive, chaotic workshop filled with thousands of workers (neurons). Each worker has a specific job: they look at a set of data points (like pictures or numbers) and decide whether to "turn on" or "turn off" based on a rule they invented. You want these workers to learn a pattern that predicts the correct answer for new data.

The paper you provided investigates what happens when you train this massive workshop using a specific method called Stochastic Gradient Descent (SGD) with a little bit of "noise" (randomness) and a rule called weight decay (which discourages workers from getting too strong).

Here is the simple breakdown of their findings:

1. The "Infinite" Workshop that Shrinks

You start with a network so wide it's practically infinite. You might expect the final solution to be a messy, complex tangle of thousands of different rules.

The Surprise: Even though you started with thousands of workers, the training process naturally forces them to collapse.

The Analogy: Imagine you have 1,000 people trying to draw a map. Instead of everyone drawing a unique, wiggly line, the training process acts like a magnet. It pulls almost everyone's drawing into just a few specific, straight lines.
The Result: The final "map" (the function the network learns) isn't a smooth, curvy blob. It becomes a piecewise affine function. In plain English, this means the final answer is made up of flat, straight-line segments joined together at sharp corners (kinks). It looks like a折线 (zigzag) in 2D or a crumpled sheet of paper in 3D.

2. The "Combinatorial" Limit

How many of these straight lines (or "kinks") does the network end up with?

The paper proves that the number of lines is not determined by how many workers you started with (which was infinite).
Instead, it is determined entirely by the geometry of your training data.
The Analogy: Think of your training data as a set of stakes driven into the ground. The number of lines the network draws is limited by how many ways you can slice the ground with a knife so that the stakes fall into different groups.
The Math: If you have $P$ ways to separate your data points with a straight line, the network will learn at most $2P - 1$ distinct directions. It's a hard cap based on the data's shape, not the network's size.

3. The "Alignment" Phenomenon

Before training, your workers (neurons) are pointing in random directions. After training, something magical happens:

The Analogy: Imagine a room full of people holding flashlights pointing in random directions. As the training progresses, the flashlights suddenly snap into alignment. They all point toward just a handful of specific directions.
The Result: The "input weights" and "biases" (the rules the neurons use) stop being unique individuals. They align along a finite number of directions. This is called Effective Width Collapse. The network effectively forgets it had thousands of neurons and behaves as if it only has a handful.

4. The "No Redundancy" Rule

The paper also found that these few remaining directions are very efficient.

The Analogy: If you have a team of experts, you don't want two experts doing the exact same job. The paper shows that every single "aligned direction" (every surviving expert) does something unique.
The Result: Each learned direction creates a distinct pattern of "on/off" signals for your training data. No two directions are redundant; no two directions are just "versions" of each other. They are all essential and distinct.

5. The Role of "Noise" and "Decay"

Why does this happen? The paper suggests it's a specific side effect (or "implicit bias") of the training algorithm:

Noise: The randomness in the training (like shaking the workshop) helps the system settle into a stable state.
Weight Decay: This is a penalty for being too "strong." It acts like a filter that prunes away unnecessary complexity.
The Outcome: Together, they push the infinite network to find the simplest possible "piecewise linear" solution that fits the data, governed strictly by the data's own geometry.

Summary

The paper claims that when you train a giant neural network with noisy SGD, the universe of infinite possibilities collapses. The network doesn't just "memorize" the data; it organizes itself into a finite, efficient structure made of straight-line segments. The complexity of this structure is dictated entirely by the shape of your data, not the size of your computer. It's as if the training algorithm is a sculptor that chips away all the excess marble until only the essential, geometrically necessary lines remain.

Technical Summary: Implicit Bias of SGD in Multivariate ReLU Networks: Effective Width Collapse

Problem Statement
Overparameterized neural networks possess infinitely many solutions that interpolate training data, yet stochastic gradient descent (SGD) consistently selects solutions with strong generalization properties. This phenomenon is attributed to the "implicit bias" of the optimization algorithm. While implicit bias is well-understood for linear models and in the "lazy" (kernel) regime of wide networks, the bias in feature-learning regimes for nonlinear networks remains incompletely understood. Specifically, there is a lack of rigorous characterization regarding how noisy SGD selects specific solutions in multivariate regression settings with ReLU activations, particularly concerning the sparsity of learned features and the geometric structure of the resulting predictor.

Methodology
The authors analyze the training dynamics of wide two-layer ReLU networks using the mean-field regime. In this framework, the network parameters are modeled as a probability measure over neurons, and the training dynamics of noisy SGD are approximated by a Wasserstein gradient flow.

Key methodological steps include:

Mean-Field Approximation: The training of a network with $N$ hidden units is described by a partial differential equation (PDE) governing the evolution of the parameter distribution $\rho_t$ . The inclusion of Gaussian noise (temperature $\beta^{-1}$ ) and weight decay ( $\lambda$ ) ensures the existence of a unique stationary measure $\rho^*$ as $t \to \infty$ .
Approximation of ReLU: Since standard ReLU neurons are unbounded and non-smooth, violating assumptions required for mean-field convergence theorems, the authors employ a bounded, smooth approximation ( $\sigma_{m,\tau}$ ) of the ReLU function. They analyze the limits as the smoothing parameter $\tau \to \infty$ and the truncation parameter $m \to \infty$ , effectively recovering the ReLU behavior.
Vanishing Temperature Limit: The study focuses on the limit where the noise level $\beta^{-1} \to 0$ (vanishing temperature), which corresponds to standard SGD with weight decay. The authors investigate the structural properties of the stationary measure and the corresponding solution function in this limit, without assuming the existence of a unique limit a priori but rather analyzing accumulation points.
Geometric Analysis: The proofs rely on analyzing the Hessian of the solution function and the geometry of the parameter space. Specifically, they utilize a "Quadratic Positivity Condition" to bound the Hessian and demonstrate that the stationary measure's support concentrates on specific rays in the parameter space.

Key Contributions and Results

Piecewise Affine Solution (Theorem 4):
The authors prove that as the network width $N \to \infty$ and noise vanishes, the solution selected by noisy SGD converges to a continuous piecewise affine function. This function is affine on the connected components of the complement of a finite set of hyperplanes.
- Complexity Bound: The number of kink hyperplanes $K$ is bounded by $K \le 2P - 1$ , where $P$ is the number of linear dichotomies realizable on the training inputs. This bound depends solely on the combinatorial geometry of the training data, not on the network width or ambient dimension.
Effective Width Collapse (Theorem 5):
The paper establishes that despite the network being infinitely overparameterized, the learned predictor admits an effectively finite representation.
- Feature Alignment: The input weights and biases of the trained network align along at most $2P - 1$ distinct directions (rays) in the parameter space.
- Collapse: Consequently, the infinite-width network collapses to an effective width of at most $2P - 1$ . The solution can be represented exactly by a shallow ReLU network with this finite number of hidden neurons.
Non-Redundancy of Features (Theorem 6):
The authors establish a strong non-redundancy property for the learned directions.
- Distinct Patterns: Each learned direction induces a unique ternary activation pattern on the training data (mapping inputs to $\{-1, 0, 1\}$ ).
- Incomparability: These patterns are incomparable under a specific partial order (where one pattern is a "degeneration" of another by turning $\pm 1$ labels to $0$). Specifically, for any binary sign pattern, at most two learned ternary patterns can degenerate from it.
Specific Characterization for Symmetric Data (Proposition 7):
For symmetric two-point datasets, the authors explicitly characterize the location of the kink hyperplanes, showing they depend on the norm of the input data and the label configuration (even vs. odd labels).

Significance and Claims
The paper claims to provide the first derivation of effective width collapse as an implicit bias of noisy SGD for multivariate regression with general training data. Unlike previous works that relied on restrictive assumptions (e.g., orthogonal or symmetric data) or analyzed static variational problems (e.g., minimum $L_1$ norm interpolation), this work derives the finite-width structure directly from the training dynamics of noisy SGD in the mean-field limit.

The results bridge the gap between static optimization characterizations and dynamic training processes, showing that the combinatorial geometry of the training data (specifically the number of linear dichotomies) governs the complexity of the learned predictor. The authors note that their analysis relies on the network being sufficiently wide to exploit the mean-field framework and requires a positive weight decay parameter ( $\lambda > 0$ ) to ensure the validity of the stationary measure analysis. They acknowledge that obtaining non-asymptotic bounds on the "failure set" (where the Hessian does not vanish) for dimensions $d > 1$ remains an open challenge compared to the univariate case.

The work extends the findings of Shevchenko et al. (2022) from univariate to multivariate settings and uniquely characterizes the structure of the stationary measure in parameter space, rather than just the function space.