Attention Visualization › Code Implementation: Systematic Head Analysis
Language AI HandbookAll chapters Attention Visualization › Code Implementation: Systematic Head AnalysisOne section · Sections in this chapter9 / 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
Rather than examining heads one at a time, a practical workflow involves computing summary statistics across the full dataset to find which heads are most distinctive. Let us walk through an example that identifies heads with the most focused attention and examines their patterns.
We first need a small dataset of sentences to analyze across. Using multiple sentences provides more reliable estimates of head behavior than any single example. The eight sentences below cover a range of syntactic structures and semantic domains, making sure that heads which show consistent patterns are general rather than fitting a specific input.
sentences = [
"The bank can guarantee deposits will eventually cover future tuition costs.",
"The cat sat on the mat and looked out the window.",
"Scientists discovered a new species of frog in the Amazon rainforest.",
"She opened the door and walked into the brightly lit room.",
"The algorithm processes each token in the sequence one step at a time.",
"His mother gave him a book about ancient Roman history for his birthday.",
"The company announced record profits despite difficult market conditions.",
"Researchers found that sleep deprivation impairs decision-making ability.",
]
# Process all sentences and collect entropy values
all_entropies = []
for sent in sentences:
enc = tokenizer(sent, return_tensors="pt")
with torch.no_grad():
out = model(**enc, output_attentions=True)
attn_weights_sent = out.attentions
sent_entropy = np.zeros((num_layers, num_heads))
for layer in range(num_layers):
for head in range(num_heads):
attn = attn_weights_sent[layer][0, head].numpy()
sent_entropy[layer, head] = attention_entropy(attn)
all_entropies.append(sent_entropy)
# Average entropy across sentences
mean_entropy = np.mean(all_entropies, axis=0) # shape: (num_layers, num_heads)5 most focused heads (lowest mean entropy): Layer 3, Head 1: entropy = 0.021 Layer 3, Head 10: entropy = 0.068 Layer 2, Head 7: entropy = 0.311 Layer 6, Head 2: entropy = 0.499 Layer 5, Head 4: entropy = 0.501 5 most diffuse heads (highest mean entropy): Layer 1, Head 1: entropy = 2.547 Layer 1, Head 7: entropy = 2.356 Layer 1, Head 5: entropy = 2.351 Layer 1, Head 9: entropy = 2.319 Layer 2, Head 9: entropy = 2.195
The heads with lowest entropy are the most selective: they consistently attend to a small number of positions across diverse sentences. These are the heads most worth examining visually, as they likely have the clearest interpretable function.
Identifying the Most Focused Head
Once you have ranked heads by entropy, you can retrieve and display the most focused one to examine what it is attending to. The minimum-entropy head across your test corpus is a strong candidate for having a clear, consistent function.
# Find the most focused head across the dataset
best_layer, best_head = np.unravel_index(
mean_entropy.argmin(), mean_entropy.shape
)
# Re-run the first sentence through this specific head for display
test_inputs = tokenizer(sentences[0], return_tensors="pt")
with torch.no_grad():
test_outputs = model(**test_inputs, output_attentions=True)
test_tokens = tokenizer.convert_ids_to_tokens(test_inputs["input_ids"][0])
focused_attn = test_outputs.attentions[best_layer][0, best_head].numpy()
Key Parameters
The key parameters in this analysis are:
- layer_idx: Which transformer layer to examine (0-indexed). Earlier layers capture surface-level patterns; later layers capture more abstract semantic information. For hypothesis generation about syntactic structure, start with layers 4-7. For semantic roles and coreference, start with layers 8-11.
- head_idx: Which attention head within the layer to examine. Heads are independently parameterized and may develop completely different attention strategies. There is no global convention for what head 0 or head 7 does across models; the identities of specialized heads vary with the random seed, model size, and pre-training data.
- entropy threshold: A quantitative cutoff for identifying focused heads. Heads below a threshold entropy value are candidates for closer qualitative inspection. A reasonable starting point is to examine heads below the 20th percentile of entropy across your dataset. This typically yields 20-30 candidates in a model like BERT-base.
- aggregation method: Whether to examine individual sentences or average across a corpus. Single examples reveal per-sentence behavior; corpus averages reveal stable head functions. For claims about model behavior in general, you need corpus-level aggregation. For debugging a specific prediction, single-example visualization is appropriate and useful.
The interaction between these parameters matters. A head may appear unfocused (high entropy) on a general corpus but become sharply focused on domain-specific text. A head that looks syntactic on standard English sentences may lose its syntactic character on code-switching text or non-standard dialects. Attention behavior is a function of both the model and the input distribution.