Understanding Transformers · Part 12 of 16

The Complete Transformer Block

Putting attention, the feedforward network, residuals, and normalization together into the unit that gets stacked.

The intuition: two moves, gather then process

A transformer block makes two moves. First attention lets each token gather information from the others (Part 7). Then the feedforward network lets each token process what it gathered, on its own (Part 10). Residual connections carry the incoming representation around each move so nothing has to be rebuilt from scratch, and normalization keeps the numbers in range. Between the embedding layer and the output projection, the model is this one block repeated, each copy with its own learned parameters.

Modern decoder-only models usually normalize before each sublayer (pre-norm): normalize x, apply the sublayer, then add the result back to x. The original transformer normalized after each sublayer instead (Vaswani et al., 2017); Xiong et al. (2020) showed that pre-norm has better-behaved gradients at initialization and trains without the learning-rate warmup post-norm needs.

import numpy as np

# Simplified: no learned gain/bias in layer_norm, no causal mask, one head.
def layer_norm(x, eps=1e-5):
    mean, std = x.mean(-1, keepdims=True), x.std(-1, keepdims=True)
    return (x - mean) / (std + eps)

def attention(x, Wq, Wk, Wv):        # Part 7
    Q, K, V = x @ Wq, x @ Wk, x @ Wv
    scores = Q @ K.T / np.sqrt(Q.shape[-1])
    weights = np.exp(scores) / np.exp(scores).sum(-1, keepdims=True)
    return weights @ V

def ffn(x, W1, W2):                  # Part 10
    return (np.maximum(0, x @ W1.T)) @ W2.T

def transformer_block(X, p):
    A = X + attention(layer_norm(X), p.Wq, p.Wk, p.Wv)
    return A + ffn(layer_norm(A), p.W1, p.W2)

flowchart TD
  X["X (input)"] --> LN1["LayerNorm"]
  LN1 --> ATT["Attention"]
  ATT --> ADD1(("+"))
  X -. residual .-> ADD1
  ADD1 --> LN2["LayerNorm"]
  LN2 --> FFN["Feedforward"]
  FFN --> ADD2(("+"))
  ADD1 -. residual .-> ADD2
  ADD2 --> XN["X_next (to the next block)"]
  

The feedforward network's place

The FFN runs at each token position independently. It expands the vector to a larger hidden dimension, applies an activation such as ReLU or GELU (or a gated variant such as SwiGLU), and projects back to the model dimension (Part 10). Attention moves information across positions; the FFN does the position-wise computation on what each position has collected.

Stacking blocks

After one block, a token's representation holds some context. After many, it can reflect more involved relationships, because each block operates on the output of the one below it. The model dimension stays constant across blocks, while attention heads and FFNs work in internal subspaces. Identical block shapes make the stack easy to scale and to split across hardware.

Pre-norm stacks apply one final layer normalization after the last block, just before the hidden vector is projected to vocabulary scores (Xiong et al., 2020). The residual stream keeps accumulating sublayer outputs, so its scale tends to grow with depth, and this last normalization puts it on a fixed scale before the output projection.