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

2026-07-03

AI Summary

This work investigates the implicit bias of noisy stochastic gradient descent in infinitely wide two-layer ReLU networks. By modeling the training dynamics as a Wasserstein gradient flow via mean-field theory, the authors prove convergence to a unique stationary measure and show that the learned predictor exhibits a continuous piecewise affine structure determined by a finite arrangement of hyperplanes. Despite the infinite network width, input weights and biases align only along finitely many directions, leading to an effective collapse of the width; each such direction induces a unique ternary activation pattern over the training data, ensuring a non-redundant representation. Furthermore, the number of affine regions of the predictor is shown to be at most \(2P - 1\), where \(P\) denotes the number of linearly realizable dichotomies on the training set, revealing that model complexity is governed by the combinatorial geometry of the data.

Abstract

We study the implicit bias of noisy stochastic gradient descent in training wide two-layer ReLU networks for multivariate regression. In a mean-field regime, the training dynamics are approximated by a Wasserstein gradient flow that converges to a unique stationary measure. We characterize the structure of this stationary measure and the predictor it represents. We show that, despite the network being infinitely overparameterized, the learned predictor admits an effectively finite representation: the input weights and biases align along finitely many directions, leading to an effective width collapse. In particular, the solution function is continuous piecewise affine, with affine regions determined by the cells of a finite hyperplane arrangement. The number of learned directions, and hence hyperplanes, is bounded above by $2\mathcal{P}-1$, where $\mathcal{P}$ denotes the number of linear dichotomies realizable on the training inputs. We further establish a non-redundancy property of the learned representation by proving that each learned direction induces a unique ternary activation pattern on the training data. Consequently, the complexity of the learned predictor is governed by the combinatorial geometry of the training data.

Problem

• implicit bias
• ReLU networks
• width collapse
• stochastic gradient descent
• piecewise affine

Innovation

• implicit bias
• effective width collapse
• ReLU networks
• Wasserstein gradient flow
• piecewise affine functions