グラフトによる拡散変圧器設計の探求

作業はスタンフォード大学、Together AI、カリフォルニア大学サンディエゴ校、ノースウェスタン大学、Google DeepMind、セールスフォース・リサーチと共同で行いました。

モデルアーキテクチャの設計は、データ、アルゴリズム、コンピューティング、ベンチマークと並んで、最新の AI システム構築の中核をなす部分です。モデルアーキテクチャは学習可能な機能を定義するもので、どの演算子 (注意、畳み込みなど) を使用するか、どのように構成するか (モデルの深さ、幅など) などの主要な選択肢が必要です。その重要な役割にもかかわらず、特に今日の基盤モデルの時代では、モデルをゼロからトレーニングすることには莫大なコストがかかるため、アーキテクチャー (何が機能し、何が機能しないか) についての洞察を得ることは困難です。そのため、特にジェネレーティブ・モデルでは、新しいアーキテクチャーの探索が依然として大きな課題となっています。

新しいソフトウェアがゼロから作成されるのではなく、既存のコードに基づいて構築されるのと同じように、 事前学習済みのモデルは、新しい建築設計を検討するための足場として役立ちますか？ 調査します 建築編集 事前トレーニング済みモデルの数。私たちは、画像やビデオの生成に広く使用されている生成トランスの一種である拡散変圧器（DIT）に焦点を当てています¹ ²³。

事前トレーニング済みのモデルは、画像や動画の生成などのタスクを実行するための計算グラフを実装します。電球を点けるように配線された電気回路のようなものだと考えてください。新しいアーキテクチャのアイデアと事前にトレーニングされたモデルをもとに、少ない計算予算で計算グラフを変更することでそのアイデアを実現できるかどうかを調べます。たとえば、DiT では、畳み込み設計がマルチヘッドアテンション (MHA) やマルチレイヤーパーセプトロン (MLP) に取って代わる可能性があるという仮説が立てられるかもしれません。この考え方を検証する簡単な方法は、モデルの品質を維持しながら MHA または MLP 演算子を畳み込み演算子と入れ替えることです。これにより、次の 2 つの重要な疑問が浮かび上がります。

• (Q1) オペレータの初期化: 新しい演算子を計算グラフに統合する前に、どのように初期化すればよいですか?
• (Q2) エラー累積: 複数の演算子が置き換えられたときのエラーの伝播をどのように軽減できますか？

これらの疑問を解決するために、アーキテクチャ編集へのシンプルな二段階アプローチであるグラフティングを紹介します。グラフトの仕組みは以下のとおりです。

• (i) 活性化蒸留：この段階では、回帰目標を使用して活性化を蒸留することにより、元のオペレーターの機能を新しいオペレーターに移します。
• (ii) 軽量な微調整:この段階では、限られたデータを使用して微調整を行うことで、複数の新しい演算子を統合することによるエラーの伝播を軽減します。

私たちは、ますます困難になる一連の生成モデリングタスクにわたってグラフトをテストします。

結果I: クラス条件付き画像生成のためのハイブリッドアーキテクチャ

まず、256×256の解像度でDIT-XL/2を使用してクラス条件付き画像生成でグラフトを検証します。このセットアップでは、ソフトマックスアテンション (MHA) を、ローカルゲート畳み込み (Hyena-SE と当社が提案した Hyena-X/Y)、ローカルアテンション (スライディングウィンドウ)、線形アテンション (Mamba-2) などの代替手法に置き換えます。MLP では、畳み込みバージョン (Hyena-X) だけでなく、膨張比が異なる (3x、6x など) バリアントをテストします。

興味深いことに、 いくつかのインターリーブハイブリッド設計が良好な発電品質を実現、FIDスコアは2.38から2.64の間です（低いほど良い。DIT-XL/2のベースライン：2.27）。グラフトは単純かつ軽量です。各実験は、元のモデルの事前トレーニング計算の 2% 未満しか使用せず、8 × H100 GPU で 24 時間以内に完了します。

グラフトモデルを使用して生成された高品質のサンプルは次のとおりです。

結果II: 効率的な高解像度のテキスト・トゥ・イメージ (T2I) 生成

Graftingは現実世界の高解像度のタスクに対応できます。これをピクサートσ (DiT) を用いた2048×2048のテキストから画像への生成に適用します。この設定は特に困難です。長いシーケンス (16,384 トークン)、テキスト条件付きのマルチモーダル設定が必要で、トレーニングデータがないためです。生成レイテンシーの 62% 以上を占めるセルフアテンションオペレーターを対象に、わずか 12,000 個の合成サンプルを使用してグラフトを適用しています。結果として得られるモデルが実行されます。 1.43倍速く、 2% 未満の低下 ジュネーブでのスコア（47.78対49.75）、接ぎ木は大規模でも効果があることを示しています。

これが描写された画像です」ピカソ、フラクタル、モザイク、モザイク模様の顔と体、美しい女性、背景の白い宮殿大ホール、フォトリアリスティック。」グラフトモデルを使用して生成されました。

結果III: グラフト処理によりモデルの奥行きを幅に変換します。

グラフトでは、オペレータの入れ替えだけでなく、より構造的な編集も可能です。MLP グラフトの結果に動機付けられて、私たちはもっと急進的なことを試します。 最新のGPUはシーケンシャル計算よりも並列計算を優先するため、トランスブロックのすべてのペアを並列化することでDIT-XL/2を再配線します。 これにより、モデルの奥行きが 2 倍 (28→14) 減少します。グラフトモデルはFID 2.77を達成し、同様の深さを持つ他のモデルよりも優れています。私たちの知る限り、これは 最初の試み 逐次トランスブロックを事前学習済みのDITで並列に変換する。グラフトによってアーキテクチャを再構築できることが示された。

深度から幅までの再構築モデルを使用して生成された高品質のサンプルは次のとおりです。

今後の展望:アプリケーション

グラフトは、モデル機能の拡張 (短編から長編動画の理解や生成など) や、画像編集などのインタラクティブなアプリケーションのための効率的な推論スタックの構築など、効率が重要な場面で可能性を秘めています。私たちの調査結果が、コミュニティがグラフトを使った新しい建築デザインを探求するきっかけになることを願っています。

自分で試してみてください！