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.)

The output head: logit scale and z-loss

The last projection is where scaling errors surface as loss spikes. The head maps [d] → [V], and when embeddings are tied the same matrix was trained as a lookup table, not as a variance-preserving projection. A residual stream that has drifted to std ≈ 7 hits that matmul and produces logits with roughly seven times the intended spread.

Large logits are dangerous twice over. Numerically, exp() overflows fp16 above about 11.09, so softmax’s max-subtraction is mandatory, not an optimization. Statistically, an over-confident softmax has a near-singular Jacobian and stops learning. The standard remedy is an auxiliary z-loss, λ · (log Z)^2 with Z = Σ_v exp(z_v) and λ ≈ 1e-4, which penalizes the log-partition function directly and keeps logits centered without touching the architecture.

Outliers, quantization, and CPU small models

For CPU inference the variance story stops being about training stability and becomes about dynamic range. Int8 quantization maps a tensor with one scale, s = max|x| / 127, so resolution is set entirely by the largest element. Transformers make this hard: a handful of residual-stream channels carry systematically huge magnitudes — emergent outlier features, tens to hundreds of times the median channel.

The arithmetic is brutal. If one channel hits 100 while typical activations sit near 1, then s ≈ 0.79 and every ordinary value collapses into one or two quantization levels. Hence the standard fixes: per-channel rather than per-tensor activation scales, mixed-precision decomposition that keeps outlier channels in fp16, and SmoothQuant’s migration of difficulty into the weights via Y = (X / s) · (s · W) — mathematically identical, but it moves the wide range into a tensor that quantizes well. For a CPU SLM that one transformation is often the difference between usable int8 and garbage.

Pitfalls that quietly break the scaling

Four failure modes account for most real incidents. Fixed init std at the wrong width: copying 0.02 into a d = 8192 model gives roughly 1.8× the variance-preserving value at every layer, compounding through depth. Forgetting the residual down-scale: omitting the 1/√(2L) factor is harmless at 6 layers and destabilizing at 60.

Weight-decaying the norm gains: g and b are scale controls, not capacity, and decaying them drags the operating point around. And changing the attention scale after training — a temperature tweak, or a kernel that folds 1/√d_k differently — silently re-tunes the sharpness of every head at once.