Activation Scaling

What activation scaling actually means

Take one hidden tensor X: [B, s, d] in the stack. Almost every stability question about it reduces to one scalar: the variance of its entries, σ^2 = E[x^2] − (E[x])^2, measured per token across the d channels. Activation scaling keeps that number roughly constant as depth increases, instead of letting it compound.

Compounding is the default. If every layer multiplies the typical magnitude by a factor γ, then after L layers the scale is γ^L. With γ = 1.2 and L = 32 that is ≈ 470×; with γ = 0.8 it is ≈ 1/1200. Nothing is wrong with the architecture in either case — the exponent simply has no mercy. The design problem is to force γ ≈ 1 at every hop, in a way that still holds once the weights have moved far from their initial values.

The variance algebra of one linear layer

Everything starts with a matrix multiply. For y = Wx with W: [d_out, d_in], each output is y_i = Σ_j W_ij · x_j: a sum of d_in products. If the W_ij and x_j are independent and zero-mean, variances add:

Var(y_i) = Σ_j Var(W_ij · x_j)
         = d_in · Var(W) · Var(x)

unit gain  ⇒  Var(W) = 1 / d_in        (LeCun / Xavier, fan-in)

That is the origin of every classical init rule. The d_in factor is why wide layers need small weights: d_in = 4096 with Var(W) = 1/4096 gives std(W) ≈ 0.0156. Put a ReLU after it and roughly half the mass is zeroed, so He init compensates with Var(W) = 2 / d_in. GPT-2’s fixed std = 0.02 is not magic either — it is about 1/√2500, a fan-in rule frozen at one width.

The residual stream is an accumulator

A transformer layer does not replace its input, it adds to it. With Pre-LN the update is x_{l+1} = x_l + F_l(LN(x_l)), and the branch F_l reads a normalized input, so its output variance is roughly independent of how big x_l has become. If each branch contributes variance v and the contributions are approximately uncorrelated, variances simply add:

Var(x_k) ≈ Var(x_0) + k · v      →   std(x_k) ≈ √k

k = number of residual sublayers = 2L (attention + FFN per layer)

So the stream grows like √depth, not exponentially — the residual connection doing its job. But it still grows, and the consequence is subtle: the relative contribution of block k is v / (k·v) = 1/k. Later layers write into a stream they can barely perturb — exactly the observed behaviour that deep Pre-LN blocks contribute progressively less and prune cheaply.

Normalization as a variance reset

Normalization is the device that stops the accumulator from feeding itself. LayerNorm standardizes across the channel axis, then reapplies a learned gain and bias; RMSNorm drops the mean subtraction and keeps only the scale:

LayerNorm:  y = g · (x − μ) / √(σ^2 + ε) + b
RMSNorm:    y = g · x / √( (1/d) Σ_i x_i^2 + ε )

both:  ∂(output scale)/∂(input scale) = 0    — scale-invariant

The important property is not the formula but its homogeneity: multiply x by any positive constant c and the output is unchanged. Whatever scale the residual stream has drifted to, the branch input is back to unit RMS. It also means the gain g is the only thing setting the branch’s input scale — which is why gains are excluded from weight decay, and why a norm whose g quietly grows is a real training signal.

Pre-LN versus Post-LN

The two placements differ by one set of parentheses and behave completely differently:

Post-LN:  x_{l+1} = LN( x_l + F(x_l) )      (original 2017 Transformer)
Pre-LN:   x_{l+1} = x_l + F( LN(x_l) )      (GPT-2 onward)

In Post-LN there is no clean identity path: every backward step passes through a LayerNorm, whose Jacobian divides by the current standard deviation. Those factors multiply across depth, so gradient magnitude at the early layers is a product of L terms — the reason Post-LN models need careful warmup and turn fragile past a few dozen layers.

Pre-LN leaves ∂x_{l+1}/∂x_l = I + ∂F/∂x_l: an unobstructed identity path, so gradients reach layer 0 essentially undamped. That is why nearly every modern decoder uses it. The price is the √k growth above, cleaned up by a single final norm before the output head.

Worked example: variance through a 24-layer stack

Take d = 2048, L = 24 layers, so k = 48 residual sublayers, Pre-LN, embeddings normalized to Var(x_0) = 1, and each branch initialized to emit unit-variance output.

naive init      v = 1        Var(x_48) ≈ 1 + 48·1     = 49    → std ≈ 7.0
                             layer 48 relative weight = 1/49  ≈ 2%

1/√(2L) init   v = 1/48     Var(x_48) ≈ 1 + 48·(1/48) = 2     → std ≈ 1.41
                             (scale the branch output projection by 1/√48 ≈ 0.144)

A 7× drift is survivable in bf16 and cosmetic after the final norm — but not free. It shifts the operating point of every downstream norm, pushes fp16 activations toward the 65504 ceiling in deeper stacks, and widens the per-layer dynamic range a quantizer must cover. The 1/√(2L) branch scale keeps the whole stream in one narrow band — precisely why GPT-2 scales residual output projections by 1/√N.

Scaling the branch: 1/sqrt(2L), ReZero, LayerScale, DeepNorm

Once you accept that the branch needs a gain, four standard answers appear, ordered from fixed to fully learned.

ReZero and LayerScale share one insight: initialize the network as a shallow one. With α = 0 every block is the identity, effective depth at step 0 is one layer, and the model grows its own depth as α lifts off zero. DeepNorm takes the other route: keep Post-LN, but put a constant α > 1 on the identity path with a matching β < 1 down-scaling of the branch init — for a decoder-only stack of N layers, α = (2N)^(1/4) and β = (2N)^(−1/4) — which bounds the per-update change in the output.

Attention logits and the 1/sqrt(d_k) factor

The most famous scaling constant in the architecture is an instance of the same variance identity. A single attention score is a dot product over d_k dimensions:

score = q · k = Σ_{i=1..d_k} q_i k_i
Var(score) = d_k · Var(q) · Var(k) = d_k      (unit-variance q, k)
std(score)  = √d_k  →  d_k = 64 gives std = 8

so:  A = softmax( QK^T / √d_k )

Why it matters is the softmax Jacobian, ∂p_i/∂z_j = p_i(δ_ij − p_j). Unscaled logits with standard deviation 8 produce gaps of 15–20 between the top scores, so p → one-hot, p_i(1 − p_i) → 0, and attention gets essentially no gradient — frozen at initialization onto arbitrary positions. Dividing by √d_k restores unit-variance logits and a well-conditioned softmax. (Maximal-update parameterization argues for 1/d_k, the right exponent when hyperparameters must transfer across widths.)