概要
画像生成AIの1つ、Stable Diffusionなどの拡散モデルの基礎となった研究。
Jonathan Ho, Ajay Jain, Pieter Abbeel「Denoising Diffusion Probabilistic Models」2020, arxiv
筆頭著者のJonathan Hoは後にIdeogramにも携わっている。
論文では、「ノイズ拡散のシミュレーション」という当時の研究者にとっても新奇なアイデアは控えめに書かれており、
当時よく知られていた画像生成AI(VAE)のフレームワークに準じて説明されている。
当時の世情
当時の主要な画像生成AIとして大きく2つの流れがあった。
拡散モデルの先行研究としてはYang Song「NCSN」やJascha Sohl-Dickstein「Deep Unsupervised Learning using Nonequilibrium Thermodynamics」などがあったが、
時期的な話もあるがそれらはDDPM論文と異なりVAEのような書き方はされていない。
VAEのフレームワークを使った書き方は功を奏し、多くのフォロワーが表れてSDの成功につながることとなった。
論文について
まず前提として、Stable Diffusion以前のVAEはSDのそれと用途がかなり異なり、乱数から直接画像を生成するAIとして研究されていた。
もともとのVAEは、乱数を潜在zとして扱い、どんなzを入力されてもデータセットの雰囲気を保った画像を出力するように訓練されるAIである。
2 Background
いきなり「latent variable(潜在変数)」と書かれているが、DDPM自体は他のVAEを使わないのでSDの潜在とは性質が異なる。
当時のVAEでは潜在は「画像を生成するための乱数入力値」だったのでそのイメージ。
現在のSDでいうと、t2iでシード値からサンプリングされる初期ノイズとか、あるいは生成中の中間ステップにおける潜在などに相当する。
さて、拡散モデルは訓練時に画像がノイズ化されるまでの過程を学習させるわけだが、元のDDPM論文では「forward process or diffusion process」と書かれている。
日本語ではこの過程は「順拡散過程」などと呼ばれているので拡散モデル特有の概念と思ってしまうが、
VAEのエンコーダによるエンコード処理なども英語では「forward pass」というので、英語で書けば当時の英語圏の研究者にもなじみのある言葉で引っかかるところはない。
もともとのVAEではエンコーダの主な目的は「学習時に教師画像を一貫性のある形で潜在表現にマッピングすること」だったので、DDPMの順拡散過程もそのアナロジーで説明されているようだ。
この拡散過程の定式化についても迷惑メールフィルタなどでよく使われているベイズの定理で説明されており、
拡散モデル特有のSDEだのランジュバン動力学だのといった呪文も出てこない。
VAEのゴールは「あらゆる入力潜在について教師画像に近い画像を出力する確率を上げること」だがこの確率は計算不可能。
なので、ELBO(変分下限の一種)という「この数字を最適化すればその確率も上がるだろう」という数字を代わりに計算してそれを最適化する。
変分下限は機械学習や計算機工学ではそれなりにお目にかかる概念で、その筋の人ならピンと来る。
このELBOはVAEやDDPMでは「KL正則化項」+「再構成誤差(教師画像と予測値との差に相当)」になる。
DDPMではKL正則化項についてさらに「ノイズ分散に依存する項」と「平均値に依存する項」に分けられる。
再構成誤差についてp(x0|x1)と書かれているが、このx1はFlow Matchingのそれと異なり最後の1ステップの潜在を精確に教師画像に復元する能力を表している。
3.1 Forward process and LT
ノイズ分散はVAEではエンコーダが潜在空間を整理するために使うのだが、
DDPMでは「ノイズ分散をモデル外部から与えるので、ELBOのうち『ノイズ分散に依存する項』は定数になるから最適化しなくて良い」としている。
もともとVAEでも誤差の分散をエンコーダに計算させるが誤差自体はそれを参照してモデル外部で計算するという「再パラメータ化トリック」を学習のために使っており、
DPPMもノイズ分散まで固定してしまうことを再パラメータ化と書いている。
ただし拡散モデルでは学習時にエンコーダを使わず、ノイズを簡単な処理で画像に合成して入力値とする。
なおDDPMは生成時の逆拡散過程においてそのステップでのノイズ分散を再現するためにノイズをまきなおすが、
VAEの再パラメータ化トリックと類似した形に見せることで読者に受け入れやすい形になっている。
3.2 Forward process and L1:T-1
KL正則化項のうち「平均値に依存する項」についての整理だが、最終的にノイズの予測誤差になる。
ただしこの段階ではステップtに依存する複雑な係数がかかっている。
係数は βt^2 / {2σt^2・αt(1-累積αt)} だが、DDPMはβt=σt^2、αt=1-βtとしているので βt / 2αt(1-累積αt) = βt / 2(1-βt)(1-累積αt) となる。
DDPMではβtは10^-4≦βt≦0.02(4 Experiments)、累積αtは0<累積αt≦1の累積信号残存率。
なので、画像があまりノイズ化されてないうちは分母が小さく勾配が巨大になるため強く学習される計算になる。
それだと勾配が不安定で学習が安定しない。
3.4 Simplified training objective
性能評価値IS、FIDを出したうえで、ELBOをまともに計算するより係数を固定して教師画像の再構成誤差もノイズ予測誤差で代用してしまう簡略化版 Lsimple で学習したほうが品質がよかったとしている。
DDPM論文では元のELBOだとノイズ量が低すぎるステップについて過剰に学習してしまうため、もっと予測の難しいステップに重みをつけたほうが良いのだろうと推定している。
FID3.17というのは当時の画像生成AIの王者であったStyleGAN2の2.67に匹敵する品質、だがCIFAR10データセットは32x32の低解像度なので実用性はまだこれからといったところ。
その後の別の研究でLsimpleは正しく、
ノイズeはxtが教師画像のいずれかである確率についての対数確率密度関数の勾配 -∇xt log pt(xt) に比例することが分かっている。