Transformers Get Stable: An End-to-End Signal Propagation Theory for Language Models

TL;DR

This work tackles the difficulty of training very deep transformers by deriving end-to-end signal propagation formulas for forward and backward passes, capturing how moments evolve through embeddings, attention, and FFN blocks under realistic non-IID inputs. It then introduces DeepScaleLM, a principled initialization and residual-scaling scheme that conserves unit output and gradient moments across hundreds of layers, enabling depth-forward improvements in NLP, speech, and vision models. The authors validate their theory numerically and empirically, showing stability across Pre-LN and Post-LN variants and across encoder-only, decoder-only, and encoder-decoder architectures, with sustained gains on pretraining perplexity and downstream tasks. The practical impact is substantial: DSLM enables very deep transformers with modest compute overhead and provides a framework for principled scaling choices to balance stability and expressivity. Overall, the paper argues and demonstrates that deeper, well-initialized transformers can outperform shallower counterparts across multiple modalities when signal propagation is controlled via analytic moment equations and residual scaling.

Abstract

In spite of their huge success, transformer models remain difficult to scale in depth. In this work, we develop a unified signal propagation theory and provide formulae that govern the moments of the forward and backward signal through the transformer model. Our framework can be used to understand and mitigate vanishing/exploding gradients, rank collapse, and instability associated with high attention scores. We also propose DeepScaleLM, an initialization and scaling scheme that conserves unit output/gradient moments throughout the model, enabling the training of very deep models with 1000 layers. We find that transformer models could be much deeper - our deep models with fewer parameters outperform shallow models in Language Modeling, Speech Translation, and Image Classification, across encoder-only, decoder-only and encoder-decoder variants, for both Pre-LN and Post-LN transformers, for multiple datasets and model sizes. These improvements also translate into improved performance on downstream Question Answering tasks and improved robustness for Image Classification.

Figures (15)

• Figure 1: Pre-LN: Variance of forward signal increases linearly across layers $N$.
• Figure 2: Pre-LN: Backward gradient variance increases hyperbolically across layers $N$.
• Figure 3: Post-LN: Backward gradient variance vanishes exponentially (y-axis log-scale).
• Figure 4: DeepScaleLM: The variances remain conserved for both forward and backward pass.
• Figure 5: Backward gradient variance increases hyperbolically after $150$k train steps.