Analyzing Vision Transformers for Image Classification in Class Embedding Space

Despite the growing use of transformer models in computer vision, a mechanistic understanding of these networks is still needed. This work introduces a method to reverse-engineer Vision Transformers trained to solve image classification tasks. Inspired by previous research in NLP, we demonstrate how the inner representations at any level of the hierarchy can be projected onto the learned class embedding space to uncover how these networks build categorical representations for their predictions. We use our framework to show how image tokens develop class-specific representations that depend on attention mechanisms and contextual information, and give insights on how self-attention and MLP layers differentially contribute to this categorical composition. We additionally demonstrate that this method (1) can be used to determine the parts of an image that would be important for detecting the class of interest, and (2) exhibits significant advantages over traditional linear probing approaches. Taken together, our results position our proposed framework as a powerful tool for mechanistic interpretability and explainability research.

Paper

References (22)

Scroll for more · 10 remaining

Similar papers

Peer review

Reviewer vRFW3/10 · confidence 4/52023-07-02

Summary

The authors propose a method that explores properties of vision transformer (ViT) features. In particular, tokens of patches at various levels of the sequence of transformer blocks are projected to class-space in the case of models pre-trained for image classification. Building of this, authors provide various insights into representations spaces of ViTs using quantitative experiments.

Strengths

1) The paper explores an interesting and useful direction on better interpretability of ViTs 2) Methods proposed can be useful to community

Weaknesses

1) Novel contribution over key related work [10,12] and additional insights in the form of quantitative experimental results seem insufficient to provide any strong information on ViTs 2) Projecting intermediate features to class space with a weight matrix from the last layer is not directly meaningful - maybe explain the reasoning / assumption why this would provide any useful information? The differences in results across layers could be due to misalignment (between feature and projection weights) and not necessarily lack of information. 3) How does explicit aligning of intermediate representations to class space (e.g. see [1]) affect findings? Maybe linear projection of those features to better verify what information is contained within them? 4) Results and discussion in section 5 are not explained well [1] Naseer, Muzammal et al. “On Improving Adversarial Transferability of Vision Transformers.” ArXiv abs/2106.04169 (2021): n. pag.

Questions

The statement “we can project the internal representations of ViT into the class embedding space to probe their categorical representations” is not well supported by explanation or results (see weaknesses above for more). This is assumption is central for almost all insights provided by paper. In fact, this leads to multiple missing links in various statements made in paper (and explanations given to results).

Rating

3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

1 poor

Presentation

3 good

Contribution

2 fair

Limitations

No clear discussion of limitations or possible impacts.

Reviewer GYsr6/10 · confidence 4/52023-07-05

Summary

The authors propose to reverse-engineer pre-trained ViTs for image classification task in order to investigate how the internal representations at different levels are projected onto the class embedding space and reveal how the models construct representations for predictions. It provides insights into the distinct contributions of self-attention and MLP layers in ViTs to the categorical composition. The proposed method can further identify important image regions for class detection as a valuable toll for achieving mechanistic interpretability and explainability in ViTs.

Strengths

1. This paper presents a pioneering approach to reverse-engineering pre-trained ViTs for image classification tasks, offering new insights into how ViTs construct representations for their predictions. While the concept of reverse-engineering is inspired by NLP research, this is the first work to apply it specifically to ViTs in computer vision tasks. 2. The authors introduce a framework that enhances mechanistic interpretability and explainability in ViTs, enabling the identification of the most relevant image regions for detecting a specific class of interest. 3. The paper emphasizes the distinct roles of self-attention and MLP layers in this process, illustrating how they contribute differently to categorical updates by utilizing compositional key-value memory pair mechanisms. 4. To evaluate their findings, the authors employ several metrics including the class-value agreement score, key-value agreement rate, class similarity change rate, and match between the top-1 prediction. 5. The paper is highly accessible with clear logical flow. The tables and figures are presented in a manner that is easy to read and understand, contributing to the overall clarity of the research.

Weaknesses

1. The findings might differ when using this method to reverse engineer transformers that are larger, have been trained with different datasets, or contain architectural modifications. For example, the paper focuses on vanilla ViTs trained on ImageNet, and it is unclear how well the proposed framework generalizes to other types of ViTs or other datasets 2. The paper does not provide a comprehensive comparison with other methods for interpretability or explainability of ViTs, although it does mention some related work in this area. 3. Certain concepts require additional investigation. The reasons behind the distinct performance of block 11 compared to other blocks, as well as the rationale for summing the gradients across the blocks, remain unclear and warrant further study.

Questions

1. Why there are performance differences betwee block 11 and other blocks? 2. Why sum the gradients over the blocks? 3. What does it mean in line 307-308?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A

Reviewer 9G5D7/10 · confidence 5/52023-07-06

Summary

This work analyzes how Vision Transformers work by analyzing the representations of individual tokens (image patch representations) and how they evolve while passing through the layers of the network. The authors also show how to use their methods to devlop an interpretability method.

Strengths

Originality: The work applies experiments originaly proposed in the NLP space to analyse Transformers to the Vision Transformer. I am not aware of anyone having done these types of experiments for ViT before, and the insights gained this way are interesting. Quality: The experiments seem solid, and while I have doubts about a small part of them (see weaknesses) they are overall well done. They are well thought out and test simple hypotheses. The value of these experiments has already been validated in the NLP space. Clarity: I found the experimental design was clear and the exposition easy to follow. Significance: I think this work contributes to our understanding of how Vision Transformers work. In themselves they don't offer completely new insights, but confirm existing knwoledge/intuition/theories and add empirical evidence that aid our understanding. I think the work is easily understandable and accessible, and gives insights that help understand the inner workings of one of the currently most-used models. I think the contribution is valuable.

Weaknesses

* The paper investigates the original ViT architecture as proposed in the 2021 ICLR paper. It completely ignores that we know now that ViT does not need a CLS token (I think this was first proposed in "Scaling Vision Transformers" by the original authors of ViT in shortly after their first paper). This paper independently confirms this finding, but simply citing that paper would have been easier. Also, by analysing a ViT model that used a Global Average Pool (or MAP) to classify would remove one additional cofounder. I'd appreciate if the authors could briefly mention/discuss this point in the next version of the paper (I feel like part of the reason people keep using CLS-token ViT is that every other paper does it, so I'd encourage the authors to point out that this isn't needed any more). * A potential relevant related work is "Understanding Robustness of Transformers for Image Classification", ICCV 2021, Bhojanapalli et al.,. The 2nd half of that paper also tries to understand how ViT work by ablating parts of the modle (e.g. by removing individual Self-Attn. and MLP layers), which feels similar to the pertubation studies presented in Table 1.

Questions

No questions come to mind.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

The authors do not discuss limitations of their work, but given the nature of the paper (empirical exploration instead of proposing a new method) this does not apply as much.

Reviewer mxXe6/10 · confidence 3/52023-07-07

Summary

Inspired by recent advancements in NLP, this paper introduces a novel framework designed to reverse engineer vision transformers for the purpose of image classification tasks. The framework focuses on analyzing the internal dynamics of Vision Transformers (ViTs) within the class-embedding space, revealing the intricate process by which ViTs construct categorical representations using self-attention and MLP layers. The empirical findings gleaned from this investigation offer valuable insights into the inner workings of ViTs. Furthermore, the proposed framework can be utilized to identify the crucial components within an image that play a significant role in class detection.

Strengths

1. The research presented in this paper enhances the current understanding of interpretability and explainability in ViTs. While previous studies have primarily examined the preservation of spatial representations in ViTs throughout the hierarchy, this work specifically investigates the construction of categorical information for the final prediction. 2. This research demonstrates the intriguing phenomenon of internal disruption in categorical representations caused by context and attention perturbations. 3. The Experiment Design in this study encompasses a comprehensive range.

Weaknesses

1. It would be beneficial to include a thorough discussion in Section 3 regarding the distinctions between the proposed framework and similar work in NLP, which also involves projecting the internal representations of these models onto the output space. Including a comparison of the technical details between the two approaches would enhance the clarity of the paper. 2. The empirical results provided in the study primarily rely on the analysis of ViTs pre-trained on ImageNet. To strengthen the findings, it would be advantageous to incorporate results obtained from ViTs pre-trained on larger datasets. This expansion would offer a broader perspective and further validate the conclusions drawn in the research.

Questions

See the weakness part.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

In Section 8, the authors discuss the limitations of the adopted model's diversity and the underexplored application for performance improvement.

Reviewer 5aAZ5/10 · confidence 3/52023-07-07

Summary

This paper utilizes a pre-trained embedding matrix to elucidate the mechanism of Vision Transformers. By using the embedding matrix, inner representations at any layer can be investigated in class spaces. Specifically, it offers a visualization of how self-attention and MLP process class categorical information. This method can also be employed to visualize the saliency map.

Strengths

1. Applying the embedding matrix is straightforward. This method does not necessitate additional training and is easy to implement. 2. Overall, the paper is well-written and well-organized

Weaknesses

1. It appears that the main point of this paper is introducing the embedding matrix for empirical analysis in vision tasks and demonstrating the usefulness, since the embedding matrix has been previously explored in NLP tasks. However, the novelty and significance of the insights obtained by using the method are somewhat limited. - Non-zero linear probing accuracy can be achieved even at the early layers [a]. - Self-attentions significantly change the representations, compared to MLPs [a]. - Self-attentions and MLPs perform complementary roles. For example, self-attentions aggregate information whereas MLPs diversify it [b]. 2. Analyses are provided only for vanilla ViT. Since modern ViTs such as the Swin Transformer utilize global average pooling instead of a CLS token, the influence of these findings might be limited. [a] Raghu, Maithra, et al. "Do vision transformers see like convolutional neural networks?." *NeurIPS* (2021). [b] Park, Namuk, and Songkuk Kim. "How do vision transformers work?." ICLR (2022).

Questions

A straightforward, and potentially more accurate, method to generate a map from token space to class space involves conducting layer-wise linear probing experiments (see, e.g., Figure 13 of [a]). When compared to linear probing, does the pre-trained embedding matrix offer any advantages? One potential advantage I can foresee is that no additional learning is required. Incorporating a comparison between linear probing and the embedding matrix might improve the manuscript.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

See the Weaknesses section for technical limitations. I cannot find any ethical issue.

Reviewer 9G5D2023-08-15

Thank you for the clarifications and the additional experiments. I stand by my original review that this work should be accepted to the conference.

Reviewer vRFW2023-08-15

Response to reviewer

I thank the authors for the rebuttal. However, my two key concerns, weaknesses 1 and 2, remain unresolved. I keep my rating as it is. On weakness 2, *"This assumes that the class prototype is encoded in the embedding weights of the class-projection matrix (an assumption that has been made before, see for example [1])."*, this assumption still is vague with references to non-peer reviewed works attempting to support it.

Authorsrebuttal2023-08-16

Thank you for your feedback. We would appreciate making the concerns concrete (please see below). Regarding weakness 2, we want to clarify that our assumption (actually, an interpretation) is both theoretically and empirically supported. It is supported theoretically by the fact that the final classification layer is only reading information from the class-embedding weight matrix, which thus has to encode the patterns that allow the network to predict each class. We assume these patterns reflect the class prototypes. This interpretation is trivially true courtesy of the image classification mechanism of the transformer architecture. The decision is a function of the dot product between the [CLS] token in the hidden state of the last layer and each row (“class prototype”) in the class embedding matrix. The probability of the image belonging to a category is proportional to this dot product. Our interpretation in terms of “class prototype” follows directly from the definition of dot product, since the latter will reflect how close the [CLS] token is to a class row. Our interpretation is further supported empirically in our experiments given that the alignment of each token to the class-embedding projection is predictive of the relevance that this token has in the categorical decision, as demonstrated in our perturbation studies. More broadly, it would be useful to better understand why the reviewer’s concerns remain unresolved. In responding to weaknesses 1 and 2, in our rebuttal and revised manuscript (1) we clarified the novelty of our work by pointing out two broad contributions that are absent in the literature, and (2) we explicated the rationale of our method, in particular, why our method is informative for our research goals and why linear probes are not appropriate for such purposes (please see rebuttal for details). To improve the quality of our work, we would appreciate it if concerns with the statements made in our rebuttal of weaknesses 1 and 2 could be made concrete.

Reviewer vRFW2023-08-18

Unsupported / Weak Claims

**Weakness 2**: `Projecting intermediate features to class space with a weight matrix from the last layer is not directly meaningful` * Consider the output space of the network's last transformer block. Within this space, the weight matrix of last linear layer (operating on the outputs of that transformer block to project to class distribution space) represents class prototype embeddings. This is what the cited supporting work claims (and can be assumed true, for what matters here). * However, in a different space, this assumption would not hold. For example, taking the last linear layer weight matrix of one network, and probing the outputs of different network (with common dimension outputs) will not give the same results. In the same way, although intermediate layers have a common dimension, these vectors lie in a different space. * Claiming that projecting these vectors to class distribution space with a mismatched weight matrix results in low performance is therefore unfair. While the original vectors may not contain class information, it could also be that the weight matrix used is mismatched with the space in which these vectors lie. * In the main rebuttal, when dismissing linear probing, the authors discuss possibility of spurious background correlations affecting linear probing, but there is no reason why the learned weight matrix of last layer would not be similar. In fact, even methods like CLIP have been shown to contain such spurious correlations. This actually raises a secondary issue of how correct (absence of spurious correlations) the final layer weight matrix is, to serve as good class prototypes. More importantly, the conflicting results of linear probing (whose theoretical weaknesses are all shared by the last-layer weights method), raise further concerns. **Weakness 1**: Given the weakness in above assumption, much of the proposed insights into network weights are rendered unsubstantiated. This leaves with minimal additional contributions over the existing works such as [10, 12]. **Final Response:** I am quite aware about contemporary work in this topic and consider myself to have a decent understanding of those works. However, in light of highly positive comments by other reviewers, and given the chance that I may have misunderstood specifics of this paper, I will change my vote to Reject (as opposed to Strong Reject). I remain at Reject because the concerns discussed above have not been sufficiently addressed or discussed in the main paper or rebuttal. Further, most results figures and tables in the paper (and even rebuttal) still appear unclear (they are not described adequately and they are not self-explanatory either).

Authorsrebuttal2023-08-21

Thank you for your clarifications. We think there are some misunderstandings about what our paper does and claims. We discuss these below. > However, in a different space, this assumption would not hold. For example, taking the last linear layer weight matrix of one network, and probing the outputs of different network (with common dimension outputs) will not give the same results. In the same way, although intermediate layers have a common dimension, these vectors lie in a different space. This implies that in our framework we need to assume that the intermediate layers lie in the same space as the class-embedding space, in order to conduct later analysis. This is not the case for our paper. Our framework does not need to assume alignment: the method quantifies this alignment. Our results show that there is an increasing alignment to the class-embedding space across the hierarchy of the network. We show this progressive alignment is meaningful through perturbation studies. The focus of our paper is to understand the mechanisms in the network (namely key-value memory pair systems) which enable this increasing alignment. > Claiming that projecting these vectors to class distribution space with a mismatched weight matrix results in low performance is therefore unfair. While the original vectors may not contain class information, it could also be that the weight matrix used is mismatched with the space in which these vectors lie. We do not make such a claim in our paper. What we claim is that we found that the degree of the activation of the key-value memory pair system in MLP layer 11 (which adds more class-aligned representations to the residual stream of ViT) correlates with performance. > In the main rebuttal, when dismissing linear probing, the authors discuss possibility of spurious background correlations affecting linear probing, but there is no reason why the learned weight matrix of last layer would not be similar. In fact, even methods like CLIP have been shown to contain such spurious correlations. This actually raises a secondary issue of how correct (absence of spurious correlations) the final layer weight matrix is, to serve as good class prototypes. More importantly, the conflicting results of linear probing (whose theoretical weaknesses are all shared by the last-layer weights method), raise further concerns. We indeed agree that is likely that the class-embedding matrix also encodes spurious correlations. However, these are the spurious correlations the network is learning to associate with the class in its final decision. In other words, these are the biases the network is learning with each class prototype. Identifying these spurious associations is important since they point to which features are relevant for improving the semantic robustness of ViT. And these features are not the ones identified with vanilla linear probing, as shown in the general rebuttal. We disagree that our results being different from linear probing is problematic since they are fundamentally tackling different research questions, as explained in the general rebuttal. > Weakness 1: Given the weakness in above assumption, much of the proposed insights into network weights are rendered unsubstantiated. This leaves with minimal additional contributions over the existing works such as [10, 12]. We respectfully disagree with the reviewer on this point, for the reasons mentioned in rebuttal 1 and in response to other reviewers. > Further, most results figures and tables in the paper (and even rebuttal) still appear unclear. As stated in the first rebuttal, we will change the captions of Table 1 to contain more details about their content. We would also add better captions to the figures and work on their visibility.

Reviewer 5aAZ2023-08-16

RE: Rebuttal by Authors

I appreciate the author's effort to further generalize the discussion and clarify its novelty. However, I still believe that the novelty is limited, and I am not fully convinced by the discussion on linear probing based on ad hoc analysis. Moreover, as reviewer vRFW pointed out, it might not be clear whether the representations are aligned. Nevertheless, as the paper has improved during this rebuttal process, I am inclined to change my recommendation to borderline accept.

Authorsrebuttal2023-08-16

Thank you for your comments and for raising the rating. We respond to each in turn. > However, I still believe that the novelty is limited, To improve the quality of our work, we would appreciate if you could further clarify why the broad and novel contributions of our work (outlined in the rebuttal) over the studies mentioned in the review are insufficient. > and I am not fully convinced by the discussion on linear probing based on ad hoc analysis. Similarly, we would appreciate it if you could clarify the meaning of “based on ad hoc analysis”. In response to the original reviewers’ comments on linear probing, we have replicated the suggested linear probe analysis of [a] and compared it with our method, reporting favorable results. Specifically, we believe our results give strong evidence that linear probing does not necessarily uncover the information driving the alignment to the class-embedding representation used in the categorical prediction, as our method does (please see rebuttal for details). Moreover, in contrast to our method, it is not possible to use linear probes to investigate the mechanisms by which categorical alignment takes place. We would appreciate concrete, actionable feedback on how to strengthen our conclusions on this point. > Moreover, as reviewer vRFW pointed out, it might not be clear whether the representations are aligned. As discussed in the response to reviewer vRFW, whether representations are aligned in absolute terms or not is orthogonal to our stated research goals. We do not claim that early layers lack information because they are not aligned to class-embedding representations. Instead, our approach aims to characterize how the alignment to the class-embedding space unfolds throughout the model’s hierarchy and the mechanisms that enable it. Thus, our approach characterizes relative alignment change by means of an identifiability measure that is continuous (as opposed to discrete measures of previous work). Moreover, we show that the alignment process takes place from very early layers (as evidenced by the significantly increased alignment of inner representations to the class embedding as compared to a random model), so we do not disregard the importance of these early (and less aligned in absolute terms) representations with our method.

Reviewer mxXe2023-08-18

Post rebuttal

Thank you for spending the time addressing my initial queries with the paper. Your inclusion of the discussions on the differences from similar work in NLP and an investigation on the generalizability of the proposed framework is appreciated. Kindly ensure that these discussions are seamlessly integrated into the revision. As of now, I would like to remain my score. Looking forward to the opinions of the other reviewers.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC