Multi-Head Attention and the Transformer Block

After this lesson, you will be able to: Assemble multi-head attention and the full transformer block (residual connections, layer normalization, feed-forward network) and explain encoder vs decoder stacks.

One attention head is a building block; a transformer stacks many heads and layers with the glue that makes deep networks trainable. This lesson assembles the whole block so a model architecture diagram becomes readable.

Prerequisites:Attention from the Math Up

Multi-head attention

Instead of one attention computation over the full d_model dimension, the transformer splits it into h heads, each operating on d_model/h dimensions with its own Q/K/V projections. The heads run in parallel and their outputs are concatenated and projected back. The point: different heads learn different relationships, one might track syntactic dependencies, another coreference, another positional patterns. More heads give the model more relationship types to attend to at once, at the same total compute as a single full-width head.

The transformer block: attention + FFN, wrapped in residuals and norm

A block is two sub-layers. First, multi-head self-attention. Second, a position-wise feed-forward network (FFN): two linear layers with a nonlinearity (GELU/SwiGLU), applied to each position independently, usually expanding to 4x d_model and back. Around each sub-layer are a residual connection (add the input back to the output) and layer normalization. Residuals let gradients flow through dozens of layers without vanishing; layer norm keeps activations well-scaled. Modern models use pre-norm (normalize before the sub-layer), which trains more stably than the original post-norm.

A transformer block in PyTorch-style pseudocode

Pre-norm decoder block, the modern default.

python

class Block(nn.Module):
    def __init__(self, d_model, n_heads):
        self.attn = MultiHeadAttention(d_model, n_heads)  # causal
        self.ffn  = nn.Sequential(
            nn.Linear(d_model, 4 * d_model), nn.GELU(),
            nn.Linear(4 * d_model, d_model),
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, x):
        # pre-norm: normalize, sub-layer, then add the residual
        x = x + self.attn(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x
# A GPT is: token embedding + positions -> N of these blocks -> final norm
# -> linear projection to vocab logits.

Common mistakes only experienced researchers catch

Forgetting the residual connection (the network becomes nearly untrainable at depth). Putting layer norm in the wrong place (post-norm vs pre-norm changes stability; report which you used). Counting heads but forgetting each head's dimension is d_model/h, not d_model. Treating the FFN as an afterthought when it holds most of the parameters and is where a lot of factual knowledge appears to live. Drawing an encoder-decoder when you mean decoder-only, which signals you have not read the architecture you are citing.

Quick Check

What is the difference between BERT and GPT at the attention level?

Pick the precise distinction.

←Attention from the Math Up

Back to Transformers and NLP Foundations

Pretraining Objectives→