Rotary Position Embedding (RoPE)

Rotate features to represent position

Rotary position embedding (RoPE) introduces position information by rotating pairs of features. In the attention pattern described here, queries and keys are rotated before their dot products are computed; values are unchanged. No position vector is added to the input.

This exercise implements the rotation alone for a matrix x of shape (T, d). Row t is treated as position t, starting at zero. It does not split attention heads, compute attention or maintain a decoding cache.

Rotate adjacent pairs

The feature count d must be positive and even. Group features as (x[0], x[1]), (x[2], x[3]), and so on. Pair index i runs from 0 through d // 2 - 1. Its frequency and position-dependent angle are:

frequency[i] = base^(-2i / d)
angle[t, i] = t * frequency[i]

Here ^ denotes exponentiation; Python uses **. Angles are in radians. For a pair (a, b) and angle phi, use:

out_even = a * cos(phi) - b * sin(phi)
out_odd  = a * sin(phi) + b * cos(phi)

Interleave these results in the original feature order. At position zero every angle is zero, so the first row is unchanged. Rotation preserves each pair's length and therefore each row's norm, up to numerical rounding.

A worked rotation

Take [1, 0, 0, 1] at position 1 with d = 4 and base = 10000. The frequencies are 1 and 10000^(-2/4) = 0.01.

• Pair (1, 0) rotates by 1 radian, giving (cos(1), sin(1)).
• Pair (0, 1) rotates by 0.01 radians, giving (-sin(0.01), cos(0.01)).

The result is approximately [0.540302, 0.841471, -0.010000, 0.999950]. This complete example checks the rotation for both position zero and one:

import jax.numpy as jnp

x = jnp.array([[1., 0., 0., 1.], [1., 0., 0., 1.]])
T, d = x.shape
base = 10000.0
pairs = jnp.arange(d // 2, dtype=jnp.float32)
frequencies = base ** (-2.0 * pairs / d)
positions = jnp.arange(T, dtype=jnp.float32)[:, None]
angles = positions * frequencies[None, :]
cosine, sine = jnp.cos(angles), jnp.sin(angles)
even, odd = x[:, 0::2], x[:, 1::2]
out = jnp.zeros_like(x)
out = out.at[:, 0::2].set(even * cosine - odd * sine)
out = out.at[:, 1::2].set(even * sine + odd * cosine)

assert jnp.array_equal(out[0], x[0])
assert jnp.allclose(
    out[1], jnp.array([0.540302, 0.841471, -0.010000, 0.999950]),
    atol=1e-6,
)
assert jnp.allclose(jnp.sum(out ** 2, axis=1), jnp.sum(x ** 2, axis=1))

What relative position means here

Write the rotation at position m as R(m). For fixed vectors q and k, using the same frequencies and pairing for both gives:

(R(m) q) dot (R(n) k) = q dot (R(n - m) k)

The explicit position dependence is through the difference n - m. The score still depends on the contents of q and k; queries and keys from different tokens need not have equal scores at equal offsets. In a complete model they may already contain information from earlier layers, so this identity is not a promise about the entire model's behavior.

Other implementations pair the first half of the features with the second half. That convention requires a corresponding feature permutation to match this one. Changing the pairing of existing weights without that conversion generally changes the result. Use adjacent pairs here.

The formula can be evaluated at positions beyond a training sequence's length. Useful predictions at those lengths are a separate question. For base > 1, increasing the base lowers frequencies for i > 0; the first frequency stays 1. That change alone does not guarantee better long-context performance.