Position-Augmented Training
If a model fails at length because it has never seen large position indices, the cheapest fix is not a longer training sequence but a shuffled or skipped one, showing the model position 90,000 inside a 2,048-token batch.
Extending a model to 128k tokens by training on 128k-token sequences costs attention compute quadratic in the length and activation memory linear in it. The observation that makes this avoidable: what fails at long context is not the model's handling of long content, it is its handling of unseen position indices. Those two things can be separated, and once separated, large positions can be trained cheaply inside short sequences.
Randomised positions
Ruoss et al. sample an ordered subset of a much larger position range and assign it to a short training sequence. A batch of 40 tokens might carry positions \(\{3, 91, 104, \ldots, 1987\}\) in increasing order, so the ordering relation is preserved while the specific indices vary every step. Across 6,000 models on 15 algorithmic tasks the method raised test accuracy on unseen lengths by 12.0 points on average (Ruoss et al., 2023, ACL, arXiv:2305.16843).
The mechanism is straightforward once stated: the model can no longer memorise a mapping from index to behaviour, because the same content appears under many indices. It is forced onto the relative structure, which is the part that transfers.
Positional skip-wise training
PoSE adapts the idea for extending real LLMs. The training window is divided into chunks; each chunk gets a skipping bias added to its position indices, with the biases resampled every example. A 2,048-token window carrying two chunks with biases 0 and 30,000 exposes the model to relative distances near 30,000 without ever materialising a 30,000-token attention matrix (Zhu et al., 2023, PoSE, arXiv:2309.10400).
The chunk boundary is doing the work. Within a chunk, distances are small and natural; across the boundary, they are large and synthetic. The model sees both in every batch.
Where it sits among the alternatives
| Approach | Training cost at target length | Needs long documents | Handles new base frequency |
|---|---|---|---|
| Full-length fine-tuning | quadratic attention, linear memory | yes | yes |
| Position interpolation | short fine-tune at target | yes | yes, by construction |
| PoSE | fixed at original window | no | yes, combines with PI or NTK |
| Randomised positions | fixed | no | pretraining-time only |
PoSE composes with interpolation rather than replacing it: the interpolation scheme decides how indices map to angles, and PoSE decides which indices the model gets to practise on. Randomised encodings are a pretraining decision and cannot be retrofitted to a released checkpoint whose weights already encode a particular index distribution.
When it breaks
Augmentation teaches positions, not long-range dependency. A model that has only ever attended across 2,048 real tokens has never had to build a retrieval circuit that spans 100,000 of them, and no amount of index skipping supplies that experience. This is the boundary these methods do not cross, and it is why Llama 3 still spent roughly 800B tokens on staged long-context pretraining across six stages from 8K to 128K rather than skipping straight to augmentation (Grattafiori et al., 2024, arXiv:2407.21783).
Synthetic distance distributions can mismatch real ones. Uniformly sampled skips give the model plenty of practice at 60,000-token separations and comparatively little at the 3,000 to 8,000 range where most genuine long-document dependencies live, and a model tuned this way can end up worse in the middle of its range than at the ends.
Evaluation is the hard part. Perplexity on long documents responds to local context and barely moves when long-range capability is broken, so a position-augmented model can look successful on the loss curve and fail retrieval. Needle-in-a-haystack alone is not enough either; it is close to saturated. Use a benchmark with multi-hop and aggregation tasks such as RULER (Hsieh et al., 2024, arXiv:2404.06654) before believing an extension worked.
10 flashcards for this concept
Click a card to reveal the answer.