Neural networks, deep learning, NLP · insight-0418

When a Transformer diverges early, check normalization placement and warmup

Situation
A Transformer trained from scratch shows loss spikes or NaNs in the first few hundred steps at a learning rate that works later in training.
Action
Check whether layers are Post-LN (normalization after the residual add) or Pre-LN (inside the residual branch). Use Pre-LN, or keep Post-LN with a learning-rate warmup.
Reason
Xiong et al. show Post-LN Transformers have large gradients near the output at initialization, which warmup avoids, while Pre-LN gradients are well-behaved and can train without warmup.
Boundary
Encoder or decoder stacks you build or modify yourself.

Symptom or risk

Early instability forces a very small learning rate for the whole run, wasting compute, or the run diverges and is restarted with ad hoc changes.

Fine-tuning a pretrained checkpoint keeps its original placement; changing it breaks the weights.

Pre-LN trains more easily; whether final quality changes depends on the setup, so compare both on validation data.

Removing warmup from a Post-LN recipe because a Pre-LN paper did.

Xiong et al. prove that Post-LN Transformers have large expected gradients near the output at initialization and that Pre-LN gradients are well-behaved, allowing warmup to be removed.
On Layer Normalization in the Transformer Architecture