Source-linked AI summary
XAI for Transformers: Better Explanations through Conservative Propagation
Ameen Ali, Thomas Schnake, Oliver Eberle, Grégoire Montavon, Klaus-Robert Müller, Lior Wolf
TL;DR
Transformers need reliable explanations because their complexity makes predictions difficult to interpret, especially in sensitive applications. The paper analyzes conservation in Transformer explanations and proposes LRP-based propagation rules for attention heads and LayerNorm. The method achieves state-of-the-art results across Transformer and Graphormer models and datasets, while the bias-analysis use case has an out-of-distribution input risk.
Problem
Transformer complexity makes predictions difficult to interpret, while existing XAI methods may not preserve conservation when applied to Transformer models.
Method
The paper extends LRP to Transformers by analyzing conservation and introducing specific propagation rules for attention heads and LayerNorm, including strategic detachment of part of the forward computation.
Results
The method systematically achieves state-of-the-art results on Transformer and Graphormer models across a broad range of datasets and outperforms existing baselines in most tasks.
Takeaways & Limitations
The explanation technique can characterize sentiment-model bias in detail without generating counterfactual examples or risking departure from the data manifold.
Takeaways & Limitations
Systematic bias analyses may produce inputs outside the training distribution, which can cause unstable predictions.
Abstract
from arXiv · showhide
Transformers have become an important workhorse of machine learning, with numerous applications. This necessitates the development of reliable methods for increasing their transparency. Multiple interpretability methods, often based on gradient information, have been proposed. We show that the gradient in a Transformer reflects the function only locally, and thus fails to reliably identify the contribution of input features to the prediction. We identify Attention Heads and LayerNorm as main reasons for such unreliable explanations and propose a more stable way for propagation through these layers. Our proposal, which can be seen as a proper extension of the well-established LRP method to Transformers, is shown both theoretically and empirically to overcome the deficiency of a simple gradient-based approach, and achieves state-of-the-art explanation performance on a broad range of Transformer models and datasets.
1. Introduction
Transformers offer strong performance across multiple domains but are difficult to interpret because of their high complexity. This work develops theoretically grounded explanations by analyzing conservation in Transformer modules and evaluates the proposed method across domains.
- Motivation: Transformer models achieve strong performance in NLP, vision, and graph applications but can contain billions of parameters, making their predictions difficult to access.Their complexity motivates methods for explaining decisions in potentially sensitive applications.
- Contribution: The paper studies how to bring theoretically sound Explainable AI to Transformers, focusing on the conservation axiom within the LRP framework.LRP enables conservation analysis at the level of individual Transformer modules and layers.
- Evaluation: The authors report qualitative and quantitative experiments spanning natural language understanding, computer vision, and graph analysis against state-of-the-art Transformer explanation baselines.Quantitative evaluation uses input perturbations to track model behavior when relevant or irrelevant features are added or removed.
- Results: The proposed method outperforms existing baseline methods in most tasks while providing theoretical soundness in its formal derivation.The paper also states that the code is publicly available.
2. Related Work
Prior Transformer explanation methods include attention extraction, gradient-based attribution, perturbation-based analysis, and LRP-based approaches. These lines of work differ in how they derive explanations from attention, gradients, relevance propagation, or input changes.
- Methods: Transformer explanation methods are grouped into attention-head extraction, gradient-based methods, and perturbation-based methods.This taxonomy organizes the principal approaches reviewed in the section.
- Attention-based methods: Attention-based approaches include raw attention extraction, Attention Rollout, and Attention Flow, but raw attention is described as myopic and unreliable.The cited work extends attention-based interpretation beyond directly visualizing attention weights.
- Gradient-based methods: Gradient-based explanations applied to Transformers include integrated gradients, input gradients, Saliency, Gradient × Input, and Guided Backpropagation.These methods have also been used across numerous models and domains.
- LRP-based methods: LRP has been implemented for Transformers and other attention-based models, including applications to NLP models such as BERT.Other work also studies LRP and gradient propagation within Transformer blocks.
- Perturbation-based methods: Perturbation-based approaches identify relevant input parts by observing changes in model confidence or using Shapley values.Related input-reduction methods have been applied to explain model predictions.
- Graph Transformers: The paper evaluates both conventional Transformers and the Graphormer model, alongside prior graph-Transformer explainability work based on attention maps.Related graph explanation methods also include gradients, decomposition, perturbations, and local surrogate models.
- Broader XAI context: The section situates these methods within broader reviews of Explainable AI.The cited reviews cover general XAI methods and applications.
3. A Theoretical View on Explaining Transformers
The paper analyzes Transformer explanations through the conservation axiom and shows that Gradient × Input can violate conservation in attention heads and LayerNorm. These violations motivate improved propagation rules that preserve conservation and perform better empirically.
- 3. A Theoretical View on Explaining Transformers: Conservation requires input relevance scores to sum to the network output, with LRP enforcing this property layer by layer.The paper frames Transformer explanation using the axiomatic approach developed for explainability in standard deep networks.
- 3. A Theoretical View on Explaining Transformers: Gradient × Input can be viewed as an LRP relevance-propagation rule, enabling layer-wise tests of whether conservation holds.This interpretation provides a basis for identifying components where Gradient × Input loses conservation and for deriving improved rules.
- 3. A Theoretical View on Explaining Transformers: Two Transformer components cause significant conservation breaks under Gradient × Input: attention heads and LayerNorm.The analysis therefore focuses on these architecture-specific components rather than treating the Transformer as uniformly linear for attribution.
- 3.1. Propagation in Attention Heads: In attention heads, conservation may fail because the gating distribution depends on the input, creating covariates between the matching function and token inputs.The proposition derives a conservation equation under centered inputs, while the paper notes that such dependence is likely because the matching function is itself input-dependent.
- 3. A Theoretical View on Explaining Transformers: The proposed alternative propagation rule addresses Gradient × Input’s lack of conservation and is intended to work better empirically while retaining conservation.For attention heads, the method locally treats gating terms as constants and applies the canonical LRP rule for linear layers.
- 3.2. Propagation in LayerNorm: LayerNorm never satisfies conservation for Gradient × Input, with especially strong breaks when ϵ is small relative to Var[x].The resulting relevance collapse is illustrated in Figure 2, including the region to the right of the peak activation.
4. Better LRP Rules for Transformers
The method extends LRP to Transformers by replacing attention-head and LayerNorm propagation with locally linear, conservative rules. An implementation using detached gating and rescaling terms computes the resulting explanations through standard Gradient × Input.
- Conservative propagation: The approach replaces Gradient × Input’s implicit attention-head and LayerNorm rules with ad-hoc propagation rules designed to conserve relevance.The method starts from the LRP interpretation of Gradient × Input and targets the two components identified as breaking conservation.
- Attention heads: Attention-head propagation treats gating terms p_ij as constants, enabling the canonical linear-layer LRP rule and assigning zero relevance to the disconnected query sequence.The locally linear view interprets the gating terms as weights mapping the input sequence to the output sequence.
- LayerNorm: LayerNorm propagation treats α = (ϵ + Var[x])−1 as constant, expressing the operation as the linear transformation αCx with centering matrix C.The matrix C centers the input sequence, whose length is N.
- Implementation: The implementation applies the modified rules by detaching gating and rescaling terms, then running standard Gradient × Input on the rewritten function.The required changes are detach() calls at appropriate network locations.
- Implementation: The implementation is at least as fast as Gradient × Input and typically faster because the modified computation simplifies gradient evaluation.The authors describe the runtime comparison as applying to the proposed improved LRP explanation.
5. Experimental Setup
The experiments benchmark the proposed LRP variants against established explanation methods across language, image, and molecular or graph Transformer tasks. The setup includes sentiment, tweet, digit, and molecular binding datasets.
- Benchmark design: The study evaluates the proposed approach on several Transformer models and datasets against explanation methods from the literature.The benchmark compares methods commonly used for Transformer-type architectures.
- Natural language processing: NLP evaluation covers binary sentiment classification on SST-2 and IMDB, plus tweet classification with TweetEval.SST-2 contains 11,844 reviews and IMDB contains 50,000 reviews.
- Graph and molecular tasks: Graph Transformer evaluation uses MNIST superpixels for digit classification, with each image patch represented as a node connected to neighboring patches.The MNIST superpixels data is extracted from 70,000 samples.
- Graph and molecular tasks: The molecular evaluation uses BACE, containing 1,522 compounds with structures and binary labels for inhibitor binding to human β-secretase 1.BACE is drawn from the MoleculeNet benchmark.
- Compared methods: Baselines include Gradient×Input, Attention-last, Rollout, attention flow, and Generic Attention Explainability.Attention-last, Rollout, and attention flow account for layer-wise Transformer structure, while GAE propagates attention gradients with gradients from other network components.
- Proposed variants: The proposed method is tested in LRP (AH), LRP (LN), and LRP (AH+LN) variants that modify attention-head propagation, LayerNorm propagation, or both.Other layers use the Gradient × Input-equivalent propagation rule in these variants.
6. Results
The evaluation validates conservation, perturbation performance, qualitative attribution quality, and a bias-analysis use case across Transformer and Graphormer models. The proposed LRP variants better conserve output relevance, improve perturbation results, focus on informative features, and support analysis of biased responses.
- 6.1. Conservation: LRP explanations tracked Transformer output scores much more closely than GI, although mild conservation breaks remained.GI explanations on MNIST were almost anticorrelated with the model output.
- 6.2. Quantitative Evaluation: Activation evaluates whether highly relevant nodes rapidly activate the correct prediction, whereas pruning measures whether removing low-relevance nodes preserves it.AUAC is higher-is-better; AU-MSE is lower-is-better.
- 6.2. Quantitative Evaluation: The proposed attention-head and LayerNorm gradient handling consistently improved activation and pruning performance across datasets.Combining both proposed mechanisms produced the greatest improvement over naive gradient implementations.
- 6.3. Qualitative Results: On SST-2, the proposed LRP methods assigned stronger relevance to sentiment-bearing words and less relevance to the entity token “eastwood” than A-last.The methods highlighted “best” and “virtues” while reducing the entity-focused attribution seen with A-last.
- 6.3. Qualitative Results: On MNIST superpixels, LRP (AH) and LRP (AH+LN) more reliably highlighted digit information than competing methods.The qualitative comparison linked raw attention methods’ weaker perturbation performance to relevance assigned to uninformative background superpixels.
- 6.4. Use Case: Analyzing Bias in Transformers: The bias-analysis use case found no consistent female-versus-male name bias but identified entity categories that strongly modulated sentiment.The attribution approach separates individual word contributions from other words, reducing confounding relative to name-occurrence correlation.
7. Conclusion
Transformers require transparent decision explanations, but common XAI methods cannot be assumed to work reliably on them. The paper identifies failures in gradient-based conservation, proposes layer-specific LRP rules, and reports state-of-the-art results across Transformer and Graphormer datasets.
- Transformers’ widespread practical use creates a need for transparency into their decisions.
- Common XAI methods do not necessarily work well on Transformers as they do on standard deep neural networks.
- Gradient × Input fails to implement conservation, while attention heads and LayerNorm require specific treatment.
- The paper proposes specific, easily implemented LRP rules for attention heads and LayerNorms.
- The method systematically achieves state-of-the-art results on Transformer and Graphormer models across a broad range of datasets.
- The explanation technique characterizes sentiment-model bias without generating counterfactual examples or risking departure from the data manifold.
A. Derivations for Attention Heads
The derivations analyze gradient propagation through attention heads and LayerNorm, including centering and rescaling, to establish the resulting conservation behavior. The combined conservation equation holds only approximately for large values of ϵ.
- The derivations begin by differentiating the softmax attention function with respect to tokens from both input sequences.
- The attention-head analysis applies the multivariate chain rule to propagate gradients into tokens of both sequences.
- Relevance scores for the two sequences are computed separately and then summed to obtain the combined result.
- Assuming zero expected values for the first token sequence and softmax input yields a simplified conservation expression.
- LayerNorm is decomposed into centering followed by rescaling, with expectation computed over a uniform distribution.
- The resulting conservation equation for LayerNorm holds only approximately for large values of ϵ.
D. Additional Conservation Experiments
Additional conservation experiments evaluate GAE, Rollout, and A-Flow on SST-2. Conservation does not hold for these baseline explanation methods.
- The experiments test conservation for three additional baselines: GAE, Rollout, and A-Flow.
- On SST-2, conservation does not hold for the evaluated baseline explanation methods.
- Figure 8 presents the conservation experiment for the additional explainability baselines.
E. Runtime Analysis
The runtime analysis compares explanation methods on SST-2 and finds the proposed method competitive with state-of-the-art explanation techniques. LRP(AH + LN) is faster than prediction in the reported measurements.
- On SST-2, LRP(AH + LN) requires 0.012 seconds per explanation, compared with 0.024 seconds for prediction.
- The proposed method is competitive in runtime with state-of-the-art explanation techniques.
H. Perturbation Experiments
The perturbation experiments reference two Twitter evaluation settings: Twitter-Sentiment and Twitter-Hate.
- Twitter-Sentiment is named as an evaluation setting.
- Twitter-Hate is named as an evaluation setting.
- The passage lists Twitter-Sentiment and Twitter-Hate together under the perturbation experiments.