Rotary Position Embeddings (RoPE)

The problem

Attention on its own ignores the order of the tokens. RoPE (rotary position embedding) adds it by rotating the queries and keys: the vector of the token at position $m$ is turned by an angle proportional to $m$. Then the score of a query and a key depends on how far apart the two tokens are (the theory page shows why). You will write the rotation.

Pairs of numbers, turned like points in the plane

A vector of $d$ numbers is read as $d/2$ pairs of neighbours: $(x_0, x_1)$, $(x_2, x_3)$, and so on. Each pair is a point in the plane, and pair $i$ turns at its own speed $\theta_i$:

At position $m$, pair $i$ is rotated by the angle $m\theta_i$:

where:

• $x_{2i}$ and $x_{2i+1}$ are the two numbers of pair $i$, and $x'$ is the rotated vector;
• $m$ is the position of the token: $0, 1, 2, \dots$;
• base is usually $10000$. The exponent makes $\theta_0 = 1$ (the first pair turns 1 radian per position) and the later pairs turn more and more slowly.

These two lines are the standard rotation of a point around the origin: its distance from the origin does not change, only its direction.

Example. $d = 4$, base $10000$: two pairs, with speeds $\theta_0 = 10000^0 = 1$ and $\theta_1 = 10000^{-2/4} = 0.01$. The same vector $(1, 0, 1, 0)$ at positions 0, 1 and 2:

• Position 0: both angles are 0, nothing turns: $(1, 0, 1, 0)$.
• Position 1: pair 0 turns by 1 radian, pair 1 by $0.01$: $(0.540, 0.841, 1.000, 0.010)$.
• Position 2: pair 0 turns by 2 radians, pair 1 by $0.02$: $(-0.416, 0.909, 1.000, 0.020)$.

For pair 0 at $m = 1$: $(1 \cdot \cos 1 - 0 \cdot \sin 1,\ 1 \cdot \sin 1 + 0 \cdot \cos 1) = (0.540, 0.841)$. Pair 0 turns fast and tells nearby positions apart; pair 1 has barely moved, and changes noticeably only over hundreds of positions.

What to implement

apply_rope(x, positions, base=10000.0): x has shape $(\dots, T, d)$ and positions is an integer array of shape $(T,)$, the position of each row. Return the rotated array, same shape.

Requirements

• Pairs are adjacent dimensions $(0, 1), (2, 3), \dots$, as in the RoFormer paper. (Some libraries pair dimension $i$ with $i + d/2$ instead; that is a different, equally valid layout, but not the one asked for.)
• Raise ValueError if $d$ is odd.
• No Python loops.

Hints

• Build the $(T, d/2)$ table of angles at once: positions as a column times the speeds as a row (broadcasting).
• x[..., 0::2] is the first number of every pair and x[..., 1::2] the second. Compute both rotated halves, then write them back into an output array with the same slices.