Rotary position embeddings: making distance affect attention

How can attention distinguish nearby and distant occurrences of the same word? RoPE rotates query and key coordinates by position. A two-number example lets you inspect the resulting score without treating position as an unexplained label.

Before you begin: You understand dot products, attention queries and keys, and basic sine and cosine.

Rotary position embeddings, usually called RoPE, introduce position by rotating pairs of query and key coordinates. The angle depends on the token position and the coordinate pair's frequency.

Begin with a two-number vector

Rotating a vector [x, y] by angle a produces [x*cos(a) - y*sin(a), x*sin(a) + y*cos(a)]. Its length stays the same. Its direction changes.

Now rotate a query by the angle for position m and a key by the angle for position n. Their dot product depends on the difference between those angles. This makes relative position participate naturally in attention scores while preserving the content-dependent query and key vectors.

# Runnable: Python 3, standard library.
import math

def rotate(vector, angle):
    x, y = vector
    c, s = math.cos(angle), math.sin(angle)
    return [x * c - y * s, x * s + y * c]

def dot(a, b):
    return sum(x * y for x, y in zip(a, b))

q, k = [1.0, 2.0], [3.0, -1.0]
m, n = 0.3, 0.8
left = dot(rotate(q, m), rotate(k, n))
right = dot(q, rotate(k, n - m))
assert math.isclose(left, right)
assert math.isclose(dot(q, q), dot(rotate(q, m), rotate(q, m)))
print(round(left, 4))

This verifies two properties: rotation preserves length, and the relative angle describes the rotated dot product. It does not show that a trained model has learned to use every distance well.

Shift both positions, then only one

Take the example's query angle 0.3 and key angle 0.8. Their difference is 0.5. Add 2 to both and the difference stays 0.5. Add 2 only to the key and it becomes 2.5. Predict which operation preserves the score under this fixed frequency rule.

Add both comparisons as assertions around the local rotation function. Use a non-equality check only for the chosen numerical example, since special vectors or angles can coincide. This experiment connects the relative-position identity to an implementation test rather than asking you to memorize a rotation formula.

Use several frequencies

A real attention head has many coordinates. RoPE rotates coordinate pairs with different frequencies, often following a geometric schedule. Faster rotations distinguish nearby positions in one way; slower rotations provide another scale. The resulting score combines these rotated content components.

Implementation conventions matter. Some code pairs adjacent coordinates; other code uses a split-half layout with an equivalent weight arrangement. A rotation helper copied from another implementation can have the right shape and still be incompatible with the checkpoint. Follow the model's convention, rotary dimension, frequency settings, and dtype behavior.

RoPE is typically applied to queries and keys, not to values in the standard construction. It changes the score calculation rather than simply adding a position vector to every token embedding.

A long address book is not a long-context skill

You can calculate rotations at positions beyond the training range. That mathematical fact does not guarantee reliable retrieval or reasoning at those positions. Frequency scaling, context extension training, and long-context evaluations address different parts of the problem.

Changing a configuration's maximum length alone can produce an endpoint that accepts longer text but answers it poorly. Test evidence at the beginning, middle, and end; include distractors and questions requiring more than one passage. Also measure cache memory and latency as the length grows.

During cached generation, use the correct position offsets for new tokens. Restarting positions at zero for every decoding step silently changes the model's computation. Padding and packed sequences need deliberate position handling as well.

Exercise: move both positions in the example forward by the same angle increment. Should their rotated dot product change?

You now have the main architectural pieces. The next lesson puts them into the larger process of producing a foundation model.

Sources

The rotation construction comes from RoFormer. LLaMA provides one architectural use of RoPE; exact settings must come from the checkpoint being used.

Continue: How Foundation Models Are Trained.