（論文紹介）Titans: Learning to Memorize at Test Time
タイタン：テスト時記憶学習
Google Research

こんにちはmakokonです。今日紹介するのはグーグルリサーチの報告したタイタンについてです。なかなかに面白い話でした。
この論文では、Titansという新しいAIモデルを紹介しています。Titansは、長い文章や情報を記憶し、理解するのが得意で、従来のAIモデルよりも長期間の文脈を考慮した対応が可能です。
例えば、過去の会話を覚えて会話したり、複雑な文章（論文など）を理解したり、過去の経験を基に判断したりすることができます。これにより、日常生活では、よりパーソナライズされたアシスタントとして、また、研究や専門分野では、高度な情報処理や意思決定のサポートが期待できます。
ただし、倫理的な判断や即時性が求められるタスクでは、AIに頼りすぎず、人間の判断も重要です。Titansは、AIの可能性を広げる画期的な技術ですが、使い方には注意が必要です。

技術的には、タイタンは、まるで人間のように驚くべき入力ほど記憶に残りやすいという長期記憶モジュールを実現しています。
タイタンの記憶機能は以下の３つです。

• 短期記憶：
  コンテキストウインドウのateentionメカニズム

• 長期記憶：
  提案されたニューラルモジュール

• 永続メモリ：
  タスクに関する記憶の保持

論文紹介

導入

この論文は、人間の記憶システムにインスパイアされた新しいニューラルメモリモジュールを設計し、長文脈を効果的に処理できる深層学習モデルの構築を目的としています。

• 既存のモデルの限界を克服する:
  Transformerモデルは、長文脈を扱う際に計算コストが大きくなるという課題があります。また、従来のモデルは短期記憶、長期記憶、メタ記憶などの重要な要素を欠いています。これらの課題を解決するために、新しい記憶メカニズムが必要です。

• 人間の記憶システムから着想を得る:
  人間の記憶システムは、短期記憶、長期記憶、ワーキングメモリなどの異なるシステムで構成され、それぞれ異なる機能を持っています。この人間の記憶システムをモデルに取り入れ、より効果的な学習パラダイムの構築を目指しています.

• 驚きの概念を利用する:
  人間は、予想に反する事象（驚き）をより記憶しやすいという性質があります。この概念をモデルに取り入れ、入力に対するモデルの勾配を驚きの指標として使用し、より重要な情報を選択的に記憶できるようにします。

• 効率的な学習と記憶を実現する:
  提案するニューラルメモリモジュールは、並列化可能な学習アルゴリズムを採用し、効率的なトレーニングが可能です。また、忘却メカニズムを取り入れることで、メモリの容量を効果的に管理しています。

• Titansアーキテクチャを提案する:
  提案するニューラルメモリモジュールを組み込んだ新しいアーキテクチャであるTitansを提案し、その有効性を検証しています。

• 長文脈での性能を向上させる:
  長文脈を扱うタスクにおいて、従来のTransformerモデルや線形回帰モデルよりも優れた性能を発揮することを目指しています。モデルの長文脈処理能力を検証しています。

• 現実世界の多様なタスクに対応する: 言語モデリング、常識推論、ゲノミクス、時系列予測など、様々なタスクで提案手法の有効性を検証し、現実世界の問題への応用を目指しています。

準備　タイタンの設計に対するモチベーション

準備 (Preliminaries)」の章では、Attentionメカニズムの基本と線形RNNのメモリに関する考察が、Titansの設計につながる重要なモチベーションとなっています。

従来技術の課題の考察

• Attentionメカニズムは、文脈長の二乗に比例する計算コストが課題で、長文脈では非効率です。

• 線形RNNは、計算コストは低いものの、長文脈情報を小さなメモリに圧縮するのが難しいという課題があります。従来のモデルは加算的にメモリに書き込むため、メモリオーバーフローを引き起こす可能性があります。また、多くのモデルは瞬間的な驚きに基づいており、トークンの流れを捉えきれていません。

タイタンが解決したい課題

• 効率的な長文脈処理:
  Attentionの課題を克服し、線形RNNの限界を超えるため、長期的なニューラルメモリモジュールを導入します。

• 人間の記憶を模倣:
  短期記憶（Attention）と長期記憶（ニューラルメモリ）を組み合わせ、驚きを指標に情報を選択的に記憶します。

• 忘却メカニズム:
  メモリを効率的に管理するため、忘却メカニズムを導入します。

• トークンの流れを考慮:
  トークンの流れを考慮した、より効果的なメモリ更新ルールを設計します。

• テスト時の学習:
  テスト時にも学習を続け、情報を記憶・更新する能力を持たせます。

• 持続的な記憶:
  タスク知識を格納する入力に依存しないパラメータを利用します。

長期記憶の設計

長期記憶の設計とその手法について説明しています。
この長期記憶ニューラルモジュールは、過去の履歴の抽象化をパラメータにエンコードするメタモデルとして機能し、オンライン学習問題として扱われます。各種技術要素を含みます。
これらの技術要素を組み合わせることで、モデルはテスト時に過去の情報を効果的に学習し、記憶し、検索することが可能になります。また、これらの手法は、現代のリカレントモデルにおける忘却メカニズムの一般化としても見なすことができます。

目的:

• 長期記憶の欠如を克服し、モデルが情報を学習、忘却、検索できるようにする。

• 訓練データへの過剰適合を避け、テスト時の一般化を向上させる。

• 人間の長期記憶システムに着想を得て、驚くべきイベントをより記憶に残るようにする。

手法:

• 驚きメトリクス: 入力に対するニューラルネットワークの勾配を驚きの尺度として使用します。勾配が大きいほど、入力データが過去のデータと異なることを示します。

• 過去の驚きと一時的な驚き: 驚きメトリクスを、過去の驚き（最近の過去の驚きの量）と一時的な驚き（入力データの驚き）に分解します。

• モメンタム: 過去の驚きを考慮することで、勾配降下法におけるモメンタムのように機能し、時間の経過とともに驚きの記憶として作用します。

• データ依存の減衰: 驚きの減衰をデータに依存させ、文脈の変化に応じて過去の驚きを無視したり、完全に組み込んだりできるようにします。

• 損失関数: 連想記憶を目的として、キーと値のペアとして過去のデータを格納し、キーと値の関連性を学習するように損失関数を定義します。

• 忘却メカニズム: メモリの容量を管理するために、不要になった情報を忘れるための適応的な忘却メカニズムを使用します。これにより、メモリの容量をより効果的に管理できます。

  • 忘却ゲート（𝛼）は、メモリを柔軟に制御し、過去の抽象化に影響を与えずにメモリを更新したり、メモリ全体をクリアしたりできます。

• 並列化可能なトレーニング: ミニバッチ勾配降下法、データ依存の学習率、および重み減衰を使用することで、トレーニングプロセスを並列化し、行列演算を最大限に活用します。

  • シーケンスをチャンクに分割し、各チャンク内で勾配降下法を計算することで、並列処理を可能にします。

  • 図1: ニューラルメモリの並列トレーニングの概念図が示されています。この図は、トレーニングプロセスがmatmul演算を使用してどのように並列化できるかを示しています。

• 深層メモリ: 長期的な履歴を効果的に格納するために、深層メモリモジュール（MLPを複数層重ねたもの）を使用します。これにより、線形モデルよりも表現力が向上します。

• メモリの検索: クエリに対応するメモリを検索するために、重みの更新なしでフォワードパス（推論）を使用します。

アーキテクチャ　３つのバリエーションを比較

ニューラル長期記憶モジュールを効果的に組み込むための3つの異なるバリエーションが提案されています。これらのバリエーションは、メモリの利用方法と、それらがアーキテクチャ内の他のモジュールとどのように相互作用するかに焦点を当てています。

バリエーションの比較一覧

1. メモリをコンテキストとして利用 (Memory as a Context, MAC)

• 概要: このアーキテクチャでは、長期記憶を現在のコンテキストに対する追加のコンテキストとして扱います。

• 処理:

  1. 現在のセグメントの入力をクエリとして使用し、長期記憶モジュールから関連情報を検索します。

  2. 永続メモリ（タスクに関する知識をエンコードする学習可能なパラメータ）と、長期記憶から検索された情報を現在のセグメントと連結します。

  3. 連結されたシーケンスに対して注意機構（Attention）を適用し、どの情報が長期記憶に格納されるべきかを決定します。

  4. 注意機構の出力を利用して長期記憶モジュールを更新し、最終的な出力を生成します。

2. メモリをゲートとして利用 (Memory as a Gate, MAG)

• 概要: このバリエーションでは、入力データを直接利用して長期記憶を更新するブランチと、スライディングウィンドウ注意（SWA）を利用するブランチを組み合わせます。

• 処理:

  1. 入力データと永続メモリを連結し、スライディングウィンドウ注意機構（SW-Attn）を適用して短期記憶を処理します。

  2. 別のブランチでは、入力データを使用して長期記憶モジュールを更新します。

  3. スライディングウィンドウ注意機構の出力と長期記憶モジュールの出力を、非線形ゲート機構を使って結合します。

3. メモリをレイヤーとして利用 (Memory as a Layer, MAL)

• 概要: このバリエーションでは、ニューラルメモリを深層ニューラルネットワークの1つのレイヤーとして利用します。

• 処理:

  1. 入力データと永続メモリを連結します。

  2. 連結されたデータは、ニューラルメモリモジュールを通過し、過去と現在のコンテキストを圧縮します。

  3. メモリモジュールの出力に対してスライディングウィンドウ注意機構を適用します。

実験　長期記憶ニューラルモジュールの性能

実験結果は、Titansが長文脈タスクにおいて、Transformerや最近の線形再帰モデルよりも効果的であることを示しています。これは、Titansのニューラル長期記憶モジュールが、過去の情報を効果的に記憶し、現在のコンテキストを理解するのに役立つためです。また、深いメモリやモメンタム、重み減衰などのメカニズムが、モデルの性能向上に重要であることが示唆されています。

言語モデリングと常識推論

• 長期記憶ニューラルモジュールの目的との関連: 長期記憶モジュールが、過去の情報を効果的に記憶し、現在のコンテキストを理解するのに役立つことを示しています。

• 概要: Titansと、Transformerやその他の線形再帰モデルなどのベースラインモデルを、言語モデリングと常識推論タスクで比較しています。

• 結果:

  • 表1: Titansのすべてのバリエーション (MAC, MAG, MAL) が、Transformer++ やその他のベースラインモデルを上回る性能を示しています。特に、Titansのニューラルメモリモジュール単体 (LMM) でも、他の非ハイブリッドモデルよりも優れた性能を示しています。

  • ハイブリッドモデルでは、Titans (MAC, MAG, MAL) が、SambaやGated DeltaNet-H2などの既存のハイブリッドモデルよりも優れた結果を示しています。特に、Titans (MAL) の優れた性能は、ニューラルメモリの強力さによるものと考えられます。

  • Titans (MAG) と (MAC) を比較すると、長い依存関係を扱う場合、MACの方が優れていることが示されています。

Needle in a Haystack (NIAH) タスク

• 長期記憶ニューラルモジュールの目的との関連: Titansの長期記憶モジュールが、非常に長い文脈でも、重要な情報を効果的に記憶し、検索できることを示しています。

• 概要: モデルが長文脈から特定の情報を抽出する能力を評価するためのタスクです。

• 結果:

  • 表2: ニューラルメモリモジュールが、すべてのシーケンス長において、ベースラインモデルよりも優れた性能を示しています。特に、シーケンス長が増加しても性能が低下せず、一貫した傾向を示しています。

  • これは、Titansが、従来のモデルと比較して、より効果的にメモリ容量を管理し、長い依存関係を捉えることができることを示唆しています。

• 長期記憶ニューラルモジュールの目的との関連: この結果は、Titansの長期記憶モジュールが、非常に長い文脈でも、重要な情報を効果的に記憶し、検索できることを示しています。

BABILongベンチマーク

• 長期記憶ニューラルモジュールの目的との関連: 長期記憶モジュールに深いアーキテクチャを使用することが、より効果的な記憶と学習につながることを示しています。

• 概要: 非常に長いドキュメントに分散した事実に基づいて推論を行う、より難しいタスクです。

• 結果:

  • 図6a: few-shot設定では、Titans (MAC) が、GPT-4を含むすべてのベースラインモデルを上回る性能を示しています。

  • 図6b: fine-tuning設定では、Titans (MAC) が、GPT-4やLlama3.1-70Bなどの非常に大規模なモデルを含む、すべてのモデルを上回る性能を示しています。特に、RMTなどのTransformerベースのメモリモデルと比較して、Titansが優れた性能を発揮するのは、モデルのパラメータに過去の情報をエンコードできる強力なメモリを持っているためです。

メモリの深さの影響

• 長期記憶ニューラルモジュールの目的との関連: 長期記憶モジュールに深いアーキテクチャを使用することが、より効果的な記憶と学習につながることを示しています。

• 概要: 長期記憶モジュールの深さ (レイヤー数) が、モデルの性能にどのように影響するかを評価しています。

• 結果:

  • 図7: メモリの深さが増すにつれて、モデルはすべてのシーケンス長でより良いperplexityを達成できます。これは、より深いメモリが、より効果的に情報を記憶し、複雑な依存関係を捉えることができることを示しています。

  • また、パラメータ数が少ないモデルでは、深いメモリを持つ方が長いシーケンスに対してよりロバストであることが示されています。

時系列予測とDNAモデリング

• 長期記憶ニューラルモジュールの目的との関連: 長期記憶ニューラルモジュールが、さまざまな種類のデータに対して、効果的に情報を記憶し、学習できる汎用的なモジュールであることを示しています。

• 概要: 長期記憶モジュールが、より幅広いタスクで有効であることを示すために、時系列予測とDNAモデリングタスクで性能を評価しています。

• 結果:

  • 表3: 時系列予測タスクにおいて、Titansのニューラルメモリモジュールは、Mambaベース、線形ベース、Transformerベースのアーキテクチャを含む、すべてのベースラインモデルよりも優れた性能を示しています。

  • 表4: DNAモデリングタスクにおいて、Titans (LMM) は、さまざまな下流のゲノミクス分類タスクにおいて、最先端のアーキテクチャと同等の性能を示しています。

アブレーションスタディ

• 概要: Titansの各コンポーネントが性能にどのように貢献しているかを評価しています。

• 結果:

  • 表5: すべてのコンポーネント（深層メモリ、畳み込み、モメンタム、重み減衰、永続メモリ）が、モデルの性能にプラスに貢献していることが示されています。特に、重み減衰とモメンタムが最も大きな貢献を示しています。

  • また、アーキテクチャ設計においては、MACとMAGが言語モデリングと常識推論タスクで同等の性能を示し、長文脈NIAHタスクではMACが優れていることが示されています。

この論文のまとめ

長文脈データ処理に特化した新しいニューラル長期記憶モジュールを提案し、それを用いたTitansという新しいアーキテクチャを開発しました。Titansは、Transformerや既存の線形再帰モデルよりも長文脈で優れた性能を発揮し、特に2M以上の文脈窓サイズでも高い精度を達成しています。
Titansは長文脈データ処理における新たな可能性を切り開くものと言えます。
主な成果。

• 長期記憶モジュール：過去の情報を効果的に記憶し、現在の文脈理解を助ける。

• Titansアーキテクチャ：このモジュールを組み込んだ3つのバリエーションを提示。

• 長文脈での有効性：言語モデリング、常識推論、NIAH、BABILongタスクで、最先端モデルを凌駕する性能。

• 深層メモリの効果：より深いメモリが、長文脈での性能向上に貢献。

• 多様なタスクでの応用：時系列予測やDNAモデリングでも有効性。

• 効率性：Titans（MAL）はベースラインモデルよりも高速なスループットを達成。

感想　makokonの考察

タイタンが搭載されたＡＩエージェントは役に立ちそうかな。

効果的な長期記憶を扱う能力は、重要性と効率を保持したまま長文脈データを処理する能力と言うことです。また、人間が管理できないほどの古い記憶、広範囲な記憶も効果的に利用することができます。この能力は従来ののＡＩエージェントと比較して広く深い応用範囲が期待できるような気がします。
単純に適応を考えるなら、以下のような分類になるでしょう。

効果的なエージェントと、あまり効果が期待できないエージェント

• 効果的なエージェント:

  • 長文脈データ処理が必要なタスク: 過去の出来事や情報を長期的に記憶する必要があるタスク（会話、意思決定、情報分析など）。

  • 複雑な推論や判断を伴うタスク: 複数の情報を統合し、長期的な視点から推論や判断を行うタスク（専門分野での意思決定など）。

  • パーソナライゼーションが求められるタスク: ユーザーの過去の行動や好みを理解し、個別のニーズに合わせた対応が求められるタスク（家庭でのアシスタントなど）。

• あまり効果が期待できないエージェント:

  • 短い文脈や単発のタスク: 過去の情報を必要としない、シンプルなタスク。

  • 倫理的判断や価値観が絡むタスク: 人間の専門家による判断や監督が必要となるタスク。

  • 即時性が求められるタスク: 迅速な対応が必要となる場面では、Titansの処理時間がネックになる可能性がある。

具体的にいくつか考えてみましょう。

家庭生活における話し相手

• 効果的なエージェント: Titansは、会話の文脈を長期にわたって記憶し、関連性のある情報を引き出す能力があるため、家庭での話し相手として非常に有効です。例えば、家族の過去の出来事や好みを覚えて、よりパーソナライズされた会話を提供できます。また、感情的なサポートが必要な場面でも、文脈を理解した上で適切な対応が可能になります。

• あまり効果が期待できないエージェント: 短い会話や単発の質問には、従来のモデルでも十分な場合があります。また、Titansの複雑な計算処理は、シンプルなタスクに対してはオーバースペックになる可能性もあります。

• さらに、会話のテンポが悪くなったり、感情的なサポートが余計なお世話になるかもしれませんね。

意思決定アシスタント

• 効果的なエージェント: Titansは、過去の意思決定とその結果を記憶し、複雑な状況下での意思決定を支援できます。例えば、過去の金融取引のデータを分析し、将来の投資戦略を提案したり、過去のプロジェクトの経験から、プロジェクトのリスクを評価したりすることができます。また、複数の情報を統合して、より多角的な視点から意思決定を支援できます。

• あまり効果が期待できないエージェント: 倫理的判断や価値観が絡む意思決定には、AIエージェント単独では限界があります。また、迅速な判断が求められる場面では、Titansの処理時間が問題になる可能性もあります。

• 同じ人が使うにしても明らかに場面を選びそうな気がしますね。一人のときにゆっくり考えるのに手伝ってもらって、現場では臨機応変に自分で判断しないとだめかも。

専門性の高い分野（医療、法律など）でのアシスタント

• 効果的なエージェント: Titansの長文脈処理能力は、専門分野での情報検索、分析、推論において非常に有効です。

  • 医療: 患者の病歴や検査結果を長期にわたって記憶し、適切な診断や治療計画を提案するのに役立ちます。

  • 法律: 過去の判例や法律文書を記憶し、法的解釈や訴訟戦略を支援するのに役立ちます。

• あまり効果が期待できないエージェント: 専門的な知識や倫理観が求められる判断には、人間専門家の監督が不可欠です。また、法律や医療分野では、AIエージェントの誤った判断が重大な結果を招く可能性があるため、注意が必要です。

• こういう専門分野について考えてみたものの、顧客の立場として考えると、ＡＩエージェントに相談しながら答えてくれる、医者や弁護士に相談したいとは思わないかなあ。実際、専門家は個人的な勉強に利用してもらって、顧客向けには、簡単な相談エージェントのようなアプリを作って、簡単な質問には答えるけど、ある程度のレベルからは専門家への相談を促すサービスを考えたほうがいいかもしれませんね。

研究助手、開発助手

研究助手としての様々なタスクにおいて非常に有望な応用可能性を感じます。もちろんあくまでも人間が主役で助手的な働きが前提ですが、Titansが効果的なエージェントとなるための条件がありそうです。

• 長文脈の理解が必要なタスク: 論文読解、論文作成、実験計画など、過去の情報や複数の情報を統合する必要があるタスク。

• 文脈の重要度に応じた適応的な記憶: どの情報を記憶し、どの情報を忘れるかを判断する能力。Titansは、驚き（surprise）の指標を用いて、より重要な情報を記憶するように設計されています。

• 反復的なプロセス: 実験計画のように、過去の経験を基に反復的に改善していくプロセス。

• 論文読解

  • 効果的な点:

    • 長文脈の理解: 論文全体の内容を一度に理解し、個々の段落や文脈の関係性を把握できます。これにより、論文全体の主題や論理構造をより深く理解できます。

    • 情報の関連付け: 複数の論文や文献間の情報を関連付け、研究分野全体の知識マップを構築するのに役立ちます。これにより、既存研究のギャップや新たな研究テーマを発見する可能性が高まります。

    • 詳細な分析: 論文中の重要な記述や実験結果を正確に抽出し、批判的な分析を支援できます。

• 論文作成

  • 効果的な点:

    • 一貫性のある文章生成: 論文全体を通して一貫性のある論理構造とスタイルを維持した文章を作成できます。これにより、質の高い論文作成を支援できます。

    • 参考文献の統合: 関連する参考文献を自動的に抽出し、適切に引用することで、研究の信頼性を向上させることができます。

    • アイデアの生成: 過去の研究や自身の思考過程を記憶し、新しい研究アイデアやアプローチを創出するのに役立ちます。

• 実験計画

  • 効果的な点:

    • 過去の実験結果の分析: 過去の実験結果を記憶し、それらの結果を基により効果的な実験計画を立案できます。

    • 実験条件の最適化: 実験条件を最適化するために、過去の成功例や失敗例を考慮に入れることができます。

    • 実験プロセスの自動化: 実験プロセスの自動化を支援し、時間と労力を節約できます。

ロングコンテキスト対応LLMで対応できない可能性のある点

こういうエージェントは最近のロングコンテキスト対応ＬＬＭならあまり気にせずに実現できるかなとは思いましたが、おそらく以下のような点で差別化されると考えています。

• 継続的な学習: Titansは、テスト時にも学習を続けることができ、長期的な記憶と適応能力においてLLMよりも優れている可能性があります。これにより、長期にわたる研究プロジェクトにおいて、より効果的なサポートが期待できます。

• 複雑な情報処理: Titansは、深層ニューラルネットワークを記憶モジュールとして使用できるため、より複雑な情報を表現し、処理できる可能性があります。これにより、LLMでは困難な高度な推論や情報統合が可能になる可能性があります。

• オンラインメタ学習: Titansは、オンラインメタ学習の枠組みで設計されており、テスト時にもデータの学習を続け、その場での状況適応能力を向上させることができます。

ハッシュタグ

#AI #機械学習 #深層学習 #自然言語処理 #長文脈 #長期記憶 　
#ニューラルネットワーク #Transformer #リカーレントモデル #TitansAI #AIアシスタント #研究支援AI #オンライン学習 #メタ学習 #効率的なAI #google