Self-Attention Attribution: Interpreting Information Interactions Inside Transformer

TL;DR

The paper introduces AttAttr, a self-attention attribution method based on integrated gradients to interpret information interactions inside Transformer models. It identifies important attention connections, constructs attribution trees to visualize information flow, and demonstrates practical uses in attention head pruning and adversarial trigger discovery on BERT. The approach reveals that attention weights do not always align with predictive contributions, enables robust head selection across homogeneous tasks, and uncovers vulnerabilities through adversarial patterns. Together, these contributions advance interpretability, model debugging, and robustness assessment for Transformer-based systems.

Abstract

The great success of Transformer-based models benefits from the powerful multi-head self-attention mechanism, which learns token dependencies and encodes contextual information from the input. Prior work strives to attribute model decisions to individual input features with different saliency measures, but they fail to explain how these input features interact with each other to reach predictions. In this paper, we propose a self-attention attribution method to interpret the information interactions inside Transformer. We take BERT as an example to conduct extensive studies. Firstly, we apply self-attention attribution to identify the important attention heads, while others can be pruned with marginal performance degradation. Furthermore, we extract the most salient dependencies in each layer to construct an attribution tree, which reveals the hierarchical interactions inside Transformer. Finally, we show that the attribution results can be used as adversarial patterns to implement non-targeted attacks towards BERT.

Figures (7)

• Figure 1: Attention score (left) and attribution score (right) of a single head in BERT. The color is darker for larger values. The model prediction for the sentence from MNLI dataset is contradiction. AttAttr tends to identify more sparse word interactions that contribute to the final model decision.
• Figure 2: Effectiveness analysis of AttAttr. The blue and red lines represent pruning attention heads according to attribution scores, and attention scores, respectively. The solid lines mean the attention heads with the smallest values are pruned first, while the dash lines mean the largest values are pruned first. The results show that AttAttr better indicates the importance of attention heads.
• Figure 3: Evaluation accuracy as a function of head pruning proportion. The attention heads are pruned according to the accuracy difference (baseline; dash yellow), the Taylor expansion method (are16headbetterthan1; solid red), and AttAttr (this work; solid blue).
• Figure 4: Correlation of attribution scores of different attention heads between datasets. Each point represents the attribution scores of a single attention head on two datasets. The datasets of homogeneous tasks are strongly correlated, which implies the same subset of attention heads are fine-tuned for similar tasks.
• Figure 5: Examples of attribution trees. (a) is from MNLI, which is predicted as entailment by BERT. (b) is from SST-2, which is predicted as positive by BERT. The grey words from the inputs do not appear in the attribution trees.