기술 요약: 아키텍처 인지 강화 학습을 통한 슬라이딩 윈도우 어텐션의 수학적 추론 경쟁력 확보

문제 정의

추론 및 에이전트형 대규모 언어 모델(LLM)의 급격한 발전은 긴 컨텍스트 추론에 대한 결정적인 수요를 창출했습니다. 그러나 표준 셀프 어텐션(Self-Attention, SA) 메커니즘은 컨텍스트 길이에 따라 이차 함수적으로($O(N^2)$) 확장되어 트랜스포머 기반 모델의 근본적인 효율성 병목 현상을 야기합니다. 슬라이딩 윈도우 어텐션(Sliding-Window Attention, SWA)과 같은 선형 복잡도 대안은 스케일링을 위한 유망한 경로를 제공하지만, 수학적 추론(짧은 입력, 긴 생성)과 같이 장거리 의존성 모델링이 필요한 작업에서는 일반적으로 SA의 성능을 따라잡는 데 어려움을 겪습니다.

이 간극을 메우려는 기존의 접근 방식들은 두 가지 주요 과제에 직면해 있습니다:

1. 기본 모델 확보: 고성능 선형 복잡도 체크포인트는 드물며, 이를 처음부터 사전 학습(Pretraining)하는 것은 비용이 지나치게 많이 듭니다.
2. 성능 격차: 이미 사전 학습된 SA 모델을 SWA로 변환하고 지도 미세 조정(Supervised Fine-Tuning, SFT)을 적용하면 상당한 성능 저하가 발생합니다. 저자들은 이것이 데이터-아키텍처 불일치(Data-architecture mismatch) 때문이라고 가설을 세웠습니다. 즉, 표준 SFT 데이터셋은 SA 모델을 위해 준비되었으며 SWA가 효과적으로 모델링할 수 없는 장거리 의존성을 포함하고 있어, 최적의 모방을 방해한다는 것입니다.

방법론: SWARR

본 논문은 수학적 추론을 위해 SWA 모델을 적응시키면서 효율성을 보존하도록 설계된 2단계 경험적 레시피인 SWARR(Sliding-Window Attention with Reinforced Adaptation for Math Reasoning)을 제안합니다.

1단계: SFT를 통한 효율적인 변환

• 초기화: 사전 학습된 SA 모델(예: Qwen2.5-Math-1.5B-Base)을 사용하며, 사전 학습된 QKV 프로젝션을 재사용하면서 전역 어텐션을 로컬 슬라이딩 윈도우로 교체하여 SWA로 변환합니다.
• 지도 미세 조정(SFT): 모델은 수학적 추론 데이터로 SFT를 거칩니다.
• 한계점: 저자들은 SFT가 "아키텍처 불가지론적(architecture-agnostic)"이라고 언급합니다. 데이터셋의 타겟 궤적(trajectory)은 SA 모델에 의해 생성되었으므로, SWA가 자연스럽게 지원할 수 없는 범위를 넘어서는 의존성을 포함할 수 있으며, 이는 부적절한 모방을 강요하게 됩니다.

2단계: RL을 통한 정책 적응

• 강화 학습(RL): SFT가 완료된 SWA 모델은 검증 가능한 보상(verifiable rewards)을 사용하는 CISPO를 통해 추가로 적응됩니다.
• 아키텍처 인지(Architecture-Awareness): 고정된 외부 궤적을 학습하는 SFT와 달리, RL은 모델 자체의 아키텍처 제약 조건 하에서 모델이 직접 생성한 궤적을 최적화합니다. 이러한 "온-폴리시(on-policy)" 최적화는 모델이 로컬 어텐션으로 생성하기 더 쉬운 패턴(즉, 더 국소적인 의존성) 쪽으로 생성 분포를 이동시킬 수 있게 하여, 데이터-아키텍처 불일치를 완화합니다.

국소성 지표(Locality Metrics)

RL이 모델을 더 "로컬(local)"하게 적응시킨다는 가설을 검증하기 위해, 저자들은 확률 기반 정보 격차 지표($g_t(w)$)와 시퀀스 수준의 국소성 점수($L(w)$)를 도입했습니다. 이 지표들은 다음 토큰에 대한 예측 정보가 전체 컨텍스트 대비 로컬 윈도우 내에 얼마나 포착되는지를 정량화합니다.

주요 기여

1. SWARR 레시피: 비용이 많이 드는 사전 학습 없이도 더 큰 윈도우를 가진 SWA 모델이 수학적 추론에서 SA와 경쟁할 수 있도록 하는 실용적인 2단계 파이프라인(SA-to-SWA 변환 + RL 적응)을 제공합니다.
2. SWA의 경험적 재평가: SFT만으로는 SWA와 SA 사이에 상당한 격차가 남지만, RL 적응은 이 격차를 크게 좁힌다는 것을 보여줌으로써 SWA가 긴 생성 작업에 부적합하다는 결론에 도전합니다.
3. 데이터-아키텍처 불일치 분석: 통제된 교차 SFT 실험을 통해, SA 생성 데이터가 SWA 생성 데이터보다 SWA 성능에 더 해롭다는 것을 보여주었습니다. 이는 RL이 학습 분포를 SWA 제약 조건에 정렬함으로써 성능을 개선한다는 가설을 뒷받침합니다.
4. 사고의 국소성(Locality of Thought): 본 연구는 효과적인 수학적 추론이 더 국소적인 토큰 의존성으로부터 발현될 수 있다는 증거를 제공하며, 이는 "사고의 국소성" 현상과 일치합니다.

실험 결과

Qwen2.5-Math-1.5B-Base 모델을 사용하여 Math500, AMC, OlympiadBench, AIME24, AIME25를 포함한 벤치마크에서 실험을 수행했습니다.

• Post-SFT 성능: 변환 및 SFT 후, SWA 모델은 SA보다 뒤처졌습니다. 윈도우 크기가 작아질수록 격차는 커졌습니다(예: SWA2k-SFT는 SA-SFT보다 17.8포인트 낮았습니다).
• Post-RL 성능:
  • 동일 단계 학습(Equal-Step Training): SWA8k-RL-900은 SA-RL-900과 대등한 성능을 달성했으며, 평균 정확도 격차는 단 -0.4%였습니다.
  • 동일 시간 학습(Equal-Time Training): SWA의 선형 복잡도 덕분에 학습이 더 빠릅니다. 동일한 실제 벽시계 시간(약 500 GPU 시간) 동안 학습했을 때, SWA4k-RL-1400은 SA-RL-900의 성능과 일치하였고(+0.1% 격차), SWA8k-RL-1200은 심지어 이를 약간 능가했습니다(+0.7% 격차, 단 통계적 불확실성 범위 내).
  • 작은 윈도우: SWA2k는 여전히 제약이 있었으나, 동일 시간 학습 조건 하에서 격차가 크게 줄어들었습니다(동일 시간 학습 시 -17.8에서 -2.5포인트로 감소).
• 효율성: SWA 모델은 상당한 효율성 이득을 보여주었습니다. 예를 들어, SWA4k는 32k 컨텍스트 길이에서 SA 대비 약 6.2배의 속도 향상을 달성했으며, 고정된 메모리 제약 하에서 훨씬 더 큰 배치 사이즈를 지원했습니다.
• 교차 SFT 검증: SWA 모델이 생성한 데이터(아키텍처 매칭)로 학습했을 때, SWA와 SA 사이의 성능 격차가 거의 사라졌으며, 이는 초기 성능 저하가 주로 SA 지향적 데이터와의 불일치 때문이었음을 확인시켜 줍니다.

의의 및 주장

본 논문의 핵심적인 기여는 경험적 발견입니다. 즉, 강화 학습(RL)이 수학적 추론을 위한 더 큰 윈도우 SWA의 생존 가능성에 대한 평가를 근본적으로 바꾼다는 것입니다.

• 결론의 변화: RL 없이 SA를 SWA로 변환한 후 SFT를 적용하는 것은 SWA가 긴 생성 작업에 열등하다는 것을 시사합니다. 그러나 RL을 적용하면 더 큰 윈도우의 SWA가 SA와 경쟁할 수 있게 됩니다.
• 메커니즘: 이러한 개선은 RL이 모델의 생성 궤적을 더 "로컬"하게 적응시켜, 정적인 SFT 데이터보다 SWA 아키텍처의 제약 조건에 더 잘 맞도록 만드는 능력 덕분입니다.
• 범위 및 한와: 저자들은 자신들의 연구 결과가 수학적 추론(짧은 입력, 긴 생성 작업)에 특화되어 있음을 명시했습니다. SWARR가 모든 작업에 적용 가능한 보편적인 레시피라고 주장하지 않으며, 다른 도메인은 다른 최적화 전략이 필요할 수 있다고 언급했습니다. 또한, 학술적 연구로서 결과는 1.5B 및 4B 모델을 기반으로 하며, 거대 상업용 모델(예: 200B+)로의 확장성은 향후 과제로 남겨두었습니다.

요약하자면, SWARR는 아키텍처 인지 RL을 활용하여 표준 미세 조정의 한계를 극복함으로써, 복잡한 추론 작업을 위해 선형 복잡도 어텐션 모델을 배포할 수 있는 실용적이고 효율적인 경로를 제공합니다.