Transformer に挑戦しようとする新しいアーキテクチャは数多くありますが、Google の「正統派」後継であるTitanアーキテクチャがより注目を集めています。

NVIDIA は、テスト時間コンピューティングを大規模モデルの 3 番目のスケーリング法則と呼んでいます。

OpenAI はこれを推論に使用し、Google は現在これを記憶に使用しています。

第一著者のアリ・ベロウズ氏は次のように述べています。

Titan は、Transformer や最新の線形 RNN よりも効率的で、200 万を超えるコンテキスト ウィンドウに効果的に拡張でき、GPT4 や Llama3 などの大規模モデルよりも優れたパフォーマンスを発揮します。

彼はまた、研究の背景にある動機についても説明した。研究チームは、トランスフォーマーの注意メカニズムは短期記憶として現れるため、ずっと昔の情報を記憶できる神経記憶モジュールも必要だと考えた。

新しい長期記憶モジュール

記憶に関して言えば、LSTM や Transformer などの古典的なモデルを思い浮かべる人もいるかもしれません。これらはさまざまな観点から人間の記憶をシミュレートしますが、それでも限界があります。

• 容量が制限された固定サイズの隠し状態にデータを圧縮します。
• 長距離の依存関係をキャプチャすることもできますが、計算コストは​​シーケンスの長さの2乗で増加します。

さらに、テスト データが分布外にある可能性があるため、トレーニング データを単に記憶するだけでは実際の使用には役立たない可能性があります。

この目的のために、Titans チームは過去の情報をニューラル ネットワークのパラメータにエンコードすることを計画し、テスト中に特定のデータを記憶/忘れる方法を学習するオンライン メタモデルをトレーニングしました。

彼らは神経心理学からインスピレーションを得て、人間の脳の原理を借用した神経長期記憶モジュールを設計しました。

• 予期しない出来事（つまり「サプライズ」）は記憶に残りやすくなります。
• 驚きのレベルは、メモリ モジュールによる入力の勾配によって測定されます。勾配が大きいほど、入力が予想外であることを示します。
• 勢いと忘却のメカニズムを導入することで、前者は短期的な驚きを蓄積して長期記憶を形成し、後者は不要になった古い記憶を消去してメモリのオーバーフローを防ぐことができます。
• メモリ モジュールは複数の MLP レイヤーで構成されており、深いデータ抽象化を保存でき、従来のマトリックス メモリよりも強力です。

このオンライン メタ学習パラダイムは、モデルが無駄なトレーニング データの詳細を記憶することを回避し、代わりに新しいデータに基づいて自身を調整する方法を学習することで、一般化能力を向上させます。

さらに、チームはこのモジュールが並列計算を実行できることを確認しました。

この強力なメモリ モジュールをディープラーニング アーキテクチャにどのように統合できるでしょうか?

この目的のために、Titans は次の 3 つのバリエーションを提案しました。

MAC、コンテキストとしてのメモリ

長期記憶と持続記憶（タスクの知識をエンコードする不変パラメータ）は、現在の入力のコンテキストとして一緒に注意に送られます。

MAG、記憶は門である

ゲーティング フュージョンは、メモリ モジュールとスライディング ウィンドウ アテンション ブランチの両方で実行されます。

MAL、層としての記憶

メモリ モジュールは別のレイヤーとして扱われ、履歴情報はアテンションに入力される前に圧縮されます。

実験の結果、それぞれの方法には長所と短所があることが明らかになりました。

Titans は、言語モデリング、常識的推論、時系列予測などのタスクにおいて、Transformer や Mamba などのさまざまなアーキテクチャの最先端 (SOTA) モデルよりも優れたパフォーマンスを発揮します。

さらに、長期記憶モジュール (LMM) のみを通じて、複数のタスクでベースラインを上回るパフォーマンスを実現します。

これは、このメカニズムが短期記憶（つまり注意）がなくても独立して学習する能力を持っていることを示しています。

長いテキストから細かい手がかりを見つける「干し草の山の中の針」テストでは、シーケンスの長さが 2k から 16k に増加しても、精度は約 90% を維持しました。

しかし、研究チームは、これらの一般的なテストでは、長文においてタイタンが優位に立つことはもはや示されていないと考えています。

非常に長い文書に分散された事実について推論する必要がある別のタスクでは、Titans は GPT4、Mamba、Llama3.1 + RAG システムよりも優れたパフォーマンスを発揮しました。

さらに、Titans は時系列予測や DNA 配列モデリングなどの特定のタスクでも優れたパフォーマンスを実現しています。

3 人の著者は、Google Research NYC アルゴリズムおよび最適化チームに所属していますが、このチームはまだ Google DeepMind に統合されていません。

第一著者はコーネル大学のインターンであるAli Behrouz 氏です。

鍾培霖は清華大学姚クラスの卒業生で、コロンビア大学で博士号を取得しました。2021年にGoogleのリサーチサイエンティストとして入社しました。

2016年、鍾培林が学部時代に筆頭著者として発表した論文がトップカンファレンスSTOC 2016に採択され、中国の学部生がSTOCで筆頭著者として論文を発表したのはこれが初めてとなった。

チームリーダーのVahab Mirrokni 氏は、 Google フェロー兼 VP です。

チームによれば、Titians は PyTorch と Jax で実装されており、モデルのトレーニングと評価用のコードを近日中に提供する予定だという。

論文の宛先:
https://arxiv.org/abs/2501.00...

参考リンク:
[1]https://x.com/behrouz\_ali/status/1878859086227255347