Positional encodings (RoPE)
Section focus: Explain every item aloud without notes
Description
Level 1 β The Simple Intuition (ELI5)
By itself, Self-Attention is permutation-equivariantβit treats a sentence like a bag of words! Without position markers,
"The dog bit the man" and "The man bit the dog" produce the exact same attention scores.Early Transformers stamped a fixed "seat number" onto each word by adding a vector (). Rotary Position Embedding (RoPE) does something much smarter: instead of adding a position offset, it rotates each Query and Key vector on a 2D clock face by an angle proportional to its position . Why? Because when two rotated vectors take a dot product, only the angle between them () matters! Whether
"dog bit man" appears at token or token , the relative distance between "dog" and "bit" is still tick of the clock.Level 2 β How It Works (The Mechanics)
- The Goal of Relative Position Encoding: We want a transformation applied to Query at position and Key at position such that their inner product depends only on their vectors and their relative distance :
- 2D Rotation Intuition: Take a 2D vector (or equivalently a complex number ) and rotate it counterclockwise by angle :
Because rotation matrices are orthogonal (), the dot product between a Query rotated by and a Key rotated by simplifies to:
The absolute positions and vanishβleaving pure relative distance !
- Extending to Dimensions (Multi-Speed Clocks):
- For a head dimension (e.g., ), RoPE pairs adjacent dimensions into independent 2D planes: .
- Each pair rotates at its own frequency (typically in LLaMA-1/2, or in LLaMA-3).
- First pairs (): High frequency (second hand of a clock) captures sharp local syntax (adjacent words).
- Last pairs (): Very low frequency (hour/year hand) captures long-range document structure.
Level 3 β Advanced Depth (Long-Context Extrapolation: PI, NTK-Aware & YaRN)
- Where RoPE Applied? RoPE is applied only to Queries () and Keys () inside every attention layer (after linear projection, before ), and never to Values ().
- Why RoPE Fails Beyond Training Context : If a model is trained up to , its slow-frequency dimensions have only seen a small arc of rotation . At token , those dimensions rotate into unseen angles, causing attention scores to go haywire.
- Context Extension Techniques:
- Position Interpolation (PI): Linearly scales positions so all angles stay inside the trained arc, but crushes high-frequency local resolution.
- NTK-Aware Scaling / RoPE Base Scaling: Increases the base wavelength (e.g., ), stretching low frequencies while keeping high frequencies intact.
- YaRN (Yet another RoPE extensioN): Splits dimensions into bandsβno interpolation on high-frequency pairs, smooth ramp interpolation on mid-frequency pairs, full interpolation on low-frequency pairs, plus an attention temperature scaling factor.
Learning notes
Visual Intuition: Multi-Speed 2D Clocks in RoPE
Loading diagramβ¦
The Efficient Vectorized Trick (No Sparse Matrix Multiplication Needed!)
Because the block-diagonal rotation matrix is mostly zeros, RoPE is computed in elementwise operations:
Evolution of Positional Encodings in Transformers
Method | Used In | How It Injects Position | Relative Distance Aware? | Length Extrapolation |
Learned Absolute | GPT-2, GPT-3, BERT | Adds learned table at layer 0 | No (absolute index only) | Impossible beyond |
Sinusoidal Absolute | Original Transformer (2017) | Adds fixed sine/cosine vector at layer 0 | Weak (mixed with content ) | Poor in practice |
ALiBi | BLOOM, MPT | Subtracts linear penalty $-m \cdot | i - j | QK^T$ |
RoPE (Rotary) | LLaMA 1β3, Mistral, Qwen, Gemma | Rotates and pairs at every attention layer | Yes (exact via orthogonal rotation) | Excellent with NTK / YaRN |
Example or implementation
Production-Style PyTorch Implementation of RoPE (LLaMA Style)
import torch def precompute_freqs_cis(dim: int, max_seq_len: int, theta_base: float = 10000.0) -> torch.Tensor: """Precomputes complex unit-circle rotations e^(i * m * theta_j) for RoPE.""" # 1. Compute frequencies theta_j = 1 / (base^(2j / dim)) for j in [0, dim/2) freqs = 1.0 / (theta_base ** (torch.arange(0, dim, 2).float() / dim)) # 2. Outer product with positions m = [0, 1, ..., max_seq_len - 1] positions = torch.arange(max_seq_len, dtype=torch.float32) angles = torch.outer(positions, freqs) # Shape: (max_seq_len, dim // 2) # 3. Convert polar (magnitude=1, angle=m*theta) to complex numbers cos(a) + i*sin(a) freqs_cis = torch.polar(torch.ones_like(angles), angles) return freqs_cis def apply_rope(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: """ Applies Rotary Position Embedding to Queries or Keys. x shape: (batch, seq_len, n_heads, head_dim) """ # Group adjacent pairs into complex numbers: (..., head_dim // 2) x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) # Reshape freqs_cis to broadcast across batch and heads: (1, seq_len, 1, head_dim // 2) freqs_cis = freqs_cis[:x.shape[1]].unsqueeze(0).unsqueeze(2) # Multiply complex numbers -> rotates phases by m * theta! x_rotated = torch.view_as_real(x_complex * freqs_cis).flatten(-2) return x_rotated.type_as(x) # Verification: Prove that <RoPE(q, m), RoPE(k, n)> depends ONLY on (m - n)! head_dim = 64 freqs = precompute_freqs_cis(dim=head_dim, max_seq_len=100) q_vec = torch.randn(1, 1, 1, head_dim) k_vec = torch.randn(1, 1, 1, head_dim) # Dot product at positions (m=5, n=2) -> relative distance = 3 q_pos5 = apply_rope(q_vec.expand(1, 10, 1, head_dim), freqs)[:, 5] k_pos2 = apply_rope(k_vec.expand(1, 10, 1, head_dim), freqs)[:, 2] dot_5_2 = (q_pos5 * k_pos2).sum() # Dot product at positions (m=45, n=42) -> relative distance = 3 q_pos45 = apply_rope(q_vec.expand(1, 50, 1, head_dim), freqs)[:, 45] k_pos42 = apply_rope(k_vec.expand(1, 50, 1, head_dim), freqs)[:, 42] dot_45_42 = (q_pos45 * k_pos42).sum() print(f"Dot(m=5, n=2) = {dot_5_2.item():.6f}") print(f"Dot(m=45, n=42) = {dot_45_42.item():.6f} <-- Identical!")
Interview-ready explanation
90-Second Spoken Pitch for Interviews
"Self-attention is permutation-equivariant, so without positional encodings it treats a prompt as an unordered bag of words. Early Transformers added absolute sinusoidal or learned vectors to token embeddings at the input layer, which polluted the semantic vector and struggled to generalize because position and position had no built-in translational invariance with position and .
Rotary Position Embedding (RoPE) solves this by encoding absolute position with a rotation matrix that naturally yields relative position in the attention dot product. Specifically, inside every attention layer, RoPE splits the -dimensional Query and Key vectors into two-dimensional pairs and rotates each pair in the 2D plane by an angle , where is the token's position and is a geometric frequency schedule ranging from fast local frequencies to slow long-range frequencies.
Because rotation matrices are orthogonal, when we compute the attention score , the rotations combine into . That means the attention score depends purely on the content vectors and their relative distance , while remaining compatible with KV caching because each Key is rotated using only its own absolute index when cached."