Info

Last Execution: 2026-09-03

Package	Version
nnsight	0.8.0
Python	3.12.13
torch	2.13.0+cu126
transformers	5.15.0.dev0

Attention Patterns and Per-Head Detection¶

Introduction¶

Every attention head in a transformer computes one matrix that says, for each query position, how much of its attention budget goes to each earlier position. That matrix is the attention pattern,

$$A = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d}}\right), \qquad A \in \mathbb{R}^{[\,\text{batch},\ \text{heads},\ \text{query},\ \text{key}\,]}$$

and it is the most direct read we have on what a head is looking at. Reading it lets us sort a model's heads into behavioural classes without running a single intervention.

In this tutorial we:

1. reach the probability matrix with nnsight's .source, which serves operations inside a module's forward — the pattern is a local variable in the attention implementation, not any module's output;
2. anchor the wiring by checking that the rows sum to 1 and that the causal mask holds;
3. score all 144 heads of GPT-2 small for three behaviours — induction, previous-token, and attention-sink — and rank them;
4. run the matching control on non-repeating text, so the effect sizes mean something;
5. check our top induction heads against the published list for GPT-2 small;
6. repeat the measurement on a grouped-query attention model, where the number of key heads and the number of query heads no longer agree.

The three behaviours we look for are all defined by a fixed offset in the pattern. Writing A[q, k] for the probability that query q puts on key k, and T for the period of a sequence that repeats itself:

head type	what it does	where the mass sits
previous-token	attends one step back	A[q, q-1]
duplicate-token	attends to the earlier copy of the current token	A[q, q-T]
induction	attends to the token after the earlier copy of the current token	A[q, q-T+1]
attention-sink	dumps mass on the first position	A[q, 0]

An induction head is the interesting one: together with a previous-token head it forms the two-head circuit behind in-context copying — "the last time I saw this token, what came next?" (Olsson et al. 2022). Because the behaviour is a single number per head, we can score all of them at once and see how concentrated it is.

📗 Prefer to use Colab? Follow the tutorial here!

Setup¶

In [1]:

try:
    import google.colab
    is_colab = True
except ImportError:
    is_colab = False

if is_colab:
    !pip install -U nnsight

try: import google.colab is_colab = True except ImportError: is_colab = False if is_colab: !pip install -U nnsight

In [2]:

from IPython.display import clear_output
import torch
torch.manual_seed(0)

import nnsight
from nnsight import TransformersModel

import plotly.express as px
import plotly.io as pio
pio.renderers.default = "colab" if is_colab else "plotly_mimetype+notebook_connected+colab+notebook"

from IPython.display import clear_output import torch torch.manual_seed(0) import nnsight from nnsight import TransformersModel import plotly.express as px import plotly.io as pio pio.renderers.default = "colab" if is_colab else "plotly_mimetype+notebook_connected+colab+notebook"

We load GPT-2 small with one extra argument: attn_implementation="eager". Eager attention is the slow, explicit implementation that actually builds the probability matrix in memory. The default on this model is sdpa, which fuses the whole operation and never writes the matrix down — we show what that costs you in a moment. For now, take the argument on faith; without it, everything below returns None.

In [3]:

model = TransformersModel(
    "openai-community/gpt2",
    task="text-generation",
    attn_implementation="eager",   # required: SDPA never materializes the probability matrix
    device_map="auto",
    dispatch=True,
)
clear_output()
print(model)

model = TransformersModel( "openai-community/gpt2", task="text-generation", attn_implementation="eager", # required: SDPA never materializes the probability matrix device_map="auto", dispatch=True, ) clear_output() print(model)

GPT2LMHeadModel(
  (transformer): GPT2Model(
    (wte): Embedding(50257, 768)
    (wpe): Embedding(1024, 768)
    (drop): Dropout(p=0.1, inplace=False)
    (h): ModuleList(
      (0-11): 12 x GPT2Block(
        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
        (attn): GPT2Attention(
          (c_attn): Conv1D()
          (c_proj): Conv1D()
          (attn_dropout): Dropout(p=0.1, inplace=False)
          (resid_dropout): Dropout(p=0.1, inplace=False)
        )
        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
        (mlp): GPT2MLP(
          (c_fc): Conv1D()
          (c_proj): Conv1D()
          (act): NewGELUActivation()
          (dropout): Dropout(p=0.1, inplace=False)
        )
      )
    )
    (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
  )
  (lm_head): Linear(in_features=768, out_features=50257, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
)

Where the attention pattern actually lives¶

The obvious place to look is the attention module's output — and it is the wrong place. model.transformer.h[0].attn.output is a tuple whose first element is the value-weighted result A V, already projected back to the residual stream. The probabilities A themselves are a local variable inside the module's forward, computed and consumed without ever becoming a submodule's input or output. A plain PyTorch forward hook structurally cannot see it.

nnsight's .source is the tool for exactly this case. It rewrites the module's forward so that every call site in it becomes an addressable operation with its own .input, .output and .skip. And because you cannot name an operation you have not seen, print(module.source) prints the forward back at you with each operation labelled in the left gutter.

In [4]:

print(model.transformer.h[0].attn.source)

print(model.transformer.h[0].attn.source)

                                             * def forward(
                                             0     self,
                                             1     hidden_states: tuple[torch.FloatTensor] | None,
                                             2     past_key_values: Cache | None = None,
                                             3     attention_mask: torch.FloatTensor | None = None,
                                             4     encoder_hidden_states: torch.Tensor | None = None,
                                             5     encoder_attention_mask: torch.FloatTensor | None = None,
                                             6     output_attentions: bool | None = False,
                                             7     **kwargs,
                                             8 ) -> tuple[torch.Tensor | tuple[torch.Tensor], ...]:
 is_cross_attention_0                    ->  9     is_cross_attention = encoder_hidden_states is not None
                                            10     if past_key_values is not None:
 isinstance_0                            -> 11         if isinstance(past_key_values, EncoderDecoderCache):
 past_key_values_is_updated_get_0        -> 12             is_updated = past_key_values.is_updated.get(self.layer_idx)
 is_updated_0                            ->  +             ...
                                            13             if is_cross_attention:
                                            14                 # after the first generated id, we can subsequently re-use all key/value_layer from cache
 curr_past_key_values_0                  -> 15                 curr_past_key_values = past_key_values.cross_attention_cache
                                            16             else:
 curr_past_key_values_1                  -> 17                 curr_past_key_values = past_key_values.self_attention_cache
                                            18         else:
 curr_past_key_values_2                  -> 19             curr_past_key_values = past_key_values
                                            20 
                                            21     if is_cross_attention:
 hasattr_0                               -> 22         if not hasattr(self, "q_attn"):
 ValueError_0                            -> 23             raise ValueError(
                                            24                 "If class is used as cross attention, the weights `q_attn` have to be defined. "
                                            25                 "Please make sure to instantiate class with `GPT2Attention(..., is_cross_attention=True)`."
                                            26             )
 self_q_attn_0                           -> 27         query_states = self.q_attn(hidden_states)
 query_states_0                          ->  +         ...
 attention_mask_0                        -> 28         attention_mask = encoder_attention_mask
                                            29 
                                            30         # Try to get key/value states from cache if possible
                                            31         if past_key_values is not None and is_updated:
 key_states_0                            -> 32             key_states = curr_past_key_values.layers[self.layer_idx].keys
 value_states_0                          -> 33             value_states = curr_past_key_values.layers[self.layer_idx].values
                                            34         else:
 self_c_attn_0                           -> 35             key_states, value_states = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2)
 split_0                                 ->  +             ...
 shape_kv_0                              -> 36             shape_kv = (*key_states.shape[:-1], -1, self.head_dim)
 key_states_view_0                       -> 37             key_states = key_states.view(shape_kv).transpose(1, 2)
 transpose_0                             ->  +             ...
 key_states_1                            ->  +             ...
 value_states_view_0                     -> 38             value_states = value_states.view(shape_kv).transpose(1, 2)
 transpose_1                             ->  +             ...
 value_states_1                          ->  +             ...
                                            39     else:
 self_c_attn_1                           -> 40         query_states, key_states, value_states = self.c_attn(hidden_states).split(self.split_size, dim=2)
 split_1                                 ->  +         ...
 shape_kv_1                              -> 41         shape_kv = (*key_states.shape[:-1], -1, self.head_dim)
 key_states_view_1                       -> 42         key_states = key_states.view(shape_kv).transpose(1, 2)
 transpose_2                             ->  +         ...
 key_states_2                            ->  +         ...
 value_states_view_1                     -> 43         value_states = value_states.view(shape_kv).transpose(1, 2)
 transpose_3                             ->  +         ...
 value_states_2                          ->  +         ...
                                            44 
 shape_q_0                               -> 45     shape_q = (*query_states.shape[:-1], -1, self.head_dim)
 query_states_view_0                     -> 46     query_states = query_states.view(shape_q).transpose(1, 2)
 transpose_4                             ->  +     ...
 query_states_1                          ->  +     ...
                                            47 
                                            48     if (past_key_values is not None and not is_cross_attention) or (
                                            49         past_key_values is not None and is_cross_attention and not is_updated
                                            50     ):
 curr_past_key_values_update_0           -> 51         key_states, value_states = curr_past_key_values.update(key_states, value_states, self.layer_idx)
                                            52         # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls
                                            53         if is_cross_attention:
 past_key_values_is_updated_0            -> 54             past_key_values.is_updated[self.layer_idx] = True
                                            55 
 using_eager_0                           -> 56     using_eager = self.config._attn_implementation == "eager"
 ALL_ATTENTION_FUNCTIONS_get_interface_0 -> 57     attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
 attention_interface_0                   ->  +     ...
                                            58         self.config._attn_implementation, eager_attention_forward
                                            59     )
                                            60 
                                            61     if using_eager and self.reorder_and_upcast_attn:
 self__upcast_and_reordered_attn_0       -> 62         attn_output, attn_weights = self._upcast_and_reordered_attn(
                                            63             query_states, key_states, value_states, attention_mask
                                            64         )
                                            65     else:
 attention_interface_1                   -> 66         attn_output, attn_weights = attention_interface(
                                            67             self,
                                            68             query_states,
                                            69             key_states,
                                            70             value_states,
                                            71             attention_mask,
                                            72             dropout=self.attn_dropout.p if self.training else 0.0,
                                            73             scaling=self.scaling,
                                            74             **kwargs,
                                            75         )
                                            76 
 attn_output_reshape_0                   -> 77     attn_output = attn_output.reshape(*attn_output.shape[:-2], -1).contiguous()
 contiguous_0                            ->  +     ...
 attn_output_0                           ->  +     ...
 self_c_proj_0                           -> 78     attn_output = self.c_proj(attn_output)
 attn_output_1                           ->  +     ...
 self_resid_dropout_0                    -> 79     attn_output = self.resid_dropout(attn_output)
 attn_output_2                           ->  +     ...
                                            80 
                                            81     return attn_output, attn_weights
                                            82 

That is transformers' real GPT2Attention.forward, annotated. Reading down the gutter we can see the whole computation: self_c_attn_1 makes QKV, the transpose_* operations reshape them into heads, and attention_interface_1 is the call that does the attention itself and returns (attn_output, attn_weights).

Two things to note before we use any of these names:

• The trailing number is an occurrence index within this forward, not a layer index, and calls and assignments share it: attention_interface_0 is the assignment a few lines up (attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(...), whose .output is the chosen function) and attention_interface_1 is the call. Every assignment in the forward is an operation named after its target, so a value that is never a call's result is still reachable.
• The numbering follows source order, and this forward has a cross-attention branch that a decoder-only GPT-2 never takes. That is why the first key transpose is transpose_2 rather than transpose_0: six operations in the dead branch have already consumed their numbers.

attention_interface_1 resolves at run time to a plain Python function, so we can chain .source a second time and read that function's operations too. Recursive .source only works inside a trace, because the call target is a local variable that only exists while the forward is running.

In [5]:

with model.trace("The Eiffel Tower is in Paris"):
    inner = nnsight.save(repr(model.transformer.h[0].attn.source.attention_interface_1.source))

print(inner)

with model.trace("The Eiffel Tower is in Paris"): inner = nnsight.save(repr(model.transformer.h[0].attn.source.attention_interface_1.source)) print(inner)

                             * def eager_attention_forward(module, query, key, value, attention_mask, scaling=None, dropout=0.0, **kwargs):
                             0     if scaling is None:
 query_size_0            ->  1         scaling = query.size(-1) ** -0.5
 scaling_0               ->  +         ...
                             2 
 key_transpose_0         ->  3     attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling
 torch_matmul_0          ->  +     ...
 attn_weights_0          ->  +     ...
                             4 
                             5     if attention_mask is not None:
 attn_weights_1          ->  6         attn_weights = attn_weights + attention_mask
                             7 
 nn_functional_softmax_0 ->  8     attn_weights = nn.functional.softmax(attn_weights, dim=-1)
 attn_weights_2          ->  +     ...
                             9 
                            10     # Downcast (if necessary) back to V's dtype (if in mixed-precision) -- No-Op otherwise
 attn_weights_type_0     -> 11     attn_weights = attn_weights.type(value.dtype)
 attn_weights_3          ->  +     ...
 nn_functional_dropout_0 -> 12     attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
 attn_weights_4          ->  +     ...
                            13 
 torch_matmul_1          -> 14     attn_output = torch.matmul(attn_weights, value)
 attn_output_0           ->  +     ...
 attn_output_transpose_0 -> 15     attn_output = attn_output.transpose(1, 2)
 attn_output_1           ->  +     ...
                            16 
                            17     return attn_output, attn_weights
                            18 

There it is, on the line labelled nn_functional_softmax_0: the probability matrix, one operation before the dtype cast and the dropout. So the handle we want is

model.transformer.h[L].attn.source.attention_interface_1.source.nn_functional_softmax_0.output

Two coarser routes give the same tensor on this model — attn.output[1] and attn.source.attention_interface_1.output[1], both of which are the weights that transformers happens to return. We use the softmax operation itself because it is the quantity in the definition: the raw softmax(QK^T/√d), before dropout, whether or not the module chooses to hand it back.

In [6]:

PROMPT = "The Eiffel Tower is in Paris. The Eiffel Tower is in"

with torch.no_grad():
    with model.trace(PROMPT):
        pattern = (
            model.transformer.h[5].attn
            .source.attention_interface_1
            .source.nn_functional_softmax_0.output.save()
        )

tokens = model.tokenizer(PROMPT)["input_ids"]
print("pattern shape :", tuple(pattern.shape), " = [batch, heads, query, key]")
print("tokens        :", [model.tokenizer.decode(t) for t in tokens])

PROMPT = "The Eiffel Tower is in Paris. The Eiffel Tower is in" with torch.no_grad(): with model.trace(PROMPT): pattern = ( model.transformer.h[5].attn .source.attention_interface_1 .source.nn_functional_softmax_0.output.save() ) tokens = model.tokenizer(PROMPT)["input_ids"] print("pattern shape :", tuple(pattern.shape), " = [batch, heads, query, key]") print("tokens :", [model.tokenizer.decode(t) for t in tokens])

pattern shape : (1, 12, 16, 16)  = [batch, heads, query, key]
tokens        : ['The', ' E', 'iff', 'el', ' Tower', ' is', ' in', ' Paris', '.', ' The', ' E', 'iff', 'el', ' Tower', ' is', ' in']

Note the torch.no_grad(). A trace runs with autograd on so that .backward() works inside it, which means every saved activation pins the forward graph alive. For a read-only capture like this one that is pure waste — on the 12-layer capture later in this notebook it is the difference between 83 MiB and 1.5 GiB.

Why attn_implementation="eager"¶

Now the argument we took on faith. Let's load the same model with its default attention implementation and ask for the weights.

In [7]:

sdpa_model = TransformersModel(
    "openai-community/gpt2", task="text-generation", device_map="auto", dispatch=True,
)
clear_output()
print("attention implementation:", sdpa_model.config._attn_implementation)

with torch.no_grad():
    with sdpa_model.trace(PROMPT):
        # ask for the inner operations first: they run before the module returns its output
        op_names = nnsight.save(
            list(sdpa_model.transformer.h[5].attn.source.attention_interface_1.source.names)
        )
        attn_output = sdpa_model.transformer.h[5].attn.output.save()

print("attn.output[0] :", tuple(attn_output[0].shape))
print("attn.output[1] :", attn_output[1])
print()
print("operations inside the SDPA implementation:")
print(op_names)

sdpa_model = TransformersModel( "openai-community/gpt2", task="text-generation", device_map="auto", dispatch=True, ) clear_output() print("attention implementation:", sdpa_model.config._attn_implementation) with torch.no_grad(): with sdpa_model.trace(PROMPT): # ask for the inner operations first: they run before the module returns its output op_names = nnsight.save( list(sdpa_model.transformer.h[5].attn.source.attention_interface_1.source.names) ) attn_output = sdpa_model.transformer.h[5].attn.output.save() print("attn.output[0] :", tuple(attn_output[0].shape)) print("attn.output[1] :", attn_output[1]) print() print("operations inside the SDPA implementation:") print(op_names)

attention implementation: sdpa
attn.output[0] : (1, 16, 768)
attn.output[1] : None

operations inside the SDPA implementation:
['kwargs_get_0', 'logger_warning_once_0', 'sdpa_kwargs_0', 'hasattr_0', 'use_gqa_in_sdpa_0', 'repeat_kv_0', 'key_0', 'repeat_kv_1', 'value_0', 'sdpa_kwargs_1', 'q_length_0', 'kv_length_0', 'getattr_0', 'is_causal_0', 'is_causal_1', 'torch_jit_is_tracing_0', 'isinstance_0', 'is_causal_item_0', 'is_causal_2', 'attention_mask_bool_0', 'torch_logical_not_0', 'to_0', 'attention_mask_0', 'key_1', 'value_1', 'position_bias_0', 'create_position_bias_mask_0', 'attention_mask_1', 'is_causal_3', 'torch_nn_functional_scaled_dot_product_attention_0', 'attn_output_0', 'attn_output_transpose_0', 'contiguous_0', 'attn_output_1']

The weights come back as None. No warning, no error — and if we had written pattern = attn.output[1] and gone straight to scoring, we would have crashed several cells later with a confusing message, or worse, silently scored zeros.

The operation list says why, and says it much more loudly than the None does: there is no softmax in it at all. SDPA calls one fused kernel, torch_nn_functional_scaled_dot_product_attention_0, which computes A V without ever writing A to memory. .source can reach any value a forward computes, but it cannot conjure a tensor that the kernel never wrote. Eager attention is the implementation that writes it down, which is why we need it — and it is also why you should not leave eager on for a production generation loop.

Compare the two lists: under eager we saw torch_matmul_0 → nn_functional_softmax_0 → nn_functional_dropout_0 → torch_matmul_1, the textbook attention. That diff is the clearest statement of what the flag changes.

In [8]:

del sdpa_model, attn_output
torch.cuda.empty_cache()

del sdpa_model, attn_output torch.cuda.empty_cache()

Anchoring the wiring¶

We have a tensor of the right shape. Before we read any science off it, we should check that it really is a matrix of attention probabilities. Two properties are cheap and both have to hold:

1. Every row sums to 1. It came out of a softmax over the key axis, so each query distributes exactly one unit of attention.
2. The causal mask holds. A decoder-only model cannot attend forward in time, so every entry strictly above the diagonal must be zero.

This is one cell, and it is what turns "the code ran" into "this is really attention".

In [9]:

n_query = pattern.shape[-1]
row_sums = pattern.sum(-1)
upper = torch.triu(torch.ones(n_query, n_query, dtype=torch.bool, device=pattern.device), diagonal=1)
above_diagonal = pattern[..., upper]

print(f"ANCHOR 1  row sums     : min {row_sums.min():.7f}  max {row_sums.max():.7f}"
      f"  max|sum-1| = {(row_sums - 1).abs().max():.2e}")
print(f"ANCHOR 2  causal mask  : {above_diagonal.numel()} entries above the diagonal,"
      f"  max value {above_diagonal.max():.3e},  all exactly zero: {bool((above_diagonal == 0).all())}")

n_query = pattern.shape[-1] row_sums = pattern.sum(-1) upper = torch.triu(torch.ones(n_query, n_query, dtype=torch.bool, device=pattern.device), diagonal=1) above_diagonal = pattern[..., upper] print(f"ANCHOR 1 row sums : min {row_sums.min():.7f} max {row_sums.max():.7f}" f" max|sum-1| = {(row_sums - 1).abs().max():.2e}") print(f"ANCHOR 2 causal mask : {above_diagonal.numel()} entries above the diagonal," f" max value {above_diagonal.max():.3e}, all exactly zero: {bool((above_diagonal == 0).all())}")

ANCHOR 1  row sums     : min 0.9999998  max 1.0000001  max|sum-1| = 1.79e-07
ANCHOR 2  causal mask  : 1440 entries above the diagonal,  max value 0.000e+00,  all exactly zero: True

Both pass. The row sums are 1 to within float32 rounding, and the upper triangle is exactly zero rather than merely small — transformers' eager path adds -inf to the masked scores before the softmax, so those entries are true zeros. A pattern with 1e-5 up there would tell us we had grabbed the pre-softmax scores, or a bidirectional model.

We can go one step further and rebuild the matrix ourselves. Q and K are also local variables of the same forward, so .source can hand them over, and attn.scaling is a plain attribute on the module. If our recomputation matches the captured tensor, there is no remaining ambiguity about what we are holding.

In [10]:

attn = model.transformer.h[5].attn

with torch.no_grad():
    with model.trace(PROMPT):
        # execution order matters: the key transpose runs before the query transpose.
        key = attn.source.transpose_2.output.save()
        query = attn.source.transpose_4.output.save()
        scaling = nnsight.save(attn.scaling)

causal = torch.triu(torch.full((n_query, n_query), float("-inf"), device=query.device), diagonal=1)
recomputed = torch.softmax(torch.matmul(query, key.transpose(-1, -2)) * scaling + causal, dim=-1)

print("query:", tuple(query.shape), " key:", tuple(key.shape), " scaling:", scaling, f"= 1/sqrt({int(1/scaling**2)})")
print(f"ANCHOR 3  max |recomputed - captured| = {(recomputed - pattern).abs().max():.3e}")

attn = model.transformer.h[5].attn with torch.no_grad(): with model.trace(PROMPT): # execution order matters: the key transpose runs before the query transpose. key = attn.source.transpose_2.output.save() query = attn.source.transpose_4.output.save() scaling = nnsight.save(attn.scaling) causal = torch.triu(torch.full((n_query, n_query), float("-inf"), device=query.device), diagonal=1) recomputed = torch.softmax(torch.matmul(query, key.transpose(-1, -2)) * scaling + causal, dim=-1) print("query:", tuple(query.shape), " key:", tuple(key.shape), " scaling:", scaling, f"= 1/sqrt({int(1/scaling**2)})") print(f"ANCHOR 3 max |recomputed - captured| = {(recomputed - pattern).abs().max():.3e}")

query: (1, 12, 16, 64)  key: (1, 12, 16, 64)  scaling: 0.125 = 1/sqrt(64)
ANCHOR 3  max |recomputed - captured| = 0.000e+00

Bit-identical. The tensor we are about to score is unambiguously $\mathrm{softmax}(QK^\top/\sqrt{d})$.

The task: sequences that repeat themselves¶

Induction is a claim about repetition: a head that, on seeing a token it has seen before, looks at what followed it last time. To measure it we need text where the model cannot possibly be using anything else — no grammar, no semantics, no memorized bigrams. The standard probe is a sequence of uniformly random tokens, repeated twice.

Each sequence is <|endoftext|> followed by T = 48 random token ids, then the same 48 ids again. If a head is doing induction, then for any query q in the second half, the token at q also appeared at q - T, and the token the head should predict next is the one at q - T + 1. So induction attention lands on offset q - T + 1, one diagonal stripe of the pattern.

We also build a control: sequences of the same length, same token distribution, but with no repeat at all. Every score we compute on the repeated batch we recompute here. Without it we cannot tell a genuine induction head from a head that likes a particular position.

In [11]:

BATCH, T = 16, 48
BOS = model.config.eos_token_id           # gpt2 has no dedicated BOS; <|endoftext|> plays the part
VOCAB = model.config.vocab_size

# avoid the very low / very high ids, which are punctuation and rarely-trained bytes
half = torch.randint(1000, VOCAB - 1000, (BATCH, T))
tokens_rep = torch.cat([torch.full((BATCH, 1), BOS), half, half], dim=1)

# control: same length, same distribution, no repeat
flat = torch.randint(1000, VOCAB - 1000, (BATCH, 2 * T))
tokens_ctl = torch.cat([torch.full((BATCH, 1), BOS), flat], dim=1)

SEQ = tokens_rep.shape[1]
print(f"batch {BATCH}, period T = {T}, sequence length {SEQ} = 1 + 2T")
print("repeat check: tokens[0, 1:6]  =", tokens_rep[0, 1:6].tolist())
print("             tokens[0, 49:54] =", tokens_rep[0, 1 + T:6 + T].tolist())
print("control     : tokens[0, 1:6]  =", tokens_ctl[0, 1:6].tolist(),
      " tokens[0, 49:54] =", tokens_ctl[0, 1 + T:6 + T].tolist())

BATCH, T = 16, 48 BOS = model.config.eos_token_id # gpt2 has no dedicated BOS; <|endoftext|> plays the part VOCAB = model.config.vocab_size # avoid the very low / very high ids, which are punctuation and rarely-trained bytes half = torch.randint(1000, VOCAB - 1000, (BATCH, T)) tokens_rep = torch.cat([torch.full((BATCH, 1), BOS), half, half], dim=1) # control: same length, same distribution, no repeat flat = torch.randint(1000, VOCAB - 1000, (BATCH, 2 * T)) tokens_ctl = torch.cat([torch.full((BATCH, 1), BOS), flat], dim=1) SEQ = tokens_rep.shape[1] print(f"batch {BATCH}, period T = {T}, sequence length {SEQ} = 1 + 2T") print("repeat check: tokens[0, 1:6] =", tokens_rep[0, 1:6].tolist()) print(" tokens[0, 49:54] =", tokens_rep[0, 1 + T:6 + T].tolist()) print("control : tokens[0, 1:6] =", tokens_ctl[0, 1:6].tolist(), " tokens[0, 49:54] =", tokens_ctl[0, 1 + T:6 + T].tolist())

batch 16, period T = 48, sequence length 97 = 1 + 2T
repeat check: tokens[0, 1:6]  = [23879, 16891, 13112, 15523, 6684]
             tokens[0, 49:54] = [23879, 16891, 13112, 15523, 6684]
control     : tokens[0, 1:6]  = [43344, 28996, 42802, 9042, 11325]  tokens[0, 49:54] = [28004, 15229, 34927, 49167, 47832]

Capturing every layer in one trace¶

We need the pattern from all 12 layers for the whole batch. The natural first instinct is a loop of forward passes, one per layer; that is 12 passes for something a single pass already computes.

The one-pass-per-layer version:

This runs the model 12 times to collect activations that a single forward already produces. It is the version to write first if you want to convince yourself the batched code below is doing the same thing, and the version to replace once you have.

patterns = []
for layer_idx in range(len(model.transformer.h)):
    with torch.no_grad():
        with model.trace({"input_ids": tokens_rep,
                          "attention_mask": torch.ones_like(tokens_rep)}):
            patterns.append(
                model.transformer.h[layer_idx].attn
                .source.attention_interface_1
                .source.nn_functional_softmax_0.output.save()
            )
patterns = torch.stack(patterns)

Inside one trace we can simply ask every layer for its pattern: the interventions are interleaved with the forward pass, so all 12 tensors come out of a single run. The whole capture goes inside torch.no_grad() — we are only reading.

In [12]:

def capture_patterns(tokens):
    """Return [layers, batch, heads, query, key] attention probabilities, on the CPU."""
    saved = {}
    with torch.no_grad():
        with model.trace({"input_ids": tokens, "attention_mask": torch.ones_like(tokens)}):
            for layer_idx, layer in enumerate(model.transformer.h):
                saved[layer_idx] = (
                    layer.attn
                    .source.attention_interface_1
                    .source.nn_functional_softmax_0.output.save()
                )

    return torch.stack([saved[i] for i in range(len(saved))]).cpu()


P_rep = capture_patterns(tokens_rep)
P_ctl = capture_patterns(tokens_ctl)

print("repeated batch :", tuple(P_rep.shape), "= [layers, batch, heads, query, key]")
print("control batch  :", tuple(P_ctl.shape))
print(f"held in memory : {P_rep.element_size() * P_rep.nelement() / 2**20:.1f} MiB per batch")
print(f"requires_grad  : {P_rep.requires_grad}   (thanks to torch.no_grad)")

def capture_patterns(tokens): """Return [layers, batch, heads, query, key] attention probabilities, on the CPU.""" saved = {} with torch.no_grad(): with model.trace({"input_ids": tokens, "attention_mask": torch.ones_like(tokens)}): for layer_idx, layer in enumerate(model.transformer.h): saved[layer_idx] = ( layer.attn .source.attention_interface_1 .source.nn_functional_softmax_0.output.save() ) return torch.stack([saved[i] for i in range(len(saved))]).cpu() P_rep = capture_patterns(tokens_rep) P_ctl = capture_patterns(tokens_ctl) print("repeated batch :", tuple(P_rep.shape), "= [layers, batch, heads, query, key]") print("control batch :", tuple(P_ctl.shape)) print(f"held in memory : {P_rep.element_size() * P_rep.nelement() / 2**20:.1f} MiB per batch") print(f"requires_grad : {P_rep.requires_grad} (thanks to torch.no_grad)")

repeated batch : (12, 16, 12, 97, 97) = [layers, batch, heads, query, key]
control batch  : (12, 16, 12, 97, 97)
held in memory : 82.7 MiB per batch
requires_grad  : False   (thanks to torch.no_grad)

Let's run the same two anchors on the full batch, since it costs nothing and the batched capture is new code.

In [13]:

upper_full = torch.triu(torch.ones(SEQ, SEQ, dtype=torch.bool), diagonal=1)
print(f"row sums   : max|sum-1| = {(P_rep.sum(-1) - 1).abs().max():.2e}")
print(f"causal mask: {P_rep[..., upper_full].numel():,} entries above the diagonal, "
      f"all exactly zero: {bool((P_rep[..., upper_full] == 0).all())}")

upper_full = torch.triu(torch.ones(SEQ, SEQ, dtype=torch.bool), diagonal=1) print(f"row sums : max|sum-1| = {(P_rep.sum(-1) - 1).abs().max():.2e}") print(f"causal mask: {P_rep[..., upper_full].numel():,} entries above the diagonal, " f"all exactly zero: {bool((P_rep[..., upper_full] == 0).all())}")

row sums   : max|sum-1| = 4.17e-07
causal mask: 10,727,424 entries above the diagonal, all exactly zero: True

Scoring every head¶

Each of our four behaviours is the average probability on one diagonal of the pattern. For a fixed offset d, the entries A[q, q-d] for all q are exactly torch.diagonal(A, offset=-d), so the whole score is one call plus a mean.

We restrict the query positions to the second repeat (q > T) for the repeat-dependent scores — before the repeat begins there is nothing for an induction head to find, and including those queries would halve every score for no reason. Previous-token and sink scores are averaged over all queries q ≥ 1.

In [14]:

def diagonal_score(P, offset, q_min):
    """Mean of P[..., q, q - offset] over batch and over query positions q >= q_min.

    Returns a [layers, heads] tensor."""
    diag = torch.diagonal(P, offset=-offset, dim1=-2, dim2=-1)   # element i is (q=i+offset, k=i)
    q_of = torch.arange(diag.shape[-1]) + offset
    return diag[..., q_of >= q_min].mean(dim=(1, -1))             # mean over batch and query


scores = {
    "induction":      diagonal_score(P_rep, T - 1, T + 1),    # A[q, q-T+1]
    "duplicate":      diagonal_score(P_rep, T,     T + 1),    # A[q, q-T]
    "previous_token": diagonal_score(P_rep, 1,     1),        # A[q, q-1]
    "sink":           P_rep[..., 1:, 0].mean(dim=(1, -1)),    # A[q, 0]
}
control = {
    "induction":      diagonal_score(P_ctl, T - 1, T + 1),
    "previous_token": diagonal_score(P_ctl, 1,     1),
    "sink":           P_ctl[..., 1:, 0].mean(dim=(1, -1)),
}

N_LAYERS, N_HEADS = scores["induction"].shape
print(f"scored {N_LAYERS} layers x {N_HEADS} heads = {N_LAYERS * N_HEADS} heads, four scores each")

def diagonal_score(P, offset, q_min): """Mean of P[..., q, q - offset] over batch and over query positions q >= q_min. Returns a [layers, heads] tensor.""" diag = torch.diagonal(P, offset=-offset, dim1=-2, dim2=-1) # element i is (q=i+offset, k=i) q_of = torch.arange(diag.shape[-1]) + offset return diag[..., q_of >= q_min].mean(dim=(1, -1)) # mean over batch and query scores = { "induction": diagonal_score(P_rep, T - 1, T + 1), # A[q, q-T+1] "duplicate": diagonal_score(P_rep, T, T + 1), # A[q, q-T] "previous_token": diagonal_score(P_rep, 1, 1), # A[q, q-1] "sink": P_rep[..., 1:, 0].mean(dim=(1, -1)), # A[q, 0] } control = { "induction": diagonal_score(P_ctl, T - 1, T + 1), "previous_token": diagonal_score(P_ctl, 1, 1), "sink": P_ctl[..., 1:, 0].mean(dim=(1, -1)), } N_LAYERS, N_HEADS = scores["induction"].shape print(f"scored {N_LAYERS} layers x {N_HEADS} heads = {N_LAYERS * N_HEADS} heads, four scores each")

scored 12 layers x 12 heads = 144 heads, four scores each

Induction heads¶

Ranked by induction score, with the control score of the same head on non-repeating text alongside — that column is what makes the first one a claim rather than a number.

In [15]:

def ranked(score, n=12):
    flat = score.flatten()
    order = flat.argsort(descending=True)[:n]
    return [(int(i) // N_HEADS, int(i) % N_HEADS, float(flat[i])) for i in order]


print(f"{'rank':>4}  {'head':>7}  {'induction':>9}  {'control':>8}  {'duplicate':>9}  {'prev':>6}  {'sink':>6}")
for rank, (l, h, s) in enumerate(ranked(scores["induction"]), start=1):
    print(f"{rank:>4}  {'L' + str(l) + 'H' + str(h):>7}  {s:>9.3f}  {control['induction'][l, h]:>8.3f}"
          f"  {scores['duplicate'][l, h]:>9.3f}  {scores['previous_token'][l, h]:>6.3f}"
          f"  {scores['sink'][l, h]:>6.3f}")

def ranked(score, n=12): flat = score.flatten() order = flat.argsort(descending=True)[:n] return [(int(i) // N_HEADS, int(i) % N_HEADS, float(flat[i])) for i in order] print(f"{'rank':>4} {'head':>7} {'induction':>9} {'control':>8} {'duplicate':>9} {'prev':>6} {'sink':>6}") for rank, (l, h, s) in enumerate(ranked(scores["induction"]), start=1): print(f"{rank:>4} {'L' + str(l) + 'H' + str(h):>7} {s:>9.3f} {control['induction'][l, h]:>8.3f}" f" {scores['duplicate'][l, h]:>9.3f} {scores['previous_token'][l, h]:>6.3f}" f" {scores['sink'][l, h]:>6.3f}")

rank     head  induction   control  duplicate    prev    sink
   1     L5H5      0.926     0.002      0.002   0.015   0.437
   2     L6H9      0.913     0.002      0.005   0.013   0.471
   3    L7H10      0.907     0.001      0.014   0.014   0.454
   4     L5H1      0.904     0.000      0.004   0.011   0.531
   5     L7H2      0.818     0.003      0.039   0.013   0.469
   6    L10H1      0.514     0.005      0.012   0.025   0.336
   7     L9H6      0.513     0.003      0.012   0.020   0.442
   8     L9H9      0.500     0.002      0.020   0.016   0.527
   9    L10H7      0.447     0.001      0.001   0.017   0.677
  10   L11H10      0.438     0.001      0.002   0.022   0.642
  11     L5H0      0.432     0.010      0.011   0.035   0.314
  12    L10H0      0.423     0.003      0.015   0.027   0.348

Five heads — L5H5, L6H9, L7H10, L5H1, L7H2 — put more than 80% of their attention on exactly the induction offset, and their control scores are all ≤ 0.003. The behaviour is caused by the repeat, not by the position: the same head on the same positions of a non-repeating sequence does nothing.

The duplicate column is worth a glance too. A head that attended to q - T (the earlier copy of the current token, rather than the one after it) would also look "repeat-aware", and it would be a different head type. None of the top five do that.

How concentrated is the behaviour across the model?

In [16]:

flat = scores["induction"].flatten()
sorted_scores = flat.sort(descending=True).values
chance = (1 / torch.arange(T + 2, SEQ + 1).float()).mean()   # uniform attention over allowed keys

print(f"median head             : {flat.median():.4f}")
print(f"control mean, all heads : {control['induction'].mean():.4f}")
print(f"uniform-attention chance: {chance:.4f}")
print(f"best head               : {sorted_scores[0]:.4f}   ({sorted_scores[0] / flat.median():.0f}x the median)")
print(f"heads above 0.10        : {int((flat > 0.10).sum())} / {flat.numel()}")
print(f"top-5 share of all induction mass: {100 * sorted_scores[:5].sum() / sorted_scores.sum():.1f}%")

flat = scores["induction"].flatten() sorted_scores = flat.sort(descending=True).values chance = (1 / torch.arange(T + 2, SEQ + 1).float()).mean() # uniform attention over allowed keys print(f"median head : {flat.median():.4f}") print(f"control mean, all heads : {control['induction'].mean():.4f}") print(f"uniform-attention chance: {chance:.4f}") print(f"best head : {sorted_scores[0]:.4f} ({sorted_scores[0] / flat.median():.0f}x the median)") print(f"heads above 0.10 : {int((flat > 0.10).sum())} / {flat.numel()}") print(f"top-5 share of all induction mass: {100 * sorted_scores[:5].sum() / sorted_scores.sum():.1f}%")

median head             : 0.0050
control mean, all heads : 0.0030
uniform-attention chance: 0.0141
best head               : 0.9259   (187x the median)
heads above 0.10        : 24 / 144
top-5 share of all induction mass: 38.6%

The median head scores 0.005 — below the 0.014 you would get from attending uniformly to every allowed position, i.e. most heads are actively looking somewhere else. Five heads out of 144 hold 39% of all the induction attention in the model. That is the shape of the result that makes induction heads worth naming: it is a property of a handful of specific heads, not a diffuse tendency of attention in general.

Previous-token heads¶

The other half of the induction circuit. A previous-token head writes the identity of token q-1 into position q, which is what lets an induction head at q match against position q-T and then read forward.

Here the control works differently, and that difference is itself a check. Attending one step back has nothing to do with repetition, so a real previous-token head must score the same on the control. If a head's previous-token score collapsed on non-repeating text, we would have mismeasured something.

In [17]:

print(f"{'rank':>4}  {'head':>7}  {'prev-token':>10}  {'control':>8}")
for rank, (l, h, s) in enumerate(ranked(scores["previous_token"], 6), start=1):
    print(f"{rank:>4}  {'L' + str(l) + 'H' + str(h):>7}  {s:>10.3f}  {control['previous_token'][l, h]:>8.3f}")

print(f"{'rank':>4} {'head':>7} {'prev-token':>10} {'control':>8}") for rank, (l, h, s) in enumerate(ranked(scores["previous_token"], 6), start=1): print(f"{rank:>4} {'L' + str(l) + 'H' + str(h):>7} {s:>10.3f} {control['previous_token'][l, h]:>8.3f}")

rank     head  prev-token   control
   1    L4H11       0.981     0.987
   2     L3H7       0.545     0.551
   3     L6H8       0.506     0.331
   4     L2H2       0.481     0.474
   5     L5H6       0.425     0.358
   6     L3H2       0.402     0.403

L4H11 puts 98% of its attention one step back and does so regardless of the input: 0.981 on the repeated batch against 0.987 on the control. The same invariance holds for L3H7, L2H2 and L3H2. It is worth reading the two columns rather than glancing at them — L6H8 and L5H6 drop noticeably on non-repeating text (0.506 → 0.331, 0.425 → 0.358), which says their attention one step back is partly conditional on the input rather than purely positional. The heads we would actually call previous-token heads are the ones whose control matches.

Note the layers: the strong previous-token heads live in layers 2–4 and the strong induction heads in layers 5–7. That ordering is what the circuit requires, and we did not put it there.

Attention sinks¶

The last class is the one that shows up whether you look for it or not. Many heads park most of their probability mass on position 0. With a softmax the row must sum to 1, so a head with nothing to do still has to attend somewhere; the first token is the conventional place to dump it (Xiao et al. 2023).

In [18]:

print(f"{'rank':>4}  {'head':>7}  {'sink':>6}  {'control':>8}  {'induction':>9}")
for rank, (l, h, s) in enumerate(ranked(scores["sink"], 6), start=1):
    print(f"{rank:>4}  {'L' + str(l) + 'H' + str(h):>7}  {s:>6.3f}  {control['sink'][l, h]:>8.3f}"
          f"  {scores['induction'][l, h]:>9.3f}")

print(f"\nmean attention to position 0 over all {flat.numel()} heads: {scores['sink'].mean():.3f}")

print(f"{'rank':>4} {'head':>7} {'sink':>6} {'control':>8} {'induction':>9}") for rank, (l, h, s) in enumerate(ranked(scores["sink"], 6), start=1): print(f"{rank:>4} {'L' + str(l) + 'H' + str(h):>7} {s:>6.3f} {control['sink'][l, h]:>8.3f}" f" {scores['induction'][l, h]:>9.3f}") print(f"\nmean attention to position 0 over all {flat.numel()} heads: {scores['sink'].mean():.3f}")

rank     head    sink   control  induction
   1     L5H7   0.936     0.939      0.000
   2     L8H4   0.919     0.919      0.000
   3     L7H8   0.898     0.891      0.000
   4     L7H4   0.882     0.909      0.000
   5     L7H9   0.878     0.961      0.000
   6    L9H10   0.874     0.851      0.000

mean attention to position 0 over all 144 heads: 0.438

The top sinks give position 0 nearly 90% of everything, on both the repeated and the control batch — unlike induction, this behaviour does not care what the input is.

A caveat that matters if you plan to use a sink score to classify heads: a high sink score is the absence of a behaviour, not the presence of one. Watch what happens on ordinary English, where there is no repeated structure for an induction head to lock onto.

In [19]:

sentences = [
    "The capital of France is Paris and the capital of Italy is Rome.",
    "She poured the coffee slowly, watching the steam rise from the cup.",
    "In 1969 the first humans walked on the surface of the moon.",
    "The theorem follows directly from the previous lemma and its corollary.",
]
# truncate to a common length instead of padding, so every position is a real token
encoded = [model.tokenizer(s)["input_ids"] for s in sentences]
n_keep = min(len(e) for e in encoded)
natural = torch.tensor([[BOS] + e[:n_keep] for e in encoded])

P_nat = capture_patterns(natural)
sink_natural = P_nat[..., 1:, 0].mean(dim=(1, -1))

print(f"{'head':>7}  {'sink (natural)':>14}  {'induction score':>15}")
for l, h, s in ranked(sink_natural, 6):
    print(f"{'L' + str(l) + 'H' + str(h):>7}  {s:>14.3f}  {scores['induction'][l, h]:>15.3f}")

sentences = [ "The capital of France is Paris and the capital of Italy is Rome.", "She poured the coffee slowly, watching the steam rise from the cup.", "In 1969 the first humans walked on the surface of the moon.", "The theorem follows directly from the previous lemma and its corollary.", ] # truncate to a common length instead of padding, so every position is a real token encoded = [model.tokenizer(s)["input_ids"] for s in sentences] n_keep = min(len(e) for e in encoded) natural = torch.tensor([[BOS] + e[:n_keep] for e in encoded]) P_nat = capture_patterns(natural) sink_natural = P_nat[..., 1:, 0].mean(dim=(1, -1)) print(f"{'head':>7} {'sink (natural)':>14} {'induction score':>15}") for l, h, s in ranked(sink_natural, 6): print(f"{'L' + str(l) + 'H' + str(h):>7} {s:>14.3f} {scores['induction'][l, h]:>15.3f}")

   head  sink (natural)  induction score
   L7H2           0.988            0.818
   L6H9           0.982            0.913
  L7H10           0.981            0.907
   L5H1           0.978            0.904
   L9H9           0.964            0.500
  L9H11           0.956            0.033

On natural text five of the six strongest sinks are induction heads from our ranking, and not one of the heads that topped the sink table on random tokens appears. An induction head with no repetition to match has nothing to point at, so it points at the first token. Rank heads by sink score on the wrong corpus and you will "discover" the induction heads and label them sinks. Only the offset-specific scores, measured on input built to elicit the behaviour, separate the classes.

Does the score measure what we think it measures?¶

The induction score reads exactly one diagonal. If a head's attention were spread over many offsets, or peaked at an offset one or two away from T-1, our number could still come out large-ish and we would have the story wrong. So let's stop assuming and sweep every offset: for each Δ from 1 to SEQ-1, the mean of A[q, q-Δ] over the second-repeat queries.

In [20]:

offsets = torch.arange(1, SEQ)
profile_heads = [(5, 5, "L5H5 (induction)"), (4, 11, "L4H11 (previous-token)"), (5, 7, "L5H7 (sink)")]

curves, labels = [], []
for l, h, name in profile_heads:
    P_head = P_rep[l:l + 1, :, h:h + 1]                       # keep the [layer, head] axes
    curves.append([float(diagonal_score(P_head, int(d), T + 1)[0, 0]) for d in offsets])
    labels.append(name)

fig = px.line(
    x=offsets.tolist(), y=curves[0], labels={"x": "offset Δ = q − k", "y": "mean attention"},
    title="Attention by offset, averaged over queries in the second repeat",
)
fig.data[0].name = labels[0]
for curve, label in zip(curves[1:], labels[1:]):
    fig.add_scatter(x=offsets.tolist(), y=curve, mode="lines", name=label)
fig.add_vline(x=T - 1, line_dash="dot", annotation_text="T−1 = induction")
fig.update_layout(showlegend=True)
fig.show()

for (l, h, name), curve in zip(profile_heads, curves):
    peak = int(torch.tensor(curve).argmax())
    print(f"{name:<26} peak at Δ = {offsets[peak]:>2}  ({max(curve):.3f});"
          f"  mass on position 0 = {float(P_rep[l, :, h, 1:, 0].mean()):.3f}")

offsets = torch.arange(1, SEQ) profile_heads = [(5, 5, "L5H5 (induction)"), (4, 11, "L4H11 (previous-token)"), (5, 7, "L5H7 (sink)")] curves, labels = [], [] for l, h, name in profile_heads: P_head = P_rep[l:l + 1, :, h:h + 1] # keep the [layer, head] axes curves.append([float(diagonal_score(P_head, int(d), T + 1)[0, 0]) for d in offsets]) labels.append(name) fig = px.line( x=offsets.tolist(), y=curves[0], labels={"x": "offset Δ = q − k", "y": "mean attention"}, title="Attention by offset, averaged over queries in the second repeat", ) fig.data[0].name = labels[0] for curve, label in zip(curves[1:], labels[1:]): fig.add_scatter(x=offsets.tolist(), y=curve, mode="lines", name=label) fig.add_vline(x=T - 1, line_dash="dot", annotation_text="T−1 = induction") fig.update_layout(showlegend=True) fig.show() for (l, h, name), curve in zip(profile_heads, curves): peak = int(torch.tensor(curve).argmax()) print(f"{name:<26} peak at Δ = {offsets[peak]:>2} ({max(curve):.3f});" f" mass on position 0 = {float(P_rep[l, :, h, 1:, 0].mean()):.3f}")

L5H5 (induction)           peak at Δ = 47  (0.926);  mass on position 0 = 0.437
L4H11 (previous-token)     peak at Δ =  1  (0.984);  mass on position 0 = 0.010
L5H7 (sink)                peak at Δ = 96  (0.959);  mass on position 0 = 0.936

L5H5 is a single spike at Δ = 47 = T − 1, exactly where an induction head should be, and essentially flat everywhere else. L4H11 spikes at Δ = 1. L5H7 has no peak at any fixed offset — instead its curve climbs at the far right of the plot, which is an artifact of the parameterisation rather than a behaviour: at Δ = q the key index is 0, so the largest offsets are just the sink column seen through an offset-shaped window. The printed "mass on position 0" is the honest description of that head.

That is the check that turns three numbers into three classifications. Notice that none of the three peaks at Δ = T = 48: that is where a duplicate-token head would sit, one position left of induction, and it would be easy to mistake for induction if we never looked. The duplicate column of the table above says the same thing for every head in the top twelve.

Visualising a head's pattern¶

The scores are summaries of matrices. Let's look at three of the matrices.

In [21]:

seq_idx = 0
show = torch.stack([P_rep[5, seq_idx, 5], P_rep[4, seq_idx, 11], P_rep[5, seq_idx, 7]])

fig = px.imshow(
    show.numpy(), facet_col=0, zmin=0, zmax=1, color_continuous_scale="Blues",
    labels={"x": "key position", "y": "query position", "color": "attention"},
    title="Attention patterns on a sequence of 48 random tokens, repeated twice",
)
for i, name in enumerate(["L5H5 — induction", "L4H11 — previous token", "L5H7 — attention sink"]):
    fig.layout.annotations[i]["text"] = name
fig.add_vline(x=T + 0.5, line_dash="dot", line_color="black")
fig.add_hline(y=T + 0.5, line_dash="dot", line_color="black")
fig.show()

seq_idx = 0 show = torch.stack([P_rep[5, seq_idx, 5], P_rep[4, seq_idx, 11], P_rep[5, seq_idx, 7]]) fig = px.imshow( show.numpy(), facet_col=0, zmin=0, zmax=1, color_continuous_scale="Blues", labels={"x": "key position", "y": "query position", "color": "attention"}, title="Attention patterns on a sequence of 48 random tokens, repeated twice", ) for i, name in enumerate(["L5H5 — induction", "L4H11 — previous token", "L5H7 — attention sink"]): fig.layout.annotations[i]["text"] = name fig.add_vline(x=T + 0.5, line_dash="dot", line_color="black") fig.add_hline(y=T + 0.5, line_dash="dot", line_color="black") fig.show()

The dotted lines mark where the second repeat begins.

L5H5 shows nothing but the position-0 column for the whole first repeat — there is nothing to match yet — and then grows a bright stripe the instant the second repeat starts. The stripe sits one position right of where a "same token" match would put it, which is the whole point: the head looks at what came next last time. L4H11 is a single sub-diagonal running the length of the sequence, identical before and after the repeat. L5H7 is one solid column at key 0.

All three have an empty upper triangle, which is anchor 2 made visible.

Now the same information for every head at once.

In [22]:

grid = torch.stack([scores["induction"], scores["previous_token"], scores["sink"]])

fig = px.imshow(
    grid.numpy(), facet_col=0, zmin=0, zmax=1, color_continuous_scale="Blues",
    labels={"x": "head", "y": "layer", "color": "score"},
    x=[str(h) for h in range(N_HEADS)], y=[str(l) for l in range(N_LAYERS)],
    title="Per-head scores across GPT-2 small",
)
for i, name in enumerate(["induction", "previous token", "attention sink"]):
    fig.layout.annotations[i]["text"] = name
fig.show()

grid = torch.stack([scores["induction"], scores["previous_token"], scores["sink"]]) fig = px.imshow( grid.numpy(), facet_col=0, zmin=0, zmax=1, color_continuous_scale="Blues", labels={"x": "head", "y": "layer", "color": "score"}, x=[str(h) for h in range(N_HEADS)], y=[str(l) for l in range(N_LAYERS)], title="Per-head scores across GPT-2 small", ) for i, name in enumerate(["induction", "previous token", "attention sink"]): fig.layout.annotations[i]["text"] = name fig.show()

Induction is five bright cells in layers 5–7 against an almost empty grid. Previous-token attention is one very bright cell at L4H11 with a scattering of weaker ones below it, all in earlier layers than the induction heads. Sink behaviour is the opposite kind of picture: broad, present in most of the second half of the model, and useless for identifying anything.

The science anchor: do we find the published heads?¶

Our numbers are internally consistent and controlled, but they could still be a consistent artifact of our own conventions. GPT-2 small is measured enough that we can check head by head. The induction-score sweep in ARENA's mechanistic interpretability curriculum — which uses the same probe, repeated random tokens, and the same diagonal convention — reports the top induction heads of GPT-2 small as 5.1, 5.5, 6.9, 7.2 and 7.10 (the layer.head notation is TransformerLens').

In [23]:

published = {(5, 1), (5, 5), (6, 9), (7, 2), (7, 10)}
ours = {(l, h) for l, h, _ in ranked(scores["induction"], 5)}

print("published top-5 :", sorted(published))
print("our top-5       :", sorted(ours))
print("set overlap     :", f"{len(published & ours)}/5")
print()
for l, h in sorted(published):
    print(f"  head {l}.{h:<2}  induction {scores['induction'][l, h]:.3f}"
          f"   control {control['induction'][l, h]:.3f}")

published = {(5, 1), (5, 5), (6, 9), (7, 2), (7, 10)} ours = {(l, h) for l, h, _ in ranked(scores["induction"], 5)} print("published top-5 :", sorted(published)) print("our top-5 :", sorted(ours)) print("set overlap :", f"{len(published & ours)}/5") print() for l, h in sorted(published): print(f" head {l}.{h:<2} induction {scores['induction'][l, h]:.3f}" f" control {control['induction'][l, h]:.3f}")

published top-5 : [(5, 1), (5, 5), (6, 9), (7, 2), (7, 10)]
our top-5       : [(5, 1), (5, 5), (6, 9), (7, 2), (7, 10)]
set overlap     : 5/5

  head 5.1   induction 0.904   control 0.000
  head 5.5   induction 0.926   control 0.002
  head 6.9   induction 0.913   control 0.002
  head 7.2   induction 0.818   control 0.003
  head 7.10  induction 0.907   control 0.001

5/5, with no head near the boundary — the sixth-place head scores 0.51 against 0.82 for the fifth. Our pipeline finds the published heads.

A note on which published list to compare against. A different set circulates for GPT-2 small: the IOI paper (Wang et al. 2022) labels 5.5, 5.8, 5.9 and 6.9 as the induction heads of its circuit. Two of those are ours; the other two are not close.

In [24]:

for l, h in [(5, 5), (5, 8), (5, 9), (6, 9)]:
    print(f"  head {l}.{h}  induction {scores['induction'][l, h]:.3f}"
          f"   sink {scores['sink'][l, h]:.3f}")

for l, h in [(5, 5), (5, 8), (5, 9), (6, 9)]: print(f" head {l}.{h} induction {scores['induction'][l, h]:.3f}" f" sink {scores['sink'][l, h]:.3f}")

  head 5.5  induction 0.926   sink 0.437
  head 5.8  induction 0.007   sink 0.724
  head 5.9  induction 0.019   sink 0.395
  head 6.9  induction 0.913   sink 0.471

This is not a contradiction, and it is worth understanding rather than averaging away. The two lists come from two different exercises. The IOI classification is behaviour on IOI prompts — natural sentences with a repeated name — and the paper is explicit that 5.8 and 5.9 are "fuzzy" induction heads, saying of 5.8 that its attention to the relevant position is "less than 0.1" and of 5.9 that it "is paying attention to S1+1 but also to tokens before S". Our numbers reproduce that caveat exactly: 5.8 scores 0.007 on the strict offset and dumps its mass on position 0 instead.

The ARENA list is the right comparator here because it is the same measurement we are making: the induction score on repeated random tokens. The lesson generalises — a head "is" an induction head only relative to a probe, and two reasonable probes disagree at the margin.

Grouped-query attention¶

Everything above assumed the number of query heads and the number of key/value heads are the same. In modern models they are not. Grouped-query attention (GQA) gives several query heads a single shared key/value head, which changes the bookkeeping — and it is easy to build a score table of the wrong shape.

The rule is: the attention pattern is indexed by query head. The key/value heads are broadcast up to the query-head count before the QK^T product, by an operation called repeat_kv. Let's watch that happen on Qwen2.5-0.5B, which has 14 query heads and 2 KV heads across 24 layers, and takes about 1 GB of GPU memory in bfloat16.

In [25]:

qwen = TransformersModel(
    "Qwen/Qwen2.5-0.5B",
    task="text-generation",
    attn_implementation="eager",
    dtype=torch.bfloat16,
    device_map="auto",
    dispatch=True,
)
clear_output()

cfg = qwen.config
head_dim = qwen.model.layers[0].self_attn._module.head_dim
print(f"layers {cfg.num_hidden_layers}, query heads {cfg.num_attention_heads}, "
      f"kv heads {cfg.num_key_value_heads}, head_dim {head_dim}")

qwen = TransformersModel( "Qwen/Qwen2.5-0.5B", task="text-generation", attn_implementation="eager", dtype=torch.bfloat16, device_map="auto", dispatch=True, ) clear_output() cfg = qwen.config head_dim = qwen.model.layers[0].self_attn._module.head_dim print(f"layers {cfg.num_hidden_layers}, query heads {cfg.num_attention_heads}, " f"kv heads {cfg.num_key_value_heads}, head_dim {head_dim}")

layers 24, query heads 14, kv heads 2, head_dim 64

In [26]:

QB, QT = 8, 48
gqa_generator = torch.Generator().manual_seed(0)
qwen_half = torch.randint(1000, 100_000, (QB, QT), generator=gqa_generator)
qwen_rep = torch.cat([torch.full((QB, 1), cfg.bos_token_id), qwen_half, qwen_half], dim=1)
qwen_flat = torch.randint(1000, 100_000, (QB, 2 * QT), generator=gqa_generator)
qwen_ctl = torch.cat([torch.full((QB, 1), cfg.bos_token_id), qwen_flat], dim=1)
QSEQ = qwen_rep.shape[1]

attn0 = qwen.model.layers[0].self_attn

with torch.no_grad():
    with qwen.trace({"input_ids": qwen_rep, "attention_mask": torch.ones_like(qwen_rep)}):
        q_states = attn0.source.transpose_0.output.save()
        k_states = attn0.source.transpose_1.output.save()
        inner_ops = nnsight.save(list(attn0.source.attention_interface_1.source.names))
        qwen_pattern = (
            attn0.source.attention_interface_1
            .source.nn_functional_softmax_0.output.save()
        )

print("query_states :", tuple(q_states.shape))
print("key_states   :", tuple(k_states.shape))
print("pattern      :", tuple(qwen_pattern.shape))
print()
print("operations inside the attention implementation:")
print(inner_ops)

QB, QT = 8, 48 gqa_generator = torch.Generator().manual_seed(0) qwen_half = torch.randint(1000, 100_000, (QB, QT), generator=gqa_generator) qwen_rep = torch.cat([torch.full((QB, 1), cfg.bos_token_id), qwen_half, qwen_half], dim=1) qwen_flat = torch.randint(1000, 100_000, (QB, 2 * QT), generator=gqa_generator) qwen_ctl = torch.cat([torch.full((QB, 1), cfg.bos_token_id), qwen_flat], dim=1) QSEQ = qwen_rep.shape[1] attn0 = qwen.model.layers[0].self_attn with torch.no_grad(): with qwen.trace({"input_ids": qwen_rep, "attention_mask": torch.ones_like(qwen_rep)}): q_states = attn0.source.transpose_0.output.save() k_states = attn0.source.transpose_1.output.save() inner_ops = nnsight.save(list(attn0.source.attention_interface_1.source.names)) qwen_pattern = ( attn0.source.attention_interface_1 .source.nn_functional_softmax_0.output.save() ) print("query_states :", tuple(q_states.shape)) print("key_states :", tuple(k_states.shape)) print("pattern :", tuple(qwen_pattern.shape)) print() print("operations inside the attention implementation:") print(inner_ops)

query_states : (8, 14, 97, 64)
key_states   : (8, 2, 97, 64)
pattern      : (8, 14, 97, 97)

operations inside the attention implementation:
['repeat_kv_0', 'key_states_0', 'repeat_kv_1', 'value_states_0', 'key_states_transpose_0', 'torch_matmul_0', 'attn_weights_0', 'attn_weights_1', 'nn_functional_softmax_0', 'to_0', 'attn_weights_2', 'nn_functional_dropout_0', 'attn_weights_3', 'torch_matmul_1', 'attn_output_0', 'attn_output_transpose_0', 'contiguous_0', 'attn_output_1']

key_states has 2 heads, query_states has 14, and the pattern has 14 — one row block per query head. The two repeat_kv operations at the top of the list are where 2 becomes 14, and they are visible as source operations rather than something you have to infer. So the score table is 24 x 14 = 336 entries, not 24 x 2.

The trap is on the key/value side: if you slice k_proj.output into heads you must use 2, and if you index a head of the pattern you are naming one of 14. Query head h reads key/value group h // (14 // 2).

Note also that the handle we used is character for character the same string as on GPT-2: self_attn.source.attention_interface_1.source.nn_functional_softmax_0.output. Both models go through transformers' eager_attention_forward, so the inner names carry over even though the module paths (transformer.h[i].attn vs model.layers[i].self_attn) do not.

Two things do not carry over, and both are worth knowing before you port anchor 3:

• transpose_0 and transpose_1 here are pre-RoPE. Rotary position embeddings are applied after them, in apply_rotary_pos_emb_0, so recomputing softmax(QK^T · scaling) from these two tensors does not reproduce the pattern — the error runs to several tenths. GPT-2 has no RoPE, which is why anchor 3 is exact there. On a rotary model, take Q and K from the output of apply_rotary_pos_emb_0 instead.
• Some families do not compute a plain softmax attention at all. Gemma-2 softcaps the attention logits (attn_logit_softcapping = 50.0) between the matmul and the softmax; its operation list carries a torch_tanh_0 that neither of the models here has. Read the list before assuming the textbook formula.

In [27]:

def capture_qwen(tokens):
    saved = {}
    with torch.no_grad():
        with qwen.trace({"input_ids": tokens, "attention_mask": torch.ones_like(tokens)}):
            for i, layer in enumerate(qwen.model.layers):
                saved[i] = (
                    layer.self_attn
                    .source.attention_interface_1
                    .source.nn_functional_softmax_0.output.save()
                )
    return torch.stack([saved[i] for i in range(len(saved))]).float().cpu()


QP_rep, QP_ctl = capture_qwen(qwen_rep), capture_qwen(qwen_ctl)
upper_q = torch.triu(torch.ones(QSEQ, QSEQ, dtype=torch.bool), diagonal=1)
print(f"shape {tuple(QP_rep.shape)}")
print(f"row sums   : max|sum-1| = {(QP_rep.sum(-1) - 1).abs().max():.2e}")
print(f"causal mask: all exactly zero above the diagonal: {bool((QP_rep[..., upper_q] == 0).all())}")

q_ind = diagonal_score(QP_rep, QT - 1, QT + 1)
q_ind_ctl = diagonal_score(QP_ctl, QT - 1, QT + 1)
Q_HEADS = q_ind.shape[1]
group = Q_HEADS // cfg.num_key_value_heads

print(f"\n{'head':>8}  {'kv group':>8}  {'induction':>9}  {'control':>8}")
flat_q = q_ind.flatten()
for i in flat_q.argsort(descending=True)[:6]:
    l, h = int(i) // Q_HEADS, int(i) % Q_HEADS
    print(f"{'L' + str(l) + 'H' + str(h):>8}  {h // group:>8}  {flat_q[i]:>9.3f}  {q_ind_ctl[l, h]:>8.3f}")

srt_q = flat_q.sort(descending=True).values
print(f"\nheads above 0.10: {int((flat_q > 0.10).sum())} / {flat_q.numel()}"
      f"   top-5 share of induction mass: {100 * srt_q[:5].sum() / srt_q.sum():.1f}%")

def capture_qwen(tokens): saved = {} with torch.no_grad(): with qwen.trace({"input_ids": tokens, "attention_mask": torch.ones_like(tokens)}): for i, layer in enumerate(qwen.model.layers): saved[i] = ( layer.self_attn .source.attention_interface_1 .source.nn_functional_softmax_0.output.save() ) return torch.stack([saved[i] for i in range(len(saved))]).float().cpu() QP_rep, QP_ctl = capture_qwen(qwen_rep), capture_qwen(qwen_ctl) upper_q = torch.triu(torch.ones(QSEQ, QSEQ, dtype=torch.bool), diagonal=1) print(f"shape {tuple(QP_rep.shape)}") print(f"row sums : max|sum-1| = {(QP_rep.sum(-1) - 1).abs().max():.2e}") print(f"causal mask: all exactly zero above the diagonal: {bool((QP_rep[..., upper_q] == 0).all())}") q_ind = diagonal_score(QP_rep, QT - 1, QT + 1) q_ind_ctl = diagonal_score(QP_ctl, QT - 1, QT + 1) Q_HEADS = q_ind.shape[1] group = Q_HEADS // cfg.num_key_value_heads print(f"\n{'head':>8} {'kv group':>8} {'induction':>9} {'control':>8}") flat_q = q_ind.flatten() for i in flat_q.argsort(descending=True)[:6]: l, h = int(i) // Q_HEADS, int(i) % Q_HEADS print(f"{'L' + str(l) + 'H' + str(h):>8} {h // group:>8} {flat_q[i]:>9.3f} {q_ind_ctl[l, h]:>8.3f}") srt_q = flat_q.sort(descending=True).values print(f"\nheads above 0.10: {int((flat_q > 0.10).sum())} / {flat_q.numel()}" f" top-5 share of induction mass: {100 * srt_q[:5].sum() / srt_q.sum():.1f}%")

shape (24, 8, 14, 97, 97)
row sums   : max|sum-1| = 3.58e-07
causal mask: all exactly zero above the diagonal: True

    head  kv group  induction   control
   L16H3         0      0.984     0.010
   L9H13         1      0.969     0.004
  L11H12         1      0.968     0.002
   L16H2         0      0.943     0.011
  L16H12         1      0.916     0.006
   L16H7         1      0.884     0.003

heads above 0.10: 56 / 336   top-5 share of induction mass: 21.8%

The anchors pass on Qwen too, and the same measurement finds the same kind of structure: several very sharp induction heads (L16H3 at 0.984 is as clean as anything in GPT-2) and controls at chance. But the behaviour is less concentrated — 56 of 336 heads clear 0.10, and the top five hold 21.8% of the induction mass against 39% for GPT-2's 144 heads. With 2.3x more heads to spread across, that is roughly what you would expect.

The kv-group column is there to be read: the six strongest heads split across both groups, so nothing about induction is confined to one shared key/value head.

Unlike GPT-2, there is no published per-head induction list for Qwen2.5-0.5B to check against, so treat this as a portability result — the same code, two lines changed, on a different architecture — rather than an anchored one.