RoPE

Take the 2D rotation from Part 3. The query of the token at position $m$ is turned by $m\theta$; the key at position $n$ by $n\theta$. Values aren’t rotated. Nothing is added to the embeddings.

The identity from Part 3: the two rotations combine into one by the difference of the angles. The same pair of tokens gets the same score at positions 3 and 7 as at 103 and 107. Attention learns relative position for free.

With $\theta = 0.4$: the query at position 1 is turned 0.4 rad, the key at position 3 is turned 1.2 rad. Move both tokens along by any number of positions and both arrows turn by the same extra angle, so the angle between them, and their dot product, stays the same.

A head of width $d_h$ is cut into $d_h/2$ pairs, and each pair gets its own speed. Pair 0 turns 1 radian per position; the last pairs turn very slowly. Like the hands of a clock, fast pairs tell nearby positions apart and slow ones cover long distances. The original paper used $b = 10{,}000$; Llama 3 uses $b = 500{,}000$ for longer contexts.

What is $\theta_1$, and how far does pair 1 turn at position $m = 10$?

def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)     # the two halves
    return torch.cat((-x2, x1), dim=-1)

# cos, sin: (T, d_h), one angle per position and pair
q = q * cos + rotate_half(q) * sin
k = k * cos + rotate_half(k) * sin

• Rotate each pair of query and key dimensions by $m\theta_i$, the position times the pair’s speed.
• $\tilde q_m \cdot \tilde k_n$ depends only on $n - m$.
• $\theta_i = b^{-2i/d_h}$; $b = 10{,}000$ originally, 500,000 in Llama 3.
• Code applies it as $x\cos + \operatorname{rotate\_half}(x)\sin$.