A simple connection from loss flatness to compressed neural representations

TL;DR

This paper tackles the unclear relationship between loss sharpness and generalization by reframing sharpness as a driver of local representation compression. It introduces three metrics—Local Volumetric Ratio (LVR), Maximum Local Sensitivity (MLS), and Local Dimensionality—and derives bounds that connect parameter-space flatness to feature-space geometry, extended to network-wide measures (NVR, NMLS) and reparametrization-invariant sharpness concepts. Through experiments on CNNs, MLPs, and Vision Transformers, the authors show that flatter minima constrain representation compression, with MLS showing particularly strong correlation with sharpness-related bounds, while local dimensionality behaves more independently. The dual perspective clarifies why sharpness can aid generalization in some settings but not others, and provides a practical framework to study robustness and compression across modern architectures.

Abstract

Sharpness, a geometric measure in the parameter space that reflects the flatness of the loss landscape, has long been studied for its potential connections to neural network behavior. While sharpness is often associated with generalization, recent work highlights inconsistencies in this relationship, leaving its true significance unclear. In this paper, we investigate how sharpness influences the local geometric features of neural representations in feature space, offering a new perspective on its role. We introduce this problem and study three measures for compression: the Local Volumetric Ratio (LVR), based on volume compression, the Maximum Local Sensitivity (MLS), based on sensitivity to input changes, and the Local Dimensionality, based on how uniform the sensitivity is on different directions. We show that LVR and MLS correlate with the flatness of the loss around the local minima; and that this correlation is predicted by a relatively simple mathematical relationship: a flatter loss corresponds to a lower upper bound on the compression metrics of neural representations. Our work builds upon the linear stability insight by Ma and Ying, deriving inequalities between various compression metrics and quantities involving sharpness. Our inequalities readily extend to reparametrization-invariant sharpness as well. Through empirical experiments on various feedforward, convolutional, and transformer architectures, we find that our inequalities predict a consistently positive correlation between local representation compression and sharpness.

Figures (11)

• Figure 1: Trends in key variables across SGD training of the VGG-11 network with fixed batch size (equal to 20) and varying learning rates (0.05, 0.1 and 0.2). After the loss is minimized (so that an approximate interpolation solution is found) sharpness and volumes decrease together. Moreover, higher learning rates lead to lower sharpness and hence stronger compression. From left to right: training loss, NMLS, sharpness (square root of \ref{['eq:zloss_sharpness']}), log volumetric ratio (\ref{['eq:logvol']}), MLS, and local dimensionality of the network output (\ref{['eq:dim']}).
• Figure 2: Trends in key variables across SGD training of the VGG-11 network with fixed learning rate size (equal to 0.1) and varying batch size (8, 20, and 32). After the loss is minimized (so that an interpolation solution is found) sharpness and volumes decrease together. Moreover, smaller batch sizes lead to lower sharpness and hence stronger compression. From left to right in row-wise order: train loss, NMLS, sharpness (square root of \ref{['eq:zloss_sharpness']}), log volumetric ratio (\ref{['eq:logvol']}), MLS, and local dimensionality of the network output (\ref{['eq:dim']}).
• Figure 3: We trained 100 different models for each combination of datasets and networks by varying learning rates, batch size, and random initializations. Pairwise scatter plots between MLS (resp. NMLS) and the sharpness-related bound on MLS (resp. NMLS) are shown here. For MLS (resp. NMLS) bound see \ref{['prop:mls']} (resp. \ref{['prop:nmls']}). The Pearson correlation coefficient $\rho$ is shown in the top-left corner for each scatter plot. See \ref{['subsec:corr']} for the full pairwise scatter matrix.
• Figure 4: Adaptive sharpness vs Normalized MLS for 181 ViT models and variants. Different colors represent different model classes. For most models, there is a positive correlation between Sharpness and MLS. However, outlier clusters also exist, for MobileViT mehta2022mobilevitlightweightgeneralpurposemobilefriendly models in the upper left corner, and two FastViT vasu2023fastvitfasthybridvision models in the lower right corner.
• Figure G.5: Empirical tightness of the bounds. We empirically verify that the inequalities in \ref{['eq:tightness']} hold and test their tightness. The results are shown for a fully connected feedforward network trained on the FashionMNIST dataset. The quantities A, B, C, and D are defined in \ref{['eq:tightness']}. We see that the gap between C and D is large compared to the gap between A and B or B and C. This indicates that partial sharpness $\lVert\nabla_{\mathbf{W}} f(\mathbf{x}_i, \boldsymbol{\theta}^*)\rVert_F$ (sensitivity of the loss w.r.t. only the input weights) is more indicative of the change in the maximum local sensitivity (A). Indeed, correlation analysis shows that bound C is positively correlated with MLS while bound D, perhaps surprisingly, is negatively correlated with MLS (\ref{['fig:fashion_corr_FNN']}).

Theorems & Definitions (30)

• Definition 3.1
• Lemma 3.2
• Definition 3.3
• Proposition 3.4
• Definition 3.5
• Proposition 3.6
• Definition 3.7
• Proposition 3.8
• Definition 3.9
• Proposition 3.10