SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Martina G. Vilas
Goethe University Frankfurt  
Ernst Strungmann Institute
Timothy Schaumlöffel
Goethe University Frankfurt The Hessian Center for AI
Gemma Roig
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.
Transformer-based deep neural networks have become one of the most popular architectures in machine learning due to their remarkable performance and adaptability to multiple domains. Consequently, recent work has focused on reverse-engineering the inner mechanisms of these networks for better understanding and control of their predictions (e.g. [2, 7, 11, 12, 18]). Of these, a series of studies in the field of NLP [2, 7, 8, 11, 12] have shed light on the predictive mechanisms of large language transformers by projecting the hidden states and parameters of intermediate processing layers onto a vocabulary space using the pre-trained output embedding matrix. This approach has enabled the analysis in human-understandable terms of how next-word predictions are built internally, introducing a simple and efficient method for the mechanistic interpretability of NLP transformers.
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 a 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 a 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
characterize the emergence of class representations in image tokens more efficiently and accurately than the commonly used linear probing approach.
Prior research has demonstrated that within ViT models, (1) the use of a 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-identifiable information. By contrast, we introduce a methodological framework that adeptly extracts categorical information from image tokens early in the processing stages, and allows us 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.
We use our framework to complement previous findings on the inner mechanisms of ViTs. Ghiasi et al. [13] 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. In 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 et 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 [3, 13, 16] 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 explanations 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.
3 Interpretation of vision transformers mechanisms in class embedding space
Figure 1: Schematic of our framework. (a) The hidden states of image tokens  $x_{n}$  in a block  $b$  are projected onto the class embedding space using the output embedding matrix  $\mathbf{E}$ . (b) A key-value memory pair system in a self-attention layer. Key vectors  $k_{j}$  belong to different attention heads  $h_f$ . The match between the hidden states  $x_{n}$  and the keys  $k_{j}$  is weighted by attention values to produce a memory coefficient. Value vectors  $v_{j}$  weighted by these memory coefficients are summed up and added to the residual stream. Adapted from Geva et al. [11].
ViT architecture. As introduced by Dosovitskiy et al. [9] and depicted in Fig. 1a, vanilla image-classification ViTs take as input a sequence of linear projections of equal-sized  $n$  image patches
with added position embeddings. We refer to these as image tokens  $\langle \mathbf{x}_1,\dots ,\mathbf{x}_n\rangle$  with  $\mathbf{x}\in \mathbb{R}^d$ . The sequence also includes a special "class token" whose initial representation is learned during training, denoted [CLS]. Hence, we have  $S = \langle \mathbf{x}_{cls}^{0},\mathbf{x}_{1}^{0},\ldots ,\mathbf{x}_{n}^{0}\rangle$  as the input of ViT (see bottom of Fig. 1a).
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. 1a and appendix for details). The  $b^{\mathrm{th}}$  block updates the inner representation (a.k.a. hidden state)  $\mathbf{x}_i^{b - 1} \in \mathbb{R}^d$  of each token  $x_{i} \in S$  indexed by  $i$ , eventually producing a new hidden state  $\mathbf{x}_i^b$ .
Given a set of classes  $C$  and a class embedding matrix  $\mathbf{E} \in \mathbb{R}^{|C| \times d}$  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:  $\mathbf{p}_{\mathrm{cls}} = \text{softmax}(\mathbf{E} \cdot \mathbf{x}_{\mathrm{cls}}^{b_n})$  (see top of Fig. 1a). We assume each row of  $\mathbf{E}$  represents a 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  $\mathbf{E}$ . The key insight is that we have  $\mathbf{x}_i^b\in \mathbb{R}^d$ . Hence, we can obtain the prediction of the class distribution  $p_i^b\in \mathbb{R}^{|C|}$  for the intermediate representation of the  $i^{\mathrm{th}}$  token in the output of the  $b^{\mathrm{th}}$  block by computing  $\mathbf{p}_i^b = \mathbf{E}\cdot \mathbf{x}_i^b$  (see Fig. 1a).
To quantify the alignment, we take inspiration from Brunner et al. [4] and adapt a measure of identifiability. Concretely, we evaluate how recoverable the correct class  $c_{j}$  of the  $j^{\mathrm{th}}$  image is from the class projection of the  $i^{\mathrm{th}}$  token using a measure of their class identifiability,  $r_{i}^{j}$ , that we define as:
$$
r _ {i} ^ {j} = 1 - \frac {\operatorname {a r g w h e r e} (\operatorname {a r g s o r t} (\mathbf {p} _ {\mathbf {i}}) = c _ {j})}{| C |} \tag {1}
$$
where  $\text{argsort}$  sorts the logits assigned to each class from higher to lower, and  $\text{argwhere}$  returns the index of the correct class in the sorted vector. We normalize and reverse  $r_i^j$  to obtain a score that ranges from 0 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 to uncover the inner mechanisms of ViT is based on previous work in NLP that takes a similar approach with autoregressive language transformers [2, 7, 8, 11, 12]. These models, trained to predict the next word of a sentence, include an output embedding matrix that translates the final hidden states of the network to a human-interpretable vocabulary space. Taking advantage of this component, Geva et al. [11] and [12] projected the outputs of intermediate MLP layers onto the vocabulary space to show that these sub-modules can be decomposed into a set of sub-mechanisms promoting human-interpretable 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 [7, 11, 12] has demonstrated 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 as systems implementing key-value memory pair mechanisms, to better understand the mappings between the inputs of a model and its predictions.
As shown in Fig. 1b, key-value memories are a set of  $j$  paired vectors  $M = \{(\mathbf{k}_1, \mathbf{v}_1), \dots, (\mathbf{k}_j, \mathbf{v}_j)\}$  where the keys  $\mathbf{k}_i$  are used to quantify the presence of a set of patterns in the input, and the values  $\mathbf{v}_i$  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-attention layers of ViT.
An MLP layer in ViT consists of:
$$
\operatorname {M L P} (\mathbf {X}) = \operatorname {G E L U} \left(\mathbf {X} \mathbf {W} _ {\text {i n p}}\right) \mathbf {W} _ {\text {o u t}} \tag {2}
$$
where  $\mathbf{X} \in \mathbb{R}^{n \times d}$  represents the  $n$ -token input, each token of dimension  $d$ , and  $\mathbf{W}_{\mathrm{inp}} \in \mathbb{R}^{d \times |M|}$  and  $\mathbf{W}_{\mathrm{out}} \in \mathbb{R}^{|M| \times d}$  represent the parameter matrices.
The main idea is that the columns of  $\mathbf{W}_{\mathrm{inp}}$  and the rows of  $\mathbf{W}_{\mathrm{out}}$  can be thought of as a set of key-value paired vectors, respectively. The result of the operation shown in blue in eq. 2 is a matrix of memory coefficients. Entry  $i,j$  in the matrix contains the coefficient resulting from the dot product of token  $i$  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  $\mathbf{x}_i$  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. 1b).
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. 1b). In transforming the hidden representation  $\mathbf{x}_i$  (row  $i$  in  $\mathbf{X}$ ), the self-attention layers can be thought of as implementing a system of key-value memory pairs where the keys not only detect the presence of a certain pattern in  $\mathbf{x}_i$  but also in the hidden states of all other tokens  $\{\mathbf{x}_j : j \neq i\}$  in the sequence  $S$ . The coefficient reflecting the match between a key and a token  $\mathbf{x}_i$  is weighted by the attention values between  $\mathbf{x}_i$  and all tokens in a self-attention head (including itself). Formally,
$$
\operatorname {M H S A} (\mathbf {X}) = \operatorname {h c o n c a t} \left[ \mathbf {A} ^ {1} \mathbf {X} \mathbf {W} _ {\text {V} _ {\mathrm {a t t n}}} ^ {1}, \dots , \mathbf {A} ^ {f} \mathbf {X} \mathbf {W} _ {\text {V} _ {\mathrm {a t t n}}} ^ {f}, \right] \mathbf {W} _ {\text {o u t}} \tag {3}
$$
where  $\mathbf{A}^h\in \mathbb{R}^{n\times n}$  are the attention weights of the  $h^{\mathrm{th}}$  head with  $f$  being the number of heads, and  $\mathbf{W}_{V_{\mathrm{attn}}}^h\in \mathbb{R}^{d\times \frac{d}{f}}$ . The result of  $\mathbf{A}^h\mathbf{X}\mathbf{W}_{V_{\mathrm{attn}}}^h$  for every  $h^{\mathrm{th}}$  head is concatenated horizontally. The output of the horizontal concatenation is a matrix of dimensions  $n\times |M|$  with  $|M| = d$ , which is then multiplied with the matrix  $\mathbf{W}_{\mathrm{out}}\in \mathbb{R}^{|M|\times d}$  containing the value vectors. Of note, we can say that the matrices  $\mathbf{W}_{V_{\mathrm{attn}}}^h$  of every  $h^{\mathrm{th}}$  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 a vocabulary space into a sequence of machine-interpretable hidden states. These hidden states are subsequently processed by a 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 human-interpretable vocabulary [11, 12]. For example, in NLP transformers some keys in MLP layers have been reported to detect thematic patterns, such as a reference to TV shows [11]. In turn, the detection of a TV show "theme" was associated with a value vector that promotes related concepts in the vocabulary space (e.g. "episode", "season").
Unlike NLP 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 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 to 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-interpretable terms.
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. Specifically, 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-1k (ViT-B/16) [9]; 2) a variant with a bigger patch size of 32 (ViT-B/32) [9]; 3) a deeper version with 24 blocks (ViT-L/16) [9]; 4) a variant fine-tuned on the CIFAR100 dataset [14]; 5) a version pre-trained on an alternative Imagenet-21K dataset with higher-quality semantic labels (MIIL) [20]; 6) a modified architecture with a refinement module
1Our code is available at https://github.com/martinagvilas/vit-cls_emb
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 [10], which consists of a 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, we projected the hidden states of the image tokens  $\mathbf{X}$  from the last block onto the class-embedding space using the output embedding matrix  $\mathbf{E}$ , by computing  $\mathbf{E} \cdot \mathbf{X}^T$ . We then measured the model's class identifiability rate, which we quantified as the percentage of image tokens that contain a class identifiability score of 1. Similarly to Ghiasi et al. [13], we found that the identifiability rate of all ViTs was significantly higher than chance, and the rate was further influenced by the variant probed (see Table 1). We additionally measured the percentage of images that contain at least one image token with an identifiability score of 1 in the last block. We found that, for all variants tested on ImageNet-S, the percentage was significantly higher than in a 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 to the correct class.
Table 1: Class identifiability rate (%) of image tokens in 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 if 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 in the appendix). This observation suggests that the development of class representations varies among image tokens.
Figure 2: Evolution of class identifiability mean scores across normalized blocks.
After establishing that image tokens can be 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.
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  $i$  we set to 0 all attention weights from each image token  $\mathbf{x}_i$  to every other image token  $\{\mathbf{x}_j : j \neq 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.12\%$ ). 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 a class representation. These results are aligned to those of Zhai et al. [21], who showed that ViTs trained without the [CLS] token can achieve similar performance to models that include it.
The impact of context on the class representations of image tokens. Class-identifiable representations may emerge earlier and more strongly in image tokens that belong to portions of an image depicting the class (e.g., striped patterns in 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. 3a 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 information 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.