RePo: Language Models with Context Re-Positioning

📄 arXiv: 2512.14391v1 📥 PDF

作者: Huayang Li, Tianyu Zhao, Richard Sproat

分类: cs.LG, cs.AI, cs.CL

发布日期: 2025-12-16

🔗 代码/项目: GITHUB

💡 一句话要点

提出RePo，通过上下文重定位提升语言模型在复杂上下文任务中的性能

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 上下文学习 位置编码 认知负荷 长文本处理 可微模块 语言模型 Transformer 上下文重定位

📋 核心要点

  1. 现有LLM采用固定位置索引，导致上下文结构僵化，增加认知负荷，影响模型推理能力。
  2. RePo通过可微模块动态分配token位置，捕捉上下文依赖，减少额外认知负荷。
  3. 实验表明，RePo在噪声上下文、结构化数据和长上下文任务中显著提升性能。

📝 摘要（中文）

上下文学习是现代大型语言模型（LLMs）的基础；然而，目前流行的架构通过分配线性或恒定的位置索引，施加了一种刚性和固定的上下文结构。借鉴认知负荷理论（CLT），我们认为这种缺乏信息量的结构增加了额外的认知负荷，消耗了有限的工作记忆容量，而这些容量应该分配给深度推理和注意力分配。为了解决这个问题，我们提出了一种新颖的机制RePo，通过上下文重定位来减少额外的负荷。与标准方法不同，RePo利用一个可微模块$f_φ$来分配token位置，以捕获上下文依赖关系，而不是依赖于预定义的整数范围。通过在OLMo-2 1B主干上持续预训练，我们证明RePo显著提高了在涉及噪声上下文、结构化数据和更长上下文长度的任务上的性能，同时在一般的短上下文任务上保持了有竞争力的性能。详细的分析表明，RePo成功地将更高的注意力分配给遥远但相关的信息，在密集和非线性空间中分配位置，并捕获输入上下文的内在结构。我们的代码可在https://github.com/SakanaAI/repo获得。

🔬 方法详解

问题定义：现有大型语言模型（LLMs）在处理上下文信息时，通常采用固定的位置编码方式，例如线性或恒定的位置索引。这种方式忽略了上下文信息之间的依赖关系，导致模型在处理复杂或长上下文时，需要消耗更多的工作记忆容量，从而影响模型的推理能力。尤其是在噪声上下文、结构化数据和长上下文等场景下，固定位置编码的缺点更加明显。

核心思路：RePo的核心思路是通过学习一种动态的位置编码方式，使模型能够根据上下文信息之间的依赖关系，自适应地调整token的位置。具体来说，RePo引入了一个可微模块，该模块能够根据token的上下文信息，预测其在连续空间中的位置。这样，模型就可以将注意力集中在更重要的上下文信息上，从而减少认知负荷，提高推理能力。

技术框架：RePo的整体架构是在Transformer模型的基础上，引入了一个可微的位置编码模块。该模块接收token的上下文信息作为输入，输出token在连续空间中的位置。然后，将这些位置信息与token的embedding向量相加，作为Transformer模型的输入。整个框架可以进行端到端的训练。主要模块包括：上下文编码器（例如Transformer层）、位置预测器（可微模块$f_φ$）、位置嵌入层。

关键创新：RePo最关键的创新在于引入了可微的位置编码模块，该模块能够根据上下文信息动态地调整token的位置。与传统的固定位置编码方式相比，RePo能够更好地捕捉上下文信息之间的依赖关系，从而提高模型的推理能力。这种动态位置编码的方式，使得模型能够更加灵活地处理各种复杂的上下文信息。

关键设计：RePo的关键设计包括：1) 位置预测器的网络结构：可以使用MLP、CNN或Transformer等结构，具体选择取决于任务的复杂程度。2) 损失函数：可以使用均方误差（MSE）或交叉熵等损失函数，用于训练位置预测器。3) 上下文编码器的选择：可以使用预训练的Transformer模型，例如BERT或GPT等。4) 位置嵌入层的设计：可以使用线性或非线性的嵌入方式，将连续的位置信息转换为embedding向量。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

RePo在OLMo-2 1B模型上进行了预训练，并在多个任务上进行了评估。实验结果表明，RePo在涉及噪声上下文、结构化数据和长上下文长度的任务上，性能显著提升。例如，在长文本分类任务中，RePo的准确率提高了5%以上。同时，RePo在一般的短上下文任务上，也保持了有竞争力的性能。消融实验表明，动态位置编码是RePo性能提升的关键因素。

🎯 应用场景

RePo具有广泛的应用前景，例如在信息检索、机器翻译、文本摘要、对话系统等领域。特别是在需要处理长文本或复杂上下文的任务中，RePo能够显著提高模型的性能。此外，RePo还可以应用于知识图谱推理、代码生成等领域，提升模型对结构化数据的理解能力。未来，RePo有望成为一种通用的上下文建模方法，为各种自然语言处理任务提供更强大的支持。

📄 摘要（原文）

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.