Published on

グーグルTitanアーキテクチャがTransformerのメモリボトルネックを打破

著者
  • avatar
    名前
    Ajax
    Twitter

Googleの新アーキテクチャ Titanの紹介

テクノロジー業界では、Googleから登場した新しいアーキテクチャであるTitanが話題になっています。これは、特にメモリの処理方法において、Transformerモデルの限界に挑戦するように設計されています。この新しいアーキテクチャは、特にGoogle内のチームによって開発されたことを考えると、Transformerの後継となる可能性として大きな注目を集めています。

既存モデルにおけるメモリの課題

LSTMやTransformerのような従来のモデルは革新的ではありますが、人間のような記憶をシミュレートする際に課題に直面しています。これらの課題には以下のようなものがあります。

  • 容量の制限: データは固定サイズの隠れ状態に圧縮されることが多く、保持できる情報量が制限されます。
  • 計算オーバーヘッド: 長距離の依存関係を捉えることは可能ですが、計算コストはシーケンス長に対して二次関数的に増加するため、非常に長いシーケンスには非効率的です。
  • 訓練データへの過度の依存: 訓練データを単に記憶するだけでは、テストデータが訓練分布外になる可能性がある実世界のアプリケーションでは必ずしも役立ちません。

Titanのアプローチ：神経に着想を得たメモリモジュール

Titanチームは、ニューラルネットワークのパラメータに情報をエンコードするという異なるアプローチを採用しました。彼らは、テスト中に特定のデータをどのように記憶し、忘れるかを学習するように設計されたオンラインメタモデルを開発しました。このモデルは、神経心理学の原則に触発されており、次の主要な要素を取り入れています。

  • トリガーとしての驚き: 予期しないイベントはより簡単に記憶されます。「驚き」は、メモリモジュールへの入力の勾配によって測定されます。勾配が大きいほど、入力はより予期しないものになります。
  • モメンタムと忘却メカニズム: モメンタムメカニズムは、短期的な驚きを長期記憶に蓄積し、忘却メカニズムは古い記憶を消去して、メモリのオーバーフローを防ぎます。
  • 多層パーセプトロン(MLP)ベースのメモリ: メモリモジュールは複数のMLPレイヤーで構成されており、データの深い抽象化を保存できるため、従来の行列ベースのメモリよりも強力です。

このオンラインメタ学習アプローチは、モデルが訓練データを単に記憶するのではなく、新しいデータに適応する方法を学習することに焦点を当てるのに役立ちます。このモジュールは並列計算用に設計されており、効率が向上します。

ディープラーニングアーキテクチャへのメモリモジュールの統合

Titanの研究チームは、メモリモジュールをディープラーニングアーキテクチャに組み込むための3つのバリエーションを提案しました。

  1. MAC (Memory as Context): この方法は、タスク知識をエンコードする長期および永続的なメモリを、注意メカニズムへの入力としてコンテキストと組み合わせます。
  2. MAG (Memory as Gate): このアプローチでは、2つのブランチにわたるスライディングウィンドウ注意メカニズムとメモリモジュールのゲート付き融合を使用します。
  3. MAL (Memory as Layer): ここでは、メモリモジュールは、注意メカニズムに供給する前に履歴情報を圧縮する独立したレイヤーとして実装されます。

チームは、各バリエーションに長所と短所があることを発見しました。

Titanのパフォーマンスと利点

Titanは、言語モデリング、常識推論、時系列予測など、さまざまなタスクで優れたパフォーマンスを発揮しています。TransformerやMambaのような最先端モデルを凌駕しています。特に、長期メモリモジュール(LMM)単独で、いくつかのタスクでベースラインモデルよりも優れた性能を発揮しており、短期メモリ(注意)なしで独立した学習能力を発揮しています。

長いテキストで微細な手がかりを見つけるように設計された「干し草の中の針」テストでは、Titanはシーケンス長が2kから16kに増加しても約90%の精度を維持しました。チームは、標準テストでは、長いテキストを処理する際のTitanの利点が十分に示されていないことを示しています。Titanは、非常に長いドキュメントにまたがる事実からの推論を必要とするタスクにおいて、GPT4、Mamba、さらにはRAGを備えたLlama3.1などのモデルよりも優れた性能を発揮しました。

Titanは、時系列予測やDNA配列モデリングなどの特定の分野でも目覚ましいパフォーマンスを示しています。

Titanを開発したチーム

この研究は、Google DeepMindには現在所属していない、Google Research NYCのアルゴリズムおよび最適化グループのチームによって実施されました。

  • コーネル大学のインターンであるAli Behrouzが論文の筆頭著者です。
  • 清華大学の卒業生であり、コロンビア大学で博士号を取得したZhong Peilinは、2021年からGoogleの研究科学者です。彼は学部生としてSTOC 2016で筆頭著者論文を発表したことで知られています。
  • Googleフェロー兼VPであるVahab Mirrokniがチームを率いています。

チームはPytorchとJaxを使用してTitanを開発し、トレーニングと評価のためのコードを近日中にリリースする予定です。