Shuffle the Context: RoPE-Perturbed Self-Distillation for Long-Context Adaptation
Abstract
Large language models (LLMs) increasingly operate in settings that require reliable long-context understanding, such as retrieval-augmented generation and multi-document reasoning. A common strategy is to fine-tune pretrained short-context models at the target sequence length. However, we find that standard long-context adaptation can remain brittle: model accuracy depends strongly on the absolute placement of relevant evidence, exhibiting high positional variance even when controlling for task format and difficulty. We propose RoPE-Perturbed Self-Distillation, a training regularizer that improves positional robustness. The core idea is to form alternative ``views'' of the same training sequence by perturbing its RoPE indices---effectively moving parts of the context to different positions---and to train the model to produce consistent predictions across views via self-distillation. This encourages reliance on semantic signals instead of brittle position dependencies. Experiments on long-context adaptation of Llama-3-8B and Qwen-3-4B demonstrate consistent gains on long-context benchmarks, including up to 12.04 percent-point improvement on RULER-64K for Llama-3-8B and 2.71 percent-point gain on RULER-256K for Qwen-3-4B after SFT, alongside improved length extrapolation beyond the training context window.
Lay Summary
Many AI assistants can now read very long inputs, such as collections of documents or large codebases, but they can still be unreliable: the same important fact may be found or missed depending on where it appears. This matters because real systems often stitch together retrieved pages, files, or conversations, so useful evidence may land near the beginning, middle, or end. We introduce a training method that makes a language model less sensitive to these positions. During training, the model sees the same text twice: once normally, and once with its internal position labels shifted. We then encourage it to make similar predictions for both versions. This teaches the model to rely more on the meaning of the text than on fragile location-specific patterns. In tests with Llama and Qwen models, the method improved long-document performance, especially when answers were hidden in difficult places or when inputs were longer than the training examples. It also preserved performance on standard short-input tasks. Overall, the approach offers a simple way to make long-context language models more dependable for question answering, multi-document analysis, and software engineering.