AI Engineer Exhaustive
MediumPlusGPT-2

Residual Weight Scaling

GPT-2

Medium

Initialization scaling for deep residual networks

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

Section 2.3 of the GPT-2 paper lists a “modified initialization which accounts for the accumulation on the residual path with model depth.” Residual branches are added (N) times along the depth of the network. If every branch is drawn from the same (\mathcal{N}(0,0.02^{2})) used for ordinary GPT linears, the residual stream’s variance grows with depth and early training becomes unstable, especially at 36–48 layers.

The fix is a one-time scale at init, not a learned gate and not a runtime multiply:

W_{\mathrm{resid}}\;\leftarrow\;\frac{W_{\mathrm{resid}}}{\sqrt{N}},

where (N) is “the number of residual layers.” Combined with pre-norm (LN at the input of each sub-block, analogous to a pre-activation ResNet), this is the paper’s recipe for training the 1.5B, 48-layer model.

The official openai/gpt-2 tree is an inference release. conv1d always constructs w with tf.random_normal_initializer(stddev=w_init_stdev) and default w_init_stdev=0.02. There is no /√N in the forward graph. Checkpoints already contain the trained values; the scale only matters if you reimplement training.

How it works

Each decoder block adds two branches to the stream (H):

H\leftarrow H+\mathrm{Attn}(\mathrm{LN}_1(H)),\qquad H\leftarrow H+\mathrm{MLP}(\mathrm{LN}_2(H)).

If every output projection is initialized at the same scale, those additions are not variance-preserving as (L) grows. A standard heuristic (used in the original GPT writeup and restated here) is to shrink the residual weights so that the sum of (N) roughly independent contributions stays (O(1)).

What counts as a residual layer is slightly underspecified. The paper says (N) is the number of residual layers, not the number of Transformer blocks. The natural reading, and the one later reimplementations adopted, is one residual add per attention projection and one per MLP projection:

N=2L,\qquad \sigma_{\mathrm{resid}}=\frac{0.02}{\sqrt{2L}}.

For the four Table 2 depths that is (\sigma\in{0.02/\sqrt{24},;0.02/\sqrt{48},;0.02/\sqrt{72},;0.02/\sqrt{96}}). Only the matrices that write back into (H) should receive this scale — in official names, attn/c_proj and mlp/c_proj. The fused QKV map (c_attn), the MLP expand (c_fc), token embeddings (wte, std (0.02)), and position embeddings (wpe, std (0.01)) stay on the unscaled initializer. Layer-norm g and b start at 1 and 0.

If someone instead sets (N=L) (one scale per block), residual writes are (\sqrt{2}) too large relative to the (2L) reading. The paper does not give a numeric example, so this is inferred from the wording plus the two-adds-per-block structure in block.

The scale is applied to the initial tensor, not divided out at every forward pass. After training, c_proj weights are just parameters; you do not keep dividing by (\sqrt{N}) at inference. Mixing a correctly scaled init with a later runtime /√N would under-scale the branches twice.

Depth and this init are coupled with the other 2.3 changes: pre-norm keeps the residual path as a clean identity, and the extra ln_f after the last block re-centers the stream before the tied LM head. Together they replace the post-norm GPT-1 block as models go from 12 to 48 layers.

Official code

Look at conv1d and block in src/model.py. Residual writes are the c_proj calls inside attn and mlp. The /√N factor is described only in the paper; it is not a runtime op in this repository. Embedding inits are the wte / wpe get_variable lines in model().

Watch-outs

Sources