AI Engineer Exhaustive
HardPlusGPT-2

Multi-Head Attention

GPT-2

Hard

Multi-head causal self-attention

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

GPT-2’s only attention is masked multi-head self-attention inside each decoder block. The paper does not re-derive Vaswani et al. (2017); section 2.3 says the stack follows GPT, with pre-norm around the sub-block. The official attn function is the concrete contract: one fused projection to queries, keys, and values, (h) heads, scaled dot-product, a causal mask that still works when a KV cache makes keys longer than queries, softmax, a value mix, and an output projection.

Default hparams use n_embd=768 and n_head=12, so each head has width (d_h=64). Larger Table 2 models keep the same pattern with more heads and a wider (d). Attention is always unidirectional: token (t) may read positions (\le t) (plus any cached past), never the future. That is what makes a single teacher-forced forward pass implement equation (1).

How it works

Input (x) has shape [B, T, d] and has already been layer-normalized (ln_1). A conv1d named c_attn maps (d\to 3d):

[Q,K,V]=\mathrm{split}(x W_{qkv}+b_{qkv},\;3)\in\mathbb{R}^{B\times T\times d}.

split_heads reshapes the last axis into [n_head, d_h] and transposes to [B, h, T, d_h]. Per head,

W=\frac{Q K^{\top}}{\sqrt{d_h}}\in\mathbb{R}^{B\times h\times T_q\times T_k}.

The official scale is tf.rsqrt of v.shape[-1], i.e. (1/\sqrt{d_h}), applied to the score tensor (not to (Q) or (K) beforehand).

The causal mask is attention_mask(nd, ns). With query length (n_d) and key length (n_s),

m_{ij}=\mathbf{1}\!\left[i \ge j-n_s+n_d\right].

When (n_d=n_s=T) this is the usual lower triangle: query (i) sees keys (0\ldots i). When a cache is present, (n_s=T_{\mathrm{past}}+T) and (n_d=T) (often (T=1) at decode), and the inequality still allows each new query to see the entire past plus its own prefix. Masked scores become (W\odot m - 10^{10}(1-m)) so softmax treats them as (-\infty) in float32.

A=\mathrm{softmax}(W)\,V\in\mathbb{R}^{B\times h\times T_q\times d_h}.

merge_heads undoes the split ([B, T, d]), and c_proj maps (d\to d). That tensor is added to the pre-norm residual in block. Softmax is the numerically stable version in softmax: subtract the per-row max, then exp/sum.

KV cache: present = stack([K, V], axis=1) has shape [B, 2, h, T, d_h] before concat. If past is not None, prior keys/values are concatenated on axis (-2) (sequence) and attention uses the longer pair. model() stacks layer presents to [B, L, 2, h, T, d_h]. This is how sample_sequence avoids quadratic re-encode of the prefix.

There is no relative position, no ALiBi, and no separate bias on the scores beyond the learned (QKV) and the causal mask. Absolute positions were already added at the embeddings. Heads do not share projections. Dropout is absent in the released inference graph.

Official code

src/model.py in openai/gpt-2: attn, attention_mask, split_states / merge_states, softmax, and the c_attn / c_proj conv1ds. past_shape documents the cache layout. The decode loop that reuses present is sample_sequence in src/sample.py.

Watch-outs

Sources