9. Rotary Positional Embeddings (RoPE) Applied to Q/K

Problem

Plain attention has no notion of token order — softmax(QKT/√d_k)V is completely permutation-invariant, so position information has to be injected some other way. Rotary positional embeddings (RoPE) do it by rotating each query and key vector by an angle proportional to its sequence position, before the dot product in attention is taken, rather than adding a separate positional vector to the token embedding.

Split each head's dimension into pairs of coordinates and treat each pair as a 2D vector; rotate pair i by an angle theta_i * position, where theta_i shrinks geometrically with i (low pairs rotate fast, high pairs rotate slowly). The payoff: the dot product q_rope(pos_a) . k_rope(pos_b) depends only on the relative offset (pos_a - pos_b), not on absolute positions.

Source: src/9_rotary_positional_embeddings.py

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

>>> x = torch.randn(1, 3, 4)
>>> positions = torch.arange(3)
>>> apply_rope(x, positions).shape
torch.Size([1, 3, 4])
>>> apply_rope(x, torch.zeros(3))[:, 0].allclose(x[:, 0])  # position 0 is unrotated
True

Step-by-Step Approach

1. Compute one inverse frequency per coordinate pair: theta_i = base ** (-2i / head_dim) for i in 0, 1, ..., head_dim/2 - 1 — the same geometric spacing scheme used for sinusoidal position embeddings.
2. Build the angle for every (position, pair) combination via an outer product: freqs = positions[:, None] * inv_freq[None, :], shape (seq_len, head_dim/2), then take cos(freqs) and sin(freqs).
3. Split x's last dimension into its even-indexed and odd-indexed coordinates: x[..., 0::2] and x[..., 1::2] — each (even_i, odd_i) pair is one 2D vector to rotate.
4. Apply the standard 2D rotation matrix to each pair: rotated_even = x_even * cos - x_odd * sin, rotated_odd = x_even * sin + x_odd * cos.
5. Interleave the rotated even/odd coordinates back into the original layout with torch.stack([rotated_even, rotated_odd], dim=-1).flatten(-2).
6. Verify two invariants directly rather than trusting the formula by eye: position 0 should leave x completely unrotated (cos(0)=1, sin(0)=0), and shifting both a query's and a key's position by the same amount should leave their dot product unchanged.

The key insight is that RoPE encodes position as a rotation, not an additive offset, and rotations compose: rotating q by pos_a and k by pos_b before their dot product is mathematically equivalent to rotating one of them by the relative angle (pos_a - pos_b) and leaving the other fixed, which is exactly why the attention score ends up depending only on relative position.

Reference solution

import torch


def apply_rope(x: torch.Tensor, positions: torch.Tensor, base: float = 10000.0) -> torch.Tensor:
    head_dim = x.shape[-1]
    assert head_dim % 2 == 0

    # inverse frequency per coordinate pair -- early pairs rotate fastest
    pair_idx = torch.arange(0, head_dim, 2, device=x.device, dtype=torch.float32)
    inv_freq = 1.0 / (base ** (pair_idx / head_dim))  # (head_dim/2,)

    freqs = positions.to(torch.float32).unsqueeze(-1) * inv_freq.unsqueeze(0)  # (seq_len, head_dim/2)
    cos = torch.cos(freqs)
    sin = torch.sin(freqs)

    x_even = x[..., 0::2]  # (batch, seq_len, head_dim/2)
    x_odd = x[..., 1::2]

    # 2D rotation matrix [[cos, -sin], [sin, cos]] applied per coordinate pair
    rotated_even = x_even * cos - x_odd * sin
    rotated_odd = x_even * sin + x_odd * cos

    return torch.stack([rotated_even, rotated_odd], dim=-1).flatten(-2)

Key Functions & Tricks

• base ** (pair_idx / head_dim) broadcasting — produces the geometric frequency spectrum in one vectorized expression, no per-pair Python loop.
• tensor.unsqueeze(-1) / unsqueeze(0) — sets up the outer-product broadcast between positions and inverse frequencies to get a full (seq_len, head_dim/2) angle grid in one multiply.
• x[..., 0::2] / x[..., 1::2] — strided slicing to split interleaved coordinate pairs without a reshape (note some real implementations instead split into first-half/second-half via x[..., :d/2]/x[..., d/2:], a different but equally valid pairing convention).
• torch.stack([a, b], dim=-1).flatten(-2) — the standard interleave-back-together idiom: stacking on a new last axis then flattening the last two axes recovers alternating a[0], b[0], a[1], b[1], ... order.
• Dot-product invariance check via direct tensor arithmetic — the strongest way to confirm a positional scheme actually encodes relative position rather than trusting the rotation math was transcribed correctly.

How to Recognize This Pattern