Technical Summary: Fisher-Geometric Sharpness and the Implicit Bias of SGD Toward Flat Minima

1. Problem Statement

The prevailing intuition in deep learning posits that Stochastic Gradient Descent (SGD) implicitly favors "flat minima" in the loss landscape, and that these flat minima generalize better than "sharp" ones. However, the theoretical foundation of this narrative has been critically undermined by Dinh et al. [2017], who demonstrated that standard Euclidean measures of flatness (e.g., the trace or maximum eigenvalue of the loss Hessian) are not invariant under reparametrization.

Specifically, for a neural network, one can rescale weights (e.g., $W_1 \to \alpha W_1, W_2 \to \alpha^{-1} W_2$) to preserve the network's input-output function exactly while arbitrarily altering the Euclidean sharpness. Since generalization is a property of the function, not the specific parameterization, a non-invariant sharpness measure cannot meaningfully explain generalization. While partial steps toward invariant measures have been taken (e.g., Jang et al. [2022]; Kristiadi et al. [2023]), a unified framework connecting invariant sharpness, the implicit bias of SGD, and generalization bounds has remained absent.

2. Methodology

2.1 Riemannian Sharpness ($S_R$)

The authors replace the Euclidean metric on the parameter space with the Fisher Information Matrix (FIM), $G(\theta)$, which induces a Riemannian structure intrinsic to the model's predictive distribution. They define Riemannian sharpness at a minimum $\theta^*$ as:
$$S_R(\theta^*) = \text{tr}\left( G(\theta^*)^{-1} \nabla^2 L(\theta^*) \right)$$
where $\nabla^2 L$ is the Hessian of the training loss. Unlike Euclidean sharpness, $S_R$ measures curvature relative to the natural geometry of the statistical manifold.

2.2 Theoretical Framework: SGD as a Stochastic Differential Equation (SDE)

The authors formalize mini-batch SGD as a continuous-time SDE. They introduce Assumption 3, positing that the covariance of the gradient noise is proportional to the FIM:
$$\text{Cov}[\nabla \tilde{L}(\theta^*)] = \frac{\eta}{B} G(\theta^*)$$
where $\eta$ is the learning rate and $B$ is the batch size. Under this assumption, they derive the stationary distribution $\pi^*$ of the SGD process.

2.3 Generalization Bounds

The authors connect the geometric bias to test performance using a PAC-Bayes framework. They construct a posterior distribution centered at the trained minimum and derive a generalization bound explicitly controlled by $S_R$.

2.4 Empirical Estimation

Since computing the full FIM is intractable, the authors utilize a diagonal empirical FIM estimator (sampling predictions from the model's own distribution) and Hutchinson's stochastic trace estimator for the Hessian. They acknowledge that while the true FIM guarantees invariance, the diagonal empirical estimator used in experiments inherits this invariance only approximately.

3. Key Contributions

1. Definition of Reparametrization-Invariant Sharpness: The paper defines $S_R$ and proves (Lemma 1) that it is invariant under smooth, function-preserving reparametrizations, directly addressing the critique of Dinh et al. [2017].
2. Stationary Distribution Analysis: The authors derive the stationary distribution of mini-batch SGD (Theorem 1) and prove that the probability mass is exponentially concentrated at minima with low Riemannian sharpness (Corollary 1). This establishes a rigorous implicit bias toward Riemannian-flat minima.
3. PAC-Bayes Generalization Bound: A generalization bound (Corollary 2) is established where the generalization gap is explicitly controlled by $S_R$, providing a theoretical link between Riemannian flatness and test performance.
4. Empirical Validation: Experiments on MNIST and CIFAR-10 confirm that $S_R$ tracks generalization across different optimizers (SGD, Adam), batch sizes, and learning rates in ways that Euclidean sharpness ($S_E$) does not. The scaling of $S_R$ with $\eta/B$ matches theoretical predictions.

4. Results

4.1 Invariance and Alignment

• Invariance: While the true FIM is strictly invariant, the empirical diagonal estimator is not perfectly invariant under arbitrary rescaling. However, it remains substantially more stable than Euclidean sharpness under weight rescaling.
• Noise Structure: Empirical tests (Section 5) show that the gradient noise covariance and the FIM are directionally aligned (cosine similarity $\approx 1.0$) across configurations. However, the scalar proportionality constant is not exactly $\eta/B$ in all regimes, particularly in small-batch, high-noise settings. The authors refine their assumption to allow for a scalar factor $\alpha(\theta)$, proving that the main theoretical conclusions (bias toward flat minima) remain robust even if the scalar claim is imperfect.

4.2 Scaling Laws and Generalization

• Batch Size ($B$): On MNIST, increasing $B$ leads to a significant increase in $S_R$ and a corresponding decrease in test accuracy, consistent with the prediction $S_R \propto B/\eta$. On CIFAR-10, the relationship is more complex due to optimization difficulties in small-batch regimes, but the trend of higher $B$ leading to higher $S_R$ holds.
• Learning Rate ($\eta$): Increasing $\eta$ generally decreases $S_R$ and improves generalization (up to a point where training becomes unstable).
• Optimizers:
  • SGD: Consistently finds minima with lower $S_R$ and higher accuracy in small-batch settings.
  • Adam: Achieves low $S_R$ and high accuracy on MNIST. On CIFAR-10, Adam finds minima with lower $S_R$ than SGD but slightly lower accuracy. The authors suggest this is because Adam's adaptive preconditioning changes the effective metric; the implicit bias of preconditioned optimizers is toward flatness relative to their own metric tensor.

4.3 Training Dynamics

Training trajectories show that small-batch SGD progressively escapes sharp basins and settles into flatter regions (decreasing $S_R$) as training continues. Large-batch SGD stabilizes at higher $S_R$ values, consistent with insufficient noise to escape sharp minima.

5. Significance and Claims

The paper claims to provide the first rigorous, reparametrization-invariant account of why flat minima generalize. By grounding sharpness in the Riemannian geometry of the statistical manifold (via the FIM), the authors resolve the fundamental objection that flatness is merely an artifact of parameterization.

The significance of the work lies in:

• Theoretical Resolution: It bridges the gap between the "flatness-generalization" intuition and the "reparametrization critique" by defining a sharpness measure that is intrinsic to the model's function.
• Mechanism of Implicit Bias: It provides a mathematical derivation showing that the noise structure of SGD (proportional to FIM) naturally drives the optimizer toward minima with low Riemannian sharpness.
• Practical Metric: It demonstrates that $S_R$ is a more reliable predictor of generalization than Euclidean sharpness, particularly when comparing models trained with different hyperparameters or optimizers.

Limitations Acknowledged by the Authors:

• The empirical diagonal FIM estimator is not perfectly reparametrization-invariant; exact invariance would require structured estimators like K-FAC.
• The scalar claim of the noise covariance ($\alpha = \eta/B$) is an approximation that holds directionally but varies in magnitude depending on the dataset and architecture.
• The relationship between $S_R$ and generalization is most informative when the optimizer has converged; in unstable regimes (e.g., very large learning rates), low $S_R$ does not guarantee good generalization.