Multi-Head Attention is an attention mechanism used in Transformer models. It performs attention through multiple parallel heads, allowing the model to capture different relationships and patterns in the input sequence.
Scaled Dot-Product Attention
In the Transformer, self-attention is implemented using scaled dot-product attention. It uses three matrices: Query (Q), Key (K), and Value (V), to calculate how strongly each token should attend to other tokens.
- Query (Q): Represents what each token is looking for.
- Key (K): Represents the features used to determine how relevant each token is to a query.
- Value (V): Contains the information that is combined according to the attention weights.
- MatMul (Matrix Multiplication): Multiplies matrices to compute attention scores and combine values.
- Scale: Divides the attention scores by the square root of the key dimension to stabilize training.
- Mask (Optional): Hides certain positions from attention, such as future tokens in causal attention.
- SoftMax: Converts attention scores into probabilities that indicate how much focus each token receives.
The self-attention is computed as:
\text{Attention}(Q, K, V) = \text{softmax} \left( \frac{QK^T}{\sqrt{d_k}} \right) V
where:
- Q = Query matrix, K = Key matrix , V = Value matrix
{d_K} = dimension of the key vectors
Drawback
- A single attention operation produces one attention pattern over the input sequence. However, different relationships may be important at the same time. For example, one pattern may focus on nearby words while another may capture relationships between words that are farther apart.
- Multi-Head Attention addresses this by using multiple attention heads. Each head learns separate projections of the queries, keys and values and performs attention independently. The outputs are then combined to capture information from different representation subspaces.
Multi-Head Attention Mechanism
Multi-head attention extends self-attention by splitting the input into multiple heads, enabling the model to capture diverse relationships and patterns. Instead of using a single set of Q, K, V matrices, the input embeddings are projected into multiple sets (heads), each with its own Q, K, V:
1. Linear Projections: The input X is projected into multiple smaller-dimensional subspaces using different weight matrices.
Q_i = XW_i^Q, \quad K_i = XW_i^K, \quad V_i = XW_i^V
where i denotes the head index.
2. Attention in Each Head: Each head independently computes its own self-attention using the scaled dot-product formula.
3. Concatenating the Head Outputs: The outputs from all heads are concatenated.
4. Final Linear Projection: A final weight matrix is applied to transform the concatenated output into the desired dimension.
If the model dimension is d_model and there are h heads, each head typically operates on a smaller dimension d_model / h.
Mathematically, multi-head attention is expressed as:
\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \text{head}_2, \dots, \text{head}_h) W^O
Where:
h= number of attention heads{W_i^Q} = query projection matrix for head i{W_i^K} = key projection matrix for head i{W_i^V} = value projection matrix for head iW^O = final output projection matrix
Multi-Head Attention in the Transformer
1. Encoder Self-Attention: This allows the encoder to learn contextual relationships within the input sequence. In the original Transformer encoder, each position can attend to all positions in the input sequence.
2. Decoder Self-Attention: In the decoder, masked self-attention prevents a position from attending to future positions, allowing autoregressive generation.
3. Cross-Attention: This layer lets the decoder attend over the encoder's output. It helps the decoder to align and focus on the appropriate input tokens when generating each output token, enabling sequence-to-sequence tasks like translation.
Note: Multi-Head Attention is the mechanism that performs multiple attention operations in parallel. It is different from Sliding Window Attention, which restricts which tokens can attend to each other.
Implementing using PyTorch
Step 1: Imports
Importing all necessary libraries for tensor manipulations and neural network building.
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
Step 2: Scaled Dot-Product Attention Function
- Self-attention results in values i.e the weighted sum for each position and head.
- Softmax ensures the attention weights sum to 1.
- When a mask is used, disallowed positions receive large negative values before softmax, making their attention weights approximately zero.
def scaled_dot_product(q, k, v, mask=None):
d_k = q.size()[-1]
scaled = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(d_k)
if mask is not None:
scaled += mask
attention = F.softmax(scaled, dim=-1)
values = torch.matmul(attention, v)
return values, attention
Step 3: Multi-Head Attention Class
The implementation below demonstrates the core Multi-Head Attention computation: projecting Q, K and V, splitting them into heads, computing attention, combining the heads and applying the final projection.
class MultiheadAttention(nn.Module):
def __init__(self, input_dim, d_model, num_heads):
super().__init__()
self.input_dim = input_dim
self.d_model = d_model
self.num_heads = num_heads
self.head_dim = d_model // num_heads
self.qkv_layer = nn.Linear(input_dim, 3 * d_model)
self.linear_layer = nn.Linear(d_model, d_model)
def forward(self, x, mask=None):
batch_size, sequence_length, input_dim = x.size()
print(f"x.size(): {x.size()}")
qkv = self.qkv_layer(x)
print(f"qkv.size(): {qkv.size()}")
qkv = qkv.reshape(batch_size, sequence_length, self.num_heads, 3 * self.head_dim)
print(f"qkv.size(): {qkv.size()}")
qkv = qkv.permute(0, 2, 1, 3)
print(f"qkv.size(): {qkv.size()}")
q, k, v = qkv.chunk(3, dim=-1)
print(f"q size: {q.size()}, k size: {k.size()}, v size: {v.size()}")
values, attention = scaled_dot_product(q, k, v, mask)
print(f"values.size(): {values.size()}, attention.size: {attention.size()}")
values = values.permute(0, 2, 1, 3)
values = values.reshape(batch_size, sequence_length, self.num_heads * self.head_dim)
print(f"values.size(): {values.size()}")
out = self.linear_layer(values)
print(f"out.size(): {out.size()}")
return out
Step 4: Example: Running Multi-Head Attention
input_dim = 1024
d_model = 512
num_heads = 8
batch_size = 30
sequence_length = 5
x = torch.randn((batch_size, sequence_length, input_dim))
model = MultiheadAttention(input_dim, d_model, num_heads)
output = model.forward(x)
Output:
x.size(): torch.Size([30, 5, 1024])
qkv.size(): torch.Size([30, 5, 1536])
qkv.size(): torch.Size([30, 5, 8, 192])
qkv.size(): torch.Size([30, 8, 5, 192])
q size: torch.Size([30, 8, 5, 64]), k size: torch.Size([30, 8, 5, 64]), v size: torch.Size([30, 8, 5, 64])
values.size(): torch.Size([30, 8, 5, 64]), attention.size: torch.Size([30, 8, 5, 5]) values.size(): torch.Size([30, 5, 512])
out.size(): torch.Size([30, 5, 512])
You can download the complete source code from here.
Applications
- Machine translation: Helps relate words across the source and target sequences.
- Text summarization: Helps identify important relationships between words and sentences.
- Question answering: Helps models relate questions to relevant parts of the input text.
- Vision Transformers: Used to model relationships between image patches.
- Speech recognition: Helps capture relationships between different parts of speech sequences.
Limitations
- Computational cost: Standard full attention still has quadratic complexity with respect to sequence length.
- Memory usage: Attention matrices can require substantial memory for long sequences.
- Additional parameters: Multiple projection matrices increase the model's parameter count compared with a single attention operation.
- Redundant heads: Different heads may sometimes learn similar attention patterns, reducing the benefit of having many separate heads.
