RePo : Modèles de langage avec repositionnement contextuel

RePo: Language Models with Context Re-Positioning

December 16, 2025

papers.authors: Huayang Li, Tianyu Zhao, Richard Sproat

cs.AI

papers.abstract

L'apprentissage en contexte est fondamental pour les modèles de langage de grande taille (LLM) modernes ; cependant, les architectures dominantes imposent une structure contextuelle rigide et fixe en attribuant des indices positionnels linéaires ou constants. En nous appuyant sur la théorie de la charge cognitive (TCC), nous soutenons que cette structure non informative augmente la charge cognitive extrinsèque, consommant une capacité de mémoire de travail limitée qui devrait être allouée au raisonnement profond et à l'allocation de l'attention. Pour résoudre ce problème, nous proposons RePo, un nouveau mécanisme qui réduit la charge extrinsèque via un repositionnement du contexte. Contrairement aux approches standard, RePo utilise un module différentiable, f_φ, pour attribuer des positions aux tokens qui capturent les dépendances contextuelles, plutôt que de s'appuyer sur un intervalle entier prédéfini. En poursuivant l'apprentissage préalable sur l'architecture de base OLMo-2 1B, nous démontrons que RePo améliore significativement les performances sur des tâches impliquant des contextes bruités, des données structurées et de plus longues longueurs de contexte, tout en maintenant des performances compétitives sur des tâches générales à contexte court. Une analyse détaillée révèle que RePo alloue avec succès une attention plus élevée aux informations distantes mais pertinentes, attribue des positions dans un espace dense et non linéaire, et capture la structure intrinsèque du contexte d'entrée. Notre code est disponible à l'adresse https://github.com/SakanaAI/repo.

English

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. Drawing on Cognitive Load Theory (CLT), we argue that this uninformative structure increases extraneous cognitive load, consuming finite working memory capacity that should be allocated to deep reasoning and attention allocation. To address this, we propose RePo, a novel mechanism that reduces extraneous load via context re-positioning. Unlike standard approaches, RePo utilizes a differentiable module, f_φ, to assign token positions that capture contextual dependencies, rather than replying on pre-defined integer range. By continually pre-training on the OLMo-2 1B backbone, we demonstrate that RePo significantly enhances performance on tasks involving noisy contexts, structured data, and longer context length, while maintaining competitive performance on general short-context tasks. Detailed analysis reveals that RePo successfully allocate higher attention to distant but relevant information, assign positions in dense and non-linear space, and capture the intrinsic structure of the input context. Our code is available at https://github.com/SakanaAI/repo.