Tutorial 05: Rotary Position Embeddings (RoPE)

Attention, by itself, is permutation-invariant — it has no idea which token came first. Without position information, the model cannot distinguish "the cat sat on the mat" from "the mat sat on the cat." Rotary Position Embeddings (RoPE) solve this by encoding position as a rotation in the query and key vectors.

---

1. The Problem: Attention Has No Sense of Order

The attention formula softmax(Q @ K^T) @ V only sees the dot products between queries and keys. If you permute the sequence, the attention weights change, but the model does not know that the permutation happened — it just sees different numbers.

Early transformers added position embeddings directly to the input:

x = token_embedding + position_embedding

This works, but has a limitation: the maximum sequence length must be fixed at training time. If the model was trained with positions 0-2047, it cannot handle position 2048 at inference time without extrapolation.

RoPE takes a fundamentally different approach: encode position into the attention computation itself by rotating the query and key vectors.

---

2. The Core Idea: Rotation as Position Encoding

Consider a 2D vector [x1, x2]. Rotating it by angle θ gives:

[x1', x2'] = [x1·cos(θ) - x2·sin(θ),  x1·sin(θ) + x2·cos(θ)]

The key insight: if you rotate query q by angle θ_m (for position m) and key k by angle θ_n (for position n), then their dot product depends only on the relative position m - n:

q · k = f(q, k, m - n)

This means the attention score between two tokens depends on how far apart they are, not their absolute positions. The model naturally learns relative position patterns.

---

3. The Mathematics

3a. Frequency Computation

For a head dimension of d, we divide it into d/2 pairs. Each pair gets a different rotation frequency:

frequency_i = 1 / (base^(2i/d))    for i = 0, 1, ..., d/2 - 1

where base is typically 10,000 (or 1,000,000 for some models like Qwen3).

Lower-indexed pairs rotate slowly (capturing long-range position differences), while higher-indexed pairs rotate quickly (capturing fine-grained position differences). This is analogous to the Fourier transform's frequency decomposition.

3b. Angle Computation

For a token at position pos, the rotation angle for pair i is:

angle[pos, i] = pos * frequency_i

This creates a 2D table of shape (max_seq_len, d/2):

         pair_0    pair_1    pair_2    ...
pos=0:   0*freq_0  0*freq_1  0*freq_2
pos=1:   1*freq_0  1*freq_1  1*freq_2
pos=2:   2*freq_0  2*freq_1  2*freq_2
...

3c. Rotation Application

Given a query (or key) vector x of dimension d, we:

1. Split it into pairs: [x_0, x_1], [x_2, x_3], ..., [x_{d-2}, x_{d-1}]
2. For each pair, apply a 2D rotation by the corresponding angle

The rotation for a pair [x_a, x_b] by angle θ is:

x_a' = x_a · cos(θ) - x_b · sin(θ)
x_b' = x_a · sin(θ) + x_b · cos(θ)

This is exactly a multiplication by the 2D rotation matrix:

[cos(θ)  -sin(θ)] [x_a]
[sin(θ)   cos(θ)] [x_b]

---

4. Precomputing the Cosine and Sine Tables

Since the frequencies and angles are deterministic (no learnable parameters), we precompute them once at initialization:

import mlx.core as mx


class RoPE:
    def __init__(self, dims: int, seq_len: int, base: int = 10000):
        """
        Args:
            dims: Head dimension (must be even — we need pairs).
            seq_len: Maximum sequence length to precompute for.
            base: Frequency base. Higher = slower rotation = better for long
                  sequences. Qwen3 uses 1,000,000; Llama uses 10,000.
        """
        assert dims % 2 == 0, "dims must be even"
        self.dims = dims
        self.seq_len = seq_len
        half_dims = dims // 2

        # Compute the frequency for each pair dimension.
        # freq_i = 1 / base^(2i / dims)
        #
        # This can be rewritten as:
        # freq_i = base^(-2i / dims)
        #
        # We use arange to get i = [0, 1, 2, ..., half_dims - 1],
        # then divide by half_dims to get exponents in [0, 1).
        inner = mx.arange(0, half_dims, dtype=mx.float32) / half_dims
        freqs = mx.power(base, -inner)

        # Compute the angle for each (position, pair) combination.
        # t is the position index: [0, 1, 2, ..., seq_len - 1]
        # angles[pos, pair] = pos * freq_pair
        t = mx.arange(seq_len)
        freqs = mx.outer(t, freqs)  # shape: (seq_len, half_dims)

        # Precompute cos and sin — the actual rotation factors
        self.cos_freqs = mx.cos(freqs)  # shape: (seq_len, half_dims)
        self.sin_freqs = mx.sin(freqs)  # shape: (seq_len, half_dims)

Why base^(-2i/dims)?

This creates a geometric progression of frequencies. Pair 0 rotates at the slowest rate (highest frequency index → smallest angle increment per position), and pair half_dims-1 rotates at the fastest rate.

The ratio between consecutive frequencies is 1/base^(2/dims), which is constant. This uniform spacing in log-frequency space ensures the model can capture position differences at all scales.

---

5. Applying the Rotation

Here is the rotation applied to a single query/key vector:

    def __call__(self, x: mx.array, offset=None) -> mx.array:
        """
        Apply rotary position embeddings to query or key tensor.

        Args:
            x: Query or key tensor. Shape: (batch, seq_len, num_heads, head_dim)
            offset: Position offset. Tells RoPE what absolute positions these
                    tokens correspond to. This is critical for KV cache — during
                    decode, we only process 1 new token at a time, but it is
                    at position (cache_len + step), not position 0.

        Returns:
            Rotated tensor. Same shape as input.
        """
        N, S, H, D = x.shape
        # N = batch, S = seq_len, H = num_heads, D = head_dim

        # Look up the cos/sin values for the relevant positions
        if offset is not None:
            cos_basis = self.cos_freqs[offset, :]
            sin_basis = self.sin_freqs[offset, :]
        else:
            cos_basis = self.cos_freqs[:S, :]
            sin_basis = self.sin_freqs[:S, :]

        # Split the head dimension into two halves: "real" and "imaginary"
        # This is the "complex number" interpretation of rotation.
        x1 = x[..., :self.half_dims]      # first half of each pair
        x2 = x[..., self.half_dims:]      # second half of each pair

        # Reshape the rotation bases to broadcast correctly:
        # cos_basis: (seq, half_dims) → (1, seq, 1, half_dims)
        cos_basis = cos_basis.reshape(-1, S, 1, self.half_dims)
        sin_basis = sin_basis.reshape(-1, S, 1, self.half_dims)

        # Apply the 2D rotation using complex number multiplication:
        #   (x1 + i·x2) × (cos + i·sin) = (x1·cos - x2·sin) + i·(x1·sin + x2·cos)
        real = mx.multiply(x1, cos_basis) - mx.multiply(x2, sin_basis)
        imag = mx.multiply(x2, cos_basis) + mx.multiply(x1, sin_basis)

        # Recombine the two halves
        y = mx.concat([real, imag], axis=-1)
        return y.reshape(N, S, H, D).astype(x.dtype)

The Complex Number Interpretation

We can think of each pair [x_a, x_b] as a complex number z = x_a + i·x_b. The rotation is then a complex multiplication:

z' = z × e^(iθ) = z × (cos θ + i sin θ)

This is why we split the vector into halves — each half is the real and imaginary part of a complex number, and the rotation is just complex multiplication.

---

6. The Offset Mechanism for KV Cache

During autoregressive generation, we process one token at a time. The first token is at position 0, the second at position 1, and so on. When we add a new token to the KV cache, it should be rotated by its absolute position, not by 0.

This is what offset handles:

# During prefill (processing the full prompt):
# positions = [0, 1, 2, 3, 4]
# offset = None → RoPE uses positions 0 through 4

# During decode (generating one token at a time):
# Step 1: new token at position 5 → offset = 5
# Step 2: new token at position 6 → offset = 6
# ...

In the attention layer:

# From Qwen3Attention.__call__ — applying offset to RoPE
rope_offsets = offsets  # could be int, list, or mx.array
projection_q = self.rope(projection_q, offset=rope_offsets)
projection_k = self.rope(projection_k, offset=rope_offsets)

Both Q and K are rotated with the same offset for a given position. This ensures the relative-position property holds: the dot product q·k depends on the difference in positions, not the absolute positions.

---

7. Why Not Learned Position Embeddings?

Approach	Pros	Cons
Absolute (add to input)	Simple	Fixed max length, no relative position
Learned absolute	Flexible	Still fixed max length, more parameters
Relative (T5-style)	Relative position aware	Complex, extra parameters
RoPE	No parameters, relative position, length generalization	Requires paired dimensions

RoPE has become the standard because: 1. Zero parameters — the rotation is purely deterministic. 2. Relative position — attention scores naturally depend on position differences. 3. Length generalization — with minor tricks (YaRN, NTK-aware scaling), RoPE can extrapolate beyond the training length. 4. Efficient — the precomputed tables are tiny; the rotation is a few multiplies.

---

8. Summary

• RoPE encodes position by rotating query and key vectors.
• The rotation is decomposed into d/2 independent 2D rotations, each at a different frequency.
• The precomputed cos/sin tables have shape (max_seq_len, head_dim/2).
• The offset mechanism allows the KV cache to apply correct rotations during autoregressive decode.
• The dot product q·k depends on relative position, enabling the model to learn distance-based patterns.

Next: Tutorial 06 — Scaled Dot-Product Attention