Attention Visualization › Extracting Attention Weights
Language AI HandbookAll chapters Attention Visualization › Extracting Attention WeightsOne section · Sections in this chapter3 / 15 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Attention weights are not outputs you receive by default when you call a language model. You need to explicitly request them and hook into the model's internals. Modern transformer libraries make this straightforward, but it is worth understanding what you are requesting and how the data is structured before you begin visualizing it.
Using Hugging Face Transformers
The Hugging Face transformers library provides a simple mechanism for extracting attention weights. When you call a model with output_attentions=True, the return value includes all attention matrices from all layers.
The attention output is a tuple of tensors, one per layer. Each tensor has shape (batch_size, num_heads, sequence_length, sequence_length). The entry at [batch, head, query_position, key_position] is the attention weight that the token at query_position places on the token at key_position, for that head.
import torch
from transformers import AutoModel, AutoTokenizer
# Load model and tokenizer
model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name, output_attentions=True)
model.eval()
# Example sentence
text = "The bank can guarantee deposits will eventually cover future tuition costs."
inputs = tokenizer(text, return_tensors="pt")
# Forward pass with attention output
with torch.no_grad():
outputs = model(**inputs, output_attentions=True)
# Extract attention weights: tuple of (num_layers,) tensors
# Each tensor: (batch=1, num_heads, seq_len, seq_len)
attention_weights = outputs.attentions
# Get tokens for display
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
num_layers = len(attention_weights)
num_heads = attention_weights[0].shape[1]
seq_len = attention_weights[0].shape[2]Number of layers: 12 Attention heads per layer: 12 Sequence length (with [CLS]/[SEP]): 14 Tokens: ['[CLS]', 'the', 'bank', 'can', 'guarantee', 'deposits', 'will', 'eventually', 'cover', 'future', 'tuition', 'costs', '.', '[SEP]']
BERT-base produces 12 layers, 12 heads each, and the sequence includes the special [CLS] token at position 0 and [SEP] at the end. These special tokens participate in attention like any other token, which is worth remembering when interpreting visualizations.
The Role of Special Tokens
The [CLS] and [SEP] tokens are not accidental participants in the attention computation. BERT was trained with [CLS] serving as the aggregate sentence representation: the vector at [CLS] after all 12 layers accumulates information from the entire sentence and is fed into the classification head for downstream tasks. This means that [CLS] must actively gather information, and many attention heads in later layers show substantial weight directed toward or from [CLS].
Similarly, [SEP] marks sentence boundaries. Research has shown that some attention heads treat [SEP] as a "null" or "no-op" target: when no other position in the sequence is highly relevant for a given query, the head routes attention to [SEP] as a way of adding minimal information. This is a learned behavior, not an architectural constraint, and it appears with striking consistency across heads.
Understanding the special token dynamics is important when reading attention heatmaps. A head that appears to strongly "attend to punctuation" may be attending to [SEP], which is functionally different. Similarly, a head whose attention is concentrated on [CLS] may be performing sentence-level aggregation rather than identifying an important word.
Accessing Individual Heads
Once you have the attention tuple, you can extract any specific layer and head combination. The convention is zero-indexed: layer 0 is the first transformer block, head 0 is the first attention head within that block.
# Extract layer 5, head 0 (zero-indexed)
layer_idx = 5
head_idx = 0
# Shape: (seq_len, seq_len)
attn_matrix = attention_weights[layer_idx][0, head_idx].numpy()Attention matrix shape: (14, 14) First row (attention FROM token 'the'): [CLS] 0.0357 the 0.0444 bank 0.2899 can 0.2238 guarantee 0.1598 deposits 0.0650 will 0.0054 eventually 0.0009 cover 0.0005 future 0.0001 tuition 0.0004 costs 0.0005 . 0.0017 [SEP] 0.1720
The first token in the sequence is always [CLS]. The output above shows which tokens the second token attends to when processing this sentence. Each row sums to 1.0 because of the softmax normalization.
The fact that every row sums to 1 has a subtle implication for interpretation. If a sequence has 14 tokens, the "average" attention weight for any given row is . A weight of 0.3 on a particular token is therefore roughly four times the average, which is meaningful. But if a sequence has 100 tokens, a weight of 0.3 on any single token stands out dramatically from the average. Interpreting attention weights without accounting for sequence length is a common mistake that can lead to comparing patterns across sentences of different lengths without adjusting for the different baselines.