Multi-Head Attention

Multi-Head Attention in Transformers

In the Transformer architecture, after input embeddings are combined with Positional Encoding, the resulting vectors are fed into the first core component: the multi-head attention mechanism. This mechanism is a significant advancement over simpler attention methods, allowing the model to process information more effectively.

The multi-head attention mechanism is a cornerstone of the Transformer model, crucial for its exceptional performance in handling sequential data, such as text in Natural Language Processing. It enables the model to simultaneously attend to information from different representation subspaces at different positions.

In this chapter, we will explore the architecture of the multi-head attention mechanism, discuss its advantages, and provide a conceptual Python implementation using NumPy to illustrate its workings.

Understanding Self-Attention (Scaled Dot-Product Attention)

Before diving into multi-head attention, it’s essential to understand its core component: self-attention, often referred to as scaled dot-product attention. This mechanism allows a model to weigh the importance of different tokens within an input sequence relative to each other when producing a representation for each token.

The Mechanics of Self-Attention

The self-attention process involves the following key steps:

• Create Query, Key, and Value Vectors: For each token in the input sequence (e.g., word embeddings), three vectors are generated by multiplying the token’s embedding by learned weight matrices: a Query vector (Q), a Key vector (K), and a Value vector (V). These vectors have dimensions d_q, d_k, and d_v respectively (often d_q = d_k).
• Calculate Attention Scores: Attention scores are computed by taking the dot product of the Query vector of a specific token with the Key vectors of all tokens in the sequence (including itself). This score determines how much focus to place on other parts of the input sequence when processing the current token.
• Scaling Scores: The scores are then scaled by dividing by the square root of the dimension of the key vectors (√d_k). This scaling prevents the dot products from growing too large, which could lead to vanishing gradients in the softmax function.
• Apply Softmax Function: A softmax function is applied to the scaled scores to obtain attention weights. These weights are positive and sum to 1, representing a probability distribution of attention over the input sequence.
• Compute Weighted Sum of Values: Finally, the output for each token is a weighted sum of the Value vectors of all tokens in the sequence, where the weights are the attention weights just computed.

Mathematically, scaled dot-product attention is defined as:

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

Where Q, K, and V are matrices packing together the query, key, and value vectors respectively.

What is Multi-Head Attention?

Multi-Head Attention enhances the self-attention mechanism by performing multiple self-attention operations in parallel, each called an “attention head.” Instead of a single set of Query, Key, and Value weight matrices, multi-head attention uses multiple sets, allowing the model to jointly attend to information from different representational subspaces at different positions.

This can be visualized as the model learning to focus on different aspects of the input sequence simultaneously. For example, one head might focus on syntactic relationships, while another focuses on semantic similarity.

Architecturally, the multi-head attention layer consists of several parallel attention heads. Each head independently performs scaled dot-product attention. The outputs of these heads are then concatenated and passed through a final linear layer to produce the output of the multi-head attention block. This structure is a key part of encoder and decoder layers in a Transformer.

Steps in Multi-Head Attention

The process for multi-head attention involves these steps:

• Linear Projections for Each Head: The input embeddings (or outputs from a previous layer) are linearly projected into multiple sets of Queries (Q_i), Keys (K_i), and Values (V_i) using different learned weight matrices (W_i^Q, W_i^K, W_i^V) for each head ‘i’. Typically, if the model dimension is d_model and there are ‘h’ heads, each head operates on vectors of dimension d_k = d_v = d_model / h.
• Parallel Scaled Dot-Product Attention: Self-attention (scaled dot-product attention) is performed in parallel for each head using its respective Q_i, K_i, and V_i projections. This results in ‘h’ separate output matrices (head_i).
• Concatenation: The output matrices from all attention heads are concatenated along the feature dimension.
• Final Linear Projection: The concatenated output is then passed through a final linear transformation (multiplying by another learned weight matrix W^O) to produce the final output of the multi-head attention layer. This projection often maps the concatenated features back to the original d_model dimension.

Mathematically, multi-head attention can be expressed as:

$$\mathrm{MultiHead(Q, K, V)} = \mathrm{Concat(head_1, \dots, head_h)}W^O$$

Where each head is calculated as:

$$\mathrm{head_i} = \mathrm{Attention(QW_i^Q, KW_i^K, VW_i^V)}$$

And W_i^Q, W_i^K, W_i^V are parameter matrices for head ‘i’, and W^O is the parameter matrix for the final linear projection.

Advantages of Multi-Head Attention

• Enhanced Representation Learning: By attending to different parts of the input sequence from different representational subspaces, the model can capture a richer variety of relationships and nuances within the data.
• Improved Ability to Focus: Different heads can learn to focus on different types of information (e.g., short-range vs. long-range dependencies, different syntactic or semantic features).
• Parallelization: The computations for each head can be performed in parallel, making the mechanism efficient on modern hardware like GPUs, contributing to faster training times compared to recurrent models for long sequences.

Conceptual Python Implementation with NumPy

The following Python script provides a conceptual implementation of the multi-head attention mechanism using NumPy. This helps in understanding the underlying mechanics. In practice, deep learning frameworks like PyTorch or TensorFlow offer optimized implementations.

import numpy as np

class MultiHeadAttention:
    def __init__(self, d_model: int, num_heads: int):
        """Initialize the MultiHeadAttention layer.

        Args:
            d_model (int): Dimensionality of the input and output (embedding dimension).
            num_heads (int): Number of attention heads.
        """
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"

        self.d_model = d_model
        self.num_heads = num_heads
        self.depth = d_model // num_heads  # Dimension of each head's key, query, value

        # Weight matrices for Q, K, V for all heads (can be seen as concatenated W_i^Q, etc.)
        # In a real implementation, these would be learnable parameters.
        # Xavier/Glorot or He initialization is common.
        self.W_q = np.random.randn(d_model, d_model) * np.sqrt(1 / d_model)
        self.W_k = np.random.randn(d_model, d_model) * np.sqrt(1 / d_model)
        self.W_v = np.random.randn(d_model, d_model) * np.sqrt(1 / d_model)

        # Weight matrix for the final linear projection
        self.W_o = np.random.randn(d_model, d_model) * np.sqrt(1 / d_model)

    def _scaled_dot_product_attention(self, q_h, k_h, v_h, mask=None):
        """Compute scaled dot product attention for a single head.

        Args:
            q_h (np.ndarray): Query tensor for a head (batch_size, seq_len_q, depth).
            k_h (np.ndarray): Key tensor for a head (batch_size, seq_len_k, depth).
            v_h (np.ndarray): Value tensor for a head (batch_size, seq_len_v, depth).
                               (seq_len_k typically equals seq_len_v)
            mask (np.ndarray, optional): Mask to apply to attention scores.

        Returns:
            tuple: (output, attention_weights)
                   output (np.ndarray): Contextualized representation (batch_size, seq_len_q, depth).
                   attention_weights (np.ndarray): Attention distribution (batch_size, seq_len_q, seq_len_k).
        """
        # Matmul Q_h and K_h.T: (batch_size, seq_len_q, depth) @ (batch_size, depth, seq_len_k)
        # -> (batch_size, seq_len_q, seq_len_k)
        matmul_qk = np.matmul(q_h, k_h.transpose(0, 2, 1))

        # Scale matmul_qk
        dk = k_h.shape[-1]
        scaled_attention_logits = matmul_qk / np.sqrt(dk)

        # Apply mask if provided (e.g., for padding or future token masking in decoders)
        if mask is not None:
            scaled_attention_logits += (mask * -1e9) # Add a large negative number where mask is 1

        # Softmax is applied on the last axis (seq_len_k) to get weights
        # Numerically stable softmax:
        scaled_attention_logits_stabilized = scaled_attention_logits - np.max(scaled_attention_logits, axis=-1, keepdims=True)
        attention_weights = np.exp(scaled_attention_logits_stabilized)
        attention_weights /= np.sum(attention_weights, axis=-1, keepdims=True)

        # Matmul attention_weights and V_h: (batch_size, seq_len_q, seq_len_k) @ (batch_size, seq_len_v, depth)
        # -> (batch_size, seq_len_q, depth)
        output = np.matmul(attention_weights, v_h)
        return output, attention_weights

    def _split_heads(self, x: np.ndarray, batch_size: int):
        """Split the last dimension into (num_heads, depth).
        Transpose the result to shape (batch_size, num_heads, seq_len, depth).
        Args:
            x (np.ndarray): Input tensor (batch_size, seq_len, d_model).
            batch_size (int): Batch size.
        Returns:
            np.ndarray: Reshaped tensor (batch_size, num_heads, seq_len, depth).
        """
        seq_len = x.shape[1]
        x = np.reshape(x, (batch_size, seq_len, self.num_heads, self.depth))
        return x.transpose(0, 2, 1, 3) # (batch_size, num_heads, seq_len, depth)

    def call(self, q_in: np.ndarray, k_in: np.ndarray, v_in: np.ndarray, mask=None):
        """Forward pass for MultiHeadAttention.

        Args:
            q_in (np.ndarray): Query input (batch_size, seq_len_q, d_model).
            k_in (np.ndarray): Key input (batch_size, seq_len_k, d_model).
            v_in (np.ndarray): Value input (batch_size, seq_len_v, d_model).
            mask (np.ndarray, optional): Mask to be applied during attention calculation.

        Returns:
            tuple: (output, attention_weights_all_heads)
                   output (np.ndarray): Final output (batch_size, seq_len_q, d_model).
                   attention_weights_all_heads (np.ndarray): Attention weights from all heads
                                                           (batch_size, num_heads, seq_len_q, seq_len_k).
        """
        batch_size = q_in.shape[0]

        # 1. Linear projections for Q, K, V
        # (batch_size, seq_len, d_model) @ (d_model, d_model) -> (batch_size, seq_len, d_model)
        q = np.dot(q_in, self.W_q)
        k = np.dot(k_in, self.W_k)
        v = np.dot(v_in, self.W_v)

        # 2. Split heads
        # (batch_size, num_heads, seq_len, depth)
        q_split = self._split_heads(q, batch_size)
        k_split = self._split_heads(k, batch_size)
        v_split = self._split_heads(v, batch_size)

        # 3. Scaled dot-product attention for each head
        # We'll store attention weights for inspection, though typically only the output is passed on.
        attention_heads_outputs = []
        attention_weights_all_heads_list = []

        # In a vectorized implementation, this loop would be avoided by operating on the num_heads dimension.
        # For clarity, we iterate here. But for efficiency, batch matrix multiply across heads is preferred.
        # Modern NumPy supports batching, so the _scaled_dot_product_attention can handle tensors like
        # (batch_size, num_heads, seq_len, depth) if implemented carefully. For this example, we'll keep it simple.
        # For a true parallel computation, think of q_split, k_split, v_split as inputs to parallel attention units.

        # Here we demonstrate a simplified version where scaled_dot_product_attention
        # is applied to the already split heads Q, K, V tensors directly if it supports batching over heads.
        # Our current _scaled_dot_product_attention expects (batch_size, seq_len, depth),
        # so we would typically loop or reshape to process all heads.
        # For this conceptual example, let's assume _scaled_dot_product_attention can handle the head dimension:
        # (batch_size, num_heads, seq_len_q, depth)
        attention_output_concat, attention_weights_concat = self._scaled_dot_product_attention(q_split, k_split, v_split, mask)
        # Note: The above assumes _scaled_dot_product_attention is modified to handle the num_heads dimension.
        # If not, a loop through heads is needed:
        # for i in range(self.num_heads):
        #     head_output, head_weights = self._scaled_dot_product_attention(q_split[:,i,:,:], k_split[:,i,:,:], v_split[:,i,:,:], mask)
        #     attention_heads_outputs.append(head_output)
        #     attention_weights_all_heads_list.append(head_weights)
        # attention_output_concat = np.stack(attention_heads_outputs, axis=1) # if looped
        # attention_weights_concat = np.stack(attention_weights_all_heads_list, axis=1) # if looped

        # 4. Concatenate heads
        # Transpose back to (batch_size, seq_len_q, num_heads, depth)
        attention_output_transposed = attention_output_concat.transpose(0, 2, 1, 3)
        # Reshape to (batch_size, seq_len_q, d_model)
        seq_len_q = q_in.shape[1]
        concat_attention = np.reshape(attention_output_transposed, (batch_size, seq_len_q, self.d_model))

        # 5. Final linear projection
        # (batch_size, seq_len_q, d_model) @ (d_model, d_model) -> (batch_size, seq_len_q, d_model)
        output = np.dot(concat_attention, self.W_o)

        return output, attention_weights_concat # attention_weights_concat will be (batch_size, num_heads, seq_len_q, seq_len_k)

# Example usage:
d_model = 512      # Dimension of embedding
num_heads = 8        # Number of attention heads
batch_size = 2     # Number of sequences in a batch
seq_len = 10       # Length of each sequence

# Create an instance of MultiHeadAttention
multi_head_attn = MultiHeadAttention(d_model, num_heads)

# Example input tensors (random data for demonstration)
# In a real scenario, Q, K, V might be the same tensor (self-attention)
# or different (e.g., encoder-decoder attention).
Q_input = np.random.randn(batch_size, seq_len, d_model)
K_input = np.random.randn(batch_size, seq_len, d_model)
V_input = np.random.randn(batch_size, seq_len, d_model)

# Perform multi-head attention
# Note: The conceptual _scaled_dot_product_attention in this class needs to be adjusted
# to properly handle the (batch_size, num_heads, seq_len, depth) input from _split_heads
# or a loop should be used. For brevity, we assume it's adapted.
# Let's refine the call to reflect how it would work if _scaled_dot_product_attention handled batched heads:

# Simplified direct call for demonstration - assuming internal handling
output, attention_weights = multi_head_attn.call(Q_input, K_input, V_input)

print("Input Query (Q_input shape):", Q_input.shape)
print("Multi-Head Attention Output (shape):", output.shape)
print("Attention Weights (shape):", attention_weights.shape)
print("\nFirst batch, first token's output vector (first 5 elements):\n", output[0, 0, :5])
print("\nAttention weights for first head, first batch, first query token (first 5 keys):\n", attention_weights[0, 0, 0, :5])

Expected Output Structure

Running the script with the example usage will produce output tensors. The shapes are important:

Input Query (Q_input shape): (2, 10, 512)
Multi-Head Attention Output (shape): (2, 10, 512)
Attention Weights (shape): (2, 8, 10, 10)

First batch, first token's output vector (first 5 elements):
 [ ... some float values ... ]

Attention weights for first head, first batch, first query token (first 5 keys):
 [ ... some float values summing to <=1 ... ]

The output tensor has the same shape as the input query, representing the contextualized embeddings. The attention weights tensor shows, for each head and each token, how much attention it paid to all other tokens in the sequence.

Conclusion

The multi-head attention mechanism is a powerful and integral part of the Transformer architecture. It allows models to effectively capture complex dependencies in sequential data by simultaneously processing information through multiple attention perspectives. Its ability to learn diverse representations and its suitability for parallel computation have made it a cornerstone of modern deep learning models, especially in NLP.

While this chapter provided a conceptual overview and a NumPy-based implementation for clarity, production models leverage optimized versions within frameworks like PyTorch and TensorFlow. Understanding these fundamentals is key to appreciating the innovations driving today’s advanced AI systems. For further study, explore the original ‘Attention Is All You Need’ paper and implementations in popular deep learning libraries.