Multi-Head Self-Attention (MHSA)

MHSA's defining feature is its use of multiple, independent attention heads that operate in parallel. Each head learns a distinct representation subspace by projecting the input sequence into unique Query (Q), Key (K), and Value (V) matrices using separate, learned linear transformations. This allows the model to simultaneously attend to different types of information—such as syntactic structure, semantic roles, or positional relationships—within the same input sequence. The outputs of all heads are concatenated and linearly projected to produce the final attended representation.

At the heart of each attention head is the scaled dot-product attention function. It computes a compatibility score between all elements in the sequence:

• Query-Key Dot Product: Measures the similarity between each query and all keys.
• Scaling: The dot product is divided by the square root of the key dimension (sqrt(d_k)). This scaling prevents the softmax function from entering regions of extremely small gradients, which stabilizes training.
• Softmax Application: A softmax is applied to the scaled scores to create a probability distribution (the attention weights).
• Weighted Sum: The final output for each position is a weighted sum of the value vectors, where the weights are these attention probabilities.

The self-attention operation is fundamentally permutation-equivariant; it treats the input as an unordered set. Without additional information, the sequence [A, B, C] would be processed identically to [C, B, A]. To inject order, positional encodings are added to the input embeddings. These can be:

• Sinusoidal Encodings: Fixed, deterministic patterns of sine and cosine waves of varying frequencies.
• Learned Positional Embeddings: Vectors learned during training, similar to word embeddings. This addition allows the model to utilize both the content (the embedding) and the absolute/relative position of each token.

MHSA excels at modeling dependencies between tokens regardless of their distance in the sequence, a key advantage over recurrent neural networks (RNNs). However, this comes with a computational cost:

• Quadratic Complexity: The need to compute attention scores between every pair of tokens results in O(n^2) time and memory complexity, where n is the sequence length. This is the primary bottleneck for processing very long documents or high-resolution images.
• Global Context: Unlike convolutional networks with limited receptive fields, each token in MHSA has direct access to the context of every other token in the sequence after a single layer, enabling immediate global information integration.

Despite its computational cost, MHSA is parameter-efficient relative to its representational capacity. The number of parameters is primarily in the linear projection layers for Q, K, V and the final output projection, which scale with model dimension and number of heads, not sequence length. The multiple heads act as a form of ensemble learning within a single layer, diversifying the types of patterns the model can capture. This structure is a key reason transformers generalize so effectively across diverse tasks with a unified architecture.

MHSA is not limited to text. In modality-specific feature extraction, it is applied to encoded representations of raw data:

• Vision Transformers (ViT): Image patches are treated as a sequence.
• Audio Spectrogram Transformers (AST): Spectrogram patches form the input sequence.
• Point Cloud Processing: Sets of 3D points can be processed with permutation-invariant attention. In these cases, MHSA operates on latent feature sequences extracted by initial convolutional or linear layers, allowing it to learn complex, long-range relationships within images, audio clips, or 3D spaces.

The seminal paper introduced MHSA as a core mechanism for machine translation. Its implementation defined the standard formula:

• Parallel Heads: Multiple attention heads operate simultaneously on linearly projected versions of the input (Query, Key, Value).
• Concatenation & Projection: Head outputs are concatenated and passed through a final linear projection layer.
• Scaled Dot-Product: Attention scores are scaled by the square root of the key dimension (d_k) to counteract vanishing gradients. This architecture demonstrated that purely attention-based models could outperform recurrent and convolutional networks on sequence tasks.

BERT (Bidirectional Encoder Representations from Transformers) uses MHSA in its encoder stack with a key modification: masked self-attention. During pre-training, the model uses a causal mask to prevent tokens from attending to future tokens, enabling next-sentence prediction and masked language modeling objectives. In its fine-tuned form, BERT uses full bidirectional attention, allowing each token to attend to all other tokens in the sequence, which is crucial for understanding context.

Generative Pre-trained Transformer (GPT) models are decoder-only architectures. Their MHSA implementation is strictly causal or unidirectional. A triangular attention mask ensures each token can only attend to previous tokens and itself, which is essential for autoregressive text generation. As context windows have grown (e.g., GPT-4), optimizations like key-value (KV) caching are critical. During generation, past KV states are cached to avoid recomputation, dramatically improving inference speed for sequential token generation.

Vision Transformers adapt MHSA for image data by treating an image as a sequence of patches. Each patch is linearly embedded and prepended with a [CLS] token. The MHSA mechanism then operates on this sequence, allowing patches to attend to each other globally. This enables the model to learn long-range dependencies across the entire image, a task traditionally handled by convolutional networks with large receptive fields. ViT demonstrated that transformers could achieve state-of-the-art results on image classification without convolution.

Self-Attention is the fundamental mechanism upon which MHSA is built. It computes a weighted sum of values for each position in a sequence, where the weights (attention scores) are determined by the compatibility between a query at that position and keys from all positions.

• Core Operation: For an input sequence, it learns to focus on different parts of the same sequence when constructing a representation for each element.
• Scaled Dot-Product: The most common implementation, using softmax(QK^T / sqrt(d_k))V, where Q, K, V are linear projections of the input.
• Purpose: Enables modeling of long-range dependencies and contextual relationships within sequences, which was a limitation of earlier recurrent architectures.

The Transformer is a neural network architecture that relies entirely on the self-attention mechanism, dispensing with recurrence and convolutions. MHSA is its central computational block.

• Encoder-Decoder: The original Transformer uses an encoder (maps input to representations) and a decoder (generates output autoregressively).
• Layer Stacking: Multiple identical layers, each containing an MHSA sub-layer and a position-wise feed-forward network, are stacked to create depth.
• Residual Connections & Layer Normalization: Each sub-layer employs residual connections followed by layer normalization to stabilize training of deep networks.
• Impact: This architecture forms the basis for models like BERT (encoder-only), GPT (decoder-only), and T5 (encoder-decoder).

Cross-Attention (or encoder-decoder attention) is a variant where the queries come from one sequence (e.g., the decoder) and the keys and values come from a different, separate sequence (e.g., the encoder output).

• Key Difference from Self-Attention: It attends across modalities or sequences rather than within a single sequence.
• Primary Use Case: Essential for sequence-to-sequence tasks like machine translation, where the decoder must attend to the encoded source sentence.
• In Multimodal Models: Used to fuse information between modalities, e.g., allowing a text query to attend to visual features from an image encoder.

Positional Encoding injects information about the order of tokens into the model, as the self-attention mechanism is inherently permutation-invariant.

• The Problem: Without it, the model has no inherent notion of sequence order; "cat dog bites" and "dog cat bites" would be identical.
• Sinusoidal Encodings: The original Transformer uses fixed, pre-defined sine and cosine functions of different frequencies.
• Learned Embeddings: An alternative where positional embeddings are learned parameters, similar to token embeddings.
• Integration: Added to the input token embeddings before the first self-attention layer.

Flash Attention is an IO-aware, exact attention algorithm that dramatically speeds up and reduces the memory footprint of standard self-attention computation.

• Core Innovation: It recomputes attention scores on-the-fly during the backward pass, avoiding the need to store the large softmax(QK^T) matrix in high-bandwidth memory (HBM).
• Key Benefit: Enables training of models with much longer context lengths (e.g., 128K tokens) by reducing memory usage from quadratic to linear in sequence length.
• Implementation: Uses tiling to keep intermediate working memory in fast SRAM, minimizing slow HBM reads/writes. It is a foundational optimization for modern large language models.

Sparse Attention refers to self-attention mechanisms that restrict the attention pattern, so each token only attends to a subset of other tokens, reducing computational complexity.

• Motivation: Standard self-attention has O(n²) complexity in sequence length, which is prohibitive for very long sequences.
• Patterns: Includes sliding window attention (local), dilated attention, block-sparse patterns, or attention based on learned or content-based clustering.
• Examples: Used in models like Longformer (sliding window + global tokens) and BigBird (random, window, and global attention).
• Trade-off: Introduces an architectural bias but enables processing of documents or books-length text.