Googleの新TITANS: RNNメモリを持つTransformer

8,465 文字

こんにちは、コミュニティの皆さん。ついに私たちは新しい形のトランスフォーマーアーキテクチャを手に入れました。それらはTitansと呼ばれ、特殊なアーキテクチャを持っています。
まずはこの素晴らしい論文から始めましょう。2週間前に発表された「Titans: テスト時の記憶学習」です。ここでは、テスト時学習、テスト時の計算、テスト時の適応、推論・推論などが全て一つにまとまっており、リカレントニューラルネットワークと大規模言語モデルの完璧な組み合わせを見出すことができます。
ここでは、言語モデルに統合できる新しいニューラル長期メモリモジュールを探しています。そしてこれがGoogleリサーチによってTitansと名付けられた新しいモデルアーキテクチャです。
ここでJaxでの言語モデル用テスト時学習（TTT）の実装に直接飛び込んでみましょう。これがGoogleによるオリジナルのコードです。TTTベースモデルがありますが、ここでは線形モデルではなくMLPベースを使用します。TTT MLPベースのクラス定義があります。
コードに興味がある方もいるでしょうが、まずは背景と実際に何をしているのかを説明させてください。これは素晴らしく見えますが、一体何をしているのでしょうか？そのアイデアは何なのでしょうか？
もちろん、RNNとは何かについて振り返ってみましょう。これらは私たちの古くからの友人です。隠れ層があり、この隠れ層の目的は連続データのためのメモリでした。2019年頃、これらのRNNの隠れ層は複数のタイムステップにわたって情報を伝達するように設計されていました。
LSTMを覚えていますか？モデルがテキスト、音声、時系列などの連続データを処理できるようにし、これらの隠れ層は将来の予測に影響を与える過去のコンテキストを保持する新しいメモリメカニズムとして機能しました。隠れ層の数学については省略しますが、素晴らしかったのは、シーケンスの長さに関係なく固定サイズだったことです。
これによりメモリと計算の面で非常に効率的でした。なぜなら、メモリの割り当て方法を正確に知ることができたからです。その主な役割は圧縮でした。隠れ状態は、それまでのコンテキストや意味情報を単一のベクトルに圧縮していました。これが非常に効果的だった理由ですが、もちろん大きな問題もありました。
特に長いシーケンスでは、この圧縮により初期の入力が忘れられてしまう可能性がありました。2019年当時のことですが、隠れ層がなければモデルは本質的に線形になってしまい、複雑な関係性を学習することが完全にできなくなってしまいます。そこで隠れ層のReLU活性化関数が非線形性を導入し、RNNモデルが複雑なパターンを学習できるようになりました。
私たちの2019年からの古い友人である隠れ層には、勾配の消失と爆発、メモリの制限、計算の逐次依存性といった問題がありました。しかし、「待って、なぜAIの石器時代の話をしているんだ？新しいTitansでの長期メモリ圧縮のためのTTT MLPの応用について理解したいだけなのに」と思われるかもしれません。
なんという偶然でしょう。以前、私たちはOpenAI o3モデルの推論について議論していました。推論時の推論について話していて、今も同じことを話しています。モデルが追加のメモリモジュールを持つ場合の推論時の推論について話しています。教師あり微調整や強化学習への厳密な依存性を覚えていますか？でも今はそれは置いておきましょう。
簡単な方法を取って、テスト時の計算、推論時の推論、そしてメモリモジュールがどのように役立つかを見てみましょう。ただし、このテスト時学習は私のチャンネルでのコメントで示唆されているような新しいものではありません。これは2019年、2020年に既にバージョン3の論文が出ていました。
この論文では、訓練データとテストデータが異なる分布から来る場合に予測モデルの性能を向上させる一般的なアプローチとしてテスト時学習を提案しています。つまり、TTTは訓練データとテストデータの分布とデータのシフトのために5〜6年前に発明されたのです。私たちの推論データは、訓練データとは異なるデータ分布を持っていたため、テスト時学習が発明されました。
では5年後、4年後に飛んでみましょう。スタンフォード大学、UCサンディエゴ、UCバークレー、そしてMetaによる「テスト時にRNNを表現力豊かな隠れ状態で学習する」という新しいアイデアです。2024年8月、自己注意機構であるTransformerは長いコンテキストでうまく機能しますが、依然として二次の計算量を持っています。
2019年の古き良き時代のRNN層を思い出してください。線形の計算量しかありませんでしたが、他の問題がありました。そこでスタンフォード大学から、RNNの隠れ状態自体を機械学習モデルにするという重要なアイデアが提案されました。更新則は自己教師あり損失関数Lを使用した自己教師あり学習のステップです。
「何を言っているんだ？」と思われるかもしれませんが、これは実際にとても強力だとわかりました。静的だった隠れ状態が今や機械学習モデル自体となり、自己教師あり損失関数を持っています。ここにそれがあります。入力トークンと出力トークンがあり、その間にRNNの隠れ状態があります。そして自己教師あり損失関数Lを参照する更新則があります。
時間依存性が明示的にあり、これを計算することができ、隠れ状態に機械学習モデル自体があります。これが一晩で起こったとは思わないでください。4年かかる長い研究プロセスでした。しかし2024年8月には、隠れ状態がテストシーケンスでも学習によって更新されるため、これらの層はTTT層と呼ばれるようになりました。
つまり、事前学習、教師あり微調整、強化学習があり、その後に線があって、その後に推論実行とテスト時があるというわけではありません。これは全て時間軸上で多かれ少なかれ同時に行われます。そして、TTTリニアとTTTMLPという2つの実装を考えています。もちろん、複雑な情報の記憶、推論、格納を可能にする非線形性関数を持つために、2層以上の多層パーセプトロンを使用します。
2019年の素朴なRNNの美しい概要があります。初期状態、更新則、出力則、固定コストを覚えていますか？その後の自己注意機構は全く固定ではありませんでした。2020年の素朴なTTTも、多かれ少なかれ固定でした。固定サイズで計算コストを抑えられるのは美しかったのですが、特定の複雑な問題には固定長では不十分かもしれないという問題に直面しました。
次に取るべきステップがすぐにわかりますね。ちなみにテスト時学習のコードを見たい方のために、2020年のものが6ヶ月前に更新されたものがあります。GoogleがTPUマシンの並列化のために数年前からJaxで実装しているテスト時学習用の言語モデルのGitHubコードがあります。完全なコードがここにあります。
さて、Titansに戻りましょう。2024年8月の状況を理解したので、2024年12月末のこのステップを理解することができます。MLPをメモリモジュールとして使用することを決定しました。このメモリモデルは追加のメモリモデルであり、大規模言語モデルの注意機構を複製しようとするものではありません。
代わりに、この論文で短期メモリと呼ばれる注意機構を完全に補完し、生のトークンを保存することを避ける方法で履歴情報を重みにエンコードし圧縮します。なぜこれが必要なのかと疑問に思われるかもしれません。2百万トークンのコンテキストウィンドウを持つ注意機構があるのに、なぜ4百万トークンのコンテキスト長に行きたいのでしょうか？
最も単純な構成で、2百万トークンの自己注意機構には含まれない履歴メモリを保存する追加のメモリモジュールを何らかの方法で計算する必要があります。この2百万トークンより前の履歴知識を持ちたい場合、トランスフォーマーアーキテクチャ内で即座にアクセスできるように、この重要な情報を全て保存するメモリモジュールが必要です。
4百万以上のトークンが新しい目標となっているわけです。論文を見ることができます。動的な重み更新はこの重みで実行され、サプライズメトリックがあり、サプライズ駆動の更新メカニズムがあります。これについては後で説明します。メモリ損失関数と、ReLUについての私の2つの動画で説明した深い非線形表現があります。
これは大体標準的な毎日の手順です。素晴らしいのは、TitanがRNNモデルから知っていることをTransformerモデルと組み合わせたことです。2種類の短期メモリがあると考えることができます。短期メモリは、私たちの古典的なTransformerの自己注意機構で、固定のコンテキストウィンドウ内の正確なトークンレベルの依存関係を処理します。
Anthropicなら約128kトークン長、Gemini 2.0なら2百万のコンテキスト長があり、即時のコンテキストが重要なタスクに理想的です。しかし、1千万トークンがある場合はどうでしょうか？そのため、長期メモリモジュールが必要になります。これは履歴コンテキストをエンコードし、注意ウィンドウより長いシーケンスのパターンを抽象化します。
LLMが長い入力を処理する際に情報を失わないようにする永続的なメモリメカニズムを提供します。単純なアイデアとして、ウィンドウに収まらない情報を全て長期メモリに入れ、この長期メモリの開発は5年前に開発された線形計算量のRNN MLPのアイデアから来ていることを理解しました。
人々が理解するのに苦労していることがあります。それは文脈内学習メカニズムと呼ばれています。論文にはありますが、これは正しくありません。私は物事を単純化して説明するのが好きですが、これは単純化が強すぎるので、なぜこれが正しくないのか説明したいと思います。
文脈内学習について考えてみてください。私のチャンネルには文脈内学習に関する2つの動画があります。これらが文脈内学習の特徴です。そして今、TitanにあるRNNの意味から学んだMLPベースの長期メモリを見てください。目的は過去のデータを利用して予測を適応し改善することです。歴史的なデータを保持し活用する、類似のコンテキスト処理、類似の動的適応、類似のタスク関連性があります。
しかし、重要な違いがあります。文脈内学習は入力トークンへの注意を通じて機能し、重みの更新はありません。一方、MLPベースの長期メモリでは、学習プロセスとテスト時学習プロセスがあるため、勾配メカニズムを通じてメモリの重みが更新されます。
範囲は、文脈内学習では提供された入力コンテキストに一時的に特化していますが、これは本当に永続的で、入力ストリーム全体から一般化したパターンを保存します。新しい未見のデータへのパラメータ更新は、文脈内学習ではありませんが、ここではテスト時にメモリの重みが更新されます。
計算の依存性は、文脈内学習ではモデルの注意ウィンドウに制限されていますが、ここでは線形計算量を持つため、非常に長いシーケンスまで効率的にスケールします。これは同じものではないことがわかります。
Titansには、テスト時計算での動的更新を持つMLPの深いメモリ層があります。MLPは本当に適応し、入力テスト時データに基づいて重みのテンソルを変更し、シーケンスが成長するにつれてメモリが進化することを保証します。学習フェーズだけでなく、テスト時推論フェーズでも更新があります。
RoPEについては2つの素晴らしい動画があります。回転位置エンコーディングと埋め込みについて学びたい方はそちらの動画をご覧ください。しかし、テスト時学習の適応について、「3週間前に動画があったよね」と思われるかもしれません。その通りです。その時、AIコミュニティの大多数はまだ注目していませんでしたが、3週間前に新しい報酬関数とMonocolotリサーチによるテスト時学習適応の兆候を既にお見せしました。
また、MITからのより複雑な、スケール推論ベースの推論によるテスト時最適化AIについての動画もありました。もしそれらの動画に戻って見たいと思われるなら、より理解が容易になるでしょう。なぜなら、今私たちは良いレベルの一般的な知識を持っているからです。
また、新しいモデルの推論教師あり微調整と推論強化学習の利点を本当に活用したい場合、古典的な教師あり微調整と強化学習の方法論を修正する必要があることを示した私の動画も覚えておいてください。
サプライズ効果について、これはとても可愛らしいものです。先ほど説明したように、MLPはサプライズメトリックに基づいて重みを動的に更新します。これは理解するのがとても簡単です。単純な数式がありますが、物語を使って説明すると、すぐに理解できるでしょう。
効果として、驚くべきデータ、つまり新しいデータが保存のために優先されます。LLMにとって新しいものが識別され、追加の重要な知識として保存されます。メアリーとジョンの短い物語を考えてみてください。メアリーはジョンの目を深く見つめ、ジョンはメアリーの目を深く見つめます。ゆっくりとメアリーの唇とジョンの唇が近づいていき、突然メアリーが言います。「ねえ、このYouTubeチャンネルには本当に面白い情報があるって知ってた？」
これは物語のサプライズ効果です。なぜなら、誰も物語がこのような展開を見せるとは予想していなかったからです。したがって、この時点でサプライズが発生し、物語の新しい結末がTitansのための多層パーセプトロンベースの適応メモリに実装されます。
Titansはかなり簡単です。私たちが既に知っているものがより大きくなり、2百万以上のコンテキスト、4百万以上のトークンを扱えるようになりました。必要な場合はこれを使用できます。しかし、小さなサプライズがあります。
なぜなら、私たちには基本的に3種類のメモリがあるからです。考えてみてください。短期メモリは、トランスフォーマーアーキテクチャからの古典的な注意機構とマルチヘッド自己注意機構です。長期メモリについては、MLPと全てのことについて完全に議論しました。そして、テスト時学習でも起こる適応メモリステップがあります。
しかし今、永続メモリも持つという素晴らしいアイデアがあります。この永続メモリは今や重要な役割を果たしています。基本的なタスクのための新しい基盤を提供します。例を見てみましょう。
ミステリー小説を読んでいると想像してください。短期メモリ、つまり自己注意機構があり、これは現在の章、現在の段落での即時の対話やアクションを追跡するのに役立ちます。これは段落の注意機構と言えるでしょう。その特定の瞬間に何が起こっているのかを正確に知っており、コンテキストウィンドウは128kトークンです。
このミステリー小説は約450ページありますが、長期メモリには、小説全体からの手がかりや登場人物、出来事に関する全ての情報が含まれています。つまり、謎を解くかもしれない早い章からの全ての情報を持っています。
さらに、永続メモリもあります。これは、いわば非常に一般的な、ジョブ固有ではないものです。これは「これはミステリー小説だ。私は既に他の10のミステリー小説を読んでいて、それらには共通のパターンがある」と教えてくれます。それらは恋愛物語から始まり、何かが起こり、幽霊のいる暗い城に行き、その幽霊は実は友好的な幽霊で、最後には完璧な幸せが訪れる、というパターンです。
これが私たちが期待できる一般的なパターンであり、これは他の全ての小説を読んだ経験から得られた情報として永続メモリとして提示されます。見てください、とても簡単です。永続メモリについて考えてみてください。固定の学習可能なパラメータが入力シーケンスに追加され、LLMが常にタスク関連の知識を持っていることを保証します。
物理学について話す場合を考えてみてください。理論物理学でタスクがある場合、理論物理学の主要な方程式を持つ永続メモリがあります。マクスウェル方程式からアインシュタイン方程式、QCD方程式まで、全ての特定のジョブ用の物理学固有の一般的な解決策があります。永続メモリにはそれらの解決策ツールボックスがあるので、新しい解決策を一から考え出す必要はありません。
3つのメモリ部分があれば、ゲーティングメカニズムを構築するのは簡単です。LSTMと古き良き時代を思い出してください。ゲートは簡単です。ゲートは特定の、例えば長期メモリからどれだけ広いコンテキストを使用するかを決定します。
ゲートは「注意の80%、長期の10%、永続の10%を使用する」と言うかもしれませんし、「ジョブ固有の解決策の情報が本当にたくさん必要だ」と言って、永続メモリに40%、長期メモリに10%を割り当て、残りを自己注意機構のマルチヘッドに割り当てるかもしれません。
ゲーティングメカニズムは特別なものではありません。これは既に知っていました。記憶として、このTitanには2種類のメモリが含まれています。RNNから知っているテスト時学習のLLM多層パーセプトロンを持つ長期メモリと、今話した永続メモリです。
長期メモリは動的で、永続メモリは静的です。理論物理学の全ての数式は最初に事前定義でき、複数のジョブで同じままです。更新は、長期メモリでは入力と勾配に基づいて継続的に更新され、MLPと損失関数で示したように勾配があります。永続メモリは推論実行中には更新されません。
目的は、長期メモリでは履歴入力依存の情報、永続メモリでは数学的な公式などの抽象的なタスク固有の知識をエンコードします。長期メモリは入力に依存し、永続メモリは入力に依存しません。理論物理学の一般的な公式はそこにあり、どのようなタスクが与えられても有効です。
Titansでは、これら2種類のメモリが補完的な役割を果たしており、LLMエージェントは知識ドメイン固有で履歴に特化したこれら2種類のメモリにアクセスできます。
医学の本を読んでいて、長期メモリに第1章から第12章までが既にエンコードされているとします。そして今、LLMの自己注意機構で第13章を読んでいます。LLMは第1章から第12章の長期メモリに即座にアクセスでき、知識と関係性は全く失われていません。
実装では、永続メモリ、長期メモリ、短期メモリの全てのメモリ要素に一度にアクセスできます。または、先ほど説明したように、ゲートが長期の10%、これの25%など、長期、短期、永続メモリの正しい量を決定します。
そして、この長期メモリをLLMの層として扱うことができ、統合が簡単になります。実装が本当に簡単です。どのような形式でも、LLMやTransformer LLMのための新しいTitanアーキテクチャを実装するためにこれら3つのアーキテクチャのいずれかを選択できます。
将来の方向性として、これを本当に実装したい場合、GoogleはTitanのコードをまもなく公開すると言っています。Jax Flagsになりますが、この動画を録画している現時点では一般に公開されていません。
代替のメモリ表現として、グラフベースの構造をすぐに考えることができます。これは非常に興味深いでしょう。複雑になり、ハードウェア構成に基づいて小さなメモリを持つ小規模言語モデルがある場合、特定のゲーティングメカニズムと忘却メカニズムの適応性を強化できます。
そして、もちろんテスト時にマルチモーダルデータを持つことができ、Titanアーキテクチャのメモリ効果を最適化するために特に適した新しい強化学習環境を持つことができます。
以上です。美しいTitan、トランスフォーマーアーキテクチャ、長期メモリ、そしてRNNとそのMLPメモリ圧縮技術から得た複数年の知識、そしてドメイン固有とタスク固有の背景情報、背景知識を提供する静的メモリとしての永続メモリという美しいアイデアを紹介しました。
楽しんでいただけたでしょうか。簡単な概要を提供できたと思います。いくつかのアイデアを提供できたと思います。Googleの完全な論文を読むことをお勧めします。とても興味深く、この15-20分の動画では紹介できない多くの追加情報があります。
これを適用したい場合は、Googleがコードを公開するか定期的にチェックしてください。オープンソースライセンスになることを願っています。その場合、次の動画でお知らせします。