U-Net から DiT へ: ドミナンス拡散モデルにおける変圧器技術の応用

近年、Transformer によって機械学習がルネッサンスを迎えています。過去 5 年間、自然言語処理、コンピューター ビジョン、その他の分野のニューラル アーキテクチャはトランスフォーマーによってほぼ独占されてきました。

しかし、この傾向の影響を受けていない画像レベルの生成モデルも数多くあります。たとえば、拡散モデルは、過去 1 年間で画像生成において驚くべき成果を達成しました。これらのほとんどすべてモデル 畳み込み U-Net をバックボーンとして使用します。これはちょっと意外ですね！過去数年間のディープ ラーニングの大きな話題は、さまざまな分野で Transformer が優勢になったことです。 U-Net や畳み込みが拡散モデルで非常に優れたパフォーマンスを発揮する特別な点はありますか?

U-Net バックボーン ネットワークを拡散モデルに初めて導入した研究は、Ho らに遡ることができます。この設計パターンは、自己回帰生成モデル PixelCNN をほんのわずかな変更を加えて継承しています。 PixelCNN は畳み込み層で構成されており、多くの ResNet ブロックが含まれています。標準の U-Net と比較すると、PixelCNN の追加の空間セルフ アテンション ブロックがトランスフォーマーの基本コンポーネントになります。他の研究とは異なり、Dhariwal 氏と Nichol 氏らは、適応正規化層を使用して条件情報とチャネル数を畳み込み層に注入するなど、U-Net のいくつかのアーキテクチャ上の選択肢を排除しています。

この記事では、UC バークレーの William Peebles とニューヨーク大学の Xie Saining が「変圧器を使用したスケーラブルな拡散モデル」を執筆しました。目標は、拡散モデルにおけるアーキテクチャ上の選択の重要性を明らかにすることです。将来の生成モデル研究のためのガイダンスを提供し、経験的なベースラインを提供します。この研究は、U-Net 誘導バイアスは拡散モデルのパフォーマンスにとって重要ではなく、変圧器などの標準設計で簡単に置き換えることができることを示しています。

この発見は、拡散モデルがアーキテクチャの統合トレンドから恩恵を受けることができることを示唆しています。たとえば、拡散モデルは他の分野からベスト プラクティスやトレーニング方法を継承し、これらのモデルのスケーラビリティを維持できます。そして効率性。標準化されたアーキテクチャは、クロスドメイン研究の新たな可能性も開きます。

• #論文アドレス: https://arxiv.org/pdf/2212.09748.pdf
• プロジェクト アドレス: https://github.com/facebookresearch/DiT
• Paper ホームページ: https:/ / www.wpeebles.com/DiT

この研究は、トランスフォーマーベースの拡散モデルの新しいクラスである拡散トランスフォーマー (略して DiT) に焦点を当てています。 DiT は、Vision Transformers (ViT) のベスト プラクティスに従い、小さいながらも重要な調整が加えられています。 DiT は、ResNet などの従来の畳み込みネットワークよりも効率的に拡張できることが示されています。

具体的には、この記事では、ネットワークの複雑さとサンプル品質の観点から、Transformer のスケーリング動作を研究します。潜在拡散モデル (LDM) フレームワークの下で DiT 設計空間を構築およびベンチマークすることにより、拡散モデルが VAE の潜在空間内でトレーニングされることで、U-Net バックボーンをトランスフォーマーに置き換えることが可能であることが示されています。この論文はさらに、DiT が拡散モデルのスケーラブルなアーキテクチャであることを示しています。ネットワークの複雑さ (Gflops で測定) とサンプルの品質 (FID で測定) の間には強い相関関係があります。 DiT を拡張し、大容量バックボーン (118.6 Gflops) で LDM をトレーニングするだけで、クラス条件付き 256 × 256 ImageNet 生成ベンチマークで 2.27 FID という最先端の結果が達成されます。

拡散トランスフォーマー

DiTs は、拡張性を維持するために標準トランス アーキテクチャに可能な限り忠実であることを目的とした拡散モデルの新しいアーキテクチャです。 DiT は ViT のベスト プラクティスの多くを保持しており、図 3 は完全な DiT アーキテクチャを示しています。

DiT への入力は空間表現 z (256 × 256 × 3 画像の場合、z の形状は 32 × 32 × 4) です。 ）。 DiT の最初の層は patchify で、各パッチを入力に線形的に埋め込むことで空間入力を T トークンのシーケンスに変換します。 patchify 後、標準の ViT 周波数ベースの位置埋め込みをすべての入力トークンに適用します。

patchify によって作成されるトークン T の数は、パッチ サイズのハイパーパラメーター p によって決まります。図 4 に示すように、p を半分にすると T が 4 倍になり、したがってトランスの Gflops が少なくとも 4 倍になります。この記事では、p = 2,4,8 を DiT 設計空間に追加します。

DiT ブロック設計: patchify の後、入力トークンは一連の Transformer ブロックによって処理されます。ノイズの多い画像入力に加えて、拡散モデルはノイズ時間ステップ t、クラス ラベル c、自然言語などの追加の条件情報を処理する場合があります。この記事では、条件付き入力をさまざまな方法で処理する 4 つのトランスフォーマー ブロックのバリエーションについて説明します。これらの設計は、標準の ViT ブロック設計にわずかながら重要な変更を加えたものです。すべてのモジュールの設計を図 3 に示します。

この記事では、モデルの深さと幅によって異なる 4 つの構成 (DiT-S、DiT-B、DiT-L、および DiT-XL) を試しました。これらのモデル構成の範囲は 33M ～ 675M のパラメーターと 0.4 ～ 119 の Gflops です。

実験

研究者らは、それぞれ異なるブロック設計 (インコンテキスト (119.4 Gflops)、クロス) を使用して、最高の Gflop で 4 つの DiT-XL/2 モデルをトレーニングしました。 -attention (137.6Gflops)、アダプティブ レイヤ ノルム (adaLN、118.6Gflops)、または adaLN-zero (118.6Gflops)。次に、トレーニング中に FID が測定されました。その結果を図 5 に示します。

モデル サイズとパッチ サイズを拡張します。図 2 (左) は、各モデルの Gflops と 400K トレーニング反復における FID の概要を示しています。モデル サイズを大きくし、パッチ サイズを小さくすると、拡散モデルが大幅に改善されることがわかります。

図 6 (上) は、パッチ サイズが一定のまま、モデル サイズが増加するにつれて FID がどのように変化するかを示しています。 Transformer をより深く、より広くすることで、トレーニングのすべての段階で 4 つの設定にわたって FID が大幅に向上しました。同様に、図 6 (下) は、パッチ サイズが縮小され、モデル サイズが一定のままの場合の FID を示しています。研究者らは、DiT によって処理されるトークンの数を単に拡大し、トレーニング プロセス全体を通じてパラメーターをほぼ固定したままにするだけで、FID が大幅に改善されたことを再度観察しました。

図 8 は、400K トレーニング ステップにおける FID-50K とモデルの Gflops の比較を示しています。

SOTA 拡散モデル 256×256 ImageNet。拡張分析の後、研究者らはステップ数 7M で最高の Gflop モデルである DiT-XL/2 のトレーニングを続けました。図 1 は、このモデルのサンプルを示し、カテゴリ条件付き生成 SOTA モデルと比較しています。結果は表 2 に示されています。

分類子ガイダンスを使用しない場合、DiT-XL/2 は以前のすべての拡散モデルよりも優れたパフォーマンスを示し、LDM が以前に達成した 3.60 を上回りました。 Best FID-50K は低下しました。 2.27まで。図 2 (右) に示すように、LDM-4 (103.6 Gflops) などの潜在空間 U-Net モデルと比較すると、DiT-XL/2 (118.6 Gflops) は ADM (1120 Gflops) よりもはるかに計算効率が高くなります。 ADM-U (742 Gflops)、ピクセル空間 U-Net モデルははるかに効率的です。

#表 3 は、SOTA アプローチとの比較を示しています。 XL/2 は、この解像度でも以前のすべての拡散モデルよりも優れたパフォーマンスを示し、ADM の以前の最高 FID 3.85 を 3.04 に改善しました。

研究の詳細については、元の論文を参照してください。

以上がU-Net から DiT へ: ドミナンス拡散モデルにおける変圧器技術の応用の詳細内容です。詳細については、PHP 中国語 Web サイトの他の関連記事を参照してください。