Rotary Position Embeddings

RMSNorm fixed the scale of every token vector, but it did the same thing to each one — it has no idea which token in the sequence it is looking at. The next sublayer, attention, is where tokens finally read each other, and it reads them with a dot product. That is the problem. A dot product does not care about order: score every query against every key and you get the same numbers whether the tokens arrived as “the cat sat” or “sat the cat.” Attention, on its own, sees a bag of tokens, not a sequence. Something has to stamp each vector with its position before attention compares them. That something is RoPE.

4.1 Why position has to be injected

Attention scores a query at one position against a key at another by taking their dot product. Nothing in that dot product knows where either token sat. Swap two tokens and their vectors swap with them, but every pairwise score is unchanged — the mechanism is blind to order by construction. So position cannot live in the attention step; it has to be written into the query and key vectors before they meet. The trick is to write it in a way that survives the dot product and turns into something useful: not the absolute position of each token, but the distance between them.

RoPE does this by rotation. It splits each vector’s features into pairs, treats each pair as a point in a 2D plane, and rotates that point by an angle proportional to the token’s position. A token at position 5 is rotated five times as far as a token at position 1. The reason this is the right move, and not just a way to mark position, is what happens when two rotated vectors meet in a dot product — which the end of this chapter returns to.

4.2 What it computes

RoPE comes in two pieces. First, a table of rotation angles, precomputed once for every position the model might see, from forge’s model.py:

def rope_tables(
    head_dim: int, base: float, context_length: int
) -> tuple[torch.Tensor, torch.Tensor]:
    # split-halves convention, the one the released weights were trained in
    inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
    pos = torch.arange(context_length, dtype=torch.float32)
    angles = pos[:, None] * inv_freq[None, :]
    angles = torch.cat([angles, angles], dim=1)
    return torch.cos(angles), torch.sin(angles)

Listing 4.1 · the angle tables. One row per position, one column per feature pair, filled with that pair’s rotation angle — then turned into cosines and sines.

Read it as three steps. inv_freq assigns each feature pair its own frequency: pair zero turns fast, and each later pair turns slower, spanning a wide range of rates. angles = pos[:, None] * inv_freq[None, :] multiplies every position by every frequency, so row p holds the angle each pair is turned through for a token at position p — angle grows with position, exactly the “further along, further rotated” idea. The cat([angles, angles]) duplicates the table across both halves of the head, because forge pairs feature i with feature i + half — the split-halves convention its comment names, the one the released weights were trained in — and both members of a pair rotate by the same angle. cos and sin of that table are all the rotation needs.

The second piece applies the rotation to an actual query or key:

def apply_rope(
    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, offset: int = 0
) -> torch.Tensor:
    _, _, t, head_dim = x.shape
    half = head_dim // 2
    c = cos[offset : offset + t].view(1, 1, t, head_dim).to(x.dtype)
    s = sin[offset : offset + t].view(1, 1, t, head_dim).to(x.dtype)
    x1, x2 = x[..., :half], x[..., half:]
    rotated = torch.cat((-x2, x1), dim=-1)
    return x * c + rotated * s

Listing 4.2 · the rotation itself. x1 and x2 are the two halves of the head; the result rotates each (x1[i], x2[i]) pair by that position’s angle.

This is one 2D rotation, done to every pair at once. For a single pair (a, b) turned by angle θ, a rotation sends it to (a·cos θ − b·sin θ, a·sin θ + b·cos θ). Line by line: x1, x2 are the two halves, so pair i is (x1[i], x2[i]). rotated = cat((-x2, x1)) builds the cross terms — the −x2 is the −sin θ term, the piece that makes this a rotation and not just a blend. Then x * cos + rotated * sin assembles both outputs: the first half becomes x1·cos − x2·sin, the second half x2·cos + x1·sin, which is exactly that rotation applied across every pair. The offset lets a token know its true position even when only the newest token is being processed — the concern of the cache chapter, not this one.

4.3 Correct, to the bit

RoPE gets the same treatment as the last chapter: an independent second implementation of the same rotation, the same angles and vectors fed to both, and a required match. The companion driver builds forge’s angle tables and rotation and the reference’s, rotates one seeded (2, 4, 9, 8) batch of query vectors with each, and prints the largest disagreement:

Exactly zero again. The rotation itself runs as the identical operation in both implementations, and the angle tables — reached by different but equivalent arithmetic — come out bit-for-bit equal for these frequencies, so there is no rounding noise to differ in. That makes the check as sharp as it can be. Any nonzero number here is a real divergence, and the next paragraph makes one appear.

The reason RoPE rotates, rather than adding a position vector, is the property it leaves behind. Rotate the query at position m and the key at position n, then take their dot product, and because a rotation composed with the reverse of another is a rotation by the difference, the score depends only on m − n — the offset between the two tokens, never their absolute places. Attention gets relative position for free, and never has to be told the length of the sequence. That dot product, where the position-stamped queries and keys finally meet, is the next box: grouped-query attention.