[論文レビュー] Exploring Vision Transformers as Diffusion Learners

He Cao, Jianan Wang+5 more|arXiv (Cornell University)|Dec 28, 2022

Generative Adversarial Networks and Image Synthesis被引用数 7

ひとこと要約

本論文では、非対称なエンコーダ・デコーダアーキテクチャを備えた、ビジョンTransformerベースの拡散バックボーンであるASCENDを提案する。このアーキテクチャにより、画像生成タスクにおける性能が顕著に向上した。U-Netに代えて強力なSwin Transformerエンコーダと軽量な畳み込みデコーダを採用することで、CIFAR-10、CelebA、LSUN、CUB-Birdの各タスクで競争力のある結果を達成した。さらに、64×64解像度を超える単一段階のテキスト到着拡散モデルを訓練した最初の例であり、エンドツーエンドの高解像度生成における新たなベンチマークを確立した。

ABSTRACT

Score-based diffusion models have captured widespread attention and funded fast progress of recent vision generative tasks. In this paper, we focus on diffusion model backbone which has been much neglected before. We systematically explore vision Transformers as diffusion learners for various generative tasks. With our improvements the performance of vanilla ViT-based backbone (IU-ViT) is boosted to be on par with traditional U-Net-based methods. We further provide a hypothesis on the implication of disentangling the generative backbone as an encoder-decoder structure and show proof-of-concept experiments verifying the effectiveness of a stronger encoder for generative tasks with ASymmetriC ENcoder Decoder (ASCEND). Our improvements achieve competitive results on CIFAR-10, CelebA, LSUN, CUB Bird and large-resolution text-to-image tasks. To the best of our knowledge, we are the first to successfully train a single diffusion model on text-to-image task beyond 64x64 resolution. We hope this will motivate people to rethink the modeling choices and the training pipelines for diffusion-based generative models.

研究の動機と目的

• ビジョンTransformerを拡散モデルバックボーンとして効果的に使用できるかを体系的に評価すること、特にU-Netとの比較を通じて低解像度および高解像度画像生成における性能を検証すること。
• U-ViTアーキテクチャの改良（IU-ViT）により、バニラViTとU-Netの低解像度画像生成における性能ギャップを埋めること。
• スケーラブルで高解像度かつマルチモーダルな拡散生成を実現する非対称エンコーダ・デコーダアーキテクチャ（ASCEND）を提案・検証すること。
• 複数段階のトレーニングパイプラインの現状を問い直すために、ビジョンTransformerバックボーンを用いたエンドツーエンドの単一段階高解像度テキスト到着画像生成を実証すること。
• 拡散ベースの生成モデリングにおけるバックボーン設計とトレーニングパラダイムの見直しを研究コミュニティに促すこと。

提案手法

• スキップ接続の強化と残留ダウンサンプリング／アップサンプリングを導入したU-ViTの改善版であるIU-ViTを提案。パッチマージング／拡張の代わりに使用する。
• Swin Transformerエンコーダと畳み込みデコーダを組み合わせた非対称エンコーダ・デコーダバックボーンであるASCENDを設計。
• 特徴階層と勾配の流れを改善するために、パッチマージング／拡張の代わりに残留ダウンサンプリング／アップサンプリングを採用。
• テキスト条件付き画像生成タスクにおけるテキストエンコーディングに、事前学習済みのT5-XXLエンコーダを統合。
• LAION-Aestheticsなどの大規模データセットを用いて、128×128および256×256解像度の画像生成をエンドツーエンドでトレーニング。
• CIFAR-10を用いたアブレーションスタディにより、スキップ接続、ダウンサンプリング手法、エンコーダ／デコーダ構成などのアーキテクチャ的選択を検証。

実験結果

リサーチクエスチョン

• RQ1 ビジョンTransformerは拡散モデルバックボーンとして効果的に使用可能か？また、低解像度および高解像度画像生成においてU-Netと比較してどのように性能を発揮するか？
• RQ2 バニラViTベースのバックボーンをU-Netと同等の性能にするために、どのようなアーキテクチャ的改良が必要か？
• RQ3 非対称なエンコーダ・デコーダ構造は、拡散ベースの画像生成における性能とスケーラビリティを向上させるか？
• RQ4 ビジョンTransformerバックボーンを用いた単一段階のエンドツーエンド拡散モデルは、64×64解像度を超える高解像度で高品質な画像を生成できるか？
• RQ5 エンコーダアーキテクチャの選択（例：Swin Transformer）が、高解像度およびテキスト到着画像生成における性能に与える影響は何か？

主な発見

• ASCENDはCIFAR-10でFréchet Inception Distance（FID）2.98を達成し、熟練したチューニング済みU-Netバックボーンと同等の性能を示した。
• ASCENDにおいてパッチマージング／拡張を残留ダウンサンプリング／アップサンプリングに置き換えることで、FIDが12.81から2.98に低下し、顕著なアーキテクチャ的向上が確認された。
• ASCENDのスキップ接続を削減するとFIDが3.54ポイント上昇し、特徴再構築におけるスキップ接続の重要性が裏付けられた。
• 590Mパラメータを持つモデルで、128×128解像度の高品質なテキスト到着画像サンプルを生成し、64×64解像度を超える単一段階トレーニングで初めて成功した。
• ASCENDは高解像度生成においてIU-ViTおよびU-Netを上回り、後期のトレーニング段階でより速い収束と優れたサンプル品質を示した。
• アブレーションスタディにより、エンコーダにSwinBlocksを、デコーダに畳み込み層を用いる構成が、対称的または他の代替構成と比較して優れた結果をもたらすことが確認された。