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

モデルアーキテクチャの設計は、データ・アルゴリズム・計算資源・ベンチマークと並んで、現代の AI システム構築の中核をなします。アーキテクチャは学習可能な関数を定義し、どの演算子を使うか（アテンション、畳み込みなど）、どう構成するか（モデルの深さ、幅など）といった重要な選択を伴います。その重要性にもかかわらず、アーキテクチャについて何がうまくいき何がうまくいかないのかという知見は得にくいままです。とりわけ現在の基盤モデルの時代では、モデルをゼロから学習させるコストが法外に高いためです。その結果、新しいアーキテクチャの探索は、特に生成モデルにおいて大きな課題であり続けています。

新しいソフトウェアがゼロから書かれるのではなく既存のコードの上に築かれるのと同じように、学習済みモデルは新しいアーキテクチャ設計を探索するための足場になり得るのでしょうか？ 私たちは、学習済みモデルに対するアーキテクチャの編集を検討します。対象とするのは、画像・動画生成で広く使われている生成 Transformer の一種、拡散 Transformer（DiT）です¹²³。

学習済みモデルは、画像や動画の生成といったタスクを実行する計算グラフを実装しています。電球を点灯させるために配線された電気回路のようなものだと考えてください。新しいアーキテクチャのアイデアと学習済みモデルが与えられたとき、少ない計算予算のもとで計算グラフを書き換えることでそのアイデアを具体化できるかを検証します。たとえば、DiT のマルチヘッドアテンション（MHA）や多層パーセプトロン（MLP）を畳み込みベースの設計に置き換えられる、という仮説を立てられます。これを検証する単純な方法は、モデルの品質を保ったまま MHA や MLP の演算子を畳み込み演算子に差し替えることです。ここで 2 つの重要な問いが生じます。

• (Q1) 演算子の初期化：新しい演算子は、計算グラフに組み込む前にどう初期化すべきか。
• (Q2) 誤差の蓄積：複数の演算子を置き換えたときの誤差の伝播をどう抑えるか。

これらの問いに答えるため、私たちはアーキテクチャ編集のためのシンプルな 2 段階手法「グラフティング（grafting）」を提案します。手順は次のとおりです。

• (i) 活性化蒸留：回帰目的関数を用いて元の演算子の活性化を蒸留し、その機能を新しい演算子へ移します。
• (ii) 軽量なファインチューニング：限られたデータでファインチューニングを行い、複数の新しい演算子を組み込んだことによる誤差の伝播を抑えます。

グラフティングを、難易度を段階的に上げた一連の生成モデリングタスクで検証しました。

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

まず、256×256 解像度の DiT-XL/2 を用いたクラス条件付き画像生成でグラフティングを検証しました。この設定では、softmax アテンション（MHA）を局所ゲート付き畳み込み（Hyena-SE と、私たちが提案する Hyena-X/Y）、局所アテンション（スライディングウィンドウ）、線形アテンション（Mamba-2）などに置き換えます。MLP については、拡張率の異なる変種（3x、6x など）と畳み込み版（Hyena-X）を試しました。

興味深いことに、いくつかの交互配置型ハイブリッド設計が良好な生成品質を達成しました。FID スコアは 2.38〜2.64（低いほど良い。DiT-XL/2 のベースラインは 2.27）でした。グラフティングはシンプルかつ軽量で、各実験は 8×H100 GPU で 24 時間以内に完了し、元モデルの事前学習計算量の 2% 未満しか使いません。

グラフティングしたモデルで生成した高品質なサンプルをいくつかご紹介します。

結果 II：高解像度 Text-to-Image（T2I）生成の効率化

グラフティングは実世界の高解像度タスクにも拡張できます。PixArt-Σ（DiT）を用いた 2048×2048 の text-to-image 生成に適用しました。この設定は特に難しく、長い系列（16,384 トークン）、テキスト条件付けを伴うマルチモーダル設定、そして学習データがない状況を扱います。生成レイテンシの 62% 以上を占める自己アテンション演算子を対象に、わずか 1.2 万件の合成サンプルでグラフティングを適用しました。得られたモデルは 1.43 倍高速に動作し、GenEval スコア（47.78 対 49.75）の低下は 2% 未満にとどまりました。大規模な設定でもグラフティングが有効であることを示しています。

以下は「ピカソ風、フラクタル、モザイク、モザイク模様の顔と身体、美しい女性、背景は白い宮殿の大広間、フォトリアリスティック。」というプロンプトから、グラフティングしたモデルで生成した画像です。

結果 III：グラフティングによるモデルの深さから幅への変換

グラフティングは演算子の差し替えにとどまらず、より構造的な編集も可能にします。MLP のグラフティング結果に着想を得て、さらに大胆な試みを行いました。現代の GPU は逐次計算よりも並列計算に適しているため、DiT-XL/2 の Transformer ブロックを 2 つずつ並列化するように配線し直しました。これによりモデルの深さは半分（28→14）になります。グラフティング後のモデルは FID 2.77 を達成し、同程度の深さを持つ他のモデルを上回りました。私たちの知る限り、これは学習済み DiT において逐次的な Transformer ブロックを並列に変換した最初の試みであり、グラフティングによってアーキテクチャを再構成できることを示しています。

深さを幅へ再構成したモデルで生成した高品質なサンプルをいくつかご紹介します。

今後の展望：応用

グラフティングは効率が重要となる場面で有望です。たとえばモデルの能力拡張（短い動画から長編動画の理解・生成へ）や、画像編集のようなインタラクティブな用途に向けた効率的な推論スタックの構築などです。今回の知見が、グラフティングを用いた新しいアーキテクチャ設計の探索をコミュニティで後押しすることを願っています。

ぜひお試しください

謝辞

スタンフォード大学、Together AI、カリフォルニア大学サンディエゴ校、ノースウェスタン大学、Google DeepMind、Salesforce Research との共同研究です。

引用

本研究を利用する場合は、以下のテクニカルレポートを引用してください。

@inproceedings{chandrasegaran2025grafting,
  title     = {Exploring Diffusion Transformer Designs via Grafting},
  author    = {Chandrasegaran, Keshigeyan and Poli, Michael and Fu, Daniel Y. and Kim, Dongjun and Hadzic, Lea M. and Li, Manling and Gupta, Agrim and Massaroli, Stefano and Mirhoseini, Azalia and Niebles, Juan Carlos and Ermon, Stefano and Li, Fei-Fei},
  booktitle = {Advances in Neural Information Processing Systems},
  volume    = {38},
  pages     = {20293--20324},
  year      = {2025},
  doi       = {10.52202/085713-0603},
}

参考文献

[1] Peebles, William, and Saining Xie. "Scalable diffusion models with transformers." Proceedings of the IEEE/CVF international conference on computer vision. 2023.

[2] Brooks, Tim, et al. “Video Generation Models as World Simulators.” OpenAI, 2024, openai.com/index/video-generation-models-as-world-simulators/

[3] Gupta, Agrim, et al. "Photorealistic video generation with diffusion models." European Conference on Computer Vision. Cham: Springer Nature Switzerland, 2024.