Prediction Stability as a Function-Space Proxy for Flat Minima

TMLR Paper8919 Authors

Abstract: Flat minima in neural network loss landscapes are widely believed to support better generalization. However, many existing definitions of flatness are poorly specified, as they are highly sensitive to reparameterization and architectural design choices. This raises a fundamental question: is flatness truly a property of the parameter space, or does it reflect stability in the learned function itself? We adopt a function-centered perspective and examine flatness through behavior in output space rather than through parameter-space geometry. From this standpoint, meaningful flatness is expressed as stability in model predictions under controlled perturbations, independent of how the weights are parameterized. To investigate this idea, we introduce the Function-Space Flatness Proxy, an analytical probe designed for both empirical and theoretical evaluation of output-space stability. The proxy measures prediction stability under perturbations and incorporates stability-guided model selection along with a stability complexity metric. It is invariant to reparameterization and does not rely on second-order quantities such as the Hessian or Fisher information matrix. Crucially, FSFP pursues flatness indirectly: it optimizes for prediction stability in output space, and flat minima are the intended destination reached through that stability objective, via a principled route that never requires measuring curvature explicitly. Using this framework, we analyze test accuracy, loss dynamics, calibration, and negative log-likelihood to explore how output-space stability relates to generalization. Experiments on CIFAR benchmarks with ResNet architectures provide empirical validation of the probe. The observed improvements in accuracy and the divergence between absolute and normalized sharpness measures are interpreted as evidence that output-space stability yields at minima as the intended outcome of a principled stability-driven approach, rather than serving as a performance comparison against existing approaches. Overall, these findings suggest that output-space stability offers a practical and interpretable lens for studying flatness (including sharpness geometry and generalization behavior) without reducing the concept to geometric properties alone. This function-centered perspective complements parameter-space approaches and broadens the understanding of generalization beyond traditional notions of sharpness.

Submission Type: Long submission (more than 12 pages of main content)

Previous TMLR Submission Url: https://openreview.net/forum?id=VvEeuu1yMT

Changes Since Last Submission:

Since the previous submission, the manuscript has been substantially revised to address presentation and structural concerns.

- The paper has been reorganized to better align with TMLR formatting.
- Core claims, contributions, and experimental findings are now presented concisely in the main body.
- The excessive use of subdivisions and sub-subdivisions has been eliminated to improve structural coherence and flow.
- Extended mathematical derivations and detailed proofs have been moved to the appendix to improve readability.
- The overall length of the main text has been reduced by approximately 10 pages.
- Redundant explanations have been removed, and several sections were rewritten for clarity and precision.

These revisions significantly improve the clarity, compactness, and editorial quality of the manuscript while preserving all technical content.

Assigned Action Editor: ~Andres_R_Masegosa1

Submission Number: 8919