ティターンズを徹底解剖：長期記憶の収束と効率的なシーケンス・モデリングへの道

巨人：テスト時の暗記学習 オリジナル：https://arxiv.org/pdf/2501.00663v1

ティターンズ・アーキテクチャ非公式実装：https://github.com/lucidrains/titans-pytorch

一人だ、 背景と動機：変圧器の限界と人間の記憶のひらめき

1. Transformerの限界：長いシーケンス処理におけるボトルネック

変圧器 導入以来、このモデルはその強力な自己注意メカニズムTransformerは、自然言語処理やコンピュータビジョンなどの分野で画期的な進歩を遂げた。しかし、タスクの複雑さが増すにつれて、Transformerは長いシーケンスを扱う際にいくつかの致命的な問題を徐々に露呈していく：

• 高い計算複雑性がモデルのスケーラビリティを制限している :.
  • 自己アテンション・メカニズムでは、次のような計算が必要となる。 トークン の時間と空間の複雑さを持つ。 O(N²)どの N はシーケンス長である。
  • これは、シーケンスの長さが長くなると、計算量とメモリ消費量が2乗ステップで増加することを意味し、長いシーケンスを処理するモデルの能力を著しく制限する。例えば、長いテキスト、ビデオ理解、長期の時系列予測などのタスクを扱う場合、Transformerはその深みにはまらないことが多い。
  

  図1：自己注意メカニズムの計算過程。

• 限定されたコンテキスト・ウィンドウで長距離の依存関係を把握する :.
  • 計算負荷を軽減するために、Transformerは通常、固定長のコンテキストウィンドウ（例えば512や1024）を使用する。
  • しかし、実世界の多くのタスクでは、長いテキストやダイアログの文脈情報を理解したり、ビデオの異なる時点の情報を統合したり、過去のデータの長期的な傾向やパターンを使って予測を行うなど、より長い時間スケールでの依存関係を捉えることができるモデルが必要とされる。

2. リニアトランスの妥協：効率対性能

Transformerの計算ボトルネックを解決するために、研究者たちは、Transformerの計算ボトルネックを解決するための方法を提案した。リニアトランス主な改善点は以下の通り：

• ソフトマックスをカーネル関数に置き換える。 自己アテンション機構のソフトマックス計算をカーネル関数に置き換えることで、計算量は以下のように減少する。 O(N).
• 並列可能な推論。 リニア・トランスフォーマーの計算過程は環状に表現することができ、より効率的な推論が可能になる。

しかし、リニア・トランスにも限界がある：

• パフォーマンスの低下。
  • カーネルのトリックは、モデルを線形リカレントネットワークに劣化させ、データは行列値の状態に圧縮される。
  • このタイプの圧縮では、複雑な非線形依存関係を効果的に捉えることは難しい。
• メモリ管理の問題。
  • 線形変換器は、過去のデータを固定サイズの行列に圧縮するが、非常に長いコンテキストを扱う場合、この圧縮の結果はメモリオーバーフローこれはモデルの性能に影響する。
  

  図2：リニア・トランスのメモリ更新プロセス。

3. ヒトの記憶システムのインスピレーション：長期記憶のより強力なメカニズムの構築

上記の課題を克服するために、著者たちは次のことを始めた。人間の記憶システムからインスピレーションを得ている。

• 記憶と学習の関係
  • 本論文では、神経心理学の文献から記憶と学習の定義を引用し、記憶を入力によって誘発される神経更新とみなし、学習を目標が与えられたときに効果的で有用な記憶を獲得する過程と定義する。
  • これはつまり、次のことを意味する。効果的な学習は、強力な記憶メカニズムなしには達成できない.
• 人間の記憶の多階層性：「記憶の多階層性」とは何か？
  • 人間の記憶システムは単一構造ではなく、短期記憶、ワーキングメモリー、長期記憶といった複数のサブシステムから構成されており、それぞれが異なる機能と組織構造を持ち、独立して作動することができる。
  • この多層的な性質が、人間が情報を効率的に保存、検索、管理することを可能にしている。
• 既存モデルの欠点。
  • 既存のニューラルネットワークアーキテクチャ（ホップフィールドネットワークからLSTM、トランスフォーマーまで）には、汎化、長さの外挿、推論を処理する上での課題がある。
  • これらのアーキテクチャーは、人間の脳からヒントを得ているとはいえ、いずれも以下のような能力に欠けていた。儚い記憶を効果的にモデル化した。記憶システムのマルチレベル性シミュレーションだ。

二人だ、 コア・イノベーション：神経長期記憶モジュールとティターンズ・アーキテクチャ

これらの反省に基づき、著者らは以下のようなイノベーションを提案する：

1. 神経長期記憶モジュール

(1) デザインコンセプト。

• メタコンテクスト学習メカニズム。
  • このモジュールはメタモデルテスト中にデータをパラメータに記憶／保存する方法を学ぶ。
  • この種のオンライン学習のアプローチは、事前学習中に固定された記憶に頼るのではなく、モデルが動的に現在の入力に記憶を適応させることを可能にする。
• サプライズ・ベース・メモリー・アップデート
  • 著者らは、「意外な出来事の方が記憶に残りやすい」という人間の記憶メカニズムを利用し、次のような方法を提案している。驚きの度合い記憶更新のメカニズム。
  • 驚きの度合い入力に対してニューラルネットワークを計算することによって勾配を測定する場合、勾配が大きければ大きいほど、入力データは過去のデータとは異なり、記憶される価値がある。
    • この方法は、データの重要な情報を捕捉し、長期記憶に保存するのに有効である。
    • 一方、リニア・トランスフォーマーは、現在の入力データに基づいて線形変換を行うことしかできないため、長距離の依存関係を効果的に捉えることが難しい。
  

  図3：サプライズベースのメモリ更新メカニズム。

(2) キーテクノロジー。

• 勢いのメカニズム。
  • 一つの意外な出来事によってモデルが過度に影響されるのを防ぐため、著者らは、次のように導入した。運動量メカニズム過去の瞬間の驚きを考慮に入れて。
  • つまり、モデルは現在の入力と過去の入力の両方の不意打ちを考慮し、よりスムーズなメモリー更新につながる。
• 減衰メカニズム。
  • メモリのオーバーフローを防ぐために、著者らは崩壊メカニズムを使用する。荷重減衰重要でない情報を徐々に忘れていく方法。
  • このメカニズムはゲート機構必要に応じてメモリを選択的に消去することができる。
  • 著者らは、この減衰メカニズムは最新のリカレントモデルにおける忘却メカニズムの一般化であり、小バッチの勾配降下、運動量、重みの減衰のもとでのメタニューラルネットワークの最適化と等価であると指摘している。

(3) メモリ構造。

• 従来の線形記憶モデルとは対照的に、著者たちは多層パーセプトロン（MLP） をメモリモジュールとして使用する。
  • MLPはより強力な非線形表現を持ち、複雑な情報をより効率的に記憶・検索することができる。
  • 対照的に、線形トランスフォーマーは、情報を格納するために行列値の状態しか使用できないため、複雑な非線形関係を捉えることが難しい。

2. タイタンズ・アーキテクチャー：長期記憶と短期記憶の統合

Neural Long-Term Memoryモジュールを設計した後、著者らはさらに、これをディープラーニング・アーキテクチャに効果的に組み込む方法を考え、以下の主要な特徴を持つTitansアーキテクチャを提案した：

(1) 3つのスーパーヘッドが連動して働く。

• コアだ。
  • 短期記憶から成り、データを扱う主要なプロセスを担う。
  • 利用する有限のウィンドウ・サイズに対する注意のメカニズム例えば、スライディング・ウィンドウ・アテンション（SWA）やフル・コネクテッド・アテンション（FCA）などである。
  • 短期記憶とは次のようなものである。短期記憶これは、現在のコンテキストにおける依存関係をキャプチャするために使用される。
• 長期記憶。
  • 過去の情報を記憶する。
  • 上記の神経長期記憶モジュールを使用する。
  • 長期記憶は次のように捉えることができる。儚い記憶これは、より長い期間にわたって情報を保存し、検索するために使用される。
• 永続的な記憶。
  • は学習可能だがデータに依存しないパラメータの集合で、タスクに関する先験的な知識をエンコードする。
  • Transformerの完全接続レイヤーのパラメーターに似ているが、機能は異なる。
  • 永続的記憶とは次のようなものである。メタメモリーこれは、文法規則や常識知識など、タスクに関連した知識を保存するために使われる。
  

  図4：ティターンズ・アーキテクチャ（MACバリアント）の概略図。

(2) 3つの異なるタイプの統合

• コンテキストとしての記憶（MAC）。
  • 長く持続的な記憶を、現在の文脈の追加情報として入力シーケンスに結びつける。
  • 注意のメカニズムは、どの情報を長期記憶に保存する必要があるかを決定する。
  • テスト時、永続記憶パラメータは固定され、注意モジュールは文脈学習にウェイトを置き、長期記憶モジュールは情報を学習／記憶し続けた。
  • この設計により、モデルは現在の入力に基づいた長期記憶情報を柔軟に利用することができる。
• ゲーティングとしての記憶（MAG）。
  • スライディングウィンドウ・アテンションを短期記憶として、神経記憶モジュールを長期記憶として使う。
  • 例えば、学習可能なベクトル値の重みを使って2つを正規化し、非線形活性化関数を適用する。
  • この設計は、異なるヘッドが異なる構造を持つマルチヘッドアーキテクチャと見なすことができる。

  
  図5：ティターンズ・アーキテクチャの異なるバリエーション（MACとMAG）。

• メモリー・アズ・ア・レイヤー（MAL）。
  • ニューラル・メモリー・モジュールをディープ・ニューラル・ネットワークのレイヤーとして使うことで、注意モジュールの前に過去と現在の文脈情報が圧縮される。
  • このデザインは、例えばH3モデルのように、文献ではより一般的である。

(3)強み。

• より柔軟なメモリ管理。
  • メモリモジュールを文脈分岐やゲート分岐として使用することで、ティターンズ・アーキテクチャは、現在の入力に基づいて長期記憶情報を動的に利用することができる。
  • これにより、メモリモジュールをレイヤーとして使用する従来のアプローチよりも柔軟性が増す。
• より強い表現。
  • 3つのスーパーヘッドの相乗効果により、ティターンズ・アーキテクチャーは長いシーケンシャル・データをより効率的に処理し、短期記憶、長期記憶、永続記憶の利点を統合することができる。
• スケーラビリティ。
  • Transformerと比較して、Titansアーキテクチャは長いシーケンスを扱う際のスケーラビリティに優れ、より大きなコンテキスト・ウィンドウで高いパフォーマンスを維持することができる。

3人だ、 実験結果と分析：ティターンズ・アーキテクチャーの検証

著者らは、Titansアーキテクチャとその亜種の性能を評価するために、いくつかのタスクで大規模な実験を行った：

1. 言語モデリングと常識的推論。

• 実験セットアップ。
  • Transformer++、RetNet、GLA、Mamba、Mamba2、DeltaNet、TTT、Gated DeltaNetを含むいくつかのベースラインモデルと同様に、異なるサイズの3つのTitansモデル（340M、400M、760Mパラメータ）が使用された。
  • 学習データにはFineWeb-Eduデータセットを使用した。
• 主な調査結果
  • 非混合モデルにおいて、ニューラル長期記憶モジュールは、プレプレキシティと精度の両指標で最高の性能を達成した。
  • Titansの3つのバリエーション（MAC、MAG、MAL）はすべて、Samba（Mamba + attention）とGated DeltaNet-H2（Gated DeltaNet + attention）を上回った。
  • MACは長距離の依存関係を扱うのに優れており、MAGとMACの両方がMALの変種を上回る。
  

  図6：言語モデリングタスクとコモンセンス推論タスクにおける、Titansとベースラインモデルの性能比較。

2. "干し草の中の針 "ミッション。

• 実験セットアップ。
  • RULERベンチマークのSingle NIAH (S-NIAH)タスクを使用し、2K、4K、8K、16Kの長さのシーケンスでモデルの検索能力を評価した。
• 主な調査結果
  • 神経長期記憶モジュールは、3つのタスクすべてで最高の結果を出した。
  • ティターンズのバリアントも好調で、MACのバリアントが最も好調だった。

3. バビロン・ベンチマーキング

• 実験セットアップ。
  • このタスクは、非常に長い文書に含まれる分散した事実情報を推論するモデルを必要とする。
  • サンプルレス設定と微調整設定がある。
• 主な調査結果
  • サンプル数の少ない設定では、Titansは、GPT-4やGPT4o-miniのようなパラメータ数の多いモデルを含むすべてのベースラインを上回る。
  • ティターンズはまた、GPT-4のような超大型モデルでさえ、微調整設定においてすべてのモデルを凌駕している。
  • トランスフォーマー・ベースのメモリー・モデル（RMT）に比べ、ティターンズは主にその強力なメモリー能力により、より優れたパフォーマンスを発揮する。
  

  図7：BABILongベンチマークにおけるTitansとベースラインモデルの性能比較。

4. 時系列予測。

• 実験セットアップ。
  • シンバ・フレームワークを使い、マンバ・モジュールを神経長期記憶モジュールに置き換えた。
  • ETT、ECL、Traffic、Weatherのベンチマークデータセットで評価。
• 主な調査結果
  • ニューラル長期記憶モジュールは、マンバ、リニアモデル、トランスフォーマーベースのアーキテクチャを含むすべてのベースラインを凌駕している。

5. DNAモデリング。

• 実験セットアップ。
  • GenomicsBenchmarksで事前学習済みモデルのダウンストリームタスク性能を評価。
• 主な調査結果
  • Titans(LMM)は、様々なダウンストリームゲノミクスタスクにおいて競争力があり、最先端の手法と同等である。

6. 効率分析。

• 主な調査結果
  • Neural Long-Term Memoryモジュールは、他のリカレントモデルと比較して学習時間が若干遅いが、これは主に、より深いメモリとより複雑な変換プロセス、およびMamba2に実装された高度に最適化されたカーネルによるものである。
  • タイタンズ（MAL）は、主にFlashAttentionの高度に最適化されたカーネルを使用しているため、メモリモジュールと同様にベースラインよりも高速である。

7. アブレーション研究。

• 主な調査結果
  • ニューラル・メモリー設計のすべてのコンポーネントがパフォーマンスにプラスに寄与しており、ウェイト減衰、運動量、畳み込み、永続メモリーが最も貢献している。
  • アーキテクチャ設計も性能に大きな影響を与え、言語モデリングとコモンセンス推論タスクではMACとMAGがほぼ同等の性能を示すが、長いコンテキストのタスクではMACの方が優れている。

4人だ、 論文の革新的な点と長所

1. 新しい神経長期記憶モジュールを提案する。
   • 驚き、勢い、忘却といった人間の記憶メカニズムの重要な要素を利用し、より効率的な記憶の更新と保存を実現する。
   • ディープニューラルネットワークを記憶モジュールとして使用し、モデルの表現力を高めている。
2. ティターンズのアーキテクチャは、長期記憶と短期記憶を組み合わせるように設計されている。
   • 3つの異なる統合アプローチが提案されており、さまざまなアプリケーションシナリオに柔軟に対応できる。
   • コア、長期メモリ、永続メモリの3つのスーパーヘッドの相乗効果により、このモデルは長いシーケンシャルデータをより効率的に処理することができる。
3. 複数の仕事をこなす。
   • 言語モデリング、コモンセンス推論、時系列予測、DNAモデリングのいずれにおいても、Titansアーキテクチャは、既存のTransformerや線形ループモデルを凌駕する強力な性能を発揮する。
4. スケーラブル。
   • より大きなコンテクスト・ウィンドウで高いパフォーマンスを維持できることで、非常に長いシーケンスを処理できる可能性が広がる。

5人だ、 将来展望

ティターンズ・アーキテクチャーは、多くの面で素晴らしい結果をもたらしたが、さらなる探求に値する以下のような方向性もある。

1. より複雑なメモリ・モジュール・アーキテクチャの探求
   • 例えば、階層的なメモリー構造を導入したり、メモリーモジュールをグラフ・ニューラル・ネットワークのような他のモデルと組み合わせたりする。
2. より効率的なメモリ更新・保存機構の開発。
   • 例えば、スパース化技術や定量化技術は、メモリ消費量や計算コストを削減するために使用される。
3. ティターンズ・アーキテクチャーをより幅広い分野に応用する。
   • 例えば、映像理解、ロボット制御、推薦システムなどである。
4. より効果的なトレーニング戦略を探る。
   • 例えば、より高度な最適化アルゴリズムの導入や、メタ学習によるモデル学習の高速化などだ。
5. ティターンズ建築の解釈可能性を探る。
   • タイタンズが長期記憶情報をどのように保存し、活用しているのかをより深く理解することで、より強力なAIシステムを構築するための新たなアイデアが生まれるかもしれない。

6人だ、 概要

本稿の中心的な貢献は以下の通りである：

• 新しい神経長期記憶モジュールを提案設計は人間の記憶システムにヒントを得ており、勾配降下、運動量、ウェイト減衰といったディープラーニングの主要概念を取り入れている。
• 巨人建築を建設長期記憶と短期記憶を有機的に結合させ、3つの異なる統合方法を模索することで、さまざまな応用シナリオに柔軟な選択肢を提供する。
• ティターンズの優れた性能は、厳密な実験によって検証されている。これは、複数のタスク、特に長いコンテキストのタスクを扱う場合に優れた性能を発揮し、強力なスケーラビリティと高い精度を実証している。