# Visualizing LLM Attention Mechanisms for Debugging: A Complete Toolkit

> Debug LLM behavior with a complete toolkit for visualizing attention mechanisms. Convert raw Transformer weights into actionable heatmaps, flow diagrams, and statistical summaries.

- Repository: [Bojie Li/ai-agent-book](https://github.com/bojieli/ai-agent-book)
- Tags: tutorial
- Published: 2026-08-01

---

**The bojieli/ai-agent-book repository provides a self-contained attention visualization toolkit in [`chapter2/attention_visualization/visualization.py`](https://github.com/bojieli/ai-agent-book/blob/main/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`](https://github.com/bojieli/ai-agent-book/blob/main/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:

```bash
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:

```python
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:

```python
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:

```python
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:

```bash
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`](https://github.com/bojieli/ai-agent-book/blob/main/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`](https://github.com/bojieli/ai-agent-book/blob/main/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.