Multi-Head Attention in Large Language Models

Matrix multiplications help implement attention computations more efficiently and compactly by replacing nested for loops.

Programming

The Problem Attention Creates for Loops

Attention mechanisms work by computing relationships between many items at once. For a sentence with several tokens, the attention computation has to be carried out for every token against every other token. Multi-head attention adds one more layer of repetition: when several attention heads are involved, the same computation is repeated for each head.

Written as a set of nested for loops, this work becomes long, hard to read, and slow to run, because each loop handles only one small piece of the collection at a time. The source for this topic makes a specific point: matrix multiplications help us implement these computations more efficiently and compactly by replacing nested for loops. The rest of this article traces how that replacement works, step by step.

Watching a Matrix Product Fill a Grid

Start with a tiny example. Suppose a sentence has three tokens, and each token is represented by a short list of two numbers. Call the query vectors q1, q2, and q3, and the key vectors k1, k2, and k3. Stack the query vectors as the rows of a matrix Q, and stack the key vectors as the rows of a matrix K. This example is generated for illustration; the numbers are not taken from the source.

python

Q = [[1, 0],
     [0, 1],
     [1, 1]]

K = [[1, 0],
     [0, 1],
     [1, 1]]

Q = [[1, 0],
     [0, 1],
     [1, 1]]

K = [[1, 0],
     [0, 1],
     [1, 1]]

Notes on the code

Each row is one token's two-number vector
Row 1 is token 1, row 2 is token 2, row 3 is token 3
K has the same rows here, so each query is compared with every key, including itself

The attention scores need a comparison of every query with every key. Computing that comparison means taking the dot product of each query row with each key row. A dot product multiplies matching entries and adds the results. Doing this for every pair produces a three-by-three grid, because there are three queries and three keys. The grid is what a single matrix multiplication produces when you multiply Q by the transpose of K, written as $S = Q K^{\top}$.

Which query and key pair produces each cell, and what value lands there, when one matrix multiplication replaces the nested loops?

Worked example

Filling the score grid one row at a time

Using the Q and K matrices above, compute the full three-by-three score matrix S = Q K^{^T}, where each entry is the dot product of one query row with one key row.

1
Row 1 of the grid (query q1 = (1, 0))
Dot product with k1 = (1, 0) gives 1 times 1 plus 0 times 0, which equals 1. Dot product with k2 = (0, 1) gives 0. Dot product with k3 = (1, 1) gives 1.
2
Row 2 of the grid (query q2 = (0, 1))
Dot product with k1 gives 0. Dot product with k2 gives 1. Dot product with k3 gives 1.
3
Row 3 of the grid (query q3 = (1, 1))
Dot product with k1 gives 1. Dot product with k2 gives 1. Dot product with k3 gives 2.
4
Assemble the output
The three rows together form the output matrix S, with nine entries. Every query has been compared with every key, and the result sits in the position that matches its row and column.

Result

S = [[1, 0, 1], [0, 1, 1], [1, 1, 2]]. Each row corresponds to one query token and each column to one key token.

Notice what the nested-loop version would have done. An outer loop would pick a query, an inner loop would pick a key, and the body would compute one dot product and store it. That is nine visits to the body for nine pairs. The matrix multiplication describes the whole grid as one operation. The program states the collection-level computation once, and the individual combinations are all covered inside that single step.

Matrix Multiplication as One Operation

A matrix is a rectangular collection of numbers arranged in rows and columns. When two matrices are multiplied, the result is another matrix whose entries are built from the rows of the first matrix and the columns of the second. The important practical consequence is that a single matrix multiplication performs many of the individual combinations that a nested loop would otherwise perform one at a time.

Tracing the state makes this concrete. Before the multiplication, the program holds two collections: the input values, and the values the multiplication will combine them with. After the multiplication, the program holds one new collection, the output. The output is computed from the input, so nothing is lost, but the work of combining the pieces has been packaged into one operation rather than spread across loop iterations.

python

# Nested-loop version: visits each (query, key) pair in turn
scores = [[0, 0, 0], [0, 0, 0], [0, 0, 0]]
for i in range(3):
    for j in range(3):
        total = 0
        for d in range(2):
            total += Q[i][d] * K[j][d]
        scores[i][j] = total

# Nested-loop version: visits each (query, key) pair in turn
scores = [[0, 0, 0], [0, 0, 0], [0, 0, 0]]
for i in range(3):
    for j in range(3):
        total = 0
        for d in range(2):
            total += Q[i][d] * K[j][d]
        scores[i][j] = total

Output (expected)

scores = [[1, 0, 1], [0, 1, 1], [1, 1, 2]]

Notes on the code

The outer loop chooses a query token i
The middle loop chooses a key token j
The innermost loop performs the dot product over the two vector entries
The result matches the matrix multiplication S = Q K^{^T} shown earlier

This code is a generated illustration of the idea, not a literal implementation from the source. Its purpose is to show what the single matrix multiplication replaces: three nested levels of iteration, where the matrix version states the whole grid computation as one operation. The source does not give code for the batched implementation, so the nested-loop version is the best concrete picture available here.

Batching the Heads Together

Now add heads. In a multi-head setting, several attention computations run side by side. A naive implementation would build and run each head separately: prepare the work for head one, run it, prepare the work for head two, run it, and then combine the results with the others. The source states that a more efficient way of creating multi-head attention modules involves batched matrix multiplications.

How does the work for each head get run in the separate-head approach, and how does one batched matrix operation cover all of the heads at once?

In the batched version, the heads are arranged so that one matrix operation covers all of them. The learning goal is the same in both cases: multi-head attention is produced. The batched version expresses the repeated work as one operation instead of several. Batching therefore reorganizes how the computation is carried out. It does not change the purpose of having multiple heads, which is still to compute several attention computations.

The source does not specify the exact tensor shapes used in batched attention. So when you study a specific library's implementation, look for the place where the head dimension is included in the matrix operation. That is the concrete sign that batching is happening. Finding that place in code is a practical way to check whether a head-by-head loop has been replaced by a batched call.

Batch, Head, and Sequence Dimensions

Three different dimensions are easy to confuse when batching is discussed. The table below gives a general description of each one, so you can tell which of them batching actually operates over. These definitions are a generated summary for teaching purposes; the source itself focuses on the head-level batching idea.

Dimension	What it indexes	Example in this article
Sequence dimension	The token positions within one sentence	Tokens 1, 2, and 3 along the rows and columns of the score grid
Head dimension	Which attention head a computation belongs to	Head 1 and Head 2, each with its own query and key projections
Batch dimension	Separate inputs processed together	Two different sentences processed in the same call

Three dimensions that are often mixed up when discussing batching

Common mistake
Assuming batched matrix multiplication is the only way to compute attention
Why it fails
The source states that matrix multiplications are not strictly required. They are an efficient and compact choice.
Do this instead
Treat matrix multiplication as the preferred practical method, and remember that the same attention computation could be written in other ways.
Common mistake
Thinking batching changes how many heads the model has
Why it fails
Batching reorganizes how multi-head attention is computed. It does not change the purpose of having multiple heads.
Do this instead
Keep the head count as a design choice, and treat batching as a separate choice about how the heads' work is executed.
Common mistake
Treating the head dimension and the batch dimension as the same thing
What it looks like
A batch can hold several separate sentences, while the head dimension separates attention heads within one computation.
Why it fails
Both can appear as extra dimensions in an array, so they look alike when reading code, but they index different things.
Do this instead
Check which index a dimension represents before deciding whether a call is batching over heads, over inputs, or over tokens.
Common mistake
Expecting a loop over tokens to disappear entirely because of matrix multiplication
Why it fails
The matrix multiplication covers every query-key pair in one step, but the sequence dimension is still the structure the multiplication describes, not something it removes.
Do this instead
Read the matrix multiplication as covering all token pairs at once, while remembering the sequence positions are still encoded in the rows and columns.

Predicting a Single Score

EASY

Using the same query and key matrices from the worked example, suppose a new key k4 = (2, 0) is added to the key matrix. Before computing anything, predict whether the score for query q2 = (0, 1) against k4 will be larger or smaller than its score against k2. Then compute the dot product and check your prediction.

Hints

A dot product multiplies matching entries and adds them. Write out the two entries of q2 and the two entries of k4 first.
Compare the second entry of q2 with the second entry of k4, and the first entry of q2 with the first entry of k4.
The answer is the sum of the two products: 0 times 2 plus 1 times 0.

Your attempt (not saved)

Key Takeaways

Attention computes relationships between every token and every other token, and multi-head attention repeats that work for each head. Written as nested for loops, this becomes long, hard to read, and slow.
Matrix multiplication replaces the nested loop structure with one operation that covers the whole collection, such as the full grid of query-key scores in $S = Q K^{\top}$.
Matrix multiplication is an efficient and compact choice, not a strict requirement, as the source explicitly notes.
Batched matrix multiplication handles the work for all heads together in one operation, which is a more efficient way to build multi-head attention than running each head separately.
Batching reorganizes how the computation runs; it does not change the purpose of having multiple heads. Keep the batch, head, and sequence dimensions distinct when reading batched code.

Related concepts

More concepts starting with M