Rotary position embeddings

Worked solutions

Problem 1

1. $R(\alpha)^\top R(\alpha) = \begin{pmatrix}\cos\alpha & \sin\alpha\\ -\sin\alpha & \cos\alpha\end{pmatrix}\begin{pmatrix}\cos\alpha & -\sin\alpha\\ \sin\alpha & \cos\alpha\end{pmatrix} = \begin{pmatrix}\cos^2\alpha + \sin^2\alpha & 0\\ 0 & \sin^2\alpha + \cos^2\alpha\end{pmatrix} = I$.Multiply out; the off-diagonal entries are $-\cos\alpha\sin\alpha + \sin\alpha\cos\alpha = 0$, and $\cos^2 + \sin^2 = 1$.
2. $\det R(\alpha) = \cos^2\alpha + \sin^2\alpha = 1$.$ad - bc$ with $b = -\sin\alpha$.
3. $R(\alpha)R(\beta) = \begin{pmatrix}\cos\alpha\cos\beta - \sin\alpha\sin\beta & -\cos\alpha\sin\beta - \sin\alpha\cos\beta\\ \sin\alpha\cos\beta + \cos\alpha\sin\beta & -\sin\alpha\sin\beta + \cos\alpha\cos\beta\end{pmatrix} = R(\alpha + \beta)$.The entries are the addition formulas: $\cos(\alpha + \beta) = \cos\alpha\cos\beta - \sin\alpha\sin\beta$ and $\sin(\alpha + \beta) = \sin\alpha\cos\beta + \cos\alpha\sin\beta$.
4. $R(\alpha)^\top = \begin{pmatrix}\cos\alpha & \sin\alpha\\ -\sin\alpha & \cos\alpha\end{pmatrix} = \begin{pmatrix}\cos(-\alpha) & -\sin(-\alpha)\\ \sin(-\alpha) & \cos(-\alpha)\end{pmatrix} = R(-\alpha)$.$\cos$ is even and $\sin$ is odd.
5. $R(\alpha)^\top R(\alpha) = I$, $\det R(\alpha) = 1$, $R(\alpha)R(\beta) = R(\alpha + \beta)$, $R(\alpha)^\top = R(-\alpha)$A rotation preserves lengths and angles (orthogonal, determinant $+1$ so no reflection), rotations compose by adding angles, and the inverse of a rotation is its transpose, which is the rotation back. Step 3 also shows rotations in the plane commute: $R(\alpha)R(\beta) = R(\beta)R(\alpha)$.

Problem 2

1. $(R(m\theta)q)^\top(R(n\theta)k) = q^\top R(m\theta)^\top R(n\theta)k$.$(Ab)^\top = b^\top A^\top$.
2. $R(m\theta)^\top = R(-m\theta)$.Problem 1, step 4.
3. $R(-m\theta)R(n\theta) = R(n\theta - m\theta) = R\big((n - m)\theta\big)$.Problem 1, step 3: angles add.
4. $(R(m\theta)q)^\top(R(n\theta)k) = q^\top R\big((n - m)\theta\big)k$The absolute positions $m$ and $n$ have cancelled: the score is the unrotated $q$ against $k$ rotated by the relative offset $n - m$ times the frequency. Nothing was assumed about $q$ or $k$, so this holds for whatever the projections produce.

Problem 3

1. $R(\varphi)k = (k_1\cos\varphi - k_2\sin\varphi,\ k_1\sin\varphi + k_2\cos\varphi)^\top$.Matrix times vector.
2. $s(\varphi) = q_1k_1\cos\varphi - q_1k_2\sin\varphi + q_2k_1\sin\varphi + q_2k_2\cos\varphi = (q_1k_1 + q_2k_2)\cos\varphi + (q_2k_1 - q_1k_2)\sin\varphi$.Dot with $q$ and collect the $\cos$ and $\sin$ terms.
3. $a = q_1k_1 + q_2k_2 = q^\top k$ and $b = q_2k_1 - q_1k_2$.Read off step 2. $a$ is the dot product; $b$ is the 2-D cross product of $k$ and $q$ (the signed area of their parallelogram).
4. $a^2 + b^2 = q_1^2k_1^2 + 2q_1k_1q_2k_2 + q_2^2k_2^2 + q_2^2k_1^2 - 2q_2k_1q_1k_2 + q_1^2k_2^2 = (q_1^2 + q_2^2)(k_1^2 + k_2^2)$.Expand both squares; the cross terms cancel and the remaining four terms factor.
5. Let $r = \|q\|\|k\|$ and choose $\psi$ with $r\cos\psi = a$ and $r\sin\psi = b$. Then $s(\varphi) = r(\cos\psi\cos\varphi + \sin\psi\sin\varphi) = r\cos(\varphi - \psi)$.Step 4 says $(a, b)$ lies on the circle of radius $r$, so such a $\psi$ exists (when $r > 0$); then the difference formula for cosine (the trig-identities page).
6. $\psi$ is the angle from $k$ to $q$.With $k$ along the first axis, $k = (\|k\|, 0)^\top$ and $q = \|q\|(\cos\psi', \sin\psi')^\top$ at angle $\psi'$: then $a = \|q\|\|k\|\cos\psi'$ and $b = q_2k_1 = \|q\|\|k\|\sin\psi'$, so $\psi = \psi'$. The dot product is unchanged by rotating both vectors, so this holds in general.
7. $s(\varphi) = (q^\top k)\cos\varphi + (q_2k_1 - q_1k_2)\sin\varphi = \|q\|\|k\|\cos(\varphi - \psi)$, $\psi$ the angle from $k$ to $q$As the offset $n - m$ grows, the score of one pair oscillates between $\pm\|q\|\|k\|$ with period $2\pi/\theta$, and is largest when the rotation by $\varphi = (n - m)\theta$ brings $k$ into line with $q$. The full score (Problem 4) is a sum of $d/2$ such sinusoids with different periods, which is what lets it single out particular offsets.

Problem 4

1. $R_m^\top R_n = \operatorname{blockdiag}\big(R(m\theta_p)^\top R(n\theta_p)\big)_p = \operatorname{blockdiag}\big(R((n - m)\theta_p)\big)_p = R_{n-m}$.Block-diagonal matrices multiply block by block, and the transpose of a block-diagonal matrix is block-diagonal with transposed blocks; each block is Problem 2 with $\theta = \theta_p$.
2. $(R_mq)^\top(R_nk) = q^\top R_m^\top R_nk = q^\top R_{n-m}k$.$(Ab)^\top = b^\top A^\top$ and step 1.
3. $q^\top R_{n-m}k = \sum_{p=1}^{d/2}q^{(p)\top}R\big((n - m)\theta_p\big)k^{(p)}$.A block-diagonal matrix acts on each pair separately, and the dot product adds up the pairs.
4. Replacing $(m, n)$ by $(m + s, n + s)$ leaves $n - m$ unchanged, so the score is unchanged.Step 3 contains $m$ and $n$ only in the combination $n - m$.
5. $\|R_mq\|^2 = q^\top R_m^\top R_mq = q^\top R_0q = q^\top q$.Step 1 with $n = m$, and $R_0 = I$ since $R(0) = I$.
6. $R_m^\top R_n = R_{n-m}$; $(R_mq)^\top(R_nk) = \sum_p q^{(p)\top}R((n - m)\theta_p)k^{(p)}$, a function of $n - m$ alone; $\|R_mq\| = \|q\|$The score depends on relative position, and so do the attention weights built from it (Problem 9): a sequence shifted by $s$ positions attends the same way. Each pair contributes $\|q^{(p)}\|\|k^{(p)}\|\cos((n - m)\theta_p - \psi_p)$ by Problem 3, and the rotation never changes the size of a query or key, so the $1/\sqrt d$ scaling argument of the attention page is untouched.

Problem 5

1. $e^{\mathrm{i}\alpha}(x + \mathrm{i}y) = (\cos\alpha + \mathrm{i}\sin\alpha)(x + \mathrm{i}y) = (x\cos\alpha - y\sin\alpha) + \mathrm{i}(x\sin\alpha + y\cos\alpha)$.Euler's formula, then multiply out using $\mathrm{i}^2 = -1$.
2. The real and imaginary parts are the two entries of $R(\alpha)(x, y)^\top$.Compare with $R(\alpha)$ applied to $(x, y)^\top$: first entry $x\cos\alpha - y\sin\alpha$, second $x\sin\alpha + y\cos\alpha$.
3. $z\bar w = (x_1 + \mathrm{i}y_1)(x_2 - \mathrm{i}y_2) = (x_1x_2 + y_1y_2) + \mathrm{i}(y_1x_2 - x_1y_2)$, so $\operatorname{Re}(z\bar w) = x_1x_2 + y_1y_2$.Multiply out; the real part is the dot product of the two pairs, and the imaginary part is Problem 3's $b$.
4. $\tilde z_p = e^{\mathrm{i}m\theta_p}z_p$ and $\tilde w_p = e^{\mathrm{i}n\theta_p}w_p$ are the rotated pairs.Step 2 with $\alpha = m\theta_p$ and $n\theta_p$.
5. $\tilde z_p\overline{\tilde w_p} = e^{\mathrm{i}m\theta_p}z_p\,e^{-\mathrm{i}n\theta_p}\bar w_p = z_p\bar w_p\,e^{\mathrm{i}(m - n)\theta_p}$.$\overline{e^{\mathrm{i}\beta}w} = e^{-\mathrm{i}\beta}\bar w$, and exponents add.
6. Rotation by $\alpha$ is multiplication by $e^{\mathrm{i}\alpha}$; $q^{(p)\top}k^{(p)} = \operatorname{Re}(z_p\bar w_p)$; score $= \operatorname{Re}\sum_p z_p\bar w_p\,e^{\mathrm{i}(m - n)\theta_p}$Step 3 applied to the rotated pairs and summed over $p$. This is how RoPE is usually implemented: view the last axis as complex, multiply by a precomputed table of $e^{\mathrm{i}m\theta_p}$, view as real again. The relative-position property is the exponent rule $e^{\mathrm{i}m\theta}e^{-\mathrm{i}n\theta} = e^{\mathrm{i}(m - n)\theta}$, which is Problem 1's $R(\alpha)R(\beta) = R(\alpha + \beta)$ in different clothes.

Problem 6

1. $\theta_1 = 10000^{0} = 1$ and $\theta_{64} = 10000^{-126/128}$.$p = 1$ gives exponent $0$; $p = 64$ gives $-2\cdot 63/128 = -126/128$.
2. $10000^{126/128} = e^{(126/128)\ln 10000} = e^{0.984375\cdot 9.2103} \approx e^{9.0664} \approx 8660$, so $\theta_{64} \approx 1.15\times10^{-4}$.$\ln 10000 = 4\ln 10 \approx 9.2103$.
3. Wavelength of pair $1$: $2\pi/1 \approx 6.28$ positions; of pair $64$: $2\pi\cdot 8660 \approx 54{,}400$ positions.A pair rotates by $\theta_p$ per position, so it completes a turn every $2\pi/\theta_p$ positions.
4. Pair $p$ completes a turn within $4096$ positions iff $2\pi/\theta_p \le 4096$ iff $10000^{2(p-1)/128} \le 4096/(2\pi)$.Substitute $\theta_p$ and invert; the inequality direction is preserved because $\theta_p > 0$.
5. $\dfrac{2(p - 1)}{128}\ln 10000 \le \ln\dfrac{4096}{2\pi} \approx \ln 651.9 \approx 6.480$, so $p - 1 \le \dfrac{128\cdot 6.480}{2\cdot 9.2103} \approx 45.03$.Take logs; $4096/(2\pi) \approx 651.9$.
6. $\theta_1 = 1$, $\theta_{64} \approx 1.15\times10^{-4}$; wavelengths $\approx 6.3$ and $\approx 54{,}400$ positions; $46$ of the $64$ pairs complete a turn within $4096$ positions$p \le 46.03$, so $p = 1, \dots, 46$. The remaining $18$ pairs turn through less than a full circle over the whole context and act like a slowly varying, almost monotone position signal; the fast pairs resolve nearby offsets. Trained at a shorter context, those slow pairs have never seen the angles a longer context produces, which is the problem position interpolation addresses (Problem 10).

Problem 7

1. $s = q^\top R_{n-m}k$.Problem 4, step 2.
2. $\nabla_qs = R_{n-m}k$.$s = q^\top c$ with $c = R_{n-m}k$ not depending on $q$, and $\nabla_q(q^\top c) = c$ (the matrix-calculus page).
3. $s = k^\top R_{n-m}^\top q = k^\top R_{m-n}q$, so $\nabla_ks = R_{m-n}q$.A scalar equals its transpose; $R_{n-m}^\top = R_{m-n}$ by Problem 1, step 4, block by block.
4. $\nabla_qL = g\,R_{n-m}k$ and $\nabla_kL = g\,R_{m-n}q$.$L$ depends on $q$ and $k$ only through $s$, so the chain rule multiplies by $g$.
5. $\tilde q = R_mq$ gives $\partial\tilde q/\partial q = R_m$, so $\nabla_qL = R_m^\top\nabla_{\tilde q}L = R_m^\top u = R_{-m}u$.A linear map's Jacobian is its matrix, and for a scalar $L$, $\nabla_qL = (\partial\tilde q/\partial q)^\top\nabla_{\tilde q}L$ (the Jacobians page); then $R_m^\top = R_{-m}$.
6. $\nabla_qL = g\,R_{n-m}k$, $\nabla_kL = g\,R_{m-n}q$; in general $\nabla_qL = R_{-m}\nabla_{\tilde q}L$: rotate the upstream gradient back by the token's own angleStep 4 is step 5 applied to $\nabla_{\tilde q}L = g\tilde k = gR_nk$: $R_{-m}R_nk = R_{n-m}k$. The backward pass of RoPE is RoPE with the sign of the position flipped, which in the complex form is multiplication by $e^{-\mathrm{i}m\theta_p}$, the conjugate table. It costs the same as the forward pass and needs no stored activations beyond the positions.

Problem 8

1. $\tfrac{d}{d\varphi}R(\varphi) = \begin{pmatrix}-\sin\varphi & -\cos\varphi\\ \cos\varphi & -\sin\varphi\end{pmatrix}$.Differentiate each entry.
2. $R(\varphi + \pi/2) = \begin{pmatrix}\cos(\varphi + \pi/2) & -\sin(\varphi + \pi/2)\\ \sin(\varphi + \pi/2) & \cos(\varphi + \pi/2)\end{pmatrix} = \begin{pmatrix}-\sin\varphi & -\cos\varphi\\ \cos\varphi & -\sin\varphi\end{pmatrix}$.$\cos(\varphi + \pi/2) = -\sin\varphi$ and $\sin(\varphi + \pi/2) = \cos\varphi$ (the trig-identities page). The two matrices agree.
3. $\dfrac{\partial}{\partial\Delta}q^{(p)\top}R(\Delta\theta_p)k^{(p)} = \theta_p\,q^{(p)\top}R(\Delta\theta_p + \pi/2)k^{(p)}$.Chain rule: the angle $\Delta\theta_p$ has derivative $\theta_p$ with respect to $\Delta$, and step 2 gives the derivative of $R$ with respect to its angle. $q^{(p)}$ and $k^{(p)}$ are constants here.
4. $\partial s/\partial\Delta = \sum_p\theta_p\,q^{(p)\top}R(\Delta\theta_p + \pi/2)k^{(p)}$.Sum step 3 over the pairs.
5. $\theta_p$ appears only in pair $p$'s term, and the angle $\Delta\theta_p$ has derivative $\Delta$ with respect to $\theta_p$.Each frequency belongs to one block.
6. $\partial s/\partial\theta_p = \Delta\,q^{(p)\top}R(\Delta\theta_p + \pi/2)k^{(p)}$.Chain rule as in step 3 with the roles of $\Delta$ and $\theta_p$ swapped.
7. $\tfrac{d}{d\varphi}R(\varphi) = R(\varphi + \pi/2)$; $\partial s/\partial\Delta = \sum_p\theta_p\,q^{(p)\top}R(\Delta\theta_p + \pi/2)k^{(p)}$; $\partial s/\partial\theta_p = \Delta\,q^{(p)\top}R(\Delta\theta_p + \pi/2)k^{(p)}$For one pair, by Problem 3 with $\varphi \to \varphi + \pi/2$, the derivative is $-a\sin\varphi + b\cos\varphi$, the derivative of $a\cos\varphi + b\sin\varphi$. The frequency gradient carries a factor $\Delta$: distant pairs of tokens push the frequencies hardest, which is why learned or fine-tuned frequencies are sensitive to the longest offsets in the training data.

Problem 9

1. $\nabla_{\tilde Q}L = \tfrac1{\sqrt d}(\nabla_SL)\tilde K$ and $\nabla_{\tilde K}L = \tfrac1{\sqrt d}(\nabla_SL)^\top\tilde Q$.The attention page's Problems 5 and 6, with $\tilde Q$, $\tilde K$ in place of $Q$, $K$: the attention computation is unchanged, it only receives rotated inputs.
2. Row $m$ of $\tilde Q$ depends on $q_m$ alone, through $\tilde q_m = R_mq_m$.Each token is rotated by its own position; no other token's query enters its row.
3. $\nabla_{q_m}L = R_m^\top\,(\text{row } m \text{ of } \nabla_{\tilde Q}L)^\top = R_{-m}\,(\nabla_{\tilde Q}L)_{m,:}^\top$.Problem 7, step 5, applied to row $m$: the upstream gradient at $\tilde q_m$ is row $m$ of $\nabla_{\tilde Q}L$, written as a column.
4. $\nabla_{k_n}L = R_{-n}\,(\nabla_{\tilde K}L)_{n,:}^\top$.The same argument for keys, with $\tilde k_n = R_nk_n$.
5. $S_{mn} = (R_mq_m)^\top(R_nk_n)/\sqrt d$ depends on $m$ and $n$ only through $n - m$.Problem 4, step 4, for every entry.
6. Shifting all positions by $s$ leaves every $S_{mn}$ unchanged, hence $A$ and $O = AV$.$(n + s) - (m + s) = n - m$; $A$ is a function of $S$ and $V$ is not rotated.
7. $\nabla_{q_m}L = R_{-m}\big(\tfrac1{\sqrt d}(\nabla_SL)\tilde K\big)_{m,:}^\top$, $\nabla_{k_n}L = R_{-n}\big(\tfrac1{\sqrt d}(\nabla_SL)^\top\tilde Q\big)_{n,:}^\top$: each row of the attention page's gradient rotated back by its own position; $A$ and $O$ are invariant to a common shift of all positionsSo RoPE is a layer wrapped around attention: forward, rotate $Q$ and $K$ row by row; backward, rotate the gradients back row by row and hand the rest to the attention page. The shift invariance means a RoPE model sees no absolute position at all; the only asymmetry it has is the causal mask, which tells a token how many keys lie before it.