Transformers & The Attention Mechanism
Introduced by Vaswani et al. in 2017 ("Attention Is All You Need"), the Transformer replaced recurrence with global self-attention, enabling massively parallelized GPU training and unlocking the modern Generative AI era.
1. Queries, Keys, and Values: Information Retrieval in Latent Space
The central innovation of Self-Attention is borrowing the concept of a soft associative database lookup. Given input representations X ∈ ℝ^(N × d_model) for N tokens, the model projects X into three distinct continuous vector spaces using learned linear weights:
Query (Q = X · W_Q)
Represents what the current token is searching for. What contextual dependencies does this token need to resolve its meaning?
Key (K = X · W_K)
Represents the address or identity of every other token in the sequence that can be searched against by the Query.
Value (V = X · W_V)
The actual semantic content and information delivered to the token once a Query-Key match is computed.
Attention(Q, K, V) = softmax( (Q · K^T) / √d_k + M ) · V
Where:
Q · K^T ∈ ℝ^(N × N)computes the pairwise semantic compatibility scores between every token pair.√d_kis the scaling factor that prevents variance explosion and vanishing gradients.Mis an upper-triangular causal attention mask (where future positions are set to-∞) to prevent tokens from cheating by peeking ahead during autoregressive decoding.softmaxconverts scores into positive weights that sum to 1.0 along rows.
2. Evolution: MHA vs. MQA vs. Grouped-Query Attention (GQA)
As context lengths expanded from 2,048 tokens in GPT-3 to 128,000+ in Llama 3 and 1,000,000 in Gemini, storing the Key-Value (KV) cache in GPU high-bandwidth memory (HBM) became the primary bottleneck. Modern architectures evolved their attention head designs:
| Attention Architecture | Head Topology | KV-Cache Memory Size | Representative Models |
|---|---|---|---|
| Multi-Head Attention (MHA) | Each Query head has its own dedicated Key and Value head (H_Q = H_K = H_V). | 100% (Baseline - highest memory consumption). | Original Transformer, GPT-3, BERT. |
| Multi-Query Attention (MQA) | All Query heads share a single Key head and single Value head. | Massive reduction (~1/H of MHA). May lose subtle semantic nuances. | Falcon-40B, PaLM. |
| Grouped-Query Attention (GQA) | Query heads are partitioned into groups (e.g. 8 groups), with each group sharing one K/V head pair. | Optimal sweet spot: 4x–8x cache reduction with near-identical MHA accuracy. | Llama 3, Mistral 7B, Gemma 2, DeepSeek-V3. |
3. Rotary Position Embeddings (RoPE)
Because Self-Attention operations are permutation-invariant (calculating dot products without inherent order), the model must be explicitly informed of token positions. Early models added static sinusoidal vectors X + P. Modern state-of-the-art models use Rotary Position Embeddings (RoPE) (Su et al., 2021).
How RoPE Operates Geometrically
Instead of adding position vectors to embeddings, RoPE rotates the Query and Key vectors in 2D coordinate planes by an angle proportional to their sequence index m · θ:
[ sin(mθ) cos(mθ) ] [ x_2 ]
Crucial Advantage: The inner product ⟨R_m Q, R_n K⟩ depends strictly on the relative distance (m - n). This allows seamless context extension through frequency scaling (YaRN, NTK-aware RoPE).
4. Hardware Acceleration: FlashAttention
Standard PyTorch attention computes S = Q · K^T, writes the large N × N matrix to slow GPU High-Bandwidth Memory (HBM), reads it back to compute P = softmax(S), writes to HBM, and reads it again to multiply with V. This causes severe memory bandwidth bottlenecks.
The FlashAttention Innovation (Dao et al.)
FlashAttention splits Q, K, V into blocks and loads them into fast on-chip SRAM (Static RAM). Using an online softmax formulation, it incrementally computes attention without ever materializing the massive N × N intermediate matrix in global HBM.
5. Scaled Dot-Product Self-Attention in PyTorch
The following clean Python/PyTorch module implements vectorized Causal Multi-Head Self-Attention from first principles:
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
class CausalSelfAttention(nn.Module):
def __init__(self, d_model: int = 768, num_heads: int = 12):
super().__init__()
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.head_dim = d_model // num_heads
# Key, Query, Value projections combined in a single linear layer
self.qkv_proj = nn.Linear(d_model, 3 * d_model, bias=False)
self.out_proj = nn.Linear(d_model, d_model, bias=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
B, T, C = x.shape # Batch size, Sequence length, Channel dimension (d_model)
# 1. Project to Q, K, V
qkv = self.qkv_proj(x)
q, k, v = qkv.chunk(3, dim=-1)
# 2. Reshape into [B, num_heads, T, head_dim]
q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
k = k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
v = v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
# 3. Scaled Dot-Product Attention with PyTorch 2.0 FlashAttention backend
# is_causal=True automatically builds and applies the triangular mask
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
# 4. Reassemble heads back to original shape: [B, T, C]
out = out.transpose(1, 2).contiguous().view(B, T, C)
return self.out_proj(out)Module Key Takeaways
- Attention maps Query tokens to weighted linear combinations of Values according to Query-Key dot product similarity.
- Grouped-Query Attention (GQA) dramatically reduces KV-cache memory pressure during autoregressive decoding.
- Rotary Position Embeddings (RoPE) represent relative token distances via continuous rotation in complex 2D vector subspaces.
- FlashAttention fuses operations into on-chip GPU SRAM, achieving 2x–4x speedups without changing the underlying mathematical output.
Explore Related Tools
Hand-picked utilities and calculators related to this tool.
