Attention Mechanism

The fundamental operation of attention is computing a weighted sum of values, where the weights are derived from a compatibility function between a query and a set of keys. This allows the model to dynamically focus on the most salient tokens in a sequence, regardless of their distance. The output for each token is a context-aware representation that aggregates information from the entire input, with the attention weights explicitly dictating the contribution of each token. This mechanism directly addresses the vanishing gradient problem of recurrent networks by providing a direct, differentiable path between any two positions in the sequence.

The standard attention function in Transformers is defined as: Attention(Q, K, V) = softmax(QK^T / √d_k)V. The query (Q) represents the token seeking information, the key (K) represents the label for each token, and the value (V) is the actual content. The dot product QK^T computes a raw relevance score. The scaling factor 1/√d_k prevents the softmax function from entering regions with extremely small gradients when the dimensionality d_k is large, ensuring stable training. This operation is computed in parallel for all tokens using matrix multiplication.

Instead of performing a single attention function, MHA linearly projects Q, K, and V into h distinct lower-dimensional subspaces. Each head performs attention in parallel, allowing the model to jointly attend to information from different representation subspaces at different positions. For example, one head might track syntactic dependencies while another resolves coreference. The outputs of all heads are concatenated and projected again. This multi-faceted view is critical for capturing the nuanced relationships in complex sequences.

Self-attention is the mechanism where Q, K, and V all originate from the same input sequence. It is the core building block of the encoder, allowing every token to build a representation informed by every other token in the input. Cross-attention is used in encoder-decoder architectures, where Q comes from the decoder's previous layer, but K and V come from the encoder's final output. This allows the decoder to focus on relevant parts of the input sequence at each generation step, aligning the output with the source.

In a decoder-only Transformer, self-attention must be causally masked to prevent the model from attending to future tokens during training. This is implemented by adding a mask matrix to the pre-softmax scores, setting the values for illegal future connections to -∞. This forces the softmax output for those positions to zero, ensuring the prediction for token t depends only on tokens 0 through t-1. This preserves the autoregressive property required for text generation.

The core limitation of standard attention is its O(n²) time and memory complexity with respect to sequence length n. The QK^T matrix multiplication produces an n x n attention matrix. For long sequences, this becomes a severe bottleneck. This has motivated a vast research field into efficient attention variants, including sparse attention (Longformer), linearized attention (Performer), and IO-aware exact attention (FlashAttention), which optimize memory access patterns rather than changing the algorithm.

The foundational building block where a sequence computes a representation of itself by relating different positions within the same sequence. Every token generates a Query, Key, and Value vector. The attention score is calculated as the dot product of the Query with all Keys, scaled, and softmaxed to produce a weighted sum of Values. This allows each token to gather context from every other token in the sequence, capturing long-range dependencies without the sequential bottleneck of RNNs.

A critical inference optimization that trades a small amount of model quality for a massive reduction in memory bandwidth. Instead of giving every query head its own key-value pair (Multi-Head Attention), GQA shares a single KV-Head across a group of query heads.

• Impact: Dramatically shrinks the KV-Cache size during decoding.
• Adoption: Used in Llama 2 70B, Llama 3, and Mistral models.
• Trade-off: Achieves quality close to MHA while approaching the memory efficiency of Multi-Query Attention (MQA).

An exact-attention algorithm that doesn't approximate; it computes the mathematically identical result by being hardware-aware. It minimizes slow High-Bandwidth Memory (HBM) reads/writes on the GPU by using two key techniques:

• Tiling: Processes attention in blocks that fit entirely in fast on-chip SRAM.
• Recomputation: Recalculates the softmax normalization factor during the backward pass instead of storing the massive attention matrix.

This reduces the memory footprint from O(N²) to O(N), enabling training on sequences of up to 8k tokens without model sharding.

The dominant position encoding method in modern LLMs (Llama, PaLM, Qwen). RoPE encodes absolute position information into the Query and Key vectors by rotating them in 2D subspaces by an angle proportional to their position index. Crucially, the dot product between a rotated Query and Key naturally depends only on their relative distance.

• Benefit: Enables superior context length extrapolation beyond training lengths.
• Extensions: NTK-aware scaling and YaRN further extend the effective context window without full retraining.

A memory buffer that stores the pre-computed Key and Value tensors from all previous tokens during autoregressive decoding. Without it, the model would redundantly recompute attention for the entire sequence at every new token generation step.

• Function: Transforms the per-step compute from quadratic to linear.
• Bottleneck: The cache size grows linearly with batch size × sequence length × layers, making it the primary memory constraint for serving long-context models.
• Optimizations: GQA, Multi-Query Attention, and quantization (KV8) directly target reducing this cache size.

The original attention mechanism from the 'Attention Is All You Need' paper. Instead of performing a single attention function, MHA projects Queries, Keys, and Values into h different lower-dimensional subspaces (heads) and runs the scaled dot-product attention in parallel.

• Purpose: Allows the model to jointly attend to information from different representation subspaces at different positions (e.g., one head for syntax, another for long-range subject-verb agreement).
• Output: The independent head outputs are concatenated and linearly projected back to the model dimension.