Depth-aware init for residual projections

Background

GPT-2 initialises most weights with std 0.02, but the projections that write into the residual stream (attn.proj, mlp.proj) get a smaller std scaled by depth:

$$\text{std} = \frac{0.02}{\sqrt{2L}}$$

where L = n_layer. Each block adds two contributions to the residual stream; scaling these down keeps the stream's variance bounded at initialisation instead of growing with depth.

Problem statement

Implement residual_init_std(n_layer, base=0.02) returning the scaled std.

Input

• n_layer — int, number of transformer blocks L.
• base — base std (default 0.02).

Output

Returns a float: base / sqrt(2 * n_layer).

Examples

Example 1 — GPT-2 small

Input:  n_layer = 12
Output: 0.004082...

(0.02 / sqrt(24).)

Constraints

• Return base / sqrt(2 * n_layer).
• Deeper models → smaller std.

Notes

• The 2 counts the two residual-writing sublayers per block (attention + MLP); the √ keeps the variance contribution per layer ∝ 1/L.