🤖 AI

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

This paper introduces SWARR, a two-stage method combining supervised fine-tuning and reinforcement learning that successfully adapts efficient sliding-window attention models to mathematical reasoning, effectively closing the performance gap with standard self-attention by aligning generated trajectories with architectural constraints.

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

Published 2026-06-11

📖 4 min read☕ Coffee break read

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Overwhelmed Librarian"

Imagine a brilliant librarian (the AI) who can answer any question. However, this librarian has a rule: to answer a question, they must read every single book in the library to find the right information.

The Good News: This librarian is incredibly smart and accurate.
The Bad News: If the library has a million books, reading them all takes forever. As the library grows, the time it takes to find an answer grows exponentially (it gets much, much slower). This is the problem with the standard AI "Self-Attention" method.

The Proposed Solution: The "Sliding Window"

To fix the slowness, researchers tried a new rule: The librarian is only allowed to look at the last 100 books on the shelf. They can't see the books from 10 years ago, only the ones right next to them.

The Good News: This is super fast! The librarian can answer questions almost instantly, no matter how big the library is.
The Bad News: The librarian starts making mistakes. If the answer to a math problem depends on a fact mentioned 500 pages ago, the librarian misses it because their "window" is too small.

The Paper's Discovery: "Training the Librarian to Think Differently"

The authors of this paper (Kai Liu and team) asked a simple question: Can we make this fast "Sliding Window" librarian just as good as the slow "Read-Everything" librarian for math problems?

They tried a two-step recipe called SWARR:

Step 1: The "Quick Switch" (Supervised Fine-Tuning)

First, they took the smart, slow librarian and simply told them, "Okay, from now on, only look at the last 100 books."

The Result: It didn't work well. The librarian was confused. They were trying to solve complex math problems using old textbooks they couldn't see anymore. They were still slower and less accurate than the original.

Step 2: The "Self-Practice" (Reinforcement Learning)

This is the magic part. Instead of giving the librarian more old textbooks to memorize, they let the librarian practice solving problems on their own using their new "small window" rule.

The Analogy: Imagine the librarian realizes, "Hey, I can't see the whole library, so I need to change how I think. Instead of looking for a fact 500 pages back, I'll break the math problem into smaller chunks that I can solve with just the last 100 pages."
The Result: The librarian learned to write their own "thought process" (reasoning steps) in a way that fits their limited view. They stopped trying to remember everything and started focusing on the immediate context.

The Surprising Conclusion

After this "Self-Practice" training, the fast librarian became almost as good as the slow librarian at solving math problems, but they did it much faster.

The paper found that:

The Gap Closed: The difference in accuracy between the slow and fast models almost disappeared.
The Reason: The fast librarian learned to be "local." They stopped trying to rely on distant information and started solving problems by keeping their thoughts tight and close to the current step.
The Lesson: You don't need to see the whole library to solve a math problem; you just need to know how to organize your thoughts so you don't need to see the whole library.

Why This Matters

The paper proves that you don't have to choose between "Smart but Slow" and "Fast but Dumb." By using a specific training method (Reinforcement Learning) that respects the AI's limitations, you can get a model that is fast, efficient, and still very smart at math.

In short: They taught the AI to stop trying to be a "memory machine" and start being a "smart problem-solver" that works within its own limits.

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

Problem Statement

The rapid advancement of reasoning and agentic Large Language Models (LLMs) has created a critical demand for long-context inference. However, standard self-attention (SA) mechanisms scale quadratically ( $O(N^2)$ ) with context length, creating a fundamental efficiency bottleneck for transformer-based models. While linear-complexity alternatives like Sliding-Window Attention (SWA) offer a promising path for scaling, they typically struggle to match SA performance on tasks requiring long-range dependency modeling, such as mathematical reasoning (short-input, long-generation).

Existing approaches to bridge this gap face two primary challenges:

Base Model Acquisition: High-performance linear-complexity checkpoints are rare, and pretraining them from scratch is prohibitively costly.
Performance Gap: Simply converting a pretrained SA model to SWA and applying Supervised Fine-Tuning (SFT) results in a significant performance drop. The authors hypothesize this is due to a data-architecture mismatch: standard SFT datasets are prepared for SA models and contain long-range dependencies that SWA cannot effectively model, leading to suboptimal imitation.

Methodology: SWARR

The paper proposes SWARR (Sliding-Window Attention with Reinforced Adaptation for Math Reasoning), a two-stage empirical recipe designed to adapt SWA models to mathematical reasoning while preserving efficiency.

Stage 1: Efficient Conversion with SFT

Initialization: A pretrained SA model (e.g., Qwen2.5-Math-1.5B-Base) is converted to SWA by replacing global attention with a local sliding window while reusing the pretrained QKV projections.
Supervised Fine-Tuning: The model undergoes SFT on mathematical reasoning data.
Limitation: The authors note that SFT is "architecture-agnostic." Since the target trajectories in the dataset were generated by SA models, they may contain dependencies extending beyond the SWA window, forcing the SWA model to imitate patterns it cannot naturally support.

Stage 2: Policy Adaptation with RL

Reinforcement Learning (RL): The SFT-converted SWA model is further adapted using RL (specifically CISPO with verifiable rewards).
Architecture-Awareness: Unlike SFT, which trains on fixed external trajectories, RL optimizes trajectories generated by the model itself under its own architectural constraints. This "on-policy" optimization allows the model to shift its generation distribution toward patterns that are easier to produce with local attention (i.e., more local dependencies), thereby mitigating the data-architecture mismatch.

Locality Metrics

To validate the hypothesis that RL adapts the model to be more "local," the authors introduce a probability-based information gap metric ( $g_t(w)$ ) and a sequence-level locality score ( $L(w)$ ). These metrics quantify how much predictive information for the next token is captured within a local window versus the full context.

Key Contributions

SWARR Recipe: A practical two-stage pipeline (SA-to-SWA conversion + RL adaptation) that enables larger-window SWA models to compete with SA in mathematical reasoning without costly pretraining.
Empirical Re-evaluation of SWA: The paper demonstrates that while SFT alone leaves a substantial gap between SWA and SA, RL adaptation significantly narrows this gap, challenging the conclusion that SWA is unsuitable for long-generation reasoning.
Data-Architecture Mismatch Analysis: Through controlled cross-SFT experiments, the authors show that SA-generated data hurts SWA performance more than architecture-matched SWA-generated data. This supports the hypothesis that RL improves performance by aligning the training distribution with the SWA constraint.
Locality of Thought: The study provides evidence that effective mathematical reasoning can emerge from more local token dependencies, aligning with the "locality of thought" phenomenon.

Experimental Results

Experiments were conducted on the Qwen2.5-Math-1.5B-Base model across benchmarks including Math500, AMC, OlympiadBench, AIME24, and AIME25.

Post-SFT Performance: After conversion and SFT, SWA models trailed SA. The gap widened as window size decreased (e.g., SWA2k-SFT was 17.8 points below SA-SFT).
Post-RL Performance:
- Equal-Step Training: SWA8k-RL-900 achieved performance comparable to SA-RL-900, with an average accuracy gap of only -0.4%.
- Equal-Time Training: Due to the linear complexity of SWA, training is faster. When trained for similar wall-clock time (~500 GPU hours), SWA4k-RL-1400 matched SA-RL-900 performance (+0.1% gap), and SWA8k-RL-1200 even slightly outperformed it (+0.7% gap, though within statistical uncertainty).
- Small Windows: SWA2k remained constrained, though the gap narrowed significantly (from -17.8 to -2.5 points under equal-time training).
Efficiency: SWA models demonstrated substantial efficiency gains. For instance, SWA4k achieved an approximately 6.2× speedup over SA at 32k context length and supported significantly larger batch sizes under fixed memory constraints.
Cross-SFT Validation: When trained on data generated by SWA models (architecture-matched), the performance gap between SWA and SA nearly disappeared, confirming that the initial drop was largely due to the mismatch with SA-oriented data.

Significance and Claims

The paper's central contribution is an empirical finding: Reinforcement Learning fundamentally changes the viability assessment of larger-window SWA for mathematical reasoning.

Shift in Conclusion: Without RL, the conversion of SA to SWA followed by SFT suggests SWA is inferior for long-generation tasks. With RL, larger-window SWA becomes competitive with SA.
Mechanism: The improvement is attributed to RL's ability to adapt the model's generation trajectory to be more "local," fitting the SWA architecture's constraints better than static SFT data can.
Scope and Limitations: The authors explicitly state that their findings are specific to mathematical reasoning (a short-input, long-generation task). They do not claim SWARR is a universal recipe for all tasks, noting that different domains may require different optimization strategies. Furthermore, as an academic study, results are based on 1.5B and 4B models; scalability to massive commercial models (e.g., 200B+) remains future work.

In summary, SWARR offers a practical, efficient pathway to deploy linear-complexity attention models for complex reasoning tasks by leveraging architecture-aware RL to overcome the limitations of standard fine-tuning.