技術要約：連想再帰メモリによるLLMコンテキストの拡張

問題提起

大規模言語モデル（LLM）は、技術レポートの分析、ソフトウェア開発、マルチドキュメント推論などのタスクにおいて、数十万または数百万トークンに及ぶ入力を処理することがますます求められています。しかし、標準的なTransformerアーキテクチャは、根本的なボトルネックに直面しています。それは、自己注意（self-attention）の計算およびメモリコストがシーケンス長に対して二次関数的（$O(N^2)$）にスケールするという点です。さらに、モデルの公称ウィンドウ内であっても、コンテキスト長が増加するにつれて性能が低下することがよくあります。再帰型アーキテクチャ（Mamba、RWKVなど）は線形スケーリングを実現していますが、これらは通常、ゼロからの学習を必要とするため、既存の学習済みLLMを活用することができず、またTransformerと比較して複雑なアルゴリズムタスクや指示への追従性に苦慮することがあります。

手法

著者らは、コンテキスト長を拡張しつつ、一定のメモリ・スケーリングと効率性を維持するための実用的な解決策として、**連想再帰メモリ・トランスフォーマー（ARMT）**を提案します。ARMTは、学習済みベースLLMのラッパーとして機能し、入力のセグメント単位での処理を可能にします。

コア・アーキテクチャ

ARMTは、長いコンテキストの入力を固定長の重複しないセグメントに分割します。各セグメント内では、フル自己注意（短期／ワーキングメモリ）を利用します。決定的なのは、セグメント間で情報を伝播させる**レイヤー単位の連想メモリ・モジュール（long-term memory）**を導入している点です。このメカニズムは、以下の3つのステージで動作します。

1. メモリ抽出（Memory Extraction）: 各Transformerレイヤーは、入力セグメントをメモリ埋め込みへと圧縮します。
2. メモリ統合（Memory Consolidation）: これらの埋め込みは、キー・バリュー（KV）ペアとして、レイヤーごとの連想行列へと統合されます。
3. 連想（Association）: 後続のセグメントにおける埋め込みは、クエリベクトルへと変換され、連想行列と乗算されることで、過去のセグメントから関連情報を検索します。

学習レシピ

論文では、学習済みLLMをARMTに適応させるための包括的な学習戦略を概説しています。

• 継続事前学習（Continued Pre-training）: 未初期化の連想メモリ・パラメータは、タスク固有のファインチューニングの前に、効果的なメモリ伝播を学習させるため、長いコンテキストを用いた教師なし言語モデリング（例：FineWeb-Eduの19Bトークン）を通じて初期化されます。
• カリキュラム学習（Curriculum Learning）: 長距離依存関係をゼロから学習することの難しさに対処するため、モデルはセグメント数を段階的に増やし（例：2から4、次に8へ）、学習率をアニーリングしながらファインチューニングされます。
• 合成データ生成（Synthetic Data Generation）: 長コンテキスト・シナリオにおけるデータの不足を克服するため、長いドキュメントから短い一節を連結し、それぞれのコンテキスト長に応じたビン（bin）を作成してQAペアを生成することで、合成訓練インスタンスを作成します。
• レイヤー・プルーニングと選択（Layer Pruning and Selection）: 連想メモリがすべてのレイヤーに必要かどうかを調査しています。特定の層（例：中間層および最終層のサブセット）にのみ連想メモリを保持する戦略を提案しており、これにより、性能の大幅な低下を招くことなく、学習パラメータ数と計算コストを削減できます。

主な貢献

1. ドメイン特化型データセット: コードにおける可変型の予測を行う**ManyTypes-long (MT)と、長文ドキュメントの質問応答を行うGovReport-long (GR)**という、現実的かつ限定的なドメインのワークロードを評価するために設計された2つの新しいデータセットを構築しました。
2. 学習レシピ: 継続事前学習、合成データ生成、カリキュラム学習、および選択的なレイヤー統合を組み合わせた、ARMTを用いたLLMコンテキスト拡張のための新しいフレームワークを提供しました。
3. 実証的検証: ARMT拡張モデルが以下のことを示す広範な実験研究を実施しました。
   • 元のコンテキスト制限を大幅に超える入力（最大64kトークン）を、インリミットのベースラインに対する性能劣化なしに処理できること。
   • ベースモデルと比較して、Long-OOD（分布外）のコンテキスト長において優れた汎化性能を示すこと。
   • ベースラインの性能を維持しつつ、元のコンテキストウィンドウ内において約30%少ないFLOPsで動作すること。

実験結果

著者らは、Gemma-3-1B-ITおよびSmolLM-2-360M-ITをバックボーンとして使用し、MTおよびGRデータセットを用いてARMTを評価しました。

• 性能: ARMTモデルは、最大65kトークンのコンテキスト長にわたって安定した性能を維持しました。対照的に、ベースモデル（ファインチューニング済みであっても）は、ネイティブのコンテキストウィンドウ（例：8kまたは32k）を超えると急激な性能低下を示しました。ARMTは、Long-OOD領域（32k–65k）においてベースモデルを大幅に上回りました。
• 効率性: ARMTは、コンテキスト長に関わらず一定のGPUメモリ使用量を示しましたが、ベースモデルのメモリ使用量は線形に増加しました。32kトークンのシーケンスにおいて、ARMTは同一のメモリ予算下で、ベースモデルの8に対し4倍のバッチサイズ（32）を可能にしました。
• FLOPsの削減: 理論的分析および経験的な推論時間の測定により、グローバルなアテンションFLOPsが $T/S$（シーケンス長をセグメントサイズで割った値）の係数で減少することが確認され、フルアテンションモデルと比較して全体で約30%の総FLOPs削減につながることが確認されました。
• アブレーション研究:
  • レイヤー・プルーニング: 特定の（事前に選択された）中間層および最終層の約20%のレイヤーにのみ連想メモリを持つモデルが、フルARMTモデルと同等またはそれ以上の性能を達成しました。
  • 事前学習: 継続事前学習は、連想メモリを初期化するために不可欠であり、インドメインおよびOODの両方の性能を大幅に向上させることが判明しました。
  • ベースライン: ARMTは、特に長コンテキストの汎化において、Mamba-2、DeltaNet、xLSTMを含む他の長コンテキスト・ベースラインを上回り、かつゼロからの学習よりも広範な事前学習を必要としませんでした。

意義と主張

本論文は、小規模から中規模のLLM（最大1Bパラメータ）において、長コンテキスト処理を可能にするための、計算効率の高い実用的なアプローチとしてARMTを位置づけています。著者らは、この手法が、Transformerの強力な短コンテキスト性能と、再帰型モデルの線形スケーリングの間の溝を埋めるものであると主張しています。

重要性に関する主な主張は以下の通りです：

• プライバシーとローカル展開: 小規模なモデルでの効率的な長コンテキスト処理を可能にすることで、ARMTは、リモートの大規模なAPIベースのLLMに依存しない、プライバシー保護型のアプリケーションを促進します。
• スケーラビリティ: この手法は、一定のメモリで任意の長さのコンテキストを処理することを可能にし、これは現実世界のドキュメント分析における重要な要件です。
• 効率性: 一定のメモリ・スケーリングとFLOPsの削減の組み合わせにより、ARMTはリソース制約のある環境における実行可能な代替手段となります。
• 汎化: 提案された学習レシピは、標準的なTransformerで見られる「Lost in the Middle（情報の埋没）」現象および、長いシーケンスを扱う際の性能低下を効果的に解決します。

著者らは、実験が最大1Bパラメータまでのモデルおよび特定のタスク（コードおよびドキュメントQA）に限定されていること、また、連想メモリとTransformer表現の相互作用の根底にあるメカニズムが部分的に未解明であることに触れ、限界を認めています。しかし、小規模モデルにおけるコンテキスト拡張の能力を示したことは、実用的なローカル長コンテキストAIに向けた重要な一歩であると主張しています。