nnterp · encyclopedia gpt2 · nnterp.families.gpt2

family gpt2

GPT-2

Llama's sequential block with LayerNorm in place of RMSNorm, a learned position embedding added to the token embeddings before block 0, and queries, keys and values cut from one fused c_attn projection.

Position embeddingsLayerNorm

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card.

embed_tokenstoken_embeddings
→
layers× 12
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsthe output
i = 0

identity layers[i].input + self_attn.attention_output + mlp.mlp_output == layer_output

The API

02 · support()

Values, by host

⚠︎ marks a value this checkpoint has only under some load: hover it for the condition.

model GPT2LMHeadModellayoutwhat it isread at
logits Logits [batch seq vocab_size] The logits the model returns, after anything it does past lm_head (a softcap, a scale)

the module's own output
token_embeddings Residual [batch seq hidden_size] The embedding module's output; layers[0].input is what enters the first block

the child module embed_tokens's output
final_norm_output Residual [batch seq hidden_size] The residual stream after the final norm, what lm_head reads

the child module norm's output
lm_head_output Logits [batch seq vocab_size] The head's logits before anything the model does past lm_head (a softcap, a scale)

the child module lm_head's output
next_token_probs NextTokenProbs [batch vocab_size] The next-token distribution at the last position; derived, read-only

the module's own output
input_ids Tokens [batch seq] The token ids the model was called with

the module's own input
attention_mask Tokens [batch seq] The attention mask the model was called with; zeros are padding

the module's own input
input_size scalar [batch, seq] of the current call; read-only

the module's own input
model.layers[i] GPT2Blocklayoutwhat it isread at
layer_input Residual [batch seq hidden_size] The residual stream entering the block

the module's own input
attention_norm_output Residual [batch seq hidden_size] The normed stream entering the attention

the child module input_layernorm's output
layer_mid Residual [batch seq hidden_size] The residual stream after the attention add

the child module post_attention_layernorm's input
mlp_norm_output Residual [batch seq hidden_size] The normed stream entering the MLP

the child module post_attention_layernorm's output
mlp_input Residual [batch seq hidden_size] What the MLP is called with

the child module mlp's input
layer_output Residual [batch seq hidden_size] The residual stream leaving the block, a tensor even when the block returns a tuple

the module's own output
model.layers[i].self_attn GPT2Attentionlayoutwhat it isread at
⚠︎attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention

inside the forward: the arguments of attention_interface_1, element 1
⚠︎attention_keys Keys [batch num_kv_heads seq qk_head_dim] The keys entering attention

inside the forward: the arguments of attention_interface_1, element 2
⚠︎attention_values Values [batch num_kv_heads seq head_dim] The values entering attention

inside the forward: the arguments of attention_interface_1, element 3
⚠︎attention_scores Pattern [batch num_heads query key] The attention scores entering the softmax, masked

inside the forward: the input of nn_functional_softmax_0 inside attention_interface_1
attention_premix HeadOutputs [batch seq num_heads head_dim] Every head's output side by side, as the output projection takes it

the child module o_proj's input
attention_output Residual [batch seq hidden_size] What the attention adds to the residual stream

the module's own output
⚠︎attention_probabilities Pattern [batch num_heads query key] The attention pattern the values are mixed with

inside the forward: the output of nn_functional_dropout_0 inside attention_interface_1
⚠︎attention_head_outputs HeadOutputs [batch seq num_heads head_dim] The per-head outputs before the output projection

inside the forward: the output of attention_interface_1, element 0
model.layers[i].mlp GPT2MLPlayoutwhat it isread at
mlp_activation Neurons [batch seq intermediate_size] The activation function's output

the child module act_fn's output
mlp_neurons Neurons [batch seq intermediate_size] What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise

the child module down_proj's input
mlp_output Residual [batch seq hidden_size] What the MLP adds to the residual stream

the module's own output

Printout

What print(model) shows on openai-community/gpt2: the native modules, standard names and every standard value with its layout.

GPT2LMHeadModel(
  (lm_head): Linear(in_features=768, out_features=50257, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(50257, 768)
  (layers): ModuleList(
    (0-11): 12 x GPT2Block(
      (input_layernorm/ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
      (self_attn/attn): GPT2Attention(
        (c_attn): Conv1D()
        (o_proj/c_proj): Conv1D()
        (attn_dropout): Dropout(p=0.1, inplace=False)
        (resid_dropout): Dropout(p=0.1, inplace=False)
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (post_attention_layernorm/ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
      (mlp): GPT2MLP(
        (c_fc): Conv1D()
        (down_proj/c_proj): Conv1D()
        (act_fn/act): NewGELUActivation()
        (dropout): Dropout(p=0.1, inplace=False)
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
  )
  (norm): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
  (logits) -> Logits [batch seq vocab_size]: The logits the model returns, after anything it does past lm_head (a softcap, a scale)
  (token_embeddings) -> Residual [batch seq hidden_size]: The embedding module's output; layers[0].input is what enters the first block
  (final_norm_output) -> Residual [batch seq hidden_size]: The residual stream after the final norm, what lm_head reads
  (lm_head_output) -> Logits [batch seq vocab_size]: The head's logits before anything the model does past lm_head (a softcap, a scale)
  (next_token_probs) -> NextTokenProbs [batch vocab_size]: The next-token distribution at the last position; derived, read-only
  (input_ids) -> Tokens [batch seq]: The token ids the model was called with
  (attention_mask) -> Tokens [batch seq]: The attention mask the model was called with; zeros are padding
  (input_size): [batch, seq] of the current call; read-only
)

Config

model.num_layers12
model.hidden_size768
model.num_heads12
model.num_kv_heads12
model.head_dim64
model.qk_head_dim64
model.vocab_size50257
model.intermediate_size3072
architecturesGPT2LMHeadModel
vocab_size50257
tie_word_embeddingsTrue
n_layer12
n_embd768
n_head12
n_positions1024
activation_functiongelu_new
layer_norm_epsilon1e-05
scale_attn_by_inverse_layer_idxFalse
reorder_and_upcast_attnFalse

The Notes

03 · what to know before you run

The block, in order

h   = x + attn(ln_1(x))      # ln_1 = input_layernorm
out = h + mlp(ln_2(h))       # ln_2 = post_attention_layernorm

The block returns a bare tensor. Nothing sits between a sublayer and the stream: attention_output is self_attn.output[0] (after c_proj) and mlp_output is mlp.output (after the MLP's c_proj), and the identity layers[i].input + attention_output + mlp_output == layer_output is exact in float32 (difference 0.0 on gpt2). Both norms are LayerNorm with a weight and a bias; ln_2 is the MLP's input norm, as Llama's post_attention_layernorm is.

Position embeddings are added after embed_tokens

token_embeddings is wte's lookup only. The learned position embedding transformer.wpe, one row per position up to n_positions (1024), is added to it before block 0, so the stream entering block 0 is the sum, and an edit to token_embeddings leaves the position term in place:

with model.trace(prompt):
    tokens = model.token_embeddings.save()
    positions = model.transformer.wpe.output.save()
    x0 = model.layers[0].input.save()

torch.equal(tokens + positions, x0)        # True

A prompt longer than n_positions fails in wpe with IndexError. Positions are absolute, so a left-padded row would read shifted rows of wpe; nnsight derives position_ids from the attention mask for a left-padded batch, and a padded prompt gives the same logits as the prompt alone (largest difference 1.3e-4 on gpt2) and the same greedy generation.

The first position is a sink with a huge norm

On gpt2 the stream at position 0 has a norm of 2,500 to 3,200 from block 2's output to block 10's, whatever the token, against medians of about 60 to 250 at the other positions. Position 0 also takes most of the attention: averaged over heads and later queries, 41% to 52% of the pattern lands on key 0 in blocks 1 and 2, and 57% to 91% in blocks 3 to 11. Leave position 0 out of activation statistics, SAE metrics and the norms you scale a steering vector by.

The tokenizer adds no BOS token, so without one the first word of the prompt becomes the sink. Writing <|endoftext|> (id 50256, the BOS and EOS token) at the start of the prompt gives the sink a fixed token; it tokenizes to that one id. TransformerLens prepends it by default, so activations recorded with its defaults have it at position 0. Words inside a sentence are tokens with their leading space (" Paris" is 6342, "Paris" is 40313), and a prompt ending in a space ends with the lone-space token 220.

Queries, keys and values are cut from one Conv1D

The projections are Conv1D, whose weight is stored [in, out] (c_attn.weight is [768, 2304] on gpt2): the module computes x @ weight + bias, the transpose of nn.Linear's layout. The queries', keys' and values' weights are column blocks of c_attn.weight, and head h is columns h * head_dim to (h + 1) * head_dim within a block:

attn = model.layers[3].self_attn
W_Q, W_K, W_V = attn.c_attn.weight.split(model.hidden_size, dim=1)
b_Q, b_K, b_V = attn.c_attn.bias.split(model.hidden_size)

with model.trace(prompt):
    h = model.layers[3].self_attn.input.save()                  # ln_1's output
    q = model.layers[3].self_attn.attention_queries.save()

heads = (model.num_heads, model.head_dim)
q_again = (h @ W_Q + b_Q).unflatten(-1, heads).transpose(1, 2)
torch.testing.assert_close(q_again, q)

The queries are a view of the split c_attn output, so with gradients on torch refuses an in-place edit of attention_queries; the keys and values take one. Edit a copy of the queries and assign it, or edit them under torch.no_grad():

with model.trace(prompt):
    q = model.layers[3].self_attn.attention_queries.clone()
    q[:, 0] = 0                                                 # head 0's queries
    model.layers[3].self_attn.attention_queries = q

Every head has its own keys and values (num_kv_heads == num_heads), and the scores are scaled by head_dim ** -0.5 (0.125).

Load with eager; two config flags change the attention

The default load runs sdpa, which has no queries, pattern or head outputs to read; attn_implementation="eager" runs transformers' eager attention function, where nnterp reads them. attention_output and mlp_output are there under either load.

The Stanford CRFM checkpoints (stanford-crfm/*-gpt2-small-* and *-gpt2-medium-*, five seeds each, with intermediate training checkpoints as revision tags such as checkpoint-99000) set reorder_and_upcast_attn and scale_attn_by_inverse_layer_idx. scale_attn_by_inverse_layer_idx divides block i's score scale by i + 1, on every attention path. reorder_and_upcast_attn makes an eager load take GPT-2's own upcast path, where nnterp reads nothing inside the attention. In float32 that path computes what the eager attention function does, up to rounding: on alias-gpt2-small-x21 the logits of the two paths agree within 2e-5. A float32 load with the flag off therefore serves the interior values and computes the same logits to within that rounding:

model = StandardizedTransformer(
    "stanford-crfm/alias-gpt2-small-x21", dispatch=True, dtype=torch.float32,
    attn_implementation="eager", reorder_and_upcast_attn=False,
)

alias-gpt2-small-x21 loads in float16 by default, and there the paths differ, since the upcast one computes the scores in float32 (largest logit difference 0.025 on GPU).

The readout: LayerNorm with a bias, tied weights

logits is lm_head.output with nothing applied, and project_on_vocab is lm_head(ln_f(x)), so a logit lens at the last block equals logits exactly. ln_f subtracts the mean over the hidden axis, so adding the same number to every coordinate of the stream changes no logit. Its bias adds lm_head(ln_f.bias) to every position's logits whatever the input: on gpt2 that vector is highest on ,, the, and, . and newline, with a standard deviation of 0.9 across the vocabulary.

Folding the final norm into the unembedding takes the gain into W_U and the bias into a vocabulary bias. With the scale frozen at the run's value the readout is linear, which is what direct logit attribution needs. Each contribution's term is that contribution, minus its mean, through the same map, and the terms of layers[0].input and of every attention_output and mlp_output, plus the bias's, sum to the logits (largest difference 1.3e-4 on gpt2, on logits of size near 100).

ln_f, W_U = model.norm._module, model.lm_head.weight
with model.trace(prompt):
    resid = model.layers[-1].layer_output[0, -1].save()
    logits = model.logits[0, -1].save()

scale = (resid.var(unbiased=False) + ln_f.eps).sqrt()          # frozen
readout = lambda v: (ln_f.weight * (v - v.mean()) / scale) @ W_U.T
logits_again = readout(resid) + ln_f.bias @ W_U.T
torch.testing.assert_close(logits_again, logits, atol=1e-3, rtol=1e-5)

lm_head has no bias and its weight is embed_tokens.weight, the same tensor, so an edit to the embedding matrix is an edit to the unembedding.

The family's checkpoints

The four OpenAI sizes are 12 × 768 (gpt2), 24 × 1024 (gpt2-medium), 36 × 1280 (gpt2-large) and 48 × 1600 (gpt2-xl, 25 heads), blocks × width; every size has head_dim 64, gelu_new, an MLP four times the width, 1024 positions and the same 50257-token vocabulary. distilgpt2 is 6 blocks of gpt2's width. Every checkpoint whose config says model_type: gpt2 takes this family (DialoGPT's and ProtGPT2's do); check its config for the two flags above.

TransformerLens hook names

Measured on gpt2 against TransformerLens 2.18's HookedTransformer.from_pretrained_no_processing, block 6, each pair equal to within 1e-4:

  • hook_embed: model.token_embeddings
  • hook_pos_embed: model.transformer.wpe.output
  • blocks.i.hook_resid_pre: model.layers[i].input
  • blocks.i.attn.hook_q, hook_k, hook_v: self_attn.attention_queries, attention_keys, attention_values, transposed (TransformerLens puts the position axis before the head axis)
  • blocks.i.attn.hook_attn_scores: self_attn.attention_scores, on the unmasked entries
  • blocks.i.attn.hook_pattern: self_attn.attention_probabilities
  • blocks.i.attn.hook_z: self_attn.attention_head_outputs
  • blocks.i.hook_attn_out: self_attn.attention_output
  • blocks.i.hook_resid_mid: model.layers[i].layer_mid
  • blocks.i.mlp.hook_post: model.layers[i].mlp.mlp_activation
  • blocks.i.hook_mlp_out: mlp.mlp_output
  • blocks.i.hook_resid_post: model.layers[i].layer_output
  • ln_final.hook_normalized: model.norm.output

HookedTransformer.from_pretrained, the default, rewrites the weights, and the hooks move with them. It centers every weight that writes to the stream, so each residual hook and hook_attn_out / hook_mlp_out equals the nnterp value minus its mean over the hidden axis. It folds ln_1's bias and the value bias out of the values, so hook_v and hook_z equal the nnterp values minus ln_1.bias @ W_V + b_V, per head. ln_final.hook_normalized drops ln_f's gain and bias, and the logits are nnterp's minus their mean over the vocabulary. The queries, keys, scores, pattern and mlp.hook_post are unchanged.

Sparse autoencoders and neuron explanations

gpt2-small-res-jb (SAELens; Hub jbloom/GPT2-Small-SAEs-Reformatted) has 24576 latents on blocks.i.hook_resid_pre for every block, plus blocks.11.hook_resid_post. They were trained on TransformerLens' centered weights (SAELens loads them with center_writing_weights), so their input is model.layers[i].input minus its mean over the hidden axis. On a 59-token text (position 0 left out), block 8's SAE leaves 15% of the variance unexplained on the centered stream and 250% on the raw one.

with model.trace(prompt):
    resid = model.layers[8].input.save()

sae_input = resid - resid.mean(-1, keepdim=True)      # gpt2-small-res-jb's input

OpenAI's v5 SAEs (openai/sparse_autoencoder; also in SAELens as gpt2-small-resid-post-v5-32k and its kin) are TopK autoencoders with 32k and 128k latents at four places per block: resid_delta_attn (attention_output), resid_post_attn (layer_mid), resid_delta_mlp (mlp_output) and resid_post_mlp (layer_output). The repository's README extracts activations with center_writing_weights=False, and the autoencoder layer-normalizes its input itself, so the nnterp value goes in as it is: block 6's resid_post_mlp SAE reconstructs layers[6].layer_output with a normalized MSE of 0.058 on the same text.

OpenAI's neuron explanations (openai/automated-interpretability) cover every MLP neuron of gpt2-xl (48 blocks × 6400) and, in a second set, of gpt2. A neuron's activation is model.layers[i].mlp.mlp_activation[..., n], the gelu_new output. The README says the published activations were computed with a different GELU and differ from the exact ones by a median of 0.009 on gpt2.

From the family module

04 · nnterp/families/gpt2.py

GPT-2 (GPT2LMHeadModel).

Its tree is transformer.{wte, wpe, drop, h[i].{ln_1, attn, ln_2, mlp}, ln_f} plus lm_head. The container-level keys are anchored at the root (transformer.h) so the aliases land on the root envoy: model.layers, not model.model.layers. wpe (learned absolute positions) and drop have no standard name; they stay reachable under their own.

A checkpoint with reorder_and_upcast_attn set takes GPT-2's own _upcast_and_reordered_attn path instead of the shared eager forward, so none of the attention's interior values is available there.