Maglev: Sliding Recurrent Memory

技術要約：Maglev – スライディング再帰的メモリ

問題提起
現在のTransformerアーキテクチャは、メモリ容量と計算効率の間の根本的なトレードオフに直面している。フル・コーザル自己注意（Full causal self-attention）は強力な非線形検索とインコンテキスト計算を可能にするが、コンテキスト長に対して線形に増加するコストを伴う。逆に、スライディングウィンドウ注意（Sliding-window attention）は注意およびKVキャッシュのコストを限定するものの、遠方の情報を破棄してしまうため、持続的な長期メモリを提供できない。古典的な再帰モデル（LSTMなど）は履歴を圧縮するが、小さな状態を逐次的に伝播させるため、容量に限界がある。メモリ・トランスフォーマーや線形注意モデルは、このギャップを埋めようと試みているが、それらは人工的なセグメント境界を課したり、構造化された線形更新に依存したり、あるいは学習時に逐次的なアンローリング（展開）を必要としたりすることが多い。核心となる課題は、並列的な事前学習を維持しつつ、有界な推論コストを実現する、実用的なトークン単位の非線形再帰メモリメカニズムをTransformerのために開発することである。

手法：Maglevアーキテクチャ
Maglevは、学習時におけるメモリ構築と推論時におけるメモリ伝播を分離することで、これに対処する。このアーキテクチャは、以下の2つの結合されたモデルで構成される：

1. プリフィラー（Prefiller, $Q$）： 観測されたシーケンスを、より大きな受容野（例：フル注意とスライディングウィンドウ注意のインターリーブ）を用いて並列に処理する、コーザルTransformer。これは、すべてのトークン位置に対してターゲットとなるメモリベクトル $m'_t$ の軌跡を生成する。
2. デコーダー（Decoder, $P$）： 次のトークンを予測し、自身のメモリベクトル $m_t$ を生成する、スライディングウィンドウTransformer。決定的なことに、$P$ は全履歴にアクセスできない。$P$ は、直近のトークンの限定されたウィンドウと、シフトされたメモリベクトルのウィンドウのみを参照する。

学習メカニズム：リフテッド・パラレル学習（Lifted Parallel Training）
再帰的な展開（これが並列化を妨げる）を行う代わりに、Maglevは2パスの学習戦略を採用する：

• パス1： プリフィラー $Q$ が、完全なターゲット・メモリ軌跡 $m'_{1:T}$ を並列に生成する。
• パス2： デコーダー $P$ は、シフトされたターゲット軌跡 $m'_{0:T-1}$ と入力トークンを併用して、次のトークンを予測し、自身のメモリ $m_{1:T}$ を生成する。
• 目的関数： モデルは以下の複合損失関数を用いて最適化される：
  • 次トークン予測損失（$L_{CE}$）： 言語モデリングのための標準的なクロスエントロピー損失。
  • メモリ一貫性損失（$L_{cons}$）： デコーダーの出力メモリ $m_t$ をプリフィラーのターゲット $m'_t$ に整合させる項。これにより、デコーダーは逐次的なアンローリングなしに、次のステップに必要なメモリ状態を再現することを学習する。

推論メカニズム
推論時には、プリフィラー $Q$ は破棄される。デコーダー $P$ は再帰的に動作し、自身が以前に生成したメモリ $m_{t-1}$ をシステムにフィードバックする。

• 再帰的K/V注入： 専用のメモリトークンを使用するのではなく、Maglevは学習された射影とゲーティング機構を介して、シフトされたメモリ状態を注意機構のKey/Value（K/V）特徴量へと注入する。
• 有界な複雑性： デコーダーは、トークンとメモリの固定サイズのスライディングウィンドウを保持する。再帰的な状態はK/Vキャッシュを通じて前方へ運ばれ、推論コスト（計算量およびキャッシュサイズ）が、標準的なスライディングウィンドウ注意のプロファイルと同様に、総シーケンス長に依存せず一定に保たれることを保証する。

主な貢献

1. トークン単位の非線形再帰： Maglevは、追加のメモリトークンや人工的なセグメント境界を必要とせず、すべてのトークンのデコーダ状態がメモリとして機能し、Transformerの完全な非線形深度を横断するメカニズムを導入する。
2. 再帰モデルのための並列学習： プリフィラーを使用してターゲットを並列に生成することで、Maglevは非線形再帰モデルを標準的なシーケンス並列スループットで学習することを可能にし、従来のRNN学習における逐次的なボトルネックを回避する。
3. パラメータ共有： 著者らは、プリフィラー $Q$ とデコーダー $P$ が、性能を大幅に低下させることなく、大部分のパラメータ（具体的にはTransformerブロック）を共有できることを示しており、これにより学習セットのパラメータメモリ・オーバーヘッドを削減している。
4. 有界な推論： デプロイされたモデルは、再帰的なK/V注入を通じて持続的なメモリ能力を獲得しながら、スライディングウィンドウ注意の効率性を維持する。

実験結果
nanochat d20アーキテクチャ（パラメータ数435M、43.52Bトークンで学習）において、Maglevは、インターリーブされたフル/スライディングウィンドウTransformer、純粋なスライディングウィンドウモデル、およびLatent Recurrent Transformer (LRT) ベースラインと比較評価された。

• 性能： 最良のMaglevバリアント（別パラメータ、$\lambda=1$）は、FineWeb-Eduの検証時のBits-Per-Byte (BPB) を、0.7413（スライディングウィンドウ・ベースライン）から0.7251へと改善した。
• ダウンストリーム・タスク： 平均ダウンストリーム精度56.4%を達成し、一致させたスライディングウィンドウTransformer（54.1%）および一致させたLRTベースラインの両方を上回った。
• 一貫性の重み： 研究では、最適な一貫性重み（$\lambda$）はパラメータ共有戦略に依存することが指摘された。高い $\lambda$ は別パラメータモデルには有益であったが、共有パラメータモデルには制約を与える可能性がある。

意義と主張
本論文は、Maglevが、トークンごとに書き換え可能な、持続的なトークン単位のメモリを備えた非線形Transformerを構築するための実用的な経路を提供すると主張している。Maglevは、メモリの生成とメモリの使用を分離することで、非線形再帰モデルの並列的な事前学習を実現しており、これはこれまで困難であった。著者らは、Maglevをこの設計空間に関する予備的な調査として位置付けており、テストされた設定においてはスライディングウィンドウや潜在再帰ベースラインよりも優れているものの、この手法のスケールアップ、デプロイのためのカーネル効率の最適化、および代替的なメモリ注入メカニズム（例：残差ストリーム注入やクロスアテンション）の探索に向けたさらなる研究が必要であると述べている。