AI Engineer Exhaustive
MediumPlusGPT-2

Scaled Dot-Product Attention

GPT-2

Medium

Core attention computation

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

Scaled dot-product attention is the scoring rule inside every GPT-2 self-attention head. The 2019 paper does not re-derive it: section 2.3 says the architecture “largely follows” the original GPT Transformer, which itself uses the attention of Vaswani et al. (2017). What GPT-2 adds is scale (more layers, wider d_model, 1024-token context) and a few residual/normalization changes around the block, not a new attention formula.

The computation turns queries, keys, and values at one head into a convex combination of value vectors. Compatibility is the inner product (QK^\top). Dividing by (\sqrt{d_k}) keeps the logits from growing with head width, so softmax does not collapse to a one-hot as (d_k) increases. The official inference graph in src/model.py implements exactly this product-then-scale form, then (in the surrounding attn wrapper) a causal mask and a value mix.

This note isolates the core score-and-mix step. Multi-head split/merge, the fused c_attn projection, and the causal mask are neighboring pieces.

How it works

Let one head have queries, keys, and values

Q,K \in \mathbb{R}^{n \times d_k},\qquad V \in \mathbb{R}^{n \times d_v}.

Vaswani et al. define

\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V.

The ((i,j)) entry of (QK^\top) is the compatibility of destination position (i) with source position (j). After scaling and a row-wise softmax, each row of the weight matrix is a distribution over sources; multiplying by (V) writes a weighted sum into that destination.

In the released GPT-2 code the same algebra appears inside multihead_attn in src/model.py. Tensors are already split per head and have shape [batch, heads, sequence, features]. The implementation is:

  1. (W \leftarrow QK^\top) via tf.matmul(q, k, transpose_b=True).
  2. (W \leftarrow W / \sqrt{d}) with tf.rsqrt on the last dimension of (V) (the head width (d_k = d_{\mathrm{model}} / n_{\mathrm{head}})).
  3. Causal masking (see the causal-attention note), then a numerically shifted softmax.
  4. (A \leftarrow WV).

For the 117M setting in Table 2, (d_{\mathrm{model}}=768) and (n_{\mathrm{head}}=12), so (d_k=64). Larger GPT-2 sizes keep the same recipe and only change width, depth, and head count.

The paper does not write the (1/\sqrt{d_k}) factor. That scale, the use of (V)’s channel size as (d_k), and the softmax axis are inferred from Vaswani plus the official graph. The public repository is an inference model: it has no attention dropout.

Softmax in src/model.py subtracts the row max before the exponential. That is a stability trick, not a change of the mathematical function.

Official code

Look at multihead_attn nested inside attn in src/model.py. The product, rsqrt scale, mask hook, softmax, and value mix all live there. Head splitting (split_heads) and the output projection c_proj sit around it. There is no separate attention training script in the public tree.

Watch-outs

Sources