An accurate flatness measure to estimate the generalization performance of CNN models
Authors: Rahman Taleghani, Maryam Mohammadi, Francesco Marchetti
Organizations: 1*Department of Mathematics ”Tullio Levi-Civita”, University of Padova, Via Trieste 63, Padova, 35121, Veneto, Italy. · Department of Computer Science, Ruhr University Bochum, University street 140, Bochum, 44801, North Rhine-Westphalia, Germany. · Department of Mathematical Sciences and Computer, Kharazmi University, Mofateh Avenue, Tehran, 15719-14911, Tehran, Iran.
Flatness measures based on the spectrum or the trace of the Hessian of the loss are widely used as proxies for the generalization ability of deep networks. However, most existing definitions are either tailored to fully connected architectures, relying on stochastic estimators of the Hessian trace, or ignore the specific geometric structure of modern Convolutional Neural Networks (CNNs). In this work, we develop a flatness measure that is both exact and architecturally faithful for a broad and practically relevant class of CNNs. We first derive a closed-form expression for the trace of the Hessian of the cross-entropy loss with respect to convolutional kernels in networks that use global average pooling followed by a linear classifier. Building on this result, we then specialize the notion of relative flatness to convolutional layers and obtain a parameterization-aware flatness measure that properly accounts for the scaling symmetries and filter interactions induced by convolution and pooling. Finally, we empirically investigate the proposed measure on families of CNNs trained on standard image-classification benchmarks. The results obtained suggest that the proposed measure can serve as a robust tool to assess and compare the generalization performance of CNN models, and to guide the design of architecture and training choices in practice.
Figures & tables
Figure 1: Last layers framework of the model.
Figure 2: a: Illustration of the convolution operation setup. The input tensor x is convolved with a kernel K . b: Example of a convolution operation. A 2×2 kernel is applied to the input feature map. c: Example of a GAP layer. A 2×2 kernel is applied to the input feature map when stride is 2.
Figure 3: Empirical analysis of 84 models. (a) Shows the predictive power of our flatness measure on the generalization gap ( R2≈0.34 ). (b) Validates that the theoretical bound ∣S∣−4+m2(2mκ+C1+ΔC2) strictly envelopes the empirical gap.
Metric
Value
Linear regression slope
6.9049×10−6
Intercept
3.5452×10−2
p (slope)
7.244×10−9
Slope standard error
1.070×10−6
R2
0.3367
Pearson r
0.5803
Table 3: Summary of the correlation between our Flatness Measure and Generalization Gap.
Figure 4: Impact of learning rate and optimizer choice on flatness and generalization.
Figure 5: Comparison of training dynamics, relative flatness, and validation accuracy for the four best parameter sets of the evaluated model.
Figure 5
Seed
Optimizer
LR
Train BS
Val Acc
Gen Gap
Flatness
a
0
sgd_mom
0.005
64
0.9687
0.0299
119.28
b
2
sgd_mom
0.010
128
0.9677
0.0308
89.85
c
0
sgd_mom
0.005
64
0.9675
0.0310
148.91
d
2
sgd_mom
0.005
64
0.9675
0.0311
127.16
Table 4: Final results of top 4 best performing models and configuration parameters
Architecture
Pearson r
Spearman ρ
R2
VGG-16
0.076
0.118
0.006
DenseNet-121
0.491
0.443
0.242
Table 5: Correlation statistics between flatness and generalization gap across different CNN architectures. The reported values reflect the raw data without outlier filtering.
Figure 6: Effect of label noise on flatness and its predictive power. Left: Flatness increases with label noise level. Right: Spearman correlation between flatness and generalization gap, indicating flatness becomes more predictive under harder generalization scenarios.
Augmentation
Test Acc. (%)
Gen. Gap
Flatness ( κl )
Corr. ρ
None
73.5 ± 0.4
0.265 ± 0.004
2.51 ± 0.42
0.584
Basic
82.1 ± 0.4
0.179 ± 0.004
4.33 ± 0.49
0.691
AutoAugment
82.2 ± 0.6
0.178 ± 0.006
4.69 ± 0.40
0.702
Mixup
80.6 ± 1.1
0.194 ± 0.011
5.04 ± 0.41
0.725
Table 6: Impact of data augmentation on flatness and generalization. Statistical values represent the mean ± standard deviation calculated across 20 independent runs per condition.
Strategy
Epochs
Accuracy (%)
Final Flatness
Time (s)
Standard
40
79.7%
454.21
578.5
Flatness-based
100
81.6%
272.84
1418.2
Combined
100
81.6%
272.84
1422.1
Table 7: Comparison of stopping strategies (averages across 40 runs).
Figure 7: Flatness dynamics during transfer learning. Top: Evolution of symbolic flatness (log scale) across strategies. Low learning rates maintain pre-trained flatness, while layer freezing leads to significant sharpness spikes. Bottom: Correlation between final accuracy and flatness (Spearman ρ=−0.68,p<10−6 ). Each diamond represents the final state of a specific strategy.
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Cross-architecture validation of flatness-generalization correlation. Each subplot shows independently trained models with diverse hyperparameters. The columns represent linear scale (left), log-scale (center), and flatness density distribution (right). All architectures exhibit a consistent correlation between the proposed flatness metric and the generalization gap.
Figure 9: A 3×3 grid analysis of active augmentation strategies. Each row represents a specific policy (Basic, AutoAugment, Mixup), while columns provide linear, log-linear, and density-based perspectives.
Flat minima are an account of why deep networks generalise. However flatness is a matter of form (parameters), while generalisation is of function. The same function can be a result of many different parameterisations. I demonstrate this by rescaling ReLU networks, changing raw Hessian trace by up to 99 times while every prediction remains fixed. Raw curvature cannot identify a function-level explanation. Previous theoretical work traced generalisation to the weakness of constraints implied by function, meaning the freedom a model retains within the bounds of what it has learned to be correct. A policy is weaker when more future commitments remain compatible with what it has learned, allowing more freedom to adapt. To measure this for neural networks, I freeze the last hidden representation and ask whether each of 512 sampled label bundles can be met by a replacement affine classifier. The resulting joint completion score is invariant under invertible linear mixing and translation of feature coordinates. Across two predeclared cohorts of 100 networks, it predicts held-out accuracy with rank correlations 0.29 and 0.47. Raw Hessian trace and relative flatness have no multiplicity-corrected association. To put it provocatively, freedom is correlated with adaptability, while flatness is a matter of description.
Michael Timothy Bennett
School of Computing, The Australian National University, Canberra, ACT 2601, Australia
A common heuristic used to explain the generalization of first-order gradient methods on non-convex neural networks is that "flat interpolators generalize well" (Hochreiter and Schmidhuber, 1994; Keskar et al., 2017), where flatness can be measured by the trace of the Hessian of the empirical loss. However, Dinh et al. 2017) showed that, using symmetry of the network that can change flatness while keeping the population and empirical losses unchanged, any interpolator can be made sharper or flatter. This result makes the earlier heuristic statement vacuous. In this paper, we show that for learning an unknown multi-index model with 2-layer non-convex homogeneous neural networks, there is a connection between flatness and generalization, despite the existence of symmetries. This connection pertains to the "flattest" interpolators, i.e., the interpolators that have orderwise minimum flatness among all interpolators. First, we show that there exists a natural class of non-generalizing interpolators whose flatness cannot be made closer to the flattest possible, even using symmetries. Second, we show that for data generated by a sum of single-index models, if the approximation error and label noise are low, any flattest interpolator achieves small population loss, i.e., the flattest interpolators always generalize. This establishes a direct link between flatness and generalization which applies to a large class of activations and realistic data distributions.
The flatness of the loss landscape at a minimizer is a widely used heuristic for reasoning about neural-network generalization, yet evidence for this relation is mostly empirical and controversial. We study this relation in a teacher-student tree committee machine, where both the ERM estimator and the Hessian spectrum are analytically tractable in the proportional high-dimensional limit. First, we use a zero-temperature Gibbs formulation to obtain predictions for the observables of the typical minimizers of the empirical loss. Secondly, we use Edwards-Jones formalism to derive the limiting Hessian resolvent around these typical minimizers. All predictions agree with finite-size gradient-descent simulations. Finally, we study three measures of flatness, namely the left and right edges and the spectral mean, and check if a decrease in generalization error as the dataset size is increased corresponds to an increase in flatness. We find that the answer strongly depends on the learning task and on the ratio of the number of parameters to the number of data points. In regression, the spectral mean and right edge correlate with the generalization error, while the left edge does so only in the overparametrized regime. In classification this correlation reliably holds only in the highly overparametrized phase, while for underparametrized networks it can even reverse.
Brandon Livio Annesi, Davide Straziota, Enrico Maria Malatesta
Department of Quantitative Life Sciences, ICTP Trieste, Italy