Visualizing Attention Patterns in Large Language Model Execution

Most developers running large language models treat inference like a opaque black box. You feed tokens into model.generate(), wait a few hundred milliseconds, and get a string back. When the output breaks, you tweak the prompt, adjust temperature, or swap models. You end up guessing because you cannot see what the model actually paid attention to during calculation.

Transformer models do not operate on magic. They run matrix multiplications where tokens explicitly vote on which previous tokens matter. Every attention layer calculates dot products between query vectors and key vectors, normalizes them with a softmax function, and uses the resulting scores to mix value vectors.

When your model hallucinates details, ignores system instructions, or gets confused by long retrieval-augmented generation (RAG) contexts, the exact failure reason sits inside those attention matrices. Inspecting these internal attention matrices transforms black-box inference into observable data flow.

The Math Behind the Attention Tensor

To extract useful insight from attention weights, you first need to track tensor shapes as tokens pass through the network. Standard self-attention processes an input sequence of length N through three projection matrices: Query (W_q), Key (W_k), and Value (W_v).

For an input matrix X with sequence length N and hidden dimension d_model:

Q = X * W_q
K = X * W_k
V = X * W_v

The core attention operation computes dot-product scores between every query token and every key token. The equation scales the dot products by the square root of key dimension d_k before applying softmax:

Attention(Q, K, V) = Softmax((Q * K^T) / sqrt(d_k)) * V

The matrix multiplication Q * K^T produces a square grid of shape [N, N]. Row i of this grid represents query token i, and column j represents key token j. The softmax function normalizes each row so the values sum to 1.0.

Modern LLMs split these projections across multiple heads. A model like Llama 3 8B features 32 attention heads per layer across 32 transformer layers. Multi-Head Attention converts that 2D matrix into a 4D tensor with shape [batch_size, num_heads, seq_len, seq_len].

Because standard causal language models use causal masking, tokens can only attend to previous positions. The upper triangle of every [seq_len, seq_len] matrix gets masked out to negative infinity before softmax, forcing those upper-right values to zero.

The resulting tensor contains millions of floats for a single forward pass. Raw float dumps tell you nothing. You need interactive visual mapping to make sense of these weights.

Extracting Attention Weights in PyTorch

Extracting raw attention matrices does not require modifying core PyTorch code or rebuilding model binaries from source. HuggingFace Transformers supports an output_attentions=True flag in model execution, but that parameter allocates large buffers that quickly run out of GPU VRAM on sequences longer than 1,000 tokens.

A cleaner, more memory-efficient approach uses PyTorch forward hooks. Hooks intercept the output of specific sub-modules right after they run, allowing you to slice, convert, or stream tensors to CPU RAM immediately without retaining full backward computation graphs. While hooks extract internal states locally, security researchers have also demonstrated risks involving stealing LLM reasoning traces through API responses.

Here is a minimal script that registers forward hooks on attention blocks to capture matrices directly:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
 
model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)
 
captured_attentions = {}
 
def get_attention_hook(layer_idx):
    def hook(module, input, output):
        # HuggingFace attention modules output tuples: (attn_output, attn_weights, ...)
        # Slice to layer index and move to CPU immediately to clear GPU memory
        if isinstance(output, tuple) and len(output) > 1:
            weights = output[1].detach().cpu().to(torch.float16)
            captured_attentions[layer_idx] = weights
    return hook
 
# Attach hook to self-attention module in each layer
for i, layer in enumerate(model.model.layers):
    layer.self_attn.register_forward_hook(get_attention_hook(i))
 
prompt = "System: You are an expert engineer.\nUser: Explain memory alignment."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
with torch.no_grad():
    outputs = model(**inputs)
 
print(f"Captured {len(captured_attentions)} layers.")
print(f"Layer 0 tensor shape: {captured_attentions[0].shape}")

Notice the call to .cpu().to(torch.float16) inside the hook. Moving tensors off the GPU instantly prevents CUDA out-of-memory errors during long trace operations.

Building a Zero-Dependency Visualizer

Static matplotlib heatmaps break down quickly when analyzing transformer runs. A model with 32 layers and 32 heads generates 1,024 separate matrix plots per token. You cannot spot broad context routing trends across 1,000 individual PNG files.

You need an interactive view that renders the sequence tokens along both axes, lets you jump between layers, and filters specific attention heads on demand. Instead of introducing heavy node servers or complex web frameworks, you can generate a single self-contained HTML file containing inline CSS and JavaScript.

The python script below takes captured tensor data, formats it into compact JSON arrays, and injects it into a lightweight D3 visualizer template:

import json
 
def export_interactive_html(tokens, attentions, output_path="attention_map.html"):
    # Convert half-precision tensors into standard Python float lists
    layers_payload = {}
    for layer_idx, tensor in attentions.items():
        # Shape: [1, heads, seq_len, seq_len] -> take batch index 0
        matrix = tensor[0].tolist()
        layers_payload[layer_idx] = matrix
 
    payload_json = json.dumps({
        "tokens": tokens,
        "layers": layers_payload
    })
 
    html_content = f"""<!DOCTYPE html>
<html>
<head>
    <meta charset="utf-8">
    <title>Attention Matrix Inspector</title>
    <script src="https://d3js.org/d3.v7.min.js"></script>
    <style>
        body {{ font-family: monospace; background: #0d1117; color: #c9d1d9; margin: 20px; }}
        .controls {{ margin-bottom: 20px; display: flex; gap: 15px; align-items: center; }}
        select, input {{ background: #161b22; color: #c9d1d9; border: 1px solid #30363d; padding: 6px; }}
        #matrix {{ display: grid; gap: 1px; margin-top: 10px; }}
        .cell {{ width: 14px; height: 14px; background: #161b22; }}
        .label {{ font-size: 11px; font-family: monospace; white-space: pre; }}
        .flex-container {{ display: flex; gap: 20px; }}
        .tooltip {{ position: absolute; padding: 6px; background: #21262d; border: 1px solid #30363d; font-size: 12px; pointer-events: none; }}
    </style>
</head>
<body>
    <h2>Attention Matrix Inspector</h2>
    <div class="controls">
        <label>Layer: <select id="layerSelect"></select></label>
        <label>Head: <select id="headSelect"></select></label>
        <span id="statInfo"></span>
    </div>
    <div class="flex-container">
        <div id="chart"></div>
    </div>
 
    <script>
        const data = {payload_json};
        const tokens = data.tokens;
        const layers = data.layers;
 
        const layerSelect = d3.select("#layerSelect");
        const headSelect = d3.select("#headSelect");
 
        Object.keys(layers).forEach(l => layerSelect.append("option").text(l).value(l));
        const numHeads = layers[0].length;
        for (let h = 0; h < numHeads; h++) {{
            headSelect.append("option").text("Head " + h).value(h);
        }}
 
        function draw() {{
            const layer = layerSelect.property("value");
            const head = headSelect.property("value");
            const matrix = layers[layer][head];
 
            d3.select("#chart").html("");
 
            const margin = {{top: 100, right: 20, bottom: 20, left: 100}};
            const cellSize = 16;
            const width = tokens.length * cellSize;
            const height = tokens.length * cellSize;
 
            const svg = d3.select("#chart").append("svg")
                .attr("width", width + margin.left + margin.right)
                .attr("height", height + margin.top + margin.bottom)
                .append("g")
                .attr("transform", `translate({{margin.left}},{{margin.top}})`);
 
            const colorScale = d3.scaleSequential(d3.interpolateInferno)
                .domain([0, 1]);
 
            // Draw Y Axis (Queries)
            svg.selectAll(".rowLabel")
                .data(tokens)
                .enter().append("text")
                .text(d => d)
                .attr("x", -6)
                .attr("y", (d, i) => i * cellSize + cellSize / 1.5)
                .attr("text-anchor", "end")
                .attr("fill", "#8b949e")
                .attr("class", "label");
 
            // Draw X Axis (Keys)
            svg.selectAll(".colLabel")
                .data(tokens)
                .enter().append("text")
                .text(d => d)
                .attr("x", 0)
                .attr("y", 0)
                .attr("transform", (d, i) => `translate({{i * cellSize + cellSize / 2}}, -6) rotate(-45)`)
                .attr("text-anchor", "start")
                .attr("fill", "#8b949e")
                .attr("class", "label");
 
            // Draw Cells
            for (let r = 0; r < matrix.length; r++) {{
                for (let c = 0; c < matrix[r].length; c++) {{
                    const val = matrix[r][c];
                    svg.append("rect")
                        .attr("x", c * cellSize)
                        .attr("y", r * cellSize)
                        .attr("width", cellSize - 1)
                        .attr("height", cellSize - 1)
                        .attr("fill", val > 0.001 ? colorScale(val) : "#161b22")
                        .append("title")
                        .text(`Query: "{{tokens[r]}}" → Key: "{{tokens[c]}}"\nWeight: {{val.toFixed(4)}}`);
                }}
            }}
        }}
 
        layerSelect.on("change", draw);
        headSelect.on("change", draw);
        draw();
    </script>
</body>
</html>
"""
    with open(output_path, "w", encoding="utf-8") as f:
        f.write(html_content)
 
token_strings = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
export_interactive_html(token_strings, captured_attentions)

Running this script exports a standalone HTML document. You open it in any web browser without needing backend execution servers or third-party visual extensions.

Standard Attention Patterns in Trained Transformers

When you start browsing through network layers, raw matrix values turn out to be structured. Model heads specialize into repeatable operational patterns. Recognizing these structures helps you identify normal execution versus routing failures.

Positional Lookback and Previous-Token Heads

Early network layers (typically Layers 0 to 4) spend most of their computation on sequence position rather than semantic context.

Token Index:    0    1    2    3    4    5
Sequence:     [BOS]  The  quick brown fox jumps
Query (fox):    0    0    0.91  0.08 0    0

Previous-token heads show up as a solid line running parallel directly under the main diagonal. Query token at position i puts almost all of its attention weight on key token i - 1. These heads pass raw local sequence information forward to build n-gram representations.

Attention Sinks

Look closely at middle layers, and you will notice vertical columns near the far left edge of the matrix. Tokens across the entire sequence will place high attention weight (sometimes 40% to 90%) on position 0, which is usually the initial [BOS] or <|begin_of_text|> token.

This pattern represents an attention sink. Transformer heads often calculate that no specific prior token in the sequence is relevant for their task. Because softmax requires every matrix row to sum to 1.0, the head needs a neutral memory slot to dump excess weight.

Model creators intentionality configure initial sequence tokens to act as baseline sinks. If your prompt truncates the system's [BOS] token during dynamic context slicing, attention sinks break, sending values back into random sequence tokens and corrupting output quality.

Induction Heads

Induction heads handle in-context learning and pattern repetition across long sequences. They operate using a two-step lookup logic:

1. Search backward for prior occurrences of the current token.
2. If token A previously appeared followed by token B, attend heavily to token B.

Sequence:  "... user_id = 42 ... set user_id = 42"
                                        ^ Query: "user_id"
                                          attends directly to Key: " 42"

If an induction head sees the sequence [foo][bar] early in the prompt and later encounters [foo], it scans back, finds [foo], shifts one position right, and fires attention at [bar].

Induction heads typically concentrate in middle network layers (around Layers 10 to 20 on a 32-layer model). If your model fails to follow few-shot instruction formats, checking whether middle layers form strong induction head diagonals reveals if the context window is failing to map repetitive patterns. Similar context window issues explain why local LLM execution yields subpar reasoning output when parameters are improperly configured.

Diagnosing Real-World RAG and Prompt Failures

Inspecting matrix visualizations pinpoints why models fail under edge cases. Two common problems in retrieval-augmented generation are prompt dilution and system prompt overwrite.

Detecting Lost in the Middle Failures

When pasting 4,000 words of retrieved documentation into a prompt, models often honor facts near the top or bottom of the document while ignoring text buried in the middle.

Inspecting attention matrices across late layers makes the cause obvious:

[System Instruction]  ████████░░░░░░░░  (High weight)
[Doc Chunk 1 (Top)]   ██████░░░░░░░░░░  (Moderate weight)
[Doc Chunk 5 (Mid)]   ░░░░░░░░░░░░░░░░  (Near-zero weight)
[Doc Chunk 10 (Bot)]  ██████████░░░░░░  (High weight)

In late transformer layers (such as Layer 28 through 32), queries generated during output synthesis query keys from the top and bottom of the context window. Middle token keys show near-zero values across almost all heads.

If your visual inspector shows that a specific document chunk receives zero incoming attention lines across all heads in the last 4 layers, the model never reads those values during forward passes. No amount of temperature tuning will fix that retrieval failure. You must restructure prompt layout or trim redundant input.

Debugging System Prompt Overrides

User inputs can hijack attention weights if formatting boundaries leak. Consider a system prompt meant to enforce strict JSON output, followed by a user prompt containing arbitrary markdown text:

prompt = """<|system|>
Return ONLY valid JSON. Do not include markdown codeblocks.
<|user|>
Here is the raw input: {user_input}"""

If user_input contains strong markdown headers (like # Instructions), you can inspect the late-layer attention rows corresponding to the generation query token.

If the model correctly maintains system constraints, query rows for generated tokens maintain high attention attention values back to the <|system|> section key tokens. If prompt injection succeeds, attention values shift entirely to the user block tokens, dropping system token weights down to zero.

Observing this shift shows you the exact layer index where control gets lost. If system attention collapses at Layer 14, you know that safety alignment or system prompts are failing mid-network, giving you a clear metric to test prompt tweaks against.

Memory Optimizations for Long Traces

Inspecting full sequences on larger models demands deliberate memory management. A sequence length of 8,192 tokens creates an attention matrix of 8192 * 8192 = 67,108,864 elements per head per layer.

At float16 precision (2 bytes per value), a single attention head consumes 134 MB of RAM per forward pass. A model with 32 layers and 32 heads stores 1,024 such matrices, ballooning total trace size to 137 GB of memory for one output step. Running large models on constrained hardware poses similar memory challenges, as explored in techniques like AirLLM running 70B parameter models on a single 4GB GPU.

To prevent system crashes during execution:

1. Downsample sequence resolution: Aggregate adjacent token cells (for instance, averaging 4x4 token blocks into 1 cell) inside the PyTorch hook before transferring data out of GPU space.
2. Filter target layers: Hook only strategic transformer blocks instead of all layers. Extracting Layer 0 (inputs), Layer 16 (middle processing), and Layer 31 (final token selection) catches structural trends using 90% less memory.
3. Slice dynamic sub-ranges: Store attention slices for key regions, such as attending from output queries back to the initial system prompt range, rather than logging the entire sequence matrix.

# Downsampling tensor inside PyTorch hook to cut footprint by 75%
def get_downsampled_hook(layer_idx, block_size=2):
    def hook(module, input, output):
        if isinstance(output, tuple) and len(output) > 1:
            weights = output[1].detach()
            # Pool spatial dimensions using avg_pool2d
            # Shape: [batch, heads, seq_len, seq_len]
            downsampled = torch.nn.functional.avg_pool2d(
                weights, 
                kernel_size=block_size, 
                stride=block_size
            )
            captured_attentions[layer_idx] = downsampled.cpu().to(torch.float16)
    return hook

Tracing attention matrices takes execution out of the dark. By capturing hooks, running lightweight local visual tools, and reading attention structures like induction heads and sinks, you can diagnose exact model failure modes directly from execution data.