Weight Initialization — Xavier, He, and μP

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 22 / 80 · Module M04 · ~3 hours

0 · Initialization sets the starting dynamics

Initialization chooses the scale and symmetry of a network before the first update. Small random weights can still erase signal; large weights can saturate activations or amplify products of Jacobians. A useful initializer follows the layer's shape and nonlinearity, not the slogan 'random is enough'.

Glorot and Bengio's 2010 analysis and He et al.'s 2015 rectifier analysis give concrete starting points. Neither proves that pretraining disappeared overnight or that every deeper network became trainable from one line. μP is a coordinated parameterization and optimizer-scaling framework; its cited transfer paper is from 2022, not a newly invented 2024 initializer. Private frontier-model adoption is not assumed.

Prerequisites

Understand variance/second moment, matrix shapes, activations and normalization. The CPU programs use generated matrices, not downloaded model weights.

1 · The variance-preservation goal

Let $y=Wx$, with W shaped (fan_out,fan_in). Assume weights are independent of x, independent across columns, zero-mean and have variance $\sigma_w^2$. For one output coordinate,

when input coordinates have a common second moment and cross terms vanish under the stated weight assumptions. Since expected y is zero over the random weights, its variance equals this second moment. If inputs also have zero mean, $\mathbb E[x_j^2]=\operatorname{Var}(x_j)$. After ReLU, inputs generally do not have zero mean; retain the second moment in the recurrence.

For a linear layer with zero-mean input, forward preservation asks $n_{in}\sigma_w^2=1$. Thus weight standard deviation is $1/\sqrt{n_{in}}$. Variance addition under independence/zero cross-covariance is enough; the central limit theorem is not needed for this identity. Mean zero alone does not imply a symmetric distribution.

1.1 The backward pass and incompatible fans

The linear pullback is $g_x=W^Tg_y$. Under analogous independence assumptions its coordinate second moment scales by $n_{out}\sigma_w^2$. Forward and backward factors cannot both be one unless the fans agree. Actual backpropagated gradients need not remain independent of weights, so this is an initialization approximation, not a guarantee after training.

2 · Xavier / Glorot initialization

For gain a,

A symmetric uniform distribution with bound $a\sqrt{6/(n_{in}+n_{out})}$ has that variance because $\operatorname{Var}(U[-b,b])=b^2/3$. It is the reciprocal of the arithmetic mean of fans, equivalently the harmonic mean of the two desired reciprocal fan variances. It is not the harmonic mean of fan_in and fan_out itself.

For fan_in=64 and fan_out=256 with gain one, weight variance is 1/160, forward factor .4 and backward factor 1.6. This compromise preserves neither exactly. nn.init.xavier_normal_ and xavier_uniform_ default to gain one; PyTorch's recommended tanh gain is 5/3. A sigmoid near zero has derivative 1/4 and nonzero output mean, so 'sigmoid is approximately identity' is false. Xavier is a useful family of recipes, not a theorem of flat variance for arbitrary tanh depth.

3 · He / Kaiming — preserve a rectifier second moment

For a symmetric zero-mean preactivation Z,

This halves the second moment, not generally the variance. For Gaussian $Z\sim N(0,\sigma^2)$,

Combining the linear second-moment recurrence and rectifier factor gives $\operatorname{Var}(W)=2/n_{in}$, hence std $\sqrt{2/n_{in}}$, not twice the old standard deviation. For a leaky ReLU with negative slope a, the second-moment factor is $(1+a^2)/2$ and the corresponding gain is $\sqrt{2/(1+a^2)}$.

kaiming_normal_(weight, mode="fan_in", nonlinearity="relu") targets the forward recurrence; fan_out targets the backward approximation. In PyTorch a Linear weight is (out,in) and used as x@W.T. If a custom tensor is stored (in,out) for x@W, pass the appropriately transposed tensor to fan-based initialization. For grouped convolutions, actual fan-in is (in_channels/groups)*kernel_h*kernel_w; blindly using all input channels is wrong.

3.1 Defaults and smooth activations are not all He-ReLU

nn.Linear.reset_parameters calls Kaiming uniform with a=sqrt(5), which yields the historical bound $1/\sqrt{fan_{in}}$ and variance $1/(3fan_{in})$. This is not ReLU He variance $2/fan_{in}$. Its bias is also initialized uniformly, not always zero. Match the chosen recipe explicitly.

GELU and SiLU have input-distribution-dependent moment factors. They are not exactly half-variance gates and there is no universal '1.7 GELU correction' absorbed harmlessly by LayerNorm. PyTorch calculate_gain has documented supported nonlinearities; it does not derive an arbitrary GELU/SiLU gain. Compare empirical or integrated second moments and actual backward behavior.

import math
import torch
from torch import nn
from scipy.integrate import quad
from scipy.special import ndtr, expit
 
phi=lambda z: math.exp(-z*z/2)/math.sqrt(2*math.pi)
relu_second=quad(lambda z:z*z*phi(z),0,math.inf)[0]
relu_mean=quad(lambda z:z*phi(z),0,math.inf)[0]
assert abs(relu_second-.5) < 1e-10
assert abs((relu_second-relu_mean**2)-(.5-1/(2*math.pi))) < 1e-10
for a in [0.,.5,1.]:
    assert abs(nn.init.calculate_gain("leaky_relu",a)**2-2/(1+a*a)) < 1e-12
moments={}
for name,f in [("gelu",lambda z:z*ndtr(z)),("silu",lambda z:z*expit(z))]:
    second=quad(lambda z:f(z)**2*phi(z),-12,12)[0]
    moments[name]=(second,1/second)
    assert 0 < second < .5
assert abs(moments["gelu"][0]-.42522148257) < 1e-9
# Default Linear's range/variance differs from He-ReLU.
torch.manual_seed(22);layer=nn.Linear(128,2048,bias=False)
assert layer.weight.abs().max() <= 1/math.sqrt(128)
assert abs(layer.weight.var(unbiased=False).item()*128-1/3) < .01
print("moment integration and Linear default", moments)

4 · Activation tracking — measure instead of copying invented output

The existing forward-inspection function is retained exactly. Its use_relu=False path is linear, not tanh. Its global activation standard deviation alone does not measure gradients, and a single random sample does not guarantee a smooth curve.

import torch, torch.nn as nn
 
def check_init(init_fn, n_layers=10, dim=512, use_relu=True):
    x = torch.randn(1, dim)
    print(f"layer -1  std {x.std().item():.4f}  mean {x.mean().item():.4f}")
    for i in range(n_layers):
        layer = nn.Linear(dim, dim, bias=False)
        init_fn(layer.weight)
        x = layer(x)
        if use_relu:
            x = torch.relu(x)
        print(f"layer {i:2d}  std {x.std().item():.4f}  mean {x.mean().item():.4f}")

# These calls execute the retained function without fabricated output comments.
torch.manual_seed(22)
check_init(lambda w: nn.init.normal_(w,std=1.),n_layers=6,dim=64)
torch.manual_seed(22)
check_init(lambda w: nn.init.normal_(w,std=.01),n_layers=6,dim=64)
torch.manual_seed(22)
check_init(nn.init.xavier_normal_,n_layers=6,dim=64)
torch.manual_seed(22)
check_init(lambda w:nn.init.kaiming_normal_(w,nonlinearity="relu"),n_layers=6,dim=64)
torch.manual_seed(22)
check_init(nn.init.xavier_normal_,n_layers=6,dim=64,use_relu=False)

For a width-512 Linear+ReLU recurrence, the second-moment factor is approximately $k=512\sigma_w^2/2$: std 1 gives k=256, std .01 gives k=.0256, square Xavier gain one gives k=.5, and He gives k=1. The old .003 factor was incorrect. After L layers it is $k^L$ under the model assumptions. $2^{100}\approx1.27e30$ is large but not automatically FP32 infinity; $2^{-100}\approx7.89e-31$ is not automatically zero. Squaring or further operations may overflow earlier than the values themselves.

Tanh outputs are bounded in magnitude by one, so 'He+tanh activations slowly explode without bound' cannot describe that tensor. They can saturate and have poorly scaled gradients. Xavier+tanh is not guaranteed flat at standard deviation one because tanh is only approximately linear near zero. Normalization can stabilize one measured scale without guaranteeing useful gradients or perfect conditioning.

5 · Biases, symmetry and legitimate zeros

Zero bias is a common choice, not an immutable law. For binary prior $\pi$, output bias $\log(\pi/(1-\pi))$ gives sigmoid output $\pi$ when the feature contribution is zero. For softmax class priors $\pi_i>0$, biases $\log\pi_i$ reproduce the normalized priors. This can improve early behavior on imbalance, but no 3–10× convergence guarantee follows.

For an LSTM, forget bias one gives $\operatorname{sigmoid}(1)\approx.731$, not perfect remembering. PyTorch has separate input-hidden and hidden-hidden biases in gate order input/forget/cell/output; setting both forget slices to one gives total two. Its default reset is not universally forget-bias-one.

All identical hidden neurons with identical outgoing roles can remain symmetric under deterministic updates. Entire all-zero deep ReLU stacks can block gradients. But zero initialization is valid for many linear/logistic models, for some output heads, residual gates and residual final projections. A zero bilinear parameter can have a nonzero derivative with respect to that parameter; 'zero times anything means zero gradient' is not a general rule.

import torch
from torch import nn
 
prior=torch.tensor([.1,.3,.6],dtype=torch.float64)
assert torch.allclose(prior.log().softmax(0),prior)
pi=torch.tensor(.01,dtype=torch.float64)
assert torch.allclose(torch.logit(pi).sigmoid(),pi)
lstm=nn.LSTM(3,4,batch_first=True)
with torch.no_grad():
    lstm.bias_ih_l0.zero_();lstm.bias_hh_l0.zero_()
    lstm.bias_ih_l0[4:8].fill_(1.)
assert torch.equal(lstm.bias_ih_l0[4:8]+lstm.bias_hh_l0[4:8],torch.ones(4))
assert abs(torch.sigmoid(torch.tensor(1.)).item()-.7310586) < 1e-7
linear=nn.Linear(2,1,bias=False)
nn.init.zeros_(linear.weight)
linear(torch.tensor([[2.,3.]])).sum().backward()
assert torch.equal(linear.weight.grad,torch.tensor([[2.,3.]]))
# ReZero: identity output, zero first branch-weight gradient, nonzero gate gradient.
w=nn.Parameter(torch.eye(2));alpha=nn.Parameter(torch.tensor(0.))
x=torch.tensor([1.,2.]);y=x+alpha*(w@x)
y.sum().backward()
assert torch.equal(y,x) and not w.grad.any() and alpha.grad.item()==3.
print("priors, LSTM total bias, zero-linear and ReZero controls passed")

6 · Initialization and normalization cover different problems

Ignoring epsilon and affine parameters, LN(cx)=LN(x) for positive c. Uniformly scaling an embedding does not change its relative coordinates, so the old claim that LN necessarily preserves a huge uniform std difference was false. With epsilon, finite precision, residual addition, learned affine parameters or nonuniform coordinate changes, scales matter again.

Three situations survive normalization: an unnormalized output/readout sets logit scale; a residual branch's magnitude relative to the skip changes the function; singular/degenerate directions and nonlinear saturation can affect Jacobians before/around the norm. Norms do not recreate erased feature distinctions or enforce all Jacobian singular values near one. Inspect both activation second moments and gradient/update statistics.

7 · Special-case initializations

7.1 Fixup, T-Fixup and ReZero are distinct recipes

Fixup uses carefully scaled residual branches, including zero final branch layers and prescribed multipliers/biases. For depth L and m linear layers per residual branch, the other branch weights use a factor $L^{-1/(2m-2)}$ for the studied architecture. It is not simply 'delete norms and zero everything'. The paper's 10,000-layer result belongs to its experiments.

T-Fixup is Huang, Pérez, Ba and Volkovs (ICML 2020), Improving Transformer Optimization Through Better Initialization, not the Xiong LayerNorm paper previously linked. Its recipe uses Xavier for non-embedding parameters, specified embedding initialization, decoder-side selected matrices/embeddings scaled by $(9N)^{-1/4}$ and encoder-side selected matrices by $.67N^{-1/4}$; it removes layer normalization and warmup in the studied encoder-decoder setup. Calling it a generic $1/\sqrt L$ post-norm recipe was incorrect. The exact selected value/output and MLP matrices matter.

ReZero uses $y=x+\alpha F(x)$ with a learned scalar gate initially zero. At that point the output/Jacobian have a clean identity path; branch-parameter gradients are initially multiplied by alpha, while the gate can have a nonzero gradient, as tested above. It does not permanently prevent branch learning.

7.2 GPT residual projection scaling

A GPT-style recipe initializes ordinary projections with a small standard deviation and scales residual-writing projections according to the number of branches/depth, often $1/\sqrt{2N}$ for two branches per block. The heuristic comes from adding approximately independent branch contributions: their variances add, so reducing branch std by the square root of branch count limits growth. Correlations and trained dynamics weaken that simple calculation. The GPT-2 paper describes residual-depth scaling; a precise 2N implementation convention should not be attributed to every later model or to private GPT-4 details.

7.3 Embedding magnitude

For D independent entries with variance $.02^2$, expected squared norm is $D(.02)^2$; root-mean-square norm is $.02\sqrt D$, about .5543 at D=768 and 1.8102 at D=8192. This is not exactly the expectation of the norm and .02 is not a universal embedding initializer. Match tied readout scale, positional representations and normalization to the actual architecture.

7.4 μP — coordinate initialization, parameter roles and updates

μP aims to keep feature-update scales useful as selected width dimensions grow. Classify a parameter by which dimensions scale: fixed input/output dimensions differ from hidden-to-hidden matrices. Coordinate initialization, forward multipliers and optimizer group learning rates together. Multiplying only the output weight std by an extra factor is not a general μP implementation, and SGD and Adam use different scaling rules.

The Microsoft implementation uses MuReadout, set_base_shapes with base and delta models, and MuAdam/MuSGD. Its MuReadout convention applies a reciprocal width-multiplier in forward with constant-variance stored readout weights; effective weights thus scale differently from the stored tensor. MuAdam divides LR by width multiplier for matrix-like parameters with two growing dimensions, leaving vector-like/input/output roles at their appropriate base convention. MuSGD has different role rules. Attention scaling may change from $1/\sqrt d$ to a constant times $1/d$ in the matching recipe; the constant preserves reference-width behavior. This is not permission to alter one attention denominator independently.

The next bounded CPU proxy explicitly implements the bias-free three-linear-layer Adam case of those role rules, with zero weight decay and fixed input/output dimensions. It records coordinate magnitudes and losses at widths 64,256,1024 and includes wrong hidden-LR/readout controls. It does not claim identical finite-width loss curves, formal μP verification or a trained language-model reproduction.

import math
import torch
from torch import nn
import torch.nn.functional as F
 
class WidthProxy(nn.Module):
    def __init__(self,width,base=64,wrong_readout=False):
        super().__init__();self.mult=width/base;self.wrong=wrong_readout
        self.input=nn.Linear(4,width,bias=False)
        self.hidden=nn.Linear(width,width,bias=False)
        self.readout=nn.Linear(width,2,bias=False)
        nn.init.normal_(self.input.weight,std=1/math.sqrt(4))
        nn.init.normal_(self.hidden.weight,std=1/math.sqrt(width))
        nn.init.normal_(self.readout.weight,std=1/math.sqrt(base))
    def forward(self,x):
        a=torch.relu(self.input(x));b=torch.relu(self.hidden(a))
        out=self.readout(b/(1 if self.wrong else self.mult))
        return out,(a.abs().mean().item(),b.abs().mean().item(),out.abs().mean().item())
 
torch.manual_seed(7);x=torch.randn(16,4);y=(x[:,0]>0).long()
rows=[]
for width in [64,256,1024]:
    torch.manual_seed(8);model=WidthProxy(width)
    opt=torch.optim.Adam([{"params":model.input.parameters(),"lr":.003},
                          {"params":model.hidden.parameters(),"lr":.003/model.mult},
                          {"params":model.readout.parameters(),"lr":.003}],weight_decay=0)
    assert opt.param_groups[1]["lr"]*model.mult == .003
    for step in range(4):
        logits,coords=model(x);loss=F.cross_entropy(logits,y)
        assert torch.isfinite(loss)
        rows.append((width,step,coords,loss.item()))
        opt.zero_grad();loss.backward();opt.step()
    # Omitting the readout multiplier changes the current function by mult.
    model.wrong=True;wrong,_=model(x);model.wrong=False;right,_=model(x)
    assert torch.allclose(wrong,right*model.mult,atol=1e-6)
    if width>64: assert not torch.allclose(wrong,right)
# Wrong hidden LR: same initial state and gradients, but first update scales by width multiplier.
import copy
torch.manual_seed(8);correct=WidthProxy(256);wrong=copy.deepcopy(correct)
for model,rate in [(correct,.003/correct.mult),(wrong,.003)]:
    opt=torch.optim.Adam(model.hidden.parameters(),lr=rate)
    logits,_=model(x);F.cross_entropy(logits,y).backward();opt.step()
assert not torch.allclose(correct.hidden.weight,wrong.hidden.weight)
print("width coordinate checks, not curve-equality guarantees",rows)

For actual transfer, verify coordinate plots across widths over several updates/seeds, then sweep base learning rate and compare optimum neighborhoods. Keep depth, tokenization/task, parameter roles and schedule policy fixed while changing intended width dimensions. Save base-shape metadata and avoid rescaling a resumed checkpoint twice. A plausible coordinate plot is a sanity check, not proof every hyperparameter transfers across depth/data/optimizer changes.

The transfer paper reports a 40M proxy to a 6.7B GPT-style target in its experiments. Cerebras-GPT and MiniCPM are public scaling case studies; MiniCPM's exact cost claims and the distinction between standard and μP model families must be read in context rather than assigning all sizes one recipe. Nemotron-H's report concerns hybrid Mamba/Transformer models, not proof that μP applies unchanged to every hybrid block. No private GPT-4 adoption, universal '90% saved' number or reason for another lab's non-adoption is asserted.

8 · Failure cases and repair decisions

• Deep tanh saturates: log bounded activations and actual gradients; test gain/scale, not a fabricated mandatory NaN at step three.
• Hardcoded .01 in a wide/deep model: compute the fan-based k first and check layer moments. A small std is not automatically safe.
• Custom layer zero parameter: calculate its derivative. A linear/bilinear weight can learn from zero; an all-zero symmetric hidden stack can fail. Use a negative control matched to the actual graph.
• He used with tanh: the preactivation scale changes and tanh may saturate; activation std cannot grow past one. Measure instead of inferring divergence from the initializer name.
• μP transfer fails: inspect stored versus effective readout scaling, base shapes, optimizer role groups, attention convention and resume rescaling. Do not declare a scaled output tensor sufficient.

10 · Try it yourself — matched matrices, real tanh and actual gradients

This experiment reuses underlying random matrices across four init/activation choices, includes a BatchNorm variant, and records both activation second moments and input-gradient norm. The SVG is a local report artifact, not a new application asset. Its logarithmic vertical axis shows relative second-moment change; it is not an accuracy benchmark.

import math
from pathlib import Path
import torch
import torch.nn.functional as F
 
torch.manual_seed(23)
base_x=torch.randn(32,64,dtype=torch.float64)
matrices=[torch.randn(64,64,dtype=torch.float64) for _ in range(20)]
curves={}
for tag,scale,act,bn in [("xavier-tanh",1/math.sqrt(64),torch.tanh,False),
                        ("he-relu",math.sqrt(2/64),torch.relu,False),
                        ("he-tanh",math.sqrt(2/64),torch.tanh,False),
                        ("xavier-relu",1/math.sqrt(64),torch.relu,False),
                        ("small-relu-bn",.01,torch.relu,True)]:
    x=base_x.clone().requires_grad_();h=x;values=[]
    for w in matrices:
        h=h@(w*scale)
        if bn: h=(h-h.mean(0))/torch.sqrt(h.var(0,unbiased=False)+1e-5)
        h=act(h);values.append(h.square().mean().item())
    h.sum().backward()
    assert torch.isfinite(x.grad).all()
    curves[tag]=values
    print(tag,"last moment",values[-1],"input grad norm",x.grad.norm().item())
assert curves["xavier-relu"][-1] < curves["he-relu"][-1]*.001
colors=["#2563eb","#059669","#dc2626","#9333ea","#d97706"]
parts=['<svg xmlns="http://www.w3.org/2000/svg" width="760" height="360"><rect width="100%" height="100%" fill="white"/><text x="20" y="22">Initialization: log10 activation second moment by layer</text>']
parts.append('<path d="M35 35 V270 H560" fill="none" stroke="gray"/><text x="220" y="300">layer (1 to 20)</text>')
for tick in [-8,-6,-4,-2,0,2]:
    parts.append(f'<text x="5" y="{99-20*tick}" font-size="12">{tick}</text>')
for i in [0,4,9,14,19]:
    parts.append(f'<text x="{30+i*27}" y="285" font-size="12">{i+1}</text>')
for (tag,values),color in zip(curves.items(),colors):
    pts=" ".join(f"{35+i*27},{95-20*max(-8,min(2,math.log10(v)))}" for i,v in enumerate(values))
    parts.append(f'<polyline points="{pts}" fill="none" stroke="{color}"/><text x="590" y="{50+colors.index(color)*25}" fill="{color}">{tag}</text>')
parts.append('</svg>')
Path('batch6-dl-training-init.svg').write_text(''.join(parts))

11 · Recap and worked recall answers

Q1 — Where does fan-in come from? The sum has fan-in independent zero-mean weight contributions; their second moments add. For zero-mean input and a linear activation, preservation gives variance 1/fan-in. State independence and equal-moment assumptions.

Q2 — Why He's factor two? Symmetry makes ReLU retain half the squared magnitude. Double weight variance, yielding std multiplied by square root two. Gaussian ReLU variance is $(1/2-1/(2\pi))\sigma^2$, not half.

Q3 — Xavier or He? Xavier/gain is a useful roughly linear/tanh recipe; explicit He targets rectifiers. Smooth gates and residual architectures require their actual recipe/moment checks. Framework defaults are not automatically He-ReLU.

Q4 — Are zero weights always broken? No. Symmetric hidden units can remain identical and all-zero deep ReLU paths can block gradients, but the executed zero-linear and ReZero examples learn appropriate gradients. Distinguish the parameter's graph role.

Q5 — Does normalization make init irrelevant? No: readout/logit scale, branch/skip ratios and conditioning/degenerate feature directions remain. Positive global scaling is approximately removed by LN away from epsilon effects, contrary to the old embedding argument.

Q6 — What is μP's benefit and setup cost? It coordinates role-dependent initialization/forward/optimizer scales to support width transfer of useful hyperparameters in studied regimes. Correct base shapes, readout scaling, optimizer convention, coordinate checks and a proxy sweep are the setup cost; exact finite-model equality is not promised.

Sources and execution boundary

Glorot and Bengio 2010, He et al. 2015, PyTorch 2.9 initialization, Linear source, forget-bias study, Fixup, T-Fixup, ReZero, μTransfer and Microsoft μP implementation support the named mechanisms. CPU exact programs, numerical integration and local proxy checks are executed; no language-model-scale transfer, GPU timing or new μP package install is claimed.

Optional pre-read resources

These inherited resource links are optional, not the primary evidence for this review. Video identity, duration and current availability are not certified.

• Weight Initialization for Deep Networks (C2W1L11)
• Building makemore Part 3: Activations, Gradients, BatchNorm
• Understanding the difficulty of training deep feedforward NNs
• Delving Deep into Rectifiers

Previous: ← DL S021 · BatchNorm and LayerNorm · Next: DL S023 · Schedulers →