LM3-L3 — Implement Causal Multi-Head Attention in PyTorch

Objective

Implement causal self-attention and a multi-head attention module whose outputs match a reference built with torch.nn.functional.scaled_dot_product_attention, whose outputs at earlier positions do not change when a later token changes, and whose tensor shapes you can state at every step.

Why a notebook? This exercise trains or runs a real PyTorch model. PyTorch does not run in the browser Python used by the site’s other labs, so you run this one on your own computer or in a hosted notebook service. The interactive browser labs of this module cover the surrounding ideas without PyTorch.

Setup and requirements

Hardware: CPU only, about 2 GB of RAM, a few minutes. CPU only; about 2 GB of RAM; the tensors are tiny (a few hundred numbers), so every cell runs in well under a second after PyTorch has loaded.

Download the notebook (.ipynb) Reference solution (try the lab first)

On your own machine (Python 3.10 or newer):

python3 -m venv llm-env
source llm-env/bin/activate        # on Windows: llm-env\Scripts\activate
pip install torch numpy jupyterlab
jupyter lab causal-multi-head-attention-in-pytorch.ipynb

In a hosted notebook service such as Google Colab or Kaggle, upload the file and run it there; PyTorch is usually preinstalled. Matplotlib is used for optional plots when it is available; the notebook runs without it.

What you will do

1. Single head. Project the input to queries, keys and values, scale the scores, mask the future with −∞ before the softmax, and mix the values.
2. Leakage test. Run a test that changes the last token and checks that earlier outputs do not move; see it catch an unmasked layer.
3. Multi-head module. Split the projections into heads with view and transpose, attend in parallel, merge the heads and apply an output projection.
4. Shape tracing. Write down the shape of every intermediate tensor and check it against the real ones.
5. Cost experiment. Read how the memory for attention scores grows with the square of the sequence length.

Cells that you must complete are marked TODO in the notebook and have a self-check right after them (an assert with a message). Run the cells from the top; a failing check tells you what is wrong.

Expected behaviour

• Self-check 1: your single-head function returns outputs of shape (B, T, d_k) and weights of shape (B, T, T) with rows that sum to 1, zeros above the diagonal, and values equal to F.scaled_dot_product_attention(..., is_causal=True).
• Self-check 2: the unmasked layer leaks (a change of the last token moves earlier outputs); yours changes them by less than 1e-6.
• Self-check 3: the multi-head module has the right output shape, does not leak, matches the PyTorch reference, and has 3·d_in·d_out + d_out² + d_out parameters regardless of the number of heads.
• Self-check 4: your shape tuples match the real tensors.

How to know you are done

• All four self-checks print “passed” when the notebook is run from the top.
• You can explain why the mask is applied before the softmax and why the number of heads does not change the parameter count.
• You can state the shape of the score tensor and why its memory grows with the square of the sequence length.

The site cannot run or grade the notebook for you. The self-checks inside the notebook are your verification. The reference solution was executed with PyTorch on a CPU when this module was prepared; your own environment may differ in version or in the last decimals of the floating-point results.

If something goes wrong

• A shape error in view: after transpose a tensor is not contiguous in memory; call .contiguous() before view, or use reshape.
• The row sums are not 1: you masked after the softmax; set the future scores to float("-inf") first.
• Not-a-number values: a row that is entirely −∞ gives NaN. In a causal mask every row keeps at least its own position, so check that the mask is diagonal=1.
• No PyTorch installed: run pip install torch numpy in your virtual environment; the CPU build is enough.

Theory for this lab

• LM3-T5: Multi-Head Attention
• LM3-T6: Tensor Shapes and Matrix Operations in Attention

Reference material

• Build a Large Language Model (From Scratch), Chapter 3: Coding Attention Mechanisms

Embeddings and the Mathematics of Attention labs

1. LM3-L1: Compute Attention Step by Step
2. LM3-L2: Implement Causal Masking
3. LM3-L3: Implement Causal Multi-Head Attention in PyTorch: this lab.

Next lab: LM4-L1: Count GPT Parameters and Tensor Shapes.