SlowGuess/NeurIPS_2023_OCR huggingface dataset:
To further explore the extent to which class-labeled tokens are needed for context-labeled tokens to build a class representation, and vice-versa, we conducted a token-perturbation study in ViT-B/32. We removed either class- or context-labeled tokens from the input (after the addition of position embeddings), and measured the class-identifiability rates of the remaining image tokens in the last block. We found that in both cases the removal of information reduced to a certain extent the class identifiability scores. Concretely, the original identifiability rate of class-labeled tokens of  $71.91\%$  decreased to  $44.70\%$ , while that of context-labeled tokens decreased from  $56.24\%$  to  $38.68\%$ . On the one hand, these results suggest that class-labeled tokens need context for building more identifiable class representations. On the other hand, they show that context-labeled tokens can build class-identifiable representations without the classes themselves being depicted in the image, which suggests that ViTs have incorporated into their class prototypes contextual information. This latter finding is in line with previous studies showing that CNNs can identify a class based on context pixels only [5], leading to impoverished performance in out-of-distribution scenarios.
6 Mechanistic interpretability applications
After establishing our ability to project ViT's internal representations onto the class embedding space to investigate the development of categorical representations, this section elaborates on how this framework can examine the involvement of self-attention and MLP layers in this process. Our findings indicate that both types of layers contribute to the building of class representations through key-value memory pair mechanisms. MLP layers leverage this system to produce strong categorical updates in late blocks that are highly predictive of the model's performance, while self-attention layers promote weaker yet more disseminated and compositional updates applied earlier in the hierarchy.
Building of class representations. To investigate the extent to which self-attention and MLP layers help build categorical representations, we measured the class similarity change rate induced by these sub-modules. Given a layer with input/output tokens and a class embedding matrix, we computed the proportion of output tokens whose correct class logits increase relative to the corresponding logits for input tokens, where all logits are obtained by projecting the tokens onto the class embedding space (see Fig. 1a). Concretely, we projected the output of each layer  $\mathbf{O}^l (\mathbf{X})$  onto the class embedding space by  $\mathbf{p}_{\mathrm{out}} = \mathbf{E}\cdot \mathbf{O}^{l}(\mathbf{X})^{T}$ , and compared it to the projection of the input itself  $\mathbf{p}_{\mathrm{inp}} = \mathbf{E}\cdot \mathbf{X}^T$ . We then quantified the proportion of tokens  $i$  where  $p_{\mathrm{out}}^i > p_{\mathrm{inp}}^i$ .
We found that, although the evolution of the class similarity change rate differed across ViT variants (see appendix and Fig. 3b for ViT-B/32), some general patterns can be observed. First, self-attention layers exhibited a higher-than-chance class similarity change rate during at least some processing stages. The increment was generally higher for the [CLS] tokens, which trivially need attention mechanisms for creating a class representation given that their initial input is identical across images. In contrast, the class similarity change of MLP layers peaked in the penultimate(ish) blocks (except for the GAP variant) but was otherwise below chance level. The next sections investigate the reasons behind these peaks.
Figure 3: Results for ViT-B/32. (a) Classifiability scores of class- and context-labeled tokens; (b) Class similarity change rate induced by self-attention and MLP layers; (c) Match between the top-1 predictions of the layers and the top-1 predictions of the residual stream.
An open question from the previous experiment is how the class representations developed by the self-attention and MLP layers are incorporated into the residual stream. To measure this, we took inspiration from Geva et al. [11] and quantified the proportion of tokens whose top-1 prediction in the residual stream equaled that of the self-attention layer, that of the MLP layer, that of the residual stream before being processed by these layers, or neither of them (indicating that the residual stream has built a new composite prediction).
We found that in early blocks the predictions of the residual stream were mostly compositional, integrating information from the self-attention and MLP layers but not being dominated by their outputs (see appendix and Fig. 3c for ViT-B/32). At later stages, the residuals incorporated more directly the predictions of the sub-modules. Results showed that the influence of self-attention layers peaked in the last block. In the MLP case, the highest values were found in block 11 (or in layers closer to the last block in the case of ViT-L/16). The exception of this latter pattern was ViT-B/16 trained with GAP, where MLP influenced the residual of the last block the most. This could be due to being the only variant where the final MLP layer can influence the classification output. In the other models, the previous to last MLP layers are those from which the [CLS] token can ultimately extract information for classification (through the self-attention layer in the last block).
Categorical updates. To investigate how self-attention and MLP layers carry out categorical updates by promoting class-prototype representations, we projected their output parameter matrices  $\mathbf{W}_{\mathrm{out}}$  of eq. 2 and eq. 3) onto the class-embedding space, by computing  $\mathbf{P}_{W_{\mathrm{out}}} = \mathbf{E} \cdot \mathbf{W}_{\mathrm{out}}^T$ . The rows of  $\mathbf{E}$  and  $\mathbf{W}_{\mathrm{out}}$  were normalized to unit length to enable comparison across ViT variants. This projection measures the extent to which each row in the output parameter matrices reflects the classes encoded in the embedding space, and thus probes whether the value vectors of the key-value memory pair system have high correspondence with class prototype representations. We were interested in evaluating the extent to which class prototypes had at least one memory value in a given layer  $l$  that resembles their representation. Hence, we extracted the maximum value of each row in  $\mathbf{P}_{W_{\mathrm{out}}}^l$  that denotes the highest similarity score obtained per class prototype in a given layer  $l$ . We call this measure class-value agreement score.
We found that, in deeper blocks, the class-value agreement scores of self-attention and MLP layers were significantly higher than those of a model initialized with random weights (Fig. 4a and appendix). This could be due to the model developing more complex and semantically meaningful keys at deeper stages. Results also showed that the peaks in MLP layers were higher than those of self-attention layers, indicating that MLP sub-modules promote stronger categorical updates. In addition, a manual inspection of the MLP layers containing high class-value agreement scores revealed that the value
vectors may promote and cluster semantically similar and human-interpretable concepts (see appendix for examples). These patterns were observed across all variants except for ViT-GAP. Given that the later model averages the representations of the image tokens for the classification task, we hypothesize that in ViT-GAP categorical representations might be distributed across value vectors instead.
Figure 4: Key-value memory pair mechanisms. (a) Class-value agreement scores, which measure the extent to which the layers promote class prototype representations; (b) Key-value agreement rates, which quantify the proportion of tokens promoting the correct class using key-value memory pair systems.
Key-value memory pairs at inference time. To investigate if self-attention and MLP layers act as key-value memory pair systems at inference time, we measured how the keys that are most activated during the processing of a sample correspond to the value vectors that promote the representation of the correct class. Concretely, for each layer, we quantified the proportion of tokens where the 5 keys with the highest coefficients are associated with value vectors whose top-5 logits (obtained by the projection  $\mathbf{E} \cdot \mathbf{W}_{\mathrm{out}}$ ) indexed the correct class. We call this metric key-value agreement rate.
As shown in Fig. 4b, in earlier blocks, the key-value agreement rates of self-attention layers were higher than those of MLP layers. The agreement rate of MLP layers peaked in the penultimate blocks, while self-attention layers peaked in the last blocks. The rates on these peaks were generally higher for self-attention layers. These findings indicate that, while self-attention layers promote weaker categorical representations (see previous section), the updates are carried out in more image tokens than in MLP layers. Of note, similarly to the class-value agreement score results, ViT-GAP exhibited a limited use of key-value memory pair systems.
We further evaluated the influence of these mechanisms on the accuracy of the model, and compared the agreement rate of correctly classified vs. misclassified samples. We found that, for all variants excluding ViT-GAP, the agreement rates in the penultimate MLP layers of correctly classified samples were significantly higher (see appendix). In addition, in the majority of ViTs, a smaller yet significant difference was also found in the deeper self-attention layers. These findings suggest that the use of key-value memory pair mechanisms has a meaningful effect on the performance of the model.
Compositionality of key-value memory pair mechanisms. The final output of a layer may combine the prediction of many key-value memory pairs, and predict a class distribution that is different from that of its most activating memories. Inspired by Geva et al. [11], we measured the compositionality of the key-value memory pair system by quantifying the number of instances where a layer's final predictions matched any of the predictions of the top-5 most activated memories for that instance.
Results showed that for all variants except ViT-GAP, the penultimate MLP layers have low compositionality scores: between  $50\%$  and  $70\%$  of instances had a final prediction that matched the prediction of one of the most activated memories (see appendix). Self-attention layers also decreased their compositionality in the last blocks, but their scores were still higher than MLP layers: more than  $70\%$  of instances had a composite prediction (except for ViT-CIFAR, see appendix).
7 Explainability application
In this section, we show that our framework can also be used to identify the parts of an image that would be the most important for detecting a class of interest, at each processing stage of ViT.
Having demonstrated that class representations develop gradually over the ViT hierarchy, and that image tokens differentially represent categorical information, we propose to use a gradient approach to quantify how much an image token of a given block would contribute to form a categorical representation in the [CLS] token. In detail, our explainability method can be applied as follows:
1. For the  $j^{\mathrm{th}}$  image and the  $b^{\mathrm{th}}$  block, project the hidden representation of the [CLS] token after the attention layer  $\mathbf{x}_{\mathrm{cls}}^{b_{\mathrm{attn}}}$  onto the class embedding space, and compute the cross-entropy loss  $L_{\mathrm{CE}}$  of the class of interest  $c_{j}$ , such that  $\ell_j^b = L_{\mathrm{CE}}(\mathbf{E} \cdot \mathbf{x}_{\mathrm{cls}}^{b_{\mathrm{attn}}}, c_j)$ . This projection quantifies how close is the representation of the [CLS] token to the prototype representation of the class of interest.
2. Compute the gradient of  $\ell_j^b$  with respect to the attention weights  $\mathbf{a}_j^b$  that the [CLS] tokens assigned to the image tokens in an attention head in the self-attention layer, such that  $\nabla \ell_j^b = -\partial \ell_j^b / \partial \mathbf{a}_j^b$ . Since we are interested in how the image tokens decrease the cross-entropy loss, we negate the gradients. In simple terms, this step estimates the rate at which an image token would increase the correct class representation of the [CLS] token if more attention were allocated to it.
The final output will consist of importance scores assigned to every image token of a block and attention head that can be visualized as a heatmap over the image (see Fig. 5a). Of note, the block- and attention-head-specific visualizations of our explainability framework differ from other widely used methods that generate global relevancy maps, for example by computing the gradients of the class logits at the final layer with respect to the input and/or aggregating the information propagated across layers up to the class embedding projection (see [6] for an overview). Instead, our method can visualize the categorical information contained in the image tokens independently for each block and attention head. This allows us to (1) better understand how categorical information is hierarchically built, and (2) characterize the importance of each block in building the class representations.
In addition, our method is class-generalizable and can be used to identify the features that would lead to the assignment of different class labels (see Fig. 5). Thus, it can be used to uncover the parts of an image that produce an incorrect prediction or trigger different predictions in a multi-label classification task. Moreover, since our framework identifies the image tokens that would increase the categorical representations of the [CLS] token, its results can be compared to the actual attention weights assigned by this type of token, to shed light on what aspects of an image the network neglects or emphasizes in relation to the class-optimal ones.
Figure 5: Examples of feature importance visualization in ViT-B/32. (1) Image example where the correct label is "Impala", but the model predicts "Ibex". (2) Image example with multiple labels.
Besides providing block- and attention-head-specific visualizations, our framework can also be used to generate a traditional global relevancy map. Concretely, we can aggregate the gradients over the blocks and attention heads, by  $\sum_{b} \nabla \ell^{b}$ , and obtain a final feature importance map that takes into account the relevance of each image token at every stage of processing (see Fig. 5b). The sum procedure also allows us to corroborate that we obtain a fair portrayal of the individual contribution of each block and attention head to the final class representation.
To validate the quality of our global relevancy maps, we compared our approach with an established explainability method [6]. Testing ViT-B/32, we separately applied both methods and: (1) Quantified
the importance of each image token; (2) Gradually removed the tokens with the least to most importance scores (negative perturbation test); (3) Gradually removed the tokens with the most to least importance (positive perturbation test); (4) Measured the accuracy of the model with each removal; (5) Computed the Area Under the Curve (AUC) of the final accuracies. Our results showed that our framework yields similar results to those of Chefer et al. [6] (see appendix), highlighting the adequacy of our approach. In addition, we found that we could maintain, and even improve, the accuracy of ViT-B/32 with the removal of up to  $60\%$  of the least important tokens. More generally, the AUC of our perturbed model was significantly higher (neg. perturbation) and lower (pos. perturbation) than that of a model whose token removal order was randomly determined.
8 Comparison with a linear probing approach
Linear probing is a widely used interpretability tool that consists of training linear classifiers to uncover the features represented in the inner layers of a model. Similarly to the goals of our interpretability framework, to shed light on the relevant factors behind category-building processes in neural networks, linear probing is sometimes used to predict from internal layers the classes represented in the output projection matrix [1]. However, previous studies (e.g. [19]) have shown that this approach can highlight features that are not relevant to the model's task. In other words, a successful probing result can reflect the separability of feature representations that are ignored by the model in the class embedding space. In contrast, our framework directly quantifies how the inner representations of ViT increasingly align with the representations encoded in the class prototypes.
To substantiate the disparity in insights derived from both methods, we conducted the following experiment. Reproducing the linear probing approach taken by Raghu et al. [18], we trained separate 10-shot linear classifiers on ImageNet-S for each token position and layer of a ViT-B/32. To test if the information learned with these probes sheds light on the categorical decisions taken by the network, we conducted negative and positive perturbation tests. Concretely, we quantified the class identifiability scores obtained from the linear probes for each token, gradually removed those with the least to most identifiable scores (for negative perturbation; vice-versa for positive perturbation), and measured the accuracy of the model with each removal. We compared these results with those of our framework's positive and negative perturbation experiments reported in the previous section. We found that even if linear probes could generally decode with better top-1 accuracy the classes in some of the inner layers, the obtained scores do not significantly predict the relevance of the image tokens in the categorical decision of the network (see appendix).
Notably, when used for the same goals, our framework (1) is more time-efficient: it comprises a one-forward pass on validation images, while linear probes additionally involve a one-forward pass over the training images and the fitting of a linear classifier for every token position and layer; and (2) can be used in low-resource scenarios, where the training of additional models is difficult due to small datasets or reduced computing capabilities.
9 Conclusions
With the increasing use of ViTs in the field of computer vision, it is necessary to build methods to understand their inner mechanisms and explain their predictions. Our work introduces an intuitive framework for such purposes that does not require optimization and provides interesting, human-interpretable insights into these networks. Concretely, we demonstrated how our method can extract class representations from the hidden states of image tokens across the hierarchy of ViT, providing insights into the category-building processes within these networks. Additionally, we used our framework to elucidate the distinct roles of self-attention and MLP layers in this process, revealing that they promote differential categorical updates that partly depend on key-value memory pair mechanisms. Lastly, we emphasized the utility of this method for explainability purposes, aiding in the identification of the most pertinent parts of an image for the detection of a class of interest.
Limitations. This work only studies ViTs trained for image classification. Future work could investigate how to adapt our framework to examine the inner representations of models with other types of output embeddings. In addition, we did not explore how our approach might be used for model editing or performance improvement purposes. Some mechanistic interpretability insights gained in this work point to aspects of ViT that could be manipulated for these goals in future studies.
10 Acknowledgments
This project was partly funded by the German Research Foundation (DFG) - DFG Research Unit FOR 5368. We are grateful to access to the computing facilities of the Center for Scientific Computing at Goethe University, and of the Ernst Strungmann Institute for Neuroscience.
[1] G. Alain and Y. Bengio. Understanding intermediate layers using linear classifier probes. arXiv preprint arXiv:1610.01644, 2016.
[2] N. Belrose, Z. Furman, L. Smith, D. Halawi, I. Ostrovsky, L. McKinney, S. Biderman, and J. Steinhardt. Eliciting latent predictions from transformers with the tuned lens. arXiv p reprint arXiv:2303.08112, 2023.
[3] S. Bhojanapalli, A. Chakrabarti, D. Glasner, D. Li, T. Unterthiner, and A. Veit. Understanding robustness of transformers for image classification. In Proceedings of the IEEE/CVF international conference on computer vision, pages 10231-10241, 2021.
[4] G. Brunner, Y. Liu, D. Pascual, O. Richter, M. Ciaramita, and R. Wattenhofer. On identifiability in transformers. arXiv preprint arXiv:1908.04211, 2019.
[5] B. Carter, S. Jain, J. W. Mueller, and D. Gifford. Overinterpretation reveals image classification model pathologies. Advances in Neural Information Processing Systems, 34:15395-15407, 2021.
[6] H. Chefer, S. Gur, and L. Wolf. Transformer interpretability beyond attention visualization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 782-791, 2021.
[7] G. Dar, M. Geva, A. Gupta, and J. Berant. Analyzing transformers in embedding space. arXiv preprint arXiv:2209.02535, 2022.
[8] A. Y. Din, T. Karidi, L. Choshen, and M. Geva. Jump to conclusions: Short-cutting transformers with linear transformations. arXiv preprint arXiv:2303.09435, 2023.
[9] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 2020.
[10] S. Gao, Z.-Y. Li, M.-H. Yang, M.-M. Cheng, J. Han, and P. Torr. Large-scale unsupervised semantic segmentation. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2022.
[11] M. Geva, R. Schuster, J. Berant, and O. Levy. Transformer feed-forward layers are key-value memories. In Empirical Methods in Natural Language Processing (EMNLP), 2021.
[12] M. Geva, A. Caciularu, K. Wang, and Y. Goldberg. Transformer feed-forward layers build predictions by promoting concepts in the vocabulary space. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics, 2022.
[13] A. Ghiasi, H. Kazemi, E. Borgnia, S. Reich, M. Shu, M. Goldblum, A. G. Wilson, and T. Goldstein. What do vision transformers learn? a visual exploration. arXiv preprint arXiv:2212.06727, 2022.
[14] A. Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
[15] M. Naseer, K. Ranasinghe, S. Khan, F. Khan, and F. Porikli. On improving adversarial transferability of vision transformers. In International Conference on Learning Representations, 2022.
[16] M. M. Naseer, K. Ranasinghe, S. H. Khan, M. Hayat, F. Shahbaz Khan, and M.-H. Yang. Intriguing properties of vision transformers. Advances in Neural Information Processing Systems, 34:23296-23308, 2021.
[17] N. Park and S. Kim. How do vision transformers work? arXiv preprint arXiv:2202.06709, 2022.
[18] M. Raghu, T. Unterthiner, S. Kornblith, C. Zhang, and A. Dosovitskiy. Do vision transformers see like convolutional neural networks? Advances in Neural Information Processing Systems, 34:12116-12128, 2021.
[19] A. Ravichander, Y. Belinkov, and E. Hovy. Probing the probing paradigm: Does probing accuracy entail task relevance? arXiv preprint arXiv:2005.00719, 2020.
[20] T. Ridnik, E. Ben-Baruch, A. Noy, and L. Zelnik-Manor. Imagenet-21k pretraining for the masses. arXiv preprint arXiv:2104.10972, 2021.
[21] X. Zhai, A. Kolesnikov, N. Houlsby, and L. Beyer. Scaling vision transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12104-12113, 2022.