Analyzing Vision Transformers for Image Classification in Class Embedding Space
Authors: Martina G. Vilas
Timothy Schaumlöffel
Gemma Roig
Venue: NeurIPS 2023 poster (2023)
Keywords: transformers
computer vision
image classification
mechanistic interpretability
explainability

Abstract: 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.

Analyzing Vision Transformers for Image Classification in Class Embedding Space
Martina G. Vilas
Timothy Schaumlöffel
Gemma Roig
Goethe University Frankfurt
Goethe University Frankfurt The Hessian Center for Al
Goethe University Frankfurt The Hessian Center for AI
Ernst Strüngmann Institute
Abstract
Despite the growing use of transformer models in computer vision, mechanistic understanding of these networks still needed. This work introduces a method to reverse-engineer Vision Transformers trained solve image classification tasks. Inspired by previous research in NLP, we demonstrate how the inner representations any level of the hierarchy can be projected space to uncover how these networks build categorical representations for their pre- dictions. We use our framework to show how image tokens develop class-specite representations that depend on attention mechanisms and contextual information, and give insights on how self-attention and MLP layers differentially contribute 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 powerful tool for mechanistic interpretability and explainability research.
the learned class embedding
Introduction
Transformer-basee deep neural networks have become one of the most popular architectures in machine learning due to their remarkable performance and adaptability to multiple domains. Conse- quently, recent work has focused reverse-engincerring the inner mechanisms of these networks for better understanding and control their predictions (e.g. 22 11 12 18]). Of these, a series of studies in the field of NLP [2] language transformers by projecting the hidden states and parameters of intermediate processing layers onto vocabulary space using the pre-trained output embedding matrix. This approach has enabled the analysis in human-unde uuman-understandable terms of how next-word predictions are built internally, introducing simple and efficient method for the mechanistic interpretability of NLP transformers.
[2] have shed light the predictive mechanisms of large
So far, to the best of our knowledge, no work has shown that a similar approach can be applied to reverse engineer Vision Transformers (ViTs) for image classification. We thus introduce method to characterize the categorical building processes of these networks by projecting their intermediate representations and parameter matrices onto the class embedding space. Our framework quantifies how the inner representations increasingly align with the class prototype learned by the model (encoded by the class projection matrix), and uncovers the factors and mechanisms behind this alignment.
In this work, we use our method to show that (1) image tokens increasingly align to class prototype representations, from early stages of the model; (2) factors such as attention mechanisms and contextual information have role in this alignment; and (3) the categorical building process partly relies on key-value memory pair mechanisms differentially imparted by the self-attention and MLP layers. In addition, we discuss how to use this framework to identify the parts of an image that would be the most informative for building a class representation. Finally, we show that our method can
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

characterize the emergence of class representations in image tokens more efficiently and accurately than the commonly used linear probing approach.
2 Related work
Prior research has demonstrated that within VIT models, (I) the use of linear probing approach in the hidden states of class tokens ([CLS) in early layers enables the decoding of class representations [18. and (2) the image tokens in the final block contain class-identifiee we introduce a methodological framework that adeptly extracts categorical information from image tokens early in the processing stages, and allows to elucidate the underlying factors and mechanisms that facilitate the development of these class representations. We also demonstrate that, unlike our method, linear probes do not necessarily uncover the features relevant to the classification task.
information. By contrast,
We use our framework to complement previous findings on the inner mechanisms of ViTs. Ghiasi et al. [3 and Raghu et al. [18] investigated how tokens preserve input spatial representations across the hierarchy. We instead analyzed how tokens increasingly represent the output classes. addition, recent work has examined and compared the role of self-attention and MLP layers: Raghu et al. [18] quantified how residual connections are differentially influenced by these sub-modules, Bhojanapalli al. 3 measured their distinctive importance in the model's performance by running perturbation studies, and Park and Kim [17] assessed their statistical properties. We supplemented these findings by carrying out a detailed analysis of how these sub-modules build class representations by exploiting mechanisms like key-value memory pair systems. Finally, previous studies B [3] [6] analyzed how the performance of VIT holds against multiple nuisance factors. We alternatively inspected if class representations in image tokens are internally disrupted by context and attention perturbations.
Besides work probing the inner mechanisms of ViTs, tools for providing human-interpretable expla- nations of the network's predictions in particular instances have been developed (6] We introduce an explainability method that follows this line of work and uncovers the parts of an image that favor the formation of meaningful class representations, independently for any block.
Interpretation of vision transformers mechanisms in class embedding space
(a)
caf
(b)
3G
Block b
MHSA
V1
V2
1.5
MLP
MLP
k
A
k
MHSA
MHSA
0.7
X
[CLS]
Figure I: Schematic our framework. (a) The hidden states of image tokens x.. in block are projected onto the class embedding space using the output embedding matrix E. (b) A key-value memory pair system in self-attention layer. Key vectors ky belong to different attention heads hy. The match between the hidden states an and the keys ky is weighted by attention values to produce a memory coefficient. Value vectors weighted by these memory coefficients are summed and added to the residual stream. Adapted from Geva et al.
VIT architecture. As introduced by Dosovitskiy et al. (9] and depicted Fig. Th vanilla image- classification ViTs take as input sequence of linear projections of cqual-sized image patches

---

with added position embeddings, We refer to these as image tokens (xI....,x,) with x E IR.. The sequence also includes a special "class token" " whose initial representation is learned during training. denoted [CLS]. Hence, we have S (x x xx)) as the input of ViT (see bottom of Fig. D
The sequence S is processed by a series of transformer blocks composed of interleaving multi-head self-attention (MHSA) and MLP layers, with residual connections between each of the layers (see Fig. and appendix for details). The 6t block updates the inner representation (a.k.a. hidden state) R& of each token 2i indexed by i, eventually producing a new hidden state x².
Given set of classes C' and a class embedding matrix E E I xd which is learned during training, the predicted class probabilities are obtained by projecting the output of the [CLS] token of the last block and applying a softmax: Pess = softmax x3( (see top of Fig. 11). We assume each row of E represents& class prototype since they encode the patterns whose detection (via matrix multiplication with the [CLS] token) determines the probability of a class being present in the image.
Activation space projection. To investigate how class representations emerge in VIT, we analyzed the alignment between the intermediate representations of the tokens with the class prototypes encoded by the final projection matrix E. The key insight is that we have x E Rd. Hence, we can obtain the prediction of the class distribution ry E IR [C] for the intermediate representation of the 7th token in the output of the 6th block by computing p = E x
(see Fig. I.)
To quantify the alignment,
take inspiration from Brunner et al. 4 and adapt measure of identifiability. Concretely, we evaluate how recoverable the correct class Cg of the jth image is from the class projection of the hth token using a measure of their class identifiability, r/, that we define as:
argwhere(argsorlpp) =cy)
[C]
where argsort sorts the logits assigned to each class from higher to lower, and argwhere returns the index of the correct class in the sorted vector. We normalize and reverse r to obtain a score that ranges from to 1, with 1 denoting that the correct class has the highest logits and 0 the lowest.
Comparison with NLP research: The idea of projecting the intermediate token representations onto the class embedding space uncover the inner mechanisms of ViT is based on previous work in
NLP that takes a similar approach with autoregressive language transformers 2 7 g 12]. These models, trained to predict the next word of sentence, include output embedding matrix that translates the final hidden states of the network to human-interpretable vocabulary space. Taking advantage of this component, Geva al. [T]] and [22 projected the outputs of intermediate MLP layers onto the vocabulary space to show that these sub-modules can be decomposed into set of nub-mechanisms promoting nummn-interppeetble concepts.
In comparison to NLP models, ViTs are trained to predict the presence in an image of a more restricted set of classes. Therefore, the semantic insights obtained with the output embedding projection differ: in the case of NLP we can uncover linguistic anticipatory mechanisms, while in the case of ViT we can investigate how categorical representations are built.
Parameter space projection and key-value memory pairs. Previous work strated that the learned parameters of transformer architectures can also be projected onto the output embedding space for reverse engineering purposes. These studies further propose that the parameter matrices can be interpreted systems implementing key-value memory pair mechanisms, to better understand the mappings between the inputs of a model and its predictions.
12
has demon-
As shown in Fig. [b, key-value memories are a set of j paired vectors M = {(k1, vI),..., (k,, V2)} where the keys k; are used to quantify the presence of set of patterns in the input, and the values v: represent how the input should be modified in response to the detection of such pattern. We next explain how this system could be implemented in the MLP and self self-attention layers of ViT.
An MI.P layer in ViT consists of:
MLP(X) GELU W imp) W out
where X E Raxd represents the n-token input, each token of dimension d, and W E Md and W out E IR X] represent the parameter matrices.
3

---

The main idea is that the columns of W imp and the rows of Wout out can be thought of as a set of key-value paired vectors, respectively. The result of the operation shown in blue in eq. [ matrix of memory coefficients. Entry i, in the matrix contains the coefficient resulting from the dot product of token with key j- This matrix is the dynamic component of the system and measures the presence of certain patterns in the hidden state. The matrix shown in red in eq. 2 (here we think of them as value vectors) encodes how the internal representations x; should change in response to the detection of such patterns. To modify the hidden states of the network, the weighted value vectors of all key-value memories are summed up and added to the residual (see top of Fig. [T).
In the case of the self-attention layers, the decomposition of the layer into keys and values is more complex. The core idea is as follows (see Fig. [Tb). In transforming the hidden representation x; (row i in X), the self-attention layers can be thought of as implementing system of key-value memory pairs where the keys not only detect the presence of a certain pattern x; but also in the hidden states of all other tokens (x3 j A i in the sequence S. The coefficient reflecting the match between a key and a token x; weighted by the attention values between x; and all tokens in a self-attention head (including itself). Formally,
M = hconcat
A A'WWL.
A A'Xw/
W
(3)
Van
out
where A Ruxx are the attention weights of the 41th head with ʃ being the number of heads, and W VA E RR 4. The result of A A Vau for every hthh head is concatenated horizontally. The output of the horizontal concatenation is a matrix of dimensions n x [M] with |M] = d, which is then multiplied with the matrix W our E RI xx containing the value vectors. Of note, we can say that the matrices Wh Van of every hthh attention head represent the key vectors of the system.
Comparison with NLP research: Autoregressive language transformers employ an input embedding matrix to convert a sequence of semantically meaningful elements within vocabulary space into a sequence of machine-interpretable hidden states. These hidden states are subsequently processed by series of transformer blocks, and the resulting output is used for prediction by projecting it back to the vocabulary space using an output embedding matrix. In many architectures, a single vocabulary- projection matrix functions as both the input and output embedding of the model, with the output embedding matrix essentially being the transpose of the input embedding matrix. Consequently, previous work in NLP examined how the keys and values of the memory system represent patterns that are translatable to a cummnn-nterrrettbble vocabulary [2]. For example, in NLP transformers some keys in MLP layers have been reported to detect thematic patterns, such reference to TV " was associated with a value vector that
shows 11I. In turn, the detection of a TV show "theme" promotes related concepts in the vocabulary space (e.g. "episode", "season").
Unlike NI. transformer models, in ViT the mappings of the input embedding matrix (projecting from image patches to hidden states) differ from those of the output matrix (projecting from hidden states to class representations). Therefore, the keys of ViT may represent patterns that are interpretable in the image input space, while the value vectors may represent updates interpretable in the class embedding space. Furthermore, interpreting the keys of ViT is not as straightforward as in the NLP case. In autoregressive language transformers, the input space retains significance throughout the network's 's hierarchy since it aligns with the output prediction space. In ViTs, however, the input space is not relevant for later projections. Given that our work aims understand how categorical representations are formed, we focus on analyzing the representations of the value vectors. We leave to future work the analysis of what the keys encode in human-interpree interpretable terms.
4 Experimental design
Vision Transformer models. We validate our approach using a variety of ViTs that differ in their patch size, layer depth, training dataset, and use of architectural modifications. [ (Seecifically, we separately probed: 1) a vanilla VIT with 12 blocks and a patch size of 16 pre-trained using the ImageNet-21k dataset and fine-tuned on ImageNet- Ik (ViT-B/16) ():2) a variant with a bigger patch size of 32 (ViT-B/32) 99: 3) deeper version with 24 blocks (ViT-L/16) [ 4) a variant fine-tuned
on the CIFAR 100 dataset [14 : 5) a version pre-trained on an alternative Imagenet-21K 21K dataset with
higher-qualits semantic labels (MIIL) 20 ; 6) modified architecture with refinement module
'Our code is available at httttttss//gtthuu httpa//owwoo.e.teaoineee/ettttaa

---

that aligns the intermediate representations of all tokens to class space (Refinement) [15] : and 7) an alternate version trained with Global Average Pooling (GAP) instead of the [CLS] token.
Image dataset. For conducting our studies we used the ImageNet-S dataset [0] which consists of sub-selection of images from ImageNet accompanied by semantic segmentation annotations. We analyzed the representations of 5 randomly sampled images of every class from the validation set.
5 Class representations in image tokens across the hierarchy
During ViT's pre-training for image classification, the only token projected onto the class-embedding space is the [CLS] token from the last block. Thus, whether image tokens across the hierarchy can be translated to the class-embedding space remains an open question. In this section, we provide an affirmative answer.
Class representations in image tokens. First, to establish whether the hidden representations of image tokens encode class-specific representations, projected the hidden states of the image tokens X from the last block onto the class-embeddingg space using the output embedding matrix E, by computing E X We then measured the model's class identifiability rale, which we quantified as the percentage of image tokens that contain a class identifiability score of Similarly to Ghiasi al. [13], found that the identifiability rate of all VITs significantly higher than chance, and the rate was further influenced by the variant probed (see Table I). We additionally measured the percentage of images that contain at least one image token with an identifiability score of in the last block. We found that, for all variants tested on ImageNet-S, the percentage was significantly higher than in model initialized with random weights, and higher than their corresponding top-1 accuracy scores in the classification task (see appendix). The latter finding implies that even misclassified samples retain information within some of their image tokens that correspond the correct class.
Table 1: Class idemifiability rate (%) of image tokens the last block.
ViT-B/32
ViT-B/16
ViT-L/16
MIIL
CIFAR100
Refinement
GAP
60.73
67.04
72.58
78.52
90.35
79.64
52.09
Evolution of class representations across the hierarchy. Second, to investigate class prototype representations can be decoded from tokens at various stages of ViT's hierarchy, we projected the hidden states of every block onto the class embedding space and examined the evolution of their class identifiability scores. Fig. 2 shows that the mean class identifiability scores of both image and [CLS] tokens increased over blocks, for all ViT variants. Additionally, the results demonstrate that image tokens exhibited greater variability in their scores in comparison to [CLS] tokens (see variant-specific plots the appendix). This observation suggests that the development of class representations varies among image tokens.
[CLS)
image
VIT-B/32
class Identifiaility
class Identifiaiiiity
VIT-L/16
0.8
0.8
VIT-B/16-MIL.
VIT-B/1E-CIFAR
0.6
0.6
VTTBB/I--Rfinemnnt
VIT-Bi1--GA
Chance Level
normalized
block
Figure Evolution of class identifubility mean scores acroxs normalized blocks.
After establishing that image tokens can effectively translated to the class embedding space across different levels of the hierarchy, we proceed to showcase the applicability of our framework in examining some underlying factors contributing to the development of class representations. We exemplify this process using the simplest model variant, ViT-B/32.
5

---

The impact of attention mechanisms on the class representations of image tokens. To investigate whether image tokens need to incorporate information from other tokens via attention mechanisms to build a class-identifiable representation, we conducted perturbation experiments. For the hidden representation of every token we set to all attention weights from each image token x; to every other image token (x, j # i]. We found that the perturbations erased the class identifiability of image tokens in the last block (mean class identifiability decreased from 60.73% to 0. 1.%). The results suggest that image tokens cannot develop class-identifiable representations in isolation.
In contrast, when removing the attention weights between the image tokens and the [CLS] token, the class identifiability rate of image tokens remained unchanged, implying that image tokens do not need to extract information from the [CLS] token to build class representation. These results
are aligned to those of Zhai et al. (2I), who showed that ViTs trained without the [CLS] token can
achieve similar performance to models that include
The impact of context the class representations of image tokens. Class-identifrable rep- resentations may emerge earlier and more strongly in image tokens that belong to portions of an image depicting the class (e.g., striped patterns an image of a zebra). Leveraging the semantic segmentation labels from the ImageNet-S, we compared the identifiability rate of class-labeled and context-labeled image tokens. Our results confirmed our hypothesis, revealing earlier and more identifiable class representations in class- than context-labeled tokens (see Fig. 33 for ViT-B/32 and appendix for the other variants). However, in deeper blocks, both types of tokens exhibited high identifiability scores, indicating that class representations also emerged in image tokens that do not depict the class. This might be the result of context tokens extracting class-relevant informa- tion via attention mechanisms from class-labeled tokens, or it might stem from the class prototype representation incorporating information about the context commonly associated with the class.
To further explore the extent to which llass-labeled tokens are needed for context-labeled tokens to build class representation, and vice-versa, we conducted a pentrreeturraaioon 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-identifiabiiity 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. 18.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 -ittr--disttrbution scenarios.
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 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 '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, measured the class similarity change rate induced by these sub-modules. Given layer with input/output tokens and a class embedding matrix, we computed the proportion of output tokens whose correct class logits increase relative the corresponding logits for input tokens, where all logits are obtained by projecting the tokens onto the class embedding space
(see Fig. [p. Concretely, we projected the output of each layer O (X) onto the class embedding
space by Pour E O' o'(X)". and compared it to the projection of the input itself Pmnp = E. x"" We then quantified the proportion of tokens i where Pour Pinp

---

We found that, although the evolution of the class similarity change rate differed across VIT variants (see appendix and Fig. Bp 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 [CI.S] 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.
blook
block
mkp composttio
of tsuera
0.0
0.
corted
33
1.0
8
0.5
0.5
25
11
10
11
dlsa
centect
Figure 3: Results for VIT-B/32. (a) Class identifiability 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. [I]] and quantified the proportion of tokens whose top-1 prediction in the residual stream equaled that of the self-attention layer, that the MLP layer, that the residual stream before being processed by these layers, 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 MIPP layers but not being dominated by their outputs (see appendix and Fig. 3 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 (Wous of eq. 2 and eq. B) onto the class tllass-embedding space, by computing Pw... = E WT our The rows of E and W 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 / that resembles their representation. Hence, we extracted the maximum value of each row in P W... that denotes the highest similarity score obtained per class prototype in a given layer (. We call this measure class-value agreement score.
We found that, in deeper blocks, the class-valus agreement scores of self att-ntttenton and MLP layers were significantly higher than those of model initialized with random weights (Fig. 4 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 MIP sub-modules promote stronger categorical updates. In addition, manual inspection of the MLP layers containing high class-valu agreement scores revealed that the value
7

---

vectors may promote and cluster semantically similar and human inteaanniterrretable 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 distributed across value vectors instead.
VIT-B/32
MT-L/18
CTTBB/B-CIFARR NT-B//B-Refinement
VIT-B/16GGAP
VIT-B/18
WT-BMBB-MIIL
attn
mlp
atin
mlp
class-value agreement
0.4
0.4
key-value agreement
75
75
0.3
0.3
50
50
0.2
0.2
0.1
0.1
0
normalized layer
normalized layer
nomalized layer
normalized layer
Figure 4: Key-value memory pair mechanisms. (a) Class-valuc 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- alff-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 W 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 the penultimate blocks, while self- celf-attention layers peaked in the last blocks. The rates on these peaks were generally higher for self-atte 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 MI.P layers. Of note, similarly to the class-value agreement score results, ViT-GAP exhibited exhibiee a limited use of key-value memory pair systems.
We further evaluated the influence of these mechanisms 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 meaningful effect on the performance of the model.
Compositionality of key-value memory pair mechanisms. The final output of 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], 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 composi- tionality 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 MIP layers: more than 70% of instances had composite prediction (except for ViT-CIFAR, 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 class of interest, at each processing stage of VIT.
8

---

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 [CI.S] token. In detail, our explainability method can be applied as follows:
1. For the 7th image and the 8t block, project the hidden representation of the [CLS] token after the attention layer x = onto the class embedding space, and compute the cross-entropy loss Len of the class of interest Cj. such that E = LCE (E- ( the representation of the [CLS] token to the prototype representation of the class of interest.
x cz). This projection quantifies how close is
2. Compute the gradient of e with respect to the attention weights aß that the [CLS] tokens assigned the image tokens in an attention head in the self- layer, such that V - /Oaj Since we are interested in how the image tokens decrease the cross-cntropp loss, we negate the gradients. In simple terms, this step estimates the rate 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 block and attention head that can be visualized as a heatmap over the image (see Fig. 5a). Of note, the block- and attention sastioon-ha--se hhad-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 66 for overview). Instead, our method can visualize the categorical information contained in the image tokens independently for each block and attention head. This allows 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 image that produce 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 emphasizes in relation to the class-optimal ones.
(1)
(a) Block- and head-specific
visualization
(1)
Impale
Inx
[2]
Meeviioon
dudo couph
(b) Sum of gradients visualization
Figure 5: Examples feature importance visualization in VIT- (1) Image example where the correct label 'IIppla", but the model predicts "Ibex". (2) Image example with multiple labels.
Besides providing block- and ditention-head-speiife visualizations, our framework can also be used to generate traditional global relevancy map. Concretely, we can aggregate the gradients over the blocks and attention heads, by Σ, VeS. and obtain a final feature importance map that takes into account the relevance of each image token at every stage of processing (see Fig. (5)). 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 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 model. Similarly to the goals of our interpretability framework, to shed light on the relevant factors behind aategory-building processes in neural networks, linear probing is sometimes used to predict from internal layers the classes represented in the output projection matrix []. 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, 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, 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 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 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- buiidiig processes within these networks. Additionally, we used our framework to elucidate the distinct roles of self-attention and MIP 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.
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.
Limitations.
10

---

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 Strüngmann Institute for Neuroscience.
References
[1] G. Alain and Y. Bengio. Understanding intermediate layers using linear classifier probes. arXiv preprint arXiv: 2016
[2] N. Belrose, Z. Furman, L. Smith, D. Halawi, I. Ostrovsky, L. McKinney, Biderman, and J. Steinhardt. Eliciting latent predictions from transformers with the tuned lens. arXiv p reprint arXiv:2303.08112 2023.
[3] Bhojanapalli, A. Chakrabarti, D. Glasner, D. Li, T. Unterthiner, and A. Veit. Understanding ro- bustness of transformers for image classification. In Proceedings of the IEEE/CVF international conference on computer vision, pages 10231-10241, 2021.
[4] G. Brunner, Liu, D. Pascual, O. Richter, M. Ciaramita, and R. Wattenhofer. On identifiability in transformers. arXiv preprint arXiv: 042//, 2019.
[5] B. Carter, Jain, J. W. Mueller, and D. Gifford. Overinterpretation reveals image classification model pathologies. Advances in Neural Information Processing Systems, 34:15395-15400 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.02533 02535, 2022.
[8] A. Din, T. Karidi, L. Choshen, and M. Geva. Jump conclusions: Short-cutting transformers with linear transformations. arXiv preprint arXiv:2303.09433 09435, 2023.
[9] A. Dosovitskiy, L.. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11990
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 Confer- ence 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. Gold- stein. What do vision transformers learn? a visual exploration. arXiv preprint at arXiv22212.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 Neural Information Processing Systems, 14:23296-23300 2021.

---

[17] N. Park and S. Kim. How do vision transformers work? arXiv preprint arXiv:2202.06700 2022.
[18] M. Raghu, T. Unterthiner, Komblith, C. Zhang, and A. Dosovitskiy. Do vision transformers see like convolutional neural networks? Advances in Neural Information Processing Systems, 1111116-11122 2021.
[19] A.
Ravichander, Belinkov, and E. Hovy. Probing the probing paradigm: Does probing accuracy entail task relevance? arXiv preprint arXiv:2005.00719, 2020.
[20] Ridnik, E. Ben-Baruch, A. Noy, and L. Zeinik-Manor, Imagenet-21k pretraining for the masses, arXiv preprint arXiv:2104.1097. 2021.
[21] X. Zhai, A. Kolesnikov. N. Houlsby, and L. Beyer. Scaling vision transformers. In Proceedings the
of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, pages 12104-12115, 13, 2022.
12


Reviews:
[{"review_id": "uDnSTcMHiQ", "reviewer": "Reviewer_5aAZ", "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.", "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.", "limitations": "See the Weaknesses section for technical limitations. I cannot find any ethical issue. ", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Applying the embedding matrix is straightforward. This method does not necessitate additional training and is easy to implement.\n2. 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. \n    - Non-zero linear probing accuracy can be achieved even at the early layers [a].\n    - Self-attentions significantly change the representations, compared to MLPs [a].\n    - Self-attentions and MLPs perform complementary roles. For example, self-attentions aggregate information whereas MLPs diversify it [b].\n2. 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.\n\n[a] Raghu, Maithra, et al. \"Do vision transformers see like convolutional neural networks?.\" *NeurIPS* (2021).\n\n[b] Park, Namuk, and Songkuk Kim. \"How do vision transformers work?.\" ICLR (2022).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback, which has helped clarify and better demonstrate the novelty, generalizability, and usefulness of our framework and findings.\n\n\n_Weakness 1_:\n\nThank you for sharing relevant previous work which we now cite and discuss in the revised manuscript. We note that our work provides additional and complementary insights to those mentioned in the reviewer's comment. Specifically:\n- [a] shows that non-zero linear probing can be achieved in the [CLS] token, but as can be seen in Fig. 11 of their work, their linear probes did not achieve good accuracy when investigating the image tokens. In contrast, our method can successfully extract categorical information from image tokens and characterizes the factors enabling the emergence of class representations (e.g. the categorical construction in class-labeled vs context-labeled tokens). Moreover, as our general rebuttal and the “linear probing” section in this response elaborates in detail, linear probing experiments do not enable the same inferences as our method does.\n- Using a Ratio of Norms analysis, [a] shows that the self-attention layers impact the residual stream more strongly than MLP layers in the [CLS] token at early layers, and the reverse is true for image tokens at later layers (MLP more than self-attention layers). Importantly, this Ratio of Norms analysis only provides an estimate of the proportion of information added to the residual stream by these layers, in comparison with that added by the skip connections. It does not characterize the amount and the mechanisms by which the information added by each layer and block updates the categorical representations of the residual, as we did in our work. The mechanistic insights that we have gained with our approach can potentially be used for model editing or performance improvements in future work.\n- The results of [b] are indeed relevant to our work and are now cited in the revised manuscript. However, our findings are complementary to those reported in [b]. Similarly to [a], [b] analyzes statistical properties of the information added by self-attention and MLP layers of ViTs. In contrast, our work aims to provide insights into the semantic characteristics of the information added by these layers. Concretely, we investigate how categorical representations are built by exploiting mechanisms like key-value memory pairs that encode semantic information, which have been reported in NLP models.\n\nMore broadly, the greatest difference between our work and those mentioned in the reviewer's comment is that we introduce a general mechanistic interpretability method that can be efficiently applied to probe the categorical representations of different and future types of ViTs, as long as a class-embedding matrix has been learned during training.\n\nIn the new version of our manuscript, we explicitly clarify the novelty and significance in the “Related Work” section; specifically, how our findings provide additional insights to those mentioned in the review.\n\n\n_Weakness 2_:\n\nThank you for pointing out this limitation, which we now overcome with additional experiments yielding favorable results. As described in detail in the general comment, we have expanded our analyses to other variants of ViTs, including those using GAP instead of the [CLS] token. This new set of experiments demonstrates that we can successfully apply our framework to other types of ViTs to gain insights into how different training variables (e.g. dataset, architecture, pooling, learning objective), affects category building mechanisms. \nFor example, we found that introducing GAP does not decrease the class identifiability scores of image tokens across layers (see Fig. 1 of rebuttal PDF), but it does decrease the class identifiability rate of the tokens in the last layer (i.e. the percentage of tokens that contain a class identifiability score of 1). In addition, as Fig. 2 and 3 of the PDF show, GAP training decreases the reliance on key-value memory pair mechanisms for building categorical representations. Taken together, these findings corroborate the idea that the categorical representations on GAP-based ViTs emerge in a more distributed fashion than in [CLS]-based ViTs. However, we also found that class and context tokens have significant differences in the evolution of their identifiability scores, which allows us to conclude that even in these networks there is a meaningful pattern of how identifiability emerges from the image tokens.\nIn summary, these new experiments show that our framework can be used to investigate modern GAP-based ViTs, which strengthens the generalizability of our work.\n\n\n_Question_:\n\nThank you for raising this point and for suggesting linear probing experiments. We have run additional experiments accordingly, with informative and favorable outcomes. Please see the general response where we describe how we empirically demonstrate the advantages of our method over linear probing.\n\nHere, we also want to add that the linear probe experiment reported in Figure 13 of [a] aggregates the representations of the tokens. Thus, this approach does not take into account the potential differences in class identifiability between image and [CLS] tokens, or between different kinds of image tokens (e.g. class-labeled vs context-labeled image tokens, as explored in our study). For that reason, we instead replicated the experiments of Figure 9 of [a]. Moreover, as described in the general comment, we obtained different results than those reported in Figure 9 of [a], which may be taken as evidence of a lack of generalizability across training conditions in linear probing experiments.\n\nWe benefited from the comments and suggestions, and believe we have addressed them successfully. If so, please consider raising the reviewer rating."}, {"review_id": "BZR5kcY6v9", "reviewer": "Reviewer_mxXe", "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.\n\n\n", "questions": "See the weakness part.", "limitations": "In Section 8, the authors discuss the limitations of the adopted model's diversity and the underexplored application for performance improvement.\n\n\n\n\n\n", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\n2. This research demonstrates the intriguing phenomenon of internal disruption in categorical representations caused by context and attention perturbations.\n\n3. 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.\n\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive remarks and suggestions for improving the clarity and generalizability of our work. Below we address these comments in detail.\n\n_Weakness 1_\n\nThank you for the concrete suggestion on how to improve the description of our method. In our revised version of the paper, we have added a subsection thoroughly comparing our approach with that of NLP.\nDue to spatial constraints, we cannot reproduce the entire subsection here. As a summary of the added information in this section:\n- We describe how the input- and output-embedding matrices of both types of networks differ. Pre-trained LLMs include a vocabulary-projection matrix that serves both as the input and output embeddings of the model. In contrast, the nature of the input and output embedding matrices differ in ViTs: the input embedding matrix is a linear projection of the image patches, while the output matrix performs a categorical projection.\n- We note that the semantic task learned by these networks is different: Pre-trained LLMs are trained to predict the next word of a sentence (only auto-regressor LLMs were investigated in similar previous NLP studies), while ViTs for image classifications are trained to predict the semantic label of an image.\n- We discuss the implications of the above differences and how they may lead to differences in the information that is encoded in key-value memory pair systems. Particularly, we give a more thorough explanation as to why interpreting the keys of ViTs is not as straightforward as in the LLM case: while the input space of LLMs continues to be relevant throughout the network’s hierarchy because it is also the space that is projected to in the output prediction (see [1]), in ViTs the input-space is no longer relevant for later projections. That is why our current work focuses on analyzing the semantics of this system’s value vectors. We however note that future work can be dedicated to finding alternative ways of investigating what is encoded in the keys of these systems. \n\n_Weakness 2_\n\nWe agree on the value of investigating the generalizability of our framework and findings by probing ViTs trained on other datasets. We have now done so more thoroughly. As described in detail in the general response, in the new version of our paper we have expanded our analyses to a ViT-B/16 fine-tuned on CIFAR100, and a ViT-B/16 trained on a higher-quality and a multi-labeled version of ImageNet-21k (MIIL) [2].\nWe found that our framework can be used to analyze both variants of ViT, by successfully translating their [CLS] and image tokens across layers into the class-embedding space to get meaningful insights into how they build categorical representations and the effect of the training dataset in these processes.\n\nConcretely, we found that fine-tuning on CIFAR100 increases the identifiability score of tokens, from the very first layers of the model (see Fig. 1 in rebuttal PDF). This result suggests that fine-tuning ViTs in small datasets creates a class-embedding representation that is potentially overfitted to detect simple patterns in the image. Training in MIIL improves the identifiability of tokens (especially the [CLS]) without showing overfitting patterns, compared with a vanilla ViT-B/16 (see Fig. 1 in rebuttal PDF).\nFig. 2 in the rebuttal PDF also shows how both variants rely on key-value memory pair mechanisms, with MIIL showing the strongest key-value agreement rates of all networks, meaning that more tokens promote categories using these mechanisms than other ViT variants.\nWe have included these experiments and interesting findings in the new version of our manuscript.\n\nWe thank you again for the concrete suggestions that we believe to have sufficiently addressed. If so, please consider raising the reviewer rating.\n\n[1] Dar et al. “Analyzing Transformers in Embedding Space.” (2023).\n[1] Ridnik et al. \"Imagenet-21k pretraining for the masses.\" (2021).\n"}, {"review_id": "gOUSvkhfW1", "reviewer": "Reviewer_9G5D", "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.", "questions": "No questions come to mind.", "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.", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "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. \n\nQuality: 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.\n\nClarity: I found the experimental design was clear and the exposition easy to follow.\n\nSignificance: 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).\n\n* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive feedback and suggestions for improving the discussion of our findings. Below we address each of these in detail.\n\n_Weakness 1_\n\nThank you for pointing out relevant previous work of [1]. We have added a reference in section 5.1 of our manuscript, stating in line 195 that “These results are aligned to those of [1], where it is shown that ViTs trained without the [CLS] token can achieve similar performance to models that include it”.\nWe also want to highlight that our experiments are not redundant but complementary to [1] in the following sense. Our results show that even in ViTs trained with the [CLS] token, the image tokens can achieve class decodable information without extracting information from the [CLS], which sheds light on the mechanisms behind the building of categorical representations.\n\nMoreover, as described in the general comment, we carried out additional experiments and used our method to probe ViTs trained with GAP. We found that the training with GAP changes the mechanisms by which the network creates categorical representations. Concretely, while introducing GAP does not decrease the class identifiability scores of image tokens across layers (see Fig. 1 of the rebuttal PDF), it decreases the class identifiability rate of the tokens in the last layer (i.e. the percentage of tokens that contain a class identifiability score of 1). In addition, as Figure 3 of the rebuttal PDF shows, GAP training decreases the reliance on key-value memory pair mechanisms for building categorical representations. Taken together, these findings corroborate the idea that the categorical representations on GAP-based ViTs emerge in a more distributed fashion than in [CLS]-based ViTs. We have added these findings and their discussion to the new version of our manuscript.\n\n_Weakness 2_\n\nThank you for highlighting the relevant work by Bhojanapalli et al. (now cited)\n\nOur analyses are complementary to those of Bhojanapalli et al. who investigate the effects of removing self-attention and MLP layers in the performance of ViTs. In contrast, we analyze in detail how the categorical representations of ViTs are modified and built by these sub-modules, regardless of performance. We concretely investigate the use of key-value memory pair systems in these processes.\n\nIn addition, we think our perturbation studies complement the findings on the robustness of ViTs in Bhojanapalli et al. We show how context-tokens are necessary for building a categorical representation in the class-labeled image tokens, which points to how much information unrelated to the class the model is including in the category representation of the embedding matrix.\n\nWe have added a discussion of the work of Bhojanapalli et al. and the mentioned differences with our study to the “Related Work” section.\n\nThanks again for the feedback. We believe we have addressed all your concerns. If so, please consider raising the reviewer rating.\n\n[1] Zhai et al. \"Scaling vision transformers.\" (2022)."}, {"review_id": "JdldVrXGTu", "reviewer": "Reviewer_GYsr", "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.", "questions": "1. Why there are performance differences betwee block 11 and other blocks? \n2. Why sum the gradients over the blocks? \n3. What does it mean in line 307-308?", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n2. 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.\n3. 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.\n4. 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.\n5. 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\n2. 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.\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the careful reading and useful feedback. We address the comments and questions below.\n\n_Weakness 1_\n\nThank you for the suggestion. We have now demonstrated the generalizability of our method to all cases mentioned by the reviewer (see general response) and expanded our analyses to other variants of ViTs, including larger models, trained on different datasets, or containing architectural modifications (e.g. use of GAP). The new experiments demonstrate that we can apply our method to investigate a wide range of ViTs. \n\n_Weakness 2_\n\nThank you for suggesting comparisons with other mechanistic interpretability and explainability methods.\n\nRegarding mechanistic interpretability, we have added a comparison with linear probing approaches (see general response), which are widely used to study if categorical information is present in the hidden states of DNNs. Our new experiments show that the insights we can obtain from linear probes are very different from those of our method, and less informative for the purposes here.\n\nRegarding explainability, our approach is complementary to existing methods and provides different kinds of insights. Specifically, SOTA methods that generate relevancy maps for a classification example compute the gradients of the class logits at the final layer with respect to the input and/or aggregate the information propagated across layers up to the class-embedding projection. Instead, our method is able to visualize the categorical information contained in the image tokens independently for each block. These block-specific visualizations allow us to (1) better understand how categorical information is hierarchically built, and (2) characterize the importance of each block in building the class representations. We clarify this conceptual difference between methods in the new version of our paper and provide some examples of layer-specific visualizations.\n\nRegarding question 2, we originally suggested aggregating the gradients over blocks to show that our framework can also provide a global relevancy map that is semantically meaningful and can be used for traditional explainability visualizations. The sum procedure also allows us to corroborate that we obtain a fair portrayal of the individual contribution of each block to the final categorical representation: if the aggregated visualization provides a coherent and useful relevancy map (measured with perturbation studies), there is evidence that the interpretation of the individual blocks is meaningful. We have explained why we carried out this sum explicitly in the new version of our paper.\n\nIn addition, we compare the quality of our global relevancy map to that of an established explainability method (reference [3] in our paper) to further prove the accuracy and quality of our proposal. Using ViT-B/32, for both methods we (1) quantified the importance of each 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 of the final accuracies. Our results show that our methods yield similar results to those of [3] (Table 2 in rebuttal PDF), highlighting the adequacy of our approach.\n\n_Q1_\n\nWe found that in VIT-B/32, MLP 11 promotes the strongest categorical updates via key-value memory pair mechanisms (section 6.2). This could explain why if the activation of the key-value memory pair system in MLP 11 is high (as measured by the key-value agreement rate), a better classification accuracy will be obtained as a result of a stronger categorical update.\nRegarding why the strongest categorical updates are encoded in MLP 11, we hypothesize it is due to this being the last layer from which the [CLS] token can extract information (through the self-attention layer in block 12) to make a categorical decision. In turn, MLP layers promoting stronger categorical updates at the very final layers could be useful because the depth allows the model to develop more complex and semantically meaningful keys (as shown in NLP [1]).\nOur new experiments provide additional evidence in favor of these hypotheses: 1) Models with deeper architectures present stronger categorical updates in the layers before the final one, not in layer 11 (Fig. 2 in reb. PDF); 2) When training explicitly aligns intermediate representations to the class embedding space in every layer, and thus the network is able to use the image token representations in MLP 12 to make categorical decisions, we found that the strength of the categorical updates in MLP 12 is even higher to those of MLP 11 (see Refinement ViT in Fig.3 of PDF).\nIn addition, we also note that in the self-attention 12 the strength of the categorical updates and the activation of key-value memory pair systems is higher than those of earlier blocks. This is further evidence that ViTs exploit key-value memory pair systems for promoting categories at the very late stages of the network so that the key representations can encode greater semantic complexity.\n\nWe have added this interpretation of the results to our paper.\n\n_Q2_\n\nSee discussion of weakness 2.\n\n_Q3_\n\nIf the question refers to our compositionality results, another way of stating these findings would be: in at least 80% of the cases, the final predictions of these layers do not exactly match the categories promoted by the most activated keys. Instead, results suggest that in the majority of cases, the final prediction is a composition of the categorical representations promoted by more than one key. The strength of this compositionality varied across blocks and layers.\n\nWe thank you for the suggestions again, which we consider to have been addressed successfully. If so, please consider raising the reviewer rating.\n\n[1] Geva et al. \"Transformer feed-forward layers are key-value memories.\" (2020)."}, {"review_id": "Qxt0ur2pBE", "reviewer": "Reviewer_vRFW", "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.", "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).", "limitations": "No clear discussion of limitations or possible impacts.", "rating": 3, "confidence": 4, "soundness": 1, "presentation": 3, "contribution": 2, "strengths": "1) The paper explores an interesting and useful direction on better interpretability of ViTs   \n2) 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\n2) 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.  \n3) 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? \n4) Results and discussion in section 5 are not explained well \n\n[1] Naseer, Muzammal et al. “On Improving Adversarial Transferability of Vision Transformers.” ArXiv abs/2106.04169 (2021): n. pag.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback which has helped 1) clarify in our manuscript the rationale, relevance, and novelty of our work, and 2) improve the generalizability and clarity of our findings. We address in detail each concern stated in the revision below.\n\n_On the rationale behind our method_\n\nWe argue that projecting intermediate representations to the class-embedding space allows us to investigate how the hidden states increasingly represent the class prototype learned by the model. 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]). The new experiments comparing our framework with linear probing studies (see general response) give evidence that our method can better detect the features that lead to a categorically meaningful representation. Note that we do not use our framework to claim that early layers lack information that, after linear and non-linear transformations, will be relevant for making categorical decisions. In fact, we show that from very early layers of the model (including layer 1 for some ViT variants) there is significant alignment to the categorical representation as compared to a random model. Instead, our framework shows how projection to the class embedding of intermediate features allows us to investigate the factors (i.e. token position) and mechanism (i.e. via key-value memory pair mechanisms) that increase the alignment with the class prototype. \n\nWe have clarified these issues in the new version of our manuscript.\n\n_On the novel contributions of our work_\n\nRegarding the novelty of our contributions over [10,12], the key additional advances of our work are the following: \n- We introduce a generalizable interpretability framework that can be used to efficiently investigate the categorical building processes of ViTs. We strengthen the generalizability and advantages of our framework by using the new experiments reported in the general response to show that: (1) We can apply our method to analyze the representations of ViTs that have learned a class-embedding matrix during training; (2) Our framework can distinguish the effects of different training variants (e.g. datasets, architectural depth or constraints) in the build-up of class representations; (3) Our method can characterize how categorical information emerges in the image tokens more efficiently and accurately than the commonly used linear probe method (see next section for additional details).\n- We use our framework to investigate how a specific mechanism not mentioned in [10,12], namely key-value memory pair systems, is used by different layers of ViTs to add semantically meaningful information to the residual stream and build categorical representations. Previous work in NLP showed that these mechanisms can be exploited for model editing, efficiency improvements at inference time, or performance improvements, and thus we consider that our study opens future avenues of research. Moreover, our new experiments (see general response) showing that these mechanisms are present across ViT variants also give evidence of the broad impact of understanding key-value memory pair systems.\n\nWe clarify and reinforce the importance of these findings in the “Introduction” section of the new version of our manuscript.\n\n_On the comparison with linear probing methods_\n\nIn the general response, we provide a comparison of our framework with that of linear probing studies. In these experiments, we specifically show how quantifying the “class-relevant” information contained in the inner representations of ViTs with linear probes does not necessarily enable insights into the features and mechanisms that the model uses for making categorical decisions. This is in contrast with our method which enables such insights.\n\n_On the analysis of ViTs with explicit intermediate alignment to class representations_\n\nThank you for suggesting applying our framework to ViTs trained by Naseer et al., as it has led to meaningful insights and has strengthened the generalizability of our work.\nAs shown in Table 1 and Fig. 1 of the rebuttal PDF, results suggest that aligning the intermediate representations of the tokens to the class embedding space leads to 1) increased class-identifiability scores in [CLS] tokens across all layers; 2) more class-identifiable image tokens in the last layer; 3) less class-identifiable tokens in middle layers. Moreover, as shown in Fig. 2 of the rebuttal PDF, these ViTs also make use of key-value memory pair systems and, unlike other networks, extend this use to block 12. The difference might be due to the use of the image tokens in MLP layer 12 to predict the correct category in the class-refinement module, compared to common ViT variants where only the [CLS] token in MLP layer 12 is used for prediction.\n\n_On the clarity of Section 5_\n\nWe have now improved the discussion of the results in Section 5, and we:\n- Explain some of our results in light of previous work, and provide more interpretation of the patterns that we observe.\n- Rephrase some of our statements to improve clarity.\n- Improve the captions of Table 1.\nIf there are other concrete concerns with the clarity of this section other than those mentioned, we will be glad to further work on them.\n\n_On the discussion of limitations_\n\nNote that we do mention the limitations of our work and its impact in the original manuscript; specifically in the original conclusion section. If a specific aspect that the reviewer considers important and was overlooked in the description of our limitations, we are glad to include it in the new version of our manuscript.\n\nWe have benefited from your comments and suggestions, and believe we have addressed them favorably. If so, please consider raising the reviewer rating.\n\n[1] Hill, F. “Why transformers are obviously good models of language.” (2023)."}]

Author rebuttal:
We thank the reviewers for feedback that helped improve our manuscript. In the general response, we address two common concerns about our work: we (1) report new experiments that provide evidence of the generalizability of our method, and (2) conceptually discuss and empirically demonstrate the advantages of our framework over linear probes. In all cases, we obtain results that are informative and favor our proposed method.

1) Generalizability

Some reviewers raised the possibility that our framework might not generalize outside vanilla ViTs. We agree that more experiments are needed to demonstrate generalizability. We have thus tested 5 additional variants of ViTs and obtained favorable results in all cases.

Specifically, we separately probed ViTs that are 1) larger (ViT-L/16); 2) fine-tuned for other datasets (CIFAR100); 3) trained on higher-quality datasets (MIIL) [1]; 4) trained with a refinement module that aligns all tokens’ intermediate representations to class space [2]; and 5) trained with Global Average Pooling (GAP) instead of the [CLS] token.

Our findings demonstrate that we can successfully translate the [CLS] and image tokens of all these variants into the class-embedding space across layers. As depicted in Table 1 (see rebuttal PDF), the identifiability rate of the image tokens in the last layer was above chance in all models. In addition, Fig. 1 shows class identifiability increases over blocks for all variants. We also found that in all networks, class-labeled tokens have higher identifiability scores than context-labeled ones, which illustrates that our method can be used to characterize the emergence of categorical representations in image tokens.

Our results also prove that we can use our framework to inspect the role of attention and MLP layers across ViT variants, and their use of key-value memory pair systems (Fig. 2). Although we observe similar patterns across networks, our method also reveals quantifiable differences indicating how the modification of vanilla ViTs affect the building of categorical representations. We discuss these differences in the responses to each reviewer comment.

We have added the above findings to our paper and: (1) swapped Fig. 2a and 4 (main paper) with Fig. 3 and 4 (PDF), (2) discussed the similarities and differences across ViT variants, and (3) included figures for each network in the supp. material.

2) Comparison to linear probing

Some reviewers point out that we do not discuss the advantages of our method over the commonly used linear probing approach. In the revised manuscript we perform experiments to demonstrate that linear probing is less informative for our research question. 

We first note that these two methods differ in the type of insights they can provide. Although both aim to characterize the information encoded in the hidden representations, only our framework can quantify how these intermediate representations increasingly align with the class prototype learned by the model (encoded by the weights of the class-embedding matrix). Furthermore, our method additionally characterizes the inner mechanisms of this alignment (e.g. by investigating key-value memory pair systems via the projection of ViT’s parameter matrices).
In contrast, linear probes measure the linear separability in the hidden representations of samples of different classes [3] but do not reveal whether the learned separability is exploited by the model in categorical decisions. Thus, linear probing does not necessarily uncover the relevant factors behind category-building processes in ViTs. In other words, the successful class probing of tokens in a given layer can rather reflect the separability of confounders whose feature representations are ignored by the model in the class-embedding space. For example, for a particular token and layer, a linear probe may use internal representations of background information that is highly correlated with the presence of a class in a specific dataset, to distinguish between categories. We show this is likely through the experiments described next.

Following [12], in the revised paper we trained separate 10-shot linear classifiers on ImageNet for each token position and layer of a VIT-B/32. To test if the information learned by these probes shed 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 the tokens 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. 

Our results demonstrate that even if linear probes can generally decode with better top-1 accuracy the classes of the image tokens, this is not informative of the relevance of each token for the categorical decision of the network (Fig. 3). Moreover, we found that our experiments achieve higher accuracy than those reported in [12], which evidences the lack of generalizability of linear probes across training conditions (e.g. different datasets).

Finally, note that even if linear probes and our method would provide similar insights (which we demonstrated is not the case), our framework is more time-efficient: it comprises a one-forward pass on validation images, and zero-pass when projecting the parameter matrices. In contrast, linear probes additionally involve 1) a one-forward pass over the training images, and 2) the fitting of a linear classifier for every token position and layer.

We add these linear probe results to a new subsection in the main manuscript, and Fig. 3 (PDF) to the supp. material.

[1] Ridnik et al. "Imagenet-21k pretraining for the masses" (2021);
[2] Naseer et al. “On Improving Adversarial Transferability of Vision Transformers” (2021);
[3] Alain & Bengio. "Understanding intermediate layers using linear classifier probes" (2016);

Meta-review: The paper performs an analysis of trained Vision Transformers. In particular, it proposes an analysis technique that uses the pre-trained network head as a probe to quantify class-alignment in different parts of the network. The paper received mixed reviews, that lean positive on average. Concerns include: the lack of actionable conclusions from the analysis, value over the standard linear probing technique, and generality of the results to other ViT flavours and datasets.

In the rebuttal the authors present a compelling experimental evidence that their class-based probe has substantial orthogonal value to linear probing, plus they provide analysis on additional networks and datasets. These address the second two concerns. Regarding "actionable conclusions", I believe that thoughtful analysis that provides new insights, such as those presented here, can still be a valuable contribution to the community, even if the take-aways for future network design/training are not immediately obvious. I do not see any blocking concerns remaining after the rebuttal, therefore I recommend acceptance.
Decision: Accept (poster)