NNX RoPE Attention

Rotate queries and keys by position

Rotary position embeddings (RoPE) add position information by rotating pairs of features in Q and K before computing attention scores. The rotation angle depends on the token position and the feature pair. V is left unchanged in this exercise.

A rotation preserves the length of each pair. Rotating Q and K by different angles changes their dot product, which changes the attention weights and therefore the weighted average of values.

Pair features and choose frequencies

Let Dh = d_model // num_heads. It must be even. This exercise pairs adjacent features: (0, 1), (2, 3), and so on. For pair index r:

frequency[r] = base ** (-2 * r / Dh)
angle[position, r] = position * frequency[r]

Angles are in radians. Positions are the row indices 0, ..., T - 1. With Dh=4 and base=10000, the two frequencies are 1 and 0.01 radians per position. Pair zero always has frequency one; changing base affects only the later pairs. For base greater than one, later pairs rotate more slowly. Base is a supplied positive number, not a learned parameter.

Some implementations pair the first half of a head with the second half. That convention can represent the same rotations after a matching feature permutation, but swapping conventions without also reordering the relevant projection features changes the model. Use adjacent pairs here.

A complete rotation example

Rotate (a, b) by angle t as (a*cos(t) - b*sin(t), a*sin(t) + b*cos(t)). Stack the two results together before flattening, so they remain interleaved.

import jax.numpy as jnp

def rotate(x, cos, sin):
    a, b = x[..., 0::2], x[..., 1::2]
    return jnp.stack((a * cos - b * sin,
                      a * sin + b * cos), axis=-1).reshape(x.shape)

x = jnp.array([[1.0, 0.0, 0.0, 2.0],
               [1.0, 0.0, 0.0, 2.0]])
frequencies = jnp.array([1.0, 0.01])
angles = jnp.arange(2)[:, None] * frequencies[None, :]
y = rotate(x, jnp.cos(angles), jnp.sin(angles))
assert jnp.allclose(y[0], x[0])
expected = jnp.array([jnp.cos(1.0), jnp.sin(1.0),
                      -2 * jnp.sin(0.01), 2 * jnp.cos(0.01)])
assert jnp.allclose(y[1], expected)
assert jnp.allclose(jnp.sum(y * y, axis=-1), jnp.sum(x * x, axis=-1))

Position zero leaves both pairs unchanged. At position one, the first pair rotates by one radian and the second by 0.01 radians. Neither pair changes length. The attention module applies this operation separately to the projected query and key arrays for every head.

Why the dot product uses relative position

Write R(t) for a 2-D rotation. For fixed content vectors q and k and frequency f, the rotated dot product satisfies:

(R(i*f) q) dot (R(j*f) k) = q dot (R((j-i)*f) k)

The rotations' position contribution depends on j - i. Shifting both positions by the same offset leaves this dot product unchanged. The result still depends on q and k; it is not a function of distance alone. The property holds pair by pair, with each pair's own frequency.

RoPE is not a mask or a monotonic distance penalty. A more distant key can have a larger score. The formula can be evaluated beyond a training sequence length, but good performance at those lengths is not guaranteed. Changing frequency scales also changes scores and needs evaluation.

The diagram uses four feature pairs (head width eight), base 10000, and initial vectors (1, 0). Move the position slider to compare their angles. Actual Q/K pairs can start at different lengths and directions.