Several heads, one shared representation.

One attention head gives each query one set of weights over the source. A single weighted mixture may have to combine several useful relationships at once. Multi-head attention allows several learned projections to construct different mixtures before those results are combined.

The key word is learned. We do not assign one head to nouns, one to verbs, and one to position. Different behaviors may emerge, but a head is defined by its projections and computation, not by a guaranteed human-readable role.

1. A head is a projected view

Let the model width be D and choose h heads. Our implementation uses head width \(d_h=D/h\), so D must be divisible by h. Each head receives its own query, key, and value projections with that smaller width.

For one head r:

The query sequence \(X_q\) and source sequence \(X_s\) are equal in self-attention. In cross-attention they come from different streams. This distinction survives unchanged when we add more heads.

The final operation concatenates the head outputs along their feature dimension and applies an output projection:

Concatenation preserves separate features from the heads. Averaging them would be a different design and would discard the head axis before the output projection had a chance to mix it.

2. Why one large projection can represent many heads

We can concatenate the per-head query matrices into one (D, D) matrix and compute a single projection. The result contains all h projected views adjacent along its final feature axis. A reshape and transpose then expose the head axis explicitly.

This does not mean we split the original input features into disjoint groups before learning the projections. Each output head can depend on all D input features through its corresponding columns of the projection matrix.

Our code uses four dense projections: Q, K, V, and output. It omits biases for these attention projections to make the correspondence between the two libraries especially clear. The feed-forward layers introduced next do have biases.

3. Shapes are the implementation

For a batch with B = 2, target length T = 4, source length S = 6, model width D = 8, and h = 2, each head has four features.

The reshape creates a head axis inside the feature layout. The transpose moves that axis before the position axis. If you reshape directly into (B, h, T, dₕ), you generally regroup the numbers incorrectly. A tensor can have the expected shape while holding the wrong arrangement of data.

After attention, transpose back from (B, h, T, dₕ) to (B, T, h, dₕ), then combine the last two axes. PyTorch’s reshape can handle non-contiguous input by copying when necessary. A blind view after a transpose can fail or require an explicit contiguous copy.

4. Implement the complete operation

The class below reuses attention from lesson 3. The NumPy object is a forward-pass implementation. The PyTorch version registers trainable linear layers so their parameters participate in optimization.

class MultiHead:
    def __init__(self, width: int, heads: int, seed: int = 0) -> None:
        if heads <= 0 or width % heads:
            raise ValueError("Model width must be divisible by head count")
        self.heads, self.head_width = heads, width // heads
        rng = np.random.default_rng(seed)
        self.wq, self.wk, self.wv, self.wo = [
            rng.normal(size=(width, width)) / np.sqrt(width) for _ in range(4)
        ]

    def split(self, x: FloatArray) -> FloatArray:
        batch, length, _ = x.shape
        return x.reshape(batch, length, self.heads, self.head_width).transpose(0, 2, 1, 3)

    def __call__(self, query: FloatArray, context: FloatArray, allowed: BoolArray | None = None) -> tuple[FloatArray, FloatArray]:
        q = self.split(query @ self.wq)
        k = self.split(context @ self.wk)
        v = self.split(context @ self.wv)
        heads, weights = attention(q, k, v, allowed)
        batch, _, length, _ = heads.shape
        joined = heads.transpose(0, 2, 1, 3).reshape(batch, length, -1)
        return joined @ self.wo, weights

class MultiHead(nn.Module):
    def __init__(self, width: int, heads: int, seed: int = 0) -> None:
        super().__init__()
        if heads <= 0 or width % heads:
            raise ValueError("Model width must be divisible by head count")
        self.heads, self.head_width = heads, width // heads
        self.wq, self.wk, self.wv, self.wo = [nn.Linear(width, width, bias=False) for _ in range(4)]
        # Shared initialization makes the two libraries directly comparable.
        rng = np.random.default_rng(seed)
        with torch.no_grad():
            for layer in [self.wq, self.wk, self.wv, self.wo]:
                w = rng.normal(size=(width, width)) / np.sqrt(width)
                layer.weight.copy_(torch.tensor(w.T, dtype=layer.weight.dtype))

    def split(self, x: Tensor) -> Tensor:
        batch, length, _ = x.shape
        return x.reshape(batch, length, self.heads, self.head_width).transpose(1, 2)

    def forward(self, query: Tensor, context: Tensor, allowed: Tensor | None = None) -> tuple[Tensor, Tensor]:
        q, k, v = self.split(self.wq(query)), self.split(self.wk(context)), self.split(self.wv(context))
        heads, weights = attention(q, k, v, allowed)
        batch, _, length, _ = heads.shape
        joined = heads.transpose(1, 2).reshape(batch, length, -1)
        return self.wo(joined), weights

Typed Python · Shapes are documented alongside the code. Open the full file for imports and dependencies, or step through a concrete example.

The code initializes both libraries from the same NumPy random stream. PyTorch’s linear layer stores weights in (out_features, in_features) order, while our NumPy multiplication uses (in_features, out_features). That is why initialization copies a transpose.

Try a cross-attention example:

import numpy as np
from numpy_core import MultiHead

rng = np.random.default_rng(4)
query = rng.normal(size=(2, 4, 8))
source = rng.normal(size=(2, 6, 8))
layer = MultiHead(width=8, heads=2)
output, weights = layer(query, source)

assert output.shape == (2, 4, 8)
assert weights.shape == (2, 2, 4, 6)
np.testing.assert_allclose(weights.sum(-1), 1)

This example also confirms that multiple heads do not force source and target lengths to match. Each head still mixes S values into T outputs.

5. What changes when the number of heads changes?

With D fixed, more heads mean narrower per-head representations in this design. The concatenated width remains D. You gain more separate attention distributions but do not simply multiply the full-width model by h.

The score calculation still has pairwise sequence interactions. A conventional materialized weight tensor has B × h × T × S entries. The total Q/K dot-product work scales with the sum of the head widths, while the number of stored attention matrices grows with h. Efficient kernels can avoid storing all those matrices; that changes memory use, not the conceptual calculation.

More heads are not automatically better. Head width, model width, task, training, and implementation efficiency interact. For this course, a small fixed choice lets us inspect the mechanics without turning architecture tuning into a second project.

6. Inspect heads without averaging them away

Our function returns separate weight matrices. Some library APIs average weights across heads by default when returning them for inspection. In PyTorch’s MultiheadAttention ↗ , average_attn_weights=False preserves that axis when need_weights=True.

Averaging for display can hide differences between heads. It also does not reproduce the model’s actual concatenation and output projection. Treat the average as a visualization choice, not as the forward computation.

Carry this forward

Multi-head attention returns a tensor with the model width D. That is exactly the shape we need to add its output back to the input. Next, we will follow that residual path and explain the normalization and feed-forward operations that turn attention into a reusable block.