RePo: Language Models with Context Re-Positioning

In-context learning is fundamental to modern Large Language Models (LLMs); however, prevailing architectures impose a rigid and fixed contextual structure by assigning linear or constant positional indices. The rigid position information poses the full burden of organizing the input structure to attention layers, thus reducing the amount of attention that could be allocated for more critical information. To address this, we propose RePo, a novel mechanism that alleviates the burden for attention layers via context re-positioning. Unlike conventional approaches, RePo utilizes a differentiable module, $f_\phi$, to assign token positions that capture contextual dependencies, rather than replying on pre-defined order. By continually pre-training on the OLMo-2 1B \& 7B models, we demonstrate that RePo consistently enhances performance on tasks involving noisy contexts, structured data, and longer context length, while maintaining competitive performance on general short-context tasks. Analysis reveals that RePo successfully allocates more attention mass to distant but relevant information, assigns positions in a dense and non-linear space, and captures the intrinsic structure of the input context.

When we read a long document, we naturally group related ideas together in our minds, even if they are physically pages apart. However, Large Language Models traditionally process text in a strict, sequential order—one word right after another. This rigid approach limits their performance on long or messy texts, as they waste processing power trying to bridge the gap between distant pieces of related information. To solve this, we developed a method called Context Re-Positioning (RePo). Instead of forcing the AI to maintain a strict reading order, RePo allows the model to dynamically reorganize its "working memory." It automatically evaluates the text and pulls highly related but distant clues closer together, while pushing irrelevant noise out of the way. This approach allows language models to organically reorganize their context to focus on what truly matters. By adopting RePo, AI models become significantly better at understanding extremely long, noisy, and complex documents, all while maintaining their high performance on standard, everyday tasks.