Architecture-Aware Reinforcement Learning Makes Sliding-Window Attention Competitive in Math Reasoning
For practitioners deploying long-context LLMs, this work shows that reinforcement learning can bridge the performance gap between efficient SWA and full self-attention, making SWA viable for math reasoning tasks.
The paper proposes SWARR, a two-stage recipe combining supervised fine-tuning and reinforcement learning to adapt sliding-window attention (SWA) models for mathematical reasoning, recovering much of the accuracy lost during SWA conversion while maintaining linear-complexity efficiency.
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.