U-ViT: 拡散モデル用ViTバックボーン
3つの要点
✔️ 画像生成タスクにおいて、拡散モデルは従来のGANを上回っている
✔️ 拡散モデルは主にCNNベースのUNetを利用しており、ViTバックボーンを導入することで性能向上
✔️ ViTベースUNetはImageNetとMS-COCO上の画像生成で最高FIDを達成
All are Worth Words: A ViT Backbone for Diffusion Models
written by Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, Jun Zhu
(Submitted on 25 Sep 2022, last revised 25 Mar 2023)
Comments: Accepted to CVPR 2023. Published on arxiv.
Subjects: omputer Vision and Pattern Recognition (cs.CV); Artificial Intelligence (cs.AI); Machine Learning (cs.LG)
code：
本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。
はじめに
拡散モデルは、高品質な画像生成のために最近登場した強力な深層生成モデルです。拡散モデルは急速に成長し、テキストから画像への生成、画像から画像への生成、ビデオ生成、音声合成、3D合成などに応用されています。
アルゴリズムの改良とともに、バックボーンの改良は拡散モデルにおいて重要な役割を果たしています。その代表例が、以前の研究で用いられてきた畳み込みニューラルネットワーク（CNN）に基づくU-Netです。CNNベースのUNetは、一連のダウンサンプリングブロック、一連のアップサンプリングブロック、そしてこれらのグループ間の長いスキップ接続によって特徴付けられています。これは、画像生成タスクの拡散モデルにおいて支配的な役割を果たしています。
一方、ビジョントランスフォーマー（ViT）は、さまざまなビジョンタスクで有望な成果を示しています。ViTは、場合によって、CNNベースのアプローチと同等またはそれ以上の性能を発揮します。このため、自然な疑問が生じます：拡散モデルにおいてCNNベースのU-Netに依存する必要性はあるのでしょうか？
今回の解説論文では、ViTをベースにしたUNet（U-ViT）を提案します。提案手法は、ImageNetとMS-COCO上の画像生成で最高FID（画質を表す評価指標）を達成しました。
提案手法
U-ViTの概要は図1に示されています。このネットワークは、拡散過程の時間 \(t\)、条件 \(c\)、およびノイズのある画像 \(x_t\) を受け取り、\(x_t\) に注入されるノイズを予測します。ViTのデザイン手法に倣い、画像はパッチに分割され、U-ViTは時間、条件、および画像パッチを含むすべての入力をトークン（単語）として扱います。
CNNベースのU-Netと同様に、U-ViTも浅い層と深い層の間の長いスキップ接続を使用します。拡散モデルの学習はピクセルレベルの予測タスクであり、低レベルの特徴に敏感です。長いスキップ接続は、低レベルの特徴に対するショートカットを提供し、ノイズ予測ネットワークのトレーニングを容易にします。
さらに、U-ViTではオプションで、出力前に3×3の畳み込みブロックを追加します。これは、トランスフォーマーによって生成される画像の潜在的なアーティファクトを防ぐためのものです。
U-ViTの各部分の詳細について、次のサブセクションで説明します。
実装詳細
本節では、CIFAR10での生成画像の画質（FID）を通じて、U-ViTの構造の最適化を行います。全体の結果の概要は図２に示されています。
・長いスキップ接続の組み合わせる方法
主ブランチと長いスキップブランチからの埋め込みを$h_m, h_s \in \mathbb{R}^{L \times D}$とします。次のトランスフォーマーブロックにフィードする前に、それらを組み合わせるいくつかの方法を検討します：
1. $h_m, h_s $を連結し、その後線形射影を行う（図1参照）：$\text{Linear}(\text{Concat}(h_m, h_s))$。
2. $h_m, h_s$ を直接加算する： $h_m + h_s$。
3. $h_s$に線形射影を行い、それを$h_m$に加算する： $h_m + \text{Linear}(h_s)$。
4. $h_m, h_s$を加算し、その後線形射影を行う： $\text{Linear}(h_m + h_s) $。
5. 長いスキップ接続を削除
図2（a）に示されているように、その中で、連結$\text{Linear} \text{Concat}(h_m, h_s)$を用いた最初の方法が最も良い結果を示しました。特に、長いスキップ接続がない時よりも、生成画質が大きく向上できました。
・時間の条件の入力方法
時間の条件$t$をネットワークに入力する2つの方法を検討します。方法(1) は、図1に示すように、トークンとして扱うことです。方法(2)は、トランスフォーマーブロック内の層正規化後に時間を組み込むことで、これはU-Netで使用されている適応的グループ正規化に似ています。2番目の方法は、適応的層正規化（AdaLN）と呼ばれます。図2（b）に示されているように、時間をトークンとして扱う方法(1) がAdaLNよりも優れた性能を示しています。
・トランスフォーマーの後に畳み込みブロックを追加する方法
トランスフォーマーの後に畳み込みブロックを追加する2つの方法があります。(1) トークンの埋め込みを画像パッチにマッピングする線形射影の後に3×3の畳み込みブロックを追加することです（図1に示すように）。 (2) この線形射影の前に3×3の畳み込みブロックを追加することです。さらに、追加の畳み込みブロックを削除した場合と比較します。図2（c）に示されているように、線形射影の後に3×3の畳み込みブロックを追加する方法(1) が、他の2つの選択肢よりもわずかに優れたパフォーマンスを示しています。
・パッチ埋め込みの方法
従来のパッチ埋め込みは、パッチをトークン埋め込みにマッピングする線形射影です（図1に示すように）。この方法のかわりに、3×3の畳み込みブロックのスタックに続いて、1×1の畳み込みブロックを使用して、画像をトークン埋め込みにマッピングという方法も検討しました。しかし、図２（ｄ）に示されているように、従来のパッチ埋め込みの方が優れたパフォーマンスを示しているので、最終のモデルはこの方法を利用します。
・位置埋め込みの方法
本論文では、元のViTで提案された1次元の学習可能な位置埋め込みを利用します。 2次元の正弦位置埋め込みという方法もありますが、図2（e）に示されているように、1次元の学習可能な位置埋め込みの方が優れたパフォーマンスを示します。また、位置埋め込みを使用しない場合も試しましたが、モデルは鮮明な画像を生成できず、これは位置情報が画像生成に重要であることを示しています。
ネットワークの深さ、幅、パッチサイズの影響
ここで、U-ViTのスケーリング特性をCIFAR10で調査し、層の数、幅、およびパッチサイズの影響を検討します。図3に示すように、層の数が9から13に増えると性能が向上しますが、17より深いモデルでは効果が見られませんでした。同様に、幅を増やすと性能が向上しますが、一定の幅を超えると効果が見られなくなります。
パッチサイズを小さくすると性能が向上しますが、一定のサイズを下回ると性能が落ちます。小さなパッチサイズは、拡散モデルの低レベルのノイズ予測タスクに適すと考えられます。一方、高解像度の画像では小さなパッチサイズの使用がコストがかかるため、まず画像を低次元の潜在表現に変換し、それらをU-ViTでモデル化する必要があります。詳しくは実験の節で説明します。
実験
データセットと設定
無条件の画像生成、クラス条件付き画像生成、およびテキストから画像生成という３つのタスクでU-ViTの効果を検証します。
無条件の画像生成の実験は、CIFAR10（50,000枚）とCelebA 64×64（162,770枚）で行われます。クラス条件付き画像生成については、1,000の異なるクラスから1,281,167枚のトレーニング画像を含む64×64、256×256のImageNetデータセット、および512×512の解像度データセットで実験を行います。テキストから画像への学習にはMS-COCO（82,783枚のトレーニング画像と40,504枚の検証画像）を利用します。
256×256および512×512の高解像度画像の生成の場合、Latent diffusion models（LDM）[Rombach et.al, 2022]が提供する事前学習済みの画像オートエンコーダを使用して、それぞれ32×32および64×64の解像度の潜在表現に変換します。その後、U-ViTを使用して、これらの潜在表現をモデル化します。
MS-COCOでテキストから画像生成の場合、離散的なテキストをCLIPテキストエンコーダを使用して埋め込みのシーケンスに変換して、これらの埋め込みをトークンのシーケンスとしてU-ViTに入力します。
無条件およびクラス条件付き画像生成
ここで、U-ViTを、以前のU-Netベースの拡散モデルやGenViTと比較します。GenViTは長いスキップ接続を持たず、正規化レイヤーの前に時間を組み込んだ、より小さなViTです。画像の品質を測定するために、FIDスコアを使用しました。
表1に示されるように、U-ViTは無条件のCIFAR10やCelebA 64×64においてU-Netと比較して同等の性能を示し、また、GenViTよりも優れたパフォーマンスを発揮しました。クラス条件付きのImageNet 64×64では、最初に131Mのパラメータを持つU-ViT-M構成を試しました。表1に示されるように、これは100Mのパラメータを持つU-Netを使用したIDDPMの6.92よりも優れた5.85のFIDを得ました。さらなる性能向上のために、287Mのパラメータを持つU-ViT-L構成を採用し、FIDが5.85から4.26に向上しました。
クラス条件付きのImageNet 256×256では、U-ViTは2.29という最高のFIDを達成し、以前の拡散モデルを上回りました。表２は、同じサンプラーを使用したさまざまなサンプリングステップで、U-ViTがLDMを上回ることを示しています。また、U-ViTは、トランスフォーマーをバックボーンとする離散拡散モデルであるVQ-Diffusionを上回りました。同様に、同じパラメータと計算コストを持つUNetと比較しても、U-ViTが優れたパフォーマンスを発揮しています。
クラス条件付きのImageNet 512×512では、U-ViTは画像のピクセルを直接モデル化するADM-Gを上回りました。図4には、ImageNet 256×256および512×512の選択されたサンプル、および他のデータセットのランダムサンプルが示されており、高品質で非常に鮮明な画像だと確認できます。
MS-COCOでテキストから画像生成
ここで、MS-COCOデータセットを使用して、テキストから画像生成のタスクで、U-ViTを評価します。また、U-ViTと同じモデルサイズを持つU-Netを採用した別の潜在拡散モデルも訓練し、U-ViTと比較します。
画像品質を測定するためにFIDスコアを利用します。MS-COCOの検証セットからランダムに30Kのプロンプトを選択し、これらのプロンプトでサンプルを生成してFIDを計算します。表３に示されているように、U-ViTは、生成モデルのトレーニング中に大規模な外部データセットにアクセスする必要がない場合でも、最先端のFIDを達成しています。レイヤーの数を13から17に増やすことで、U-ViT-S（Deep）はさらに優れたFIDを達成できます。
図6は、定性的な比較のために、同じランダムシードを使用してU-NetとU-ViTから生成されたサンプルを示しています。U-ViTはより高品質なサンプルを生成し、画像の内容がテキストとよりよく一致しています。
例えば、「a baseball player swinging a bat at a ball」というテキストが与えられた場合、U-Netはバットやボールを生成しませんが、U-ViTはボールを生成し、U-ViT-S（Deep）はさらにバットを生成しています。これは、U-ViTでのテキストと画像の各層でのより頻繁な相互作用が、U-Netよりもテキストと画像の間で発生しているためだと考えられます。
結論
今回の記事では、拡散モデルを用いた画像生成のためのシンプルで汎用性のあるViTベースのアーキテクチャ、U-ViTを提案します。U-ViTは、すべての入力（時間、条件、ノイズのある画像パッチ）をトークンとして扱い、浅い層と深い層の間に長いスキップ接続を採用しています。U-ViTは、無条件および条件付きの画像生成、テキストから画像への生成などのタスクで評価されました。
U-ViTは、同様のサイズのCNNベースのU-Netと比較して、同等以上の性能を発揮しています。これらの結果から、拡散ベースの画像モデリングでは、長いスキップ接続が重要であり、CNNベースのU-Netのダウン/アップサンプリング演算子が常に必要とされるわけではないことが示唆されます。
U-ViTは、将来の拡散モデルのバックボーンに関する研究に示唆を与え、多様なモダリティを持つ大規模なデータセットにおける生成モデリングに利益をもたらす可能性があります。
この記事に関するカテゴリー
