PYTORCH · BUILD IT YOURSELF

Lab: implement multi-head attention

Introduction

One query compares with two keys, producing softmax shares near 0.67 and 0.33. Those shares mix two value vectors into one context vector before the same calculation expands to batches and heads.

Learning goal

Implement and test multi-head attention from score calculation through masking, value mixing, and final output projection.

Before you start

Dot products, softmax, tensor shapes, matrix multiplication, and basic PyTorch functions.

Lesson plan

  1. Work one query-key-value example by hand and rewrite it as a function.
  2. Add batch and head axes while reading transpose, reshape, and matrix multiplication carefully.
  3. Separate mask meanings, connect the function to a module, and run failure tests.

Full build path and setup · Download transformer parts · Download training runner · Download correctness tests

Start with one question and two values

Predict the context: if normalized weights are [0.2, 0.8] and scalar values are [10, 5], the result is 0.2 × 10 + 0.8 × 5 = 6. The larger weight gives the second value more influence.

Attention combines information from several positions into one answer.

Query — what am I looking for?
A vector from the position that needs an answer.
Key — how well do I match?
A vector from each source position. Compare it with the query.
Value — what information can I offer?
The information to mix after calculating the matches.

Work through one small example

Use one query, two keys, and two values. These are chosen numbers, not a trained model's predictions.

  1. Choose the inputs.
    Query: [1,0]
    Keys: [1,0] and [0,1]
    Values: [10,0] and [0,20]
  2. Compare query and keys.
    The dot products are [1,0].
  3. Scale the scores.
    Each key has two features, so divide by √2: [0.707,0].
  4. Convert scores into shares.
    Softmax gives approximately [0.670,0.330]. The shares add to one.
  5. Mix the values.
    0.670 × [10,0] + 0.330 × [0,20] ≈ [6.698,6.605]

Write the calculation as a function

import math
import torch

def attention(q, k, v, allowed):
    scores = q @ k.transpose(-2, -1) / math.sqrt(q.shape[-1])
    if not allowed.any(dim=-1).all():
        raise ValueError("Each query needs at least one allowed key")
    scores = scores.masked_fill(~allowed, float("-inf"))
    weights = scores.softmax(dim=-1)
    return weights @ v, weights

q = torch.tensor([[1., 0.]])
k = torch.eye(2)
v = torch.tensor([[10., 0.], [0., 20.]])
out, weights = attention(q, k, v, torch.ones(1, 2, dtype=torch.bool))
torch.testing.assert_close(weights.sum(-1), torch.ones(1))
print(out)  # approximately tensor([[6.6976, 6.6048]])

Read the three important operations

The scale uses the width of one key, not the sequence length. With many roughly independent components, unscaled dot products can grow large. Softmax then becomes very concentrated. Scaling helps control this effect. It does not guarantee good training.

Add batches and heads without changing the idea

QuantityShapeMeaning
InputB × T × CB sequences, T positions, C features
QueriesB × H × T × DH heads, D features per head
Keys and valuesB × H × S × DS available source positions
WeightsB × H × T × SOne distribution over sources per query
Joined outputB × T × CC = H × D

A projection computes a weighted sum plus an optional bias. Three learned projections create Q, K, and V.

  1. Reshape (B,T,C) into (B,T,H,D).
  2. Transpose positions and heads to get (B,H,T,D).
  3. Calculate attention separately within each head.
  4. Reverse the transpose, then reshape to join the heads.
  5. Apply the output projection.

Common mistake: reshaping without transposing mixes positions and heads incorrectly.

Keep three different mask ideas separate

Inspect a single query row. Its allowed weights should sum to 1. Its forbidden weights should be 0. If a future or padding position receives weight, the mask is wrong even when the final tensor shape is correct.

Combine the allowed positions with logical AND. A position must pass every applicable rule.

In this course's attention function, True means allowed. Do not transfer that meaning blindly to another API. PyTorch's scaled-dot-product function uses this meaning, while the boolean masks in MultiheadAttention use a blocking convention. Read the relevant API before replacing the educational implementation.

Check the order

  1. At query position 2, allow keys 0, 1, and 2. Block key 3.
  2. Apply the mask before softmax. Masking the final output cannot undo leaked information.
  3. Check that every row has an allowed key. Softmax over only negative infinity produces invalid values.

A row with no allowed keys needs an error or a deliberately designed fallback.

Connect the function to a model

The numbers in the first example were chosen by us. In a model, training learns projection weights. Attention itself has no separate target. The final prediction loss sends gradients through the value mixture, softmax, and projections. The complete transformer runner below supplies that loss and optimizer.

New Python syntax in the full source

A class groups learned tensors with the function that uses them. You do not need to know object-oriented design before this lesson.

Break it, then prove the repair

  1. Normalize the wrong axis. Replace dim=-1 with dim=-2. Weights now normalize across queries, not sources. Check row sums to catch this.
  2. Remove the causal mask. Change the last input. If an earlier output changes, future information has leaked into the prediction.
Exercise: what happens if only the first key is allowed?

The weights become [1,0]. The answer is exactly the first value, [10,0]. Blocking a source changes the weight distribution before mixing, rather than setting one final feature to zero.

Run the attention, causal-mask, and gradient tests linked below. Then train the complete model. Passing an attention unit test proves local arithmetic. It does not prove that a model learned a useful task.

Read the executable source

Keep these three ideas

These source listings are copied from the runnable files by scripts/sync-build-code.mjs. The training runner and module belong in the same folder. Read a section, run the tests, then make one small change.

ManualMultiHeadAttention — executable Python source
class ManualMultiHeadAttention(nn.Module):
    """Multi-head attention with explicit score, scale, mask, and softmax.

    ``query`` has shape ``(B, Tq, D)``. ``key_value`` has shape ``(B, Tk, D)``
    and defaults to ``query`` for self-attention.  Normally the return value is
    a tensor ``(B, Tq, D)``.  With ``use_cache=True`` the return value is
    ``(output, new_cache)``.
    """

    def __init__(
        self,
        d_model: int,
        num_heads: int,
        *,
        dropout: float = 0.0,
        use_rope: bool = False,
        bias: bool = True,
    ) -> None:
        super().__init__()
        if d_model % num_heads:
            raise ValueError("d_model must be divisible by num_heads")
        head_dim = d_model // num_heads
        if use_rope and head_dim % 2:
            raise ValueError("RoPE requires an even head dimension")
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.scale = head_dim**-0.5
        self.use_rope = use_rope
        self.q_proj = nn.Linear(d_model, d_model, bias=bias)
        self.k_proj = nn.Linear(d_model, d_model, bias=bias)
        self.v_proj = nn.Linear(d_model, d_model, bias=bias)
        self.out_proj = nn.Linear(d_model, d_model, bias=bias)
        self.attention_dropout = nn.Dropout(dropout)
        self.output_dropout = nn.Dropout(dropout)
        self.rope = RotaryEmbedding()

    def split_heads(self, x: torch.Tensor) -> torch.Tensor:
        """Convert ``(B, T, D)`` to ``(B, H, T, Dh)``."""
        batch, length, width = x.shape
        if width != self.d_model:
            raise ValueError(f"last dimension must be d_model={self.d_model}")
        return x.reshape(batch, length, self.num_heads, self.head_dim).transpose(1, 2)

    def merge_heads(self, x: torch.Tensor) -> torch.Tensor:
        """Convert ``(B, H, T, Dh)`` back to ``(B, T, D)``."""
        batch, heads, length, head_dim = x.shape
        if heads != self.num_heads or head_dim != self.head_dim:
            raise ValueError("head dimensions do not match this attention module")
        return x.transpose(1, 2).contiguous().reshape(batch, length, self.d_model)

    def _expand_attention_mask(
        self,
        mask: torch.Tensor,
        batch: int,
        query_length: int,
        key_length: int,
    ) -> torch.Tensor:
        if mask.dtype != torch.bool:
            raise TypeError("attention_mask must be bool (True means allowed)")
        if mask.ndim == 2:
            if tuple(mask.shape) != (query_length, key_length):
                raise ValueError("2D attention_mask must have shape (Tq, Tk)")
            return mask[None, None, :, :]
        if mask.ndim == 3:
            if tuple(mask.shape) != (batch, query_length, key_length):
                raise ValueError("3D attention_mask must have shape (B, Tq, Tk)")
            return mask[:, None, :, :]
        if mask.ndim == 4 and mask.shape[0] in (1, batch) and mask.shape[1] in (
            1,
            self.num_heads,
        ):
            if tuple(mask.shape[-2:]) != (query_length, key_length):
                raise ValueError("4D attention_mask has incorrect sequence dimensions")
            return mask
        raise ValueError("attention_mask must have 2, 3, or 4 dimensions")

    def forward(
        self,
        query: torch.Tensor,
        key_value: torch.Tensor | None = None,
        *,
        attention_mask: torch.Tensor | None = None,
        key_padding_mask: torch.Tensor | None = None,
        query_padding_mask: torch.Tensor | None = None,
        causal: bool = False,
        cache: KVCache | None = None,
        use_cache: bool = False,
        query_position_offset: int | None = None,
        key_position_offset: int = 0,
    ) -> torch.Tensor | tuple[torch.Tensor, KVCache]:
        if query.ndim != 3 or query.shape[-1] != self.d_model:
            raise ValueError("query must have shape (B, Tq, d_model)")
        if query.shape[1] == 0:
            raise ValueError("attention needs at least one query token")
        if cache is not None and not use_cache:
            raise ValueError("cache input requires use_cache=True so the updated cache is returned")
        if use_cache:
            expected_offset = cache.length if cache is not None else 0
            if key_position_offset != 0 or query_position_offset not in (None, expected_offset):
                raise ValueError("cached positions must start at zero and continue from cache length")
        if cache is not None and key_value is not None and key_value is not query:
            raise ValueError("KV caching is only supported for self-attention")
        if use_cache and key_value is not None and key_value is not query:
            raise ValueError("KV caching is only supported for self-attention")
        if use_cache and self.training and (
            self.attention_dropout.p > 0.0 or self.output_dropout.p > 0.0
        ):
            raise ValueError("KV caching with dropout is inference-only; call eval() first")

        is_self_attention = key_value is None or key_value is query
        source = query if key_value is None else key_value
        if source.ndim != 3 or source.shape[0] != query.shape[0] or source.shape[-1] != self.d_model:
            raise ValueError("key_value must have shape (B, Tk, d_model)")

        batch, query_length, _ = query.shape
        current_key_length = source.shape[1]
        if key_padding_mask is None:
            current_valid = torch.ones(
                (batch, current_key_length), dtype=torch.bool, device=query.device
            )
        else:
            _check_valid_mask(
                key_padding_mask, (batch, current_key_length), "key_padding_mask"
            )
            current_valid = key_padding_mask
        if query_padding_mask is not None:
            _check_valid_mask(
                query_padding_mask, (batch, query_length), "query_padding_mask"
            )

        q = self.split_heads(self.q_proj(query))
        new_k = self.split_heads(self.k_proj(source))
        new_v = self.split_heads(self.v_proj(source))

        inferred_offset = cache.length if cache is not None else (
            key_position_offset if is_self_attention else 0
        )
        q_offset = inferred_offset if query_position_offset is None else query_position_offset
        if q_offset < 0 or key_position_offset < 0:
            raise ValueError("position offsets must be non-negative")
        if self.use_rope:
            q = self.rope(q, offset=q_offset)
            new_k_offset = inferred_offset if cache is not None else key_position_offset
            new_k = self.rope(new_k, offset=new_k_offset)

        if cache is not None:
            if cache.key.shape[:2] != (batch, self.num_heads):
                raise ValueError("cache batch/head dimensions do not match")
            if cache.key.shape[-1] != self.head_dim:
                raise ValueError("cache head width does not match")
            k = torch.cat((cache.key, new_k), dim=2)
            v = torch.cat((cache.value, new_v), dim=2)
            valid_keys = torch.cat((cache.key_padding_mask, current_valid), dim=1)
        else:
            k, v, valid_keys = new_k, new_v, current_valid

        key_length = k.shape[2]
        allowed = valid_keys[:, None, None, :]  # (B, 1, 1, Tk)
        if causal:
            query_positions = torch.arange(
                q_offset, q_offset + query_length, device=query.device
            )
            first_key_position = 0 if cache is not None else key_position_offset
            key_positions = torch.arange(
                first_key_position, first_key_position + key_length, device=query.device
            )
            causal_mask = key_positions[None, :] <= query_positions[:, None]
            allowed = allowed & causal_mask[None, None, :, :]
        if attention_mask is not None:
            expanded = self._expand_attention_mask(
                attention_mask, batch, query_length, key_length
            )
            allowed = allowed & expanded

        # Expand before checking so each batch/head/query row is checked.
        allowed = allowed.expand(batch, self.num_heads, query_length, key_length)
        if not bool(allowed.any(dim=-1).all()):
            raise ValueError("attention mask contains an all-masked query row")

        # (B,H,Tq,Dh) @ (B,H,Dh,Tk) -> one score per query/key pair.
        scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
        scores = scores.masked_fill(~allowed, torch.finfo(scores.dtype).min)
        weights = torch.softmax(scores, dim=-1)
        weights = self.attention_dropout(weights)
        mixed_values = torch.matmul(weights, v)  # (B, H, Tq, Dh)
        output = self.out_proj(self.merge_heads(mixed_values))
        output = self.output_dropout(output)
        if query_padding_mask is not None:
            output = output * query_padding_mask.unsqueeze(-1).to(output.dtype)

        if use_cache:
            return output, KVCache(k, v, valid_keys)
        return output

Further reading

Primary references: Attention Is All You Need, RoFormer, and PyTorch attention API and mask conventions. This implementation is a small teaching model, not a reproduction of a pretrained model.