Multi-Head Attention Layer

Hello! Welcome back to our deep dive into the mechanisms that power modern AI.

In our last lesson, we derived and implemented scaled dot-product self-attention. We established the Query, Key, and Value (Q, K, V) framework and understood the critical role of scaling to ensure stable training. This gave us a powerful tool for a sequence to "look" at itself and build context-rich representations.

However, a single attention mechanism might be forced to average its focus, potentially missing out on diverse and subtle relationships within the text. For example, in the sentence "The tired robot dropped the heavy gear because it broke," we might want to simultaneously understand that "it" refers to "gear" (co-reference) and that "broke" is linked to "robot dropped" (causal relationship). A single attention head might struggle to capture both.

This brings us to today's topic: Multi-Head Attention. Your learning outcome for this lesson is to build a multi-head attention layer, the core component of the Transformer. We will see how running multiple attention mechanisms in parallel allows the model to capture a richer variety of linguistic features and dependencies.

1. The Rationale: From Single Head to Many

The core idea behind multi-head attention is to run the scaled dot-product attention process multiple times in parallel with different, learned linear projections. This allows the model to jointly attend to information from different "representation subspaces" at different positions.

You can think of it as giving the model multiple "perspectives" on the input sequence. Each "head" can learn to focus on a different type of relationship:

• One head might track syntactic dependencies, like subject-verb agreement.
• Another might track co-reference, like which pronoun refers to which noun.
• A third might focus on positional information, like which words are nearby.

By combining the outputs of these heads, the model can construct a far more nuanced and comprehensive representation of each word in the sequence.

2. The Architecture of Multi-Head Attention

Let's break down the architecture. At a high level, the process is:

1. Project: For each of the h heads, create separate, learned weight matrices to project the input Queries, Keys, and Values into a lower-dimensional space.
2. Attend in Parallel: Perform scaled dot-product attention for each head simultaneously, using its projected Q, K, and V. This yields h separate output matrices.
3. Concatenate & Combine: Concatenate the h output matrices back together and pass them through a final, learned linear projection to produce the layer's final output.

This architecture is elegantly captured in the following diagram.

The mathematical formulation from the original "Attention Is All You Need" paper is as follows:
$\mathrm{MultiHead}(Q,K,V)=\mathrm{Concat}(head\_1,...,head\_h)W^O$
$\text{where}head\_i=\mathrm{Attention}(Q{W}^Q\_i,K{W}^K\_i,V{W}^V\_i)$

Let's dissect this:

• The inputs are the full-dimensional Q, K, V matrices (e.g., dimension d_model = 512).
• For each head i, we have unique weight matrices $W_i^Q,W_i^K,W_i^V$. These project the inputs down to the head's dimension (d_k, e.g., 64).
• Attention(...) is the scaled dot-product attention function we learned in the last lesson.
• The results from all heads are concatenated.
• A final weight matrix, $W^O$, is applied to combine the information from all heads and project the result back to the original d_model dimension.

To build a strong visual intuition for this process, the following resource is excellent.

3. Implementation in PyTorch

Now, let's translate this architecture into code. Your background in Python and software engineering will be valuable here as we construct a MultiHeadAttention module. The key challenge lies in efficiently managing the tensor shapes as we split into heads and then merge back.

The most efficient implementation doesn't create h separate linear layers. Instead, it uses single, larger linear layers and then reshapes the tensors. The process for the forward pass is:

1. Define linear layers to map the input d_model to d_model for Q, K, and V.
2. Pass the input query, key, and value tensors through these layers.
3. Reshape the resulting (batch_size, seq_len, d_model) tensors into (batch_size, seq_len, num_heads, head_dim).
4. Permute the dimensions to (batch_size, num_heads, seq_len, head_dim) to facilitate batch matrix multiplication across all heads at once.
5. Apply the scaled dot-product attention formula using the reshaped tensors.
6. Permute and reshape the output back to (batch_size, seq_len, d_model).
7. Pass this through a final linear output layer.

The following video provides a complete, line-by-line implementation of this process. It's an excellent resource for seeing exactly how the tensor manipulations work in PyTorch.

While the CodeEmporium video uses a combined linear layer for Q, K, and V, an alternative and very common approach is to use three separate linear layers. Aladdin Persson's video on building a Transformer from scratch demonstrates this, and also introduces the torch.einsum function—a powerful tool for expressing complex tensor operations that might appeal to your programming background.

Conclusion

Congratulations! You have now constructed the single most important component of the Transformer architecture. By understanding how to split the model's representation into parallel "heads," you've unlocked the mechanism that allows Transformers to capture rich and diverse relationships within data.

Key Takeaways:

• Multi-Head Attention improves on single-head attention by allowing the model to focus on different information from different representational subspaces simultaneously.
• The implementation involves projecting the Q, K, V inputs for each head, running scaled dot-product attention in parallel, and then concatenating and projecting the results.
• Efficient implementation relies on clever tensor reshaping and permutation to perform calculations for all heads at once using batched matrix multiplications.

Preview of the Next Lesson:
We are now ready to assemble a full layer of the Transformer. In the next lesson, we will take the MultiHeadAttention layer we've just built and wrap it with the other essential components: residual connections and layer normalization, followed by a position-wise feed-forward network. This will complete our construction of a Transformer encoder block.