Rotary Position Embeddings (RoPE)

Overview

Rotary Position Embeddings (RoPE) is a positional encoding method that encodes position information by rotating query and key vectors using rotation matrices. It provides better length extrapolation, simpler implementation, and stronger position representation compared to standard position embeddings.

• Paper: "RoFormer: Enhanced Transformer with Rotary Position Embedding" (Su et al., 2021)
• Key Innovation: Use complex number rotations to encode position
• Adoption: LLaMA, Mistral, Qwen, GPT-4, most modern LLMs
• Advantage: Works better for long contexts than absolute/relative position embeddings
• Foundation: Based on complex number rotations and geometric principles

---

Problem: Why Standard Position Embeddings Fail

Absolute Position Embeddings (APE)

Approach: Add learned position vectors to token embeddings
- ┌──────────────────────────────────────────┐
    - "The cat sat on the mat"                 │
    - pos1 pos2 pos3 pos4 pos5 pos6           │
    - ↓    ↓    ↓    ↓    ↓    ↓             │
    - [ embed + pos_1, embed + pos_2, ... ]   │
    - Each position has fixed learned vector   │
  - ┘

Problems:
❌ Requires retraining for longer sequences
❌ Cannot generalize to unseen positions
❌ Position information lost quickly in deep layers
❌ No connection between position and semantic distance

Relative Position Embeddings (RPE)

Approach: Encode relative distance between tokens
- ┌──────────────────────────────────────────┐
    - Attention(i, j) includes distance (j - i)│
    - Distance = 1: "adjacent tokens"          │
    - Distance = 5: "5 tokens apart"           │
  - ┘

Problems:
❌ Still has maximum distance limit
❌ Requires bucketing large distances
❌ Complex implementation in transformers
❌ Loses absolute position information
❌ Difficult to extrapolate beyond training length

Issues Both Face

Core Problem:
─────────────
Position information should:
1. Encode absolute position (where am I?)
2. Encode relative distance (how far is this token?)
3. Generalize to longer sequences
4. Scale with distance naturally

Standard embeddings:
  - Treat each position independently (APE)
  - Or only encode pairwise distance (RPE)

Need: Unified, scalable, extrapolatable approach!

---

Rotary Position Embeddings: The Solution

Core Intuition

Key Insight: Use rotation angles to encode position
             Position = Rotation angle

If we rotate vectors by angle θ:
  - θ = 0° → no rotation (position 0)
  - θ = π/4 → rotated by 45° (position 1)
  - θ = π/2 → rotated by 90° (position 2)
  - θ = π → rotated by 180° (position N)
  - Each position has unique rotation!

Benefit:
  - Relative position → relative rotation angle difference
  - Absolute position → absolute angle
  - Smooth interpolation for unseen positions
  - Geometric property: rotation preserves distances!

Mathematical Foundation

Complex Number Representation:
─────────────────────────────
Position m and dimension j:
θ_{m,j} = m × Θ^{-2j/d}   where Θ = 10,000

For 2D (complex number representation):
z = x + iy = (x + iy) × e^{iθ_m}

Rotation matrix form (2D):
[cos(θ_m)  -sin(θ_m)] [x]
[sin(θ_m)   cos(θ_m)] [y]

Attention with RoPE:
<q_m, k_n> = Re(e^{imθ} q^* × e^{inθ} k)
            = Re(q^* × e^{i(m-n)θ} × k)

Key property:
  - Only depends on (m - n): relative position!
  - Absolute positions cancel out
  - Automatic length extrapolation!

Implementation

class RotaryPositionalEmbeddings(nn.Module):
    def __init__(self, hidden_dim, base=10000, device=None):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.base = base
        self.device = device

        # Pre-compute rotation angles
        # inv_freq = 1.0 / (base^(2j/d)) for j = 0..d/2
        inv_freq = 1.0 / (base ** (torch.arange(0, hidden_dim, 2).float() / hidden_dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    @torch.no_grad()
    def forward(self, seq_len, device=None):
        """
        Compute rotation matrices for all sequence positions

        Output: (seq_len, hidden_dim)
        Each position has its rotation angle
        """
        device = device or self.device or self.inv_freq.device

        # t = [0, 1, 2, ..., seq_len-1]
        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)

        # freqs = outer(t, inv_freq)
        # freqs[m, j] = m × inv_freq[j]
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)

        # Complex representation: e^{iθ}
        # cos and sin pairs
        emb = torch.cat([freqs, freqs], dim=-1)  # Duplicate for cos/sin

        # Return cos and sin for efficient application
        return emb.cos(), emb.sin()

def apply_rotary_embeddings(q, k, cos, sin):
    """
    Apply rotary embeddings to query and key vectors

    Args:
        q: (batch, num_heads, seq_len, head_dim)
        k: (batch, num_heads, seq_len, head_dim)
        cos: (seq_len, head_dim)
        sin: (seq_len, head_dim)
    """
    # Reshape q and k to process rotation
    # q = [q_1, q_2, q_3, q_4, ...] where pairs are (x, y) coordinates
    q_complex = q.view(*q.shape[:-1], -1, 2)  # (..., head_dim//2, 2)
    k_complex = k.view(*k.shape[:-1], -1, 2)

    cos = cos.view(-1, 1, -1, 1)  # (seq_len, 1, head_dim//2, 1)
    sin = sin.view(-1, 1, -1, 1)

    # Apply rotation matrix: [cos -sin; sin cos]
    # [x', y'] = [cos -sin; sin cos] @ [x, y]
    q_rotated = torch.cat([
        q_complex[..., 0] * cos - q_complex[..., 1] * sin,  # x' = x*cos - y*sin
        q_complex[..., 0] * sin + q_complex[..., 1] * cos,  # y' = x*sin + y*cos
    ], dim=-1)

    k_rotated = torch.cat([
        k_complex[..., 0] * cos - k_complex[..., 1] * sin,
        k_complex[..., 0] * sin + k_complex[..., 1] * cos,
    ], dim=-1)

    # Reshape back to original form
    q_rotated = q_rotated.view(*q.shape)
    k_rotated = k_rotated.view(*k.shape)

    return q_rotated, k_rotated

# Usage in attention:
class RoPEAttention(nn.Module):
    def __init__(self, hidden_dim, num_heads):
        super().__init__()
        self.rope = RotaryPositionalEmbeddings(hidden_dim // num_heads)
        self.num_heads = num_heads

    def forward(self, q, k, v):
        # q, k, v: (batch, seq_len, hidden_dim)
        seq_len = q.shape[1]

        # Reshape to (batch, seq_len, num_heads, head_dim)
        batch_size = q.shape[0]
        head_dim = self.hidden_dim // self.num_heads
        q = q.view(batch_size, seq_len, self.num_heads, head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.num_heads, head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.num_heads, head_dim).transpose(1, 2)

        # Get rotation matrices
        cos, sin = self.rope(seq_len)

        # Apply RoPE
        q, k = apply_rotary_embeddings(q, k, cos, sin)

        # Standard attention
        scores = q @ k.transpose(-2, -1) / sqrt(head_dim)
        weights = softmax(scores, dim=-1)
        output = weights @ v

        return output

---

Why RoPE is Better

Length Extrapolation

Experiment: Train on 2K context, test on 4K

Traditional APE:
  - Position 1-2000: Seen during training
  - Position 2001-4000: Never seen
  - Performance: Severe degradation (~50% worse)

Relative Position Embeddings:
  - Trained with max_distance=100 buckets
  - Beyond 100: All treated as "very far"
  - Performance: Moderate degradation (~30% worse)

RoPE:
  - No seen/unseen distinction
  - Rotation angles interpolate smoothly
  - Extrapolation is mathematically grounded
  - Performance: Minor degradation (~5-10% worse)
  - Much better than alternatives!

Theoretical Properties

1. Absolute Position Encoding
  - Each position has unique angle
  - Can recover absolute position from rotation
  - Preserves absolute information

2. Relative Position Encoding
  - Dot product only depends on (m - n)
  - Attention score depends on relative distance
  - Automatic relative encoding

3. Distance Preservation
  - Rotation preserves Euclidean distance
  - Cosine similarity patterns preserved
  - Semantic relationships maintained

4. Inductive Bias
  - Geometrically grounded in rotations
  - Not arbitrary learned vectors
  - Stronger generalization

---

Practical Considerations

Frequency Choices

Standard Base=10,000:
  - θ_j = 10,000^{-2j/d}
  - For hidden_dim=4096:
    - j=0: θ₀ = 1.0
    - j=64: θ₆₄ = 0.1
    - j=1024: θ₁₀₂₄ = 0.00001
    - j=2048: θ₂₀₄₈ ≈ 10^-7
  - Covers frequencies from 1 Hz to 10^-7 Hz

Intuition:
  - Low frequencies: Slow oscillation (long-range position)
  - High frequencies: Fast oscillation (short-range position)
  - Mixture: Captures all distance scales

Typical adjustment for longer contexts:
  - base = 100,000 (for 100K+ contexts)
  - base = 1,000,000 (for 1M+ contexts)
  - Stretches frequency spectrum for longer sequences

Position Interpolation vs Extrapolation

Linear Interpolation (ALiBi-style):
  - Shrink position indices for unseen lengths
  - If trained on 2K, test on 4K:
    - Use positions [0, 0.5, 1, 1.5, ..., 1999]
  - Forces model to use learned frequencies only
  - Works well: 5-10% degradation for 2x length

Frequency Scaling:
  - Increase base inversely with scaling
  - For 2x length: use base = 10000 × 2
  - Stretches frequency spectrum
  - Also works: 5-10% degradation for 2x length

Natural Extrapolation:
  - Don't change anything
  - RoPE inherently extrapolates
  - Smoothly handles novel positions
  - Best option if model trained well

---

Comparison with Alternatives

Position Encoding Method Comparison:

                    APE     RPE     ALiBi   RoPE
─────────────────────────────────────────────────
Extrapolation       ❌      ⚠️      ✅      ✅
Simplicity         ✅      ❌      ✅      ✅
Memory overhead    ✅      ✅      ✅      ✅
Length 4K (train 2K):
  Performance      ⚠️      ⚠️      ✅      ✅
Adoption          ❌      ❌      ⚠️      ✅✅✅

Winner: RoPE (industry standard now)
Reason: Better extrapolation + simple + works great

---

Key Takeaways

🌀 RoPE uses rotations to encode position — elegant and geometric
📈 Excellent extrapolation: 2K→4K training generalizes well
🎯 Only depends on relative distance in attention computation
⚡ Simple to implement and computationally efficient
🚀 Industry standard in modern LLMs (LLaMA, Mistral, GPT-4)

---

Related Notes

• 00 Attention Mechanisms - Where RoPE is applied
• Length Extrapolation & Long Context - Related challenge
• Flash Attention - Works together with RoPE
• Llm Inference Optimization - Complete inference stack