SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In Figure 15, we confirm empirically that using only the first-order terms in the SAM update—which is equivalent to the gradient norm regularization  $\sum_{i=1}^{n} \|\nabla \ell_i(\pmb{\theta})\|_2$ —leads to the same effect in all key metrics including the feature rank and the number of active RLUs. This empirical validation supports the argument outlined in Proposition 1 which is based only on the first-order terms that empirically appear to be dominant for the low-rank effect of SAM.
Training statistics for SAM (ReLU activation)
Training statistics for SAM (tanh activation)
Training statistics for SAM (absolute value activation)
Figure 14: Two-layer ReLU/tanh/abs networks in the teacher-student setup. We can see that models trained with higher  $\rho$  of SAM have a smaller feature rank, smaller number of active tanh/abs units, and higher weight norm.
Training statistics for gradient regularization  $\sum_{i=1}^{n} \|\nabla \ell_i(\pmb{\theta})\|_2$  (ReLU activation)
Figure 15: Two-layer ReLU networks in the teacher-student setup. We can see that models trained with higher  $\rho$  of gradient regularization (which can be seen as the first-order approximation of SAM) have a smaller feature rank, smaller number of active ReLUs, and higher weight norm.
E Additional results for investigations of the low-rank mechanism on deep networks
In Figure 16, we show more detailed plots compared to Figure 7 regarding the behavior of feature ranks at different layers of the vision transformer on MS-COCO. First, we note that the feature rank grows over the residual blocks starting from very small values in the first blocks due to the identical positional encoding applied to all inputs at the first layer. Second, we do not observe a clear trend in the reduction of rank for larger  $\rho$  after applying attention or GeLU activations in the residual branches. This is in contrast to the ranks in the residual stream (i.e., the values right after adding the identity) after attention subblocks which, as highlighted in Figure 7 in the main part, increase less for models with higher  $\rho$ . For the sake of completeness, the absolute values of ranks in the residual stream are shown in Figure 16.
Feature ranks before and after attention/GELU in the residual branches
Feature ranks for the residual stream
Figure 16: Contrastive image-text training on MS-COCO (image encoder). A more detailed look on the feature rank trend at different blocks depending on the  $\rho$  of SAM.
Finally, in Figure 17, we show the feature ranks for the last layer of the text encoder with unfrozen and frozen text encoders. Note that a single linear layer is used on top of the frozen text features to match the dimensions of the image and text encoders within the same feature space. We conclude that even a single linear layer trained on top of the text encoder can be sufficient for the low-rank effect of SAM.
Unfrozen BERT text encoder
Frozen BERT text encoder
Figure 17: Contrastive image-text training on MS-COCO. We show the feature ranks for the last layer of the text encoder with unfrozen and frozen text encoders. Matrix multiplication denotes the last linear layer used to match the dimensions of the image and text encoders within the same feature space.