Simplifying Transformer Blocks

TL;DR

<3-5 sentence high-level summary> The paper addresses the complexity of standard Transformer blocks and investigates how far they can be simplified without sacrificing training speed. It combines signal propagation theory with empirical observations to remove components such as skip connections, value/projection matrices, sequential sub-blocks, and normalisation, resulting in SAS and SAS-P blocks. Across autoregressive and encoder-only setups (including BERT), the simplified blocks match or exceed the training speed and downstream performance of standard blocks while reducing parameters and increasing throughput by about 15–16%. The work highlights practical pathways to more efficient transformers and provides a framework for further exploration of how architectural choices shape training dynamics.

Abstract

A simple design recipe for deep Transformers is to compose identical building blocks. But standard transformer blocks are far from simple, interweaving attention and MLP sub-blocks with skip connections & normalisation layers in precise arrangements. This complexity leads to brittle architectures, where seemingly minor changes can significantly reduce training speed, or render models untrainable. In this work, we ask to what extent the standard transformer block can be simplified? Combining signal propagation theory and empirical observations, we motivate modifications that allow many block components to be removed with no loss of training speed, including skip connections, projection or value parameters, sequential sub-blocks and normalisation layers. In experiments on both autoregressive decoder-only and BERT encoder-only models, our simplified transformers emulate the per-update training speed and performance of standard transformers, while enjoying 15% faster training throughput, and using 15% fewer parameters.

Figures (28)

• Figure 1: Comparison between different Transformer blocks. (Left) The standard Pre-LN block. (Top Right) Our most simplified block. (Bottom Right) The parallel block zhao2019musegpt-j. Like the parallel block, our block eschews the need for sequential sub-blocks, but we additionally remove all skip connections and normalisation layers, as well as value and projection parameters. Here, $\otimes$ denotes a matrix multiplication, and $\oplus$ denotes a (potentially weighted) sum.
• Figure 2: Loss of training speed in transformers without attention sub-block skip he2023deep, even with Shaped Attention, \ref{['eq:shaped_attn']}, and MLP skips ($\alpha_{\text{FF}}=1$).
• Figure 3: Restricting updates to ${\mathbf{W}}^V, {\mathbf{W}}^P$, through smaller $\beta_V, \beta_P$, recovers training speed in skipless transformers ($\alpha_{\text{SA}}=0$).
• Figure 4: Residual-skip gain ratios $\frac{\beta_V}{\alpha_V},\frac{\beta_P}{\alpha_P}$ converge to 0 during training.
• Figure 5: Training speed in terms of runtime. We see our models match (or even slightly outperform) the Pre-LN block.