Technical Summary: Flatness and Generalization in Multi-Index Models with Homogeneous Neural Networks

1. Problem Statement and Motivation

The paper addresses the long-standing heuristic in deep learning that "flat interpolators generalize well," where flatness is typically measured by the trace of the Hessian of the empirical loss. While empirical evidence suggests a correlation between flatness and generalization, theoretical counter-examples exist. Specifically, Dinh et al. (2017) demonstrated that for neural networks with homogeneous activations, rescaling symmetries allow any interpolator to be made arbitrarily sharp or flat without changing its input-output relationship or its population/empirical loss. This renders the general claim "flat interpolators generalize" vacuous, as "bad" interpolators (those with high population loss) can theoretically be made flat via symmetry.

The authors investigate whether a meaningful connection between flatness and generalization can be restored by focusing on the flattest interpolators—those achieving the orderwise minimum flatness among all interpolators. The study focuses on learning unknown multi-index models using 2-layer non-convex homogeneous neural networks. The specific questions posed are:

1. Are there "bad" interpolators that cannot be made the flattest interpolators using symmetries?
2. Do all flattest interpolators generalize?

2. Methodology and Setup

Model and Data

• Data Distribution: The true labels $y$ are generated from a multi-index model: $y = \sigma^\star(\Theta^\star x) + \xi$, where $x \sim \mathcal{N}(0, I_d)$, $\Theta^\star \in \mathbb{R}^{m^\star \times d}$ is an orthonormal matrix defining the true directions, $\sigma^\star$ is an unknown link function, and $\xi$ is label noise. The setting includes single-index models ($m^\star=1$) and sums of single-index models.
• Neural Network: A 2-layer network with $m$ hidden neurons, homogeneous activation $\sigma$, and weights $w = (a, \{\theta_j\}_{j=1}^m)$. The output is $h(w, x) = \sum_{j=1}^m a_j \sigma(\langle \theta_j, x \rangle)$.
• Overparameterization: The setting assumes $m \ge n$ (number of neurons $\ge$ number of samples), ensuring the existence of interpolators ($F_S(w) = 0$).
• Symmetry: The network uses $c$-homogeneous activations (e.g., ReLU, LeakyReLU, Quadratic), satisfying $\sigma(\alpha b) = \alpha^c \sigma(b)$. This induces a rescaling symmetry: weights can be transformed as $a_j \to \alpha_j^{-c} a_j$ and $\theta_j \to \alpha_j \theta_j$ without changing the model output.

Flatness Definition

Flatness is defined as the trace of the Hessian of the empirical loss, $\text{Tr}(\nabla^2 F_S(w))$. Due to the rescaling symmetry, the flatness of an interpolator $w$ can be minimized by optimizing the scaling factors $\{\alpha_j\}$. The authors define $\Upsilon(w)$ as the minimum flatness achievable within the equivalence class of $w$ under rescaling.

Key Technical Tools

• Convex Duality: Used to derive lower bounds on the flatness $\Upsilon(w)$ by relating it to the $\ell_p$ norm of the outer-layer weights $a$ (where $p = \frac{2c}{c+1}$).
• Concentration of Measure: Utilized to bound the eigenvalues of the activation matrix $U$ and the behavior of Sub-Weibull random variables (arising from the link functions and activations).
• Hanson-Wright Inequality: Applied to analyze quadratic forms of Sub-Weibull variables to bound the population loss and flatness.

3. Key Contributions and Results

Contribution 1: Bad Interpolators Cannot Be the Flattest

The paper characterizes a class of "bad" interpolators where the inner-layer weights $\{\theta_j\}$ are misaligned with the true directions $\Theta^\star$ (specifically, the span of $\{\theta_j\}$ has low correlation with the span of $\Theta^\star$).

• Theorem 1 & 2: The authors prove that for such bad interpolators, even after applying optimal rescaling symmetry, their minimum flatness $\Upsilon(w)$ is strictly larger than the global minimum flatness $\Upsilon^\star$.
• Quantitative Gap: Specifically, $\Upsilon(w) \gtrsim n^{\frac{c}{c+1}} \Upsilon^\star$. Since $n \ge 1$, this implies that bad interpolators are inherently "sharper" than the flattest possible interpolators.
• Implication: This establishes a necessary condition for the claim "flat interpolators generalize": symmetry cannot arbitrarily flatten all bad interpolators. There exists a fundamental barrier preventing certain bad interpolators from achieving the minimum flatness.

Contribution 2: Flattest Interpolators Generalize

The paper provides sufficient conditions under which the set of flattest interpolators (those achieving $\Upsilon(w) \approx \Upsilon^\star$) consists entirely of generalizing models.

• Assumptions: The result requires low approximation error (the link function $\sigma^\star$ is well-approximated by the network architecture) and low label noise.
• Theorem 3: Under these conditions, any interpolator $w$ that achieves the minimum flatness (up to constants) satisfies a small population loss: $F(w) - F^\star \lesssim n^{-\min\{1/2, \epsilon_1, \epsilon_2\}}$.
• Mechanism: The condition for achieving minimum flatness imposes a constraint on the outer-layer weights $a$, specifically $\|a\|_\infty \approx O(m^{-1})$. This "balancedness" property implicitly forces the network to align with the true signal structure, thereby ensuring generalization without explicitly requiring the training algorithm to find the true directions.

4. Experimental Verification

The authors empirically verify their theoretical findings using synthetic data for single-index and sum-of-single-index models with various activations (ReLU, LeakyReLU, Quadratic).

• Setup: They construct interpolators with varying degrees of alignment ($\rho^\star$) between inner weights and true directions. They solve a convex optimization problem to find the outer weights $a$ that minimize flatness for fixed inner weights.
• Results:
  • For Sum of Single-Index models, a strong correlation is observed: interpolators with low flatness consistently exhibit low population loss, regardless of the activation function.
  • For Single-Index models with ReLU, the connection is weaker due to conditioning issues, but the flattest solutions still tend to generalize.
  • Special Cases: The experiments confirm that when approximation error is high (e.g., learning a complex multi-index function with a 2-layer network that cannot represent it), the flattest interpolator fails to generalize, consistent with the theoretical necessity of low approximation error.

5. Significance and Claims

The paper claims to establish a direct link between flatness and generalization in the presence of rescaling symmetries, a setting where previous general claims were refuted.

• Refinement of the Heuristic: The authors argue that the original heuristic "flat interpolators generalize" is imprecise. Instead, the correct statement is: "The flattest interpolators generalize."
• Resolution of Symmetry Paradox: They resolve the paradox posed by Dinh et al. (2017) by showing that while symmetry can make some bad interpolators flatter, it cannot make all of them the flattest. The set of global minimizers of flatness excludes a specific class of bad interpolators.
• Scope: The results apply to a broad class of homogeneous activations (including ReLU and Quadratic) and realistic data distributions (Gaussian features, multi-index models), extending beyond specific counter-examples found in prior literature (e.g., discrete distributions or specific XOR problems).

The paper concludes that while gradient-based methods are not explicitly analyzed for convergence to these flattest minima, the theoretical landscape suggests that if such minima are reached, they possess strong generalization properties. The work highlights the importance of considering the global minimum of flatness rather than just local flatness or arbitrary flatness.