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.