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

AI summary

Overview

• Research area: Statistical machine learning theory — specifically, generalization theory for overparameterized neural networks, implicit regularization, and the training dynamics of gradient descent in the Edge-of-Stability (EoS) regime.
• Technical level: Advanced. The paper is a theory contribution built on half-space depth (Tukey depth), metric entropy, weighted path norms, and empirical process arguments.
• Scope: The paper theoretically characterizes how data geometry — summarized by a quantity the authors call data shatterability — controls the implicit regularization and generalization of two-layer ReLU networks whose parameters satisfy a Below-Edge-of-Stability (BEoS) curvature condition.

What This Paper Is About

Overparameterized neural networks have far more capacity than they need to memorize their training data, yet gradient descent often finds solutions that generalize well. This paper asks which data distributions lead to that good outcome and which lead to memorization, by analyzing the generalization gap of any two-layer ReLU network that satisfies the BEoS condition λ_max(∇²ℒ(θ)) ≤ 2/η (Definition 2.1), without assuming the solution is optimal or stationary. The authors' answer is that the deciding factor is how easily the data distribution can be partitioned into many disjoint small regions by ReLU half-spaces — data that is hard to shatter is implicitly regularized strongly, while data that is easy to shatter (such as data on the unit sphere) permits memorization.

Key Contributions

1. A spectrum of generalization bounds for isotropic data. For a one-parameter family of isotropic Beta(α)-radial distributions — defined by radial profile h(r) = 1 − (1 − r)^(1/α) with R ~ Uniform[0,1] and U ~ Uniform(S^{d−1}) — the authors derive generalization upper bounds (Theorem 3.4) and matching lower bounds (Theorem 3.5) that depend smoothly on the radial concentration parameter α. As α decreases, mass moves toward the boundary and the guarantee degrades.

2. Explicit spherical interpolation / memorization result. In the limiting case where the support collapses to the unit sphere, the authors construct width-K ≤ n networks that interpolate the data while satisfying the stability condition (Theorem 3.6), with λ_max(∇²ℒ) ≤ 1 + (D² + 2)/n, or ≤ (D² + 2)/n if the output bias β in Equation (1) is removed.

3. Provable adaptation to intrinsic low-dimensionality. Under a mixture-of-subspaces assumption where inputs lie on a union of m-dimensional balls in R^d with m < d, the authors prove that all BEoS-stable solutions enjoy a generalization rate Õ(n^{−1/(2m+4)}) depending on the intrinsic dimension m rather than the ambient dimension d (Theorem 3.10), with the dependence on the number of mixture components at most polynomial. Synthetic experiments confirm gradient descent behaves according to this intrinsic-dimension law.

4. A unifying "data shatterability" principle and a new proof technique. The authors introduce a half-space-depth quantile partition that bypasses global metric-entropy control, decoupling the analysis into a "good" T-deep region where implicit regularization enforces complexity control and a "bad" shallow region controlled only by its probability mass.

Main Findings

• The BEoS condition bounds a data-dependent weighted path norm. Proposition 2.2 states that for any θ ∈ Θ_BEoS(η, 𝒟), ‖f_θ‖_{path,g_𝒟} ≤ 1/η − 1/2 + (R + 1)√(2ℒ(θ)). The weight function g_𝒟(u, t) = min{g̃_𝒟(u, t), g̃_𝒟(−u, −t)}, with g̃_𝒟(u, t) = P(Xᵀu > t)² · E[Xᵀu − t | Xᵀu > t] · √(1 + ‖E[X | Xᵀu > t]‖²), measures how hard it is for gradient descent to place a ReLU ridge with normalized direction u and threshold t.

• A generic geometry-driven generalization decomposition. With high probability, the worst-case gap over Θ_BEoS(η, 𝒟) is bounded by the sum of two terms (Equation 6): Õ(P_X(X ∉ Ω_T(𝒫_X))) for the shallow region, plus Õ(g_𝒫^min(T)^{−d/(2d+3)} n^{−(d+3)/(4d+6)}) for the T-deep region, where g_𝒫^min(T) is the minimum weight-function value among ReLU boundaries intersecting the T-deep region Ω_T(𝒫_X).

• Beta(α)-radial specialization. For isotropic Beta(α)-radial distributions, the shallow mass scales as P_X(𝔸_ε^d) ≍ ε^α and the regularization strength as g_𝒫(1 − ε) ≍ ε^{d+2α} (Lemma G.2 and Proposition G.6), reducing Equation (7) to Õ(ε^α) + Õ(ε^{−d(d+2α)/(2d+3)} n^{−(d+3)/(4d+6)}) before optimizing over ε.

• Two-case upper bound with a threshold at α = 3d/(2d − 3). Theorem 3.4 gives, with probability at least 1 − δ, rates of n^{−α(d+3)/(2(d²+4αd+3α))} when α ≥ 3d/(2d − 3) and n^{−α/(2d+4α)} when α < 3d/(2d − 3), with prefactors involving the step size η and M := max{D, ‖f_θ|{𝔹₁^d}‖{L^∞}, 1}. The authors note a two-case discussion is necessary because the empirical process only guarantees sup |g_𝒟(u,t) − g_𝒫(u,t)| ≤ Õ(n^{−1/2}).

• A matching lower bound. Theorem 3.5 shows that over distributions with marginal 𝒫_X(α) and labels supported on [−1, 1], the minimax expected generalization gap over the class Θ_g(𝔹₁^d; 1, 1) is ≳_{d,α} n^{−2α/(d−1+2α)}. Smaller α means more mass in the annulus 𝔸_ε^d and more disjoint spherical caps carrying comparable probability mass, increasing memorization capacity.

• The spherical limit interpolates perfectly within the stability regime. Theorem 3.6 constructs width-K ≤ n interpolating networks for distinct inputs on 𝕊^{d−1} with λ_max(∇²_θℒ) ≤ 1 + (D² + 2)/n (or ≤ (D² + 2)/n without β). Remark 3.7 notes that removing the hidden bias b and assuming a uniform spherical distribution yields a bound of approximately Õ(n^{−1/4}), compatible with the rates in Wu and Su (2023).

• The concentration index captures shatterability. Definition 3.8 introduces the half-space-depth quantile function Ψ_{𝒫_X}(T) = P(depth(X, 𝒫_X) ≥ T) and the concentration index S_DQ(𝒫_X) = (∫0^{1/2} Ψ{𝒫_X}(T) dT)^{−1}. Large index (small α) implies boundary-concentrated mass that sits in pre-existing packing slots and facilitates shattering; small index (large α) starves those regions of data. For distributions on a sphere, depth is zero everywhere, the quantile function collapses, and the index diverges — matching Theorem 3.6.

• Neural shattering on the uniform ball is one special point of the spectrum. The "neural shattering" phenomenon identified by Liang et al. (2025) for the uniform ball distribution is shown to be a single special case of the broader generalization spectrum the paper uncovers.

• Global metric entropy is infinite, so distribution-agnostic bounds are impossible. The authors state that the function class induced by the EoS data-dependent regularity has infinite L^∞ metric entropy, and that the spherical interpolation result shows nontrivial distribution-agnostic bounds cannot hold. This motivates the partition-based technique.

• Anisotropic data: intrinsic depth governs, but the index becomes conservative. For a finite mixture 𝒫_X = Σ_j π_j 𝒫_{X,j}, the law of total probability gives Ψ_{𝒫_X}(T) = Σ_j π_j P(depth(X, 𝒫_X) ≥ T | Z = j), and depth(x, 𝒫_X) ≥ π_j · depth(x, 𝒫_{X,j}) for each fixed j. On low-dimensional structures such as lines, high-dimensional spherical caps degenerate into intervals or knots (Figure 2a), so the concentration index becomes a conservative estimate of difficulty rather than a precise one.

• Upper bounds are not tight, and the authors say why. Remark 3.9 explains that the analysis exploits only the area of a single inscribed rectangle under T ↦ Ψ_{𝒫_X}(T), because choosing the optimal T* in Equation (6) is equivalent to picking the largest such rectangle and discarding the remaining area under the curve.

Methodology in Plain English

The authors do not track the full trajectory of gradient descent. Instead, they characterize the set of parameter vectors that could plausibly be reached and maintained by training: those satisfying the BEoS curvature condition λ_max(∇²ℒ(θ)) ≤ 2/η. They then relate this curvature constraint to a weighted path norm whose weights are computed from the training data itself, which turns the question of implicit regularization into a concrete, data-dependent complexity measure.

The central difficulty is that this complexity measure is extremely uneven across the input space. Under standard covering-number arguments, the function class has infinite metric entropy, so no uniform, distribution-independent bound is possible. The authors' workaround is a geometric partition. Using Tukey's half-space depth, they identify a T-deep region where every hyperplane through a point leaves at least a T-fraction of data on each side. Any ReLU whose activation boundary crosses that region must therefore have activation probability at least T, which forces the data-dependent weight function to be bounded below by a positive quantity g_𝒫^min(T). Inside the deep region, the network therefore behaves like an ordinary unweighted path-norm class with radius proportional to (g_𝒫^min(T))^{−1}, for which standard generalization bounds apply. Outside the deep region, where regularization is weak, the authors abandon function-space covering entirely and simply bound the error by the worst-case L^∞ amplitude times the probability mass of that region. This ties the final error directly to the shape of the data distribution.

They then evaluate the resulting trade-off on two families of distributions. For isotropic distributions, depth depends only on radius, so the deep region is a ball of radius 1 − ε and the trade-off reduces to the mass in the ε-annulus versus the regularization strength at offset 1 − ε. For Beta(α)-radial distributions, both quantities are computed in closed form (ε^α and ε^{d+2α}), and optimizing over ε yields the stated rates. Matching lower bounds are constructed by placing localized ReLU atoms on disjoint spherical caps in the shallow shell, where the weight function is small, so that these neurons can accumulate large path norm at little cost and form families of hard-to-distinguish networks. For anisotropic data on mixtures of low-dimensional balls, the same principle is applied using the intrinsic dimension of each subspace rather than the ambient dimension.

Why This Matters

• Impact on research. The paper provides end-to-end, distribution-dependent generalization bounds for the EoS/BEoS regime, moving beyond distribution-agnostic uniform-convergence arguments that the authors argue cannot apply here. It abstracts the neural-shattering analysis of Liang et al. (2025) for the uniform ball into a general depth-based framework, and it offers a theoretical justification for phenomena reported by Zhang et al. (2017) — for instance, why real data are harder to overfit than random Gaussian data. Conceptually, it inverts the classical VC-dimension viewpoint: VC dimension describes a model's capacity to shatter arbitrary data, while "data shatterability" describes the feasibility of shattering a specific dataset by the gradient-descent-trained network.

• Real-world applications. The paper connects its theory to two practical techniques:

  • Mix-up data augmentation (Zhang et al., 2018; Zhang et al., 2021): the authors state they provide theoretical insight into how it works.
  • Activation-based pruning (Hu et al., 2016; Ganguli and Chong, 2024): likewise cited as a technique their results illuminate.
  • Understanding dataset-dependent overfitting behavior, such as why models fit random labels easily but resist overfitting natural data — relevant to diagnosing when memorization is likely.
  • Architecture and training-regime choices that interact with the EoS condition, since the bounds depend on the learning rate η and on the data-dependent regularity.

  The paper does not report experiments beyond stating that synthetic experiments confirm gradient descent follows the intrinsic-dimension law.

• Industry relevance. The results suggest that data geometry, not just model size or explicit regularization, determines effective capacity, which is directly relevant to practitioners choosing learning-rate schedules (the bounds depend explicitly on η) and to those reasoning about when large models will memorize versus generalize. The intrinsic-dimension result implies the relevant complexity lives on the data manifold rather than in the ambient input dimension.

Future Directions

• Tightening the upper bounds. Remark 3.9 identifies exactly where slack enters: the analysis exploits only the largest inscribed rectangle under the depth-quantile curve Ψ_{𝒫_X}. Closing the gap between the upper bounds and the lower bounds of Theorem 3.5 would require using the full area under that curve.

• Extending beyond the isotropic and mixture settings. The concentration index S_DQ(𝒫_X) is described as a precise proxy for shatterability in the isotropic case but only a conservative estimate for anisotropic structured data, where partitioning directions are constrained. A sharper anisotropic index would be a natural next step.

• Broadening the architecture and dynamics scope. The authors state their framework is situated in the feature-learning ("rich") regime of overparameterized networks and analyzes BEoS solutions without assuming optimality or stationarity. Appendix B.2 is described as containing only informal, heuristic discussion of gradient dynamics, leaving a full dynamical account open.

• Connecting more directly to overfitting phenomena and augmentation. The paper claims new theoretical insight into mix-up augmentation and activation-based pruning; turning those connections into quantitative predictions would be a follow-on direction.

Target Audience

This paper is aimed at researchers in statistical machine learning theory and deep learning theory — particularly those working on generalization bounds, implicit regularization, the edge-of-stability phenomenon, and neural network approximation theory. It will be most useful to readers comfortable with high-dimensional probability, half-space depth, path norms, and empirical process arguments. The high-level "data shatterability" message and the intrinsic-dimension result are accessible to a broader machine learning audience, but the theorems, exponents, and proof strategy require an advanced theoretical background.

Authors’ 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 overparametrized 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.