Generalization Below the Edge of Stability: The Role of Data Geometry

TL;DR

The paper addresses how data geometry governs generalization of overparameterized two-layer ReLU networks trained below the edge of stability. It introduces data shatterability as a unifying principle and develops three main results: (i) generalization bounds that adapt to intrinsic low-dimensionality via mixtures of subspaces, (ii) a spectrum of generalization bounds for isotropic Beta-radial data controlled by α, and (iii) a flat interpolation construction on spherical data showing memorization is possible under BEoS. The approach combines a data-dependent weighted path norm framework with BEoS, yielding bounds that reflect data geometry rather than ambient dimension, supported by experiments on synthetic data and real datasets (e.g., MNIST). These findings illuminate when implicit regularization via stability aligns with robust generalization and suggest directions for leveraging data structure in designing architectures and training regimes with improved generalization and controlled memorization. The work provides a principled link between data shatterability, optimization dynamics, and representation learning, with implications for understanding generalization in deep learning beyond conventional capacity-control arguments.

Abstract

Understanding generalization in overparameterized neural networks hinges on the interplay between the data geometry, neural architecture, and training dynamics. In this paper, we theoretically explore how data geometry controls this implicit bias. This paper presents theoretical results for overparameterized two-layer ReLU networks trained below the edge of stability. First, for data distributions supported on a mixture of low-dimensional balls, we derive generalization bounds that provably adapt to the intrinsic dimension. Second, for a family of isotropic distributions that vary in how strongly probability mass concentrates toward the unit sphere, we derive a spectrum of bounds showing that rates deteriorate as the mass concentrates toward the sphere. These results instantiate a unifying principle: When the data is harder to "shatter" with respect to the activation thresholds of the ReLU neurons, gradient descent tends to learn representations that capture shared patterns and thus finds solutions that generalize well. On the other hand, for data that is easily shattered (e.g., data supported on the sphere) gradient descent favors memorization. Our theoretical results consolidate disparate empirical findings that have appeared in the literature.

Figures (8)

• Figure 1: Visualization of Data Shatterability.(a) Knots as learned representations: When data lies on a mixture of lines (black lines), the decision boundaries of ReLU neurons (translucent planes) reduce to a finite set of knots (colored dots). (b & c) More caps v.s. Less caps: Red dots depict data points and dashed lines represent neuron boundaries that partition the data support into non-overlapping caps (shaded regions), each containing equal probability mass.
• Figure 2: How data geometry controls generalization.(a) Union of $J=20$ lines ($m=1$) embedded in $\R^d$ with $d\in\{10,50,100,500\}$. The regression slopes remain nearly constant across different $d$, showing that generalization adapts to intrinsic rather than ambient dimension. (b) Fixed ambient dimension $d=5$ with isotropic Beta-radial distributions (Definition \ref{['def:isotropic_beta_distribution']}) for $\alpha\in\{1,5,10,50\}$. Larger $\alpha$ yields steeper slopes in the log--log error curve, consistent with improved rates as probability mass concentrates away from the boundary.
• Figure 3: Neuron activation statistics under different geometries. (a) On the uniform sphere, most neurons fire on less than $10\%$ of the data, indicating highly specialized ReLUs as we predict in Theorem \ref{['thm:gap_lower_bound']} and Theorem \ref{['thm:flat_interpolation']}. On the low-dimensional mixture, many neurons fire on $10$-$40\%$ of the data, reflecting broader feature reuse. (b) Scatter of weight coefficient magnitude versus activation rate. On the sphere, GD produces many low-activation neurons with large coefficients. On the low-dimensional mixture, neurons spread to medium activation rates with moderate coefficients.
• Figure 4: Data geometry and memorization on MNIST.Left panel: Comparison of training curves under the same ground-truth predictor with Gaussian inputs versus MNIST inputs ($n=30000$). GD on the Gaussian data set quickly interpolates, while MNIST resists overfitting for tens of thousands of steps. Right panel: Prediction error against Tukey depth for MNIST samples. Shallow points (low depth) exhibit larger errors. This region refers to "highly shatterable region".
• Figure 5: Training curves on different geometries. Training loss and clean MSE on Sphere vs. Low-dimensional mixture. We can see GD on sphere interpolate very quickly (before the 2000-th epoch) while the mixed low-dimensional data resist to overfitting.

Theorems & Definitions (71)

• Definition 2.1: Below Edge of Stability qiao2024stable
• Proposition 2.2
• Theorem 3.2: Generalization Bound for Mixture Models
• Definition 3.3: Isotropic Beta-radial distributions
• Theorem 3.5
• Theorem 3.6: Generalization Gap Lower Bound
• Theorem 3.7: Flat interpolation with width $\le n$
• Remark 3.8
• Definition D.1
• Definition D.2