技術要約：アーキテクチャを意識した強化学習により、スライディングウィンドウ・アテンションが数学的推論において競争力を獲得

問題提起

推論およびエージェント型の大規模言語モデル（LLM）の急速な進展に伴い、長文脈推論への切実な需要が生じている。しかし、標準的な自己アテンション（SA）メカニズムは、コンテキスト長に対して二次関数的（$O(N^2)$）にスケールするため、トランスフォーマーベースのモデルにおける根本的な効率性のボトルネックとなっている。スライディングウィンドウ・アテンション（SWA）のような線形計算量の代替手法は、スケーリングへの有望な道筋を示す一方で、数学的推論（短い入力、長い生成）のような長距離依存性のモデリングを必要とするタスクにおいては、通常、SAの性能に追いつくことに苦慮している。

このギャップを埋めるための既存のアプローチは、主に2つの課題に直面している：

1. ベースモデルの獲得： 高性能な線形計算量モデルのチェックポイントは稀であり、それらをゼロから事前学習することは極めてコストが高い。
2. 性能の差： 事前学習済みのSAモデルを単にSWAに変換し、教師あり微調整（SFT）を適用すると、大幅な性能低下を招く。著者らは、これがデータとアーキテクチャのミスマッチによるものであると仮定している。標準的な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）」と指摘している。データセット内のターゲット・トラジェトリはSAモデルによって生成されたものであるため、SWAのウィンドウを超える依存関係が含まれている可能性があり、SWAモデルが自然にサポートできないパターンを模倣することを強いてしまう。

ステージ2：RLによる方策適応

• 強化学習（RL）： SFT後のSWAモデルは、RL（具体的には検証可能な報酬を用いたCISPO）を用いてさらに適応される。
• アーキテクチャへの意識（Architecture-Awareness）： 固定された外部のトラジェトリで学習するSFTとは異なり、RLはモデル自身のアーキテクチャの制約下で生成されたトラジェトリを最適化する。この「オンポリシー」の最適化により、モデルは自身の生成分布を、ローカル・アテンションで生成しやすいパターン（すなわち、より局所的な依存関係）へとシフトさせることができ、それによってデータとアーキテクチャのミスマッチを緩和できる。

局所性指標

RLがモデルをより「ローカル」に適応させるという仮説を検証するために、著者らは確率ベースの情報ギャップ指標（$g_t(w)$）と、シーケンスレベルの局所性スコア（$L(w)$）を導入している。これらの指標は、次のトークンの予測情報が、フルコンテキストと比較してローカルウィンドウ内にどれだけ捉えられているかを定量化する。

主な貢献

1. SWARRレシピ： コストのかかる事前学習を行うことなく、より大きなウィンドウを持つSWAモデルを数学的推論においてSAと競合させることを可能にする、実用的な2段階のパイプライン（SAから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指向のデータとのミスマッチに起因していたことが確認された。

意義と主張

本論文の中心的な貢献は、経験的な発見である。すなわち、強化学習は、数学的推論における大きなウィンドウを持つSWAの生存可能性に関する評価を根本的に変えるということである。

• 結論の転換： RLなしでは、SAからSWAへの変換とそれに続くSFTは、SWAが長文生成タスクにおいて劣っていることを示唆する。しかし、RLを用いることで、より大きなウィンドウを持つSWAはSAと競争力を持つようになる。
• メカニズム： この改善は、RLがモデルの生成トラジェトリを、静的なSFTデータよりもSWAアーキテクチャの制約により適合するように、より「ローカル」なものへと適応させる能力に起因している。
• 範囲と限界： 著者らは、彼らの知見が数学的推論（短い入力、長い生成というタスク）に特有のものであることを明示している。SWARRがすべてのタスクに対する普遍的なレシピであるとは主張しておらず、異なるドメインでは異なる最適化戦略が必要になる可能性があると述べている。また、学術的研究であるため、結果は1.5Bおよび4Bモデルに基づいている。大規模な商用モデル（例：200B以上）へのスケーラビリティについては今後の課題である。

要約すると、SWARRは、アーキテクチャを意識したRLを活用することで、標準的なファインチューニングの限界を克服し、複雑な推論タスクに対して線形計算量のアテンションモデルをデプロイするための、実用的かつ効率的な経路を提供するものである。