ヤン・シュイチェン/チェン・ミンミンの新作！ Sora のコアコンポーネントである DiT トレーニングは 10 倍高速化され、Masked Diffusion Transformer V2 はオープンソースです

Sora の魅力的なコア テクノロジーの 1 つとして、DiT は拡散トランスフォーマーを利用して生成モデルを大規模に拡張し、優れた画像生成効果を実現します。

ただし、モデルのサイズが大きくなると、トレーニングのコストが急増します。

南開大学Sea AI LabのYan Shuicheng氏とCheng Mingming氏の研究チーム、およびKunlun Wanwei 2050 Research Instituteは、ICCV 2023カンファレンスでマスク拡散トランスと呼ばれる新しいモデルを提案しました。このモデルは、マスク モデリング技術を使用して、意味表現情報を学習することで拡散トランスフォーマーのトレーニングを高速化し、画像生成分野で SoTA 効果を実現します。このイノベーションは、画像生成モデルの開発に新たなブレークスルーをもたらし、研究者により効率的なトレーニング方法を提供します。研究チームは、さまざまな分野の専門知識とテクノロジーを組み合わせることで、トレーニング速度を向上させ、生成結果を向上させるソリューションを提案することに成功しました。彼らの研究は、人工知能分野の発展に重要な革新的なアイデアに貢献し、将来の研究と実践に有益なインスピレーションを提供しました

写真

論文アドレス: https://arxiv.org/abs/2303.14389

GitHub アドレス: https://github.com/sail-sg/MDT

##最近、Masked Diffusion Transformer V2 が再び SoTA を更新し、DiT と比較してトレーニング速度が 10 倍以上向上し、ImageNet ベンチマークで 1.58 の FID スコアを達成しました。

論文とコードの最新バージョンはオープンソースです。

背景

DiT に代表される拡散モデルは画像生成の分野で大きな成功を収めてきましたが、研究者らは、拡散モデルは多くの場合、画像内のオブジェクトの部分間の意味的関係を効率的に学習することは困難であり、この制限がトレーニング プロセスの収束効率の低下につながります。

図

たとえば、上の図に示すように、DiT は 50k 番目で学習しました。トレーニング ステップ 犬の髪のテクスチャを生成し、200k トレーニング ステップで犬の目と口の 1 つを生成する方法を学習しますが、もう 1 つの目は見逃します。

300k のトレーニング ステップでも、DiT によって生成された犬の 2 つの耳の相対位置はあまり正確ではありません。

このトレーニングと学習のプロセスにより、拡散モデルは画像内のオブジェクトのさまざまな部分間の意味的関係を効率的に学習できず、各オブジェクトの意味的情報を個別に学習するだけであることがわかります。 。

研究者らは、この現象の理由は、拡散モデルが各ピクセルの予測損失を最小限に抑えることによって実際の画像データの分布を学習するためであると推測しています。このプロセスでは、さまざまなピクセル間の関係が無視されます。画像内のオブジェクトの各部分間の意味的な相対関係により、モデルの収束速度が遅くなります。

方法: マスクされた拡散トランスフォーマー

上記の観察に触発されて、研究者らは、トレーニングを改善するためにマスクされた拡散トランス (MDT) を提案しました。拡散モデルの効率とビルド品質。

MDT は、拡散トランスフォーマーの文脈的意味情報の学習能力を明示的に強化し、画像内のオブジェクト間の関係を強化する意味情報の連想学習を目的として、拡散トランスフォーマー用に設計されたマスク モデリング表現学習戦略を提案します。

写真

上図に示すように、MDT は拡散トレーニングを維持しながらマスク モデリングを導入します。学習戦略のプロセス。ノイズを含む画像トークンをマスクすることで、MDT は非対称拡散変換器 (非対称拡散変換器) アーキテクチャを使用して、マスクされていないノイズを含む画像トークンからマスクされた画像トークンを予測し、それによってマスク モデリングと拡散トレーニング プロセスを同時に実現します。

推論プロセス中、MDT は標準の拡散生成プロセスを維持します。 MDT の設計により、Difffusion Transformer は、マスク モデリング表現の学習によってもたらされる意味情報表現能力と、画像の詳細を生成する拡散モデルの能力の両方を得ることができます。

具体的には、MDT は VAE エンコーダーを通じて画像を潜在空間にマッピングし、それらを潜在空間で処理してコンピューティング コストを節約します。

トレーニング プロセス中、MDT はまずノイズが追加された画像トークンの一部をマスクし、残りのトークンを非対称拡散変換器に送信して、ノイズ除去後のすべての画像トークンを予測します。

#非対称拡散トランス アーキテクチャ

##写真

として上の図に示されているように、非対称拡散トランスのアーキテクチャには、エンコーダ、サイド補間器 (補助補間器)、およびデコーダが含まれています。

図

トレーニング プロセス中、Encoder はマスクされていないトークンのみを処理します。推論では、マスク ステップがないため、すべてのトークンが処理されます。

したがって、デコーダーがトレーニングまたは推論フェーズ中に常にすべてのトークンを処理できるようにするために、研究者らは解決策を提案しました。トレーニング プロセス中に、以下で構成される DiT ブロックを使用するというものです。補助補間器 (上の図に示す) は、エンコーダーの出力からマスクされたトークンを補間および予測し、推論のオーバーヘッドを追加することなく推論段階でそれを削除します。

MDT のエンコーダとデコーダは、グローバルおよびローカル位置エンコード情報を標準 DiT ブロックに挿入して、マスク部分のトークンの予測を支援します。

#非対称拡散トランス V2

##写真

として上の図に示すように、MDTv2 は、マスク拡散プロセス用に設計されたより効率的なマクロ ネットワーク構造を導入することにより、拡散とマスク モデリングの学習プロセスをさらに最適化します。

これには、エンコーダでの U-Net スタイルのロング ショートカットとデコーダでの高密度入力ショートカットの統合が含まれます。

このうち、dense input-shortcut は、マスクされたトークンにノイズを追加してデコーダーに送信し、マスクされたトークンに対応するノイズ情報を保持するため、拡散のトレーニングが容易になります。プロセス。 。

さらに、MDT は、より高速な Adan オプティマイザー、タイムステップ関連の損失重み、拡散モデルのマスクされたトレーニング プロセスをさらに加速する拡張マスク比など、より優れたトレーニング戦略も導入しました。 。

#実験結果

ImageNet 256 ベンチマーク生成の品質比較

Image

#上の表は、ImageNet 256 ベンチマークにおけるさまざまなモデル サイズでの MDT と DiT のパフォーマンスを比較しています。

MDT は、すべてのモデル サイズにおいて、より少ないトレーニング コストでより高い FID スコアを達成することは明らかです。

MDT のパラメータと推論コストは基本的に DiT と同じです。これは、前述したように、MDT の推論プロセス中に DiT と一致する標準の拡散プロセスが維持されるためです。

最大の XL モデルの場合、400k ステップでトレーニングされた MDTv2-XL/2 は、7000k ステップでトレーニングされた DiT-XL/2 を大幅に上回り、FID スコアは 1.92 向上しました。この設定では、MDT が DiT よりも約 18 倍高速にトレーニングできることが結果からわかります。

小規模モデルの場合、MDTv2-S/2 は、大幅に少ないトレーニング ステップで、DiT-S/2 よりも大幅に優れたパフォーマンスを実現します。たとえば、400k ステップの同じトレーニングでは、MDTv2 の FID インデックスは 39.50 で、これは DiT の FID インデックス 68.40 を大幅に上回っています。

さらに重要なのは、この結果は、400k トレーニング ステップでのより大きなモデル DiT-B/2 のパフォーマンスも上回っていることです (39.50 対 43.47)。

ImageNet 256 ベンチマーク CFG 生成の品質比較

Image

我々はまだ上の表は、分類子を使用しないガイダンスの下で、MDT と既存の方法の画像生成パフォーマンスを比較しています。

MDT は、FID スコア 1.79 で、以前の SOTA DiT や他の手法を上回ります。 MDTv2 はパフォーマンスをさらに向上させ、少ないトレーニング ステップで画像生成の SOTA FID スコアを新たな最低値の 1.58 に押し上げます。

DiT と同様に、トレーニングを継続しても、トレーニング中にモデルの FID スコアの飽和は観察されませんでした。

#MDT が PaperWithCode のリーダーボードで SoTA を更新

収束速度の比較

図

上の図は、ImageNet 256 ベンチマークにおける 8×A100 GPU 上の DiT-S/2 ベースライン、MDT-S/2、および MDTv2 を比較しています。 . - 異なるトレーニング ステップ/トレーニング時間における S/2 の FID パフォーマンス。

MDT は、より優れたコンテキスト学習機能のおかげで、パフォーマンスと生成速度の両方で DiT を上回っています。 MDTv2 のトレーニング収束速度は、DiT の 10 倍以上です。

MDT は、トレーニング ステップとトレーニング時間の点で DiT よりも約 3 倍高速です。 MDTv2 は、MDT と比較してトレーニング速度をさらに約 5 倍向上させます。

たとえば、MDTv2-S/2 は、トレーニングに約 100 時間 (1500k ステップ) かかる DiT-S/2 よりも、わずか 13 時間 (15k ステップ) で優れたパフォーマンスを示します。は、コンテキスト表現の学習が拡散モデルの生成学習を高速化するために重要であることを明らかにしています。

要約と考察

MDT は、拡散学習に MAE と同様のマスク モデリング表現学習スキームを導入することで、画像オブジェクトの特性を利用できます。コンテキスト情報は、不完全な入力画像の完全な情報を再構築することで、画像内の意味部分間の相関関係を学習し、画像生成の品質と学習速度を向上させます。

研究者らは、視覚表現学習を通じて物理世界の意味的理解を強化することで、物理世界に対する生成モデルのシミュレーション効果を向上できると考えています。これは、生成モデルを通じて物理世界シミュレーターを構築するという Sora のビジョンと一致します。この研究が、表現学習と生成学習の統合に関するさらなる研究のきっかけとなることを願っています。

参考:

https://arxiv.org/abs/2303.14389

以上がヤン・シュイチェン/チェン・ミンミンの新作！ Sora のコアコンポーネントである DiT トレーニングは 10 倍高速化され、Masked Diffusion Transformer V2 はオープンソースですの詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。