LayerScale: A Small Change That Stabilizes Deep Vision Transformers
· 2 min read · 379 words
LayerScale is a learned, per-channel multiplier placed at the end of each residual branch in a Transformer block. Its purpose is simple: let a deep network begin training close to the identity function, then learn how strongly each channel should modify the residual stream.
For a pre-normalization block, the two residual updates become
Each is learned. Unlike a single scalar residual weight, LayerScale can strengthen or suppress individual feature channels. The parameters are initialized to a small positive value, so the block initially makes only a small perturbation to its input.
Why it helps
Residual connections make deep networks trainable, but the accumulated updates from many randomly initialized branches can still destabilize optimization. LayerScale controls those updates at initialization without permanently restricting the network: the gains remain trainable.
It is related to ReZero and other residual-scaling methods, but the original LayerScale design uses one gain per channel rather than one gain per block. The paper introduced it while studying deep image Transformers and combined it with careful stochastic-depth scheduling.
Practical guidance
- Use LayerScale when increasing Transformer depth causes unstable training or diminishing returns.
- Treat the initial value as a depth- and recipe-dependent hyperparameter; deeper models generally start with smaller residual gains.
- LayerScale and stochastic depth solve different problems and can be used together, as in the original paper.
- Do not expect it to repair an unstable optimizer, bad normalization, or incorrect learning-rate schedule by itself.
LayerScale adds only one learned vector per residual branch and one elementwise scaling operation.
Primary source
- Hugo Touvron et al., Going deeper with Image Transformers, 2021.
