The idea
When I implemented my decoder-only Transformer from scratch, the architecture looked right and the training pipeline was working, but the model started with a much higher loss than I expected.
Instead of immediately changing the optimizer or learning rate, I traced the model back to its starting point: parameter initialization.
Before training, the model has no reason to strongly prefer one token over another. We therefore don't want arbitrary large biases in its initial predictions. With cross-entropy, assigning very little probability to the correct token produces a large loss.
But the output distribution 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 simplified Transformer block can be viewed as:
Each attention and MLP sublayer adds a new contribution to the same residual stream. My model has six Transformer layers, so these contributions accumulate repeatedly through the network.
If the residual branches start with weights that are too large, their combined contribution can make the representations poorly scaled before training has meaningfully changed anything.
My initial implementation relied on PyTorch's default initialization. The model trained, but its starting loss was unexpectedly high. The issue was that the initialization did not account for the depth of the residual network.
I switched to GPT-2-style residual-aware initialization.
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 outputs of the attention and MLP branches back into the residual stream, the initialization is scaled approximately as:
where (N) is the number of Transformer layers.
The goal isn't to make every parameter small. It is to control how the contributions from repeated residual branches accumulate with depth.
This changed how I think about initialization. I initially thought of it mainly as giving every parameter a sensible random starting value. But in a deep architecture, initialization is also about how computation behaves before learning starts.
The most useful lesson wasn't memorizing the (1/\sqrt{2N}) scaling. 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 at all. For my Transformer, looking at initialization led to a much better understanding of how residual connections, network depth, and parameter scale interact.
What to take away
Initialization isn't just a generic first step before training. In a deep residual architecture, it is closely connected to how signals accumulate through the network.
For this Transformer, GPT-style residual-aware initialization gave me a concrete example of how a seemingly small implementation detail can affect the behavior of the entire model, and why understanding the architecture matters when choosing how to initialize it.