AtMan: Understanding Transformer Predictions Through Memory Efficient Attention Manipulation

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 extra memory, since they rely on backpropagation which allocates almost twice as much GPU memory as the forward pass. This makes it difficult, if not impossible, to use them 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 based 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 while being computationally efficient. As such, AtMan is suitable for use in large model inference deployments.

Paper

References (35)

Scroll for more · 23 remaining

Similar papers

Peer review

Reviewer wnDz5/10 · confidence 4/52023-07-04

Summary

The paper introduces AtMan, a method that manipulates the attention mechanism of generative transformer models, to provide explanations. More specifically, it perturbs raw attention scores instead of the token inputs. The authors provide extensive experiments for different modalities to back their claims regarding better results to established baselines.

Strengths

1. The paper introduces a new approach towards Explainable AI, that does not require calculation of gradients through the network, thus enabling evaluation of larger models with the same hardware. 2. The idea of manipulating attention scores is novel and interesting. 3. The authors embark on an extensive evaluation of their method on different modalities. The universality of the Transformer architecture, enabled them to seamlessly switch between modalities. The methods performs well on a range of architectures (encoder, decoder).

Weaknesses

1. One of the main motivation of this work is to enable more widespread adoption of explainable AI methods, via reducing required computational resources. The issue of memory is discussed extensively. Inference speed is only mildly explored, in Figure 5. For large sequence lengths, the method requires upwards of 2 orders of magnitude more time to make the final predictions. 2. The method requires one to set new hyperpatameres, namely $f$ and $\kappa$. The effect of those with respect to e.g. sequence length or model architecture, are not discussed. For example, different sequence lengths may lead to different softmax entropy and thus different $f$ values might be required. 3. According to Eq. (4), tokens are masked by setting the corresponding $H$ values to 0. This does by no means lead to masking of the tokens, as briefly explained in Appendix A5. Setting the values to the value $0$ constitutes another hyperparameter, that needs to be set for the proposed method. In principle any real value could have been chosen (In Appendix A5 the authors explore the values $0$ and $-\infty$). The range of the possible values was not explored at all. 4. In general a lot of approximations are made for the derivations of the last method, with possible alternatives not discussed. Why was Eq. (4) chosen to mask tokens? Why is the cosine similarity of the embeddings a good measure to suppress common tokens? In [1] for example, the cosine similarity between keys is used as a more suitable measure. The cosine similarity is calculated based on the embeddings of what layer? Is there a benefit from calculating this similarity on different layers? There is a lot of work on what type of information token embeddings carry for different layers in a Vision Transformer. [1] Bolya, Daniel, et al. "Token merging: Your vit but faster." arXiv preprint arXiv:2210.09461 (2022).

Questions

1. Line 62 states that the source code is released although I could not find it in the supplementary. 2. Line 33-34 states that memory consumption leads to uneconomical productive deployment. This is a overstatement, especially given the computational requirements of the new proposed method. 3. A lot of literature has been exploring the effect of attention scores in the context of Stable Diffusion, e.g. [2] and follow-ups. 4. Is there a minus missing before $L^{\text{target}}$ in Eq. (2)? 5. Line 245, it is not clear to me how the mean of generated scores affects the final predictions. Are some tokens in the target label more important than others as expected? 6. Is there a way to prioritize checking some tokens first instead of exhaustively trying out all tokens to find the ones that explain better the target? This might help to reduce the inference speed. You can think about having a method to rank tokens to check, and then create a plot of computational budget vs performance, where you check iteratively more tokens. You could also think of doing this in a more hierarchical way, similar to what is proposed in Appendix A8. [2] Hertz, Amir, et al. "Prompt-to-prompt image editing with cross attention control." arXiv preprint arXiv:2208.01626 (2022).

Rating

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

Confidence

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

Soundness

3 good

Presentation

4 excellent

Contribution

2 fair

Limitations

The paper proposes a way to manipulate attention scores to explain prediction of Transformer models. Manipulating these scores is done a token at a time, leading to inference speed that scales with the sequence length, and is orders of magnitude bigger than previous approaches. Although a lot of experiments are presented, some of the design choices are not adequately justified.

Reviewer FKKY6/10 · confidence 3/52023-07-06

Summary

The authors proposed a new resource-friendly modality-agnostic perturbation method named AtMan to explain the predictions of generative transformers. AtMan suppresses or amplifies the attention entry relying on the cosine similarity neighborhood to evaluate the influence of each token, and it saves many memory resources compared with current gradient-based methods, which means that it's more suitable for large model inference developments. Empirical results show that AtMan can do language/visual reasoning better than the current state-of-the-art gradient-based methods.

Strengths

1. AtMan is simple but useful, and the idea that changing perturbation space from the raw input to the embedding space is important. It can perform on decoder as well as encoder architectures and the empirical results are supportive. 2. The core contribution of AtMan is that it reduce the memory cost and can be applied to the large models, which will be important for the future works of explainability of large models. Now it can be applied to 30B multi-modal model but previous methods fail. 3. The paper is well written and easy to follow.

Weaknesses

1. The paper mentions that the $\kappa$ parameter in AtMan is set to be $0.7$, but it may be important for the performance of AtMan, so it is better to do additional experiments to show the effect of different kappa on the experimental results on different datasets. Furthermore, the choice of (6) is not unique, there are some other ways like using the softmax value of $s_{i,k}$ to be $f_{k,*}^i$ and using temperature as a hyper-parameter. It's better to compare these ways. 2. The increase in the time cost of AtMan is significant compared with gradient-based methods, so it's better to show the time-cost figure after pipeline-parallel execution. 3. The results of the locality are important, maybe the authors can try to add them (like figure 10 in Appendix) into the main paper.

Questions

Similar to Weakness-1 and Weakness-2.

Rating

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

Confidence

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

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The author mentions that their work provides reference methods for future work, including explanatory studies of current generative models impacting our society. The author doesn't mention other potential negative societal impacts.

Reviewer sfjw6/10 · confidence 4/52023-07-07

Summary

This paper proposed ATMAN, a modality and architecture-agnostic perturbation-based XAI method for generative transformer models. Specifically, ATMAN evaluates the sensitivity of model prediction w.r.t. the perturbation of attention scores on a set of cosine-distance neighboring tokens. Empirical results suggest that ATMAN can effectively identify the relevant input tokens for target prediction, and outperforms the baseline methods in terms of precision and recall.

Strengths

As far as I know, this paper is the first work that identifies the crucial tokens/pixels by attention perturbation. The idea is novel and looks interesting to me. The authors further apply the cosine similarity to collect the related neighboring tokens and reduce the searching redundancy. It helps us to understand how transformer models process input signals and identify crucial ones.

Weaknesses

1. As perturbing every token to evaluate their sensitivity, ATMAN can incur high time complexity to generate the explanation for transformer models. As suggested by Figure 5, the runtime is much higher than the baseline. Therefore, to reduce the time complexity, would there be any solution to evaluate the token/pixel sensitivity in a higher granularity, like dividing the tokens/pixels into clusters? 2. The paper mentioned steering model prediction by applying the attention scores of identified tokens but I did not find any empirical results to claim the advantages of this approach. Could you provide more details and results or correct me if I am wrong?

Questions

Can you provide more model details about Section 3.1? Which black-box model do you evaluate?

Rating

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

Confidence

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

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The limitation is the high runtime cost when evaluating every input token/pixel. It would be great if the authors can extend the analysis into higher granularity.

Reviewer YuQG7/10 · confidence 2/52023-07-08

Summary

This paper proposed a new permutation method for Transformer explanation, named ATMAN, which directly manipulates attention scores in the hidden layers of the Transformer model rather than permuting the model inputs. Importantly, the authors established a bridge between ATMAN and the relevance between inputs and output predictions. Experimental results demonstrate the effectiveness and efficiency of ATMAN on interpreting Transformer model on different types of data and different model scales.

Strengths

The proposed ATMAN method is novel, and the experimental results shows that ATMAN significantly improve the effectiveness and efficiency over existing perturbation and propagation methods on interpreting Transformer model on different types of data and different model scales.

Weaknesses

NA.

Questions

NA.

Rating

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

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

NA.

Reviewer HHf83/10 · confidence 4/52023-07-13

Summary

This paper proposes atman, a method for understanding the predictions of transformers by masking out the tokenized inputs. The paper evaluates atman on explainability tasks in text pipelines and vision pipelines. The authors argue that atman is superior to gradient-based explanation methods since it uses less memory.

Strengths

The paper is well-written and the contributions are easily understood. The method is simple to implement and understand, and does take less memory footprint than gradient-based methods. The evaluation on explainability benchmarks seem to show that it behaves at par with gradient-based methods. The figures make the method very easy to understand.

Weaknesses

The novelty of the method seems questionable, and it doesn't seem necessary to create an entire "method" with a name to describe the observation that masking out input tokens can be useful in tracing which parts of the input sequence are important. This type of analysis is standard in analyzing the outputs of Transformers, e.g. it is a standard technique in Olsson et al 2022 (https://arxiv.org/abs/2209.11895). To claim that it is a novel method is a bit much in terms of framing. There may be some novelty in the vision domain in looking at the correlated outputs instead of masking individual patches of the input image, but as far as I can tell this methodological change is not ablated. These novelty questions point to a weakness in the evaluation and the discussion of related work. I don't have any problem with a simple method (in fact I like it) - but if the method is a first standard trick that most ML engineers would use to understand what is happening when a Transformer makes predictions, then it suggests a different framing and evaluation. For example: What other "surprisingly simple" interventions can be used to explain predictions? Can you look at embeddings of tokens in the final output? What are the differences (in NLP) between generative-models (BERT) and encoder-only models (GPT)? How are the explanations for vision affected when the correlations are not taken into account? What happens if this intervention is done partway through the backbone instead of at the beginning?

Questions

Suggested list of questions to improve the work: 1. What other "surprisingly simple" interventions can be used to explain predictions? 2. Can you look at embeddings of tokens in the final output? 3. What are the differences (in NLP) between generative-models (BERT) and encoder-only models (GPT)? 4. How are the explanations for vision affected when the correlations are not taken into account? 5. What happens if this intervention is done partway through the backbone instead of at the beginning?

Rating

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

Confidence

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

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

N/a

Reviewer wnDz2023-08-14

Thank you for the additional details and the clarity in the text, these were really helpful. Figure 17 adds to the understanding of the hyperparameters and I think would help to be discussed in the paper itself. Further limitations are further discussed by the authors, namely 1. Inference speed and 2. Numerous simplifications made, guided namely by preliminary empirical results. Regarding (1) authors provide Figure 16 which I find particularly misleading, given that multiple workers and batch-decoding is applied only for their method. Regarding (2) there are numerous choices made are not always intuitive or sufficiently explained. I do however find the merit in the approach and believe the new results have demonstrated some potential for applicability and some benefits compared to baselines. I have updated by score.

Authorsrebuttal2023-08-21

Thank you again, we indeed will add a more sophisticated discussion of Fig 17 to the paper itself. W.r.t. (1), Fig 16: Note that (pipe-)parallel execution on gradient based methods is not applicable at all, as it _multiplies_ the already exceeding memory consumption ---the primary measurement criterion for this study--- and further requires a 'warm-up' phase filling the execution pipeline, to further reduce the computation time. The latter is usually not fully filled by practical model queries. The same technique, however, is inherently applied to plain-forward passes. It as such compares 'practical query execution times' as it was requested, and further demonstrates the computability of the proposed method (at all). Other parallel execution schemes like model-parallelism, on the other hand, would improve both methods equally. We will further clarify the caption. W.r.t. (2), sorry for not being able to fully meet your expectations. Some parts indeed are directed by empirical assessments and provide scope for subsequent research. Your review already helped us rendering the paper more complete. Please let us know about any concrete aspects we could improve our work further with. all the best the authors

Reviewer HHf82023-08-16

Thank you for your detailed response to my questions and concerns, I appreciate the effort. I will be keeping them in mind as we discuss the paper.

Reviewer FKKY2023-08-17

Thank you for your feedback. I'm glad to see that Figures 16 and 17 effectively address my concerns, highlighting an acceptable time-cost trade-off and the reason for choosing the value of kappa. I suggest the authors include these in the revised paper. However, I believe the caption of Figure 16 could be made more understandable by providing clearer descriptions. Besides, the exploration of the token attention score suppression/facilitation discussion is indeed intriguing. I will maintain my positive score on your work and encourage the author to test other relevant algorithms in the following works.

Authorsrebuttal2023-08-21

Thank you again for your review, we will add the discussions to the paper. We will further clarify, i.p. given more space, the caption of Fig16. Kindly find a final remark to Fig 16 as follows: Note that pipe-parallel execution on gradient based methods is not applicable at all, as it multiplies the already exceeding memory consumption ---the criterion in focus of this study--- and further requires a 'warm-up' phase filling the execution pipeline, to further reduce the overall computation time. The latter is usually not fully filled by practical inference queries. The same technique, however, is inherently applied to plain-forward passes, not affected by these limitations. It, as such, compares 'practical query execution times' as it was requested, and further demonstrates the computability of the proposed method (at all). Other parallel execution schemes like model-parallelism, on the other hand, would improve both methods equally. all the best, the authors

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC