Language Models Need Sleep

技術的サマリー：「言語モデルは睡眠を必要とする」

問題定義
トランスフォーマーベースの大規模言語モデル（LLM）は、アテンション機構がコンテキスト長に対して二次的にスケーリングし、キー・バリュー（KV）キャッシュメモリが線形的に増大するため、長視野タスクに苦戦する。近年の効率的なハイブリッドモデル（アテンションと状態空間モデル（SSM）を組み合わせる）は、固定サイズの「ファストウェイト」を通じてメモリコストを軽減するが、推論深度において重大なボトルネックに直面している。

著者らは、純粋な SSM-アテンションハイブリッドモデルが、KV キャッシュから排除されたトークンに対する深い逐次計算を実行できないことを実証した。メモリ容量が情報を保存するのに十分であっても、必要な推論深度が増加するにつれて（例えば、細胞オートマトンを多数ステップにわたってシミュレートする場合や、多段のグラフ探索を行う場合）、モデルのパフォーマンスは低下する。ボトルネックは単なる記憶容量ではなく、単一のフォワードパス内で排除されたコンテキストを有用な内部状態に変換するために十分な計算リソースが不足していることにある。

手法：LLM 睡眠
生物学的な記憶固定（睡眠中の海馬での再生）に着想を得て、本論文は LLM 向けの「睡眠に似た」メカニズムを提案する。このアプローチは、追加の計算をオフラインフェーズへ移行させ、オンライン予測フェーズの遅延を維持する。

1. アーキテクチャ: モデルは、アテンションブロック（$B^{attn}$）とファストウェイトメモリ（$S_t$）を備えた SSM ブロック（$B^{ssm}$）を交互に配置したハイブリッドアーキテクチャを採用する。
2. 睡眠フェーズ: コンテキストウィンドウが満杯になった際（例えば、$L$ トークンごと）、モデルは「睡眠」フェーズに入る。このフェーズ中：
   • モデルは外部入力トークンを受信しない。
   • 蓄積されたコンテキストに対して$N$ 回のオフライン再帰パスを実行する。
   • 各パスにおいて、モデルは学習された局所ルール（例えば、ゲート付きデルタルール更新）を用いて、SSM ブロック内のファストウェイトを再帰的に更新する。
   • このプロセスにより、モデルは排除されたコンテキストの圧縮表現を反復的に洗練させることができる。
3. 排除と予測: $N$ 回のパス後、KV キャッシュはクリアされる（スライディングウィンドウ変種では部分的にクリアされる）、そしてモデルは更新されたファストウェイトを用いて操作を再開する。
4. 推論制約: 決定的なことに、実際の予測フェーズ（クエリへの回答）において、モデルは単一のフォワードパスのみを実行する。すべての深い推論とコンテキストの固定は、先行する「睡眠」フェーズ中に完了していなければならない。
5. 学習: モデルは、コンテキストが排除された後の予測誤差を最小化するため、再帰的な睡眠パスを含む計算グラフ全体を通じたバックプロパゲーションによりエンドツーエンドで学習される。

主要な貢献

• 推論ボトルネックの特定: 著者らは、制御された設定において、純粋な SSM-アテンションハイブリッドモデルは、メモリ容量が固定かつ十分であっても、深い推論を必要とするタスク（細胞オートマトンにおける高い $t$、グラフ検索における高い $k$）で失敗することを示した。これにより、ボトルネックがメモリだけでなく、計算深度にあることが特定された。
• 睡眠アーキテクチャ: 彼らは、記憶固定のために再帰的計算を統合するアーキテクチャを提案する。睡眠パスの数（$N$）を増やすことで、推論遅延を増加させることなく、モデルはより深い推論タスクを処理するために計算リソースをスケーリングできる。
• 実証的検証: この手法は以下のタスクで検証された：
  • ルール 110（細胞オートマトン）: 深い逐次シミュレーションを必要とするタスク。
  • Depo: 断片化されたコンテキスト上のグラフ探索を必要とする、多段の知識検索タスク。
  • GSM-Infinite: 長いコンテキストとダミートークンを含む現実的な数学推論ベンチマークで、事前学習済み LLM（Jet-Nemotron 2B および Ouro 1.4B）上で評価された。

結果

• 合成タスク: 睡眠時間（$N$）の増加はパフォーマンスを著しく向上させた。$t=32$ のルール 110 タスクにおいて、非再帰モデル（$N=1$）はランダム推測に近い性能（約 10% の精度）を示したが、3〜4 回の睡眠パスを持つモデルは 30% 以上の精度を達成した。同様に、Depo タスクでは、$N$ を増加させることで学習が加速し、特に単一パスモデルが停滞する高ホップクエリ（例えば 16 ホップ）において精度が向上した。
• 現実的な推論（GSM-Infinite）: 事前学習済みモデルにおいて、$N$ を増加させることは、より困難な問題（6〜8 回の算術演算を必要とする）で大幅な改善をもたらした。例えば、Ouro 1.4B モデルにおいて、睡眠パスを 1 から 4 に増加させることで、6 演算問題の精度は 41.9% から 61.5% に向上した。
• スライディングウィンドウ: モデルが最近のトークンを保持しつつも古いコンテキストの圧縮に睡眠に依存するスライディングウィンドウ排除戦略であっても、その恩恵は維持された。
• 学習スループット: 睡眠パスの逐次的性質により、学習スループットは $N$ にほぼ比例して減少するが、著者らは、大きなコンテキストウィンドウの場合、このオーバーヘッドは GPU が飽和した状態であるため、完全に並列なベースラインと比較してウォールクロック学習時間を著しく妨げないことを指摘している。

意義と主張
本論文は、オフライン再帰が、排除されたコンテキストをファストウェイトに組織化する viable なメカニズムであり、モデルが直接注意を払うことのできない情報について深く推論することを可能にすると主張している。

• 限定的な範囲: 著者らは、彼らの発見が制御された合成タスクと小規模な事前学習済みモデルに基づいていることを強調している。彼らはすべての長コンテキスト推論問題を解決するとは主張せず、むしろ「睡眠」（オフライン固定）が記憶容量と推論深度の間のギャップを埋める特定のアーキテクチャ経路を実証している。
• 生物学的着想: この研究は、短期コンテキストを長期ウェイトへ転送するには、反復的計算（再生）が必要であり、これはリアルタイム予測の厳格な遅延制約から切り離し得るという仮説を検証するものである。
• 将来の方向性: 本論文は、これをより能力のある長コンテキストシステムを構築するための方法的貢献として位置づけ、推論時の推論深度に対する学習時の逐次性のトレードオフは、将来の研究における有望な道筋であると示唆している。