Activations and initialization from variance propagation

Before this lesson: lesson 01 (the MLP), lesson 03 (the backward pass, $G_l$ and $\sigma'$). From other tracks: the variance of a sum of independent random variables and expectations under a normal distribution (probability-and-statistics track); singular values (linear-algebra lesson 08) for one advanced box. Time: about 65 minutes, plus the notebook.

Why this matters #

• Deep networks train or fail on these choices. Before 2010, deep networks with sigmoid units and badly scaled weights did not train well. ReLU activations and initialization derived from the variance argument (Glorot and Bengio 2010; He et al. 2015) were key steps towards very deep networks.
• The same variance argument appears everywhere. It explains the factor $1/\sqrt{d_k}$ in attention (lesson 09), GPT-2's scaled residual weights, and the role of normalization (lesson 05).
• Modern architectures. GELU is the activation of BERT, GPT-2 and ViT. SwiGLU is the feed-forward layer of PaLM and the LLaMA family, and "why is its width $8d/3$?" is a standard interview question.
• Hyperparameter transfer. μP (maximal update parametrization) extends the same reasoning to the size of the updates, so that a learning rate tuned on a small model works on a large one.

1. Predict first #

2. The picture #

One unit: variances add #

Take one unit with $n$ inputs: $z = \sum_{i=1}^{n} a_i w_i$. Assume the weights $w_i$ are independent, with mean 0 and variance $v$, and independent of the inputs $a_i$. Then each term $a_i w_i$ has mean 0 and variance $v\,\mathbb{E}[a_i^2]$, and the terms are uncorrelated. Variances of uncorrelated terms add:

So each layer multiplies the size of the signal by a factor. With factor 1 the signal keeps its size through any depth. With factor $1/2$ it halves in every layer. With factor 2 it doubles. Over $L$ layers the factor is raised to the power $L$: the change is exponential in depth.

The activation enters this factor through $\mathbb{E}[\sigma(z)^2]$. For ReLU, half of the inputs are set to 0, and $\mathbb{E}[\operatorname{relu}(z)^2] = \tfrac12 \operatorname{Var}(z)$. So with $v = 1/n$ the factor is $n \cdot \tfrac1n \cdot \tfrac12 = \tfrac12$ per layer, and with $v = 2/n$ it is exactly 1. That is He initialization.

The backward pass has the same structure. The gradient goes through $W^\top$ and is multiplied by $\sigma'(z)$, so its factor contains $\mathbb{E}[\sigma'(z)^2]$. A saturating activation has $\sigma'(z) \approx 0$ for large $|z|$, so the gradient can vanish even when the forward signal is fine.

3. Definitions and main results #

3.1 Activation functions #

The activation function $\sigma$ is applied entry by entry after each linear layer. The table uses $s(x) = 1/(1 + e^{-x})$ for the logistic sigmoid, and $\Phi$, $\varphi$ for the CDF and the density of the standard normal distribution.

• ReLU passes gradient 1 for $x > 0$, so it does not saturate for positive inputs. It gives sparse activations (many exact zeros). It is positively homogeneous: $\operatorname{relu}(cx) = c\operatorname{relu}(x)$ for $c > 0$. So the factor $\tfrac12$ below holds at every scale of $z$. For tanh the factor depends on $\operatorname{Var}(z)$.
• Saturation. For large $|x|$ the sigmoid and tanh are flat, so $\sigma'(x) \approx 0$. Through $L$ sigmoid layers the gradient is multiplied by $L$ factors $s'(z_l) \le 1/4$ (times the weights), and $(1/4)^{10} \approx 10^{-6}$. Sigmoid outputs are also not centered at zero. This vanishing gradient problem made early deep networks hard to train; recurrent networks have the same problem over time (lesson 08).
• Dead ReLU. A ReLU unit whose pre-activation is negative for every training input outputs 0 and gets zero gradient for its incoming weights and bias. So these weights stop changing, and the unit usually stays dead (it can come back only if the layers below change its inputs). Causes: a large update (a high learning rate, a loss spike) that pushes its weights or bias far negative, or a bad initialization. Fixes: a lower learning rate or a warmup, a leaky ReLU, GELU or SiLU, or normalization before the activation.
• GELU (Gaussian error linear unit; Hendrycks and Gimpel 2016) multiplies its input by the probability that a standard normal variable is below it: $x \cdot P(X \le x)$. It is a smooth gate. PyTorch also offers a tanh approximation, approximate="tanh".
• SiLU (sigmoid linear unit, also called Swish) is $x\,s(x)$. GELU and SiLU are smooth and not monotone. They have a nonzero gradient for moderately negative inputs and are close to ReLU for large $|x|$.
• Smoothness matters in ML for science. A ReLU network is piecewise linear, so its second derivative is zero almost everywhere. A physics-informed loss that contains $u_{xx}$ then gets no useful signal. Use tanh, SiLU, GELU or sine activations there.

3.2 Gated linear units and SwiGLU #

A gated linear unit (GLU; Dauphin et al. 2017) multiplies a linear branch entry by entry by a gate: $(xW + b) \odot s(xV + c)$. Shazeer (2020) put variants of it into the transformer feed-forward network (FFN), without biases:

Here $d$ is the model width and $d_{ff}$ the hidden width of the FFN.

• Why it may help. The product creates multiplicative (second-order) interactions between features, and the gate can switch features on and off for each input.
• The evidence is empirical. At equal parameters and compute, GEGLU and SwiGLU had the best held-out log-perplexity in Shazeer's T5 experiments (1.942 and 1.944 after 65,536 steps, against 1.997 for the ReLU FFN) and better fine-tuning scores. PaLM and LLaMA adopted SwiGLU. There is no accepted theory for why it works better.
• Parameter matching. A GELU FFN of width $4d$ has $2 \cdot d \cdot 4d = 8d^2$ weights. SwiGLU has three matrices, $3\, d\, d_{ff}$ weights. Equal counts give $3 d\, d_{ff} = 8 d^2$, so $d_{ff} = 8d/3$. The FLOPs match too, because a matmul costs 2 FLOPs per weight per token. Implementations round up: LLaMA takes $\lfloor 8d/3 \rfloor$ and rounds it up to a multiple of 256 (for $d = 4096$: $10922 \to 11008$).

3.3 The variance recursion #

Setting. Use the row convention of this track: layer $l+1$ computes $z_{l+1} = a_l W_{l+1}$, with $W_{l+1} \in \mathbb{R}^{n_l \times n_{l+1}}$ and $a_l = \sigma(z_l)$, where $n_l$ is the width of layer $l$ (the fan-in of $W_{l+1}$; $n_{l+1}$ is its fan-out). Unit $j$ computes $z_{l+1,j} = \sum_{i=1}^{n_l} a_{l,i}\, W_{ij}$. Assume, at initialization:

1. the $W_{ij}$ are independent, with mean 0 and variance $v_{l+1}$;
2. they are independent of $a_l$ (true at initialization, because $a_l$ depends only on earlier layers);
3. the biases are 0;
4. the units of layer $l$ have the same second moment $\mathbb{E}[a_l^2] := \mathbb{E}[a_{l,i}^2]$ for every $i$ (true for $l \ge 1$ because the units are exchangeable; otherwise read it as the average over $i$);
5. $z_l$ is symmetric around 0, that is $z_l$ and $-z_l$ have the same distribution (true when the weights are symmetric). Only the ReLU step below uses this.

Theorem (forward variance). $\operatorname{Var}(z_{l+1,j}) = n_l\, v_{l+1}\, \mathbb{E}[a_l^2]$.

Proof.

1. $\mathbb{E}[z_{l+1,j}] = \sum_i \mathbb{E}[a_{l,i}]\,\mathbb{E}[W_{ij}] = 0$, by independence (2) and $\mathbb{E}[W] = 0$ (1).
2. So, again by (2),

3. For $i \ne k$, $\mathbb{E}[W_{ij} W_{kj}] = \mathbb{E}[W_{ij}]\,\mathbb{E}[W_{kj}] = 0$ by (1), so the cross terms vanish. For $i = k$ it is $v_{l+1}$.
4. Hence $\operatorname{Var}(z_{l+1,j}) = v_{l+1} \sum_i \mathbb{E}[a_{l,i}^2] = n_l\, v_{l+1}\, \mathbb{E}[a_l^2]$ by (4). $\blacksquare$

Note that the formula has $\mathbb{E}[a^2]$, not $\operatorname{Var}(a)$. ReLU outputs have a positive mean, and it counts.

He (Kaiming) initialization, for ReLU. Let $z = z_l$. Since $z$ is symmetric (5), $\mathbb{E}[z^2 \mathbf{1}\{z > 0\}] = \mathbb{E}[z^2 \mathbf{1}\{z < 0\}]$, and the two add up to $\mathbb{E}[z^2]$. So

using $\mathbb{E}[z] = 0$ (step 1 of the proof). Then $\operatorname{Var}(z_{l+1}) = \tfrac12\, n_l\, v_{l+1} \operatorname{Var}(z_l)$, and the variance stays the same if and only if $v_{l+1} = 2/n_l$. Over depth,

Any constant factor $c \ne 1$ per layer gives $c^{L-1}$: exponential in depth.

LeCun initialization, for linear units (or tanh near 0). If $\sigma(z) \approx z$, then $\mathbb{E}[\sigma(z)^2] \approx \operatorname{Var}(z)$, and the condition is $n_l v = 1$: $v = 1/\mathrm{fan}_{\mathrm{in}}$.

Xavier (Glorot) initialization. The forward condition wants $v = 1/\mathrm{fan}_{\mathrm{in}}$ and the backward condition wants $v = 1/\mathrm{fan}_{\mathrm{out}}$ (for linear or tanh units near 0). Both hold only if $\mathrm{fan}_{\mathrm{in}} = \mathrm{fan}_{\mathrm{out}}$. Glorot and Bengio (2010) take the compromise

which is the harmonic mean of the two conditions. (The ReLU analogue would be $4/(\mathrm{fan}_{\mathrm{in}} + \mathrm{fan}_{\mathrm{out}})$.)

One condition is enough. With the forward condition $v_{l+1} = 2/n_l$ (ReLU), the backward factor per layer is $n_{l+1} \cdot \tfrac{2}{n_l} \cdot \tfrac12 = n_{l+1}/n_l$. Over the whole network the product telescopes:

a bounded change, not an exponential one (He et al. 2015, §2.2).

Uniform versions and library defaults. A uniform distribution $U(-a, a)$ has variance $a^2/3$, so $a = \sqrt{3v}$. Xavier uniform: $a = \sqrt{6/(\mathrm{fan}_{\mathrm{in}} + \mathrm{fan}_{\mathrm{out}})}$. He uniform: $a = \sqrt{6/\mathrm{fan}_{\mathrm{in}}}$. For a convolution, the fan-in is $C_{\text{in}} k_h k_w$ (input channels times kernel size). PyTorch's default for nn.Linear and nn.Conv* is $U(\pm 1/\sqrt{\mathrm{fan}_{\mathrm{in}}})$, a variance of $1/(3\,\mathrm{fan}_{\mathrm{in}})$: neither LeCun nor He.

3.4 What goes wrong #

Zero or constant initialization. Suppose all units of a hidden layer start with the same incoming weights and bias and the same outgoing weights. Then they compute the same function, so they get the same gradient: a unit's gradient depends on the unit only through these weights. After the update they are still equal. This is true after every step, so the layer acts like one unit. This is the symmetry problem, and random initialization is how we break the symmetry.

All-zero weights and biases in a ReLU or tanh MLP are even worse. Every hidden activation is $\sigma(0) = 0$. So $\bar W_L = A_{L-1}^\top G_L = 0$, and $G_{L-1} = (G_L W_L^\top) \odot \sigma' = 0$ because $W_L = 0$. Only the output bias gets a gradient, and the network stays in this state. Biases can start at 0, because random weights already make the units different.

Xavier in a deep ReLU network. For square layers Xavier gives $v = 1/n$. The factor is $\tfrac12 \cdot n \cdot \tfrac1n = \tfrac12$ per layer, forward and backward. After 50 layers the variance is about $2^{-49} \approx 1.8 \cdot 10^{-15}$ of its first value.

LeCun with tanh decays slowly. With $q_l = \operatorname{Var}(z_l)$ and $v = 1/n$, the recursion is $q_{l+1} = \mathbb{E}[\tanh(z_l)^2]$ for $z_l \sim N(0, q_l)$. For small $q_l$, $\tanh(z)^2 \approx z^2 - \tfrac23 z^4$, so $q_{l+1} \approx q_l - 2q_l^2$ (using $\mathbb{E}[z^4] = 3q_l^2$). Then $q_l$ falls like $1/(2l)$, not geometrically: after 50 layers it is about $0.01$.

Second moments for $z \sim N(0, 1)$: $\mathbb{E}[\sigma(z)^2] = 0.5$ (ReLU), $0.425$ (GELU), $0.356$ (SiLU), $0.394$ (tanh). A weight variance $v = 1/(n\, \mathbb{E}[\sigma(z)^2])$ keeps unit variance: $2/n$ for ReLU, about $2.35/n$ for GELU.

Answers to section 1. (1) About $10^{-15}$. The first layer gives $\operatorname{Var}(z_1) = 1000 \cdot \tfrac{1}{1000} \cdot 1 = 1$, and then each ReLU layer multiplies by $\tfrac12$, so $\operatorname{Var}(z_{50}) = 2^{-49} \approx 1.8 \cdot 10^{-15}$. (2) $8d/3$.

4. Worked example #

(a) Activations at $x = -1$. We need $\Phi(-1) = 0.1587$, $\varphi(-1) = 0.2420$ and $s(-1) = 0.2689$.

• GELU: value $-1 \cdot \Phi(-1) = -0.1587$. Derivative $\Phi(-1) + (-1)\varphi(-1)$ $= 0.1587 - 0.2420 = -0.0833$. It is negative because GELU is not monotone there.
• SiLU: value $-1 \cdot s(-1) = -0.2689$. Derivative $s(-1)\,(1 + (-1)(1 - s(-1)))$ $= s(-1) \cdot s(-1) = 0.2689^2 = 0.0723$.
• ReLU: value 0 and derivative 0. At $x = -1$ a ReLU unit passes no gradient; GELU and SiLU do.

(b) Initialization of a BERT-base FFN layer, fan-in 768, fan-out 3072.

• LeCun std: $\sqrt{1/768} = 0.0361$. He std: $\sqrt{2/768} = 0.0510$.
• Xavier normal std: $\sqrt{2/(768 + 3072)} = \sqrt{2/3840} = 0.0228$. Xavier uniform bound: $\sqrt{6/3840} = 0.0395$.
• PyTorch default: bound $1/\sqrt{768} = 0.0361$, so std $0.0361/\sqrt{3} = 0.0208$.

(c) Depth. 50 square ReLU layers with Xavier: factor $\tfrac12$ per layer, so $\operatorname{Var}(z_{50})/\operatorname{Var}(z_1) = 2^{-49} \approx 1.8 \cdot 10^{-15}$. The gradients vanish in the same way. With He the factor is 1.

(d) SwiGLU at $d = 768$. The GELU FFN of width $4d = 3072$ has $2 \cdot 768 \cdot 3072 = 8 \cdot 768^2 = 4{,}718{,}592$ weights. SwiGLU with $d_{ff} = 8 \cdot 768/3 = 2048$ has $3 \cdot 768 \cdot 2048 = 4{,}718{,}592$ weights. ✓ This is exactly Shazeer's setup: $d_{ff} = 3072$ for the two-matrix FFNs and $2048$ for the GLU variants.

5. In code #

NumPy and SciPy check the numbers of the worked example and the depth effect. The notebook has these cells and the exercises.

import numpy as np
from scipy.stats import norm

s = lambda x: 1 / (1 + np.exp(-x))
x = -1.0
print(round(x * norm.cdf(x), 4), round(norm.cdf(x) + x * norm.pdf(x), 4))   # GELU: -0.1587 -0.0833
print(round(x * s(x), 4), round(s(x) * (1 + x * (1 - s(x))), 4))              # SiLU: -0.2689 0.0723

rng = np.random.default_rng(0)
z = rng.standard_normal(10**6)
print(round(np.mean(np.maximum(z, 0) ** 2), 2))           # 0.5 = E[relu(z)^2] for z ~ N(0, 1)

# (b) standard deviations for fan-in 768, fan-out 3072
fi, fo = 768, 3072
print([round(float(v), 4) for v in (np.sqrt(1 / fi), np.sqrt(2 / fi), np.sqrt(2 / (fi + fo)), 1 / np.sqrt(3 * fi))])
# [0.0361, 0.051, 0.0228, 0.0208]   LeCun, He, Xavier, PyTorch default

# (c) 50 ReLU layers of width 500: variance of the last pre-activation
for name, v in [("Xavier", 1 / 500), ("He", 2 / 500)]:
    A = rng.standard_normal((200, 500))
    for _ in range(50):
        Z = A @ (rng.standard_normal((500, 500)) * np.sqrt(v))
        A = np.maximum(Z, 0)
    print(name, f"{np.mean(Z**2):.1e}")    # Xavier 1.5e-15 (about 2^-49), He 1.8e+00 (about Var(z_1) = 2)

# (d) SwiGLU width at the parameter count of a 4d FFN
d = 768
print(2 * d * 4 * d, 3 * d * (8 * d // 3))                # 4718592 4718592

In PyTorch, nn.init.calculate_gain("relu") returns $\sqrt 2$, the factor in He's standard deviation $\sqrt 2/\sqrt{\mathrm{fan}_{\mathrm{in}}}$, and a fresh nn.Linear(768, 3072) has weight standard deviation about $0.0208$ (the default, not He). The notebook checks both when PyTorch is installed.

6. Where it shows up #

• Language and vision models. GELU in BERT, GPT-2 and ViT; SwiGLU in PaLM and the LLaMA family.
• Transformer initialization. GPT-2 scales the weights of the residual output projections at initialization by $1/\sqrt{N}$ for $N$ residual layers (nanoGPT uses standard deviation $0.02/\sqrt{2 n_{\text{layer}}}$ for them). So the variance of the residual stream does not grow with depth (lesson 05). DiT's adaLN-Zero starts each block as the identity map.
• Attention. The scores $q^\top k$ are sums of $d_k$ products, so their variance is $d_k$ when $q$ and $k$ have unit-variance entries. Dividing by $\sqrt{d_k}$ brings it back to 1 (lesson 09).
• ML for science. Losses of physics-informed networks need second derivatives, hence tanh, SiLU or sine activations. SIREN (a network with sine activations) draws its weights from a uniform distribution scaled with $\sqrt{6/n}$, so that the input of each sine is close to $N(0, 1)$: the same argument with a different $\sigma$.
• Porting code between frameworks. Defaults differ: PyTorch's nn.Linear uses $U(\pm 1/\sqrt{\mathrm{fan}_{\mathrm{in}}})$, Flax's nn.Dense uses LeCun normal. A model moved from one to the other starts from a different scale unless you set the initializer explicitly.

7. Common mistakes #

• $\operatorname{Var}(a)$ instead of $\mathbb{E}[a^2]$ in the recursion. For ReLU the mean is not zero.
• Standard deviation vs variance. He std is $\sqrt{2/\mathrm{fan}_{\mathrm{in}}}$; $2/\mathrm{fan}_{\mathrm{in}}$ is the variance. torch.randn(...) * std needs the standard deviation.
• Xavier for deep ReLU networks. Xavier is derived for linear or tanh units; with ReLU it halves the variance in every layer.
• Assuming PyTorch's default is He. nn.Linear uses $U(\pm 1/\sqrt{\mathrm{fan}_{\mathrm{in}}})$, variance $1/(3\,\mathrm{fan}_{\mathrm{in}})$.
• Conv fan-in $= C_{\text{in}}$. It is $C_{\text{in}} k_h k_w$.
• "A good initialization is enough." The size of the updates also depends on the width: that is the point of μP.
• Comparing SwiGLU and GELU FFNs at the same $d_{ff}$. SwiGLU then has 1.5 times more parameters.

8. Check yourself #

1. Why does a 10-layer sigmoid network train badly even with good weight scales?
   Answer

   Each layer multiplies the backpropagated gradient by $s'(z) \le 1/4$, so the factor from the activations alone is at most $(1/4)^{10} \approx 10^{-6}$. Saturated units give $s' \approx 0$, which is even worse.
2. Which weight variance keeps unit pre-activation variance for GELU when $\operatorname{Var}(z) = 1$?
   Answer

   $v = 1/(n\, \mathbb{E}[\operatorname{GELU}(z)^2]) = 1/(0.425\, n) \approx 2.35/n$.
3. Why does $\mathbb{E}[\operatorname{relu}(z)^2] = \tfrac12\operatorname{Var}(z)$ need $z$ to be symmetric?
   Answer

   The proof splits $\mathbb{E}[z^2]$ into the parts with $z > 0$ and $z < 0$ and uses that they are equal. If $z$ has a positive mean or is skewed (for example with a large positive bias), more than half of the second moment lies on the positive side, and the factor is larger than $\tfrac12$.
4. All weights of an MLP with a softmax cross-entropy output are set to 0.5. At step 0 the first-layer gradient is exactly zero. Why?
   Answer

   All entries of $W_L$ are 0.5, so the entries of $\bar A_{L-1} = G_L W_L^\top$ are $0.5 \sum_c G_{bc} = 0.5 \sum_c (p_c - y_c)/B = 0$, because both $p$ and $y$ sum to 1. Even after the output layer moves, the hidden units stay identical (symmetry).
5. He initialization (forward condition) is used in a network whose widths shrink from $n_1 = 4096$ to $n_L = 256$. How does the gradient variance change over the whole network?
   Answer

   The backward factor per layer is $n_{l+1}/n_l$, and the product telescopes to $n_L/n_1 = 1/16$. A bounded change, not an exponential one.
6. You double the width of a hidden layer trained with Adam. In μP, how do its initialization variance and its learning rate change, and why?
   Answer

   The initialization variance $1/\mathrm{fan}_{\mathrm{in}}$ halves, and the learning rate also halves ($\propto 1/\mathrm{fan}_{\mathrm{in}}$), because the change $a\,\Delta W$ sums twice as many coherent terms.

9. Interview angle #

• "Why ReLU rather than sigmoid? What is GELU? What is a dead ReLU and how do you fix it?" Sigmoid saturates and its derivative is at most $1/4$; ReLU passes gradient 1 for positive inputs. GELU is $x\Phi(x)$, a smooth version. A dead ReLU is negative for all inputs; fix it with a lower learning rate, warmup, leaky ReLU, GELU or normalization.
• "Derive He initialization." Variances add: $\operatorname{Var}(z_{l+1}) = n_l\, v\, \mathbb{E}[a^2]$, and $\mathbb{E}[\operatorname{relu}(z)^2] = \tfrac12\operatorname{Var}(z)$, so $v = 2/n_l$.
• "Why not initialize all weights to zero?" Symmetry: equal units stay equal. With zero weights the hidden activations and the hidden gradients are exactly zero.
• "Why is LLaMA-7B's FFN width 11008?" $8d/3$ for $d = 4096$ is $10922.7$, rounded up to a multiple of 256.
• "How would you tune a model you can train only once?" μP and μTransfer from a narrow proxy model.

10. Practice #

The notebook for this lesson runs in Colab. Its exercises use NumPy; PyTorch is used as a reference when it is installed.

1. Write ReLU, sigmoid, tanh, GELU and SiLU with their derivatives, and check the derivatives with central differences.
2. Estimate $\mathbb{E}[\operatorname{relu}(z)^2]$ and $\mathbb{E}[\operatorname{relu}'(z)^2]$ by Monte Carlo and compare with $\tfrac12$.
3. Write init_std(name, fan_in, fan_out) and preactivation_variances(...), and plot the variance through 50 layers for LeCun, Xavier and He with ReLU and tanh.
4. Train a small network from a constant initialization and show that its hidden units stay identical; then break the symmetry.
5. Write a SwiGLU FFN in NumPy with the parameter count of a GELU FFN of width $4d$.

11. Watch and read #

• Simon Prince, Understanding Deep Learning (free PDF): chapter 3 (activation functions) and chapter 7, §7.5 (parameter initialization, forward and backward).
• Zhang, Lipton, Li and Smola, Dive into Deep Learning, §5.4 Numerical stability and initialization.
• Glorot and Bengio, Understanding the difficulty of training deep feedforward neural networks (2010).
• He, Zhang, Ren and Sun, Delving deep into rectifiers (2015): He initialization, §2.2.
• Hendrycks and Gimpel, Gaussian error linear units (GELUs) (2016).
• Shazeer, GLU variants improve Transformer (2020), §2–3.
• Yang et al., Tensor Programs V: Tuning large neural networks via zero-shot hyperparameter transfer (2022), §2–4 and Table 3.