Attention visualization and limitations

Overview

Attention visualization is the process of inspecting attention weights from transformer models to understand which input tokens the model focuses on when making predictions. While intuitively appealing—"let's see what the model looks at"—attention visualization has significant limitations as an explanation tool that every practitioner must understand.

Think of it like watching someone read: where their eyes linger might tell you what they find important. But just because they looked at a word doesn't mean it caused their conclusion.

Core Concepts

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

The attention weights are the matrix $A = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)$ where $A_{ij}$ represents how much position $i$ attends to position $j$.

WHY softmax? It normalizes scores into a probability distribution: all weights sum to 1, making them interpretable as "importance scores."

WHAT does $A_{ij}$ mean? It's the fraction of information from token $j$ that flows into the representation of token $i$ at this layer.

HOW to extract? After computing attention in your model, simply save the softmax output before multiplying with $V$.

Deriving Attention from First Principles

Let's build up why attention takes this form:

1. Goal: Combine information from multiple input tokens, with learned importance weights
2. Similarity scoring: For query vector $q_i$ and key vector $k_j$, compute similarity via dot product: $q_i \cdot k_j$
   • Why dot product? It measures alignment: high when vectors point in the same direction
   • In high dimensions, raw dot products grow with $d_k$, so we scale: $\frac{q_i \cdot k_j}{\sqrt{d_k}}$
3. Normalization: Convert scores to probabilities using softmax: $\alpha_{ij} = \frac{\exp(q_i \cdot k_j / \sqrt{d_k})}{\sum_{j'} \exp(q_i \cdot k_{j'} / \sqrt{d_k})}$
   • Why softmax? Ensures weights are positive and sum to 1 (valid probability distribution)
4. Weighted combination: Output is $\sum_j \alpha_{ij} v_j$—each value vector weighted by its attention score

This gives us the formula above. No magic, just learned similarity → normalized weights → weighted average.

$\text{MultiHead}(Q,K,V) = \text{Concat}(\text{head}_1, .., \text{head}_H)W^O$

where $\text{head}_h = \text{Attention}(QW^Q_h, KW^K_h, VW^V_h)$

Each head has its own attention matrix $A^{(h)}$. Visualization typically shows:

• Per-head heatmaps: $A^{(h)}_{ij}$ for each head $h$
• Averaged attention: $\bar{A}_{ij} = \frac{1}{H}\sum_h A^{(h)}_{ij}$
• Max attention: $\max_h A^{(h)}_{ij}$ (which head attends most)

Visualization Methods

Input: "The cat sat on the mat" → Predict next word

Visualization: 7×7 heatmap where cell $(i,j)$ shows how much token $i$ attends to token $j$.

Interpretation attempt: If "mat" attends strongly to "on" and "the", we might think the model is using prepositional context.

Code structure:

import matplotlib.pyplot as plt
import numpy as np
 
# attention_weights shape: [seq_len, seq_len]
attention = model.get_attention_weights(input_ids, layer=6, head=3)
plt.imshow(attention, cmap='viridis')
plt.xlabel('Key position')
plt.ylabel('Query position')

Why this step? The heatmap shows the raw $A$ matrix, making patterns visually obvious.

Task: Sentiment classification with BERT

Approach: Visualize all 12 heads in layer 8 side-by-side

Observation:

• Head 2: Attends to punctuation (exclamation marks, question marks)
• Head 7: Attends to negation words ("not", "never")
• Head 10: Dispersed attention (uninformative)

Interpretation: Different heads specialize in different linguistic features.

Critical question: But do these patterns cause the prediction? (Spoiler: Not necessarily)

Why this step? Comparing heads reveals specialization, but we must test causality separately.

Problem: Single-layer attention doesn't show information flow through the entire model.

Solution: Attention rollout computes cumulative attention across layers:

$A_{\text{rollout}}^{(l)} = A^{(l)} \cdot A^{(l-1)} \cdot ... \cdot A^{(1)}$

Derivation: Information flows from layer to layer. If layer1 moves information from token $j$ to $k$, and layer 2 moves from $k$ to $i$, then effectively $j$ influences $i$ through $k$.

Matrix multiplication $A^{(2)} \cdot A^{(1)}$ captures this transitive flow: entry $(i,j)$ in the product sums over all intermediate tokens $k$.

Why this step? Captures deep information flow, not just local attention.

Limitation: Assumes attention is the only information pathway (it's not—residual connections and MLPs matter too).

Fundamental Limitations

The wrong intuition: If a token gets high attention weight, it must be important for the prediction.

Why it feels right: Attention literally means "paying attention to"—more attention should mean more influence, right?

The reality: Attention weights don't directly reveal causal importance. Research by Jain & Wallace (2019) and Serano & Smith (2019) showed:

1. Adversarial examples: You can find completely different attention distributions that produce the same output
2. Attention is not explanation: High attention on token doesn't mean ablating it would change the prediction
3. Value vectors matter: Even with high attention weight $\alpha_{ij}$, if value vector $v_j$ is in the null space of downstream layers, it contributes nothing

Mathematical fix: $\text{Contribution}_{ij} = \alpha_{ij} \cdot ||v_j|| \cdot \text{downstream}_{\text{sensitivity}}$

True importance requires measuring gradient or ablation impact, not just attention weights.

Why this happens: Attention is only one operation in a deep network. The value vectors pass through LayerNorm, residual additions, and MLPs that can amplify or nullify their contribution.

The wrong approach: Average all attention heads and interpret the result.

Why it feels right: Averaging smooths out noise and gives a "consensus" view.

The reality: Heads can be:

• Redundant: Multiple heads learning the same pattern (wasteful but harmless)
• Antagonistic: Some heads attend to confounders, others to true signal—averaging mixes them
• Specialized: One head carries all information for a task; averaging dilutes it

The fix: Analyze heads individually with gradient-based attribution to identify which heads are actually used by the classification layer.

Attention flow: Forward-pass information routing through attention weights

Gradient flow: Backward-pass credit assignment showing which inputs affect the loss

Attention shows correlation, gradients show causation.

For token $x_j$ and output $y$:

• Attention: "How much did the model look at $x_j$?"
• Gradient $\frac{\partial \mathcal{L}}{\partial x_j}$: "How much does changing $x_j$ change the loss?"

These can disagree! A token might receive high attention but have near-zero gradient (it was examined but not used), or low attention but high gradient (indirect influence through other tokens).

Practical Visualization Guidelines

For $L$ layers with attention matrices $A^{(1)}, .., A^{(L)}$:

$\tilde{A}^{(l)} = A^{(l)} + I$ $A_{\text{rollout}} = \tilde{A}^{(L)} \cdot \tilde{A}^{(L-1)} \cdot ... \cdot \tilde{A}^{(1)}$

Why add identity $I$? Residual connections mean each layer both reads attention and passes through raw input. Adding $I$ models the residual stream.

Normalization: Row-normalize after each product to maintain probability interpretation: $\tilde{A}_{\text{norm}} = \text{diag}((\tilde{A} \cdot \mathbf{1})^{-1}) \cdot \tilde{A}$

Why this step? Without normalization, values explode geometrically through multiplication.

When to use: When you care about end-to-end attribution from input to output, not just local patterns.

To incorporate gradient information:

$\text{Attribution}_{ij} = A_{ij} \cdot \frac{\partial y}{\partial A_{ij}}$

Derivation: By chain rule, the influence of attention weight $A_{ij}$ on output $y$ is: $\frac{\partial y}{\partial x_j} = \sum_i \frac{\partial y}{\partial A_{ij}} \cdot \frac{\partial A_{ij}}{\partial x_j}$

The term $\frac{\partial y}{\partial A_{ij}}$ is the gradient of the output w.r.t. the attention weight—it measures sensitivity of the prediction to that attention connection.

Multiplying $A_{ij}$ (forward importance) by this gradient (backward sensitivity) gives a better attribution score than attention alone.

Why this step? Combines "what the model looked at" (attention) with "what mattered for the output" (gradient).

When to Trust Attention Visualization

Trust it for: ✓ Exploratory analysis: "What patterns does the model learn?" ✓ Debugging: "Is the model attending to padding tokens?" (bug indicator) ✓ Linguistic analysis: "Do heads capture syntactic structures?"

Don't trust it for: ✗ Causal explanations: "This token caused the prediction" ✗ Model debugging: "The model failed because it didn't attend to X" ✗ Fairness audits: "The model discriminates because of attention pattern Y"

Gold standard: Combine attention visualization with:

• Ablation studies: Remove tokens and measure output change
• Gradient-based attribution: Integrated Gradients, Attention×Gradient
• Probing: Train classifiers on representations to test what information is encoded

Now here's the catch: just because you looked at something doesn't mean it helped you answer! Maybe you read the wrong section, or you already knew the answer from memory. Watching where you looked tells us something about your strategy, but it doesn't prove why you got the question right or wrong.

That's the limitation: attention maps show what the AI "looked at," but not what actually caused its decision. To find that out, we need to do experiments—like covering up parts of the notes and seeing if the answer changes (that's ablation testing).

Connections

• Transformer Architecture - Where attention comes from
• Integrated Gradients -based attribution alternative
• Attention Mechanisms - Mathematical foundation
• SHAP for Deep Learning - Model-agnostic explanation method
• Layer-wise Relevance Propagation - Another gradient-based approach
• Probing Classifiers - Testing what representations encode
• BertViz Tool - Practical visualization library
• Adversarial Attention - Attention distribution robustness

#flashcards/ai-ml

What are attention weights in a transformer? :: The matrix $A = \text{softmax}(QK^T/\sqrt{d_k})$ where $A_{ij}$ represents how much position $i$ attends to position $j$—the fraction of information from token $j$ flowing into token $i$'s representation.

Concept Map

Hinglish (regional understanding)