Positional encoding

Positional encoding is how a language model keeps track of word order. The model looks at all the words at once, so without extra information "the dog bites the man" and "the man bites the dog" would look much alike. Positional encoding adds each word's place, and the method affects how well a model handles texts longer than its training texts, such as a long contract.

Why does a transformer need it?

The transformer has no recurrence, so Vaswani et al. (2017) had to "inject some information about the relative or absolute position of the tokens in the sequence". Attention compares every token with every other token, and that comparison by itself does not look at where the tokens are. Models that write left to right are a partial exception: each token only sees the tokens before it, and Haviv et al. (2022) found that such models pick up a sense of position even without any positional encoding and stay competitive.

Which methods are there?

• Sinusoidal encoding. The original transformer adds a fixed pattern of sine and cosine waves to each token's embedding, with wavelengths from 2π to 10,000 × 2π. The authors chose it because it "may allow the model to extrapolate" to longer texts; learned position vectors gave nearly identical results in their tests.
• Learned position embeddings. The model learns one vector per position. BERT (Google, 2018) trained its position embeddings on sequences of up to 512 tokens.
• Rotary position embedding (RoPE). Su et al. (2021) rotate the query and key vectors by an angle that depends on the position, so the attention score between two tokens depends on their distance. Meta's LLaMA (2023) replaced absolute position embeddings with RoPE in every layer.
• ALiBi. Press, Smith and Lewis (2021) add no position vectors at all: each attention score gets a penalty proportional to the distance between the two tokens. A 1.3-billion-parameter model trained on 1,024 tokens handled 2,048 as well as a sinusoidal model trained on 2,048, and trained 11% faster. BLOOM (BigScience, 2022) uses it and found that it also made training smoother.

How do models reach longer contexts?

Running a model on texts longer than in training mostly fails. Press et al. found that a sinusoidal model trained on 512 tokens improved for up to 20 extra tokens and got worse after 50; rotary models degraded too. For RoPE models the usual route is to rescale the positions and train briefly on long texts. Position Interpolation (Chen et al., Meta, 2023) squeezes the positions into the trained range and extended LLaMA to 32,768 tokens within 1,000 fine-tuning steps. YaRN (Peng et al., 2023) extended Llama 2 to a context window of 128K tokens with a few hundred steps. Meta raised the RoPE base frequency of Llama 3 to 500,000 and grew its context from 8K to 128K tokens in six stages, with about 800 billion training tokens.

Example

RoPE works on pairs of numbers in the query and key vectors. This snippet turns one pair per position and shows that, for the same query and key, the score depends only on the distance between the two tokens:

1
import math
2
 
3
def rotate(x, y, pos, theta=0.5):
4
    # RoPE for one pair of dimensions: turn the vector by pos * theta
5
    a = pos * theta
6
    return x * math.cos(a) - y * math.sin(a), x * math.sin(a) + y * math.cos(a)
7
 
8
def score(q, k, q_pos, k_pos):
9
    qx, qy = rotate(*q, q_pos)
10
    kx, ky = rotate(*k, k_pos)
11
    return round(qx * kx + qy * ky, 3)
12
 
13
q, k = (1.0, 0.5), (0.3, 2.0)
14
print(score(q, k, 2, 5), score(q, k, 10, 13))  # distance 3 twice: -1.753 -1.753
15
print(score(q, k, 2, 6))                       # distance 4: -2.223