AI Engineer Exhaustive
MediumPlusLLaMA

Rotary Positional Embeddings

LLaMA

Medium

RoPE position encoding via rotation

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Rotary positional embeddings (RoPE; Su et al., 2021) encode absolute position by rotating pairs of features, so that a query–key inner product depends on relative offset. Llama 3 uses RoPE on every attention layer. The herd paper’s only new knob versus Llama 2 is the base frequency (\theta=500{,}000) (Table 3; §3.2 cites Xiong et al. 2023 for long-context behavior up to 32K). The rotation algebra is Su’s; the official apply step is apply_rotary_emb in llama/model.py.

Values are not rotated. Positions come from the precomputed freqs_cis table (separate note), sliced at start_pos.

How it works

Split each head vector of width (d_h) into (d_h/2) two-dimensional pairs. Llama’s pairing is consecutive coordinates ((x_0,x_1),;(x_2,x_3),;\ldots), implemented by reshape(..., -1, 2) then view_as_complex. Pair (i) at position (m) is multiplied by a unit complex number:

$$ \theta_i = \theta^{-2i/d_h},\qquad \begin{pmatrix} x'{2i} \ x'{2i+1} \end{pmatrix}

\begin{pmatrix} \cos(m\theta_i) & -\sin(m\theta_i) \ \sin(m\theta_i) & \cos(m\theta_i) \end{pmatrix} \begin{pmatrix} x_{2i} \ x_{2i+1} \end{pmatrix}. $$

Equivalently, if (z_i = x_{2i} + \mathrm{i}, x_{2i+1}) and (\omega_{m,i}=e^{\mathrm{i} m\theta_i}),

z'_i = z_i \cdot \omega_{m,i}.

apply_rotary_emb does exactly that product in (\mathbb{C}^{d_h/2}), then view_as_real and flatten back to (d_h). Queries and keys use the same (\omega) slice. The broadcast helper (reshape_for_broadcast) expects freqs_cis to already be length (S) on the sequence axis and (d_h/2) on the complex axis, and inserts singleton dims so it multiplies ((B,S,n_{\mathrm{heads}},d_h/2)) without a Python loop.

After rotation, the real inner product (\langle q'_m, k'_n\rangle) equals the pairwise-rotated product of the unrotated vectors, which depends on (m-n) (Su et al., §3). That is why RoPE can replace learned (P) added to tokens: relative phase is baked into (QK^\top).

Llama 3 specifics, specified vs inferred:

A common alternate layout (“rotate-half”) takes the first (d_h/2) dims as the real parts and the last (d_h/2) as the imaginary parts. That is a different pairing than reshape(..., 2). Checkpoints trained with the official pairing will not match rotate-half without a shuffle.

Official code

llama/model.py — apply_rotary_emb and reshape_for_broadcast. Called from Attention.forward on xq and xk after the QKV view and before the cache write. The table itself is precompute_freqs_cis (see the frequency-table note). Hugging Face conversions sometimes reimplement this with explicit cos/sin and a rotate_half; those are equivalent only if the pair layout matches.

Watch-outs

Sources