Why Initialization Matters in a Transformer

When my from-scratch Transformer started training with an unexpectedly high loss, I traced the problem back to parameter initialization. This write-up covers the choices that mattered: small normal weight initialization, a tied output head with final LayerNorm, and depth-aware scaling of the residual branches.

The idea

When I implemented a decoder-only Transformer from scratch, the architecture looked correct and the training pipeline worked, but the model started with a much higher loss than I expected.

Before changing the optimizer or the learning rate, I traced the problem back to the model's starting point: parameter initialization.

Initialization affects three parts of the model: the scale of the weights, the normalization layers, and the residual connections. The sections below take each in turn.

What a fresh model should do

Before training, the model has no reason to prefer one token over another. Ideally every token in the vocabulary is roughly equally likely, so the starting prediction should be close to uniform. With a vocabulary of size $V$, a perfectly uniform prediction gives a cross-entropy loss of:

$$\mathcal{L}_0 = -\ln\frac{1}{V} = \ln V$$

For example, with a vocabulary of 50,257 tokens, this baseline is about 10.8. It also gives a check that can be run before any training: the loss at step zero should be close to $\ln V$. A much higher value means the initial predictions are confidently wrong, and with cross-entropy, assigning very little probability to the correct token produces a large loss.

Small, normal weights

The output head in this model shares its weights with the token embedding, so the scale of the embedding also sets the scale of the logits. PyTorch's default nn.Embedding initialization draws from a standard normal distribution, with a standard deviation of 1. Following GPT-2, I instead draw every linear and embedding weight from a normal distribution with a standard deviation of 0.02 and set all biases to zero:

def _init_weights(self, module):
    if isinstance(module, nn.Linear):
        nn.init.normal_(module.weight, mean=0.0, std=0.02)
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Embedding):
        nn.init.normal_(module.weight, mean=0.0, std=0.02)
    elif isinstance(module, nn.LayerNorm):
        nn.init.ones_(module.weight)
        nn.init.zeros_(module.bias)

With weights this small, the logits start small, and a softmax over small logits is close to uniform. This matches the starting point described above.

Normalization

LayerNorm has two learned parameters for each hidden dimension: a scale $\gamma$ and a bias $\beta$. I initialize them explicitly as:

$$\gamma = 1, \qquad \beta = 0$$

These are also PyTorch's defaults, so this does not fundamentally change LayerNorm. It makes the intended starting behavior explicit: before training, LayerNorm performs ordinary standardization without adding a learned rescaling or offset.

For each token independently, LayerNorm subtracts the mean of that token's hidden vector and divides by its standard deviation:

$$\operatorname{LN}(x_i) = \frac{x_i - \mu_i}{\sqrt{\sigma_i^2 + \epsilon}}$$

This operation happens during every forward pass, not only at initialization. The mean and standard deviation are calculated across the hidden dimensions of one token; LayerNorm does not normalize all tokens in the sequence together.

The Transformer uses a Pre-Norm structure. Each sublayer receives a normalized input, while the original residual stream bypasses the normalization:

$$x' = x + \operatorname{Attention}(\operatorname{LN}_1(x))$$ $$x_{l+1} = x' + \operatorname{MLP}(\operatorname{LN}_2(x'))$$

This gives the attention and MLP branches inputs with a predictable scale while preserving a direct path through the residual stream. LayerNorm does not itself normalize the residual stream after every addition; the stream can still accumulate information from the branch outputs.

Just before the output head, the final LayerNorm normalizes the accumulated residual representation. It controls the overall mean and scale of the vector entering the head, even if the residual stream has grown across layers. It does not erase the relative pattern of features: which hidden dimensions are relatively high or low still carries information about the context.

The logits are produced by multiplying this normalized representation by the tied embedding matrix:

$$\text{logits} = \operatorname{LN}(h) E^\top$$

LayerNorm therefore stabilizes the hidden activation, but it does not normalize the output-head matrix $E$. The shared embedding weights still help determine the logit scale and which token directions receive higher scores. This is why their 0.02 initialization matters even after the final LayerNorm.

Scaling the residual branches

Weight scale is only part of the problem. The deeper question is:

What happens to the scale of the representations as they pass through a deep Transformer?

A single Pre-Norm residual branch can be written as:

$$x_{k+1} = x_k + F_k(\operatorname{LN}_k(x_k))$$

Here, $F_k$ represents either an attention or an MLP branch, and $k$ counts residual branches rather than Transformer layers. Each attention and MLP sublayer adds a new contribution to the same residual stream. The model has six Transformer layers, so there are twelve such branches.

If the residual branches start with weights that are too large, their combined contribution can leave the representations poorly scaled before training has changed anything meaningful.

The key idea is simple:

As the residual network gets deeper, each residual branch should make a smaller initial contribution.

For the projections that write the attention and MLP outputs back into the residual stream, GPT-2-style initialization scales the standard deviation approximately as:

$$\sigma_{\text{residual}} = \frac{\sigma_{\text{base}}}{\sqrt{2N}}$$

where $N$ is the number of Transformer layers.

To see where this comes from, consider a simplified residual stream at initialization:

$$x_{\text{final}} = x_0 + F_1(\operatorname{LN}_1(x_0)) + F_2(\operatorname{LN}_2(x_1)) + \cdots + F_{2N}(\operatorname{LN}_{2N}(x_{2N-1}))$$

There are $2N$ branch additions because every Transformer layer writes to the residual stream twice: once from attention and once from the MLP. If the branch outputs are treated as roughly independent and each has variance $v$, their variances add approximately:

$$\operatorname{Var}(x_{\text{final}}) \approx \operatorname{Var}(x_0) + 2N v$$

Without scaling, the accumulated branch variance therefore grows with depth. To keep the total contribution from all branches at the same order as one unscaled branch, we choose:

$$2N\,\sigma_{\text{residual}}^2 = \sigma_{\text{base}}^2$$

Solving for the standard deviation gives:

$$\sigma_{\text{residual}} = \frac{\sigma_{\text{base}}}{\sqrt{2N}}$$

The factor of 2 counts the two residual-writing sublayers per Transformer layer. The square root appears because variance is proportional to the square of the standard deviation. Thus, $1 / \sqrt{2N}$ is a variance-control rule, not an arbitrary extra multiplier.

This is an initialization heuristic rather than an exact prediction of the model's variance. In practice, branch outputs are not perfectly independent, and their variance depends on the activations and architecture. The goal is to start the residual stream at a controlled scale as depth increases.

The scaling applies only to the projections that write back into the stream: the attention output projection and the second linear layer of the MLP. It is applied to both projections in every Transformer layer during model construction:

def _scale_residual_projections(self):
    scale = (2 * self.num_layers) ** -0.5
    for block in self.blocks:
        nn.init.normal_(
            block.attention_module.linear_out.weight,
            mean=0.0,
            std=0.02 * scale,
        )
        nn.init.normal_(
            block.mlp[2].weight,
            mean=0.0,
            std=0.02 * scale,
        )

With a base standard deviation of 0.02 and six layers, these projections start at about 0.006. The goal is not to make every parameter small. It is to control how the contributions from repeated residual branches accumulate with depth.

This is initialization-time scaling, not a permanent cap. During training, the projection weights are updated and can grow. The scaling gives every residual branch in the stack a conservative starting contribution; it does not force the branches to remain that small forever.

The final LayerNorm rescales the stream before the output head, so this scaling is mainly about how much each branch contributes as depth grows, not about setting the initial logit scale.

What to take away

Initialization is not just a generic first step before training. In this model, three things are at work. Small normal weights keep the starting predictions close to uniform. The final LayerNorm rescales the residual stream before the output head, while the tied embedding matrix still plays an important role in setting the initial logit scale. Residual projection initialization controls how contributions from repeated branches accumulate with depth.

The LayerNorm parameters themselves did not need a special treatment beyond the neutral initialization $\gamma=1$ and $\beta=0$. The more important question was whether the activations and residual updates entering those normalization layers started at a sensible scale.

The most useful lesson was not the $1 / \sqrt{2N}$ scaling itself. It was learning to ask:

What is the model doing before it has learned anything?

When training starts unexpectedly badly, the problem may exist before the optimizer has had a chance to do much. Looking at initialization gave me a clearer understanding of how weight scale, normalization, residual connections, and network depth interact, and why the architecture matters when choosing how to initialize it.