Multi-Head Attention & the Full Transformer Block

nn1-9 and nlp1-8 each previewed self-attention and each explicitly deferred "the full multi-head, multi-layer Transformer" to this course. This is that chapter — the single largest deferred-technical-depth thread either course carried.

Why One Attention Computation Isn't Enough

llm1-4's own self-attention computes one set of relationships per sequence — one Query, Key, and Value projection, one resulting pattern of attention weights. Real language has many kinds of relationships happening at once inside the same sentence: which word a pronoun refers to, which verb a subject belongs to, which adjective modifies which noun. Forcing a single attention computation to capture all of these simultaneously is a real constraint — exactly the kind llm1-4's own finding-box already flagged for forcing Query and Key into one shared vector.

Multi-Head Attention

The fix: split the embedding dimension into several smaller heads, each with its own independent set of learned W_Q/W_K/W_V matrices, each computing llm1-4's own scaled dot-product attention independently and in parallel.

heads = []
for i in range(num_heads):
    Q_i = X @ W_Q[i]      # each head gets its own learned projections
    K_i = X @ W_K[i]
    V_i = X @ W_V[i]
    heads.append(scaled_dot_product_attention(Q_i, K_i, V_i))

multi_head_output = concat(heads) @ W_O    # concatenate, then project back down

Each head is free to specialize — one might learn to track subject-verb agreement, another might learn coreference (which noun a pronoun points back to), another might attend mostly to nearby, local context. Concatenating every head's own output and projecting it back through W_O combines all these specialized views into a single vector per token, richer than any one head could produce alone.

The Feed-Forward Sublayer

Attention's own job is mixing information across tokens — deciding what each token should attend to. After attention, every token's own resulting vector independently passes through a small two-layer network, applied identically and separately to each position:

FFN(x) = max(0, x @ W1 + b1) @ W2 + b2   # (or GELU instead of the ReLU shown here)

This is where genuine non-linear transformation of each token's own representation happens, independent of every other token — attention decides what to look at; the feed-forward layer decides what to do with it, position by position.

Residual Connections

Real Transformers stack dozens of these blocks — GPT-3 uses 96 layers. Very deep stacks risk exactly the kind of degraded gradient flow nn1-6's own training material warned about for deep networks generally. The fix here: each sublayer's output is added to its own input, rather than replacing it outright.

x = x + MultiHeadAttention(x)
x = x + FeedForward(x)

Layer Normalization

Each sublayer's input is also normalized — rescaled to have stable mean and variance across its own feature dimension, per token — keeping training numerically well-behaved at depth. Most modern models apply this normalization before each sublayer (Pre-LN) rather than after (Post-LN, the original Transformer paper's own choice), since Pre-LN has proven more stable to train at very large depth.

Assembling One Full Transformer Block

def transformer_block(x):
    x = x + multi_head_attention(layer_norm(x))
    x = x + feed_forward(layer_norm(x))
    return x

for _ in range(num_layers):        # dozens of these, stacked
    x = transformer_block(x)

This is the complete picture: llm1-3's own embedded, position-encoded input vectors enter at the bottom, pass through this exact block repeated dozens of times, and emerge the other end having been repeatedly reshaped by multi-head attention and per-token feed-forward transformation, with residual connections and layer norm keeping the whole stack trainable.

Hands-On Exercises