想定する読者:
2017年にGoogleの研究者らが発表したTransformerモデルは、自然言語処理（NLP）の世界に革命をもたらしました。その鍵となったのが、Attentionメカニズムです。Attentionは、文中の各単語が他のすべての単語とどのように関連しているかを捉え、文脈をより深く理解することを可能にしました。しかし、Transformerにも限界があります。それは、計算コストが系列長（文の長さ）の二乗に比例して増加することです。このため、非常に長い文やドキュメントを扱うことが困難でした。 この問題を解決するため、多くの研究者が「線形Transformer」と呼ばれる効率的なAttentionメカニズムの開発に取り組んできました。これらのモデルは、計算コストを系列長に比例するように抑えることで、より長い系列を扱えるようにしました。しかし、線形Transformerは、情報を固定サイズの「状態」に圧縮するため、長い系列の情報を効果的に保持できないという新たな課題に直面しました。 人間の脳は、短期記憶と長期記憶を使い分け、さらにそれらを連携させることで、複雑な情報を処理しています。従来の深層学習モデルは、この人間の脳の記憶システムを十分に模倣できていませんでした。 そこで、Google Researchの研究チームは、Titansと呼ばれる新しいアーキテクチャを提案しました。Titansは、TransformerのAttentionメカニズム（短期記憶）と、新たに開発された「ニューラル長期記憶」（長期記憶）を組み合わせることで、長い系列の情報を効率的に保持し、高い精度で処理することを可能にしました。さらに、Titansは、テスト時に記憶を学習するという革新的な機能を備えています。 本稿では、Titansの論文「Titans: Learning to Memorize at Test Time」を詳細に解説し、その革新性と可能性を探ります。
TransformerのAttentionメカニズムは、以下の式で表されます。
Titansの核となるのが、ニューラル長期記憶（Neural Long-term Memory） モジュールです。このモジュールは、テスト時にデータを記憶するように学習する、メタ・インコンテキストモデルです。
人間の長期記憶は、期待に反する出来事（驚き） をより強く記憶する傾向があります。Titansのニューラル長期記憶は、このメカニズムを模倣しています。具体的には、入力に対するニューラルネットワークの勾配を「驚き」の指標として利用します。 Mₜ = Mₜ₋₁ - θₜ ∇ℓ(Mₜ₋₁; xₜ) ここで、ℓ(Mₜ₋₁; xₜ) は損失関数、θₜ は学習率です。 しかし、この単純な驚き指標では、大きな驚きの後に続く重要な情報を見逃す可能性があります。そこで、Titansは、過去の驚きと瞬間的な驚きを組み合わせた、より洗練された驚き指標を導入します。 Mₜ = Mₜ₋₁ + Sₜ Sₜ = ηₜSₜ₋₁ - θₜ ∇ℓ(Mₜ₋₁; xₜ) ここで、ηₜ は過去の驚きの減衰率、θₜ は瞬間的な驚きの重みを制御するデータ依存のパラメータです。この式は、モーメンタム付き勾配降下法と類似しており、Sₜ がモーメンタム項に相当します。
ニューラル長期記憶の目的は、過去のデータをキーとバリューのペアとして記憶することです。そのため、損失関数は以下のように定義されます。 ℓ(Mₜ₋₁; xₜ) = ||Mₜ₋₁(kₜ) - vₜ||²₂ ここで、kₜ = xₜWₖ, vₜ = xₜWᵥ は、入力 xₜ を線形変換して得られるキーとバリューです。この損失関数を最小化することで、ニューラル長期記憶は、キーとバリューの間の関連付けを学習します。
長い系列を扱う際には、記憶容量の制約から、不要な情報を忘却するメカニズムが重要になります。Titansは、以下の適応的な忘却メカニズムを導入します。 Mₜ = (1 - αₜ)Mₜ₋₁ + Sₜ Sₜ = ηₜSₜ₋₁ - θₜ ∇ℓ(Mₜ₋₁; xₜ) ここで、αₜ ∈ [0, 1] は、忘却の度合いを制御するゲーティングメカニズムです。このメカニズムは、現代のRNNにおけるゲーティングメカニズムの一般化と見なすことができます。
Titansのニューラル長期記憶は、単純な多層パーセプトロン（MLP）で構成されます。MLPの層数（Lₘ）は、少なくとも2以上であることが推奨されます。これは、深層のMLPが線形モデルよりも表現力が高いという理論的結果に基づいています。
記憶された情報を検索するには、重みを更新せずにニューラル長期記憶を順伝播させます。 yₜ = M*ₜ(qₜ) ここで、qₜ = xₜWᵩ は、入力 xₜ を線形変換して得られるクエリです。
ニューラル長期記憶の学習は、本質的に逐次的な処理ですが、バッチ処理とテンソル演算を活用することで、並列化が可能です。具体的には、以下の手順で計算を行います。
ニューラル長期記憶は、文脈に依存した記憶を形成します。しかし、タスクに関する一般的な知識を記憶するためには、入力に依存しないパラメータも必要です。Titansでは、これを永続メモリ（Persistent Memory） と呼びます。 永続メモリは、学習可能なパラメータ P = [p₁, p₂, …, pₙₚ] として実装され、入力系列の先頭に連結されます。 xₙₑ𝓌 = [p₁, p₂, …, pₙₚ] || x 永続メモリは、以下の3つの観点から、その有用性が説明されます。
Titansは、以下の3つの主要なコンポーネントで構成されます。
このアーキテクチャでは、記憶を現在の情報に対する文脈として扱います。
このアーキテクチャでは、スライディングウィンドウAttention（短期記憶）とニューラル長期記憶（長期記憶）を、ゲーティングメカニズムで組み合わせます。
このアーキテクチャでは、ニューラル長期記憶を深層ニューラルネットワークの層として使用します。
Titansは、以下のタスクで評価されました。
Titansは、Transformer++、RetNet、GLA、Mamba、Mamba2、DeltaNet、TTT、Gated DeltaNetなどの既存モデルと比較して、Perplexityと精度で優れた性能を示しました。特に、ニューラル長期記憶モジュール単体でも、他の非ハイブリッドモデルを上回る性能を達成しました。
Titansは、長い文脈から情報を検索する能力を評価するNIAHタスクにおいて、既存モデルを大幅に上回る性能を示しました。これは、Titansが長い系列の情報を効果的に保持できることを示しています。
非常に長い文書に分散した情報から推論するタスクであるBABILongベンチマークにおいて、TitansはGPT-4などの大規模モデルを含む既存のモデルを上回る性能を示しました。これはTitansが長期的な依存関係を捉える能力が非常に高いことを示しています。
ニューラル長期記憶の深さ（層数）を変化させた実験では、深い記憶ほど長い系列に対して高い性能を示すことが確認されました。ただし、記憶が深いほど計算コストが増加するため、性能と効率のトレードオフを考慮する必要があります。
Titansは、時系列予測とDNAモデリングのタスクにおいても、既存モデルと同等以上の性能を示しました。これは、Titansが自然言語処理以外の分野にも応用できる可能性を示唆しています。
Titansのニューラル長期記憶モジュールは、Mamba2やGated DeltaNetと比較して、わずかに計算コストが高いことが示されました。これは、より表現力の高い記憶更新メカニズムを採用しているためと考えられます。一方、Titans (MAL) は、FlashAttentionによる最適化の恩恵を受け、より高いスループットを達成しました。
Titansの各コンポーネント（線形メモリ、畳み込み、モーメンタム、重み減衰、永続メモリ）の寄与を評価するアブレーションスタディでは、すべてのコンポーネントが性能向上に貢献していることが確認されました。
Titansは、TransformerのAttentionメカニズムと、新たに開発されたニューラル長期記憶を組み合わせることで、長い系列の情報を効率的に保持し、高い精度で処理することを可能にしました。さらに、Titansは、テスト時に記憶を学習するという革新的な機能を備えています。 Titansは、自然言語処理の分野に新たな可能性をもたらすだけでなく、深層学習モデルの設計に新たな視点を提供します。今後の研究により、Titansのさらなる発展と応用が期待されます。