Positional encoding

Problem

"The dog bit the man" and "the man bit the dog" contain the same words. A model that cannot tell them apart is useless for language, yet that is exactly what plain attention is. Each token's output is a weighted average of value vectors, with weights from query–key dot products. Shuffle the input tokens and every output is shuffled the same way, with identical values: attention is permutation-equivariant. The embedding table does not help either, because it gives "dog" the same vector wherever it appears.

So order has to be injected. The question is how: add a position signal to each token's vector, learn one, or change the attention scores themselves so they depend on how far apart two tokens are. The choice decides whether a model can read sequences longer than it was trained on, how its KV cache works, and how many parameters position costs.

Intuition

Think of a wall of clocks, each hand turning at a different speed. The fastest hand goes round every few steps, the slowest takes thousands. Read all the hands at once and you know exactly what time it is, and two nearby moments have similar readings. The sinusoidal encoding is that wall of clocks: each pair of dimensions is one hand, a sine and a cosine at its own speed, and the reading is added to the token's embedding.

A learned table is the bluntest fix: give each slot 0, 1, 2, … its own trainable vector, like seat numbers in a theatre. It works, but there are no seats past the last row, so it cannot handle longer inputs.

Rotary embeddings use the clocks differently. Instead of adding a reading, they turn the query and key vectors by an angle proportional to their position. When a query at position m meets a key at position n, both have been turned, and the dot product only sees the difference between the two turns: n − m. Attention then naturally asks "how far back is this word?" rather than "which absolute slot is it in?". ALiBi asks the same question more bluntly: it subtracts a penalty proportional to distance from each attention score, so far-away tokens get less weight.

Each pair of dimensions is a clock hand at its own speed; the sine is shown, and the cosine partner is the same curve shifted a quarter turn.

Mechanics

Why attention ignores order

Let X be the n × d matrix of token vectors and P any n × n permutation matrix. Self-attention computes softmax(XW_Q(XW_K)ᵀ/√d_k) XW_V. Replacing X by PX permutes the rows and columns of the score matrix together, and softmax works row by row, so

Attention(PX) = P · Attention(X)

The outputs move with the inputs and nothing else changes. A causal mask breaks the symmetry partly, since a token can count how many tokens it may see, and Haviv et al. found causal language models without any position signal still learn some positional information that way. Every mainstream architecture adds an explicit signal anyway.

Sinusoidal encoding

Vaswani et al. add a fixed vector to each token embedding (which they first scale by √d):

PE(pos, 2i) = sin(pos · ω_i)    PE(pos, 2i+1) = cos(pos · ω_i)    ω_i = 10000^−2i/d

The wavelengths 2π/ω_i form a geometric progression from 2π up to nearly 10000 · 2π. The paper chose sinusoids over a learned table, which it found performed about the same, hoping for two properties. The first is exact: for any fixed offset k, each (sin, cos) pair at pos + k is the pair at pos rotated by the angle k·ω_i, by the angle-addition identities:

[ sin((pos+k)ω) , cos((pos+k)ω) ] = [ [cos kω, sin kω], [−sin kω, cos kω] ] · [ sin(pos·ω) , cos(pos·ω) ]

So PE(pos + k) is a linear function of PE(pos) that depends only on k, which the paper hoped would let attention learn relative offsets. A consequence is that PE(p) · PE(q) = Σ_i cos(ω_i(p − q)) depends only on the distance between the positions. The second hope, extrapolating to longer sequences than seen in training, did not hold up well: Press et al. measured sinusoidal models degrading quickly past their training length. Because the signal is added to the token vector, query–key products mix content and position terms, and the content terms see absolute position.

Learned absolute positions

BERT and GPT-2 instead learn a second embedding table with one row per position, nn.Embedding(max_positions, d), and add row t to the token at position t. It costs max_positions × d parameters, 786,432 for GPT-2 small's 1,024 × 768 table, and positions past max_positions simply do not exist. Its rows are free to learn whatever the data rewards, which is why it matched sinusoids in the original paper.

Rotary position embedding (RoPE)

Su et al. leave the token vectors alone and rotate the queries and keys inside every attention layer. Split a head's d_k dimensions into d_k/2 pairs and rotate pair i of the vector at position m by the angle m·θ_i, with θ_i = base^−2i/d_k (base 10000 in the paper):

RoPE(x, m)_{2i, 2i+1} = [ [cos mθ_i, −sin mθ_i], [sin mθ_i, cos mθ_i] ] · [ x_2i , x_2i+1 ]

Write R_m for that block-diagonal rotation. Rotations are orthogonal, so R_mᵀ R_n = R_n−m, and the attention score becomes

(R_m q)ᵀ (R_n k) = qᵀ R_mᵀ R_n k = qᵀ R_n−m k

The score depends on the content of q and k and on the offset n − m only, never on m or n alone. Rotation preserves length, so RoPE changes no vector's norm, and it adds no parameters. Values are not rotated, so the output mixes content only. Because each key is rotated by its own position once, a KV cache stores keys already rotated and never touches them again. Some implementations pair dimension i with dimension i + d_k/2 instead of 2i with 2i + 1; that is the same method on a fixed permutation of the dimensions, but weights trained with one layout need the same layout at inference.

Shifting both positions by the same amount spins the whole picture; the angle between q and k, and so their dot product, stays put.

ALiBi: a bias, not an embedding

Press, Smith and Lewis add no position vector at all. Each head h gets a fixed slope s_h, and before the softmax the score of query i against key j ≤ i is lowered in proportion to the distance:

score_ij = q_i · k_j / √d_k − s_h · (i − j)

The slopes are not learned. For 8 heads they are the geometric sequence 1/2, 1/4, …, 1/256; in general, for a power-of-two head count H, the sequence starts at 2^−8/H and has that same ratio. Steep heads look locally; gentle heads see far back. The paper's headline result is training on short sequences and evaluating on longer ones with little loss, since a longer distance just means a larger penalty. In PyTorch it is a float attn_mask for scaled_dot_product_attention, which adds it to the scores, or a score_mod in FlexAttention.

Worked example

The same matrix shape serves every head; only the slope changes, from 1/2 for head 1 down to 1/256 for head 8.

Implementation

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

def sinusoidal_table(max_len, d, base=10000.0):            # (max_len, d), no parameters
    pos = torch.arange(max_len).float()[:, None]
    freq = base ** (-torch.arange(0, d, 2).float() / d)     # (d/2,)
    pe = torch.zeros(max_len, d)
    pe[:, 0::2] = torch.sin(pos * freq)
    pe[:, 1::2] = torch.cos(pos * freq)
    return pe

learned = nn.Embedding(1024, 768)                           # GPT-2-style table: x = tok + learned(arange(T))

def rope_angles(seq_len, head_dim, base=10000.0):
    freq = base ** (-torch.arange(0, head_dim, 2).float() / head_dim)
    angles = torch.outer(torch.arange(seq_len).float(), freq)   # (T, head_dim/2)
    return angles.cos(), angles.sin()

def apply_rope(x, cos, sin):                                # x: (B, H, T, head_dim)
    x1, x2 = x[..., 0::2], x[..., 1::2]                     # the (2i, 2i+1) pairs
    out = torch.stack((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1)
    return out.flatten(-2)                                  # back to (B, H, T, head_dim)

def alibi_bias(n_heads, T):                                 # (H, T, T), for a power-of-two n_heads
    slopes = torch.tensor([2 ** (-8 * (h + 1) / n_heads) for h in range(n_heads)])
    i = torch.arange(T)[:, None]
    j = torch.arange(T)[None, :]
    bias = -slopes[:, None, None] * (i - j)
    return bias.masked_fill(j > i, float("-inf"))           # causal mask folded in

B, H, T, Dh = 2, 8, 16, 64
q, k, v = (torch.randn(B, H, T, Dh) for _ in range(3))
cos, sin = rope_angles(T, Dh)
out_rope = F.scaled_dot_product_attention(apply_rope(q, cos, sin), apply_rope(k, cos, sin), v, is_causal=True)
out_alibi = F.scaled_dot_product_attention(q, k, v, attn_mask=alibi_bias(H, T))  # float mask is added

apply_rope rotates the pairs exactly as the formula above; with a KV cache, slice cos and sin at the new token's position and rotate only the new query and key. The ALiBi mask is a float tensor, so scaled_dot_product_attention adds it to the scores (its reference implementation, linked in the table, does attn_bias = attn_mask + attn_bias); a boolean mask would instead mean keep-or-drop. For long sequences, a FlexAttention score_mod such as lambda score, b, h, q_idx, kv_idx: score - slope[h] * (q_idx - kv_idx) computes the same bias without materializing an L × S tensor.

Tradeoffs

• Default choice today. RoPE dominates open decoder LLMs (the Llama 3 report is one example) because it is relative, parameter-free, cache-friendly and works with fused attention kernels, since it only touches q and k before the kernel.
• Extending RoPE's context. Past the trained length, the low-frequency pairs reach angles never seen in training. Position interpolation (Chen et al.) scales positions down so they stay inside the trained range, then fine-tunes briefly; YaRN (Peng et al.) rescales frequency bands differently. Raising the base, as some long-context models do, slows every rotation.
• ALiBi's bias toward recency is built in: every head penalizes distance, so it suits tasks where nearby context matters most, and it materializes or computes a bias per score, which some fused kernels support only through a mask.
• Learned tables are simple and fine for fixed-length inputs (image patches, a fixed context window), but cost parameters and cannot generalize past their last row.
• Common bugs. Mixing the interleaved and split-half RoPE layouts between training and inference; applying RoPE to values; forgetting to offset positions by the cache length during generation; adding positions after, rather than before, the first layer.

Connections

The Transformer page shows where the encoding enters the original model, the decoder-only LLM page counts GPT-2's learned position table among its parameters, and the KV cache page explains why RoPE'd keys are stored after rotation.

References

• Vaswani et al. Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762
• Su, Lu, Pan, Murtadha, Wen and Liu. RoFormer: Enhanced Transformer with Rotary Position Embedding. 2021. arXiv:2104.09864
• Press, Smith and Lewis. Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. ICLR 2022. arXiv:2108.12409
• Shaw, Uszkoreit and Vaswani. Self-Attention with Relative Position Representations. NAACL 2018. arXiv:1803.02155
• Raffel et al. Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. JMLR 2020. arXiv:1910.10683
• Devlin, Chang, Lee and Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. NAACL 2019. arXiv:1810.04805
• Haviv, Ram, Press, Izsak and Levy. Transformer Language Models without Positional Encodings Still Learn Positional Information. Findings of EMNLP 2022. arXiv:2203.16634
• Chen, Wong, Chen and Tian. Extending Context Window of Large Language Models via Positional Interpolation. 2023. arXiv:2306.15595
• Peng, Quesnelle, Fan and Shippole. YaRN: Efficient Context Window Extension of Large Language Models. ICLR 2024. arXiv:2309.00071
• Llama Team. The Llama 3 Herd of Models. 2024. arXiv:2407.21783
• PyTorch v2.15.0: torch/nn/modules/sparse.py, scaled_dot_product_attention, flex_attention, torch/nn/modules/transformer.py.