AttnLRP: Attention-Aware Layer-Wise Relevance Propagation for Transformers

Reduan Achtibat; Sayed Mohammad Vakilzadeh Hatefi; Maximilian Dreyer; Aakriti Jain; Thomas Wiegand; Sebastian Lapuschkin; Wojciech Samek

Source-linked AI summary

AttnLRP: Attention-Aware Layer-Wise Relevance Propagation for Transformers

Reduan Achtibat, Sayed Mohammad Vakilzadeh Hatefi, Maximilian Dreyer, Aakriti Jain, Thomas Wiegand, Sebastian Lapuschkin, Wojciech Samek

arXiv:2402.05602v2Feb 8, 2024cs.CLcs.AIcs.CVcs.LG

TL;DR

Transformer models can be biased and hallucinate, creating a need for faithful and efficient explanations of their internal reasoning. AttnLRP extends LRP to nonlinear attention and latent representations, and it outperforms existing methods in faithfulness while enabling latent-neuron interpretation and interaction.

Problem
Faithfully attributing the full behavior of black-box transformer models, including latent representations, remains difficult while maintaining computational efficiency.
Method
AttnLRP extends Layer-wise Relevance Propagation within the Deep Taylor framework with rules for nonlinear transformer attention and latent neurons.
Results
AttnLRP consistently outperforms state-of-the-art methods in faithfulness, including a 46% top-1 accuracy improvement over CP-LRP on Mixtral 8x7b.
Takeaways & Limitations
AttnLRP supports accurate attribution of latent neurons inside and outside attention modules, enabling interpretation and real-time interaction with transformer representations.
Takeaways & Limitations
Accurate ViT attributions remain sensitive to the γ-parameter, while quantization and custom GPU kernels require further investigation for reducing memory consumption.

Abstract

from arXiv · show

Large Language Models are prone to biased predictions and hallucinations, underlining the paramount importance of understanding their model-internal reasoning process. However, achieving faithful attributions for the entirety of a black-box transformer model and maintaining computational efficiency is an unsolved challenge. By extending the Layer-wise Relevance Propagation attribution method to handle attention layers, we address these challenges effectively. While partial solutions exist, our method is the first to faithfully and holistically attribute not only input but also latent representations of transformer models with the computational efficiency similar to a single backward pass. Through extensive evaluations against existing methods on LLaMa 2, Mixtral 8x7b, Flan-T5 and vision transformer architectures, we demonstrate that our proposed approach surpasses alternative methods in terms of faithfulness and enables the understanding of latent representations, opening up the door for concept-based explanations. We provide an LRP library at https://github.com/rachtibat/LRP-eXplains-Transformers.

1. Introduction

AttnLRP addresses the difficulty of faithfully and efficiently explaining transformer models by extending LRP to nonlinear attention and latent neurons. The method is presented as outperforming existing approaches while supporting interaction with latent representations.

1. Introduction: Transformer models can produce biased predictions and hallucinations, motivating methods that reveal their latent reasoning process.Attention maps provide useful interaction information but do not explain the full model or other latent components.
1. Introduction: AttnLRP outperforms existing state-of-the-art methods in explanation faithfulness and computational efficiency.The introduction describes this comparison across transformer architectures, including language and vision models.
1. Introduction: AttnLRP extends LRP within the Deep Taylor framework to derive efficient, faithful rules for nonlinear transformer attention.The approach targets transformer components that standard attribution rules do not adequately handle.
1. Introduction: The method attributes latent neurons inside and outside attention modules, enabling insights into their encodings and interaction with model representations.The paper illustrates this capability by identifying relevant neurons and manipulating latent representations.
1. Introduction: The authors provide an efficient, ready-to-use open-source implementation of AttnLRP for transformer models.The implementation is presented as part of the paper’s contributions.

2. Related Work

Prior attribution methods offer partial explanations of transformers but face trade-offs among computational cost, output specificity, latent coverage, and attribution faithfulness. AttnLRP is positioned against these limitations by targeting holistic, efficient attribution across transformer components.

2. Related Work: Perturbation and local-surrogate methods require many forward passes and cannot explain the original model’s hidden representations.Layer-wise latent attribution can require separate computations for each layer, further increasing cost.
2. Related Work: Attention maps capture aspects of data distribution but are not class-specific and do not meaningfully explain final predictions on their own.Subsequent rollout and gradient-weighted variants address class specificity only through approximate procedures.
2. Related Work: Attention-rollout methods are computationally efficient but produce limited-resolution, noisy input maps and cannot attribute hidden latent features beyond softmax outputs.These limitations restrict their coverage of the model and the fidelity of hidden-space explanations.
2. Related Work: Gradient-based methods can suffer from gradient shattering, producing noisy attributions in large or deep transformer models.SmoothGrad and Integrated Gradients are described as unsuccessful at resolving this issue for large transformers.
2. Related Work: Earlier LRP variants modify backpropagation for more reliable attributions but can be numerically unstable for softmax and fail to satisfy relevance conservation in matrix multiplication.These limitations motivate deriving rules specifically for nonlinear transformer operations.

3. Attention-Aware LRP for Transformers

The method builds faithful relevance propagation for transformer operations by decomposing layer functions and extending LRP to nonlinear attention. It combines conservation-oriented attribution with efficient backward propagation across linear, nonlinear, and attention components.

3. Attention-Aware LRP for Transformers: AttnLRP applies Deep Taylor Decomposition successively to transformer operations, deriving efficient and faithful propagation rules for the full architecture.The approach locally linearizes operations into independent contributions and applies the resulting rules throughout the transformer.
3. Attention-Aware LRP for Transformers: LRP decomposes each output into input relevances and propagates them backward through the layered computation graph while preserving relevance across layers.This provides importance values linked to the model output and supports attribution from outputs back to inputs.
3. Attention-Aware LRP for Transformers: Linear layers use the ε-LRP rule, whose stabilizing term absorbs a small amount of relevance and handles cases where the layer output approaches zero.Bias terms and ε can absorb relevance that is not assigned directly to input variables.
3. Attention-Aware LRP for Transformers: The γ-LRP rule is applied to linear layers in vision transformers because gradient shattering produces pronounced attribution noise there, whereas LLMs show no visible noise.This model-specific choice improves the signal-to-noise ratio where the authors observed the problem.
3. Attention-Aware LRP for Transformers: Element-wise nonlinearities use the identity rule because each operation has one input and one output, so incoming relevance is assigned to that single input.The decomposition is therefore the operation itself rather than a multi-input redistribution.
3. Attention-Aware LRP for Transformers: Attention requires dedicated rules because softmax is highly nonlinear and matrix multiplication is bilinear, with relevance propagated through queries, keys, values, and outputs.The attention formulation distinguishes batch and head dimensions, token counts, and embedding dimensions before deriving operation-specific rules.

4. Experiments

The experiments evaluate AttnLRP for faithfulness, plausibility, efficiency, and latent-neuron understanding across language and vision transformers. AttnLRP consistently outperforms comparison methods, scales more efficiently than perturbation, and identifies neurons whose manipulation changes predictions.

4. Experiments: The evaluation measures faithfulness through ordered input perturbations and plausibility through top-1 answer-token accuracy and IoU against SQuAD v2 answer masks.Faithfulness is assessed on ImageNet classification, IMDB classification, and Wikipedia next-word prediction; plausibility uses Mixtral 8x7b and Flan-T5-XL on SQuAD v2.
4.2. Computational Complexity and Memory Consumption (Q2): LRP requires two forward and one backward pass with checkpointing, whereas linear-time perturbation requires NT forward passes.With checkpointing, LRP memory scales logarithmically with the number of layers, while perturbation uses O(1) memory.
4.3. Understanding & Manipulating Neurons (Q3): Manipulating identified neurons changes model outputs toward targeted concepts, including a shift from ice-bear context to sweet, sugary treats.Deactivating neuron #3948 and amplifying neuron #5687 produces the reported prediction change.

5. Conclusion

The conclusion presents AttnLRP as an efficient and faithful extension of LRP for nonlinear attention in language and vision transformers. It emphasizes accurate latent-neuron attribution as a basis for model interaction and interpretation.

5. Conclusion: AttnLRP extends LRP to nonlinear attention with new propagation rules for softmax and matrix multiplication.The method is interpreted through Deep Taylor Decomposition and applies to both LLMs and ViTs.
5. Conclusion: AttnLRP combines simplicity, faithfulness, and efficiency while accurately attributing neurons in latent space, including within attention modules.The conclusion links latent attribution to real-time transformer interaction and interpretation.

Limitations & Open Problems

The paper identifies ViT attribution stability and memory efficiency as open problems. Accurate ViT attributions depend on the γ-parameter, while quantization and custom kernels remain to be investigated.

Limitations & Open Problems: Accurate ViT attributions remain sensitive to the choice of the γ-parameter.The paper identifies this parameter adjustment as crucial for attribution accuracy.
Limitations & Open Problems: The impact of quantization and custom GPU kernels on LRP rules remains open for reducing memory consumption.These are proposed directions for improving the method's memory efficiency.

Impact Statement

The paper frames transformer attribution as a route toward understanding and debugging opaque AI systems, while contrasting LRP with perturbation and gradient-based alternatives. Its broader motivation connects interpretable model behavior with reducing misleading outputs and biases.

Impact Statement: AttnLRP is positioned as a foundation for analyzing and debugging transformer systems to help reduce hallucinations and mitigate biases.The paper identifies healthcare and finance as potential future critical-domain applications where explainability is important.
Impact Statement: Input × Gradient locally linearizes the model by weighting gradients with input features, but gradient shattering can produce noisy heatmaps.The paper notes that this noisiness can make explanations not meaningful in many cases.
Impact Statement: Self-attention rollout propagates token relevance through averaged attention matrices and adds an identity matrix to account for residual connections.The method uses iterative left multiplication to relate hidden-layer tokens to original input tokens.

A.2.1. CONSERVATION & NUMERICAL STABILITY OF BIAS TERMS

This section examines how bias-term handling affects relevance conservation and numerical stability in LRP. It compares distributing, retaining, or omitting bias relevance and identifies instability risks.

Bias-term strategies: Bias relevance can be distributed uniformly across inputs, handled with the identity rule, or omitted to enforce strict conservation.These choices determine whether the bias absorbs relevance and whether relevance is propagated only through input variables.
Stability condition: For functions with f(0) = 0, zero output activations receive zero relevance and avoid the described instability.The instability is instead associated with assigning non-zero relevance to zero-valued inputs.
Bias omission: Omitting the bias term can produce sign flips when the bias-free linearization has a different sign from the original function.This preserves a strict conservation formulation but requires checking the signs of the resulting relevance contributions.
Numerical stability: Both uniform bias distribution and identity-rule handling can cause severe numerical instabilities during relevance propagation.The issue arises when non-zero relevance is propagated to an input variable whose activation is zero.

A.2.2. HIGHLIGHTING THE DIFFERENCE BETWEEN VARIOUS LRP METHODS

This section contrasts LRP variants for transformer components, focusing on conservation, attribution quality, and numerical behavior. AttnLRP introduces component-specific rules intended to improve these properties.

Softmax attribution comparison: AttnLRP highlights the correct token most strongly, keeps relevances between [−4, 4], and captures model reasoning most accurately in the faithfulness experiments.Alternative bias-handling rules produce relevance ranges between [−10^15, 10^15], making their heatmaps resemble random noise.
Matrix multiplication: Applying the ε-rule to bi-linear matrix multiplication violates conservation, whereas AttnLRP’s sequential ε-rule and uniform-rule treatment avoids that violation.The standard ε-rule and related variants therefore differ from AttnLRP in how relevance is distributed through multiplication.
Matrix multiplication: AttnLRP combines the ε-rule for bi-linear matrix multiplication with a novel uniform rule that ensures conservation and high faithfulness.The uniform rule is derived from the DTD framework and addresses the conservation problem of applying the ε-rule alone.
Attention rollout: Attention-rollout methods are computationally efficient but provide low-resolution input maps and cannot attribute hidden latent features beyond the softmax output.These limitations restrict explanations to only part of the model and can produce checkerboard artifacts.
Gradient handling: Vision transformers can suffer gradient shattering, so AttnLRP applies the γ-rule selectively in convolutional and linear layers outside attention.The paper reports sparse, low-noise LLM attributions but susceptibility to gradient shattering in ViTs.

A.2.4. IMPACT OF TEMPERATURE SCALING ON THE SOFTMAX RULE

This section analyzes softmax attribution under temperature scaling and vanishing gradients. It explains why saturated softmax outputs can interrupt relevance flow and motivates higher temperatures for classification explanations.

Temperature scaling: Higher temperature values flatten softmax distributions, whereas lower temperatures sharpen them and reduce output randomness.Temperature scaling divides the softmax input x by T before applying the softmax function.
Vanishing gradients: Low temperatures or highly confident classifications can make the softmax derivative vanish when s_j ≈ 1.Because DTD uses the Jacobian to calculate attributions, the bias term can then capture all relevance and stop relevance flow.
Attention versus classification: Within attention, multiplying the softmax output by the value path transmits relevance through the uniform rule, but standalone softmax attribution remains problematic.The attention pathway therefore circumvents the saturation limitation more effectively than an independently used classification softmax.
Scope limitation: The paper recommends increased temperature scaling when explaining classification outputs, but does not investigate classification-output attribution because experiments remove the softmax and attribute logits.The effects of this setting remain identified as future work.
Related conservation result: For bi-linear matrix multiplication, the standard ε-rule violates conservation, while the sequential ε-rule and uniform rule preserve it by dividing with factor 2.The ε term absorbs only a negligible amount of relevance to maintain numerical stability.

B.1. Models and Datasets

The evaluations span vision transformers, language models, and encoder-decoder or mixture-of-experts models, using perturbation-based faithfulness and plausibility analyses. They also specify baselines, prompts, attribution targets, and implementation settings for reproducible comparisons.

Datasets and sampling: The ViT perturbation experiments use 3200 fixed ImageNet samples, while Wikipedia and IMDB faithfulness experiments use 4000 fixed validation samples.Wikipedia next-word prediction uses a context size of 512, whereas SQuAD context sizes vary.
Evaluation setup: SQuAD v2 plausibility is assessed with top-1 accuracy of the most relevant token and IoU between positive attributions and the ground-truth answer mask.Only correctly answered instances are retained, and the evaluation uses a fixed prompt format.
Evaluation setup: Perturbation faithfulness is measured by iteratively replacing the most- or least-relevant input features with a baseline and tracking the model’s output confidence.For LLMs, entire token embeddings are flipped; for ViTs, input pixels are flipped, using zero as the baseline.
Model architecture: AttnLRP rules improve performance across all evaluated layer types, with larger gains as model nonlinearity increases.The ViT configuration uses different γ values across layer types, and the most faithful AttnLRP and CP-LRP composites are reported separately.

B.6. Additional Perturbation Evaluations on Vision Transformers

Additional ImageNet perturbation evaluations extend the vision-transformer comparison to ViT-L-16 and ViT-L-32. AttnLRP surpasses the comparative baselines, while γCP-LRP remains highly competitive.

Additional ViT evaluations: AttnLRP surpasses all comparative baselines on ViT-L-16 and ViT-L-32 ImageNet perturbation evaluations.The results are reported as faithfulness scores based on perturbation curves.
Additional ViT evaluations: γCP-LRP remains highly competitive, but greater architectural nonlinearity increases AttnLRP’s superiority over competing methods.The same tuned composite was obtained for both models, although the authors provide no assurance that other models share these γ parameters.

B.7. Attributions on SQuAD v2

SQuAD v2 heatmaps compare attribution methods on Mixtral 8x7b and Flan-T5-XL answer tokens. AttnLRP most consistently identifies the correct answer tokens, with its advantage larger on models containing additional nonlinear routing or weighting components.

Cross-model comparison: AttnLRP has a small but consistent top-1 advantage over CP-LRP on Flan-T5-XL, while substantially outperforming CP-LRP on Mixtral 8x7b and LLaMa 2.The authors associate this pattern with the additional nonlinear weighting or routing layers in Mixtral and LLaMa 2.
Baseline behavior: AtMan can be excellent in some cases but fails in others, while gradient-based methods, Attention Rollout, and Grad×Attention Rollout often produce noisy or uninformative maps.This variability is visible across the supplied Mixtral and Flan-T5 examples.
Mixtral 8x7b: On Mixtral 8x7b, AttnLRP is the exception that most strongly highlights the correct answer tokens in the “France” and “Ibn Sina” examples.Other methods show noise, background relevance, sparsity, or missed answer tokens.
Flan-T5-XL: In the Flan-T5 example, AttnLRP and CP-LRP identify the answer token and relate it to other numerical context through positive and negative relevance.The authors conjecture that this targeted contrast may reflect distinctions needed between related questions, but leave systematic analysis for future work.

B.8. Benchmarking Cost, Time and Memory Consumption

The benchmarking section compares attribution cost, runtime, and peak memory for LLaMa 2 and illustrates qualitative attributions on several language-model examples. LRP-based methods offer efficient backpropagation, but checkpointed LRP can exceed available memory for the largest model.

Knowledge neurons: Knowledge-neuron visualizations show apparent disentanglement, with individual neurons associated with concepts such as railways and teachers.AttnLRP highlights the tokens in reference samples that maximally activate these neurons.
Qualitative attributions: On Mixtral 8x7b, AttnLRP highlights “France” and the complete “Ibn Sina” answer, whereas competing methods show noise, background relevance, or incomplete token coverage.These qualitative examples correspond to SQuAD v2 attribution analyses.
Computational benchmarking: For LLaMa 2-70b, AttnLRP curves stop because memory consumption exceeds 160 GB, and checkpointed LRP requires more memory than the benchmarking node supplies.The benchmark uses four A100-SXM4 40GB GPUs and measures cost, runtime, and peak GPU memory across context sizes.