Research map / Architectures and efficiency / Attention and transformers
Transformers Get Stable: An End-to-End Signal Propagation Theory for Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Develops an end-to-end signal propagation theory giving closed-form forward/backward moment formulas for transformer components and blocks, and uses it to design DeepScaleLM, an initialization/residual-scaling scheme that trains very deep (up to ~192-1000 layer) transformers that beat shallower baselines.
Problem
Deep transformers suffer from vanishing/exploding gradients, rank collapse, and instability from large QK values, and prior signal-propagation analyses rely on simplifying assumptions (IID inputs, uncorrelated outputs, ignored QK init, simplified nonlinearities).
Approach
Derive closed-form mean/variance/correlation propagation formulas for each transformer component (including input token correlation via Zipf, dropout, ReLU/GeLU, softmax, LayerNorm) and blocks, then use them to set per-layer initialization and residual scaling (lambda, beta with beta^2 = k/N) so forward and backward signal variances stay near unit across depth.
Key results
- A 192-layer 256-d BERT with DeepScaleLM reaches 12.9 perplexity vs 14.2 for the 12-layer 1024-d baseline and outperforms the 24-layer 336M baseline with more than twice the parameters.
- BERT-base trained for 64B tokens with DSLM (48/384) improves MNLI from 82.4 to 83.7 and RACE-High from 63.7 to 65.7.
- DSLM on ViT improves ImageNet top-1 from 79.8 to 80.3 at 300 epochs with a 4x deeper model of equal parameters.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply DeepScaleLM to the deeper 48/512 model | BERT, Post-LN · Pile-CC MLM, 48/512 | perplexity: 14.8 -> 13.1 | Helped |
| Apply DeepScaleLM to the deeper 192/256 model | BERT, Post-LN · Pile-CC MLM, 192/256 | perplexity: 17.2 -> 12.9 | Helped |
| Apply DeepScaleLM to the shallow 12/1024 model | BERT, Post-LN · Pile-CC MLM, 12/1024 | perplexity: 14.2 -> 15.5 | Hurt |
| Apply DeepScaleLM to the deeper Pre-LN model | BERT · Pre-LN BERT MLM, 96/512 | perplexity: 20.6 -> 15.4 | Helped |
| Use the 4x deeper DSLM model | BERT-base · MNLI | accuracy: 82.4 -> 83.7 | Helped |
| Use the 4x deeper DSLM model | BERT-base · QQP | accuracy: 90.8 -> 91.1 | Helped |
| Apply DeepScaleLM | GPT, Pre-LN · GPT, 48/512 | perplexity: Baseline 12/1024: 11.6; DSLM 48/512: 11.2 | Helped |
| Apply DeepScaleLM | Speech encoder-decoder Transformer, Pre-LN · MuST-C en→de | BLEU: 24.9 -> 25.6 | Helped |
| Apply DeepScaleLM | ViT · ImageNet, 90-epoch | accuracy: 76.5 -> 77.2 | Helped |
| Apply DeepScaleLM | ViT · ImageNet-v2, 90-epoch | accuracy: 63.7 -> 65.2 | Helped |
| Remove residual scaling from DSLM (ablation) | 96-layer Post-LN BERT · BERT MLM, 96/512 | perplexity: 11.7 -> diverge | Hurt |
| Make β learnable instead of fixed (ablation) | DSLM-Init + β²=2/N, 96-layer Post-LN BERT · BERT MLM, 96/512 | performance: 11.7 -> 12.2 | Hurt |
Surprising findings
- Xavier initialization with QK values scaled 2x can blow up backward gradients roughly 10000x through a 192-layer model.
- Post-LN models with standard init diverge when deepened, while DSLM allows a 768-layer Post-LN model to converge.
Limitations
- Theory relies on assumptions such as Gaussian inputs and positive covariance; negative covariance is not analyzed in full.
- Experiments are mostly limited to Chinchilla-scale training due to compute, with some models trained with sequence-length limits.
- Analysis of gradient and forward growth is at initialization, with post-training validity shown empirically only for selected models.
Experimental setup
Datasets: Pile-CC Wikipedia RACE-M RACE-H MNLI QQP MuST-C ImageNet-1k
Benchmarks: MNLI QQP RACE-Middle RACE-High MuST-C en-de ImageNet ImageNet-v2 ImageNet-R ImageNet-Sketch
Metrics: perplexity accuracy BLEU forward/backward variance relative error
Baselines: Vanilla Pre-LN Vanilla Post-LN DSInit ADMIN SkipInit ReZero LayerScale DeepNet Bamboo Value-SkipInit muP
Builds on: Xavier initialization DeepNet DSInit ADMIN Fixup
Closest papers by meaning
- The Curse of Depth in Large Language Models (NeurIPS 2025 · similarity 0.89)
- Building on Efficient Foundations: Effective Training of LLMs with Structured Feedforward Layers (NeurIPS 2024 · similarity 0.89)
- A Multi-Level Framework for Accelerating Training Transformer Models (ICLR 2024 · similarity 0.88)
- AttnLRP: Attention-Aware Layer-Wise Relevance Propagation for Transformers (ICML 2024 · similarity 0.88)
- DenseFormer: Enhancing Information Flow in Transformers via Depth Weighted Averaging (NeurIPS 2024 · similarity 0.88)
- Exploring the Role of Large Language Models in Prompt Encoding for Diffusion Models (NeurIPS 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.