[論文レビュー] All are Worth Words: A ViT Backbone for Diffusion Models

Fan Bao, Nie, Shen+5 more|arXiv (Cornell University)|Sep 25, 2022

Generative Adversarial Networks and Image Synthesis被引用数 7

ひとこと要約

この論文は、時間、条件、ノイズの画像パッチをすべてトークンとして扱い、浅い層と深い層の間に長いスキップ接続を採用する、拡散モデル用のビジョントランスフォーマー（ViT）ベースのバックボーン、U-ViTを提案する。U-ViTは、ImageNet 256×256のクラス条件付き生成でFIDスコア2.29、MS-COCOのテキストから画像への生成で5.48を達成し、ダウンサンプリング／アップサンプリング操作を必要としないにもかかわらず、CNNベースのU-Netを上回る性能を示した。

ABSTRACT

Vision transformers (ViT) have shown promise in various vision tasks while the U-Net based on a convolutional neural network (CNN) remains dominant in diffusion models. We design a simple and general ViT-based architecture (named U-ViT) for image generation with diffusion models. U-ViT is characterized by treating all inputs including the time, condition and noisy image patches as tokens and employing long skip connections between shallow and deep layers. We evaluate U-ViT in unconditional and class-conditional image generation, as well as text-to-image generation tasks, where U-ViT is comparable if not superior to a CNN-based U-Net of a similar size. In particular, latent diffusion models with U-ViT achieve record-breaking FID scores of 2.29 in class-conditional image generation on ImageNet 256x256, and 5.48 in text-to-image generation on MS-COCO, among methods without accessing large external datasets during the training of generative models. Our results suggest that, for diffusion-based image modeling, the long skip connection is crucial while the down-sampling and up-sampling operators in CNN-based U-Net are not always necessary. We believe that U-ViT can provide insights for future research on backbones in diffusion models and benefit generative modeling on large scale cross-modality datasets.

研究の動機と目的

• ビジョントランスフォーマー（ViT）が、拡散モデルにおけるCNNベースのU-Netバックボーンに効果的に置き換え可能かどうかを調査すること。
• U-Netアーキテクチャにおけるダウンサンプリングおよびアップサンプリング操作が、拡散ベースの画像生成において必須かどうかを評価すること。
• ViTベースの拡散モデルバックボーンにおける長距離スキップ接続の影響を特定すること。
• シンプルなViTベースのアーキテクチャを用いて、非条件生成、クラス条件付き生成、テキストから画像への生成タスクでSOTAの性能を達成すること。

提案手法

• U-ViTは、ノイズの画像パッチ、タイムステップ埋め込み、条件埋め込みをすべてトークンとして扱い、統一的なトークンベースの入力表現を可能にする。
• 拡散モデルにおけるピクセルレベルの再構成に不可欠な低レベル特徴を保持するために、浅い層と深い層の間に長いスキップ接続を採用する。
• クラストークンと位置埋め込みを備えた標準的なViTトランスフォーマーエンコーダーを用い、その後に学習可能な位置埋め込みを備えたデコーダーヘッドを接続する。
• 最終出力の前に、視覚的品質を向上させるためにオプションの3×3畳み込みブロックを追加する。
• ピクセル空間および潜在空間の両方の拡散モデルにこのアーキテクチャを適用し、ノイズ予測目的を用いて学習する。
• FIDやその他の指標を用いて、CIFAR-10、ImageNet 256×256、ImageNet 512×512、MS-COCOの複数のベンチマークで評価する。

実験結果

リサーチクエスチョン

• RQ1 CNNベースのU-Netを置き換えるために、ビジョントランスフォーマーバックボーンが拡散モデルで効果的に機能するか。
• RQ2 U-Netにおけるダウンサンプリングおよびアップサンプリング機構が、拡散モデルにおける高品質な画像生成に不可欠かどうか。
• RQ3 標準的なViTアーキテクチャと比較して、ViTベースの拡散モデルバックボーンにおける長距離スキップ接続の重要性はどの程度か。
• RQ4 大規模な外部データセットに依存せずに、ViTベースのアーキテクチャがテキストから画像への生成およびクラス条件付き画像生成でSOTAの性能を達成できるか。

主な発見

• U-ViTは、大規模な外部データセットにアクセスしない状況でも、クラス条件付きImageNet 256×256生成でSOTAのFIDスコア2.29を達成した。
• MS-COCOにおけるテキストから画像への生成では、U-ViT-S（Deep）がFID 5.48という新しいSOTAを達成し、優れた意味的整合性と画像品質を示した。
• U-ViTは、サイズとFLOPsが同程度のCNNベースのU-Netよりも、非条件および条件付き画像生成タスクの両方で優れた性能を示した。
• アブレーションスタディの結果、長距離スキップ接続が性能に不可欠であることが確認された一方で、U-Netにおけるダウンサンプリング／アップサンプリング操作は、高品質な生成に常に必要ではないことが示された。
• 特にテキストから画像への生成では、U-ViTはU-Netよりもより意味的に正確なサンプルを生成し、プロンプトに記載されたバットやボールなどのオブジェクトをよりよく捉えた。
• サンプリングステップやインセプションスコア、精度、再現率といった複数の指標において、U-ViTは強固で高品質な生成性能を維持していた。