大型モデルの性能が125Mから1.3Bに向上しました。
新しいラージ言語モデル (LLM) アーキテクチャは、これまで AI 分野で普及してきた Transformer に代わるものとして期待されており、そのパフォーマンスは Mamba よりも優れています。月曜日、テストタイム トレーニング (TTT) に関する論文が人工知能コミュニティで話題になりました。 
論文リンク: https://arxiv.org/abs/2407.04620 この研究の著者は、スタンフォード大学、カリフォルニア大学バークレー校、カリフォルニア大学サンディエゴ校、およびメタの出身です。彼らは、RNN の隠れた状態を機械学習モデルに置き換えた新しいアーキテクチャ TTT を設計しました。このモデルは、入力トークンの実際の勾配降下を通じてコンテキストを圧縮します。 研究著者の一人であるKaran Dalal氏は、これが言語モデルのアプローチを根本的に変えることになると信じていると述べた。 機械学習モデルでは、TTT レイヤーが Attendance を直接置き換え、表現力豊かなメモリを通じて線形複雑性アーキテクチャのロックを解除し、コンテキスト内で数百万 (場合によっては数十億) のトークンを使用して LLM をトレーニングできるようにします。 著者は、パラメーター サイズが 125M から 1.3B までの大規模モデルで一連の比較を実施し、TTT-Linear と TTT-MLP の両方が、最も強力な Transformers および Mamba アーキテクチャ手法に匹敵するか、またはそれに勝つことができることを発見しました。 新しい情報圧縮およびモデル記憶メカニズムとして、TTT レイヤーは、Transformer のセルフアテンション レイヤーを簡単かつ直接置き換えることができます。 
Mamba と比較すると、TTT-Linear は、複雑さが低く、FLOP が少なく (左)、長いコンテキストをより有効に利用できます (右): 
これは、理論の複雑さ、および実際の実行において線形であるだけではありません。時間も早くなります。 
- 論文がオンラインになった後、著者は人々がトレーニングしてテストできるようにコードとjaxを公開しました: https://github.com/test-time-training/ttt-lm-jax
- PyTorch 推論コード: https://github.com/test-time-training/ttt-lm-pytorch
長いコンテキストの課題は、 RNN 層の性質: セルフ アテンション メカニズムとは異なり、RNN 層はコンテキストを固定サイズの隠れた状態に圧縮する必要があり、更新ルールは基礎となる構造と、数千または数百万のトークン間の関係を検出する必要があります。 研究チームは、自己教師あり学習が大規模なトレーニング セットを LLM などのモデルの重みに圧縮できること、および LLM モデルがトレーニング データ間の意味論的なつながりを深く理解していることが多いことを初めて観察しました。 この観察に触発されて、研究チームは、隠れ状態がモデルであり、更新ルールが自己教師あり学習のステップである新しいクラスのシーケンス モデリング レイヤーを設計しました。テスト シーケンスの隠れ状態を更新するプロセスは、テスト時にモデルをトレーニングすることと同等であるため、研究チームはこの新しい層をテスト時トレーニング (TTT) 層と呼んでいます。 
研究チームは、TTT-Linear と TTT-MLP という 2 つの単純な例を紹介します。ここで、隠れ状態はそれぞれ線形モデルと 2 層 MLP です。 TTT レイヤーは、RNN レイヤーやセルフ アテンションと同様に、あらゆるネットワーク アーキテクチャに統合でき、エンドツーエンドで最適化できます。 
TTT 層をより効率的にするために、この研究では TTT 層を改善するためのいくつかのトリックを採用しました: まず、より良い並列性を得るために通常のトレーニング中にミニバッチ シーケンスに勾配ステップを行うのと同様に、この研究ではTTT 中には少量のトークンを使用します。 
次に、この研究では、最新の GPU と TPU をより有効に活用するために、各 TTT ミニバッチ内の操作の二重形式を開発します。双対形式の出力は単純な実装と同等ですが、トレーニングは 5 倍以上高速です。図 3 に示すように、TTT-Linear は Transformer よりも高速で、8k のコンテキストでは Mamba に匹敵します。 研究チームは、図 4 に示すように、すべてのシーケンス モデリング レイヤーは歴史的コンテキストを非表示の状態に保存していると見なすことができると考えています。 
たとえば、LSTM、RWKV、Mamba レイヤーなどの RNN レイヤーは、時間の経過とともにコンテキストを固定サイズの状態に圧縮します。この圧縮には 2 つの結果があります。1 つは、各トークンの更新ルールと出力ルールに一定の時間が必要なため、入力トークン x_t を出力トークン z_t にマッピングすることは効率的です。一方、長いコンテキストにおける RNN 層のパフォーマンスは、その隠れ状態 s_t の表現力によって制限されます。 セルフアテンションは、その隠れ状態 (多くの場合、Key-Value キャッシュと呼ばれます) が t とともに線形に増加するリストであることを除いて、上記の観点から見ることもできます。その更新ルールは単に現在の KV タプルをこのリストに追加するだけですが、その出力ルールは t より前のすべてのタプルをスキャンしてアテンション マトリックスを形成します。隠蔽状態では、すべての履歴コンテキストが圧縮なしで明示的に保存されるため、長いコンテキストに対するセルフアテンションが RNN 層よりも表現力豊かになります。ただし、この直線的に増加する隠れ状態をスキャンするのに必要な時間も直線的に増加します。長いコンテキストを効率的かつ表現力豊かに保つために、研究者はより優れた圧縮ヒューリスティックを必要としています。具体的には、数千、場合によっては数百万のトークンを、その基礎となる構造と関係を効果的にキャプチャする隠された状態に圧縮する必要があります。難しく聞こえるかもしれませんが、実際には多くの人がこのヒューリスティックに精通しています。 バックボーンアーキテクチャ。 RNN レイヤーをより大きなアーキテクチャに統合する最もクリーンな方法は、Transformer のセルフアテンション (ここではバックボーンと呼ばれます) を直接置き換えることです。ただし、既存の RNN (Mamba や Griffin など) は、Transformer とは異なるバックボーン層を使用します。最も注目すべき点は、バックボーン層には RNN 層の前に時間的畳み込みが含まれており、これは時間の経過に伴うローカル情報の収集に役立つ可能性があります。研究者らは、Mamba バックボーンを実験した結果、TTT 層の複雑性も改善できることが判明したため、図 16 に示すように、提案手法に組み込まれました。 
実験では、研究者らはTTT-LinearとTTT-MLPを、2つのベースラインであるTransformerとMambaと比較しました。 図 11 から、次の結論を引き出すことができます:
- 2k コンテキスト、TTT-Linear (M)、Mamba、および Transformer のパフォーマンスは同等であるため、行のほとんどが重なっています。 TTT-MLP (M) は、FLOP バジェットが大きいほどパフォーマンスがわずかに低下します。 TTT-MLP はさまざまなモデル サイズで TTT-Linear よりも複雑性が優れていますが、FLOP の追加コストがこの利点を相殺します。
- 8k コンテキストでは、TTT-Linear (M) と TTT-MLP (M) の両方が Mamba よりも大幅に優れたパフォーマンスを示しますが、これは 2k コンテキストでの観察とはまったく異なります。 Transformer バックボーン ネットワークを使用する TTT-MLP (T) でさえ、約 1.3B で Mamba よりわずかに優れています。重要な現象は、コンテキストの長さが増加するにつれて、Mamba 層に対する TTT 層の利点も拡大することです。
- コンテキストの長さが 8k に達しても、Transformer は各モデル サイズの下でパープレキシティにおいて依然として良好なパフォーマンスを発揮しますが、FLOP のコストにより競争力がなくなりました。
上記の結果は、TTT レイヤーを Mamba バックボーン ネットワークから Transformer バックボーン ネットワークに切り替えた場合の影響を示しています。研究者らは、シーケンス モデリング層の隠れ状態の表現力が低い場合、Mamba バックボーン ネットワークの時間畳み込みがより役立つと仮説を立てました。線形モデルは MLP よりも表現力が低いため、畳み込みのメリットが大きくなります。 長いコンテキストの能力を評価するために、Pile の人気のあるサブセットである Books3 を使用して、コンテキストの長さを 1k から 32k まで 2 倍の増分で実験しました。ここでのトレーニング方法は Pile と同じで、TTT 層のすべての実験は 1 回のトレーニング実行で実行されます。図 12 の結果のサブセットから、彼らは以下の観察を行いました: 
Books 2k のコンテキストでは、Mamba のパフォーマンスが TTT-Linear よりわずかに優れていることを除いて、Pile 2k の観察はすべて依然として当てはまります ( Pile 2k ではラインがほぼ重なっています)。 32k のコンテキストでは、Pile 8k の観察と同様に、TTT-Linear (M) と TTT-MLP (M) の両方が Mamba よりも優れたパフォーマンスを示します。 Transformer バックボーンを備えた TTT-MLP (T) でさえ、32k コンテキストでは Mamba よりもわずかに優れたパフォーマンスを発揮します。 TTT-MLP (T) は、1.3B スケールでは TTT-MLP (M) よりわずかに悪いだけです。上で述べたように、明確な線形近似が欠如しているため、経験に基づくスケーリング則を導き出すことは困難です。ただし、TTT-MLP (T) の強い傾向は、Transformer バックボーンが、評価の範囲を超えて、より大規模なモデルやより長いコンテキストに適している可能性があることを示唆しています。 LLM のトレーニングと推論は、前方、後方、生成に分解できます。推論中のキュー ワード処理 (事前入力とも呼ばれる) は、逆方向演算が中間アクティベーション値の格納を必要としないことを除いて、トレーニング中の順方向演算と同じです。 前方（訓練中および推論中）と後方の両方を並列処理できるため、ここでは双対形式が使用されます。新しいトークンの生成 (デコードとも呼ばれます) は本質的に逐次的なものであるため、ここでは生の形式が使用されます。 研究者は、リソースの制限により、この記事の実験は JAX で書かれ、TPU で実行されたと述べました。 v5e-256 TPU ポッドでは、2k コンテキストでのトレーニングに Transformer ベースラインは反復あたり 0.30 秒かかりますが、TTT-Linear は反復あたり 0.27 秒かかり、システムを最適化しない場合は 10% 高速になります。 Mamba (PyTorch、Triton、CUDA で実装) は GPU 上でのみ実行できるため、公平な比較を行うために、研究者らは GPU 上で実行できるようにこのメソッドの予備的なシステム最適化を実施しました。 図 15 の左側は、バッチ サイズ 16 での各モデルのフォワード カーネルのレイテンシーを示しています。全モデル1.3B（マンバは1.4B）。ここでは HuggingFace Transformer の代わりに vLLM が使用されているため、ここでの Transformer ベースラインは Mamba 論文のものよりもはるかに高速であることに注意してください。 
さらに、研究者らは生成用に別の GPU カーネルを作成し、図 15 の右側にあるバッチ サイズ 512 でその速度をベンチマークしました。一般的に使用されるもう 1 つの実時間メトリクスはスループットです。これは、より大きなバッチ サイズを使用することによる潜在的な利点を考慮しています。スループットについては、上記のすべての観察とメソッド間の順序付けが依然として当てはまります。 TTT研究が提出された後、論文著者の1人であるUCSD助教授Xiaolong Wang氏がツイートで祝福の意を表した。 TTTの研究には1年半かかったそうですが、テストタイムトレーニング(TTT)のアイデアが生まれてから実は5年も経っています。当初のアイデアと現在の結果はまったく異なりますが。 
TTT 論文の主な著者 3 人は、それぞれスタンフォード大学、UC バークレー校、UCSD の出身です。 その中で、Yu Sun はスタンフォード大学の博士研究員であり、カリフォルニア大学バークレー校 EECS で博士号を取得し、長期的な研究の方向性を目指しています。 
Xinhao Li は UCSD の博士課程候補者で、中国電子科学技術大学を卒業しました。 
Karan Dalal は、カリフォルニア大学バークレー校の博士課程候補者で、高校在学中に Otto という獣医遠隔医療スタートアップを共同設立しました。 
上記の 3 人は全員、研究の方向性を紹介する個人 Web サイトの最初の行に、テスト時のトレーニングについて書いています。 研究の詳細については、元の論文を参照してください。 以上が言語モデルを完全に変更します。新しいアーキテクチャ TTT は Transformer を上回り、ML モデルは RNN 隠れ状態を置き換えます。の詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。
