Where Am I? Position in Attention, from Sinusoids to RoPE

Attention, as The Transformer’s Three Attention Patterns derives it, compares every token with every other through dot products of queries and keys. Nothing in that computation knows where a token is. For language that’s a serious gap: “the model trains the data” and “the data trains the model” contain the same tokens. This post is about how position gets into a Transformer, and about the one method nearly every current language model uses, rotary position embeddings (RoPE), derived from the property that makes it work.

Table of contents

Attention can’t see order

Take attention as The Transformer’s Three Attention Patterns defines it, with no position information at all, and permute its input tokens. The query–key scores are permuted the same way in both directions, the softmax is taken row by row, and the weighted sums follow: every output is exactly the output the token had before, moved to its new place. Attention is permutation-equivariant. Checked, along with the other facts this post relies on:

import numpy as np

# 1. Attention without positions can't tell order apart: permute the tokens,
#    and every output is permuted the same way; nothing else changes. With a
#    causal mask that's no longer true: each token sees a different prefix.
rng = np.random.default_rng(0)
T, H = 5, 8
X = rng.normal(size=(T, H))
Wq, Wk, Wv = (rng.normal(size=(H, H)) / np.sqrt(H) for _ in range(3))
def attend(X, causal=False):
    s = (X @ Wq) @ (X @ Wk).T / np.sqrt(H)
    if causal:
        s = np.where(np.tri(T, dtype=bool), s, -np.inf)
    w = np.exp(s - s.max(1, keepdims=True)); w /= w.sum(1, keepdims=True)
    return w @ (X @ Wv)
perm = rng.permutation(T)
print(f"attention(permuted input) == permuted attention(input): "
      f"{np.allclose(attend(X[perm]), attend(X)[perm])}")
print(f"the same, with a causal mask: "
      f"{np.allclose(attend(X[perm], causal=True), attend(X, causal=True)[perm])}")

# 2. RoPE: rotate each pair of dimensions (2i, 2i+1) by angle pos * theta_i.
def rope(x, pos, base=10_000.0):
    half = x.shape[-1] // 2
    theta = base ** (-np.arange(half) / half)             # one frequency per pair
    ang = pos * theta
    x1, x2 = x[..., 0::2], x[..., 1::2]
    out = np.empty_like(x)
    out[..., 0::2] = x1 * np.cos(ang) - x2 * np.sin(ang)
    out[..., 1::2] = x1 * np.sin(ang) + x2 * np.cos(ang)
    return out

q, k = rng.normal(size=H), rng.normal(size=H)
for m, n in [(3, 5), (10, 12), (100, 102), (3, 10)]:
    print(f"query at {m:3d}, key at {n:3d} (distance {n - m:2d}): score {rope(q, m) @ rope(k, n):+.4f}")
print(f"rotation keeps length: |q| = {np.linalg.norm(q):.4f}, |rope(q, 57)| = {np.linalg.norm(rope(q, 57)):.4f}")

Full script: permutation_rope.py

attention(permuted input) == permuted attention(input): True
the same, with a causal mask: False
query at   3, key at   5 (distance  2): score -6.3341
query at  10, key at  12 (distance  2): score -6.3341
query at 100, key at 102 (distance  2): score -6.3341
query at   3, key at  10 (distance  7): score -5.9618
rotation keeps length: |q| = 4.4482, |rope(q, 57)| = 4.4482

The first line is the problem, for attention that sees the whole sequence. The second is a twist: a language model’s attention is causal, each token sees only the tokens before it, and that breaks the symmetry. Each token sees a different number of predecessors, and a model can roughly count them, for instance by averaging evenly over everything before it, a signal that gets weaker as the prefix grows. Language models trained with no position encoding at all do learn position this way (Haviv et al., 2022), and the model with no positions below uses it. The rest of the output is the solution the rest of this post builds up to.

Adding position to the tokens

The original Transformer (Vaswani et al., 2017) adds a position vector to each token’s embedding before the first layer. Theirs was fixed, made of sinusoids at different frequencies:

so that each position gets a unique pattern, from fast-changing dimensions that tell neighbors apart to slow ones that change over thousands of positions. GPT-2 instead learns one vector per position, a second embedding table. Both are absolute: they label each token with where it is, and the model has to work out from two labels how far apart two tokens are. And both only cover the positions seen in training. A learned table simply has no trained row for position 5,000 if training never went past 4,096.

Penalizing distance: ALiBi

Most of what attention needs is relative: how far back is that word? Learned relative-position terms inside attention go back to Shaw, Uszkoreit and Vaswani (2018), and T5 uses a simple version. ALiBi (Press, Smith and Lewis, 2022) is simpler still, and adds no position vectors at all. For a query at position $m$ and a key at position $n \le m$, it subtracts a penalty in proportion to their distance from the attention score, $s - a\,(m - n)$, with a different slope $a$ for each head: a geometric sequence, $\tfrac12, \tfrac14, \dots, \tfrac{1}{256}$ for 8 heads, so that some heads look nearby and others far. The penalty is defined at any distance, so the model runs on longer sequences without retraining, though its steepest heads effectively ignore distant tokens. And the penalty ignores content and only grows with distance: it can make attention prefer recent tokens, but on its own it can’t point at one exact distance.

RoPE: rotate, so only the difference counts

RoPE (Su et al., 2021) asks for a different property: a relative-position effect that interacts with what the query and key contain, so that the query–key score depends on the two positions only through their difference, but in a way the model can shape. Here’s the construction. Split the query and key vectors into pairs of dimensions, and treat each pair as a point in a plane. Before taking the dot product, rotate the query at position $m$ by an angle $m\theta$, and the key at position $n$ by $n\theta$. For one pair, with $R(\alpha)$ the rotation by $\alpha$:

because rotations compose by adding angles and the transpose of a rotation is the rotation backwards. The absolute positions $m$ and $n$ have disappeared; only $n - m$ is left. The output above shows it: a query and key two positions apart score exactly $-6.3341$ whether they sit at 3 and 5, 10 and 12, or 100 and 102. And a rotation doesn’t change a vector’s length, so RoPE changes only the angle between query and key, not their sizes.

Each pair $i = 0, \dots, H/2 - 1$ rotates at its own frequency, $\theta_i = 10{,}000^{-2i/H}$ for head dimension $H$, the same range as the sinusoids: the first pairs turn fast and distinguish neighbors, the last turn slowly and measure long distances. Only queries and keys are rotated, never the values. It adds no parameters, and costs a few multiplications per number in each query and key.

One trap for anyone loading real weights: which dimensions form a pair is a convention. The code here pairs $2i$ with $2i + 1$, as the RoFormer paper does; many implementations, including Hugging Face’s Llama (rotate_half), pair dimension $i$ with $i + H/2$. The two are the same up to a fixed reordering of dimensions, which the trained weights absorb, but mixing the conventions with pretrained weights silently breaks the model.

Drag the two positions below, then press “Move both by +5”: the score changes only when the distance does.

A test: two tasks, four times the length

To see what these choices do, train the same tiny Transformer with each of them on sequences of 32 tokens, then test on 32, 64 and 128. The model: 2 layers of causal attention with 1 head, width $D = 64$, an MLP and LayerNorm, AdamW at $10^{-3}$, batch 64, a fixed 2,000 steps, 2 seeds, tested on 256 fresh sequences. With one head, ALiBi gets one slope, so it runs twice: at $\tfrac12$, the steepest of the usual eight, which keeps attention close to recent tokens, and at $\tfrac{1}{16}$. RoPE also runs with a base of 1,000,000, for the section after this one. Two tasks:

• lag: at every position, output the token from three positions back, from a vocabulary of 16 (chance: 1 in 16, 6%). It needs exact relative position, always the same short distance.
• seen: at every position, say whether this token appeared anywhere earlier, from a vocabulary of 96. It needs to look back over every distance, including distances longer than any training sequence.

Each run is one call of train_position.py with a scheme, a seed and a task, on the CPU; all_runs.py makes the 28 calls, which take about 6 minutes in all on an Apple M4 laptop. Besides accuracy at each length, each line gives accuracy at length 128 on the positions past 32 only, where the model is beyond anything it trained on, and a test-time tweak explained below.

The schemes enter the model in two places. The absolute encodings are added to the token embeddings, in the model’s forward:

    def forward(self, x):
        h = self.emb(x)
        T = x.shape[1]
        if self.kind == "sinusoid":
            h = h + sinusoids(T, D)
        if self.kind == "learned":
            h = h + self.pos(torch.arange(T))                   # rows past 32 are never trained
        for layer in self.layers:
            h = layer(h)
        return self.out(h)

RoPE and ALiBi act inside each attention layer, on the scores:

    def forward(self, x):
        T = x.shape[1]
        q, k, v = self.qkv(self.n1(x)).chunk(3, -1)
        if self.kind == "rope":
            q, k = rope(q, self.arg), rope(k, self.arg)
        s = SCALE * q @ k.transpose(1, 2) / math.sqrt(D)
        if self.kind == "alibi":                                # a penalty growing with distance
            dist = torch.arange(T)[:, None] - torch.arange(T)[None]
            s = s - self.arg * dist.clamp(min=0)
        s = s.masked_fill(torch.ones(T, T).triu(1).bool(), float("-inf"))
        x = x + self.o(torch.softmax(s, -1) @ v)
        return x + self.mlp(self.n2(x))

Full script: train_position.py

all_runs.py calls it for both tasks, all seven schemes and both seeds:

Full script: all_runs.py

lag  none         seed 0: length 32 / 64 / 128: 0.522 / 0.339 / 0.240; at 128, positions past 32: 0.158; scores x log128/log32: 0.239
lag  none         seed 1: length 32 / 64 / 128: 0.526 / 0.341 / 0.243; at 128, positions past 32: 0.163; scores x log128/log32: 0.243
lag  sinusoid     seed 0: length 32 / 64 / 128: 1.000 / 0.539 / 0.315; at 128, positions past 32: 0.113; scores x log128/log32: 0.317
lag  sinusoid     seed 1: length 32 / 64 / 128: 1.000 / 0.531 / 0.307; at 128, positions past 32: 0.095; scores x log128/log32: 0.301
lag  learned      seed 0: length 32 / 64 / 128: 1.000 / 0.520 / 0.292; at 128, positions past 32: 0.078; scores x log128/log32: 0.291
lag  learned      seed 1: length 32 / 64 / 128: 1.000 / 0.516 / 0.293; at 128, positions past 32: 0.079; scores x log128/log32: 0.291
lag  rope         seed 0: length 32 / 64 / 128: 1.000 / 1.000 / 1.000; at 128, positions past 32: 1.000; scores x log128/log32: 1.000
lag  rope         seed 1: length 32 / 64 / 128: 1.000 / 1.000 / 0.999; at 128, positions past 32: 0.999; scores x log128/log32: 1.000
lag  rope:1000000 seed 0: length 32 / 64 / 128: 1.000 / 1.000 / 1.000; at 128, positions past 32: 1.000; scores x log128/log32: 1.000
lag  rope:1000000 seed 1: length 32 / 64 / 128: 1.000 / 1.000 / 1.000; at 128, positions past 32: 1.000; scores x log128/log32: 1.000
lag  alibi:0.5    seed 0: length 32 / 64 / 128: 0.940 / 0.939 / 0.939; at 128, positions past 32: 0.938; scores x log128/log32: 0.878
lag  alibi:0.5    seed 1: length 32 / 64 / 128: 0.942 / 0.936 / 0.936; at 128, positions past 32: 0.938; scores x log128/log32: 0.867
lag  alibi:0.0625 seed 0: length 32 / 64 / 128: 0.698 / 0.634 / 0.605; at 128, positions past 32: 0.587; scores x log128/log32: 0.645
lag  alibi:0.0625 seed 1: length 32 / 64 / 128: 0.699 / 0.636 / 0.602; at 128, positions past 32: 0.582; scores x log128/log32: 0.643
seen none         seed 0: length 32 / 64 / 128: 0.997 / 0.959 / 0.808; at 128, positions past 32: 0.747; scores x log128/log32: 0.960
seen none         seed 1: length 32 / 64 / 128: 0.997 / 0.963 / 0.813; at 128, positions past 32: 0.754; scores x log128/log32: 0.960
seen sinusoid     seed 0: length 32 / 64 / 128: 1.000 / 0.901 / 0.725; at 128, positions past 32: 0.638; scores x log128/log32: 0.715
seen sinusoid     seed 1: length 32 / 64 / 128: 1.000 / 0.892 / 0.723; at 128, positions past 32: 0.633; scores x log128/log32: 0.715
seen learned      seed 0: length 32 / 64 / 128: 0.999 / 0.760 / 0.638; at 128, positions past 32: 0.524; scores x log128/log32: 0.641
seen learned      seed 1: length 32 / 64 / 128: 0.999 / 0.727 / 0.617; at 128, positions past 32: 0.498; scores x log128/log32: 0.627
seen rope         seed 0: length 32 / 64 / 128: 0.999 / 0.936 / 0.820; at 128, positions past 32: 0.762; scores x log128/log32: 0.844
seen rope         seed 1: length 32 / 64 / 128: 1.000 / 0.942 / 0.813; at 128, positions past 32: 0.753; scores x log128/log32: 0.845
seen rope:1000000 seed 0: length 32 / 64 / 128: 0.999 / 0.968 / 0.870; at 128, positions past 32: 0.827; scores x log128/log32: 0.892
seen rope:1000000 seed 1: length 32 / 64 / 128: 0.999 / 0.967 / 0.837; at 128, positions past 32: 0.785; scores x log128/log32: 0.877
seen alibi:0.5    seed 0: length 32 / 64 / 128: 0.933 / 0.877 / 0.752; at 128, positions past 32: 0.689; scores x log128/log32: 0.785
seen alibi:0.5    seed 1: length 32 / 64 / 128: 0.938 / 0.881 / 0.755; at 128, positions past 32: 0.698; scores x log128/log32: 0.791
seen alibi:0.0625 seed 0: length 32 / 64 / 128: 0.991 / 0.942 / 0.839; at 128, positions past 32: 0.790; scores x log128/log32: 0.652
seen alibi:0.0625 seed 1: length 32 / 64 / 128: 0.993 / 0.940 / 0.831; at 128, positions past 32: 0.780; scores x log128/log32: 0.663

(For “seen”, always answering “no” would score 85% at length 32, 73% at 64 and 55% at 128, and 45% on the positions past 32 at length 128, since more earlier tokens mean more repeats. These follow from the vocabulary: a token at position $k$ is new with probability $(95/96)^k$, averaged over the positions scored.)

On lag, the relative task, the approaches separate cleanly.

• No position information: 52% at the training length, far above the 6% of chance, then falling with length, to 16% past position 32. That’s the causal mask’s counting signal: sharp near the start, blurring as the prefix grows.
• Sinusoidal and learned, the absolute encodings: perfect at length 32, then 52–54% at 64 and 29–32% at 128. Those totals are a mixture. A causal model computes the first 32 positions of any sequence exactly as in training, so it gets them right; past 32 it’s at 8–11%, near the 6% of chance. For the learned table the reason is plain: rows 32 to 127 were never used in training and are still their random initial values. The sinusoids fail just as badly, although their rows are well defined: the model has learned what the labels it saw mean, and nothing about the ones it didn’t.
• RoPE: perfect at every length, at either base. The task only needs a distance of 3, and RoPE’s score depends on nothing but the distance, which is the same at position 100 as at position 10.
• ALiBi: 94% at every length with the steep slope, since its penalty is defined at any distance and doesn’t degrade. But a penalty that only grows with distance can’t single out “exactly three back”. With the gentle slope, which barely separates distance 3 from 2 or 4, it gets only 70% even at the training length.

On seen, every method degrades with length, and the no-position model degrades as much as RoPE or less: 81% at 128 against 81–82%, and 84–88% against 81–82% on a rerun on an NVIDIA RTX 3090 (its saved outputs). It has no angles to go out of range, so most of this drop isn’t about position at all. With four times as many earlier tokens, the one matching token gets a smaller share of the softmax, and its signal is diluted. The last column tests that: it multiplies every attention score by $\log 128 / \log 32$ at test time, changing nothing learned (the factor is 1 at the training length), which keeps the softmax as sharp over 128 tokens as it was over 32. That idea, scaling attention scores by the log of the length, is Chiang and Cholak’s (2022). It brings the no-position model back to 96% (98% on the rerun). RoPE recovers much less, to 84%, and 88–89% with the larger base: part of RoPE’s loss is position, slow pairs reaching angles training never showed. The rest:

• ALiBi at the steep slope is at 93% even at the training length, only 8 points above always answering “no”, because its penalty hides repeats that are far back. With the gentle slope it’s among the best at 128 without rescaling, 83–84%, level with RoPE at the larger base (84–87%); rescaling, which sharpens the content scores against a fixed penalty, hurts it.
• The absolute encodings do worst, 62–73% at 128. Past position 32 the learned table is at 50–53%, barely above always “no” (45%).

So extending context is two problems. One is position: angles or labels the model never saw. The other is attention over more tokens than training ever showed it, which hits even a model with no positions. Decoder models with no position encoding generalize to longer lengths surprisingly well in careful comparisons (Kazemnejad et al., 2023), and this task is an example; on “lag” they fail.

The frequencies, and longer contexts

How far each pair of dimensions can “see” is set by its wavelength, the number of positions for one full turn, $2\pi/\theta_i$:

import numpy as np

# Each pair of dimensions rotates at its own speed; its wavelength is how many
# positions a full turn takes. Head dimension 128 (64 pairs), for two bases:
# 10,000 (the original) and 500,000 (Llama 3).
H = 128
i = np.arange(H // 2)
print(f"{'pair':>5s} " + "".join(f"{f'base {b:,}':>18s}" for b in (10_000, 500_000)))
for p in (0, 8, 16, 32, 48, 63):
    print(f"{p:5d} " + "".join(f"{2 * np.pi * b ** (p / (H // 2)):>18,.1f}" for b in (10_000, 500_000)))
for b in (10_000, 500_000):
    wl = 2 * np.pi * b ** (i / (H // 2))
    print(f"base {b:,}: {np.sum(wl > 8192):2d} of 64 pairs never finish one turn within 8,192 positions, "
          f"{np.sum(wl > 131_072):2d} within 131,072")