what is the enhancement of multi-head attention compared to self-attention

Theoretical Background:

Self-attention is a mechanism where each element of an input sequence pays attention to every other element and computes a weighted average of all elements. This allows the model to capture dependencies regardless of their distance in the input sequence. However, using a single attention mechanism can limit the model's ability to capture diverse patterns.

Multi-head attention, introduced in the Transformer model, extends this by using multiple attention mechanisms (or 'heads'). Each head learns different attention patterns, enabling a more nuanced understanding of the input sequence. Formally, for an input sequence $X$, each attention head computes an output as follows:

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

Where $Q$, $K$, and $V$ are the query, key, and value matrices, respectively, and $d_k$ is the dimensionality of the key vectors. Multi-head attention applies this mechanism multiple times in parallel and concatenates the results:

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

Practical Applications:

The enhanced representation from multi-head attention is crucial for tasks like machine translation, text summarization, and other NLP tasks where understanding context and dependencies is essential. By capturing different aspects of the data through multiple heads, models can achieve state-of-the-art performance in these applications.

Code Example:

In practice, multi-head attention is implemented in deep learning libraries like TensorFlow and PyTorch using built-in functions. For example:

import torch
from torch.nn import MultiheadAttention

# Define a multi-head attention module
multihead_attn = MultiheadAttention(embed_dim=512, num_heads=8)

# Dummy input
query = torch.rand(10, 32, 512)  # (sequence_length, batch_size, embed_dim)
key = torch.rand(10, 32, 512)
value = torch.rand(10, 32, 512)

# Apply multi-head attention
attn_output, attn_output_weights = multihead_attn(query, key, value)

External References:

• Attention is All You Need: The original paper introducing Transformers and multi-head attention.
• The Illustrated Transformer: A visual and intuitive explanation of Transformers and multi-head attention.

Diagram:

graph TD;
    A(Input Sequence) --> B[Split into Heads];
    B --> C[Head 1];
    B --> D[Head 2];
    B --> E[Head h];
    C --> F[Concatenate Heads];
    D --> F;
    E --> F;
    F --> G[Linear Transformation];
    G --> H[Output];

The diagram above illustrates how an input sequence is processed through multiple heads in a multi-head attention mechanism, each head contributing to a richer final output.