Core ML

Attention is Not Explanation

Just because a Transformer's attention mechanism highlights a word does not mean that word caused the final prediction. Attention measures data routing, not mathematical contribution.

Attention weights measure how data is routed between tokens, but deep non-linear layers can completely erase that data before the final prediction.

Why Does This Exist?

When the Transformer architecture (the "T" in ChatGPT) took over the world in 2017, the NLP community thought they had solved the explainability problem for free.

A Transformer works by calculating attention weights. When it reads a sentence, it assigns a percentage (between 0 and 1) representing how much one word "attends" to another word. Researchers immediately began visualizing these attention weights as heatmaps. If the model predicted that a movie review was "Negative," and the attention heatmap brightly highlighted the word "Bad," everyone assumed the problem was solved: The model predicted Negative because it paid 90% attention to the word Bad.

In 2019, a famous paper titled "Attention is Not Explanation" (by Jain and Wallace) dropped a bombshell on the community. They proved mathematically and empirically that looking at attention weights to explain a prediction is deeply flawed and often completely wrong.

Think of It Like This

Think of a Transformer like a massive corporate bureaucracy.

The Attention Mechanism is the mailroom. If the mailroom routes 90% of the daily mail (data) to Bob's desk, it looks like Bob is the most important person in the company.

However, Bob is just an intern. After he receives the mail, he passes it through the Deep Layers (the middle managers). The middle managers completely ignore Bob's work, shred his reports, and instead base their final decision on a tiny sticky note they got from Alice (who only received 1% of the mail).

If you only look at the mailroom (Attention), you assume Bob drove the decision. But if you track the actual mathematical impact on the final output, Alice drove the decision.

How It Actually Works

To understand the fallacy, you have to look at what happens after the attention mechanism in a Transformer block.

1. The Attention Layer (Routing)

Attention simply moves information from one token to another. If the word "Bad" has an attention weight of 0.9, it successfully transfers a massive amount of vector data into the next layer.

2. The Feed-Forward Network (Squashing)

Right after the attention layer, the vector passes through a dense Feed-Forward Network (FFN) with non-linear activation functions (like ReLU or GELU).

This is where the deception happens. A ReLU function famously turns any negative number into a flat zero. The FFN can take that massive, highly-attended vector from the word "Bad", multiply it by negative weights, and ReLU it into oblivion. The mathematical impact of the word "Bad" is now exactly zero.

3. Layer Normalization (Scaling)

Furthermore, Transformers use Layer Normalization. If the attention mechanism makes the "Bad" vector huge, LayerNorm might mathematically divide it, scaling it down so it doesn't overpower the other tokens.

4. The Result

By the time the data reaches the final classification head, the word with the highest attention weight might have literally zero impact on the final prediction. Conversely, a word with an attention weight of 0.05 might bypass the ReLUs perfectly and completely drive the final output.

The Counter-Argument

It is worth noting that a few months after "Attention is Not Explanation" was published, another paper titled "Attention is Not Not Explanation" was published (by Wiegreffe and Pinter).

This paper argued that while attention weights do not equal strict mathematical feature importance (like shap values do), they are still a plausible, human-interpretable lens into the model's internal routing. They argued that discarding attention entirely is an overreaction.

Today, the consensus is: Attention is useful for debugging routing, but it must never be presented to a user as a mathematically rigorous explanation for a prediction.

Show Me the Code

If you want a true explanation for a Transformer prediction, you should not extract the attention weights. You should use a gradient-based method like Integrated Gradients (which tracks the math all the way to the output) or an occlusion-based method like LIME.

Here is an example of what not to do (extracting raw attention) vs what to do (using Captum for Integrated Gradients).

1# ❌ THE FLAWED APPROACH (Extracting Attention)2# This only tells you how the mailroom routed the data.3outputs = transformer_model(input_ids, output_attentions=True)4attention_weights = outputs.attentions[-1] # Grabbing the last layer's attention5plot_heatmap(attention_weights) # DO NOT use this to explain the final prediction!6
7# ✅ THE RIGOROUS APPROACH (Integrated Gradients)8# This tracks the actual mathematical impact on the final output.9from captum.attr import IntegratedGradients10
11ig = IntegratedGradients(transformer_model)12attributions = ig.attribute(inputs=input_embeddings, target=1)13plot_heatmap(attributions) # This is a mathematically sound explanation.

Watch Out For

The Multi-Head Trap

Even if attention was explanation, modern Transformers use Multi-Head Attention (e.g., 96 different attention heads running in parallel in GPT-3). Which head do you visualize? If Head 12 highlights the word "Bad," but the other 95 heads ignore it, is "Bad" actually important? Averaging the heads destroys the nuance, and picking one head is cherry-picking. This further invalidates attention as an explainability tool.

The Quick Version

It is incredibly tempting to use a Transformer's attention weights to explain its predictions, but it is mathematically flawed.
Attention measures how information is routed between words, not how that information affects the final output.
Deep non-linear layers (like FFNs and LayerNorm) exist after the attention mechanism. They can easily squash a highly-attended word to zero, or amplify a mostly-ignored word to drive the prediction.
To rigorously explain a Transformer, use gradient-based techniques (like Integrated Gradients) or game-theory techniques (like SHAP), which track the true mathematical impact on the final output.