RePo: Language Models with Context Re-Positioning

上下文学习是现代大语言模型（LLM）的核心机制；然而，当前主流架构通过分配线性或固定的常数位置索引来强加一种僵化且固定不变的上下文结构。基于认知负荷理论（Cognitive Load Theory, CLT），我们认为这种缺乏信息性的结构会增加外在认知负荷，消耗本应用于深度推理和注意力分配的有限工作记忆资源。为解决这一问题，我们提出了一种名为RePo的新机制，通过重新调整上下文的位置来降低不必要的认知负担。与标准方法不同，RePo采用一个可微模块 $f_φ$ 来动态分配能够反映上下文依赖关系的标记位置，而非依赖预定义的整数范围。我们在OLMo-2 1B模型基础上持续进行预训练，结果表明，RePo在处理包含噪声上下文、结构化数据以及长上下文的任务中显著提升了性能，同时在通用的短上下文任务上仍保持具有竞争力的表现。详细分析显示，RePo能够更有效地关注距离较远但相关的信息，将位置分配到密集且非线性的空间中，并捕捉输入上下文的内在结构。我们的代码已公开于 https://github.com/SakanaAI/repo。

现有的大语言模型（LLM）通常使用线性或固定的整数位置编码来处理上下文，这种刚性的位置结构无法有效反映输入中真实的语义和依赖关系，导致模型在处理长上下文、噪声环境或结构化数据时效率低下。这会引入不必要的‘外在认知负荷’（extraneous cognitive load），占用本应用于深度推理的有限注意力资源。论文提出：是否可以通过更智能的位置分配机制来缓解这一问题？这是一个尚未被充分探索的重要方向，尤其在提升模型对复杂上下文的理解能力方面具有现实意义。

受认知负荷理论（Cognitive Load Theory, CLT）启发，论文提出RePo机制，摒弃传统的固定整数位置索引，转而使用一个可学习、可微分的模块$f_φ$动态地为token分配更合理的‘语义位置’。这些位置不再是线性的，而是密集且非线性的，能够反映上下文中token之间的实际依赖关系，从而降低模型处理无关信息的认知负担，提升对关键远距离信息的关注能力。其核心创新在于将位置建模从预定义规则转变为可学习的上下文感知结构。

作者在OLMo-2 1B模型上持续预训练RePo，并在多种任务上验证其有效性：包括含噪声上下文、结构化输入以及长文本理解等场景，均表现出显著性能提升，同时在标准短上下文任务上保持竞争力。分析表明，RePo能更有效地引导注意力关注遥远但相关的信息，学习到更紧凑和语义合理的位置分布。代码已开源（https://github.com/SakanaAI/repo），增强了可复现性。未来可探索将其扩展至多模态情境，或结合稀疏注意力进一步优化计算效率。

1. Rethinking Attention with Performers (Choromanski et al., 2021) 2. LongNet: Scaling Transformers to 1,000,000,000 Tokens (Tay et al., 2023) 3. Ring Attention for Long Context Generative Modeling (Google DeepMind, 2024) 4. Llama-3: A Meta AI Approach to Scaling LLMs (Meta, 2024) 5. Retentive Network: A Scalable Alternative to Transformer for Large Language Models (Ren et al., 2023)