Initialization, Loss Functions, and Stable Outputs

For a linear layer \(y = Wx\) with i.i.d. zero-mean weights of variance \(\sigma_w^2\) and inputs of variance \(\sigma_x^2\), each output has variance \(\text{fan\_in}\cdot\sigma_w^2\sigma_x^2\). Stack \(L\) such layers and the activation variance is multiplied by \((\text{fan\_in}\,\sigma_w^2 \cdot c)^L\), where \(c\) is the nonlinearity's effect on the second moment. Above 1 explodes, below 1 vanishes, and backward repeats the product with fan-out. The derivation is in the activations and initialization notes; here we cover the API, the defaults, and measurement.

Karpathy's makemore part 3 sets the habit: before training, inspect per-layer activation and gradient statistics, and check that the step-0 loss is \(\ln V\) for a \(V\)-way classifier. The hooks are from module 10; module 18 turns them into a debugging routine. Residuals and normalization reduce init sensitivity; 13.3 shows they do not remove it.

Every function in torch.nn.init is in-place, runs under no_grad, and returns its tensor. In 2.8 the random ones accept generator= (checked against torch/nn/init.py at v2.8.0), so you can init from a private RNG without touching the global seed.

Trap: trunc_normal_'s a=-2, b=2 are values, not multiples of std, so trunc_normal_(w, std=0.02) truncates at \(\pm 100\sigma\), i.e. never (largest of 10,000 draws measured: 0.077). For \(\pm 2\sigma\) pass a=-2*std, b=2*std.

Gain cancels the nonlinearity's effect on the second moment; calculate_gain(nonlinearity, param=None):

Fan follows layout: _calculate_fan_in_and_fan_out reads dim 1 as inputs, dim 0 as outputs, times the receptive field. A Conv2d weight (64, 3, 7, 7) gives fan_in = 3·49 = 147, fan_out = 64·49 = 3136; nn.Linear's (out, in) gives fan_in = in_features. A weight stored (in, out) (Hugging Face GPT-2 Conv1D) has its fans swapped. mode="fan_in" preserves forward variance, "fan_out" backward; He et al. (arXiv:1502.01852) show either suffices for a plain ReLU net.

nn.Linear.reset_parameters in torch/nn/modules/linear.py calls init.kaiming_uniform_(self.weight, a=math.sqrt(5)) and then draws the bias from \(U(-1/\sqrt{\text{fan\_in}}, 1/\sqrt{\text{fan\_in}})\). Work through the algebra. With leaky-ReLU slope \(a=\sqrt5\), the gain is \(\sqrt{2/(1+5)} = 1/\sqrt3\). kaiming_uniform_ uses the bound \(\text{gain}\cdot\sqrt{3/\text{fan\_in}}\), so

That is the legacy Torch7 U(±1/√fan_in) rule in Kaiming's clothing (the source comment cites issue #57109): one third of the variance-preserving linear value, one sixth of what ReLU needs. Conv*d does the same with \(\text{fan\_in} = C_{in}\prod k\).

The dangerous row is Embedding at \(N(0,1)\) under weight tying (lm_head.weight = wte.weight, logit \(j\) is \(h\cdot e_j\)). At init the final LayerNorm output \(h\) is roughly the input token's own embedding \(e_x\), so that token's logit is \(\lVert e_x\rVert^2 \approx d_{model}\) while the others are \(N(0, d_{model})\): the model confidently predicts "next token = this token". A minimal CPU probe (embedding, LayerNorm, tied head; \(V=512\), \(d_{model}=128\), random targets) measured an initial loss of 128.3 against \(\ln 512 = 6.24\); with normal_(std=0.02) it gave 6.26. The capstone model shows the same effect through its blocks (module 32), which is why GPT-2 and nanoGPT use N(0, 0.02).

Every parametrized module defines reset_parameters(), called by its constructor. To re-init a model, model.apply(fn) visits every submodule (module 10):

def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.normal_(m.weight, std=0.02)
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.Embedding):
        nn.init.normal_(m.weight, std=0.02)

model.apply(init_weights)

apply overwrites everything it matches, so per-name special cases go after it.

A 20-layer, width-256 MLP with a forward hook recording each activation's std and retaining its grad; backward uses a unit-variance random upstream gradient:

def probe(scheme, act=nn.Tanh, depth=20, width=256):
    torch.manual_seed(0)
    layers = []
    for _ in range(depth):
        lin = nn.Linear(width, width)
        if scheme is not None:          # e.g. lambda w: nn.init.xavier_uniform_(w)
            scheme(lin.weight); nn.init.zeros_(lin.bias)
        layers += [lin, act()]
    net, stds, acts = nn.Sequential(*layers), [], []
    def hook(mod, inp, out):
        stds.append(out.std().item()); out.retain_grad(); acts.append(out)
    for m in net:
        if isinstance(m, act):
            m.register_forward_hook(hook)   # returns None: output unchanged
    torch.manual_seed(1); out = net(torch.randn(512, width))
    out.backward(torch.randn_like(out))
    return stds, [a.grad.std().item() for a in acts]

One local PyTorch 2.8 CPU run on an Apple M-series laptop produced:

The default collapses activations to ~0.04 (held up only by the \(U(\pm 1/16)\) biases) and shrinks gradients by \(10^4\) (tanh) to \(10^7\) (ReLU). Xavier with gain 1 is too small for tanh and far too small for ReLU. Matching the gain keeps both statistics within an order of magnitude over 20 layers. std=1 saturates tanh: forward std looks perfect at 0.97 while backward explodes, so forward statistics alone do not prove trainability.

Data-dependent init (LSUV). Instead of deriving a gain, measure it. Mishkin & Matas's layer-sequential unit-variance init (arXiv:1511.06422), which fast.ai's course rebuilds with hooks, starts each weight orthogonal, then walks the layers in order on one real batch and divides each weight by its output's std until that std is 1:

with torch.no_grad():
    h = x_batch                                       # one real (B, d) batch
    for m in net:
        if isinstance(m, nn.Linear):
            nn.init.orthogonal_(m.weight); nn.init.zeros_(m.bias)
            for _ in range(5):                        # converges in a few passes
                m.weight.div_(m(h).std())
        h = m(h)

On a 10-layer width-64 ReLU MLP this ran locally and left every Linear output at std 1. It adapts to any nonlinearity or odd layer without a gain table, at the cost of depending on the batch you happened to use.

Residuals and pre-LN make transformers far more forgiving; what remains is residual-stream variance growing with depth (13.5).

Each pre-LN block adds two branch outputs (attention, MLP) to the residual stream, so with branch variance \(\sigma^2\) the stream reaches \(1 + 2L\sigma^2\). GPT-2 scales residual output projections by \(1/\sqrt{N}\), \(N\) the number of residual layers; nanoGPT sets std = 0.02 / sqrt(2 * n_layer) on every c_proj.weight after self.apply(self._init_weights) (nanoGPT model.py), as does ./code/llm/model.py.

Zero-init goes further: a branch whose last layer is zero is the identity at step 0. Goyal et al. (arXiv:1706.02677) zero the last BatchNorm \(\gamma\) in each ResNet block (torchvision's zero_init_residual=True, not executed here). LoRA zeroes \(B\) so the adapted layer starts equal to the pretrained one (module 36). A gradient path must survive: \(A\) is random so \(\partial L/\partial B \neq 0\); zero both and neither ever moves. For recurrences, orthogonal_ sets every singular value of \(W_{hh}\) to 1 (Saxe et al., arXiv:1312.6120).

nn.CrossEntropyLoss is LogSoftmax followed by NLLLoss in one numerically fused op. For logits \(z\) and target class \(y\),

Log-sum-exp subtracts \(\max_k z_k\) first (module 05), and the gradient \(p - y\) is bounded in \([-1, 1]\) (backpropagation notes). Two recurring mistakes: softmax then CrossEntropyLoss applies a second softmax to values in \([0,1]\), so no class can beat another by more than \(e^1 : e^0\), the loss cannot approach 0, and training silently plateaus; and torch.log(torch.softmax(z)) then NLLLoss is correct algebra that hits -inf once a probability underflows. Use F.log_softmax.

Binary is worse. nn.BCELoss clamps each log at \(-100\); at logit 100 with target 0, sigmoid(100) is exactly 1.0 in float32, the loss is the clamp value 100, and the gradient to the logit is exactly zero: maximally wrong, no signal. nn.BCEWithLogitsLoss computes \(\max(z,0) - zy + \log(1+e^{-|z|})\) with gradient \(\sigma(z) - y = 1\) (docs; fence 4).

nn.CrossEntropyLoss always treats dim 1 as the class dim. Input is (N, C) or (N, C, d1, d2, ...), and index targets are (N) or (N, d1, ...). LM logits are (B, T, V) with classes last, so you either flatten or transpose:

loss = F.cross_entropy(logits.reshape(-1, V), labels.reshape(-1))   # (B*T, V) vs (B*T,)
loss = F.cross_entropy(logits.transpose(1, 2), labels)              # (B, V, T) vs (B, T)

Passing (B, T, V) directly treats \(T\) as classes: a shape error, or silent garbage if \(T = V\).

Label smoothing (Szegedy et al., arXiv:1512.00567) stops logits growing without bound; its minimum is no longer 0, so do not compare smoothed and unsmoothed curves. The weight and smoothing formulas (with ignored positions excluded) matched hand computations to ~1e-7 on CPU.

For LMs you want mean NLL per valid token:

ignore_index with reduction="mean" computes exactly this; a mean of per-sequence means weights a 3-token sequence like a 300-token one (the distinction in the masking module's lab).

Accumulation has the same bias one level up. With \(K\) micro-batches of token-loss sums \(S_k\) and valid counts \(n_k\), the full-batch loss is \(\sum_k S_k / \sum_k n_k\), but backward on each micro-batch mean divided by \(K\) gives the gradient of

Each token in micro-batch \(k\) is weighted \(1/(K n_k)\) instead of \(1/\sum_j n_j\), over-weighting short micro-batches. Fix: count valid labels across all \(K\) first, then backward each micro-batch's sum over that global count:

n_total = sum((mb["labels"] != -100).sum() for mb in micro_batches)  # count first
for mb in micro_batches:
    logits = model(mb["input_ids"])
    s = F.cross_entropy(logits.transpose(1, 2), mb["labels"], reduction="sum")
    (s / n_total).backward()                                         # grads now sum to the full-batch gradient

Under data parallelism, also all-reduce n_total across ranks. Module 16 builds this into the loop; module 33 uses it.

All take reduction (ignore the deprecated size_average/reduce). CE and KL are defined in the information theory notes; the full list is in the nn docs.

Huber vs SmoothL1: huber_loss(x, y, delta=d) == d * smooth_l1_loss(x, y, beta=d) (checked for \(d \in \{0.5, 1, 2\}\); identical at 1); Huber's linear-region gradient is \(\delta\), SmoothL1's is 1. KLDivLoss reverses the math notation: input is the model's log-probs, target the reference, result \(\mathrm{KL}(\text{target}\,\|\,\text{model})\). Default "mean" divides by \(B\cdot C\) and warns in 2.8 ("'batchmean' divides only by the batch size, and aligns with the KL div math definition"); use "batchmean", plus log_target=True for teacher log-probs to avoid log(0) (docs). CTCLoss wants time-major (T, N, C) log-probs; infeasible alignments give inf unless zero_infinity=True (docs). GaussianNLLLoss takes variance, not std, clamped at eps=1e-6.

A loss is any function returning a scalar from differentiable ops; make it an nn.Module only for state such as class-weight buffers that must follow .to(). Work in log space and gather the target column. Focal loss (Lin et al., arXiv:1708.02002) down-weights easy examples by \((1-p_y)^\gamma\):

def focal_loss(logits, target, gamma=2.0, ignore_index=-100):
    # logits (N, C), target (N,) int64
    valid = target != ignore_index
    logp = F.log_softmax(logits[valid], dim=-1)                       # (n, C), stable
    logp_y = logp.gather(1, target[valid].unsqueeze(1)).squeeze(1)    # (n,)
    p_y = logp_y.exp()
    return (-(1 - p_y) ** gamma * logp_y).sum() / valid.sum().clamp(min=1)

At gamma=0 it must equal F.cross_entropy, ignored targets included (checked locally); always test custom losses against such a coinciding reference. clamp(min=1) returns 0, not nan, for an all-padding batch.

Fence 1 proves the API matches 13.2–13.3 (default bound, gains, kaiming_normal_ std, conv fans) on this build, not that any init suits your model.

import math
import torch
import torch.nn as nn

torch.manual_seed(0)
lin = nn.Linear(256, 64)
bound = 1 / math.sqrt(256)
assert lin.weight.abs().max().item() <= bound + 1e-7
assert lin.bias.abs().max().item() <= bound + 1e-7
assert math.isclose(nn.init.calculate_gain("relu"), math.sqrt(2))
assert math.isclose(nn.init.calculate_gain("tanh"), 5 / 3)

g = torch.Generator().manual_seed(1)
w = torch.empty(512, 1024)
nn.init.kaiming_normal_(w, mode="fan_in", nonlinearity="relu", generator=g)
target = math.sqrt(2 / 1024)
rel = abs(w.std().item() - target) / target
assert rel < 0.03, rel
fan_in, fan_out = nn.init._calculate_fan_in_and_fan_out(torch.empty(64, 3, 7, 7))
assert (fan_in, fan_out) == (147, 3136)
print(f"default bound {bound:.4f}, kaiming std rel err {rel:.4f}")

Fence 2 proves CE = log-softmax + NLL, gradient \((p - \text{onehot})/N\), one-hot probability targets = index targets, and the smoothing formula, for (N, C) inputs only.

import torch
import torch.nn.functional as F

torch.manual_seed(0)
N, C = 8, 5
logits = torch.randn(N, C, requires_grad=True)
t = torch.randint(0, C, (N,))

ce = F.cross_entropy(logits, t)
manual = -F.log_softmax(logits, dim=1).gather(1, t[:, None]).squeeze(1).mean()
torch.testing.assert_close(ce, manual, atol=1e-6, rtol=1e-6)

ce.backward()
expected = (torch.softmax(logits.detach(), 1) - F.one_hot(t, C).float()) / N
torch.testing.assert_close(logits.grad, expected, atol=1e-6, rtol=1e-5)

onehot = F.one_hot(t, C).float()
torch.testing.assert_close(F.cross_entropy(logits.detach(), onehot),
                           ce.detach(), atol=1e-6, rtol=1e-6)

eps = 0.1
lp = F.log_softmax(logits.detach(), dim=1)
hand = ((1 - eps) * -lp.gather(1, t[:, None]).squeeze(1) + eps * -lp.mean(1)).mean()
torch.testing.assert_close(F.cross_entropy(logits.detach(), t, label_smoothing=eps),
                           hand, atol=1e-6, rtol=1e-6)
print(f"CE {ce.item():.4f} == logsoftmax+NLL; grad == (p - y)/N")

Fence 3 proves that on (B, T, V) logits with -100 padding the loss is the valid-token mean and ignored positions affect neither loss nor gradient. It does not test the 13.8 accumulation fix (module 16 does).

import torch
import torch.nn.functional as F

torch.manual_seed(0)
B, T, V = 3, 6, 11
logits = torch.randn(B, T, V, requires_grad=True)
labels = torch.randint(0, V, (B, T))
labels[0, 4:] = -100
labels[2, 1:] = -100
valid = labels != -100

loss = F.cross_entropy(logits.transpose(1, 2), labels)
lp = F.log_softmax(logits, dim=-1)
tok = -lp.gather(-1, labels.clamp(min=0).unsqueeze(-1)).squeeze(-1)   # (B, T)
manual = (tok * valid).sum() / valid.sum()
torch.testing.assert_close(loss, manual, atol=1e-6, rtol=1e-6)

perturbed = logits.detach().clone()
perturbed[~valid] += 50.0 * torch.randn(int((~valid).sum()), V)
torch.testing.assert_close(F.cross_entropy(perturbed.transpose(1, 2), labels),
                           loss.detach())
loss.backward()
assert logits.grad[~valid].abs().max().item() == 0.0
print(f"masked CE {loss.item():.4f} over {int(valid.sum())} valid tokens")

Fence 4 proves the BCE logit-100 behaviour of 13.6 and that batchmean KL equals the hand KL while "mean" is \(C\) times smaller and warns, for these 2.8 CPU kernels, not every backend.

import warnings
import torch
import torch.nn as nn
import torch.nn.functional as F

z1 = torch.tensor([100.0], requires_grad=True)
y = torch.tensor([0.0])
l1 = nn.BCEWithLogitsLoss()(z1, y)
l1.backward()
assert torch.isfinite(l1) and abs(z1.grad.item() - 1.0) < 1e-6

z2 = torch.tensor([100.0], requires_grad=True)
l2 = nn.BCELoss()(torch.sigmoid(z2), y)       # sigmoid(100) == 1.0 in float32
l2.backward()
assert z2.grad.item() == 0.0                   # clamped log: no learning signal

torch.manual_seed(0)
Bk, C = 4, 7
student_logp = F.log_softmax(torch.randn(Bk, C), dim=-1)
teacher_logp = F.log_softmax(torch.randn(Bk, C), dim=-1)
p = teacher_logp.exp()
hand = (p * (teacher_logp - student_logp)).sum() / Bk
kl = nn.KLDivLoss(reduction="batchmean", log_target=True)(student_logp, teacher_logp)
torch.testing.assert_close(kl, hand, atol=1e-6, rtol=1e-5)
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    kl_mean = nn.KLDivLoss(reduction="mean", log_target=True)(student_logp, teacher_logp)
assert any("batchmean" in str(w.message) for w in caught)
torch.testing.assert_close(kl_mean * C, kl, atol=1e-6, rtol=1e-5)
print(f"BCE grads logits={z1.grad.item():.1f} probs={z2.grad.item():.1f}; KL {kl.item():.4f}")