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.
Keep reading
- Squared ReLU: The Tiny Activation Change Used by Primer
Squared ReLU computes the square of a rectified activation. This note explains why Primer used it, how it changes gradients, and when not to adopt it blindly.
- Mix-FFN in SegFormer: Adding Local Context to Transformer MLPs
SegFormer Mix-FFN inserts a depthwise 3×3 convolution between two feed-forward projections, giving image tokens local spatial context without explicit positional embeddings.
- Varifocal Loss: Learning IoU-Aware Classification Scores for Detection
Varifocal Loss trains dense detectors to predict a score that combines object confidence with localization quality, while focusing training on difficult negative examples.
