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

背景

随着推理型（Reasoning）和智能体（Agentic）大型语言模型（LLMs）的快速发展，对长上下文推理（Long-context Inference）的需求急剧增加。然而，标准的自注意力机制（Self-Attention, SA）的计算复杂度随上下文长度呈二次方增长（Quadratic Scaling），这成为了限制模型处理超长文本的主要瓶颈。

为了突破这一限制，滑动窗口注意力（Sliding-Window Attention, SWA）因其线性复杂度而受到关注，它通过限制注意力范围来降低计算开销。然而，直接将预训练的 SA 模型转换为 SWA 架构通常会导致性能显著下降，特别是在需要捕捉长距离依赖关系的数学推理任务中。现有的微调方法往往未能有效解决这种架构转换带来的性能损耗，导致 SWA 模型在数学推理能力上远逊于 SA 模型。

核心内容

本文提出了一种名为 SWARR（Sliding-Window Attention with Reinforced Adaptation for Math Reasoning，用于数学推理的滑动窗口注意力强化适应）的实用方案，旨在将 SWA 模型有效地适配到数学推理任务中。该方案包含两个关键阶段：

1. 高效转换与监督微调（SFT）： 首先，通过监督微调（Supervised Fine-Tuning, SFT）将预训练的 SA 模型高效地转换为 SWA 架构。这一阶段避免了从头预训练一个新基座模型的高昂成本，实现了架构的快速迁移。

2. 基于强化学习的策略适应（RL）： 在 SFT 之后，研究团队发现 SWA 模型的表现仍然落后于 SA 模型。他们假设这种差距部分源于“数据-架构不匹配”（Data-Architecture Mismatch）：大多数 SFT 数据是为 SA 模型准备的，其中包含大量长距离依赖关系，而 SWA 难以有效建模这些依赖。

   为了解决这一问题，SWARR 引入了强化学习（Reinforcement Learning, RL）进行策略适应。由于在线强化学习（On-policy RL）优化的是在 SWA 约束下自我生成的轨迹（Trajectories），它能够自适应地调整生成过程，使轨迹更好地匹配 SWA 的架构特性，从而弥补长距离建模能力的不足。

实验结果显示，这种“SFT + RL”的组合策略显著缩小了 SWA 与 SA 之间的性能差距。SWARR 不仅恢复了在 SWA 转换过程中丢失的大部分准确率，还保留了线性复杂度注意力机制带来的效率优势。

关键要点

• 架构转换的局限性：仅通过监督微调（SFT）将 SA 模型转换为 SWA 模型，无法完全保留其在数学推理上的性能，主要原因在于训练数据中的长距离依赖与 SWA 的局部注意力机制存在不匹配。
• RL 的关键作用：强化学习（RL）通过优化模型自我生成的推理轨迹，使模型能够适应 SWA 的架构约束。这种“架构感知”的适应过程是提升 SWA 性能的核心。
• 性能与效率的平衡：SWARR 方案证明了 SWA 在数学推理任务中可以变得具有竞争力（Competitive），既恢复了接近 SA 模型的准确率，又维持了线性计算复杂度带来的推理效率优势。
• 方法论贡献：本文的核心贡献在于实证发现，仅依靠转换和 SFT 得出的关于 SWA 可行性的结论是不完整的；引入 RL 后，SWA 在数学推理中的可行性得到了实质性提升。

意义与影响

这项研究对大型语言模型的架构选择和部署策略具有重要的指导意义：

1. 降低长上下文推理成本：通过证明 SWA 经过适当优化后可在复杂推理任务中媲美 SA，该研究为部署高效、低成本的长上下文模型提供了可行路径。这对于资源受限的环境或需要高频推理的应用场景尤为关键。
2. 重新评估架构迁移策略：研究揭示了“数据-架构不匹配”这一潜在问题，并指出简单的微调不足以解决架构变更带来的性能损失。这提示未来在模型架构迁移时，必须考虑引入能够适应新架构特性的训练阶段（如强化学习）。
3. 推动高效推理模型的发展：SWARR 方案为开发兼具高精度和高效率的推理模型提供了新的技术路线，有助于推动更广泛的 LLM 应用落地，特别是在需要处理长文本和复杂逻辑的领域。