AI Engineer Exhaustive
HardPlusGPT-2

Decoder Block

GPT-2

Hard

Full pre-norm transformer decoder block

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

A GPT-2 stack is a chain of identical decoder-only Transformer blocks. There is no encoder and no cross-attention: each block is masked self-attention plus a position-wise MLP, with residual connections. Section 2.3 lists the GPT-2-specific layout changes relative to GPT / Vaswani:

  1. LayerNorm moves to the input of each sub-block (pre-norm), “similar to a pre-activation residual network” (He et al., 2016).
  2. One extra LayerNorm is applied after the last block (ln_f in code), not inside the block itself.
  3. Residual-path weights are scaled at init by (1/\sqrt{N}) ((N) = number of residual layers). That init is a sibling note; the forward block only consumes already-scaled weights.

Pre-norm is the load-bearing change. Post-norm (LayerNorm after the residual add) is what the original Transformer used; GPT-2 normalizes, then transforms, then adds back to the unnormalized residual stream. That keeps the residual path a clean identity and is what the official block function implements.

How it works

Let (x\in\mathbb{R}^{B\times T\times d}) with (d=) n_embd (768 for the 117M model in Table 2). One block, matching block() in src/model.py, is

\begin{aligned} a &= \mathrm{Attn}\!\left(\mathrm{LN}_1(x)\right),\\ x &\leftarrow x + a,\\ m &= \mathrm{MLP}\!\left(\mathrm{LN}_2(x)\right),\\ x &\leftarrow x + m. \end{aligned}

(\mathrm{LN}_1) and (\mathrm{LN}_2) are independent LayerNorms (ln_1, ln_2): mean/variance over the feature axis, then learned gain (g) and bias (b), (\varepsilon=10^{-5}). They do not share parameters.

(\mathrm{Attn}) is the causal multi-head module: fused projection to (Q,K,V), scaled dot-product with the causal mask, output projection. It also returns present, the new (K,V) for that layer, used as a cache. (\mathrm{MLP}) is the 4(d)-wide GELU feed-forward (mlp(..., nx*4)). Both submodules see normalized activations; both residual adds use the pre-norm (x).

Shapes stay (B\times T\times d) across the block. The MLP expands to (4d) internally and projects back. Head count is n_head; head width is (d/n_{\mathrm{head}}).

A full model is (L) such blocks (h0 … h{L-1}) then ln_f. Table 2: (L\in{12,24,36,48}) for the four published sizes. The extra final norm is outside this block. Putting a third LayerNorm inside block would double-normalize the last layer relative to the released graph.

The paper specifies pre-norm and the final extra LN, and cites GPT for the rest (decoder-only, GELU MLP, learned positions). It does not write the two-add equations above; those are read off block. Dropout is absent from the public inference model.

Init scaling ((1/\sqrt{N}) on residual outgoing weights) is not part of the forward expression. If you only re-implement block you can ignore it; if you train from scratch you cannot.

Official code

src/model.py: block is the whole decoder block; it calls norm, attn, and mlp. model loops block(h, 'h%d' % layer, ...) and then norm(h, 'ln_f'). Layer scopes ln_1, attn, ln_2, mlp match checkpoint variable names. There is no separate block file.

Watch-outs

Sources