Architecture-Aware Reinforcement Learning Makes Sliding-Window Attention Competitive in Math Reasoning

📄 arXiv: 2606.11634v1 📥 PDF

作者: Kai Liu, Peijie Dong, Xinchen Xie, Jianfei Gao, Qipeng Guo, Xiaowen Chu, Shaoting Zhang, Kai Chen

分类: cs.AI

发布日期: 2026-06-10

💡 一句话要点

提出SWARR以解决长上下文推理中的自注意力效率问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 滑动窗口注意力 强化学习 数学推理 自注意力 长上下文推理

📋 核心要点

  1. 现有的自注意力模型在处理长上下文时效率低下，导致在数学推理任务中表现不佳。
  2. 提出SWARR方法，通过监督微调将预训练的SA模型转换为SWA，并利用强化学习进行策略适应。
  3. 实验结果显示，SWARR显著提高了SWA的性能，缩小了与SA模型的准确性差距。

📝 摘要（中文）

随着推理和智能大型语言模型（LLMs）的快速发展，对长上下文推理的需求不断增加，但自注意力（SA）在上下文长度上呈现平方级别的扩展。为此，本文研究了SWARR（滑动窗口注意力与强化适应的数学推理），提出了一种将预训练的SA模型高效转换为SWA的方案，并通过强化学习进行策略适应。实验表明，该方法显著缩小了SWA与SA之间的性能差距，同时保留了线性复杂度注意力的效率优势。我们的核心贡献在于发现强化学习改变了对SWA在数学推理中可行性的结论。

🔬 方法详解

问题定义：本文旨在解决滑动窗口注意力（SWA）在数学推理任务中性能不足的问题。现有的自注意力（SA）模型在长上下文推理中效率低下，导致SWA在处理长范围依赖时表现不佳。

核心思路：SWARR方法通过两阶段的流程来优化SWA模型：首先是通过监督微调（SFT）将预训练的SA模型转换为SWA，其次是通过强化学习（RL）对策略进行适应，以更好地匹配SWA的特性。

技术框架：SWARR的整体架构包括两个主要阶段：第一阶段是高效的模型转换，第二阶段是基于RL的策略优化。通过这种方式，SWARR能够在不重新预训练新模型的情况下，提升SWA的推理能力。

关键创新：SWARR的创新在于通过强化学习优化自生成轨迹，使得SWA能够更好地适应长范围依赖，从而显著提高了其在数学推理任务中的表现。与传统的SFT方法相比，SWARR能够有效缩小SWA与SA之间的性能差距。

关键设计：在SWARR中，关键设计包括选择合适的损失函数以优化SWA的性能，以及在RL阶段中使用的策略更新机制。这些设计确保了模型在保持效率的同时，能够有效捕捉长范围依赖特征。

📊 实验亮点

实验结果表明，SWARR在数学推理基准测试中显著缩小了SWA与SA之间的性能差距，恢复了在SWA转换过程中损失的准确性，同时保持了线性复杂度注意力的效率。具体性能提升数据未知。

🎯 应用场景

该研究的潜在应用领域包括教育、科学计算和智能助手等需要处理复杂数学推理的场景。通过提升SWA在长上下文推理中的能力，SWARR能够为这些领域提供更高效的推理工具，推动智能系统的进一步发展。

📄 摘要（原文）

The rapid progress of reasoning and agentic large language models (LLMs) has increased the demand for long-context inference, but self-attention (SA) scales quadratically with context length. To address this, we study SWARR (Sliding-Window Attention with Reinforced Adaptation for Math Reasoning), a practical recipe for adapting SWA models to mathematical reasoning. SWARR has two stages: (1) efficient conversion from a pretrained SA model to SWA with supervised fine-tuning (SFT), which avoids pretraining a new base model, and (2) policy adaptation with reinforcement learning (RL). We find that SWA still underperforms SA after SFT, and we hypothesize that this gap is caused in part by a data-architecture mismatch: most SFT data are prepared for SA models and may contain long-range dependencies that are difficult for SWA to model. Because on-policy RL optimizes self-generated trajectories under the SWA constraint, it can adapt trajectories to better match SWA. Experiments on mathematical reasoning benchmarks show that this recipe substantially narrows the gap between SWA and SA, recovering much of the accuracy lost during SWA conversion while preserving the efficiency benefits of linear-complexity attention. Our central contribution is the empirical finding that RL changes the conclusion one would draw from conversion and SFT alone about SWA's viability for math reasoning.