Combine attention heads

Learn it

Words to know

attention head

One attention calculation. A layer runs several side by side, each with its own queries, keys and values.

attention

The step in which each token gathers information from other tokens, taking more from those that match what it is looking for.

tensor

PyTorch's block of numbers. It can be one number, a list, a table, or have more dimensions.

How it works

1. The inputs have shape [heads, time, width]: for each head, one row per token.
2. torch.nn.functional.scaled_dot_product_attention(q, k, v) runs attention on the last two dimensions. The leading dimension, the heads, is handled as a batch: each head is computed on its own.
3. The result has shape [heads, time, width]. Swap the first two dimensions with .transpose(0, 1) to get [time, heads, width].
4. .reshape(time, -1) joins the heads of each token into one row of length heads * width. The -1 asks PyTorch to work out that length.

Worked example

Input
Q = [[[0], [0]], [[1], [1]]]
K = [[[0], [0]], [[0], [5]]]
V = [[[2], [4]], [[10], [30]]]

Output
[[3, 29.8661], [3, 29.8661]]

Head 0 looks at both tokens equally. Head 1 looks almost only at the second token. Each token's final row holds both views, side by side.

Why it matters

The workshop's GPT uses 4 heads, and large models use dozens. The KV cache and grouped-query lessons in chapter 7 are about making all those heads cheaper to run.

Your turn

Using PyTorch, run attention separately for each head and join the results for each token.

Q

queries of shape [heads, time, width]

K

keys of the same shape

V

values of the same shape

Returns

A nested Python list of shape [time, heads * width]: for each token, the output of head 0, then head 1, and so on. No mask is used.

This lesson needs PyTorch, which does not run in the browser. The app gives you the exercise as a file with its tests, to run on your own computer.

Input
Q = [[[0], [0]], [[1], [1]]]
K = [[[0], [0]], [[0], [5]]]
V = [[[2], [4]], [[10], [30]]]

Output
[[3, 29.8661], [3, 29.8661]]

Input
Q = [[[0]]]
K = [[[0]]]
V = [[[5]]]

Output
[[5]]

Input
Q = [[[0, 0], [0, 0]], [[2, 0], [2, 0]]]
K = [[[0, 0], [0, 0]], [[0, 0], [2, 0]]]
V = [[[1, 2], [3, 4]], [[10, 20], [30, 40]]]

Output
[[2, 3, 28.8839, 38.8839], [2, 3, 28.8839, 38.8839]]

Input
Q = [[[0]], [[0]]]
K = [[[0]], [[0]]]
V = [[[2]], [[3]]]

Output
[[2, 3]]

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Where to start

   PyTorch already has the attention function. Your work is to call it and then put the dimensions in the right order.

2. Hint 2: The key step

   The function returns [heads, time, width]. You want each token's heads next to each other, so move time to the front before you flatten.

3. Hint 3: The outline

   Make float64 tensors q, k, v. Call torch.nn.functional.scaled_dot_product_attention(q, k, v). Then .transpose(0, 1).reshape(q.shape[1], -1).tolist().

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Step by step

1. Attention for every head at once: [heads, time, width]. The function works on the last two dimensions and treats anything in front as a batch, so no head sees another head's keys or values.
2. Put time first: [time, heads, width]. Reshape reads numbers in order. With time first, the numbers that end up next to each other belong to the same token.
3. Join each token's heads into one row: [time, heads * width]. -1 lets PyTorch work out the second size from the number of elements.

The same numbered steps are the comments in the solution.

import torch


def solve(Q, K, V):
    q = torch.tensor(Q, dtype=torch.float64)
    k = torch.tensor(K, dtype=torch.float64)
    v = torch.tensor(V, dtype=torch.float64)
    # 1. Attention for every head at once: [heads, time, width]
    out = torch.nn.functional.scaled_dot_product_attention(q, k, v)
    # 2. Put time first: [time, heads, width]
    out = out.transpose(0, 1)
    # 3. Join each token's heads into one row: [time, heads * width]
    time = q.shape[1]
    return out.reshape(time, -1).tolist()

Watch out

• Reshaping without the transpose gives the right shape with the wrong contents: the numbers stay in head-by-head order, so a row mixes different tokens instead of holding the heads of one token.
• Do not pass is_causal=True here. This exercise asks for attention with no mask.

Check yourself