Multi-Head Attention

Multi-Head Attention is a pivotal technique that underpins the prowess of Transformer architectures. This mechanism amplifies the self-attention capability, enabling the model to simultaneously focus on diverse segments of the input sequence, thereby capturing a richer array of dependencies.

The Concept of Multi-Head Attention

The core idea behind Multi-Head Attention is straightforward: instead of relying on a single attention mechanism, why not employ multiple attention mechanisms operating in parallel? Each of these attention mechanisms, termed as "heads," can concentrate on distinct facets of the input data. By doing so, the model can capture diverse linguistic phenomena that single-headed attention might overlook.

In a typical Transformer network, the input data is transformed into three key matrices: Query (Q), Key (K), and Value (V). These matrices are essential for computing attention scores. The multi-head attention mechanism allows the model to project Q, K, and V into multiple subspaces, each corresponding to a different attention head. This is achieved by applying different sets of learned linear transformations to the input data.

Illustration of the Multi-Head Attention mechanism, where the input Query (Q), Key (K), and Value (V) matrices are projected into multiple subspaces, each corresponding to a different attention head. The outputs from all heads are then concatenated to produce the final output.

Mathematical Formulation

To delve deeper, consider that for a given input, multi-head attention performs the following operations:

1. Linear Transformation: For each head, the input is linearly transformed into different Q, K, and V matrices using learned weight matrices, often denoted as $W_i^Q$, $W_i^K$, and $W_i^V$.

2. Scaled Dot-Product Attention: Each head computes attention scores using the scaled dot-product attention mechanism. For a single head, the attention output is computed as:

   $$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

   Here, $d_k$ is the dimensionality of the key vectors, which serves as a scaling factor to prevent large dot-product values that could push the softmax function into regions with extremely small gradients.

3. Concatenation and Linear Transformation: The outputs from all heads are concatenated and then linearly transformed to produce the final output of the multi-head attention layer:

   $$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h)W^O$$

   Here, $W^O$ is the learned weight matrix that combines the outputs of all attention heads.

Code Implementation

To better understand how multi-head attention is implemented, consider the following simplified code snippet using PyTorch, a popular deep learning library:

import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads
        self.d_model = d_model
        self.depth = d_model // num_heads

        self.wq = nn.Linear(d_model, d_model)
        self.wk = nn.Linear(d_model, d_model)
        self.wv = nn.Linear(d_model, d_model)
        self.dense = nn.Linear(d_model, d_model)

    def split_heads(self, x, batch_size):
        x = x.view(batch_size, -1, self.num_heads, self.depth)
        return x.transpose(1, 2)

    def forward(self, v, k, q, mask=None):
        batch_size = q.size(0)

        q = self.split_heads(self.wq(q), batch_size)
        k = self.split_heads(self.wk(k), batch_size)
        v = self.split_heads(self.wv(v), batch_size)

        scaled_attention = self.scaled_dot_product_attention(q, k, v, mask)
        scaled_attention = scaled_attention.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)

        return self.dense(scaled_attention)

    def scaled_dot_product_attention(self, q, k, v, mask):
        matmul_qk = torch.matmul(q, k.transpose(-2, -1))
        dk = k.size(-1)
        scaled_attention_logits = matmul_qk / torch.sqrt(torch.tensor(dk, dtype=torch.float32))

        if mask is not None:
            scaled_attention_logits += (mask * -1e9)

        attention_weights = torch.nn.functional.softmax(scaled_attention_logits, dim=-1)
        output = torch.matmul(attention_weights, v)
        
        return output

In this implementation, each input vector undergoes a series of transformations to produce Q, K, and V matrices. The data is split into multiple heads, processed independently, and then recombined. This allows the model to attend to information across different representation subspaces.

Advantages and Applications

The multi-head attention mechanism provides several advantages:

1. Parallelization: By processing multiple attention heads in parallel, the model can simultaneously focus on different parts of the input, making it more efficient and effective.

2. Expressive Power: With multiple heads, the model can capture a wider range of dependencies and complex interactions within the data.

3. Enhanced Feature Representation: Each head contributes a unique perspective, allowing for a more nuanced representation of the input data.

These attributes make multi-head attention indispensable in tasks such as machine translation, text summarization, and beyond, where understanding contextual relationships is paramount.

As we continue to explore the intricacies of Transformer architectures, appreciating the role of multi-head attention is crucial. It is not just an enhancement but a fundamental shift in how models perceive and process sequential data, contributing significantly to the robustness and adaptability of modern AI systems.