MoNe、長文コンテキストをテスト時ニューラルメモリに書き込み、128K推論の総FLOPsを81%削減

Qualcomm AI Researchが提案したMoNeは、長文コンテキスト処理を「メモリへの書き込み」と「クエリへの回答」という2段階に分割する。システムは文書を512-tokenのセグメントに分割し、凍結したTransformerへ順番に入力する。各層にはSwiGLU高速重みネットワークが追加され、その層で生成されたkey-valueの関連付けと局所勾配を使ってメモリを更新するため、モデル全体を通したバックプロパゲーションは不要だ。書き込みの完了後、クエリtokenから高速重みを介して固定数のメモリtokenを直接生成し、それらをセルフアテンションのkeyおよびvalueとして使用する。元のコンテキストをKV cacheへ再度読み込む必要はない。

この設計では、前処理コストはコンテキスト長に対してO(N)で増加する一方、1回のクエリコストはNに依存しなくなる。保存した高速重みは複数のクエリで再利用でき、新しいデータを使って逐次更新することも可能だ。論文では、凍結したQwen2.5-0.5B-Instructを用いて評価している。128K tokenでは、書き込み全体と推論を合わせたピークメモリ使用量が1.41GB、計算量が149.61T FLOPsだった。直接ICLの7.07GB、786.33T FLOPsと比べ、それぞれ約80%と81%の削減となる。プラグインの訓練に使用したコンテキスト長は最大4Kにすぎないが、RULERのsingle-needle retrieval、multi-key retrieval、frequent word extractionでは、128Kでそれぞれ0.96、0.94、0.96を記録した。直接ICLのスコアは0.28、0、0.23だった。

エンジニアリング上の価値は、既存Transformerのバックボーン重みを維持でき、線形スケーリングする長文コンテキスト処理のためにアーキテクチャ全体を再事前学習する必要がない点にある。ただし、6.4%の追加パラメータは依然としてオフラインで訓練する必要があり、テスト時の書き込みにも勾配更新が含まれる。また、論文が報告しているのはFLOPsとピークメモリ使用量であり、異なるハードウェア上でのエンドツーエンドレイテンシではない。評価対象は1種類の0.5Bモデルと、高度に構造化された3つのタスクに限られている。長文書を対象とした質問応答、コードベース分析、マルチターンエージェントのメモリにおいて、詳細な意味情報を保持できることはまだ証明されていない。今後は、より大規模なモデル、実際のワークロード、メモリ汚染、長期的な逐次更新の安定性を検証する必要がある。