Lost in Latent Space

Chapter: Rotary Positional Embeddings

1. Motivation

The attention dot product $\langle q, k \rangle$ is permutation-invariant: shuffling the input tokens does not change which keys match which queries. Transformers need some signal of position. Three approaches have been tried:

1. Absolute positional embeddings. Add a learned or sinusoidal vector $p_i$ to the input embedding at position $i$. Used in the original transformer (Vaswani et al., 2017).
2. Relative positional embeddings. Bias the attention logits by a learned function of the relative offset $i - j$. Used in T5 and Transformer-XL.
3. Rotary positional embeddings (RoPE). Rotate the query and key vectors by an angle that depends on position. Used in PaLM, LLaMA, GPT-NeoX, and most current open-weight models.

This chapter develops RoPE.

2. Definition

Let $q, k \in \mathbb{R}^{d}$ with $d$ even. Partition each into $d/2$ adjacent pairs:

\[q = (q_0, q_1, q_2, q_3, \dots), \quad k = (k_0, k_1, k_2, k_3, \dots)\]

Treat each pair $(q_{2m}, q_{2m+1})$ as a 2D vector. For position $i$, choose a frequency $\theta_m = 10000^{-2m/d}$, and rotate the $m$-th pair by angle $i \theta_m$:

\[R_{i,m} = \begin{pmatrix} \cos(i\theta_m) & -\sin(i\theta_m) \\ \sin(i\theta_m) & \cos(i\theta_m) \end{pmatrix}\]

Apply this rotation to every pair to get the position-encoded query $q^{(i)}$. Do the same to $k$ at position $j$ to get $k^{(j)}$.

3. The key property

Compute the dot product between $q^{(i)}$ and $k^{(j)}$:

\[\langle q^{(i)}, k^{(j)} \rangle = \sum_m (R_{i,m} q_m)^T (R_{j,m} k_m) = \sum_m q_m^T R_{i,m}^T R_{j,m} k_m\]

Because rotations satisfy $R_a^T R_b = R_{b-a}$:

\[\langle q^{(i)}, k^{(j)} \rangle = \sum_m q_m^T R_{j-i,m} k_m\]

The result depends only on the relative offset $j - i$, not on the absolute positions $i$ or $j$. RoPE delivers relative positional information through a transformation that is applied independently at each position — no quadratic-size bias matrix needed.

4. Worked example

Let $d = 2$, so there is one pair. Set $\theta_0 = 1$. Take $q = (1, 0)$, $k = (1, 0)$. At positions $i = 0$ and $j = 2$:

• $q^{(0)} = (1, 0)$ (no rotation).
• $k^{(2)} = (\cos 2, \sin 2) \approx (-0.416, 0.909)$.
• $\langle q^{(0)}, k^{(2)} \rangle = -0.416$.

At positions $i = 5$ and $j = 7$ (same offset of 2):

• $q^{(5)} = (\cos 5, \sin 5) \approx (0.284, -0.959)$.
• $k^{(7)} = (\cos 7, \sin 7) \approx (0.754, 0.657)$.
• $\langle q^{(5)}, k^{(7)} \rangle = 0.284 \cdot 0.754 + (-0.959)(0.657) \approx -0.416$.

Identical, as the theorem promised.

5. Why the geometric frequencies?

The choice $\theta_m = 10000^{-2m/d}$ mirrors the sinusoidal scheme of Vaswani et al. Each pair rotates at a different rate: low-index pairs spin fast (encoding fine positional differences), high-index pairs spin slowly (encoding long-range structure). One model, multiple positional scales, no extra parameters.

6. Test your understanding

1. Show that $R_{i,m}$ is orthogonal, so RoPE preserves vector norms. Why does that matter for attention stability?
2. If you set every $\theta_m = 0$, what happens to the attention pattern? What if every $\theta_m = \pi$?
3. RoPE applied to queries and keys does not mix them with the value vectors. Why is that a deliberate choice?

7. What’s next

The natural extension is RoPE scaling — methods like NTK-aware interpolation, YaRN, and LongRoPE that let a model trained at, say, 4K context generalize to 32K or 128K by remapping the frequencies $\theta_m$ at inference time. Read Peng et al.’s YaRN paper alongside the Su et al. original (2021).