Self-Attention

Unlike recurrent neural networks (RNNs) that process tokens sequentially, self-attention computes pairwise relationships between all tokens simultaneously. This parallelization enables the model to capture long-range dependencies without the bottleneck of a hidden state that must be passed step-by-step.

• Mechanism: Every token generates a Query, Key, and Value vector. The attention matrix is computed as a scaled dot-product between all Queries and Keys.
• Result: A token at position 500 can directly attend to a token at position 1 with the same computational path length as position 499.
• Training Efficiency: Full parallelization over the sequence dimension dramatically reduces training time on modern GPU hardware.

The fundamental mathematical operation at the heart of self-attention. It computes a weighted sum of Value vectors, where the weight assigned to each Value is determined by the compatibility of the corresponding Query and Key.

• Formula: Attention(Q, K, V) = softmax(QK^T / √d_k) * V
• Scaling Factor (√d_k): The dot products are scaled down by the square root of the key dimension to prevent the softmax function from entering regions of extremely small gradients, which would occur with large magnitude dot products.
• Softmax: Converts the raw compatibility scores into a probability distribution, ensuring all attention weights sum to 1.

Instead of performing a single attention function, the model projects the Queries, Keys, and Values into multiple, lower-dimensional subspaces and runs the scaled dot-product attention in parallel on each. This allows the model to jointly attend to information from different representation subspaces at different positions.

• Diverse Focus: One head might track syntactic dependencies, another might resolve pronoun references, and a third might focus on semantic topic continuity.
• Concatenation: The outputs of all heads are concatenated and projected linearly back to the model's hidden dimension.
• Example: In a standard Transformer base model, h=8 or h=12 heads are used, with each head operating on a d_model/h dimensional subspace.

Because the self-attention operation is inherently permutation-invariant—it has no built-in notion of token order—positional information must be explicitly injected. This is done by adding a positional encoding vector to the input embeddings before the first attention layer.

• Sinusoidal Encoding: The original Transformer uses fixed sine and cosine functions of varying frequencies, allowing the model to learn relative positions and theoretically extrapolate to sequence lengths unseen during training.
• Learned Embeddings: Many modern implementations use a standard embedding lookup table for absolute positions, which is learned during training.
• Relative Position: Advanced methods like RoPE and ALiBi encode position directly into the attention score calculation, often improving length generalization.

The primary computational limitation of standard self-attention is its O(n²) time and memory complexity with respect to sequence length n. The QK^T matrix multiplication creates an n x n attention matrix.

• Memory Bottleneck: For a sequence of 8,192 tokens, the attention matrix alone requires significant GPU memory, growing quadratically.
• Mitigation Strategies: This constraint has motivated efficient variants like FlashAttention (hardware-aware exact attention), Sparse Attention (Longformer, BigBird), and Linear Attention approximations.
• KV-Cache: During autoregressive decoding, the Key and Value tensors for all previous tokens are cached to avoid recomputation, but this cache also grows linearly with sequence length.

In the original Transformer architecture, a critical distinction exists between self-attention and cross-attention, particularly in the Decoder stack.

• Self-Attention: Queries, Keys, and Values all originate from the same source sequence. Used in the Encoder (to build contextual representations of the input) and in the Decoder's masked self-attention (to process the output sequence).
• Cross-Attention: Queries come from the Decoder's current state, while Keys and Values come from the Encoder's final output. This is the mechanism by which the Decoder aligns the generated output with the input source, crucial for tasks like machine translation.
• Causal Masking: Decoder self-attention uses a mask to prevent a token from attending to future tokens, preserving the autoregressive property.

The broader family of techniques that allow a model to dynamically weigh the importance of different parts of an input. Self-attention is a specific case where the query, key, and value all originate from the same sequence. The core operation computes a weighted sum of values, where the weight is a compatibility function between a query and a key.

• Cross-attention: Queries from one sequence attend to keys/values from another (e.g., decoder attending to encoder output).
• Scaled Dot-Product: The standard score calculation: softmax(QK^T / √d_k).
• Computational Complexity: Standard self-attention scales quadratically with sequence length, O(n²).

The standard practice of running multiple self-attention operations in parallel. Instead of a single attention function, the query, key, and value are linearly projected h times. This allows the model to jointly attend to information from different representation subspaces at different positions.

• Parallel Heads: Each head captures a different syntactic or semantic relationship.
• Concatenation: Outputs of all heads are concatenated and projected back to the model dimension.
• Trade-off: High expressivity comes at the cost of a large KV-Cache memory footprint during inference.

An interpolation between Multi-Head Attention and Multi-Query Attention. GQA divides query heads into groups, and each group shares a single key-value head. This drastically reduces the size of the KV-Cache during autoregressive decoding.

• Memory Savings: Reduces KV-Cache size by a factor equal to the number of query heads per group.
• Quality Retention: Maintains model quality much closer to MHA than MQA, making it the standard for large models like Llama 2 70B.
• Inference Speed: Lower memory bandwidth requirements accelerate token generation.

An IO-aware exact attention algorithm that avoids materializing the full N×N attention matrix in slow GPU High Bandwidth Memory (HBM). It uses tiling to compute softmax in blocks and recomputation during the backward pass.

• Speedup: Achieves 2-4x wall-clock speedup over standard PyTorch attention.
• Memory Efficiency: Reduces memory footprint from O(n²) to O(n), enabling longer context lengths.
• Exact Computation: Not an approximation; the output is mathematically identical to standard attention.

A position encoding method that encodes absolute position with a rotation matrix and naturally incorporates explicit relative position dependency in the self-attention formulation. RoPE rotates the query and key vectors by an angle proportional to their position index.

• Relative Decay: The dot product between a query and key naturally decays with their relative distance.
• Long-Range Extrapolation: Supports context length extension beyond training lengths via techniques like NTK-aware scaling.
• Adoption: Used in Llama, PaLM, and Qwen architectures.

A memory buffer that stores the pre-computed key and value tensors for all previously generated tokens during autoregressive decoding. Without it, the model would recompute attention for the entire sequence at each step.

• Prefill Phase: The prompt is processed in parallel to populate the initial cache.
• Decode Phase: Only the new token's key/value is computed and appended; previous tensors are read from cache.
• Bottleneck: For long sequences, the cache size becomes the primary memory constraint, motivating GQA and Multi-Query Attention.