How to Visualize Attention Weights for Debugging

How to Visualize Attention Weights for Debugging

Transformer-based models have become the backbone of modern NLP, but their inner workings can feel like a black box. Attention weights — the scores that determine how much each token "looks at" every other token — offer a window into the model's reasoning. Visualizing these weights is one of the most effective debugging techniques available to ML engineers working with attention-based architectures.

What Are Attention Weights?

In a transformer, each attention layer computes a distribution over input tokens for every query position. These distributions, produced by applying softmax to scaled dot-product scores, are the attention weights. They sum to 1 across the key dimension and indicate the relative importance of each token when producing a given output representation.

A typical multi-head attention layer with H heads, sequence length L, produces a tensor of shape (H, L, L). Stacked across layers, you get a full picture of how information flows through the network. Crucially, these weights are different from the final hidden states — they are interpretable probabilities, which makes them ideal for inspection.

Why Visualizing Attention Matters

Attention visualization is not just for research papers. It serves several practical debugging purposes:

• Detecting degenerate attention: Some heads collapse to attending only to the first token, the EOS token, or a single position. This wastes representational capacity.
• Validating prompt engineering: When a model ignores part of your prompt, attention maps can confirm whether the relevant tokens are being attended to at the right layers.
• Diagnosing long-context failures: In long sequences, attention may dilute across many tokens. Visualizations reveal whether the model focuses on the right context window.
• Comparing fine-tuned vs. base models: After fine-tuning, attention patterns often shift. Visual diffs help you understand what changed.
• Debugging custom architectures: If you implement a new attention variant, visualizations quickly reveal whether it behaves sensibly.

Extracting Attention Weights

The first step is obtaining the weights from your model. Most transformer libraries expose them when requested. Here is how to do it with Hugging Face Transformers:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    output_attentions=True,   # request attention weights
    return_dict=True,
)

text = "The cat sat on the mat because it was tired."
inputs = tokenizer(text, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)

# outputs.attentions is a tuple of length num_layers
# each element has shape (batch, num_heads, seq_len, seq_len)
attentions = outputs.attentions
print(f"Number of layers: {len(attentions)}")
print(f"Shape of layer 0: {attentions[0].shape}")

For PyTorch-native implementations, you can register forward hooks on the attention modules to capture the softmax outputs directly:

captured = {}

def make_hook(name):
    def hook(module, inp, out):
        # assumes out is (attn_output, attn_weights) or similar
        if isinstance(out, tuple) and len(out) > 1:
            captured[name] = out[1].detach().cpu()
    return hook

for block_idx, block in enumerate(model.transformer.h):
    block.attn.register_forward_hook(make_hook(f"layer_{block_idx}"))

Building a Basic Heatmap

The simplest visualization is a per-head heatmap using matplotlib. The following function plots attention from a single layer and head:

import matplotlib.pyplot as plt
import numpy as np

def plot_attention_head(attn, tokens, layer, head):
    # attn: tensor of shape (seq_len, seq_len)
    matrix = attn.numpy()
    fig, ax = plt.subplots(figsize=(8, 8))
    im = ax.imshow(matrix, cmap="viridis")

    ax.set_xticks(range(len(tokens)))
    ax.set_yticks(range(len(tokens)))
    ax.set_xticklabels(tokens, rotation=90)
    ax.set_yticklabels(tokens)
    ax.set_xlabel("Key (attended to)")
    ax.set_ylabel("Query (attending from)")
    ax.set_title(f"Layer {layer}, Head {head}")
    fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    plt.tight_layout()
    plt.savefig(f"attention_layer{layer}_head{head}.png", dpi=150)
    plt.show()

tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
plot_attention_head(attentions[6][0, 4], tokens, layer=6, head=4)

For a multi-head overview, you can tile all heads of a single layer into a grid:

def plot_layer_heads(layer_attn, tokens, layer_idx):
    # layer_attn: shape (num_heads, seq_len, seq_len)
    num_heads = layer_attn.shape[0]
    cols = 4
    rows = (num_heads + cols - 1) // cols
    fig, axes = plt.subplots(rows, cols, figsize=(cols * 3, rows * 3))

    for h in range(num_heads):
        ax = axes[h // cols, h % cols]
        ax.imshow(layer_attn[h].numpy(), cmap="viridis")
        ax.set_title(f"Head {h}")
        ax.set_xticks([])
        ax.set_yticks([])

    for h in range(num_heads, rows * cols):
        axes[h // cols, h % cols].axis("off")

    fig.suptitle(f"Layer {layer_idx} — all heads")
    plt.tight_layout()
    plt.savefig(f"layer_{layer_idx}_all_heads.png", dpi=150)
    plt.show()

plot_layer_heads(attentions[6][0], tokens, layer_idx=6)

Interactive Visualization with BertVista-style HTML

Static plots are useful, but interactive HTML views let you hover over tokens and inspect attention across layers and heads. A lightweight approach is to render an attention matrix as an HTML table with cell opacity proportional to the weight:

def attention_to_html(attn, tokens, layer, head):
    # attn: (seq_len, seq_len) numpy array
    rows = []
    rows.append("<table style='border-collapse:collapse;'>")
    rows.append("<tr><th></th>" +
                "".join(f"<th style='font-size:10px;'>{t}</th>" for t in tokens) +
                "</tr>")
    for i, tok in enumerate(tokens):
        cells = []
        for j in range(len(tokens)):
            w = float(attn[i, j])
            alpha = int(w * 255)
            cells.append(
                f"<td style='background:rgba(100,149,237,{w:.3f});"
                f"width:30px;height:30px;text-align:center;font-size:9px;'>"
                f"{w:.2f}</td>"
            )
        rows.append(f"<tr><th style='font-size:10px;'>{tok}</th>" +
                    "".join(cells) + "</tr>")
    rows.append("</table>")
    html = "\n".join(rows)
    with open(f"attention_L{layer}_H{head}.html", "w") as f:
        f.write(html)

attention_to_html(attentions[6][0, 4].numpy(), tokens, layer=6, head=4)

For production-grade interactive views, consider established tools:

• BertViz: Provides head view, model view, and neuron view for transformers.
• Tensor2Tensor's attention visualization: A classic interactive tool.
• Captum: PyTorch's interpretability library, useful for layer-wise attribution.
• TransformerLens: A mechanistic interpretability library with built-in attention plotting.

Here is a quick BertViz example:

from bertviz import head_view

# outputs.attentions must be enabled, and you need both
# input_ids and the attention tuple.
head_view(
    attentions=outputs.attentions,
    tokens=tokens,
    layer=6,
    heads=[4, 7],
)

Aggregating Across Heads and Layers

Looking at every head individually is overwhelming. Aggregation helps you spot global patterns. A common technique is to average across heads within a layer:

def aggregate_layer(layer_attn):
    # (num_heads, L, L) -> (L, L)
    return layer_attn.mean(dim=0)

def aggregate_model(attentions):
    # list of (1, H, L, L) -> (num_layers, L, L)
    return torch.stack([a[0].mean(dim=0) for a in attentions])

model_avg = aggregate_model(attentions)
print(model_avg.shape)  # (num_layers, L, L)

You can also compute per-head entropy to detect degenerate heads. A head that always attends to a single token has near-zero entropy:

def head_entropy(attn):
    # attn: (L, L), rows are distributions over keys
    eps = 1e-12
    return -(attn * (attn + eps).log()).sum(dim=-1).mean().item()

for layer_idx, layer_attn in enumerate(attentions):
    for head in range(layer_attn.shape[1]):
        ent = head_entropy(layer_attn[0, head])
        print(f"Layer {layer_idx} Head {head}: entropy={ent:.3f}")

Low entropy across many positions suggests the head is essentially a hard lookup, which may or may not be desirable depending on your task.

Best Practices for Attention Debugging

• Do not over-interpret a single example. Attention patterns vary widely across inputs. Always inspect multiple samples before drawing conclusions.
• Remember attention is not explanation. Attention weights reflect where the model looks, but they are not causal attributions. Use them as a hint, not proof.
• Check the right layers. Early layers often capture syntactic patterns, while later layers reflect task-specific behavior. Focus your debugging on the layers most relevant to the failure mode.
• Account for masking. If you use padding or causal masks, ensure your visualization reflects the effective attention distribution, not the raw softmax over padded positions.
• Normalize for comparison. When comparing models, normalize attention maps or use rank-based metrics, since raw weight scales can differ.
• Log attention during training. Periodically saving attention snapshots helps track how patterns evolve, which is invaluable for diagnosing training instabilities.
• Combine with other tools. Pair attention visualization with gradient-based attribution or activation patching for a fuller picture.

Common Pitfalls

One frequent mistake is treating averaged attention as the model's "true" focus. Averaging across heads can wash out meaningful specialization — one head might attend to the previous token, another to the subject noun, and averaging produces a meaningless blur. Always inspect individual heads before aggregating.

Another pitfall is ignoring the effect of positional encodings. In models with relative positional biases, attention patterns can look uniform even when the model is using position information effectively. Cross-reference attention maps with the positional encoding scheme to avoid misinterpretation.

Finally, be cautious with decoder-only models and causal masking. The lower-triangular structure means each query only attends to earlier keys, so the upper half of the matrix is always zero. Visualizations that do not reflect this can be misleading.

Conclusion

Visualizing attention weights is a powerful, accessible technique for debugging transformer models. By extracting the weights, rendering them as heatmaps or interactive views, and applying aggregation and entropy metrics, you can uncover why a model behaves the way it does — whether it is ignoring part of a prompt, collapsing to a degenerate head, or failing on long contexts. While attention is not a complete explanation of model behavior, it is an indispensable first diagnostic. Combined with disciplined inspection across multiple examples and layers, it turns the opaque internals of transformers into something you can actually reason about and improve.