Vol. VI · Positions and Long Context Plate 6 of 10context extension, position interpolation, NTK scaling, YaRN, RoPE base, long-context fine-tuningAL

Stretching the Context.

Slow the rotations so long texts look like short ones, then train briefly on long texts. Four times the context for a sliver of the original compute.

RoPE fails past its training length because its rotations reach angles the model never saw. The family of fixes keeps the angles familiar. Position interpolation (Chen and colleagues, 2023) divides every position by a factor s, so that position 256 is rotated like position 64 used to be: nothing is out of range, but neighbouring tokens are now only a quarter-step apart. NTK-aware scaling instead raises RoPE’s base, which slows the slow planes a lot and the fast ones hardly at all, keeping fine local resolution. YaRN (Peng and colleagues, 2023) combines the two with a temperature correction.

Each works somewhat with no training at all, and much better after a short fine-tune on long sequences. The workshop uses the course’s RoPE model from plate VI.5, trained at 64 tokens, and shows four ways of running it on longer texts: untouched, interpolated by 4, NTK-scaled by 4, and interpolated then fine-tuned for 400 steps at 256 tokens.

Below: held-out loss against evaluation length for the course’s RoPE model, untouched and with three ways of stretching it four-fold. The fine-tuned version cost 400 steps — about a tenth of the original training tokens.

What you will learn
Workshop
The course’s RoPE model (trained at 64 tokens) evaluated at 64–512 tokens: untouched (blue), positions divided by 4 (ochre), base raised for 4× (green), and divided by 4 then fine-tuned for 400 steps at 256 tokens (red).

Slow the clock, then practise

At 256 tokens, four times its training length, the untouched model scores 4.702 nats. Interpolating positions by 4 with no training brings every angle back into range but crushes local resolution — neighbouring tokens are now a quarter-step apart — and it pays for that at short lengths (4.784 at 64 against 4.363 untouched) while helping at long ones (4.620 at 256). NTK scaling, which leaves the fast planes nearly alone, is the best training-free option: 4.246 at 128 against 4.391 untouched, and 4.444 at 256. Interpolation followed by 400 steps of fine-tuning at 256 tokens gives 4.067 at 256 — better than any training-free option — for 10% of the original training tokens.

That is the recipe large models follow, at larger scale. Llama 3 was pretrained at 8 192 tokens and then extended to 128 000 in a dedicated stage with a raised RoPE base and long documents (plate XII.4); YaRN and its relatives are used to stretch open models further at serving time. Extension has limits: a model can be made to accept a long context much more easily than it can be made to use one well — the subject of plate VI.8 — and every extra token adds to the cost measured in plate VI.9.

Position interpolation was found almost simultaneously in 2023 by a Meta team and by an independent developer, who posted on a forum that dividing RoPE positions by 4 let a Llama model read 8 000 tokens after brief fine-tuning.

Terms

Position interpolation
Dividing positions by a factor so a longer text’s angles stay within the trained range.
NTK-aware scaling
Raising RoPE’s base so slow planes stretch while fast planes keep their resolution.
YaRN
A refined context-extension method combining per-frequency scaling with an attention temperature.
Context extension
Training or adjusting a model to accept longer contexts than it was pretrained on.
RoPE models are stretched by slowing their rotations and fine-tuning briefly at the new length: four times the context here for about 10% of the original training tokens.

Sources

  1. Chen, Wong, Chen and Tian, “Extending Context Window of Large Language Models via Positional Interpolation”, arXiv:2306.15595, 2023.
  2. Peng, Quesnelle, Fan and Shippole, “YaRN: Efficient Context Window Extension of Large Language Models”, ICLR 2024 (arXiv 2023).