Saved in:
Bibliographic Details
Main Authors: Li, Huayang, Zhao, Tianyu, Cai, Deng, Sproat, Richard
Format: Preprint
Published: 2025
Subjects:
Online Access:https://arxiv.org/abs/2512.14391
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866912944027074560
author Li, Huayang
Zhao, Tianyu
Cai, Deng
Sproat, Richard
author_facet Li, Huayang
Zhao, Tianyu
Cai, Deng
Sproat, Richard
contents 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 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. 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. We will open-source the code and model weights. Our code is at https://github.com/SakanaAI/repo.
format Preprint
id arxiv_https___arxiv_org_abs_2512_14391
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle RePo: Language Models with Context Re-Positioning
Li, Huayang
Zhao, Tianyu
Cai, Deng
Sproat, Richard
Machine Learning
Artificial Intelligence
Computation and Language
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 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. 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. We will open-source the code and model weights. Our code is at https://github.com/SakanaAI/repo.
title RePo: Language Models with Context Re-Positioning
topic Machine Learning
Artificial Intelligence
Computation and Language
url https://arxiv.org/abs/2512.14391