この記事は中国のNLP研究者Jianlin Su氏が運営するブログ「科学空間」で掲載された解説記事の日本語訳です。
原文の掲載日は2022/6/13です。
苏剑林. (Jun. 13, 2022). 《生成扩散模型漫谈（一）：DDPM = 拆楼 + 建楼 》[Blog post]. Retrieved from https://spaces.ac.cn/archives/9119
拡散生成モデル漫談（一）：DDPM＝解体＋建設
生成モデルというと、VAEやGANといった名前をよく耳にするだろう。このブログでも幾度も取り上げている。加えて、フローモデルやVQ-VAEといった、少々ニッチだが人気を集めている選択肢もある。特にVQ-VAEやその亜種であるVQ-GANは、近頃は「画像のトークナイザ」の地位にまで発展しつつあり、NLPで使われる様々な学習手法が応用されている。
これらに加えて、本来更にニッチな手法だった「拡散モデル(Diffusion Models)」が、生成モデル分野で突如台頭しつつある。現在最も先進的なtext-to-imageの生成手法である、OpenAIのDALL・E2とGoogleのImagenは、いずれも拡散モデルに基づいて作られたものだ。
そこでこの記事から、新たな試みを始めてみたいと思う。ここ数年における拡散生成モデルの進展を、じっくり紹介していきたい。拡散生成モデルは数学的に複雑だと思われているらしい。何ならVAEやGANよりも難解だというが、本当にそうだろうか？拡散モデルは「平易な言葉」で理解することはできないだろうか？ぜひ刮目いただきたい。
新たな起点
このブログでは以前、「エネルギー視点からみたGAN（三）：生成モデル＝エネルギーベースモデル」「ノイズ除去オートエンコーダーから生成モデルへ」と題した記事でも拡散モデルを簡単に紹介したことがあった。拡散モデルというと、多くの文章はエネルギーベースモデル(EBM)、スコアマッチング、Langevin方程式といった概念に言及するだろう。簡単に言えば、スコアマッチングなどの技術を使ってEBMを訓練し、Langevin方程式でサンプリングを行う、という流れだ。
これは理論的に成熟した方針である。原理上、この方法に従えばいかなる連続的な対象（音声、画像など）の生成やサンプリングも可能だ。しかし現実的には、エネルギー関数の訓練はとても困難である。特にデータの次元数が大きい場合（高解像度の画像など）、完全なエネルギー関数を得ることはなかなか難しい。一方で、Langevin方程式によるサンプリングにも大きな不確実性があり、往々にしてノイジーな結果を出力してしまう。ゆえにこの伝統的な方針に基づいた拡散モデルは、長い間にわたり低解像度な画像でしか実験を行われて来なかった。
現在起きている拡散生成モデルのブレイクは、2020年に提案されたDDPM（Denoisiong Diffusion Probabilistic Model）が起点になっている。「拡散モデル」という名前こそ付いているものの、実際にはサンプリング過程が形式的に若干類似性がある以外は、伝統的なLangevinサンプリングに基づく拡散モデルとは全く別物だと言っていい。完全に新たな起点である。
DDPMは「グラデーションモデル」と呼んだ方がより正確で、拡散モデルという名前は逆に誤解を生みやすい。伝統的な拡散モデルにおけるEBM・スコアマッチング・Langevin方程式といった概念は、実はDDPMおよびその後続亜種とはあまり関係がない。興味深いことに、実はDDPMの数学的な枠組みはICML2015の論文「Deep Unsupervised Leraning using Nonequilibrium Thermodynamics」で既に完成されている。DDPMは初めてその手法で高解像度画像生成を成功させたことで、後のブレイクに繋がったのである。あるモデルの誕生や流行には、タイミングや機運も大事であることが、ここでも示されている。
解体と建設
多くの文章はDDPMを紹介する際、初っ端から分布変換を導入したり、変分推定を始めたりして、一気に数学記号を羅列した結果多くの人の出鼻をくじいてしまう（もちろんこの手の紹介から、改めてDDPMは本質的に拡散モデルではなくVAEであることが分かる）。そこに伝統的な拡散モデルに対する印象も加わることで、「高い数学力が求められる」という錯覚が生まれたのであろう。実はDDPMも、もっと平易な喩えで理解することができ、よく「偽物鑑定」に喩えられるGANよりも難しいということはない。
まずは、GANみたいに生成モデルをイメージしてみよう。このモデルは、ランダムなノイズ$z$をデータサンプル$x$に変換する過程である。
\frac{ランダムノイズz \rightarrow^{変換}サンプルデータx}{建築資材\rightarrow^{建設}高層ビル}\qquad (1)
この過程は「建設工事」と想像することができる。ランダムノイズ$z$はレンガやコンクリートなどの建築資材で、データサンプル$x$は高層ビルだとすると、生成モデルは即ち建築資材から高層ビルを建てる建設業者である。
この建設過程をモデル化するのは当然難しい。だからこそ、これまで多くの生成モデルが研究されてきたのだ。ただ、何事も「作るのは大変、壊すのは一瞬」と言うように、建てることはできなくても、壊すことくらいはできるだろう。そこで、逆に高層ビルを徐々にレンガとコンクリートに解体する作業を考えてみよう。$x_0$を完成した高層ビル（データサンプル）、$x_T$を解体されたレンガ（ランダムノイズ）とし、「解体工事」には$T$ステップ必要だとすると、この過程はこう書くことができる。
x=x_0\rightarrow x_1 \rightarrow x_2 \rightarrow \cdots \rightarrow x_{T-1} \rightarrow x_T = z \qquad (2)
高層ビルの建設が難しいのは、原材料$x_T$から高層ビル$x_0$までの変化が大きすぎて、$x_T$から$x_0$に変化する仕組みが理解し難いからである。しかし、ビルを解体する中間過程$x_1, x_2,\cdots,x_T$が分かれば、$x_{t-1}\rightarrow x_t$がビル解体作業の一歩であるなら、その逆である$x_t \rightarrow x_{t-1}$がビル建設作業の一歩と見なせるのではないだろうか。もしも両者の変換関係$x_{t-1}=\mu(x_t)$を知ることができれば、$x_T$から始めて、$x_{T-1}=\mu(x_T), x_{T-2}=\mu(x_{T-1}),\cdots$を繰り返し実行すれば、最終的に高層ビル$x_0$を建てることができるはずだ。
解体の仕方
飯を一口ずつ食べるように、高層ビルも一歩ずつ建てなければならない。DDPMが生成モデルを定義する過程は、まさに先ほどの「解体-建設」の喩えと完全に一致する。すなわち、まずはデータサンプルが徐々にランダムノイズに変化する過程を構築してから、その逆変換を考え、逆変換を繰り返し実行することでデータサンプルの生成を実現するわけだ。冒頭でDDPMを「拡散モデル」ではなく「グラデーションモデル」と呼ぶべきだと言ったのも、それが理由である。
具体的に、DDPMは「解体工事」の過程を以下のようにモデル化した。
x_t=\alpha_tx_{t-1} + \beta_t\varepsilon_t, \quad \varepsilon_t\sim\mathcal N(0,I) \qquad (3)
ここで$\alpha_t, \beta_t>0$かつ$\alpha_t^2+\beta_t^2=1$である。$\beta_t$は0にとても近い数値で、毎回の「解体工事」で高層ビルを壊す程度を表している。ノイズ$\varepsilon_t$は元信号への破壊を表している。あるいは高層ビルの「建築資材」と考えることもできるだろう。つまり毎回の「解体工事」では、$x_t$を「$\alpha_tx_{t-1}$の高層ビル+$\beta_t\varepsilon_t$の建築資材」に分解したとみなすわけだ。
本解説における$\alpha_t$と$\beta_t$の定義は、DDPMの元論文と異なる。
この「解体工事」を繰り返し実行することで、以下の式が得られる。
\begin{split}
x_t&=\alpha_tx_{t-1}+\beta_t\varepsilon_t \\
&=\alpha_t(\alpha_{t-1}x_{t-2}+\beta_{t-1}\varepsilon_{t-1})+\beta_t\varepsilon_t \\
&=\cdots \\
&=(\alpha_t\cdots\alpha_1)x_0+\underbrace{(\alpha_t\cdots\alpha_2)\beta_1\varepsilon_1+(\alpha_t\cdots\alpha_3)\beta_2\varepsilon_2+\cdots+\alpha_t\beta_{t-1}\varepsilon_{t-1}+\beta_t\varepsilon_t}_{互いに独立したガウスノイズの和} \quad (4)
\end{split}
先ほどの説明で、なぜ係数は$\alpha_t^2+\beta_t^2=1$を満たさねばならないのか気になった読者もいたかもしれないが、ここでその問いに答えることができる。式$(4)$の括弧部分は、互いに独立したガウスノイズの和である。ノイズの平均値は0で、分散はそれぞれ$(\alpha_t\cdots\alpha_2)^2\beta_1^2$、$(\alpha_t\cdots\alpha_3)^2\beta_2^2$、$\cdots$、$\alpha_t^2\beta_{t-1}^2$、$\beta_t^2$である。
ここで確率論の知識である「正規分布の再生性」を使えば、これらの独立したガウスノイズが従う分布は、平均値が0、分散が$(\alpha_t\cdots\alpha_2)^2\beta_1^2+(\alpha_t\cdots\alpha_3)^2\beta_2^2+\cdots+\alpha_t^2\beta_{t-1}^2+\beta_t^2$の正規分布であることが分かる。最後に$\alpha_t^2+\beta_t^2=1$の条件を加えると、式$(4)$の各係数の二乗和が1になる。つまり、
(\alpha_t\cdots\alpha_1)^2+(\alpha_t\cdots\alpha_2)^2\beta_1^2+(\alpha_t\cdots\alpha_3)^2\beta_2^2+\cdots+\alpha_t^2\beta_{t-1}^2+\beta_t^2=1\qquad(5)
すると式$(4)$は
x_t=\underbrace{(\alpha_t\cdots\alpha_1)}_{\overline\alpha_tと記す}x_0+\underbrace{\sqrt{1-(\alpha_t\cdots\alpha_1)^2}}_{\overline\beta_tと記す}\overline\varepsilon_t, \quad \overline\varepsilon_t\sim\mathcal{N}(0,I)\qquad\qquad(6)
となる。こうすることで、$x_t$の計算が大きく単純化される。
ちなみに、DDPMでは適切な$\alpha_t$を選択することで、$\overline\alpha_T\approx0$となるようにしている。これはつまり、$T$ステップの解体作業を行ったあと、高層ビルはほとんど残っておらず、ほぼすべて建築資材$\varepsilon$に変換されたことを表している。
本解説における$\overline\alpha_t$の定義は、DDPMの元論文と異なる。
建設の仕方
$x_{t-1}\rightarrow x_t$という「解体作業」をたくさんこなすことで、大量のペアデータ$(x_{t-1}, x_t)$を得ることができた。「建設作業」とは即ち、これらのペアデータから$x_t\rightarrow x_{t-1}$のモデルを学習することである。このモデルを$\mu(x_t)$とすると、両者のユークリッド距離を最小化するという目標関数を容易に思い付くだろう。
||x_{t-1}-\mu(x_t)||^2 \qquad(7)
実はこの時点で、既に最終的なDDPMの手法にかなり近づいている。続いてこの過程をより丁寧に考えてみよう。まず、「解体作業」の式は$x_{t-1}=\frac{1}{\alpha_t}(x_t-\beta_t \varepsilon_t)$に変形させることができる。であれば、「建設作業」のモデル$\mu(x_t)$も以下のように書けるのではないだろうか。
\mu(x_t)=\frac{1}{\alpha_t}(x_t-\beta_t\epsilon_\theta(x_t,t)) \qquad(8)
ここで$\theta$は学習可能なパラメータである。これを先ほどの目標関数に代入すると、
||x_{t-1}-\mu(x_t)||^2=\frac{\beta_t^2}{\alpha_t^2}||\varepsilon_t-\epsilon_\theta(x_t,t)||^2\qquad(9)
$\frac{\beta_t^2}{\alpha_t^2}$は損失の重みで、ひとまず放っておく。最後に式$(3)$と式$(6)$が規定した$x_t$の式
x_t=\alpha_tx_{t-1}+\beta_t\varepsilon_t=\alpha_t(\overline\alpha_{t-1}x_0+\overline\beta_{t-1}\overline\varepsilon_{t-1})+\beta_t\varepsilon_t=\overline\alpha_tx_0+\alpha_t\overline\beta_{t-1}\overline\varepsilon_{t-1}+\beta_t\varepsilon_t\quad(10)
これを代入すると、損失関数は以下の形式になる。
||\varepsilon_t-\epsilon_\theta(\overline\alpha_tx_0+\alpha_t\overline\beta_{t-1}\overline\varepsilon_{t-1}+\beta_t\varepsilon_t, t)||^2\qquad(11)
なぜわざわざ式$(3)$まで戻って$x_t$を求め直したのか疑問に思った読者もいるかもしれない。直接式$(6)$の$x_t$を使えばいいのではないだろうか？答えは否である。ここでは先に$\varepsilon_t$をサンプリングしており、$\overline\varepsilon_t$は$\varepsilon_t$と独立ではないことに注意する必要がある。$\varepsilon_t$が決まった後、$\overline\varepsilon_t$を独立的にサンプリングしてはいけないのだ。
分散を抑える
原理的には、もう損失関数$(11)$でDDPMの学習ができるようになっているが、実際には損失関数の分散が大きくなり、収束が遅くなってしまうリスクがある。この現象を理解するのは難しくない。式$(11)$を観察すると分かる通り、この式は4つの確率変数を含んでいる。
- 学習データセットからランダムに選ばれた$x_0$、
- 標準正規分布$\mathcal N(0,I)$からサンプリングされた$\overline\varepsilon_{t-1},\varepsilon$（2つの異なるサンプリング結果）、
- $1\sim T$からランダムに選ばれた$t$
サンプリングが必要な確率変数が多いほど、損失関数の正確な近似は難しくなる。すると毎回計算される損失関数の振れ幅（分散）が大きくなってしまうのだ。幸い、とある積分テクニックを使い$\overline\varepsilon_{t-1},\varepsilon$を一つの正規分布確率変数に合体させることで、分散が大きくなる問題を改善することができる。
この積分テクニックは正直少々トリッキーではあるが、それほど複雑ではない。正規分布の再生性により、$\alpha_t\overline\beta_{t-1}\overline\varepsilon_{t-1}+\beta_t\varepsilon_t$は単一の確率変数$\overline\beta_t\varepsilon|\varepsilon\sim\mathcal{N}(0,I)$に等しい。同じく$\beta_t\overline\varepsilon_{t-1}-\alpha_t\overline\beta_{t-1}\varepsilon_t$も単一の確率変数$\overline\beta_t\omega|\omega\sim\mathcal{N}(0,I)$に等しく、かつ$\mathbb{E}[\varepsilon\omega^\mathsf{T}]=0$が確認できるので、これらは互いに独立した確率変数である。
続いて、$\varepsilon_t$を逆に$\varepsilon$と$\omega$で書き直してみる。
\varepsilon_t=\frac{(\beta_t\varepsilon-\alpha_t\overline\beta_{t-1}\omega)\overline\beta_t}
{\beta_t^2+\alpha_t^2\overline\beta_{t-1}^2}=\frac{\beta_t\varepsilon-\alpha_t\overline\beta_{t-1}\omega}{\overline\beta_t}\qquad (12)
式$(11)$に代入すると、
\begin{align}
&\mathbb{E}_{\overline\varepsilon_{t-1},\varepsilon_t\sim\mathcal{N}(0,I)}[||\varepsilon_t-\epsilon_\theta(\overline\alpha_tx_0+\alpha_t\overline\beta_{t-1}\overline\varepsilon_{t-1}+\beta_t\varepsilon_t,t))||^2] \\
=&\mathbb{E}_{\omega,\varepsilon\sim\mathcal{N}(0,I)}[||\frac{\beta_t\varepsilon-\alpha_t\overline\beta_{t-1}\omega}{\overline\beta_t}-\epsilon_\theta(\overline\alpha_tx_0+\overline\beta_t\varepsilon, t)||^2]
\qquad\qquad(13)
\end{align}
ここでは$(11)$に代入するついでに$(11)$の期待値計算に転向している。
損失関数は$\omega$に関してはただの2次関数なので、$\omega$の部分は直接計算して期待値の外に出すことができる。するとこう書ける。
\frac{\beta_t^2}{\overline\beta_t^2}\mathbb{E}_{\varepsilon\sim\mathcal{N}(0,I)}[||\varepsilon-\frac{\overline\beta_t}{\beta_t}\epsilon_\theta(\overline\alpha_tx_0+\overline\beta_t\varepsilon, t)||^2]+定数\qquad(14)
定数と損失関数の重みを無視すれば、最終的なDDPMの損失関数になる。
||\varepsilon-\frac{\overline\beta_t}{\beta_t}\epsilon_\theta(\overline\alpha_tx_0+\overline\beta_t\varepsilon, t)||^2\qquad(15)
DDPM論文の$\epsilon_\theta$がここでの$\frac{\overline\beta_t}{\beta_t}\epsilon_\theta$にあたる。
再帰的な生成
ここまで、DDPMの学習過程を全体的に整理することができた。色々書いたし、超簡単だとまでは言えないが、特に難しい所もほとんど無かったはずだ。伝統的なエネルギー関数やスコアマッチングみたいな道具は使っていないし、変分推定の知識すら触れていない。単純に「解体-建設」のたとえ話と基本的な確率論の知識だけで、まったく同様の結論に導くことができた。だから、DDPMに代表される新興の拡散生成モデルは、実は多くの読者が想像するような複雑なものではなく、我々が「解体-建設」をしながら新たな知識を学ぶ過程を分かりやすくモデル化したに過ぎないのだ。
モデル学習が完了した後、ランダムなノイズ$x_T\sim\mathcal{N}(0,I)$から式$(8)$を$T$回繰り返して生成を行う。
x_{t-1}=\frac{1}{\alpha_t}(x_t-\beta_t\epsilon_\theta(x_t,t))\qquad(16)
この生成方式は自己回帰モデルにおけるGreedy Search方式にあたる。ランダム性を導入したい場合は、ノイズ項を導入すればいい。
x_{t-1}=\frac{1}{\alpha_t}(x_t-\beta_t\epsilon_\theta(x_t,t))+\sigma_tz,\quad z\sim\mathcal{N}(0,I)\qquad(17)
ここでは$\sigma_t=\beta_t$とする、つまり双方向の分散を一致させることが普通である。このサンプリング方法が伝統的な拡散モデルにおけるLangevinサンプリングと異なるのは、DDPMは毎回ランダムなノイズから出発し、$T$回繰り返すことでサンプルを得るのに対して、Langevinサンプルは任意の点から出発し、無限回繰り返す。理論上、無限回繰り返した後、すべてのデータサンプルが生成される。なので両者は形式的に類似性がある以外は、本質的に全く異なるモデルである。
以上の生成過程を見ると、実はSeq2Seqモデルのデコード過程と同じような、順序的な自己回帰的生成過程であることを感じ取っただろう。ゆえに生成速度はネックである。DDPMは$T=1000$としているので、画像を一つ生成するたびに$\epsilon_\theta(x_t,t)$を1000回実行することになる。DDPMの最大の弱点はまさにサンプリングの遅さで、後続の研究はこれを改善することを目指したものが多い。
「画像生成＋自己回帰モデル＋遅い」と聞いて、PixelRNNやPixelCNNといったモデルを思い浮かべた読者もいるかもしれない。これらの手法は画像生成を言語モデルのタスクに置き換え、同じく再帰的にサンプリングや生成を行い、また同じく遅かった。ではDDPMのような自己回帰的な生成手法は、PixelRNN/PixelCNNの自己回帰的生成と比べて、どこが違うのだろうか？なぜPixelRNN/PixelCNNは流行らず、DDPMがその役を担うことになったのだろうか？
PixelRNN/PixelCNNを理解している人なら分かると思うが、この類の生成モデルはピクセルごとに画像を生成するものである。自己回帰モデルは順序的な過程なので、つまり生成する画像のピクセルの「順番」を予め決めておかねばならず、最終的な生成結果もその順番に大きく左右される。しかし「順番」は人間が経験的に設計することしかできず（このような経験的に設計された機構はInductive Biasと呼ばれる）、当面は理論的な最適解を見つけられていない。言い換えれば、PixelRNN/PixelCNNの生成結果はInductive Biasに大きく左右されてしまうのだ。
しかしDDPMは違う。DDPMは「解体工事」の要領で新しい自己回帰の方向を定義した。この自己回帰はすべてのピクセルにとって平等なものなので、Inductive Biasの影響が軽減され、性能が向上した。加えて、DDPM生成の繰り返し回数は固定値の$T$である一方、PixelRNN/PixelCNNは画像の解像度（幅×高×チャンネル数）に等しいので、DDPMが高解像度画像を生成する速度はPixelRNN/PixelCNNより遥かに速い。
ハイパラ設定
ハイパラの設定についても議論してみたいと思う。
DDPMでは$T=1000$としている。多くの読者にとって、思ったよりも大きな数値に感じるだろう。なぜこんなに大きな$T$を設定するのだろうか。一方で、$\alpha_t$に関しても、論文の設定をこの記事の記号に変換すると、
\alpha_t = \sqrt{1-\frac{0.02t}{T}}
となる。これは単調減少する関数だが、なぜ単調減少するような$\alpha_t$を選ぶのだろうか？
実はこの二つの問いは似たような答えを持っており、実際のデータ形式が背景にある。我々は簡単のために、ユークリッド距離$(7)$を損失関数に用いた。以前から画像生成をやったことがある読者なら分かると思うが、ユークリッド距離は画像のリアルさを測るのに適した基準ではない。VAEの学習でユークリッド距離を使うと、ぼやけた画像を出力するようになってしまう。入力と出力画像が元々とても近い場合に限り、ユークリッド距離を使っても比較的高品質な結果を得られることが分かっている。大きな$T$を使うのは、まさに入力と出力を可能な限り近くなるようにして、ユークリッド距離がもたらす「ぼやけ問題」を抑えるためである。
単調減少する$\alpha_t$を選ぶことも似たような考慮から来ている。$t$が比較的小さい時、$x_t$はまだ本物の画像に近い状態なので、$x_{t-1}$と$x_t$の差を抑えて、ユークリッド距離が使いやすいようにしたい。よって大きめな$\alpha_t$を使いたい。一方$t$が大きい時、$x_t$はただのノイズに近い状態なので、$x_{t-1}$と$x_t$の差は少し大きくてもかまわない、すなわち$\alpha_t$を小さくしてもいいのだ。
では、ずっと大きな$\alpha_t$を使えばいいのではないだろうか？できなくはないが、$T$を更に大きくしなければならない。$(6)$を導いた際、$\overline\alpha_T\approx0$にすべきだと言ったのを思い出してほしい。この$\overline\alpha_T$の値は近似的にこう求めることができる。
\log\overline\alpha_T=\sum_{t=1}^T\log\alpha_t=\frac{1}{2}\sum_{t=1}^T\log(1-\frac{0.02t}{T})<\frac{1}{2}\sum_{t=1}^T(1-\frac{0.02t}{T})=-0.005(T+1)\quad(19)
$T=1000$を代入すると、$\overline\alpha_T\approx e^{-5}$くらいになる。これはちょうど$\approx 0$を満たす水準である。もし最初から最後まで大きな$\alpha_t$を使うと、更に大きな$T$でないと$\overline\alpha_T\approx0$が満たせなくなる。
最後に、「建設」モデルにおける$\epsilon_\theta(\overline\alpha_tx_0+\overline\beta_t\varepsilon,t)$では、明示的に$t$を入力として記していることに留意してほしい。原理上、異なる$t$では異なるレベルの対象を処理しているので、異なるモデルを使わなければならない、つまり$T$個の異なるモデルを用意しなければならない。ここではすべてのモデルにパラメーターを共有させ、$t$を条件として入力することで要件を満たしている。DDPM論文の付録によれば、$t$は本ブログの「Transformer進化の道筋：1、Sinusoidal位置エンコーディングの根源を遡る」で紹介した位置エンコーディングに変換した後、直接差分モジュールに加算されているらしい。
まとめ
本記事では「解体-建設」の喩えから最新の拡散生成モデルDDPMを紹介した。この視点から、比較的平易な説明と最小限の数学的推論を通じて、元の論文と一致する結論を導くことができた。DDPMもGANのような分かりやすい喩えで説明することができ、VAEの「変分推定」、GANの「確率分布距離」や「最適輸送」といった概念に頼る必要もない。そう考えると、DDPMはある意味VAEやGANよりも簡単だとすら言えるかもしれない。