Google Titans と MIRAS：200万トークン超の長期 AI メモリのための新アーキテクチャ

Google は長大なシーケンスを処理するための新しいアーキテクチャ Titans と、再帰型ネットワークの高速性とトランスフォーマーの精度を統合する理論的フレームワーク MIRAS を発表した。同社によれば、このアプローチにより、モデルは膨大な文書からゲノムデータに至るまで、極めて長いコンテキストにわたって本質的な詳細を保持できるようになるという。

古典的なトランスフォーマーは、モデルが過去を「振り返り」、テキストの中で最も意味のある部分を切り出せるようにするアテンション機構によって、この分野に革命をもたらした。しかしこの能力には代償がある。アテンションの計算量はシーケンス長に対して二乗で増加するため、数百万トークン規模のコンテキストへトランスフォーマーを拡張することは、計算コストの面で極めて困難になる。

研究コミュニティはすでに、高速な線形時間アーキテクチャ ― 効率的な RNN や Mamba-2 のような状態空間モデル ― によってこの制約を回避しようと試みてきた。これらのシステムは過去のコンテキストを固定サイズのベクトルに圧縮し、線形スケーリングを可能にする。しかしこのメモリ制約のため、事実間の微妙な関係を保持しなければならない、真に豊かで長大なシーケンスでは苦戦しがちだ。

Google は 2 本の新しい論文で、Titans と MIRAS を組み合わせて用いることを提案している。Titans は具体的なアーキテクチャであり、MIRAS はそのようなシステムを連想記憶の一形態として解釈する統一的な理論フレームワークだ。両者は「推論時学習」というアイデアを前進させる。すなわちモデルは、パラメータから情報を取り出すだけでなく、新しいデータが到着するたびに長期メモリを更新しながら、その場で自らを洗練させることができる。

Titans の中心的な発想は、人間の短期記憶と長期記憶を思わせるメモリ階層である。短期記憶は依然としてアテンションが担い、局所的なコンテキストに対して優れた性能を発揮する。一方、長期記憶は単一の固定ベクトルや行列としてではなく、多層パーセプトロンという完全なディープニューラルネットワークとして実装される。これによりモデルは過去のコンテキストをはるかに豊かに符号化し、単に断片的なメモを保存するのではなく、物語構造そのものを「理解」できるようになる。

もう一つの決定的な概念が「サプライズ」機構である。人間は日常的な細部はすぐに忘れる一方で、予想を裏切る出来事は鮮明に記憶する。Titans も同じ原理で動作する。モデルは現在のメモリ状態と新しい入力を比較し、それがどれほど「予想外」かを測定する。差異が小さければ、長期的に保存する必要はない。入力が期待から大きく外れていれば、保存すべき重要または異常な事実であることを示すシグナルとなる。

システムが単発のスパイクにだけ反応してしまうのを防ぐため、Titans には 2 つの安定化コンポーネントが組み込まれている。第一にモメンタムである。モデルは現在のサプライズだけでなく直近の履歴も考慮し、関連する出来事の連鎖を捉えられるようにする。第二に適応的ウェイト減衰であり、極端に長いシーケンスにおいて、古い情報を徐々にフェードアウトさせることでメモリのあふれを防ぐ「逃し弁」として機能する。

MIRAS はこのクラス全体のモデルを、より高い抽象レベルで記述する。MIRAS の枠組みでは、あらゆるシーケンスアーキテクチャは、キーを値に写像することを学習する連想記憶として捉えられ、新しい情報を取り込みつつ、これまで学んだことを破壊しないという 2 つの相反する力のバランスを取るものとみなされる。MIRAS は、メモリ自体の構造、何に注意を向けるかを選ぶメカニズム、保持と忘却のプロセス、そしてメモリを更新する最適化ルールという 4 つの中核要素を特定している。

Google は特に、MIRAS が従来の平均二乗誤差やドット積類似度といったユークリッド的な指標から離れる点を強調している。ほとんどの現代的システムでは、これらの指標がモデルの好みやメモリの正則化の基盤となっている。これらは扱いやすい一方で外れ値に敏感であり、取り得る解の空間を制限してしまう。MIRAS は、非ユークリッドな損失関数や、最適化理論や統計学から引き出された、より多様な正則化手法に基づいて動作するアーキテクチャへの道を開く。

このフレームワークのもとで、Google の研究者たちは YAAD、MONETA、MEMORA という 3 つの明示的なアテンションフリーモデルを開発した。

• YAAD は誤差に対してより穏やかな Huber 損失を用い、誤字や散発的なアーティファクトを含むノイズの多いデータでより良い性能を発揮する。
• MONETA は高度なノルムを試し、記憶と忘却に対してより厳密な数学的「法則」を課すことで、長期メモリの安定性を高められるかどうかを探る。
• MEMORA はメモリが確率分布のように振る舞うことを強制することで、各更新を制御された精密なものとし、メモリを極限まで安定化させようとする。

実験では、Titans と MIRAS ベースのモデルは、Transformer++、Mamba-2、Gated DeltaNet といった最新アーキテクチャと比較された。テストには標準的な言語モデリングデータセット（C4、WikiText）と、HellaSwag や PIQA を含むゼロショット常識推論ベンチマークが用いられた。Google によれば、新しいモデルは同程度の規模のベースラインよりも低いパープレキシティと高い精度を達成しつつ、線形スケーラビリティと並列学習能力を維持しているという。

Titans の利点は、極端に長いコンテキストを必要とするタスクで最も顕著に現れる。非常に長い文書全体に散在する事実をまたいで推論することが求められる BABILong ベンチマークにおいて、Titans は GPT-4 のような大規模システムを含むすべてのベースラインモデルを、はるかに少ないパラメータ数で上回った。著者らはまた、Titans がパフォーマンスを維持したまま、コンテキストウィンドウを 200 万トークン超まで拡張できることも報告している。

追加の研究では、メモリモジュールの深さが極めて重要であることが示された。メモリの「サイズ」を同一に保ったまま深さだけを変えた場合でも、より深いバリアントは一貫してパープレキシティが低く、シーケンス長が増加してもより滑らかにスケールした。これは、豊かで深い長期メモリが、広大なコンテキストをモデルがうまく扱ううえで実際に役立つことのさらなる証拠である。

最終的に Google は、Titans と MIRAS を、その場で「今この瞬間」に学習し、膨大なコンテキストにわたって本質的な詳細を保持しつつ、実運用に耐える効率性を保つ新世代モデルへの一歩として位置づけている。MIRAS はオンライン最適化、連想記憶、アーキテクチャ設計を統合し、Titans は長コンテキスト AI の時代において、RNN の速度とトランスフォーマーの表現力をどのように調和させうるかを示している。