Lecture 21: Transformer Implementation

What this lecture covers

This lecture answers how the Transformer block described conceptually in the previous lecture translates into actual code. It builds a NumPy implementation of self-attention, extends it to mini-batches, adds multi-head attention, and finally assembles a full Transformer encoder layer, checking each step against PyTorch's reference implementation.

After watching, you can implement self-attention, batched self-attention, multi-head attention, and a Transformer encoder block in plain array code, explain why Transformers need a genuine batch matrix multiplication rather than a reshape-based trick, and describe how PyTorch packs the key/query/value weights and multiple heads into single matrices internally.

Key ideas

• Combined K, Q, V weights: rather than three separate projections, WK, WQ, WV are concatenated into one matrix so X @ W_KQV is computed in a single, larger matrix multiplication, then split into three components; this mirrors how PyTorch stores its in_proj_weight.
• Output projection: after the attention-weighted sum of V, a further linear layer W_out is applied; this is not strictly necessary for single-head attention but matters once multiple heads are combined.
• Mask by addition: the causal mask is added (with -inf on masked entries) rather than subtracted, matching PyTorch's convention.
• Batch matrix multiplication is genuinely different: unlike convolution batching, which can be reduced to one big 2D matrix multiply via reshaping, batched self-attention requires computing K_i @ Q_i.T independently for each batch element, which is a true batch matrix multiply (bmm), not a reshape trick.
• Batch-first layout for Transformers: unlike RNNs, which need (time, batch, hidden) for contiguous per-timestep slices, Transformers should use (batch, time, hidden) since attention multiplies over the trailing two dimensions.
• Multi-head attention: instead of one large attention computation per layer, K, Q, V are split along the feature dimension into H heads, attention is computed independently per head with the score scaled by sqrt(D/H), and the outputs are concatenated back together; the intuition is that a single large softmax wastes the non-linearity's capacity.
• PyTorch quirk: nn.MultiheadAttention returns the average attention matrix across heads, not each head's individual attention matrix.
• Transformer block in code: self-attention, then residual add and layer norm, then a two-layer ReLU feed-forward network, then another residual add and layer norm, implementable in under 20 lines of NumPy.

Walkthrough

Basic self-attention as a layer (0:00)

The lecture reframes self-attention as a module with its own weights WK, WQ, WV, and an output projection W_out, computing softmax(X@WK @ (X@WQ).T / sqrt(D)) @ (X@WV) @ W_out. Biases are set to zero for simplicity. A NumPy softmax helper is written first since NumPy has no built-in.

Implementing and testing single-head attention (2:00)

WK, WQ, WV are combined into one W_KQV matrix so the three projections happen in a single matrix multiplication, then split via np.split along the last axis. The resulting attention() function is compared against PyTorch's nn.MultiheadAttention with one head; the weights are pulled from PyTorch's in_proj_weight and out_proj.weight, and the outputs match to numerical precision.

Mini-batching and why it needs a real batch matmul (15:10)

The lecture argues Transformers should use (batch, T, D) layout, unlike the (T, batch, D) layout used for RNNs, since attention multiplies over the trailing two dimensions and this keeps memory contiguous. It then demonstrates in NumPy that ordinary matrix multiplication of higher-rank tensors (as used for convolution batching) is not the same operation as batch matrix multiplication: multiplying a tensor by a plain 2D matrix flattens the leading dimensions, while true batch matmul computes an independent matrix product per batch element. Self-attention needs the latter. The attention function is then generalized to work in both batched and unbatched form by splitting and transposing on the last axes.

Multi-head attention (30:31)

The motivation given is that a single large dot product per position wastes the softmax non-linearity's expressive power, so K, Q, V are split along the feature dimension into H heads, each of size D/H, attention is computed per head, and results are concatenated. The implementation reshapes each of K, Q, V from (B, T, D) to (B, H, T, D/H) using reshape and swapaxes, runs the same attention computation batched over heads, then swaps back and reshapes to (B, T, D) before applying W_out. This is checked against PyTorch with 4 heads, matching to numerical precision; the lecture notes PyTorch's returned attention matrix is actually the average across heads, unlike this implementation which returns every head's matrix.

Assembling the Transformer block (44:38)

layer_norm and relu helpers are defined, then a transformer function combines them: Z = layer_norm(X + multi_head_attention(X, mask, W_KQV, W_out, heads)), followed by layer_norm(Z + relu(Z @ W_ff1) @ W_ff2). This is compared against PyTorch's nn.TransformerEncoderLayer (noting TransformerEncoderLayer, not the decoder variant, is what's normally wanted outside sequence-to-sequence translation setups), with weights copied over from the PyTorch module's attention and linear layers, and the outputs agree to within floating-point precision.

Before you watch

• Watch the previous lecture on self-attention and Transformer architecture, since this lecture implements exactly the equations it derives.
• Recall the earlier LSTM implementation lecture's discussion of batching and contiguous memory, which this lecture extends and contrasts with true batch matrix multiplication.
• Be comfortable reading and writing NumPy array reshaping and axis operations (reshape, swapaxes, split).

Check your understanding

1. Why can convolution batching be implemented with an ordinary 2D matrix multiply, while self-attention batching requires a genuine batch matrix multiply?
2. Why does the lecture recommend (batch, T, D) layout for Transformers instead of the (T, batch, D) layout used for RNNs?
3. What is the purpose of splitting K, Q, V into multiple heads rather than using one large attention computation?
4. What does PyTorch's nn.MultiheadAttention return as its attention matrix when using multiple heads, and how does that differ from this lecture's implementation?