8. Multi-Head Attention From Scratch

Problem

A single attention head can only learn one notion of "relevance" between tokens. Splitting the model dimension into H parallel heads, each running scaled dot-product attention on a d_model / H-sized slice, lets the model attend to different relations at once (e.g. one head tracking syntactic adjacency, another tracking coreference) for roughly the same compute as one big head, since the per-head matmuls are smaller but there are more of them.

Implement MultiHeadAttention.forward: project the input into per-head queries/keys/values by hand (reshape + transpose, not nn.MultiheadAttention), run causal or non-causal scaled dot-product attention independently per head, then merge the heads back and apply the output projection. __init__ already builds the four linear projections — your job is the reshape-attend-merge logic in forward.

Source: src/8_multi_head_attention.py

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int): ...
    def forward(self, x: torch.Tensor, causal: bool = False) -> torch.Tensor: ...
    # x: (B, T, d_model) -> (B, T, d_model), self-attention only

>>> torch.manual_seed(0)
>>> mha = MultiHeadAttention(d_model=16, num_heads=4)
>>> x = torch.randn(2, 5, 16)
>>> mha(x, causal=True).shape
torch.Size([2, 5, 16])

Step-by-Step Approach

1. Project x through q_proj, k_proj, v_proj — each still shape (B, T, d_model) at this point.
2. Split each projection into heads: view(B, T, num_heads, d_head) then transpose(1, 2) to get (B, H, T, d_head), so each head becomes an independent leading-batch-like dimension.
3. Compute per-head scaled dot-product attention scores: q @ k.transpose(-1, -2) / sqrt(d_head), shape (B, H, T, T).
4. If causal, add a (T, T) upper-triangular -inf mask before softmax so token i can't attend to tokens after it.
5. Softmax over the last axis, then multiply by v to get (B, H, T, d_head) per-head outputs.
6. Merge heads back: transpose(1, 2) then .contiguous().view(B, T, d_model) — the .contiguous() is required because transpose leaves the tensor non-contiguous in memory.
7. Apply out_proj to the merged result.

The key insight is that all H heads run in exact parallel through the same batched matmul, using the head dimension as an extra leading axis — there's no Python loop over heads anywhere; the reshape is what turns "one big attention" into "H independent smaller attentions" for free.

Reference solution

import math
import torch
import torch.nn as nn


class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_head = d_model // num_heads
        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)

    def _split_heads(self, x: torch.Tensor) -> torch.Tensor:
        # (B, T, d_model) -> (B, num_heads, T, d_head)
        B, T, _ = x.shape
        return x.view(B, T, self.num_heads, self.d_head).transpose(1, 2)

    def forward(self, x: torch.Tensor, causal: bool = False) -> torch.Tensor:
        B, T, _ = x.shape
        q = self._split_heads(self.q_proj(x))
        k = self._split_heads(self.k_proj(x))
        v = self._split_heads(self.v_proj(x))

        # (B, H, T, d_head) @ (B, H, d_head, T) -> (B, H, T, T)
        scores = q @ k.transpose(-1, -2) / math.sqrt(self.d_head)
        if causal:
            causal_mask = torch.triu(torch.full((T, T), float("-inf"), device=x.device), diagonal=1)
            scores = scores + causal_mask
        weights = torch.softmax(scores, dim=-1)
        out = weights @ v  # (B, H, T, d_head)

        # merge heads: (B, H, T, d_head) -> (B, T, d_model)
        out = out.transpose(1, 2).contiguous().view(B, T, self.d_model)
        return self.out_proj(out)

Key Functions & Tricks

• tensor.view(B, T, num_heads, d_head) — splits the feature axis into a heads axis and a per-head axis, cheap because it's a reshape, not a copy.
• tensor.transpose(1, 2) — moves the heads axis next to the batch axis so batched matmul treats (B, H) together as independent leading dims.
• tensor.contiguous() — required before the merge view(), since transpose only changes stride metadata and leaves memory non-contiguous.
• torch.triu(..., diagonal=1) — builds the strictly-upper-triangular causal mask (diagonal itself stays unmasked, since a token may attend to itself).
• nn.MultiheadAttention(..., batch_first=True) with in_proj_weight/out_proj copied in — the oracle used to verify a from-scratch implementation matches PyTorch's own reference exactly.

How to Recognize This Pattern