Learning to Explain: Supervised Token Attribution from Transformer Attention Patterns

TL;DR

This work addresses the opacity of transformer models in high-stakes domains by reframing explanation as a learnable mapping from attention patterns to token-level rationales. ExpNet uses a frozen BERT-base as a feature extractor and trains a lightweight MLP to predict token importance from bidirectional attention signals, supervised by human rationales. It demonstrates strong cross-task generalization across SST-2, CoLA, and HateXplain, outperforming a broad suite of baselines while offering substantial computational efficiency. The results indicate that attention-based explanations can be learned to align with human reasoning and deployed in real-time, though future work is needed to scale to larger models and decoder-based architectures and to validate user impact. Overall, ExpNet provides a principled, efficient pathway to human-aligned explanations in transformer-based NLP systems.

Abstract

Explainable AI (XAI) has become critical as transformer-based models are deployed in high-stakes applications including healthcare, legal systems, and financial services, where opacity hinders trust and accountability. Transformers self-attention mechanisms have proven valuable for model interpretability, with attention weights successfully used to understand model focus and behavior (Xu et al., 2015); (Wiegreffe and Pinter, 2019). However, existing attention-based explanation methods rely on manually defined aggregation strategies and fixed attribution rules (Abnar and Zuidema, 2020a); (Chefer et al., 2021), while model-agnostic approaches (LIME, SHAP) treat the model as a black box and incur significant computational costs through input perturbation. We introduce Explanation Network (ExpNet), a lightweight neural network that learns an explicit mapping from transformer attention patterns to token-level importance scores. Unlike prior methods, ExpNet discovers optimal attention feature combinations automatically rather than relying on predetermined rules. We evaluate ExpNet in a challenging cross-task setting and benchmark it against a broad spectrum of model-agnostic methods and attention-based techniques spanning four methodological families.

Figures (6)

• Figure 1: ExpNet complete pipeline from BERT's attention patterns to importance predictions.
• Figure 2: AUROC values across datasets show ExpNet consistently achieves competitive ranking performance (often above 0.7), generally assigning higher scores to important tokens and lower scores to unimportant ones more effectively than most baselines.
• Figure 3: AUPR values across datasets
• Figure 4: Example from CoLA highlighting grammatical acceptability cues. The sentence contains an incomplete verb phrase that determines grammaticality. ExpNet closely matches the gold standard by precisely identifying "hoped" as the critical token. Many baselines either over-highlight the full sentence (LIME, SHAP, Integrated Gradient, LRP, Full LRP, GAE, CAM, GradCAM, AttCAT, MGAE) or fail to identify any specific tokens (RawAt, Rollout), failing to isolate the specific grammatical violation that human annotators identified.
• Figure 5: Word importance attribution comparisons on a toxic example from the HateXplain dataset. Gold-standard human annotations highlight slurs and hate-inducing terms that constitute the offensive content. ExpNet aligns well with human judgments, correctly identifying the key toxic tokens. Some gradient-based methods (MGAE, LRP, Integrated Gradient) also capture relevant terms, though with varying precision. Model-agnostic methods (SHAP, LIME, Full LRP) highlight irrelevant tokens or produce overly diffuse attributions across most of the sentence, failing to isolate the specific language making the content hateful.