Visualizing LLM Attention Mechanisms for Debugging: A Complete Toolkit

The bojieli/ai-agent-book repository provides a self-contained attention visualization toolkit in chapter2/attention_visualization/visualization.py that converts raw Transformer attention weights into actionable heatmaps, flow diagrams, and statistical summaries for debugging LLM behavior.

Visualizing LLM attention mechanisms for debugging reveals exactly where a model focuses during generation, exposing patterns like attention drift or sink saturation that raw logits cannot show. The ai-agent-book project ships with a dedicated Python module that transforms JSON-serialized attention weights from inference runs into publication-ready figures, requiring only matplotlib, seaborn, and numpy.

Core Visualization Components

The toolkit implements six distinct views, each targeting a specific debugging scenario. All functions reside in chapter2/attention_visualization/visualization.py and operate on a standardized JSON schema containing input_tokens, output_tokens, attention_steps, and context_length.

Attention Heatmaps

The create_attention_heatmap function renders a 2-D matrix of shape [output_len × total_len] that maps each generated token back to its attended context【†L41-L65】. This view answers the fundamental question: Which input tokens influenced this specific output?

Key parameters include context_boundary to visually separate the prompt from generated text, and save_path to export PNG files directly.

Flow Diagrams

To observe how attention evolves across generation steps, create_attention_flow_diagram produces per-step bar charts【†L27-L48】. This visualization exposes attention drift—phenomena where the model progressively ignores earlier context—or sudden shifts when encountering tool-use triggers.

Token Attention Summaries

The create_token_attention_summary function aggregates statistics across all generation steps, outputting average attention weights, maximum values, and distribution histograms【†L19-L45】. Use this for high-level validation: if the summary shows uniform attention where you expect spikes, your model may have attention-head saturation issues.

Layer-Level Heatmaps

For deep architectural debugging, create_layer_attention_heatmap visualizes the full [seq × seq] self-attention matrix for a specific layer and head【†L72-L84】. The function supports optional "sink" annotations, highlighting the attention-sink phenomenon discussed in Chapter 2 of the book where token 0 disproportionately accumulates attention mass.

Side-by-Side Comparisons

When A/B testing configurations—such as tool-use versus no-tool prompts—create_attention_comparison renders multiple attention matrices in a single figure【†L56-L70】. Pass a list of matrices and corresponding titles to identify behavioral divergences at a glance.

Sink Statistics

The attention_sink_stats helper quantifies the attention-sink effect, computing mean and maximum attention shares for a designated sink token (typically index 0)【†L39-L49】. This provides numerical grounding for the visual patterns observed in layer-level heatmaps.

Running the Visualization Pipeline

The module exposes visualize_results as its primary entry point【†L48-L66】. This function accepts a JSON file path produced by LLM inference scripts and automatically generates heatmaps, flow diagrams, and summary figures for every entry in the dataset.

Execute the full pipeline via the command line:

python -m chapter2.attention_visualization.visualization \
    attention_results.json

The script’s __main__ block handles file I/O and delegates to the individual visualization functions, writing figures to disk as specified in the JSON metadata【†L61-L73】.

Practical Code Examples

Generate a Heatmap for a Single Generation

Load inference results and render a 2-D attention matrix, averaging multi-head weights where necessary:

import json
import numpy as np
from chapter2.attention_visualization.visualization import (
    create_attention_heatmap,
)

with open("sample_result.json") as f:
    res = json.load(f)[0]

attention_matrix = []
for step in res["attention_steps"]:
    w = step["attention_weights"]
    if isinstance(w[0], list):  # Multi-head case

        w = np.mean(w, axis=0).tolist()
    attention_matrix.append(w)

fig = create_attention_heatmap(
    attention_weights=attention_matrix,
    input_tokens=res["input_tokens"],
    output_tokens=res["output_tokens"],
    context_boundary=res["context_length"],
    title="LLM Attention Heatmap – Example 1",
    save_path="heatmap_example1.png",
)
plt.close(fig)

Plot Attention Flow Across Steps

Visualize the evolution of attention weights during the first eight generation steps:

from chapter2.attention_visualization.visualization import (
    create_attention_flow_diagram,
)

fig = create_attention_flow_diagram(
    attention_steps=res["attention_steps"],
    input_tokens=res["input_tokens"],
    context_length=res["context_length"],
    max_steps=8,
    save_path="flow_example1.png",
)
plt.close(fig)

Compare Two Layers Side-by-Side

Load NumPy arrays containing per-layer attention weights and contrast their patterns:

from chapter2.attention_visualization.visualization import (
    create_attention_comparison,
)

layer1 = np.load("layer1.npy")   # Shape: [seq, seq]

layer2 = np.load("layer2.npy")

fig = create_attention_comparison(
    matrices=[layer1, layer2],
    tokens_list=[res["input_tokens"] + res["output_tokens"]] * 2,
    titles=["Layer 1", "Layer 2"],
    suptitle="Layer-wise Attention Comparison",
    save_path="layer_comparison.png",
)
plt.close(fig)

Batch Process Results Files

For automated debugging pipelines, invoke the module directly against JSON output from your inference engine:

python -m chapter2.attention_visualization.visualization \
    attention_results.json

This executes visualize_results, generating all three standard visualizations (heatmap, flow, summary) for each record in the batch.

Summary

• The chapter2/attention_visualization/visualization.py module in bojieli/ai-agent-book provides six specialized functions for visualizing LLM attention mechanisms for debugging: heatmaps, flow diagrams, summaries, layer views, comparisons, and sink statistics.
• All functions consume a standard JSON schema produced by LLM inference runs, requiring no proprietary dependencies beyond matplotlib, seaborn, and numpy.
• The visualize_results entry point automates batch processing, while individual functions like create_attention_heatmap offer granular control for interactive debugging.
• The toolkit specifically addresses the "attention sink" phenomenon through both visual annotations in create_layer_attention_heatmap and quantitative metrics in attention_sink_stats.

Frequently Asked Questions

What input format does the toolkit expect?

The visualization functions expect JSON objects containing four required keys: input_tokens (list of strings), output_tokens (list of strings), attention_steps (list of dictionaries with attention_weights), and context_length (integer marking the prompt boundary). The attention_weights field may be a 1-D array for single-head attention or a 2-D array for multi-head attention, which you should average before passing to create_attention_heatmap.

How do I visualize attention for multi-head attention?

When attention_weights contains per-head arrays (shape [num_heads, seq_len]), compute the mean across the head dimension using numpy.mean(w, axis=0) before constructing the attention matrix. The code example in the heatmap section demonstrates this preprocessing step explicitly.

What is the "attention sink" phenomenon?

Attention sink refers to the tendency of certain tokens—usually the first token in the sequence (index 0)—to accumulate disproportionately high attention weights across many layers and heads. The toolkit exposes this through attention_sink_stats, which calculates the mean and maximum attention share for the sink token, and through create_layer_attention_heatmap, which can annotate sink regions visually. This phenomenon is documented in book/chapter2.md of the repository.

Can I compare attention patterns across different model configurations?

Yes. Use create_attention_comparison to render side-by-side heatmaps of attention matrices from different layers, prompts, or experimental conditions (e.g., with-tool versus without-tool). Pass the matrices as a list along with corresponding tokens_list and titles parameters to generate comparative figures that highlight behavioral differences.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}