SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Maksym Andriushchenko
maksym.andriushchenko@epfl.ch
Dara Bahri
dbahri@google.com
Nicolas Flammarion
Sharpness-aware minimization (SAM) is a recently proposed method that minimizes the sharpness of the training loss of a neural network. While its generalization improvement is well-known and is the primary motivation, we uncover an additional intriguing effect of SAM: reduction of the feature rank which happens at different layers of a neural network. We show that this low-rank effect occurs very broadly: for different architectures such as fully-connected networks, convolutional networks, vision transformers and for different objectives such as regression, classification, language-image contrastive training. To better understand this phenomenon, we provide a mechanistic understanding of how low-rank features arise in a simple two-layer network. We observe that a significant number of activations gets entirely pruned by SAM which directly contributes to the rank reduction. We confirm this effect theoretically and check that it can also occur in deep networks, although the overall rank reduction mechanism can be more complex, especially for deep networks with preactivation skip connections and self-attention layers. We make our code available at https://github.com/tml-epfl/sam-low-rank-features.
Understanding generalization and features learned by overparametrized deep networks is the key topic of modern machine learning. The training objective of deep networks typically has many global optima where the training points are perfectly fitted (Zhang et al., 2017), but very different features and generalization performance are obtained (Chizat et al., 2019; Liu et al., 2019). It has been observed recently that the sharpness of a minimum—i.e. how quickly the loss can change in some neighborhood around a minimum in the parameter space—correlates with the generalization error (Keskar et al., 2017; Jiang et al., 2020). The idea of minimizing the sharpness during training to improve generalization has motivated recent works that propose to use worst-case perturbations of the weights on every iteration of training (Foret et al., 2021; Zheng et al., 2021; Wu et al., 2020). We refer to this family of methods collectively as sharpness-aware minimization (SAM) and focus on the version proposed in Foret et al. (2021) that uses a single step of normalized gradient ascent to approximate the worst-case weight perturbation.
Contributions. In our paper, we are interested in shedding more light on the effect of SAM on the structure of the features learned by the network. In particular, our main observation is that:
Sharpness-aware minimization on overparametrized neural networks leads to low-rank features.
While standard training techniques can already produce features of reduced rank (Huh et al., 2021), SAM significantly amplifies this effect. This effect is of interest when the goal is to identify a low-dimensional structure in the data, as well as for more efficient feature quantization and nearest neighbor retrieval based on the learned features.
We make the following contributions in our paper:
- In Section 3, we present extensive empirical evidence of low-rank features for various models (ResNets, ViTs, MLP-Mixers) trained with SAM on four classification tasks (CIFAR-10/100, Tiny ImageNet, ImageNet-1k) as well as for contrastive text-image training (MS-COCO).
- In Section 4, we provide a mechanistic understanding of how low-rank features arise in a simple two-layer ReLU network. We observe that a significant number of activations gets entirely pruned by SAM, which directly contributes to the rank reduction. Furthermore, we provide a theoretical argument supporting this effect of SAM.
- In Section 5, we confirm that this effect can also occur in deep networks, although the overall rank reduction mechanism can be more complex, especially for deep networks with preactivation skip connections and self-attention layers.
- Finally, in Section 6, we show that directly inducing low-rank features does not improve generalization on natural data. Thus, we conclude that SAM is unique in its ability to both improve generalization and decrease the feature rank in a data-adaptive fashion.
2 Related work and background knowledge on SAM
In this section, we first present the most relevant previous works and then cover background on SAM.
2.1 Related work
Sharpness-aware minimization. The idea behind SAM introduced by Foret et al. (2021) is to minimize the sharpness during training to improve generalization. SAM modifies stochastic gradient descent (SGD) such that on every iteration of training, the gradient is taken not at the current iterate but rather at an approximate worst-case point in its vicinity. Zheng et al. (2021) concurrently propose a similar weight perturbation method which also successfully improves standard generalization on multiple deep learning benchmarks. Wu et al. (2020) also propose an almost identical algorithm with the same motivation, but with the focus on improving robust generalization of adversarial training. Chen et al. (2022) discover that SAM is particularly helpful to improve generalization for new architectures like vision transformers (Dosovitskiy et al., 2021) and MLP-Mixers (Tolstikhin et al., 2021). Andriushchenko and Flammarion (2022) study the reasons behind the success of SAM and characterize its effect on simple diagonal linear networks. Bartlett et al. (2022) and Wen et al. (2023) theoretically analyze the regularization effect of SAM close to a minimum.
Implicit sharpness minimization. There are multiple works that suggest that sharpness can also be minimized implicitly by standard training algorithms. For example, Cohen et al. (2021) formulate the dependency between the learning rate used for training and the sharpness given by the maximum eigenvalue of the Hessian. Barrett and Dherin (2021) and Smith et al. (2021) derive a gradient regularization term that is related to SAM by analyzing SGD and GD with finite step sizes. Damian et al. (2021) suggest a connection between label-noise SGD which implicitly minimizes the trace of the Hessian and SAM whose variations can also minimize the same quantity (Wen et al., 2023). Mulayoff et al. (2021) and Nacson et al. (2023) suggest a relation between sharpness and smoothness for two-layer networks depending on the learning rate of gradient descent.
Low-rank and sparse features. The most related work to ours is the one by Ramesh et al. (2022). They briefly note that contrastive language-image pretraining (CLIP) (Radford et al., 2021) with SAM leads to a reduction of the feature rank which they leverage for more efficient feature quantization. Huh et al. (2021) observe that existing training techniques such as SGD on neural networks can already produce features of reduced rank. On a similar note, Ethayarajh (2019) and Cai et al. (2021) observe an extreme anisotropy of the feature space for language models trained with standard methods, including an interesting low-rank cluster structure. Galanti et al. (2022) and Timor et al. (2023) discuss a low-rank effect in the weight matrices due to a small weight norm which is achieved either with weight decay or with a small initialization scale. Na et al. (2022) suggest that models trained with SAM are more compressible via weight pruning and quantization. Gulcehra et al. (2022) observe that
using SAM with dropout can lead to a lower rank in reinforcement learning tasks. Andriushchenko et al. (2023) study the effect of large step size SGD on the Jacobian of the network and activation sparsity. Li et al. (2023) observe that standard training leads to extreme activation sparsity in the MLP blocks of transformers.
2.2 Background on sharpness-aware minimization
Let  $\mathcal{I}_{\mathrm{train}} = \{1,\dots ,n\}$  be the indices of the training set  $\{\pmb {x}_i,y_i\}_{i = 1}^n$  and  $\ell_i(\theta)$  be the loss of a model parametrized by weights  $\pmb {\theta}\in \mathbb{R}^{|\pmb {\theta}|}$  and evaluated at point  $(\pmb {x}_i,\pmb {y}_i)$ . Foret et al. (2021) propose to minimize the following worst-case objective instead of standard average loss minimization:
$$
\text {S A M o b j e c t i v e :} \min  _ {\boldsymbol {\theta} \in \mathbb {R} ^ {| \boldsymbol {\theta} |}} \underset {\mathcal {I} \subset \mathcal {I} _ {\text {t r a i n}}} {\mathbb {E}} \max  _ {\| \boldsymbol {\varepsilon} \| _ {2} \leq \rho} \frac {1}{| \mathcal {I} |} \sum_ {i \in \mathcal {I}} \ell_ {i} (\boldsymbol {\theta} + \boldsymbol {\varepsilon}), \tag {1}
$$
where  $\mathcal{I}$  is a random subset of  $m$  training points. We note this objective is based on the maximization of the sum of losses over batches of  $m$  points each. To make SAM practical, Foret et al. (2021) propose to minimize the objective (1) with stochastic gradients. Denoting the batch indices at time  $t$  by  $\mathcal{I}_t$ , this leads to the following update rule on each iteration of training:
$$
\text {S A M} \quad \theta_ {t + 1} := \theta_ {t} - \frac {\gamma_ {t}}{\left| \mathcal {I} _ {t} \right|} \sum_ {i \in \mathcal {I} _ {t}} \nabla \ell_ {i} \left(\theta_ {t} + \frac {\rho_ {t}}{\left| \mathcal {I} _ {t} \right|} \sum_ {j \in \mathcal {I} _ {t}} \nabla \ell_ {j} \left(\theta_ {t}\right)\right). \tag {2}
$$
We note that the same batch  $\mathcal{I}_t$  is used for the inner and outer gradient steps, and  $\rho_t$  typically includes the gradient normalization, i.e.,  $\rho_t \coloneqq \rho / \left\| \frac{1}{|\mathcal{I}_t|} \sum_{j \in \mathcal{I}_t} \nabla \ell_j(\pmb{\theta}_t) \right\|_2$ . The worst-case perturbations and the use of a small  $m$  in SAM are essential for the generalization improvement which depends continuously on  $m$  as noted in Foret et al. (2021).
3 Experimental evidence of low-rank features due to SAM
We first discuss how we measure the feature rank and then present strong empirical evidence for low-rank features for models trained with SAM. We consider first classification tasks on CIFAR-10, CIFAR-100 (Krizhevsky and Hinton, 2009), Tiny ImageNet (Le and Yang, 2015), and ImageNet-1k (Deng et al., 2009), and then contrastive learning on MS-COCO (Lin et al., 2014).
How we measure the feature rank. Consider a neural network  $f: \mathcal{X} \to \mathbb{R}^{D}$  which maps inputs from a set  $\mathcal{X}$  (e.g., images) to a  $D$ -dimensional vector (e.g., logits for classification tasks or embeddings for contrastive learning). We assume that  $f$  can be decomposed on  $B$  blocks, i.e.,  $f(\pmb{x}) = f_{B} \circ f_{B-1} \circ \ldots \circ f_{1}(\pmb{x})$  where a single block  $f_{b}: \mathbb{R}^{d_{b-1}} \to \mathbb{R}^{d_{b}}$  can consist of multiple layers. By features we refer to the intermediate values  $f_{b}(\pmb{x}) \circ \ldots \circ f_{1}(\pmb{x}) \in \mathbb{R}^{d_{b}}$  at some block  $b$ . To assess their rank, we first do a PCA on the feature matrix  $[f_{b}(\pmb{x}_{i})^{\top}]_{i=1}^{d_{b}} \in \mathbb{R}^{n \times d_{b}}$  and then select the minimal number of principal components that span 99% of the variance in the data. We refer to this measure as the feature rank, as it captures the notion of the most informative features in an interpretable manner. We note that this is not equivalent to the matrix rank in a strict mathematical sense. However, it provides a practical alternative to selecting a fixed threshold on singular values, which can often be challenging to determine.
3.1 Low-rank features for ResNets on standard classification tasks
Setup. We train a PreAct ResNet-18 (He et al., 2016b) with standard augmentations on standard deep learning datasets: CIFAR-10, CIFAR-100, and Tiny ImageNet. We consider both (1) a minimal setting with a small learning rate, no momentum, and no weight decay, and (2) a state-of-the-art setting which includes large learning rates, momentum, and weight decay. Since we observe that neural collapse—convergence of the feature rank of the penultimate layer to the number of classes (Papyan et al., 2020)—occurs in a cascading fashion (Hui et al., 2022) and interferes with the low-rank trend of SAM, we exclude the last residual superblock and instead report the rank at the third superblock.
Observations. We plot the main metrics for these three datasets in Figure 1, 2, and 3. We observe that the models trained with SAM achieve a substantially smaller feature rank which occurs not
Further details can be found in our code https://github.com/tml-epfl/sam-low-rank-features.
Minimal setting: small learning rate, no momentum, no weight decay
State-of-the-art setting: large learning rate, momentum, weight decay
Figure 1: ResNet-18 on CIFAR-10. SAM improves test error (left), leads to more generalizable features (middle), and noticeably reduces the feature rank at the intermediate ResNet block (right). Note that the test error improvement is U-shaped in  $\rho$  unlike the rank reduction which is monotonic.
Figure 2: ResNet-18 on CIFAR-100. SAM improves test error (left), leads to more generalizable features (middle), and noticeably reduces the feature rank at the intermediate ResNet block (right). Note that the test error improvement is U-shaped in  $\rho$  unlike the rank reduction which is monotonic.
only at the end but also throughout the training. We also note that the generalization improvement is  $U$ -shaped in  $\rho$  (i.e., too large  $\rho$  is harmful) unlike the rank reduction which is monotonic. We note that for the state-of-the-art setting, the rank exhibits a sudden drop at the beginning and then a gradual increase, both for standard training and SAM. We verified that such behavior originates from the usage of initial large learning rates and is not specific to SAM. Overall, the rank reduction effect is very consistent over the three datasets and over both minimal and state-of-the-art settings. We also observe that augmentations play a crucial role in revealing the low-rank trend with respect to the  $\rho$  of SAM, whereas the addition of only weight decay is insufficient for revealing a similar trend. Additionally, to confirm the generalizability of the low-rank features taken at an intermediate layer, we also include transfer learning performance using k-nearest neighbors classification with  $k = 10$ . For this purpose, we compute (1) the k-NN classification error on CIFAR-10 for models trained on
State-of-the-art setting: large learning rate, momentum, weight decay
Figure 3: ResNet-18 on Tiny ImageNet. SAM improves test error (left), leads to more generalizable features (middle), and noticeably reduces the feature rank at the intermediate ResNet block (right). Note that the test error improvement is U-shaped in  $\rho$  unlike the rank reduction which is monotonic.
CIFAR-100 and (2) the k-NN classification error on CIFAR-100 for the models trained on CIFAR-10 and Tiny ImageNet. Without this experiment, one could assume that since these features are not from the penultimate layer, they can be of limited use for downstream tasks. However, this experiment highlights that it is not the case: the block-3 features are still suitable for transfer learning and show completely non-trivial performance.
Additional experiments. We report a few more related experiments in Appendix B. First, we note that the feature rank shows the same trend on both training and test sets and is not sensitive to the chosen variance threshold (see Figure 9 and 10 in Appendix). Moreover, the same picture holds also for different batch sizes (see Figure 11 in Appendix). Additionally, we also report results for more recent variants of SAM such as ASAM from Kwon et al. (2021) and GAM from Zhang et al. (2023) (see Table 2 in Appendix). Finally, we discuss the behavior of the feature rank at the penultimate layer instead of the intermediate ResNet block (see Figure 12 in Appendix).
3.2 Low-rank features for ViTs and MLP-Mixers on ImageNet-1k
Setup. We take publicly available models trained on ImageNet-1k from the vision_transformer library<sup>2</sup> which contains models from Dosovitskiy et al. (2021), Tolstikhin et al. (2021), and Chen et al. (2022). We evaluate the feature rank on 12800 examples using the zeroth token for ViTs and average features over all tokens for MLP-Mixers.
Observations. We report the results in Table 1. We note that neural collapse is not an issue here since the feature dimension is smaller than the number of classes on ImageNet-1k. We see a very consistent rank reduction trend from SAM for each setting, with the only exception of MLP-Mixers after  $3/4$  of the total number of blocks. Finally, we note that these models are trained with rather small  $\rho$  of SAM (e.g., 0.2 for ViT-B/16), and we can expect a stronger low-rank effect for higher  $\rho$ .
3.3 Low-rank features in contrastive language-image training on MS-COCO
Setup. Here we use a training setting similar to CLIP (Radford et al., 2021) but we start from pretrained image and language models. We fine-tune the  $\mathrm{R + Ti / 16}$  vision transformer from Steiner et al. (2021) and BERT from Devlin et al. (2018) on MS-COCO using the InfoNCE contrastive loss (Oord et al., 2018). We add a linear head to each of the encoders to match the dimension of both pre-trained encoders. We note that for contrastive training, unlike for classification, neural collapse is
$^{2}$ https://github.com/google-research/vision_transformer
Table 1: ViTs and MLP-Mixers on ImageNet-1k. We compute the feature ranks for publicly available models from the vision_transformer library.
Training	Block	ViT-B/16	ViT-B/32	ViT-L/16	ViT-L/16	Mixer-B/16
Standard	Last	680	672	903	898	695
SAM	Last	617	654	820	844	620
Standard	After 3/4 blocks	467	484	595	595	301
SAM	After 3/4 blocks	426	440	314	442	477
Standard	After 1/2 blocks	412	390	387	469	425
SAM	After 1/2 blocks	346	362	250	318	369
Figure 4: Contrastive image-text training on MS-COCO. SAM improves the retrieval error (top) and reduces the feature rank at the last layer of both image and text encoders (bottom).
not an issue, so we report the feature rank directly at the last layer. We measure the retrieval error within each batch of size 128.
Observations. In Figure 4, we observe that SAM both substantially improves the retrieval error and leads to features of much smaller rank. We note that having features of reduced rank may contradict the common intuition that low rank is typically harmful in contrastive learning (Chen and He, 2021; Hua et al., 2021), however, we are still in the regime that the remaining dimensions are expressive enough to preserve (and even improve) the retrieval performance. In addition, we note that the low-rank features are observed both in the image and text encoders suggesting that this effect of SAM is not specific to image data. Moreover, we observe the low-rank effect also when the text encoder is frozen during fine-tuning (Figure 13 in Appendix). Thus, even a single linear layer on the side of text features can be sufficient to get the low-rank effect.
4 How SAM can induce low-rank features: insights from simple models
In this section, we dissect the source of the low-rank effect of SAM on two-layer neural networks. We first provide clear empirical evidence that this effect can be attributed to zero activations and then confirm it theoretically.
4.1 Mechanistic understanding of low-rank features for a two-layer ReLU network
Setup. We consider a ReLU network  $f(\pmb{\theta}) = \langle \pmb{a}, \sigma(\pmb{W}\pmb{x}) \rangle$  trained with the squared loss and parametrized by  $\pmb{\theta} = [\mathrm{vec}(\pmb{W}), \pmb{a}]$  where  $\pmb{W} \in \mathbb{R}^{m \times d}$  and  $\pmb{a} \in \mathbb{R}^m$ . We use the teacher-student
Generalization improvement and low feature rank
Understanding the low-rank effect
Figure 5: Two-layer ReLU networks in the teacher-student setup. The models trained with a higher  $\rho$  of SAM generalize better, have a significantly smaller feature rank, smaller number of active ReLUs, and higher  $\ell_2$  weight norm. Note that the generalization improvement is U-shaped in  $\rho$  unlike the rank reduction which is monotonic.
setup with 3 teacher neurons,  $m = 100$  student neurons, and input dimension  $d = 3$  inspired by the setup of Chizat et al. (2019). The goal of the student network is to recover the teacher network from a finite set of training points which are sampled from the Gaussian distribution and labeled by the teacher network. We use SAM for the first  $50\%$  of iterations and then disable it to achieve full convergence as we observed that running SAM, especially with high  $\rho$ , hinders convergence. To obtain a smooth trend of rank reduction and clear separation between different  $\rho$ , we exceptionally use the PCA variance threshold of  $99.99\%$ , which is higher than the  $99\%$  used in all other experiments.
Low-rank features are due to zero activations. We provide a mechanistic understanding of how low-rank features emerge in Figure 5. First, we observe that even in this simple setting, SAM improves generalization and learns features of substantially smaller rank (4 instead of 19). We investigate the origin of low-rank features by examining the number of active ReLU units, which noticeably decreases with increasing  $\rho$ . Although sparse activations do not necessarily imply low-rank features (e.g., the identity matrix is sparse but full-rank), we observe that SAM sets entire activations to zero during training. For the largest  $\rho = 0.6$ , only 10 out of 100 neurons are activated on at least one training example, and even fewer of them contribute to the  $99.99\%$  PCA variance in the data. Similar results are shown for other activations, such as tanh and absolute value (see Figure 14 in Appendix), where the low-rank effect can be similarly attributed to zero activations.
SAM  $\neq$  weight decay. Interestingly, SAM reduces the number of active RLUs not by shrinking to zero the weight vectors associated with ReLU units, but by leveraging the negative part of ReLU to effectively prune neurons. As suggested in Figure 5 (right), the weight norms only increase with  $\rho$ , and the effect of SAM is opposite to that of weight decay. Thus, theoretical results demonstrating that weight decay leads to low-rank features (Galanti et al., 2022; Timor et al., 2023) are not directly applicable to studying the low-rank effect of SAM. This observation highlights the difference between SAM and classical regularization methods and explains why we cannot achieve the same regularization effect with weight decay. To better understand this phenomenon for this simple model, we need to find a theoretical argument specific to SAM.
4.2 Theoretical argument for low-rank features in a two-layer ReLU network
Let  $\ell(\pmb{\theta})$  be the loss on example  $(x, y)$  sampled by SGD and assume the squared loss for concreteness. We formulate our theoretical argument in the following proposition.
Proposition 1. Every update of SAM contains a component that decreases all pre-activation values  $\{\langle w_j,x\rangle \}_{j = 1}^m$  of a two-layer ReLU network trained with the squared loss by a non-negative amount equal to  $\eta \rho \sqrt{\ell(\theta)} /\| \nabla f(\theta)\| _2\sigma (\langle w_j,x\rangle)\| x\| _2^2$
Proof. We have for the update of SAM:
$$
\nabla \ell \left(\boldsymbol {\theta} + \rho \frac {\nabla \ell (\boldsymbol {\theta})}{\| \nabla \ell (\boldsymbol {\theta}) \| _ {2}}\right) = \nabla \ell (\boldsymbol {\theta}) + \rho \nabla^ {2} \ell (\boldsymbol {\theta}) \frac {\nabla \ell (\boldsymbol {\theta})}{\| \nabla \ell (\boldsymbol {\theta}) \| _ {2}} + \mathcal {O} \left(\rho^ {2}\right) = \nabla \left[ \ell (\boldsymbol {\theta}) + \rho \| \nabla \ell (\boldsymbol {\theta}) \| _ {2} + \mathcal {O} \left(\rho^ {2}\right) \right].
$$
Thus, under the first-order approximation, a step of SAM corresponds to a gradient update on the regularized objective  $\ell (\pmb {\theta}) + \rho \| \nabla \ell (\pmb {\theta})\|_{2}$ . Now recall that the layerwise gradients of a two-layer
ReLU network can be written as: