Multi-Head Attention: The Power of Multiple Perspectives in LLMs

If there's one mechanism that truly defines the revolutionary power of Large Language Models (LLMs) and the Transformer architecture they're built upon, it's attention. And within the attention mechanism, the Multi-Head Attention (MHA) operator is the star player. Instead of a single focus, MHA allows the model to simultaneously look at different aspects of the input sequence, leading to a much richer and more nuanced understanding of language.

Multi-Head Attention: The Power of Multiple Perspectives in LLMs

Revisiting Self-Attention: The Foundation

Before diving into MHA, let's quickly recall self-attention. For each word (or token) in an input sequence, self-attention calculates how relevant every other word in that same sequence is to it. This is done by computing three learned vectors for each token:

• Query (Q): Represents what the current word is "looking for."

• Key (K): Represents what each other word "offers."

• Value (V): Contains the actual information of each word.

The attention score between a Query and a Key (typically via a dot product) determines their similarity. These scores are then normalized (using softmax) to get attention weights, which are used to take a weighted sum of the Value vectors. This weighted sum becomes the new, context-aware representation of the Query word.

The output of a single self-attention "head" would look like this:

Attention(Q,K,V)=softmax(dk​​QKT​)V

Where dk​ is the dimension of the Key vectors, used for scaling to prevent very large dot products from pushing softmax into saturation.

The Problem with Single-Headed Attention

While powerful, a single self-attention mechanism has a limitation: it can only learn one type of relationship or focus on one aspect of the data at a time. Language is complex; words can have multiple meanings, participate in different grammatical roles, and form various semantic relationships. A single "eye" might miss crucial nuances.

For example, in the sentence "The bank had strong currents," one attention head might focus on "bank" as a financial institution, while another might interpret it as a river bank. A single head would struggle to capture both simultaneously.

The Solution: Multi-Head Attention

Multi-Head Attention addresses this limitation by running multiple self-attention operations in parallel, each with its own independent set of learned Query, Key, and Value projection matrices.

Here's how it works:

1. Linear Projections: The input embeddings (representing the words) are first linearly transformed h times, where h is the number of "heads." Each transformation uses a different set of learned weight matrices (WiQ​,WiK​,WiV​) for Query, Key, and Value, creating h distinct sets of Qi​,Ki​,Vi​ for each head i.

   • Qi​=XWiQ​

   • Ki​=XWiK​

   • Vi​=XWiV​

     (where X is the input embedding matrix, and W are the projection matrices)

2. Parallel Attention Computation: Each of these h sets (Qi​,Ki​,Vi​) is then fed into an independent scaled dot-product attention function. Each "head" computes its own attention output, focusing on different relationships or "representation subspaces" of the input.

   • headi​=Attention(Qi​,Ki​,Vi​)

3. Concatenation: The outputs from all h attention heads are then concatenated back together into a single, larger matrix.

4. Final Linear Transformation: This concatenated output undergoes one last linear transformation (with a learned weight matrix WO) to project it back into the desired output dimension, which matches the model's hidden size for subsequent layers.

   • MultiHead(Q,K,V)=Concat(head1​,…,headh​)WO

Why Multiple Heads are Better than One

Multi-Head Attention offers several significant advantages for LLMs:

1. Diverse Relationship Capture: Each head can learn to focus on different types of relationships. For instance, one head might specialize in capturing syntactic dependencies (like subject-verb agreement), another on semantic relationships (e.g., "bank" as a financial institution vs. river bank), and yet another on coreference resolution (e.g., linking pronouns to their antecedents). This provides a comprehensive understanding of the input.

2. Enhanced Representational Capacity: By combining multiple perspectives, the model can create richer and more nuanced representations of each token, leading to a deeper understanding of the entire sequence.

3. Improved Robustness: Having multiple heads means the model's understanding doesn't rely on a single attention pattern. If one head occasionally misinterprets a relationship, other heads can compensate, making the model more robust to noise and ambiguities.

4. Parallel Computation: The independent computations for each head can be performed in parallel, making the MHA operation computationally efficient despite its apparent complexity.

The Multi-Head Attention operator is the ingenious mechanism that allows LLMs to process information from various angles simultaneously, much like how a human might consider multiple interpretations of a sentence. It's a cornerstone of the Transformer architecture, enabling LLMs to grasp the intricate nuances of human language and deliver their impressive performance across a wide array of tasks.