【DL輪読会】 Titans: Learning to Memorize at Test Time

DEEP LEARNING JP [DL Papers] Titans: Learning to Memorize at Test Time Yuya IMAI, Matsuo Iwasawa Lab http://deeplearning.jp/ 1

• http://deeplearning.jp/

書誌情報 Titans: Learning to Memorize at Test Time Authors: Ali Behrouz, Peilin Zhong, and Vahab Mirrokni Affiliation: Google Research TL;DR 推論時にも新しい情報を効率的に記憶し続けることができる、長期記憶 モジュール「Neural Memory」を開発し、それを組み込んだ「Titans」 という高性能なモデルを作った 2

背景 既存の系列モデルのジレンマ • Transformer – 長所: Attention機構により、文脈内のトークン間の依存関係を正確に捉え（Incontext Learning）大規模なモデルで驚異的な性能を発揮。 – 短所: 計算量とメモリ使用量がコンテキスト長の 二乗 (O(N^2)) で増加。 • 線形RNN – 例: RetNet(2023), Mamba(2024), Gated DeltaNet(2024) – 試み: 計算量を 線形 (O(N)) に削減し、スケーラビリティを改善。 – 限界: 情報を固定サイズの隠れ状態に圧縮するため、長いコンテキストの情報 を正確に保持することが難しく、性能が低下しやすい。 ⇒ 高速・高効率でありながら、超長文の情報を正確に扱えるモデルが求められてい る。 3

背景 人間の記憶システムからの着想 • 人間の脳は、単一の記憶システムを持っているわけではない。 – 短期記憶 (Short-term Memory): 現在のタスクに関連する少量の情報を一時的に保持。 – 長期記憶 (Long-term Memory): 膨大な情報を長期間にわたり保持し、必要に応じて 取り出す。 – これらは独立しつつも相互に連携するシステムとして機能する。 • 既存のAIアーキテクチャの課題 – これらの複雑な記憶メカニズムを十分に模倣できていない。 – 単一のベクトルや行列に過去の情報を詰め込むのは、長期記憶の表現としては単純す ぎるのではないか？ • 本研究の問い – 効果的な「記憶」の構造とは？ – 適切な「記憶」の更新・検索メカニズムは？ – 複数の記憶モジュールをどう効率的に統合するか？ 4

手法 提案手法: Titans アーキテクチャの全体像 Titansの3つの主要コンポーネント: • Core（コア）: 短期記憶 – 限定的なウィンドウサイズの注意（Attention）機構を使用。 – 現在のコンテキスト内の正確な依存関係を捉える役割。 • Neural Long-term Memory（ニューラル長期記憶） – 本研究で提案する新しい記憶モジュール。 – 過去の膨大な履歴をそのパラメータ内に符号化し、テスト時にも学習・記憶を続ける。 • Persistent Memory（持続的記憶） – 学習可能だが、入力データに依存しないパラメータ群。 – タスクに関する普遍的な知識を保持する。 5

手法 コアコンセプト①: ニューラル長期記憶モジュール(1/3) - 学習プロセス • 目的: 過去の情報を「キー (k_t)」と「バリュー (v_t)」のペアとして、連想記憶のように パラメータ内に保存すること。 • 学習の目標: 記憶モジュール M が、与えられたキー k_t から対応するバリュー v_t を正 しく予測できるように以下の損失関数で評価する。 • 更新の基本形: この損失を最小化するように、勾配降下法で記憶モジュール M のパラ メータを更新する。この勾配が、モデルにとっての「驚き (Surprise)」の指標となる。 6

手法 コアコンセプト①: ニューラル長期記憶モジュール(2/3) - Momentum の導入 • 課題: 瞬間的な「驚き」だけでは、大きな出来事の後に続く重要な情報を見逃す可能性が ある（勾配消失）。 • 解決策: Momentumを導入し、「過去の驚きの記憶」と「現在の驚き」を統合すること で、文脈の流れを捉えた、より持続的な記憶が可能になる。 – 「驚きの記憶」の更新: – 記憶モジュールの更新: 7

手法 コアコンセプト①: ニューラル長期記憶モジュール(3/3) - 忘却と検索 • 忘却の必要性: 長大なシーケンスを処理する際、限られたメモリ容量を効率的に管理する ためには、不要になった情報を「忘れる」能力が不可欠。 – 忘却メカニズムの導入: 適応的な忘却ゲート α_t を導入し、過去の記憶をどの程度保 持するかを動的に制御する。これは、実質的に重み減衰 (Weight Decay) として機能 する。 – 最終的な更新式: • 記憶の検索 (Retrieval): 学習・更新された記憶モジュールから情報を取り出す際は、新 しいクエリ q_t を入力し、単純な順伝播（重みを更新しない推論）を行う。 8

手法 (再掲) Titans アーキテクチャの全体像 Titansの3つの主要コンポーネント: • Core（コア）: 短期記憶 – 限定的なウィンドウサイズの注意（Attention）機構を使用。 – 現在のコンテキスト内の正確な依存関係を捉える役割。 • Neural Long-term Memory（ニューラル長期記憶） – 本研究で提案する新しい記憶モジュール。 – 過去の膨大な履歴をそのパラメータ内に符号化し、テスト時にも学習・記憶を続け る。 • Persistent Memory（持続的記憶） – 学習可能だが、入力データに依存しないパラメータ群。 – タスクに関する普遍的な知識を保持する。 9

手法 コアコンセプト②: 持続的記憶モジュール(1/2) - 役割と実装 タスクの「知識」を格納する、文脈非依存のアンカー • 役割: 長期記憶が文脈に沿った「エピソード記憶」だとすれば、持続的記憶はタ スクの解き方など、文脈に依存しない普遍的な知識を保持する。 • 実装: 学習可能なパラメータ群 P を用意し、入力シーケンス x の先頭に連結する。 10

手法 コアコンセプト②: 持続的記憶モジュール(2/2) - 導入の動機 なぜ持続的記憶が必要なのか？ • 記憶の視点から: タスク遂行に必要な抽象的な知識（例:「要約せよ」という指示 の意味）を保持するため。 • FFNの視点から: TransformerのFFN層のように、データに依存しないキー・バ リューとして機能させるため。 • 技術的な視点から: Attentionがシーケンスの初期トークンに過度に集中するバイ アスを緩和し、注意の重みをより効果的に再分配するため。 11

手法 (再掲) Titans アーキテクチャの全体像 Titansの3つの主要コンポーネント: • Core（コア）: 短期記憶 – 限定的なウィンドウサイズの注意（Attention）機構を使用。 – 現在のコンテキスト内の正確な依存関係を捉える役割。 • Neural Long-term Memory（ニューラル長期記憶） – 本研究で提案する新しい記憶モジュール。 – 過去の膨大な履歴をそのパラメータ内に符号化し、テスト時にも学習・記憶を続け る。 • Persistent Memory（持続的記憶） – 学習可能だが、入力データに依存しないパラメータ群。 – タスクに関する普遍的な知識を保持する。 ⇒3種類のアーキテクチャーを提案 12

手法 アーキテクチャ①：MAC (Memory as a Context) 長期記憶を、アテンションが処理するための追加の文脈（コンテキスト）として扱 う。 過去の記憶から関連情報を検索し、それを現在の入力に連結してからアテンション に渡す。 13

手法 アーキテクチャ②：MAG (Gated Memory) 短期記憶（アテンション）と長期記憶（ニューラルメモリ）を並列に処理し、それ ぞれの出力をゲート機構で動的に混ぜ合わせる。 14

手法 アーキテクチャ③：MAL (Memory as a Layer) 長期記憶モジュールを、アテンション層の前段に置かれた一つの層（レイヤー）と して扱う。 データをまず記憶層に通し、その出力を次にアテンション層に渡すという直列的な 構成。 15

実験 実験①：基本性能の評価（言語モデリング・常識推論） 実験設定の概要 • 評価モデル: Titansの3つの派生形（MAC, MAG, MAL）および長期記憶モジュー ル（LMM）単体 。 • 比較対象: Transformer++, Mamba, Mamba2, Gated DeltaNetなどの最新リカレ ントモデルおよびハイブリッドモデル 。 • データセット: FineWeb-Eduデータセットで事前学習し、Wikitextや各種常識推 論タスクで評価 16

実験 実験①：基本性能の評価（言語モデ リング・常識推論） 結果 • LMM単体の性能: 長期記憶モジュール （LMM）単体でも、Transformer++を含む 他の非ハイブリッドモデルを上回る性能を 達成 。 • ハイブリッドモデルの性能: Titansのハイブ リッド版（MAC, MAG, MAL）は、Samba やGated DeltaNet-H2といった他のハイブ リッドモデルよりも優れた結果を示した 。 • アーキテクチャの優位性: 特にMACとMAG の性能が高く、文脈やゲートとして記憶を 統合するアーキテクチャの有効性が示され た。 17

実験 実験②：有効文脈長の検証 (Needle-in-a-Haystack) タスクの概要 • Needle-in-a-Haystack (NIAH) は、長いダミーテキスト（Haystack）の中から特定の情 報（Needle）を検索させることで、モデルの真の有効文脈長を測定するタスク 。 結果 18

実験 その他 言語モデル以外のモデル（時系列予測やDNAモデリング）としても高い性能を発揮 19

まとめ 結論と貢献 • 新しい記憶パラダイムの提示: テスト時にも学習・記憶・忘却を行う「ニューラ ル長期記憶モジュール」を提案。 • 革新的なアーキテクチャの設計: 注意機構（短期記憶）とニューラル記憶（長期 記憶）を効果的に統合するアーキテクチャファミリー「Titans」（MAC, MAG, MAL）を提示。 • 優れた性能の実証: 言語モデリング、常識推論、ゲノム解析、時系列予測など、 多様なタスクで最先端の性能を達成。 • 超長文コンテキストのブレークスルー: 特に200万トークンを超える超長文の読 解・推論タスクにおいて、既存の巨大言語モデルを凌駕する能力を実証。 20

まとめ 今後の展望 • より効率的で表現力の高いニューラル記憶アーキテクチャの探求。 • 異なる記憶モジュールの最適な組み合わせ方のさらなる研究。 • Titansのアーキテクチャを、マルチモーダルなど、より多様な領域へ 応用。 • 後継研究 – ATLAS: Learning to Optimally Memorize the Context at Test Time – https://arxiv.org/abs/2505.23735 21

• https://arxiv.org/abs/2505.23735