LM10-T2 — Rotary Position Embeddings

Learning objectives

After this lesson, you should be able to:

• Explain why attention needs position information and what absolute learned embeddings cannot do.
• Derive that the dot product of two rotated vectors depends only on their position difference.
• Write the rotation for a head of width d_h with the frequency schedule θ_i = base^(−2i/d_h).
• Implement and test RoPE in PyTorch, and name its common bugs.

The idea: position as an angle

Self-attention computes a score q_m · k_n for every pair of positions m (the query) and n (the key). Nothing in that product knows where m and n are, so without help the model cannot tell “the dog bit the man” from “the man bit the dog”. Position must be injected.

The simplest injection, used in Module 4, adds a learned vector p_m to the token embedding once at the input. This works, but it ties the model to absolute positions: the vector for position 1,000 was never trained if the training windows ended at 1,024 tokens of every document, and what the network needs is usually relative: “the word two places back”, not “the word at position 817”.

Rotary embeddings make position relative by construction. Look at the plane in the figure. A two-dimensional vector rotated by angle α and another rotated by angle β have a dot product that depends on the difference α − β (and on their lengths), not on the angles themselves. So if the query at position m is rotated by mθ and the key at position n by nθ, their score depends on (m − n)θ: the offset. Shift the whole sentence by 100 places and every score is unchanged.

Many planes, many speeds

One angle per token would only encode distance modulo a full turn, and a single speed is too coarse or too fine. RoPE therefore cuts the head vector of width d_h into d_h/2 two-dimensional planes and rotates plane i at its own speed θ_i. The first plane turns fast (one radian per token) and tells neighbours apart; the last turns very slowly (with the usual base of 10,000, in the order of 10⁻⁴ rad per token) and distinguishes tokens that are thousands of positions apart. It is the same trick as the digits of a clock: seconds, minutes and hours together give a unique time.

Where it is applied

Only to queries and keys, in every layer, immediately before the dot product. Values are not rotated: their job is to carry content, not to find it. The token embeddings carry no position at all, and there is no position table, hence no position parameters.

Why rotate and not add?

The mathematics: rotation, relativity and frequencies

Take one plane: a pair (x₁, x₂) of entries of a query or key. Rotating it by an angle φ is multiplication by a matrix:

Rotations preserve length (|x′| = |x|) and compose by adding angles: R(α)^T R(β) = R(β − α). Now rotate the query at position m by mθ and the key at position n by nθ:

This is the whole proof. The score between positions depends on the content (q, k) and on the offset (n − m)θ, never on m or n alone.

The full head

For a head of width d_h (even), plane i = 0 … d_h/2 − 1 has frequency

d_h

head width, for example 128; the head has d_h/2 planes.

θ_i

rotation per token of plane i; θ₀ = 1 radian, θ_{d_h/2 − 1} ≈ base^(−1).

m

absolute position of the token (0, 1, 2, …).

cos(mθ_i), sin(mθ_i)

the tables computed once for all positions: shape (n_pos, d_h/2).

Implementation convention. Pair entry i with entry i + d_h/2 (the “half-split” layout, used in the code below): with x₁ = x[…, :d_h/2] and x₂ = x[…, d_h/2:],

Another layout pairs neighbouring entries (2i, 2i + 1). Both are rotations; they differ by a fixed permutation of the entries. A model trained with one layout must be run with the same one, because the query and key projections have learned to put information in the matching places. This is the most common source of conversion bugs.

Worked numbers. With d_h = 4 the two frequencies are θ₀ = 1 and θ₁ = 0.01. At position 2 the angles are 2 and 0.02 rad, the cosines −0.4161 and 0.9998, the sines 0.9093 and 0.0200. A vector (1, 1, 0, 0) at position 0 is unchanged; at position 2 its first plane (1, 0) becomes (cos 2, sin 2) = (−0.4161, 0.9093).

Long contexts. At a position never seen in training the slow planes show angles the model never met, and quality degrades. Practical remedies (position interpolation, a larger base, other rescalings, usually with some fine-tuning) rescale positions or frequencies so that long inputs map into familiar angle ranges. They change the tables, not the architecture, and the extended behaviour must be evaluated, not assumed.

A worked example: implement and test RoPE

The code builds the tables, rotates queries and keys and checks the three properties that define the method: position 0 is the identity, lengths are preserved, and scores depend only on the offset. It also shows that without RoPE the score between identical content does not change with position, which is exactly the blindness described above.

import torch

def rope_tables(hd, npos, base=10000.0):
    theta = base ** (-torch.arange(0, hd, 2).float() / hd)            # (hd/2,)
    angles = torch.arange(npos).float()[:, None] * theta[None, :]    # (npos, hd/2)
    return torch.cos(angles), torch.sin(angles)

def apply_rope(x, cos, sin):                                          # x: (B, H, T, hd)
    half = x.shape[-1] // 2
    x1, x2 = x[..., :half], x[..., half:]
    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)

torch.manual_seed(0)
hd, T = 8, 16
cos, sin = rope_tables(hd, T)
print([round(v, 4) for v in cos[2, :2].tolist()], [round(v, 4) for v in sin[2, :2].tolist()])   # [-0.4161, 0.9998] [0.9093, 0.02]

q = torch.randn(1, 1, 1, hd).expand(1, 1, T, hd)        # the SAME query content at every position
k = torch.randn(1, 1, 1, hd).expand(1, 1, T, hd)        # the SAME key content at every position

plain = (q @ k.transpose(2, 3))[0, 0]
print(torch.allclose(plain[5, 3], plain[9, 3]))          # True: without positions the score ignores where tokens are

qr, kr = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
scores = (qr @ kr.transpose(2, 3))[0, 0]
print(torch.allclose(qr.norm(dim=-1), q.norm(dim=-1), atol=1e-5))     # True: rotations keep lengths
print(torch.allclose(qr[:, :, 0], q[:, :, 0]))                        # True: position 0 is the identity
print(torch.allclose(scores[5, 3], scores[9, 7], atol=1e-5))          # True: both offsets are 2
print(torch.allclose(scores[5, 3], scores[5, 2], atol=1e-3))          # False: offset 2 versus 3

The first check passes because the content is the same at every position and no position enters the product; this is why the model needs position information at all. The next three confirm that the rotation adds position without changing lengths, and that the score is a function of the offset. Notebook LM10-L3 asks you to write these functions yourself and then use them inside an attention layer with a KV cache.

What the frequencies look like

Engineering implications and limitations

Match the layout exactly. Half-split or interleaved pairs, the base, and the head width must equal those used in training. A silent mismatch gives a model that runs and produces plausible but degraded text; test against reference logits from the original implementation on a short prompt.

Rotate before caching. In decoding, rotate each new key by its own absolute position once and store the rotated key (Lesson LM10-T5). Rotating cached keys again, or rotating the new token at position 0, are classic bugs; Notebook LM10-L3 includes a test that catches both.

Compute the tables once. The cos and sin tables depend only on the positions and the head width. Precompute them for the maximum context in the working precision and slice them with the start position.

Mind the precision. At large positions the angle mθ is a large number, and in 16-bit floating point its rounding error becomes a significant angle error. Compute the angles in 32 bits and cast the tables afterwards.

Treat longer-than-trained context as a new experiment. After rescaling the positions, evaluate on long-document tasks and also re-check short-context quality; some rescalings trade one for the other.

No new parameters. RoPE changes nothing about parameter counts or fine-tuning; LoRA adapters on the query and key projections work as before.

Common misconceptions

Engineering takeaways

1. Attention without position information cannot see order; RoPE adds it by rotating q and k in d_h/2 planes by mθ_i.
2. The score ⟨R(mθ)q, R(nθ)k⟩ = qᵀR((n−m)θ)k depends only on the offset n−m.
3. Frequencies θ_i = base^(−2i/d_h) give fast planes for neighbours and slow planes for long ranges; there are no position parameters.
4. Apply it to queries and keys only, in every layer, with the layout and base used in training; cache the rotated keys.
5. Longer contexts than trained require rescaled positions and fresh evaluation.

References and further reading

Lesson LM3-T4 introduced absolute and sinusoidal positions. The numbers in the table above were computed from the formula with base 10,000 and d_h = 128.

• Su, J. et al. (2021). RoFormer: Enhanced transformer with rotary position embedding. arXiv:2104.09864.
• Vaswani, A. et al. (2017). Attention is all you need (sinusoidal positions). arXiv:1706.03762.
• Touvron, H. et al. (2023). Llama 2: Open foundation and fine-tuned chat models. arXiv:2307.09288.
• Su et al. (2021), RoFormer (research paper)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: multi-query and grouped-query attention, sharing keys and values to shrink the cache.

Related content

Nothing below is required to finish this lesson.

• Before
  LM10-T1: Decoder-Only Models Then and Now

  Compare a minimal GPT with modern open models and separate what was built earlier from what is new.

• Lab
  LM10-L3: Implement a KV Cache and Rotary Embeddings in PyTorch

  Add rotary embeddings and a KV cache to attention and verify that cached decoding equals full recomputation.