More than 1 year has passed since last update.

拡散モデルの整理

Posted at 2022-09-02

はじめに

拡散モデルは現在Text-to-imageの分野で主流なモデルである。
ところが初期や最新の拡散モデルに関する論文を一個読んだだけでは途中の拡散モデルの歴史や変化を追えないし、把握できない。そこで拡散モデルにおける変更を自分の理解を分野ごとに分けて整理してみた。
また、多段の拡散モデルの構成に関しては自記事「Diffusion modelは何故1個しかないのか？」も参考。

1.拡散過程

1.1.Denoising Diffusion Probabilistic Models(DDPM)

Probabilistic(確率論的)

1.2.Denoising Diffusion Implicit Models(DDIM)

Implicit(暗黙)
DDPMは左のように徐々にノイズを与えて拡散過程を作る。DDIMは右のように一気にノイズを与えて拡散過程を作るという違いであると自分は理解をしている。

1.3.Noise Schedule

DDPMでは$T=1000,\beta_1=10^{-4},\beta_T=0.02$で拡散過程のノイズを線形に増加させている。
Improved DDPMの論文にノイズを線形に増加させるのではなくcosine状に増加させた結果が示されている。

1.4.Variational Diffusion Models(VDM)

ノイズスケジュールを学習して最適化している。
その時、信号とノイズ比のSNRのlogを取ったのを指標にしている。

SNR(t)=\alpha_t^2/\sigma_t^2

1.5.Distilled model

拡散モデルはサンプリングを行うのに250stepとか1000stepとか何度も拡散モデルにデータを流す必要がある。かといってstep数を減らすと性能が落ちる。このサンプリング回数を減らしつつ精度を維持する課題にDistilled modelがある。

2.拡散モデルの損失

2.1.variational lower bound (VLB)

拡散モデルが$T$段ありそれぞれの入力が$x=x_1,x_2,…,x_T$であるとする。
これに拡散モデルを通した結果$y_{pred}=y_1,y_2,…,y_T$とし、近づける出力を拡散過程の入力を1段ずらした入力値$y_{true}=x_0,x_1,…,x_{T-1}$とする。この分布のKLダイバージェンスを取る。

x=x_1,x_2,…,x_T\\
y_{pred}=y_1,y_2,…,y_T\\
y_{true}=x_0,x_1,…,x_{T-1}\\

L_{vlb}=L_0+L_1+L_2+...+L_T\\
L_t=D_{KL}(y_{pred},y_{true})

2.2.simple

DDPMの論文でLossをシンプルにするのが書かれている。KLダイバージェンスの代わりに$\alpha$が任意の場合のMSEを使うというのが自分の理解。

L_{simple}(\theta)=E_{t,x,\epsilon}[\begin{Vmatrix}\epsilon - \epsilon_\theta(\sqrt{\alpha}x+\sqrt{1-\alpha}\epsilon, t)\end{Vmatrix}^2]\\

L_{simple}=MSE(y_{pred},y_{true})

2.3.hybrid

Improved DDPMの論文にシンプルlossと変分下限lossの両方使うのが示される。ただし結果を見るにNoise Schedule次第だし、必ずしもhybridがもっとも有利というわけではない。

L_{hybrid} = L_{simple} + \lambda L_{vlb}

3.拡散の種類

3.1.Gaussian Noise

画像にガウスノイズを加えていく。拡散モデルならこれが標準。

3.2.Mask

画像を徐々にMaskで隠していく。後に述べるVQ-Diffusionなどで見られる。

3.3.Colorization,JPEG restoration

モノクロの画像を復元する。
Paletteにみられる。画像をガウスノイズではなく何らかの手法によって劣化させていくものと思われる。

3.4.Gaussian Blur

Cold Diffusionの手法にある。
低解像度から徐々に高解像度の画像をつくるならStyleGANと似た傾向があるようにも思える。

3.5.Discrete Denoising Diffusion Probabilistic Models(D3PM)

Uniform、Gaussian、Absorbingの3パターンを考えている。
この内、Uniformが通常のガウスノイズでAbsorbingは[MASK]を掛けるのと等しい。

4.Super-Resolution

拡散モデルで大きな出力を作成するにはテクニックがいる。

4.1.one-stage Diffusion Model

基本的な拡散モデルである。
拡散モデルは低解像度に比べ高解像度の画像を生成するのは一般に苦手である。

4.2.two-stage Diffusion Model

Improved DDPM、ADM、SR3などで見られる。
最初に低解像度の$z$のみを求める拡散モデルを作り、次に低解像度$z$と高解像度$x$を入力とする別の拡散モデルを作る。

4.3.Cascaded Diffusion Models(CDM)

低解像度$z$と高解像度$x$とした時、入力が$(x_i,z)$で出力が$(x_{i-1})$タイプの拡散モデル。
個人的な推測を語ればCDMは解像度の異なる拡散モデルを作るではなく高解像度のUnetと低解像度のUnetの両方を内包しているのかと考えられる。つまり二個の拡散モデルを一個にまとめているのかと思われる。
入力画像の条件によって切り替えられる構造で入力にlowres_cond(低解像の画像)があればタイムステップに余分のembedding成分を足す。

5.Guidance

Sampling(画像生成)時に勾配を足す。

5.1.Classifier Guidance

初出はADMの論文。
画像を入力とする分類器を作り、これに適切な分類yを与え、このモデルの勾配を使用することで画像生成の精度を向上させる。個人的なイメージを語れば理科の実験でミョウバンの結晶を作る時、種結晶を与えてやると大きな結晶を作りやすいというイメージである。ホコリが多かったり種結晶を与えないと無数の小さな結晶が多く発生し、大きな結晶を作るのは困難である。
これは拡散モデルにおける画像生成は完全なガウスノイズから始め、拡散モデルのデノイズによって画像を生成するが、完全なノイズから生成するより、初期段階において何らかの種ノイズを与えた方が画像生成は安定化されるのではと考えられる。

5.2.Classifier free Guidance

Classifier-Free Guidanceの論文。
別名、unconditional guidanceの名で呼ばれる。
さて、上述のClassifier Guidanceには欠点がある。
ガイダンスを作る分類器は汎用の分類器を使えず、入力はノイズ交じりの画像と拡散過程のtimestep。出力はclassとする特殊なモデルを学習する必要があるからである。さて、Classifier-Free Guidanceは分類器の勾配は上述の特殊な分類器が無くても計算できることを示している。

すなわち、勾配を$\nabla\epsilon=(y_{true}-y_{pred})$とすれば勾配は予測出力と正解出力の差の比例である。
拡散モデルの入力を$z_t$、拡散モデルの出力$y_{pred}=\epsilon(z_t,c)$、拡散モデルの正解出力を$y_{true}=z_{t-1}$とすれば、拡散モデルにラベルを与えた時の変化$\nabla\epsilon(z_t,c)=(z_{t-1}-\epsilon(z_t,c))$と拡散モデルにラベルを与えない時の変化$\nabla\epsilon(z_t,none)=(z_{t-1}-\epsilon(z_t,none))$の差分を取ると正解出力は消え$\nabla\epsilon(z_t,c)-\nabla\epsilon(z_t,none)=\epsilon(z_t,c)-\epsilon(z_t,none)$である。こうすると分類器は必要でないし、逆伝播である勾配すら計算する必要がない。拡散モデルの順伝播の差分の計算でClassifier Guidanceの勾配に相当する変化が計算できる。
ただし、この拡散モデルを学習する際にラベルを与えない場合も学習する必要があり、これは拡散モデルに低確率(例えば10%)でラベルを与えないで学習する変更が必要である。
とはいえ分類器を別途学習する手間に比べたらだいぶ楽である。

5.3.ガイダンス重み

Classifier GuidanceやClassifier free Guidanceに掛ける係数$w$。
$w$が小さければ画像の多様性は高いが、画像の質は落ちる。
逆に$w$を大きくすれば画像の質は上がるが、多様性が失われて似た画像ばかり出力する。
この係数はFID性能とIS性能のトレードオフの関係がある。

5.4.Classifier free Guidance in text

Classifier free Guidanceは拡散モデル出力の付与クラスありと付与クラスなしの差分である。
拡散モデルの入力条件が画像クラスではなく、テキスト潜在ベクトルの場合も、拡散モデル出力の付与テキストありと付与テキストなしの差分でClassifier free Guidanceは計算できる。

ところで拡散モデルの入力がもうClassではなくTextなのにClassifier-free(分類器フリー)というのは何か違和感を感じる。

5.5.CLIP

CLIPはファインチューニングが不要な一種の分類器の一種である。
これは、TextEncoderとImageEncoderの対照学習で事前学習をし、例えばラベルが「飛行機」「猫」「犬」…「鳥」ならTextEncoderに「飛行機の写真」、「猫の写真」「犬の写真」…「鳥の写真」と入れた各ベクトルと画像のベクトルの最も近いものを使って分類器とする。
例えば犬と同定出来たあとも「白い犬の写真」「黒い犬の写真」「青い犬の写真」…「茶色い犬の写真」のベクトル距離を比べれば「黒い犬の写真」と同定できるし、「草原にいる黒い犬の写真」「サンゴ礁にいる黒い犬の写真」…「宇宙にいる黒い犬の写真」のベクトル比較すればその気になればいくらでも細かい分類ラベルを生成できる。これは「サンゴ礁にいる黒い犬の写真」の学習データが1例も存在しなくても分類を作成できる。

5.6.CLIP Guidance

GLIDEの論文中に出てくる。
成分としてはCLIPを分類器とみなした場合のClassifier Guidanceである事が分かる。
$\nabla(f(x)\cdot g(c))$
しかし、一般にCLIPは拡散過程のノイズを加えた画像に対して弱い。これはCLIPの学習時にノイズを付与した画像で学習されていないからである。これに対してNoised CLIPはCLIPの学習にノイズを加えて学習させる。
しかし、正直CLIP Guidanceは勾配を使わない(Unguided)よりはましだが、Classifier-Free Guidanceより優れているかと言われれば、そうでもない。
後続のMake-A-Scene、DALL-E2、ImagenでCLIP GuidanceではなくいずれもClassifier free Guidanceが使われる。

5.7.Large Guidance weight

Imagenではガイダンス重みを大きくする。
従来$w=4$程度だったが$w=8$程度まで大きくすることができる。

Sampling(画像生成)時、ガイダンス重みを増加させた方が画像の質が改善される。
しかし、大きければ大きいほど良いかといえばそうではなく、これが大きすぎると却って質が悪化する問題がある。これはガイダンス寄与により拡散モデルの訓練データの分布より外側に行ってしまうと拡散モデル自体のデノイズ性能が悪化してしまうためである。
この対策に閾値を用い、静的閾値はDDPM、動的閾値はImagenで見られる。動的というのは拡散モデルのタイムステップによって閾値が異なるという事であろう。

また、ガイダンス重みを増加させると多様性と忠実度のトレードオフが発生するジレンマが存在するが、ガイダンス重みを大きくした上で、低解像度のモデルでノイズ(aug_level)を加え、超解像モデルではノイズを加えない。すると低解像度で加えたノイズによって破損するので大きなガイダンス重みでも多様性が維持され、超解像(高解像度)では大きなガイダンス重みで理論上は忠実度を上げられる。とはいえ実際には中々複雑である。

6.Latent Diffusion

潜在空間で拡散モデルを考える。

6.1.Latent Diffusion

最初にこのモデルは別途VAEが必要である。
このVAEの潜在変数層で拡散プロセスを通じて拡散モデルを作る。
Sampling時はノイズから拡散モデルを通して潜在変数を作成し、これをVAEのDecoderを通して画像を生成する。

6.2.VQ-Diffusion

VQ-Diffusionでは上記のLatent DiffusionのVAEの代わりにVQ-VAEを使う。
しかし、VQによって離散化(量子化)してる潜在変数に拡散を使ってノイズを足すと意味が急変してしまうので、ガウスノイズではなくマスク([mask])を使って拡散過程を作る。

6.3.CLIP & unCLIP

DALL-E2では「CLIPのimg encoder」の逆の「img decoder」を作り、これを「unCLIP」と呼ぶ。
CLIP & unCLIPの組み合わせはencoder+decoderであるからVAEと構成が近く、そういう意味ではLatent Diffusionの一種といえる。

6.4.prior

DALL-E2では拡散モデルにCLIP text embeddingとCLIP image embeddingの両方を渡すモデルを学習させる実験をしている。
以下の図の例では画像生成時には拡散モデルの入力にimage latentは得られないから
一行目はtext embeddingはtext latentを渡し、image embeddingにはゼロベクトルを渡す。
二行目はtext embeddingはtext latentを渡し、image embeddingにも同じtext latentを渡す。
三行目はtext embeddingはtext latentを渡し、image embeddingに疑似image latentを渡す。
priorモデルは本来の拡散モデルとは別でtext latentを入力として、疑似image latentが出力である。
prior自体がunCLIP+CLIP Image encoderに相当するからか、unCLIPスタックと呼ばれる？

7.自己回帰モデル：Autoregressive(AR) model

Partiは拡散モデルではなく自己回帰モデルである。
DALL-E2のpriorでも自己回帰モデルも検討されたが、こちらでは自己回帰モデルよりも拡散モデルの方が良かった。

CogViewの図が参考になる。

また、Partiでの図は以下である。
いずれにしてもTransfomerがAR modelである。

自分のイメージ図を書くと以下の様になる。
TransformerはSource-Target Attentionがあるので異なる入力を2個持つ。

一般に拡散モデルはUnetだが、自己回帰モデルはTransformerなので(自社開発TPUがある)Googleさんなら多分楽に(楽ではない)、巨大なパラメータを持つモデルを構築可能なのだと思われる。

Stable Diffusionは860M UNetと123M text encoder(これでもGPUメモリ10GB必要)。
Imagenでは300M,500M,1B,2B UnetとT5-small(60M), T5-Base(220M),T5-Large(770M),T5-XL(3B),T5-XXL(11B) text encoderである。
一方、Partiは以下の様に20Bでそもそもモデルサイズが桁違いである。
これはfp32で80GBのモデルサイズで一枚のGPUではload_weightすら無理である。

8.画像参照

8.1.Artistic Images

検索拡張拡散モデル(RDM)というモデルの一種である。
またはKNN-Diffusionも概念的に近い。
詳しくは読めてないが、ArtBenchから類似スタイル画像を検索すれば、類似画像群の画像ベクトル(CLIP image encoderの出力)を用いて拡散モデルにおいて芸術的な画像の生成が比較的容易になるという事なのだろうか？つまり、同様の拡散モデルでは通常text latentのみを拡散モデルに入力しているが、類似スタイル画像におけるimage latentを入力に追加すると芸術スタイルの復元精度が上がるのだと思われる。
しかし、もし自分の推定が正しいと仮定すれば、これは芸術的な出力＝高性能と人間が錯覚する誤謬を利用しているように思える。

Midjourneyでは--stylizeというオプションがあり、デフォルトでも芸術的な出力に寄っている。昨今、Midjourneyがもてはやされているが、芸術的スタイル出力に特化しただけで、このモデルのテキスト理解性能はImagenやPartiよりも必ずしも高くないのではないかと思わなくもない。(とはいえ公開されないモデルより好感度は高いが)
また、外部記憶(database)を使っているという事に対して一般ユーザーがこの拡散モデルが既存の画像から盗作しているという誤解を与えないかも心配である。

8.2.pseudo-word(疑似単語)

サンプル画像から特定の物体を表現する疑似単語を探す。
text embeddingは単語が例えば4000単語ならガウス球上に存在する高々4000通りのベクトルしか表現できない。しかし、ガウス球上のベクトルは4000よりもずっと多い。(ほぼ無限)
このガウス球上を探索すればInput Sampleが言葉で指定しにくい物体であっても近い疑似単語ベクトルを探すことができる(かもしれない)。この疑似単語ベクトルを使う方が言葉で表現しづらいInput Sampleの物体の合成が楽になる。
8.1はImage Encoderで画像ベクトル、8.2はText Encoderで単語ベクトルで探すベクトル自体は違うが、どちらも事前類似サンプル画像群から対象に相当するベクトルを探索するという意味合いでは似ている。

8.3.blended Diffusion

blended DiffusionやRePaintが例。
マスクを使ってマスク部分だけ画像をSamplingする。他は元画像の出力を採用する。
また、マスクのアルファチャンネル(透明度)も調整すればラフ画像を元に生成できる。

Stable Diffusionなどではラフなスケッチ入力画像を元に画像生成(img2img)ができる。