29. Rotary Positional Embeddings (RoPE)

Problem

Instead of adding a learned or sinusoidal position vector to the input embedding, RoPE encodes position by rotating pairs of dimensions within each query/key vector by an angle proportional to that token's position. The rotation angle grows with position and shrinks across dimension pairs (via a geometrically decaying frequency), so the dot product between a rotated query and rotated key ends up depending only on their relative position i - j, not their absolute positions — exactly what a causal, position-aware attention score needs. Mistral, LLaMA, and effectively every modern open-weight decoder-only model use RoPE instead of learned absolute embeddings for this relative-position property, and because it needs no extra learned parameters.

Given a tensor already split into heads, apply RoPE to each position along the sequence axis using the standard "rotate half" formulation: x_rotated = x * cos + rotate_half(x) * sin, where rotate_half splits the head dimension in half [x1, x2] and returns [-x2, x1].

Source: src/29_rotary_positional_embeddings.py

def apply_rotary_pos_emb(x: torch.Tensor, base: float = 10000.0) -> torch.Tensor: ...

>>> x = torch.randn(1, 2, 5, 4)
>>> apply_rotary_pos_emb(x).shape
torch.Size([1, 2, 5, 4])

Step-by-Step Approach

1. Compute the per-pair frequencies: freqs = 1.0 / (base ** (arange(0, half) / half)) where half = head_dim // 2 — low-index pairs rotate fast, high-index pairs rotate slowly.
2. angles = torch.outer(positions, freqs) gives an (seq_len, half) grid of rotation angles, one per (position, dimension-pair) combination.
3. Duplicate the angles across both halves of head_dim with torch.cat([angles, angles], dim=-1) before taking cos/sin, so they broadcast directly against the full-width input.
4. Build rotate_half(x) by splitting the last dimension in half, [x1, x2], and returning torch.cat([-x2, x1], dim=-1) — this is the "90°-rotated" companion vector the rotation formula needs.
5. Apply the rotation-matrix identity elementwise: x * cos + rotate_half(x) * sin.
6. Sanity-check position 0: its angle is always 0 for every dimension pair, so cos = 1, sin = 0 and the output at position 0 must equal the input exactly (no rotation).

The key insight is that RoPE never touches the attention score formula itself — it rotates q and k before the dot product, and the rotation angles are constructed so that rotate(q_i) · rotate(k_j) algebraically depends only on i - j, giving relative-position awareness for free.

Reference solution

import torch


def apply_rotary_pos_emb(x: torch.Tensor, base: float = 10000.0) -> torch.Tensor:
    batch, heads, seq, hd = x.shape
    half = hd // 2

    # geometrically decaying frequencies: dimension pair p rotates at
    # base^(-2p/hd), so low pairs rotate fast (fine-grained position) and
    # high pairs rotate slowly (coarse position) -- like sinusoidal encodings
    # applied as a rotation instead of an additive vector
    freqs = 1.0 / (base ** (torch.arange(0, half, dtype=torch.float32, device=x.device) / half))
    pos = torch.arange(seq, dtype=torch.float32, device=x.device)
    angles = torch.outer(pos, freqs)  # (seq, half)

    # duplicate each angle across both halves of head_dim so cos/sin
    # broadcast directly against the full-width x and rotate_half(x)
    cos = torch.cat([angles.cos(), angles.cos()], dim=-1)  # (seq, hd)
    sin = torch.cat([angles.sin(), angles.sin()], dim=-1)  # (seq, hd)

    x1, x2 = x[..., :half], x[..., half:]
    rotate_half = torch.cat([-x2, x1], dim=-1)

    # standard rotation-matrix identity applied elementwise per dimension pair
    return x * cos + rotate_half * sin

Key Functions & Tricks

• base ** (torch.arange(0, half) / half) — produces the geometric frequency schedule; a vector op, not a python loop over dimension pairs.
• torch.outer(pos, freqs) — forms the full (seq_len, half) angle grid via an outer product in one call instead of broadcasting by hand.
• torch.cat([angles, angles], dim=-1) — the trick that lets one (seq, half) angle grid drive a full (seq, head_dim) elementwise multiply without a second frequency computation.
• torch.cat([-x2, x1], dim=-1) — the rotate_half primitive; note the sign flip is on the first half of the output, not the second.
• torch.polar(abs, angle) / torch.complex(real, imag) — used only in the independent reference to cross-check via complex-number rotation, an alternative formulation worth knowing about even though the rotate-half version is what production code typically uses.

How to Recognize This Pattern