Section 08

A full transformer block

Attention, feed-forward layers, and residual connections

Sources: Attention Is All You Need — Vaswani et al., 2017; Root Mean Square Layer Normalization — Zhang & Sennrich, 2019; The Llama 3 Herd of Models — Grattafiori et al., 2024

Attention and the MLP are the two main computations in a standard dense transformer block. A layerlayerOne stage of a neural network: a group of units that transform the values from the previous stage before passing them on.See in glossary → (or transformer block, the names are often interchangeable) combines them with two supporting ideas: residual connections and normalization. These make deep stacks easier to optimize by preserving information paths and keeping activation scales under control.

A modern transformer block, in pseudocode

Here’s the canonical pre-norm structure used in Llama, Mistral, Qwen, and related models:

def transformer_block(x):
    # x: (seq, d_model)
    h = x + attention(rmsnorm(x))    # block A: communication
    y = h + mlp(rmsnorm(h))          # block B: computation
    return y
A modern (pre-norm) transformer block
Data flows top to bottom. The two side loops are residual connections — each sub-block's output is added back to the pre-block input rather than replacing it.
residualresidualinput xRMSNormnormalize before sub-blockMulti-Head AttentionQ, K, V projections · scaled dot-productcross-token mixing+RMSNormnormalize before sub-blockMLP (SwiGLU)up · nonlinearity · downper-token transformation+output yblock Acommunicationblock Bcomputation

Two halves, both following the same shape:

  1. Normalize the input.
  2. Apply some operation (attention, then MLP).
  3. Add the operation’s output to the unnormalized input.

That + is the residual connectionresidual connectionoutput = x + f(x). Lets gradients flow through deep stacks and means each block adds a refinement rather than rewriting.See in glossary →, and the normalize step is the RMSNormRMSNormRoot Mean Square Normalization — a normalization layer that divides each activation by the root-mean-square (√(mean(x²))) of the whole vector, then multiplies by a learned per-dimension scale. Cheaper than LayerNorm (no mean subtraction, no learned bias) and empirically just as good. Standard in Llama-class models.See in glossary →. Together they let a block refine its input without having to reconstruct it from scratch.

Residual connections: out=x+f(x)\text{out} = x + f(x)

Without residuals, a deep stack of layers means the output is block_N(block_{N-1}(... block_1(x) ...)). Every step has to fully reproduce whatever it wants to preserve from earlier layers. Information and gradients must pass through every operation, which makes optimization increasingly difficult as depth grows.

With residuals, each block computes a delta, or change, to add to the running representation. The “main stream” of information is the residual itself, sometimes called the residual stream. Each block reads from it (via the normalize-then-apply path), computes some refinement, and adds it back. If a block has nothing useful to contribute, it can output approximately zero and the residual passes through unchanged.

This is also why models can survive layer pruninglayer pruningRemoving selected layers from a trained network to reduce its computation or memory requirements.See in glossary →, the removal of selected layers, surprisingly well: many layers contribute small refinements, and removing one degrades quality smoothly rather than catastrophically.

RMSNorm: Root Mean Square Normalization

The other helper is normalization. Without it, the magnitudes of the residual stream can grow or shrink as the stack gets deeper, and activations explode or vanish. The fix is to renormalize the vector before each operation.

The original Transformer used LayerNorm:

LayerNorm(x)=γ⋅x−μσ2+ϵ+β\text{LayerNorm}(x) = \gamma \cdot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta

where μ,σ\mu, \sigma are the mean and standard deviation across the vector’s dimensions, and γ,β\gamma, \beta are learned per-dimension scale and shift.

Modern models use RMSNorm, which drops the mean-centering step:

RMSNorm(x)=γ⋅xRMS(x)2+ϵ,RMS(x)=1d∑ixi2\text{RMSNorm}(x) = \gamma \cdot \frac{x}{\sqrt{\text{RMS}(x)^2 + \epsilon}}, \quad \text{RMS}(x) = \sqrt{\tfrac{1}{d} \sum_i x_i^2}

In many studied settings, RMSNorm matches LayerNorm while requiring less computation because it does not calculate the mean. Its learned shift is commonly omitted, though implementations can vary. Llama, Mistral, Qwen, and Gemma use RMSNorm variants.

Pre-norm vs post-norm

There’s one more architectural detail: where you normalize.

The original Transformer was post-norm:

out = norm(x + sublayer(x))

Modern models are pre-norm:

out = x + sublayer(norm(x))

Pre-norm keeps the residual stream unnormalized: only the input to each block is normalized. This often improves optimization at depth because the residual path gives gradients a more direct route backward. It is widespread in decoder-only LLMs, although other normalization schemes also exist.

Parameter count for one layer

For Llama-3-8B, one transformer block contains, roughly:

  • Two RMSNorm scale vectors (4096 params each, tiny).
  • Attention: WQ,WK,WVW_Q, W_K, W_V, but the K and V are smaller because of GQA, see §12. Plus WOW_O. ~42M parameters total.
  • SwiGLU MLP: Wup,Wgate,WdownW_{\text{up}}, W_{\text{gate}}, W_{\text{down}}, each (4096 × 14336). ~176M parameters.

So roughly 218 million parameters per layer, dominated by the MLP. Multiply by 32 layers and add the embedding and output layers, and you have the model.

A language model stacks these blocks between the input embedding table and the output projection. The embeddings turn token IDs into vectors; the output projection turns the final vectors into logits.