SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Björn Deiseroth $^{1,2,3*}$  Mayukh Deb $^{1*}$  Samuel Weinbach $^{1*}$  Manuel Brack $^{2,4}$  Patrick Schramowski $^{2,3,4,5}$  Kristian Kersting $^{2,3,4}$
$^{1}$ Aleph Alpha  $^{2}$ Technical University Darmstadt
$^{3}$ Hessian Center for Artificial Intelligence (hessian.AI)
4German Center for Artificial Intelligence (DFKI) 5LAION
{bjoern.deiseroth, mayukh.deb, samuel.weinbach}@aleph-alpha.com
{manuel.brack, patrick.schramowski}@dfki.de
{kersting}@cs.tu-darmstadt.de
Generative transformer models have become increasingly complex, with large numbers of parameters and the ability to process multiple input modalities. Current methods for explaining their predictions are resource-intensive. Most crucially, they require prohibitively large amounts of additional memory since they rely on backpropagation which allocates almost twice as much GPU memory as the forward pass. This renders it difficult, if not impossible, to use explanations in production. We present ATMAN that provides explanations of generative transformer models at almost no extra cost. Specifically, ATMAN is a modality-agnostic perturbation method that manipulates the attention mechanisms of transformers to produce relevance maps for the input with respect to the output prediction. Instead of using backpropagation, ATMAN applies a parallelizable token-based search method relying on cosine similarity neighborhood in the embedding space. Our exhaustive experiments on text and image-text benchmarks demonstrate that ATMAN outperforms current state-of-the-art gradient-based methods on several metrics and models while being computationally efficient. As such, ATMAN is suitable for use in large model inference deployments.
1 Explainability through attention maps
Generalizing beyond single-task solutions using large-scale transformer-based language models has gained increasing attention from the community. In particular, the switch to open-vocabulary predictions promises AI systems capable of generalizing to new tasks. Arguably, transformers are the state-of-the-art method in Natural Language Processing (NLP) and Computer Vision. They have shown exceptional performance in multi-modal scenarios, e.g., bridging Computer Vision (CV) capabilities with text understanding to solve Visual Question Answering (VQA) scenarios [10, 18, 31, 30]. However, the increasing adoption of transformers also raises the need to understand their otherwise black-box predictions better. Unfortunately, the "scale is all you need" assumption of recent transformer models results in severely large and complex architectures. This renders tasks such as their training, inference deployment, and explaining of their outputs resource-intensive, requiring multiple enterprise-grade GPUs or even entire computing nodes, along with prolonged runtimes.
arXiv:2301.08110v5 [cs.LG] 5 Nov 2023
* These authors contributed equally to this work
Source code: https://github.com/Aleph-Alpha/AtMan
Atman (Sanskrit) "essence"
(a) "What am I looking at?"
(b) ATMAN in the transformer architecture.
Figure 1: (a) The proposed explainability method ATMAN visualizes the most important aspects of the given image while completing the sequence, displayed above the relevance maps. The generative multi-modal model MAGMA is prompted to describe the shown image with: “<Image> This is a painting of”. (b) The integration of ATMAN into the transformer architecture. We multiply the modifier factors and the attention scores before applying the diagonal causal attention mask as depicted on the right-hand side. Red hollow boxes  $(\square)$  indicate one-values, and green ones  $(\square)$  -infinity. (Best viewed in color.)
Explainable AI (XAI) methods target to elucidate the decision-making processes and internal workings of AI models to humans. State-of-the-art methods for transformers focus on propagating gradients back through the model. This backpropagation allows for the accumulation of influence of each input feature on the output by utilizing stored activations during the forward pass [6, 27]. Unfortunately, this caching also leads to significant overhead in memory consumption, which renders their productive deployment to be uneconomical, if not impossible. Often half of the available memory of the GPU has to remain unused at inference, or it requires an entirely separate deployment of the XAI pipeline.
Fortunately, another popular XAI idea, namely perturbation [19, 24], is much more memory-efficient. However, it has not yet been proven beneficial for explaining the predictions of transformers, as the immense number of necessary forward trials to derive the explanation accumulate impractical computation time.
To tackle these issues and in turn, scale explanations with the size of transformers, we propose to bridge relevance propagation and perturbations. In contrast to existing perturbation methods—executing perturbations directly in the input space—ATMAN pushes the perturbation steps via attention manipulations throughout the latent layers of the transformer during the forward pass. This enables us—as we will show—to produce state interpolations, token-based similarity measures and to accurately steer the model predictions. Our explanation method leverages these predictions to compute relevancy values for transformer networks. Our experimental results demonstrate that ATMAN significantly reduces the number of required perturbations, making them applicable at deployment time, and does not require additional memory compared to backpropagation methods. In short, ATMAN can scale with transformers. Our exhaustive experiments on the text and image-text benchmarks and models demonstrate that ATMAN outperforms current state-of-the-art gradient-based methods while being more computationally efficient. For the first time, ATMAN allows one to study generative model predictions as visualized in Fig. 1a without extra deployment costs. During the sequence generation with large multi-modal models, ATMAN is able to additionally highlight relevant features wrt. the input providing novel insights on the generation process.
Contributions. In summary, our contributions are: (i) An examination of the effects of token-based attention score manipulation on generative transformer models. (ii) The introduction of a novel and memory-efficient XAI perturbation method for large-scale transformer models, called ATMAN, which reduces the number of required iterations to a computable amount by correlating tokens in the embedding space. (iii) Exhaustive multi-modal evaluations of XAI methods on several text and image-text benchmarks and autoregressive (AR) transformers.
We proceed as follows. We start by discussing related work in Sec. 2. In Sec. 3, we derive ATMAN and explain its attention manipulation as a perturbation technique. Before concluding and discussing the benefits as well as limitations in Sec. 5, we touch upon our experimental evaluation Sec. 4,
showing that ATMAN not only nullifies memory overhead and proves robust to other models but also outperforms competitors on several visual and textual reasoning benchmarks.
Explainability in CV and NLP. The explainability of AI systems is still ambiguously defined [8]. XAI methods are expected to show some level of relevance to the input with respect to the computed result of an algorithm. This task is usually tackled by constructing an input relevance map given the model's prediction. The nature of relevance can be class-specific, e.g., depending on specific target instances of a task and showing a "local solution" [26, 27], or class-agnostic, i.e., depending on the overall "global behavior" of the model only [1, 4]. The level of fine granularity in the achieved explanation thus depends on the selected method, model, and evaluation benchmark. Explainability in CV is usually achieved by mapping the relevance maps to a pixel level and treating the evaluation as a weak segmentation task [25, 20, 27]. On the other hand, NLP explanations are much more vaguely defined and usually mixed with more complex philosophical interpretations, such as labeling a given text to a certain sentiment category [8].
The majority of XAI methods can be divided into the classes of perturbation and gradient analysis. Perturbations treat the model as a black box and attempt to derive knowledge of the model's behavior by studying changes in input-output pairs only. Gradient-based methods, on the other hand, execute a backpropagation step towards a target and aggregate the model's parameter adoptions to derive insights. Most of these XAI methods are not motivated by a specific discipline, e.g., neither by NLP nor CV. They are so generic that they can be applied to both disciplines to some extent. However, architecture-specific XAI methods exist as well, such as GradCAM [25], leveraging convolutional neural networks' spatial input aggregation in the deepest layers to increase efficiency.
Explainability in transformers. Through their increasing size, transformers are particularly challenging for explainability methods, especially for architecture-agnostic ones. Transformers' core components include an embedding layer followed by multiple layers of alternating attention and feed-forward blocks Fig. 1b. The attention blocks map the input into separate "query", "key", and "value" matrices and are split into an array of "heads". As with convolutions in CNN networks, separated heads are believed to relate to specific learned features or tasks [13]. Further, the attention mechanism correlates tokens wrt. the input sequence dimension. Care needs to be taken here as GPT-style generative models, in contrast to BERT-style embedding models, apply a causal masking of activations as we will highlight further below. This in particular might obscure results for gradient-based methods.
Consequently, most explainability adoptions to transformers focus on the attention mechanism. Rollout [1] assumes that activations in attention layers are combined linearly and consider paths along the pairwise attention graph. However, while being efficient, it often emphasizes irrelevant tokens, in particular, due to its class-agnostic nature. The authors also propose attention flow, which, however, is unfeasible to compute as it constructs exhaustive graphs. More recently, Chefer et al. (2021) proposed to aggregate backward gradients and LRP [20] throughout all layers and heads of the attention modules in order to derive explanation relevancy. Their introduced method outperforms previous transformer-specific and unspecific XAI methods on several benchmarks and transformer models. It is further extended to multimodal transformers [7] by studying other variations of attention. However, they evaluated only on classification tasks, despite autoregressive transformers' remarkable generative performance, e.g., utilizing InstructGPT [21] or multimodal transformers such as MAGMA [10], BLIP [18] and OFA [31].
Multimodal transformers. Contrarily to these explainability studies evaluated on object-detection models like DETR and ViT [5, 9], we study explainability on open-vocabulary tasks, i.p., generated text tokens of a language model, and not specifically trained classifiers. Due to the multimodality, the XAI method should produce output relevancy either on the input text or the input image as depicted in Fig. 1a. Specifically, to obtain image modality in MAGMA [10], the authors propose to fine-tune a frozen pre-trained language model by adding sequential adapters to the layers, leaving the attention mechanism untouched. It uses a CLIP [22] vision encoder to produce image embeddings. These embeddings are afterward treated as equal during model execution, particularly regarding
other modality input tokens. This multi-modal open-vocabulary methodology has shown competitive performance compared to single-task solutions.
3 ATMAN: Attention Manipulation
We formulate finding the best explainability estimator of a model as solving the following question: What is the most important part of the input, annotated by the explanator, to produce the model's output? In the following, we derive our perturbation probe mathematically through studies of influence functions and embedding layer updates on autoregressive (AR) models [14, 2].
Then we show how attention manipulation on single tokens can be used in NLP tasks to steer a model's prediction in directions found within the prompt. Finally, we derive our multi-modal XAI method ATMAN by extending this concept to the cosine neighborhood in the embedding space.
3.1 Influence functions as explainability estimators
Transformer-based language models are probability distribution estimators. They map from some input space  $\mathcal{X}$ , e.g., text or image embeddings, to an output space  $\mathcal{Y}$ , e.g., language token probabilities. Let  $\mathcal{E}$  be the space of all explanations, i.e., binary labels, over  $\mathcal{X}$ . An explainer function can then be defined as  $e: (\mathcal{X} \to \mathcal{Y}) \times (\mathcal{X} \times \mathcal{Y}) \to \mathcal{E}$ , in words, given a model, an input, and a target, derive a label on the input.
Given a sequence of words  $\mathbf{w} = [w_1, \dots, w_N] \in \mathcal{X}^N$ , an AR language model assigns a probability to that sequence  $p(\mathbf{w})$  by applying factorization  $p(\mathbf{w}) = \prod_t p(w_t | w_{<t})$ . The loss optimization during training can then be formalized as solving:
$$
\begin{array}{l} \max  _ {\theta} \log p _ {\theta} (\mathbf {w}) _ {\text {t a r g e t}} = \sum_ {t} \log p _ {\theta} \left(w _ {t} \mid w _ {<   t}\right) _ {\text {t a r g e t} ^ {t}} = \sum_ {t} \log \operatorname {s o f t m a x} \left(h _ {\theta} \left(w _ {<   t}\right) W _ {\theta} ^ {T}\right) _ {\text {t a r g e t} ^ {t}} (1) \\ =: - \sum_ {t} L ^ {\text {t a r g e t}} (\mathbf {w} _ {<   t}, \theta) =: - L ^ {\text {t a r g e t}} (\mathbf {w}, \theta). (2) \\ \end{array}
$$
Here  $h_\theta$  denotes the model,  $W_\theta$  the learned embedding matrix, and  $target^t$  the vocabulary index of the  $t - th$  target token, with length  $|target| = N$ . Eq. 1 is derived by integrating the cross-entropy loss, commonly used during language model training with  $target = \mathbf{w}$ . Finally,  $L$  denotes our loss function.
Perturbation methods study the influence of the model's predictions by adding small noise  $\epsilon$  to the input and measuring the prediction change. We follow the results of the studies [14, 2] to approximate the perturbation effect directly through the model's parameters when executing Leaving-One-Out experiments on the input. The influence function estimating the perturbation  $\epsilon$  of an input  $z$  is then derived as:
$$
\mathcal {L} ^ {\text {t a r g e t}} \left(z _ {\epsilon}, z\right) = \frac {d L ^ {\text {t a r g e t}} (z , \theta_ {\epsilon})}{d \epsilon} | _ {\epsilon = 0} \approx L ^ {\text {t a r g e t}} \left(z, \theta_ {- z _ {\epsilon}}\right) - L ^ {\text {t a r g e t}} (z, \theta). \tag {3}
$$
Here  $\theta_{-z_{\epsilon}}$  denotes the set of model parameters in which  $z_{\epsilon}$  would not have been seen during training. In the following, we further show how to approximate  $\theta_{-z_{\epsilon}}$ .
3.2 Single token attention manipulation
The core idea of ATMAN is the shift of the perturbation space from the raw input to the embedded token space. This allows us to reduce the dimensionality of possible perturbations down to a single scaling factor per input token. Moreover, we do not manipulate the value matrix of attention blocks and therewith do not introduce the otherwise inherent input-distribution shift of obfuscation methods. By manipulating the attention scores at the positions of the corresponding input sequence tokens, we are able to interpolate the focus of the prediction distribution of the model—amplifying or suppressing concepts of the prompt resulting in an XAI method as described in the following.
Attention was introduced in [28] as:  $\mathbf{O} = \mathrm{softmax}(\mathbf{H})\cdot \mathbf{V}$ , where  $\cdot$  denotes matrix multiplication. The pre-softmax query-key attention scores are defined as
$$
\mathbf {H} = \mathbf {Q} \cdot \mathbf {K} ^ {T} / \sqrt {d _ {h}}.
$$
In the case of autoregression, a lower left triangular unit mask  $\mathbf{M}$  is applied to these scores as  $\mathbf{H}_M = \mathbf{H} \circ \mathbf{M}$ , with  $\circ$  the Hadamard product. The output of the self-attention module is  $\mathbf{O} \in$
(a) Illustration of the perturbation method.
(b) Steering through ATMAN.
Figure 2: (a) Illustration of the proposed explainability method. First, we collect the original cross-entropy score of the target tokens (1). Then we iterate and suppress one token at a time, indicated by the red box  $(\square)$ , and track changes in the cross-entropy score of the target token (2). (b) Manipulation of the attention scores, highlighted in blue, steers the model's prediction into a different contextual direction. Note that we found measuring of such proper generative directions to perform better than radical token-masking, as we show later (c.f. Fig. 11). (Best viewed in color.)
$\mathbb{R}^{h\times s\times d_h}$ , the query matrix is  $\mathbf{Q}\in \mathbb{R}^{h\times s\times d_h}$  and  $\mathbf{K}$ ,  $\mathbf{V}\in \mathbb{R}^{h\times s\times d_h}$  the keys and values matrices. Finally  $\mathbf{H},\mathbf{M},\mathbf{H}_M\in \mathbb{R}^{h\times s\times s}$ . The number of heads is denoted as  $h$ , and  $d_{h}$  is the embedding dimension of the model. Finally,  $s$  is the length input-sequence.
The perturbation approximation  $\theta_{-z_{\epsilon}}$  required by Sec. 3.1 can now be approximated through attention score manipulation as follows: Let  $\mathbf{w}$  be an input token sequence of length  $|\mathbf{w}| = n$ . Let  $i$  be a token index within this sequence to be perturbed by a factor  $f$ . For all layers and all heads  $\mathbf{H}_u$  we modify the pre-softmax attention scores as:
$$
\widetilde {\mathbf {H}} _ {u, *, *} = \mathbf {H} _ {u, *, *} \circ \left(\mathbf {1} - \mathbf {f} ^ {i}\right), \tag {4}
$$
where  $\mathbf{1} \in [1]^{s \times s}$  denotes the matrix containing only ones and  $\mathbf{f}^i$  the suppression factor matrix for token  $i$ . In this section we set  $\mathbf{f}_{k,*}^i = f$ , for  $k = i$  and  $f \in \mathbb{R}$  and 0 elsewhere. As depicted in Fig. 1b, we thus only amplify the  $i-th$  column of the attention scores of  $H$  by a factor  $(1 - f)$ —this, however, equally for all heads. We denote this modification to the model as  $\theta_{-i}$  and assume a fixed factor  $f$ .<sup>1</sup> Note that this position in the architecture is somewhat predestined for this kind of modification, as it already applies the previously mentioned causal token-masking.
We define the explanation for a class label target as the vector of the positional influence functions:
$$
\mathcal {E} (\mathbf {w}, \text {t a r g e t}) := \left(\mathcal {I} ^ {\text {t a r g e t}} \left(\mathbf {w} _ {1}, \mathbf {w}\right), \dots , \mathcal {I} ^ {\text {t a r g e t}} \left(\mathbf {w} _ {n}, \mathbf {w}\right)\right), \tag {5}
$$
with  $\mathcal{I}^{target}$  derived by Eq. 2,3 as  $\mathcal{I}^{target}(\mathbf{w}_i,\mathbf{w})\coloneqq L^{target}(\mathbf{w},\theta_{-i}) - L^{target}(\mathbf{w},\theta)$ . In words, we average the cross-entropy of the AR input sequence wrt. all target tokens and measure the change when suppressing token index  $i$  to the unmodified one. The explanation becomes this difference vector of all possible sequence position perturbations and thus requires  $n$  forward passes.
Fig. 2a illustrates this algorithm. The original input prompt is the text "Ben likes to eat burgers and" for which we want to extract the most valuable token for the completion of the target token "fries". Initially, the model predicts the target token with a cross-entropy score of 1.72. Iterated suppression over the input tokens, as described, leads to "burgers" being the most-influential input token with the highest score of 14.2 for the completion.
Token attention suppression steers the model's prediction. Intuitively, for factors  $0 < f < 1$ , we call the modifications "suppression", as we find the model's output now relatively less influenced by the token at the position of the respective manipulated attention scores. Contrarily,  $f < 0$  "amplifies" the influence of the manipulated input token on the output.
An example of the varying continuations, when a single token manipulation is applied, can be seen in Fig. 2b. We provide the model a prompt in which the focus of continuation largely depends on two tokens, namely "soccer" and "math". We show how suppressing and amplifying them alters the prediction distributions away from or towards those concepts. It is precisely this distribution shift we measure and visualize as our explainability.
<sup>1</sup>We ran a parameter sweep once (c.f. App. A.4) and fixed this parameter to 0.9 for this work.
(i)
Select
(a) Correlated token suppression on images.
Figure 3: (a) Correlated token suppression of ATMAN enhances explainability in the image domain. i) Shows an input image along with three perturbation examples  $(A, B, C)$ . In  $A$ , we only suppress a single image token (blue). In  $B$ , the same token with its relative cosine neighborhood (yellow), and in  $C$ , a non-related token with its neighborhood. Below depicted are the changes in the cross-entropy loss. The original score for the target token "panda" is denoted by  $c_{0}$  and the loss change by  $\Delta$ . ii) Shows the resulting explanation without Cosine Similarity (CS) and with CS. We evaluated the influence of the CS quantitatively in Fig. 12a. (b) An example of the SQuAD dataset with ATMAN explanations. The instance contains three questions for a given context, each with a labeled answer pointing to a fragment of the context. ATMAN is used to highlight the corresponding fragments of the text responsible for the answer. It can be observed that the green example is full, the blue in part, and the yellow is not at all recovered according to the given labels. However, the yellow highlight seems at least related to the label. (Best viewed in color.)
Cosine Sin
${13}/{14}$
cross entropy
#
$\Delta_{1} + \Delta_{2}$
Without CS
Context with AtMan Explanation
Question
With CS
From which countries did the Norse originate?
When were the Normans in Normandy?
Denmark, Iceland and Norway
10th and 11th centuries
In what country is Normandy located?
France
(b) SQuAD example with ATMAN Explanations.
3.3 Correlated token attention manipulation
Suppressing single tokens works well when the entire entropy responsible for producing the target token occurs only once. However, for inputs with redundant information, this approach would often fail. This issue is especially prominent in Computer Vision (CV), where information like objects in an image is often spread across multiple tokens—as the image is usually divided into patches, each of which is separately embedded. It is a common finding that applied cosine similarity in the embedding space, i.p., for CLIP-like encodings as used in MAGMA, gives a good correlation estimator [17, 2]. This is in particular due to the training by a contrastive loss. We integrate this finding into ATMAN in order to suppress all redundant information corresponding to a particular input token at once, which we refer to as correlated token suppression.
Fig. 3a summarizes the correlated token suppression visually. For  $n$  input tokens and embedding dimension  $d$ , the embedded tokens result in a matrix  $T = (t_i) \in \mathbb{R}^{n \times d}$ . The cosine similarity is computed from the normalized embeddings  $\widetilde{T} = (\widetilde{t}_i)$ , with  $\widetilde{t}_i = t_i / ||t_i||$ , for  $i \in 1, \ldots, n$ , as  $S = (s_i) = \widetilde{T} \cdot \widetilde{T}^\top \in [-1,1]^{n \times n}$ . Note that the index  $i$  denotes a column corresponding to the respective input token index. Intuitively, the vector  $s_i$  then contains similarity scores to all input tokens. In order to suppress the correlated neighborhood of a specific token with the index  $i$ , we, therefore, adjust the suppression factor matrix for Eq. 4 as