The Modern Decoder Block: Pre-Norm, RMSNorm, SwiGLU

Same block, three edits#

The Transformer block of Chapter 7 is still there: attention, then a feed-forward layer, each wrapped in a residual connection with normalization. Llama 2 describes its architecture as the standard Transformer with pre-normalization using RMSNorm, the SwiGLU activation and rotary positional embeddings Established, which is a fair summary of what most open models changed.

1. Normalize first (pre-norm)#

The 2017 block normalized after adding the residual: $\text{Norm}(x + \text{Sublayer}(x))$. Most models now normalize inside the branch: $x + \text{Sublayer}(\text{Norm}(x))$. The residual path from input to output is then an untouched sum.

Xiong and colleagues showed that at initialization the post-norm Transformer has large gradients near the output, which is why it needs a learning-rate warmup, while pre-norm's gradients are well-behaved and it can train without warmup Established. Large runs still warm up (learning-rate schedules), but they are less fragile.

2. RMSNorm: skip the mean#

LayerNorm subtracts each vector's mean and divides by its standard deviation. RMSNorm only divides by the root mean square:

$$\text{RMSNorm}(a)_i = \frac{a_i}{\sqrt{\tfrac{1}{n}\sum_j a_j^2}} \, g_i .$$

Tiny example. $a = [2, -1, 3, 0]$. LayerNorm: mean 1, standard deviation 1.58, giving $[0.63, -1.26, 1.26, -0.63]$. RMSNorm: root mean square $\sqrt{14/4} = 1.87$, giving $[1.07, -0.53, 1.60, 0]$. Both bring the vector to a standard scale; only LayerNorm also centres it.

Zhang and Sennrich hypothesized that LayerNorm's re-centring is dispensable; RMSNorm matched LayerNorm's quality while reducing running time by 7–64% across the models they tested Established.

3. A gated feed-forward layer (SwiGLU)#

The original feed-forward layer is $\max(0, xW_1)W_2$. A gated linear unit multiplies one projection by an activated copy of another:

$$\text{FFN}_{\text{SwiGLU}}(x) = \big(\text{SiLU}(xW) \odot xV\big) W_2, \qquad \text{SiLU}(z) = z\,\sigma(z).$$

Shazeer tested several such variants in the Transformer's feed-forward layer, shrinking the hidden width so the parameter count stayed constant, and found some of them, SwiGLU among them, improved quality over ReLU and GELU Established. He offered no theory for why. Llama 3 8B uses SwiGLU with a model dimension of 4,096 and a feed-forward dimension of 14,336 Established.