2025-11-20T08:25:14.880374

Titans: Learning to Memorize at Test Time

Behrouz, Zhong, Mirrokni
Over more than a decade there has been an extensive research effort on how to effectively utilize recurrent models and attention. While recurrent models aim to compress the data into a fixed-size memory (called hidden state), attention allows attending to the entire context window, capturing the direct dependencies of all tokens. This more accurate modeling of dependencies, however, comes with a quadratic cost, limiting the model to a fixed-length context. We present a new neural long-term memory module that learns to memorize historical context and helps attention to attend to the current context while utilizing long past information. We show that this neural memory has the advantage of fast parallelizable training while maintaining a fast inference. From a memory perspective, we argue that attention due to its limited context but accurate dependency modeling performs as a short-term memory, while neural memory due to its ability to memorize the data, acts as a long-term, more persistent, memory. Based on these two modules, we introduce a new family of architectures, called Titans, and present three variants to address how one can effectively incorporate memory into this architecture. Our experimental results on language modeling, common-sense reasoning, genomics, and time series tasks show that Titans are more effective than Transformers and recent modern linear recurrent models. They further can effectively scale to larger than 2M context window size with higher accuracy in needle-in-haystack tasks compared to baselines.
academic

Titans: テスト時の記憶学習

基本情報

  • 論文ID: 2501.00663
  • タイトル: Titans: Learning to Memorize at Test Time
  • 著者: Ali Behrouz, Peilin Zhong, Vahab Mirrokni (Google Research)
  • 分類: cs.LG cs.AI cs.CL
  • 発表日: 2024年12月31日
  • 論文リンク: https://arxiv.org/abs/2501.00663

要約

本論文は、履歴コンテキストを記憶することを学習し、注意機構が長期の過去情報を活用しながら現在のコンテキストに焦点を当てるのを支援する新しいニューラル長期記憶モジュールを提案している。著者は記憶の観点から、注意機構は限定されたコンテキストながら正確な依存関係モデリングにより短期記憶として機能し、ニューラル記憶はデータ記憶能力により長期でより永続的な記憶として機能すると主張している。これら2つのモジュールに基づき、著者は新しいアーキテクチャファミリーTitansを導入し、記憶を効率的にアーキテクチャに統合する3つの変種を提案している。実験結果は、Titansが言語モデリング、常識推論、ゲノミクス、時系列タスクにおいてTransformersおよび現代的な線形循環モデルより有効であり、2M以上のコンテキストウィンドウサイズまで効率的にスケーリング可能であることを示している。

研究背景と動機

中核的な問題

既存のシーケンスモデリングアーキテクチャは効率性と性能のトレードオフに直面している：

  1. Transformers: 依存関係の正確なモデリングが可能だが、計算複雑度がO(n²)であり、コンテキスト長を制限している
  2. 線形Transformers/RNN: 効率的だが、情報を固定サイズの状態に圧縮するため、長シーケンスの性能が低下する
  3. 記憶システムの欠如: 既存アーキテクチャは人間の脳のような多層的記憶システム（短期記憶、長期記憶、メタ記憶など）を欠いている

研究動機

著者は人間の記憶システムに着想を得て、効果的な学習パラダイムには以下が必要であると考えている：

  • 学習プロセスの重要な構成要素を各々が担当する異なるが相互に接続されたモジュール
  • データから能動的に学習し、過去の履歴の抽象化を記憶する能力
  • テスト時にも継続的に学習と適応が可能なメカニズム

核心的な貢献

  1. ニューラル長期記憶モジュール: テスト時にデータをパラメータに記憶/保存する方法を学習するメタモデルとしての深いニューラルネットワークを提案
  2. 記憶管理メカニズム: 「驚き度」に基づく記憶更新メカニズムと適応的忘却メカニズムを設計
  3. Titansアーキテクチャファミリー: 記憶を深層学習アーキテクチャに統合する3つの方法を提案：Memory as Context (MAC)、Memory as Gate (MAG)、Memory as Layer (MAL)
  4. 並列化訓練アルゴリズム: 深い記憶モジュールの効率的な訓練を可能にする高速並列化訓練アルゴリズムを提供
  5. 広範な実験検証: 言語モデリング、常識推論、ゲノミクス、時系列予測を含む複数のタスクでTitansの有効性を検証

方法の詳細

タスク定義

本論文はシーケンスモデリングタスクを研究し、入力はシーケンス x∈RN×dinx \in \mathbb{R}^{N \times d_{in}} であり、目標は長いシーケンスを効果的に処理できるモデルを学習することである。このモデルは以下を必要とする：

  • テスト時に継続的に学習と記憶が可能
  • 短期記憶と長期記憶の使用のバランスを取る
  • 線形複雑度を持ちながら高い表現力を維持する

ニューラル長期記憶モジュール

中核的な設計思想

人間の長期記憶に着想を得て、予期に反する（驚くべき）イベントはより記憶されやすい。著者は入力に対するニューラルネットワークの勾配を使用して「驚き度」を測定する。

記憶更新メカニズム

基本的な更新規則：

M_t = M_{t-1} - θ_t ∇ℓ(M_{t-1}; x_t)

改善された更新規則（モーメンタムの導入）：

M_t = M_{t-1} + S_t
S_t = η_t S_{t-1} - θ_t ∇ℓ(M_{t-1}; x_t)

ここで：

  • S_t: 驚き度モーメンタム、過去の驚き度と瞬間的驚き度を含む
  • η_t: データ依存の驚き度減衰パラメータ
  • θ_t: 瞬間的驚き度の統合程度を制御するパラメータ

忘却メカニズム

長いシーケンスを処理するため、適応的忘却メカニズムを導入：

M_t = (1 - α_t)M_{t-1} + S_t

ここで α_t ∈ [0,1] はゲート機構であり、忘却の程度を制御する。

目的関数

連想記憶損失関数を使用：

ℓ(M_{t-1}; x_t) = ||M_{t-1}(k_t) - v_t||²₂

ここで k_t = x_t W_K, v_t = x_t W_V

Titansアーキテクチャの変種

1. Memory as Context (MAC)

  • 記憶を現在の情報のコンテキストとして機能させる
  • シーケンスをチャンク単位で処理し、各チャンクが長期記憶をクエリして関連する履歴情報を取得
  • 注意機構が長期記憶情報の必要性を判定

2. Memory as Gate (MAG)

  • 一つのブランチはスライディングウィンドウ注意を短期記憶として使用
  • もう一つのブランチはニューラル記憶モジュールを長期記憶として使用
  • ゲート機構を通じて2つのブランチの出力を統合

3. Memory as Layer (MAL)

  • ニューラル記憶を深いネットワークの一層として機能させる
  • 順序処理：まず記憶層を通過し、その後注意層を通過
  • 既存のハイブリッドモデルの層級設計に類似

並列化訓練

訓練プロセスを行列乗算と合計操作を使用して再定式化することにより、効率的な並列化訓練を実現：

  • シーケンスをサイズbのチャンクに分割
  • 並列連想スキャンを使用してモーメンタム項を計算
  • テンソル化ミニバッチ勾配降下法により高速訓練を実現

実験設定

データセット

  • 言語モデリング: FineWeb-Eduデータセット、15B/30B トークン
  • 常識推論: PIQA, HellaSwag, WinoGrande, ARC-easy/challenge, SIQA, BoolQ
  • 長コンテキストタスク: RULER benchmark (S-NIAH), BABILong benchmark
  • 時系列: ETT, ECL, Traffic, Weather データセット
  • ゲノミクス: GenomicsBenchmarks データセット

モデルスケール

  • 170M、340M、400M、760Mパラメータのモデル
  • 訓練長：4K トークン
  • コンテキストウィンドウ：2M+ トークンまでスケーリング可能

比較手法

  • Transformers: Transformer++
  • 線形循環モデル: RetNet, GLA, Mamba, Mamba2, DeltaNet, TTT, Gated DeltaNet
  • ハイブリッドモデル: Samba, Gated DeltaNet-H2
  • 大規模モデル: GPT-4, Llama3, RecurrentGemma, Mistral

実験結果

言語モデリング性能

340Mパラメータモデルにおいて：

  • Titans (LMM): パープレキシティ26.18 (Wiki), 29.97 (LMB)
  • 最良ベースラインTTT: パープレキシティ27.44 (Wiki), 34.19 (LMB)
  • ハイブリッドモデル中のTitans (MAG) が最良：パープレキシティ25.07 (Wiki), 28.72 (LMB)

長コンテキストタスク

S-NIAHタスク（16Kシーケンス長）において：

  • Titans (MAC): S-NIAH-PK 98.4%, S-NIAH-N 97.4%, S-NIAH-W 95.2%
  • Mamba2: S-NIAH-PK 5.4%, S-NIAH-N 0.0%, S-NIAH-W 0.0%
  • TTT: S-NIAH-PK 88.4%, S-NIAH-N 4.4%, S-NIAH-W 0.0%

BABILongベンチマーク

  • Titansはfew-shot設定下でGPT-4を含むすべてのベースラインを上回る
  • fine-tuning設定下では、小規模Titansはパラメータ数が70倍大きいGPT-4を上回る

時系列予測

ニューラル記憶モジュールはすべてのデータセットにおいてベースライン手法を上回り、Mamba、Transformer、線形モデルベースの手法を含む。

アブレーション実験

各構成要素の貢献度（重要度順）：

  1. 重み減衰（忘却メカニズム）
  2. モーメンタムメカニズム
  3. 畳み込み層
  4. 永続記憶
  5. 深い記憶 vs 線形記憶

関連研究

線形循環モデル

  • 第1世代: RetNet, LRU, RWKV, S4/S5 - データ非依存の遷移行列を使用
  • 第2世代: Griffin, Mamba系列 - ゲート機構を導入
  • 第3世代: DeltaNet, TTT, Longhorn - メタ学習/オンライン学習に基づく更新規則

Transformer変種

  • 効率最適化: スパース注意、線形注意、I/O認識実装
  • セグメント化Transformer: RMTなどがチャンク間で情報を伝達するシンプルなベクトル記憶を使用

テスト時訓練

  • 初期の局所学習アルゴリズムに着想を得ている
  • MNMおよびTTT-layerと最も関連があるが、Titansは忘却メカニズムとモーメンタム更新を備えている

結論と考察

主要な結論

  1. 記憶システムの重要性: 多層的記憶システム（短期+長期+永続記憶）はシーケンスモデリングに不可欠である
  2. テスト時学習の有効性: テスト時の継続的学習と記憶は長シーケンス処理能力を大幅に向上させることができる
  3. アーキテクチャ設計の影響: MACおよびMAGアーキテクチャは従来のMAL層級設計より優れている
  4. スケーラビリティの検証: Titansは2M+ コンテキストウィンドウまで効率的にスケーリング可能である

制限事項

  1. 計算オーバーヘッド: 深い記憶モジュールはシンプルな行列状態より多くの計算リソースを必要とする
  2. 記憶深度のトレードオフ: より深い記憶モジュールはより良い効果をもたらすが、訓練がより遅い
  3. パラメータ感度: 驚き度関連パラメータの慎重な調整が必要である
  4. 理論分析の不足: 記憶容量と忘却戦略に関する理論的保証が不足している

今後の方向性

  1. 記憶アーキテクチャの最適化: より効率的なニューラル記憶アーキテクチャ設計の探索
  2. 理論分析: 記憶容量、忘却戦略の理論分析の提供
  3. 大規模検証: より大規模なモデルでの方法の有効性の検証
  4. 応用拡張: より多くの領域での応用可能性の探索

深い評価

利点

  1. 概念の革新性が強い: 人間の記憶システムの観点からシーケンスモデリングを再考し、新しい多層的記憶アーキテクチャを提案している
  2. 技術貢献が包括的: ニューラル記憶モジュールの提案だけでなく、3つのアーキテクチャ統合方法と効率的な並列化アルゴリズムも設計している
  3. 実験検証が充分: 複数の領域（NLP、時系列、ゲノミクス）で包括的な実験を実施し、説得力のある結果を得ている
  4. 理論基礎が堅実: 記憶更新を勾配降下法、モーメンタム、重み減衰と関連付け、理論的説明を提供している

不足点

  1. 計算複雑度分析が不十分: 線形複雑度を主張しているが、深い記憶モジュールの実際の計算オーバーヘッド分析が詳細でない
  2. 超パラメータ感度: 複数のデータ依存パラメータ(α_t, θ_t, η_t)の設定が複雑である可能性がある
  3. 人間の記憶との類比が限定的: 人間の記憶に着想を得ているが、「驚き度」の定義は相対的にシンプルである
  4. 大規模モデル検証の不足: 最大モデルは760Mパラメータのみであり、十億規模パラメータモデルの検証が不足している

影響力

  1. 学術的価値: シーケンスモデリングに新しい記憶の観点を提供し、関連研究を刺激する可能性がある
  2. 実用的価値: 長シーケンス処理タスクで優れた性能を示し、実用的応用の可能性がある
  3. 再現性: 著者がコードのオープンソース化を約束しており、方法の普及と検証を支援する

適用シーン

  1. 長文書処理: 長文書を処理する必要があるNLPタスクに適用可能
  2. 時系列分析: 長期の履歴情報が必要な予測タスクに特に適している
  3. オンライン学習シーン: テスト時に継続的に適応する必要があるアプリケーションに適用可能
  4. 記憶集約的タスク: 質問応答システム、対話システムなど大量の情報を記憶する必要があるタスク

参考文献

論文は138篇の関連文献を引用しており、Transformer、循環ニューラルネットワーク、注意機構、記憶ネットワーク、テスト時訓練など複数の関連領域の重要な研究をカバーしており、本研究に堅実な理論基礎を提供している。