Parallel Processing of Attention Heads

The purpose of the wrapper stays the same in both versions: it combines multiple single-head attention modules. What changes is how the computation is organized, not which quantity the layer produces.

Programming

Why the Wrapper Feels Correct but Slow

A MultiHeadAttentionWrapper is built by combining several single-head attention modules. Each of those modules is a complete attention unit that can process an input on its own. The wrapper gathers them so that one call produces the output of many heads. This article looks at one feature of how that wrapper reaches its result: its heads are evaluated one after another. Understanding that sequential structure is the starting point for the improvement the source proposes, which is to process the heads in parallel.

What do you think happens?

When the wrapper's forward method runs several heads on the same input, do the heads run at the same time or one after another?

All heads run at the same time
Each head runs one after another in turn

Reveal answer

Answer:

Each head runs one after another in turn

The current forward method passes the input through each head in turn using a list comprehension. Every head receives the same input x, but the heads are evaluated sequentially.

Tracing the Sequential Forward Pass

The current wrapper's forward method uses a list comprehension of the form head(x) for head in self.heads. Read step by step, the input x is handed to the first head, that head produces its result, then the input is handed to the next head, and so on until every head has been evaluated. The result is correct, because each head sees the same input and contributes its own output. The important feature is the order of work: one head finishes before the next one begins.

python

class MultiHeadAttentionWrapper:
    def __init__(self, heads):
        self.heads = heads

    def forward(self, x):
        # Each head receives the same input x.
        # The heads are evaluated one after another.
        outputs = [head(x) for head in self.heads]
        return combine(outputs)  # combine is a placeholder for the gathering step

class MultiHeadAttentionWrapper:
    def __init__(self, heads):
        self.heads = heads

    def forward(self, x):
        # Each head receives the same input x.
        # The heads are evaluated one after another.
        outputs = [head(x) for head in self.heads]
        return combine(outputs)  # combine is a placeholder for the gathering step

Notes on the code

The list comprehension evaluates head(x) once per head, in order
combine is a generic placeholder; the source does not give the gathering code

What runs one head after another in the loop version, and how does a single matrix operation change the order of work in the parallel version?

How Limiting the Loop Shapes the Wrapper

Because the wrapper calls each head in turn, its design is organized around separate head objects. Each head is a self-contained module, and the wrapper's job is to call them and gather their results. That structure is clear and correct, but it means the wrapper depends on evaluating every head as its own step. The source proposes a different route: compute the outputs for all attention heads simultaneously by way of matrix multiplication, so that a single operation produces the results for every head together.

Key point

The source describes the parallel version as treating the heads as a single batch of work. In the sequential version, the input passes through one head and then the next. In the parallel version, the heads are handled together in one computation, so the forward pass no longer needs a separate step per head.

Implementing Multi-Head Attention With Weight Splits

The next step in the source moves from the idea of parallel computation to a specific method: implementing multi-head attention with weight splits. The name carries the main point. Instead of keeping a separate, unrelated set of weights for every head as independent objects, the weights are organized so that they can be divided into portions, one portion for each head. Those portions can then be handled together in the matrix multiplication described above.

When one large projection matrix replaces separate per-head matrices, which columns belong to which head, and how do the index ranges map to each head?

The source names the method but does not walk through its details in this section, so treat the slices as an organizing principle rather than a finished implementation. The core idea is that a single set of weights is partitioned by head so that all heads can be computed in one pass.

Worked Example: Splitting Six Columns Across Three Heads

Worked example

Assigning columns of one weight matrix to three heads

A single projection weight matrix has six columns, numbered 0 through 5. It must serve three heads, and each head should own an equal portion. Which columns belong to each head?

1
Count the columns
The matrix has 6 columns in total.
2
Count the heads
There are 3 heads, so each head should receive an equal share of the columns.
3
Divide the columns
6 divided by 3 gives 2 columns per head.
4
Assign index ranges
Head 1 takes columns 0 and 1, head 2 takes columns 2 and 3, and head 3 takes columns 4 and 5.
5
Check coverage
Every column appears in exactly one head's range, so no column is left unused and none is shared.

Result

Head 1 owns columns 0 to 1, head 2 owns columns 2 to 3, and head 3 owns columns 4 to 5. The index ranges map each head to its own portion of the single weight matrix.

Changing the Combined Output Without Changing the Heads

The exercise behind this article asks you to change the combined output of the wrapper so that it returns two-dimensional embedding vectors. The source gives a clear rule for doing this: when modifying the wrapper's output, change how the combined result is computed and keep the role of each head unchanged. The change belongs to the gathering step at the end of the forward pass, not to the definition of what a head does.

Mistakes When Modifying the Wrapper Output

Common mistake
Changing what each head computes when only the output format needs to change
Why it fails
The source says the purpose of each head must stay the same. Altering a head's internals changes the layer rather than reshaping its result.
Do this instead
Leave each head untouched and modify only how the combined result is computed.
Common mistake
Treating the parallel idea as a change in the mathematical quantity produced
Why it fails
The source states that the parallel version changes how the computation is organized, not which quantity the layer produces.
Do this instead
Keep the goal fixed: the wrapper still combines multiple single-head attention modules into one output.
Common mistake
Assuming the weight split is a finished implementation
Why it fails
The source names the weight-split method but does not walk through its details, and the parallel branch in its sketch is shown only as a comment.
Do this instead
Treat the slices as an organizing principle and work out the exact code for your own implementation, verifying the shapes as you go.
Common mistake
Ordering the combined outputs inconsistently, so each result lands in an unexpected place in the final output
Common mistake
Letting the combined output's shape drift from what downstream code expects after the gathering step changes

Practice Tracing the Two Forward Passes

EASY

Take the sequential forward method in the code above and list, in order, what happens to the input x from the moment forward is called until the combined output is returned. Then describe in one sentence how the parallel version would change that list.

Hints

Start with the list comprehension and note which head is evaluated first.
In the parallel version, ask how many separate steps the heads require before the outputs are ready.

Your attempt (not saved)

MEDIUM

A weight matrix has 8 columns and must serve 4 heads with equal shares. Write the column range for each head, then state which step of the forward pass you would change if the combined output needed a different shape.