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

2026-06-09

AI Summary

This work addresses the significant performance gap between sliding window attention (SWA) and standard self-attention (SA) in mathematical reasoning tasks, where SWA struggles to model long-range dependencies. The authors propose SWARR, a novel approach that first efficiently transfers a pretrained SA model to an SWA architecture via supervised fine-tuning and then employs an architecture-aware reinforcement learning strategy to optimize reasoning trajectories tailored to SWA’s localized attention mechanism. SWARR is the first method to demonstrate that reinforcement learning can effectively mitigate SWA’s performance deficit in mathematical reasoning. By doing so, it substantially narrows the accuracy gap between SWA and SA on established mathematical reasoning benchmarks while preserving SWA’s linear computational complexity, thereby challenging the prevailing assumption that SWA is inherently unsuitable for complex reasoning tasks.

Abstract

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.

Problem

• sliding-window attention
• mathematical reasoning
• long-context inference
• data-architecture mismatch
• self-attention

Innovation

• Sliding-Window Attention
• Reinforcement Learning
• Mathematical Reasoning
• Architecture-Aware Adaptation
• Long-Context Efficiency