用强化学习让滑窗注意力在数学推理中逼近自注意力表现
Architecture-Aware Reinforcement Learning Makes Sliding-Window Attention Competitive in Math Reasoning

- 先微调后强化学习,分两阶段优化滑窗注意力模型
- 在数学推理任务上恢复了转换损失的大部分准确率
- 适合关注长文本效率与性能平衡的研究者
随着推理型大模型发展,对长上下文推理的需求激增,但自注意力(SA)随上下文长度呈平方增长。为应对这一问题,本文提出SWARR(滑窗注意力结合强化适应的数学推理方法),通过两阶段优化:(1) 用监督微调将预训练的SA模型高效转换为滑窗注意力(SWA),避免重新预训练;(2) 采用强化学习进行策略适应。实验发现,仅经微调的SWA仍显著落后于SA,我们推测原因在于数据-架构不匹配:多数微调数据针对SA设计,包含长程依赖,难以被SWA建模。而基于SWA约束生成轨迹的在线强化学习能有效调整推理路径,使其更适配滑窗结构。在多个数学推理基准测试中,该方法显著缩小了SWA与SA之间的差距,在保持线性复杂度优势的同时,恢复了大部分精度损失。核心贡献在于实证表明:仅靠转换与微调会低估SWA在数学推理中的潜力,强化学习可改变这一判断。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。