Research map / Architectures and efficiency / Attention and transformers
Simplifying Transformer Blocks
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.
Shows that a Transformer block can be simplified by removing skip connections, value/projection parameters, sequential sub-blocks, and normalisation layers while matching Pre-LN training speed per update and gaining ~16% throughput with ~15% fewer parameters.
Problem
Standard Pre-LN transformer blocks are complex and brittle; it is unclear which components (skip connections, value/projection matrices, sequential sub-blocks, normalisation) are necessary for fast training.
Approach
Combine signal propagation theory with empirical findings to build skipless attention (Shaped Attention with identity-initialised, fixed value/projection matrices), then use a parallel attention/MLP block with no skip connections, downweighted MLP residual, and optionally no normalisation, matching Pre-LN per-update speed.
Key results
- SAS-P matches Parallel block training speed with 13% fewer parameters on CodeParrot, and SAS-P trains 16% faster per iteration than Pre-LN in the crammed BERT setting.
- On GLUE, SAS-P (78.3) and SAS (78.4) are within statistical noise of Pre-LN Crammed BERT (78.9) while using 16% fewer parameters (101M vs 120M).
- Removing normalisation layers from SAS-P gives a 1.20x throughput speedup over Pre-LN in the crammed BERT setting while matching training speed in runtime, though with slight per-step degradation.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace the Pre-LN block with SAS | BERT · GLUE | GLUE: 78.9_{\pm.7} -> 78.4_{\pm.8} | No clear effect |
| replace the Pre-LN block with SAS-P | BERT · GLUE | Speed: 1 -> 1.16 | Helped |
| replace the Pre-LN block with SAS | BERT · GLUE | Params: 120M -> 101M | Helped |
| replace the Parallel block with SAS-P | BERT · GLUE | GLUE: 78.5_{\pm.6} -> 78.3_{\pm.4} | No clear effect |
| replace the Parallel block with SAS-P | BERT · GLUE | Speed: 1.05 -> 1.16 | Helped |
| remove normalisation layers from SAS-P (ablation) | SAS-P · GLUE | GLUE | Hurt |
| remove value and projection parameters by fixing them to identity (ablation) | skipless attention block · CodeParrot | per-update training speed | Helped |
| add trainable value and projection parameters initialized to identity (ablation) | Pre-LN block · CodeParrot | loss reduction per update | Hurt |
| add Shaped Attention instead of Value-SkipInit attention (ablation) | SAS and SAS-P · CodeParrot | training performance | Helped |
| remove the MLP skip connection (ablation) | SAS block · CodeParrot | training speed | Hurt |
| linearise ReLU using Leaky ReLU (ablation) | skipless MLP block · CodeParrot | training speed | Hurt |
| replace the linear learning-rate decay schedule with cosine decay (ablation) | SAS and baseline models · CodeParrot | final performance | Hurt |
Surprising findings
- Setting value and projection matrices to identity improves per-update training speed in skipless blocks but slightly hurts them in standard Pre-LN blocks with attention skips.
- Learned residual-skip gain ratios for value/projection matrices converge to zero, so trainable value/projection capacity is unused in skipless blocks.
Limitations
- Models studied are small (100-300M parameters) relative to the largest transformers.
- Removing normalisation layers caused instabilities in fine-tuning and degraded downstream performance.
- Hyperparameters such as the AdamW optimiser and fine-tuning protocol may be overfit to the default Pre-LN block.
Experimental setup
Datasets: CodeParrot The Pile GLUE Languini Books
Benchmarks: GLUE CodeParrot next-token prediction loss Pile masked language modelling loss Languini books perplexity
Metrics: training loss per step wall-clock runtime speed GLUE score parameter count throughput
Baselines: Pre-LN Parallel block Value-SkipInit Crammed BERT
Builds on: Value-SkipInit Shaped Attention Pre-LN transformer parallel transformer block Crammed BERT
Closest papers by meaning
- Transformer Layers as Painters (AAAI 2025 · similarity 0.87)
- HybridNorm: Towards Stable and Efficient Transformer Training via Hybrid Normalization (NeurIPS 2025 · similarity 0.87)
- Unraveling the Gradient Descent Dynamics of Transformers (NeurIPS 2024 · similarity 0.87)
- What Does It Mean to Be a Transformer? Insights from a Theoretical Hessian Analysis (ICLR 2025 · similarity 0.87)
- FrameQuant: Flexible Low-Bit Quantization for Transformers (ICML 2024 · similarity 0.86)
- A General and Efficient Training for Transformer via Token Expansion (CVPR 2024 · similarity 0.86)
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.