아키텍처 인지 강화 학습: 슬라이딩 윈도우 어텐션을 활용한 수학적 추론 성능 향상

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

추론 및 에이전트 기반 대규모 언어 모델(LLM)의 빠른 발전으로 인해 긴 문맥을 처리하는 능력에 대한 요구가 증가하고 있지만, 셀프 어텐션(SA)은 문맥 길이에 따라 2차적으로 성능 저하가 발생합니다. 본 연구에서는 수학적 추론에 슬라이딩 윈도우 어텐션(SWA) 모델을 적용하기 위한 실용적인 방법인 SWARR (Sliding-Window Attention with Reinforced Adaptation for Math Reasoning)을 제안합니다. SWARR은 다음과 같은 두 단계로 구성됩니다: (1) 지도 학습(Supervised Fine-tuning, SFT)을 통해 사전 훈련된 SA 모델을 효율적으로 SWA 모델로 변환하여 새로운 기본 모델을 처음부터 훈련하는 것을 방지하고, (2) 강화 학습(Reinforcement Learning, RL)을 활용한 정책 적응입니다. 실험 결과, SFT 후에도 SWA는 여전히 SA에 비해 성능이 낮았으며, 이는 데이터와 아키텍처 간의 불일치 때문이라고 판단됩니다. 대부분의 SFT 데이터는 SA 모델을 위해 준비되었으며, SWA가 모델링하기 어려운 장거리 의존성을 포함할 수 있기 때문입니다. 반면, 온-정책 RL은 SWA 제약 조건 하에서 자체적으로 생성된 경로를 최적화하여 SWA에 더 적합한 경로로 조정할 수 있습니다. 수학적 추론 벤치마크 실험 결과, SWARR은 SWA와 SA 간의 성능 격차를 크게 줄이며, SWA 변환 과정에서 손실되는 정확도를 상당 부분 회복하는 동시에 선형 복잡도의 효율적인 어텐션 이점을 유지합니다. 본 연구의 핵심 기여는 강화 학습이 SWA 모델의 수학적 추론 적용 가능성에 대한 기존의 판단을 변화시킨다는 경험적 증거를 제시했다는 점입니다.

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.