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}")