💬 NLP

Learning to Explain: Supervised Token Attribution from Transformer Attention Patterns

This paper introduces ExpNet, a lightweight neural network that automatically learns to map transformer attention patterns to token-level importance scores, offering a more adaptive and computationally efficient alternative to existing rule-based attention aggregation and model-agnostic explanation methods.

Original authors: George Mihaila

Published 2026-01-22

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: George Mihaila

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a super-smart robot (a "Transformer" model) that reads text and makes decisions, like telling if a movie review is positive or if a sentence is grammatically correct. The problem is, this robot is a black box. It gives you an answer, but you have no idea why it chose that answer. In high-stakes situations (like healthcare or law), we can't just trust a black box; we need to know its reasoning.

This paper introduces a new tool called ExpNet (Explanation Network) to solve this mystery. Here is how it works, explained simply:

The Problem: The Robot's "Internal Monologue"

Inside the robot, there is a mechanism called Self-Attention. You can think of this as the robot's internal "gaze." When the robot reads a sentence, it looks at different words and decides how much to focus on each one.

The Old Way: Previously, scientists tried to guess which words were important by using fixed rules (like "always look at the words the robot stared at the longest"). It's like trying to guess a person's favorite food just by watching them stare at a menu for a long time. Sometimes it works, but often the rules are too rigid and miss the nuance.
The Other Way: Other methods treat the robot like a mystery box, poking it with different inputs to see how it reacts. This is slow, expensive, and often creates nonsense sentences that the robot doesn't understand.

The Solution: ExpNet (The "Translator")

The authors built ExpNet, which is like a translator or a coach that learns to read the robot's "gaze" (attention patterns) and translate it into a clear explanation.

Instead of using fixed rules, ExpNet learns from humans.

Training: The researchers showed ExpNet examples where humans had already highlighted the important words (the "rationales").
Learning: ExpNet looked at the robot's internal "gaze" during those examples and figured out: "Ah, when the robot looks at the word 'terrible' in this specific way, humans think that word is the key reason for the negative review."
The Result: ExpNet becomes a lightweight, fast neural network that can look at the robot's attention patterns and say, "These are the words that matter," with high accuracy.

How It Works (The Analogy)

Imagine the robot is a detective solving a case.

The Robot's Attention: The detective has a magnifying glass. He moves it over different clues (words).
The Old Methods: They just assumed that whatever clue the detective held the magnifying glass over the longest was the most important.
ExpNet: ExpNet is a trainee detective who watched the master detective solve many cases where a human expert pointed out the actual smoking gun. The trainee learned to recognize the pattern of the master's gaze. Now, when the master looks at a new case, the trainee can instantly point out the clues, even if the master hasn't been trained on that specific type of case before.

The Big Test: Can It Learn General Rules?

The authors tested ExpNet in a very tough scenario: Cross-Task Generalization.

They trained ExpNet on two different types of problems (e.g., detecting hate speech and analyzing movie reviews).
Then, they tested it on a completely different problem (e.g., checking grammar) that it had never seen before.

The Result: ExpNet didn't just memorize the answers; it learned the language of importance. It performed better than 13 other existing methods, even on the new, unseen task. It proved that the "gaze" of the robot contains universal clues about what humans consider important, and ExpNet is the best at decoding them.

Why Is This Better?

It's Fast: Unlike other methods that have to poke and prod the robot thousands of times, ExpNet just takes a quick look at the robot's internal data and gives an answer instantly.
It's Accurate: It matched human reasoning better than any other method tested.
It's Flexible: It works across different types of text (sentiment, grammar, hate speech) without needing to be retrained for every single new topic.

The Catch (Limitations)

It needs a teacher: To learn, ExpNet needs examples where humans have already highlighted the important words. If you have a brand new task with no human examples, you can't train it yet (though the paper shows it can still guess well if trained on other tasks).
It sees words, not phrases: It explains things word-by-word. It might miss the nuance of a phrase like "not bad" (which means good) because it looks at "not" and "bad" separately.
It reflects human bias: Since it learns from human annotations, if the humans who trained it were biased, ExpNet will learn those biases too. It reflects human reasoning, not necessarily absolute truth.

Summary

The paper argues that instead of guessing how AI models think using rigid rules or slow experiments, we should teach a small, smart helper (ExpNet) to read the model's internal attention patterns. By learning from human examples, this helper can quickly and accurately explain why the AI made a decision, making these powerful tools more trustworthy and transparent.

Technical Summary: Learning to Explain (ExpNet)

Problem Statement

The deployment of transformer-based models in high-stakes domains (healthcare, legal, finance) is hindered by their opacity, which undermines trust and accountability. While self-attention mechanisms in transformers offer a potential source of interpretability, existing methods for extracting explanations suffer from significant limitations:

Model-Agnostic Methods (e.g., LIME, SHAP): Treat the model as a black box, ignoring internal reasoning. They incur high computational costs and, in NLP, often generate semantically invalid inputs through perturbation (masking/replacing words), leading to out-of-distribution examples that confound causal interpretation.
Attention-Based Methods: While they utilize internal signals, they rely on manually defined aggregation strategies and fixed heuristic rules (e.g., averaging attention weights across heads or layers). These predetermined rules may not capture the optimal combinations of attention features required to align with human reasoning across different tasks.

The core problem addressed is the lack of a method that can automatically learn to map transformer attention patterns to human-judged token importance without relying on rigid, hand-crafted rules or expensive black-box perturbations.

Methodology

The paper introduces ExpNet (Explanation Network), a lightweight neural network designed to learn an explicit mapping from transformer attention patterns to token-level importance scores.

1. Base Model and Feature Extraction

Base Model: The approach uses a frozen BERT-base model as a feature extractor. The model is not modified; only its attention patterns are utilized.
Feature Design: ExpNet extracts features from the final self-attention layer (Layer 12) of BERT, which contains 12 attention heads. For each token $j$ $j$ , it captures two distinct, non-symmetric perspectives of attention relative to the [CLS] token:
- Task-to-Token Attention ( $A_{[CLS],j}$ ): Measures how strongly the aggregation mechanism ([CLS]) attends to token $j$ .
- Token-to-Task Attention ( $A_{j,[CLS]}$ ): Measures how much token $j$ attends to the task representation.
Input Vector: These features are concatenated to form a 24-dimensional vector ( $2 \times 12$ heads) for each token, preserving individual head contributions from both directions rather than averaging them.

2. Network Architecture

ExpNet is a simple feed-forward network consisting of:

Input Layer: 24 dimensions.
Hidden Layer: 16 units with ReLU activation.
Output Layer: A single sigmoid unit producing a probability $\hat{y}_j \in [0, 1]$ representing the likelihood that token $j$ is important.

3. Training Procedure

Supervision: The network is trained as a token-level binary classifier using human-annotated rationales. Labels are projected from word-level annotations to subword tokens (if any subtoken of a word is marked important, all subtokens receive a positive label).
Data Filtering: Training is performed exclusively on instances where the underlying BERT classifier made a correct prediction. This ensures the explainer learns attention patterns associated with successful reasoning, avoiding noise from error patterns.
Loss Function: To address the class imbalance (most tokens are unimportant), Focal Loss is used instead of cross-entropy. This down-weights easy examples and focuses learning on hard-to-classify tokens.
Inference: A threshold of 0.5 is applied to generate binary predictions. A fallback mechanism ensures at least one token is marked important per example to prevent empty explanations.

4. Experimental Design

The evaluation employs a cross-task generalization protocol. ExpNet is trained on the combined rationale data of two NLP tasks and evaluated on a held-out third task. The tasks used are:

HateXplain: Hate speech detection (Human rationales).
SST-2: Sentiment analysis (LLM + Human rationales).
CoLA: Grammatical acceptability (LLM + Human rationales).

This setup tests whether ExpNet learns generalizable explanation patterns rather than task-specific heuristics.

Key Contributions

Novel Framework: A supervised learning approach that discovers optimal attention feature combinations automatically, replacing fixed aggregation rules with a learnable mapping from attention patterns to human rationales.
Cross-Task Generalization: Demonstration that attention-based explanations can generalize across distinct NLP domains (sentiment, grammar, hate speech) when trained on a combination of tasks.
Architectural Validation: Comprehensive ablation studies justifying the use of bidirectional attention features (task-to-token and token-to-task) and the specific network design.
Empirical Performance: Extensive evaluation showing state-of-the-art performance against a broad spectrum of baselines, including model-agnostic methods (LIME, SHAP), gradient-based methods (Integrated Gradients, Grad-CAM), and attention-based techniques (Rollout, GAE, AttCAT).

Results

ExpNet was evaluated against 13 baseline methods across three datasets using Token-level F1-score and AUROC.

F1-Score Performance: ExpNet achieved the highest F1 scores across all three tasks in the cross-task setting:
- SST-2: 0.398 (vs. best baseline GAE/MGAE at 0.350).
- CoLA: 0.468 (vs. best baseline Grad-CAM at 0.356), a 31% relative improvement.
- HateXplain: 0.473 (vs. best baseline Grad-CAM at 0.396), a 19% relative improvement.
AUROC Performance: ExpNet achieved competitive AUROC values (consistently above 0.7), though it did not always top this metric. The authors note a distinction: ExpNet excels at selecting the correct set of important tokens (high binary classification accuracy), whereas some gradient-based methods may preserve relative ranking quality across all tokens more effectively.
Efficiency: ExpNet is computationally efficient, requiring only a single forward pass of a lightweight MLP over pre-computed attention features, avoiding the sampling and perturbation costs of model-agnostic methods.

Significance and Claims

The paper claims that ExpNet demonstrates that transformer attention can be systematically refined into human-aligned explanations through supervised learning. Its significance lies in:

Superior Alignment: It consistently outperforms established baselines in aligning with human rationales, suggesting that learned mappings are superior to hand-crafted heuristics for extracting explanations.
Generalizability: The ability to train on two tasks and successfully explain a third indicates that the model captures fundamental patterns of how human-important features manifest in transformer attention, independent of specific classification objectives.
Practicality: By maintaining computational efficiency while improving explanation quality, ExpNet offers a viable path for deploying interpretable models in real-time, high-stakes applications.

The authors acknowledge limitations, including the reliance on human-annotated rationales for training (though cross-dataset results suggest potential for zero-shot application to unannotated tasks), the focus on encoder-only models, and the token-level granularity which may miss phrase-level semantics. They position these not as fundamental flaws but as areas for future extension to decoder-only models and broader domain coverage.