[2006.11239] Denoising Diffusion Probabilistic Models [paper-reading]

[yoheikikuta]

論文リンク

https://arxiv.org/abs/2006.11239

公開日（yyyy/mm/dd）

2020/06/19

概要

新しいタイプの画像生成モデルである denosing diffusion probabilistic models を提案した論文。
diffusion model は画像にノイズを載せていく forward process とそれを denoise していく reverse process からなり、それらを対応づけながら解析的に扱いやすい形で variational bound を用いた objective function を定式化。gaussian distribution と markov process をベースにしており解析的に綺麗に扱える部分が多く、さらにモデルとしての表現力も高い。
実験的にも既存手法と同等以上の性能を発揮することを示した。
Langevin dynamics や denoising score matching との類似性を示しており、今後より研究が広がっていきそうな題材である（この論文を読んだタイミングではむしろ色々研究が広まった後であり、その辺を理解していくためにまずはこれを読んでみることにしたという次第）。

公式実装: https://github.com/hojonathanho/diffusion

[yoheikikuta]

diffusion process で生成モデル！みたいな話をよく聞くようになったが、内容を全然知らないので一つ論文を読んでみることにした。
最初にどれを読むべきか分かっていないが、結構引用されているこの論文を選んでみる。

[yoheikikuta]

この論文は diffusion probabilistic model https://arxiv.org/abs/1503.03585 の発展である。
diffusion probabilistic model という名前は元の論文では使われてなくて、非平衡熱力学を使った unsupervised learning と言っている。

非平衡熱力学というとなんか凄そうな話に聞こえるが、実態は画像にステップごとに gaussian noise をのせていくものを拡散過程と呼び、最終的に平行状態（標準正規分布からサンプルした noise ）になるところを持って非平衡熱力学から inspire されたと言ってるのみ。それくらいなので別に非平衡熱力学とかを殊更に意識する必要はない。

この拡散過程を forward process と呼び、対応づける形で reverse process を学習することで、表現力と tractability を兼ね備えたモデルが作れるというのが元の論文の主張である。

この仕組みで high quality な画像生成をやったという話はこれまでなかったが、それをやったというのがこの論文の主張（元論文は CIFAR10 とかの結果で今から見るとかなり primitive な感じ）。
それと diffusion model のある parametrization がある種の denoising score matching と等価であると示した。

この辺の意味するところを読み解いていく。

[yoheikikuta]

この論文の基本となる diffusion process についてみていく。
ある画像に noise を乗せていくのが forward process (= diffusion process) であり、逆に noise を除去して復元するのが reverse process である。

reverse process は Markov chain で定義されていて、標準正規分布（これからサンプルしたものを initial の noise とする）から画像を生成していくものとなっている。これは $ t → t-1 $ の transition を gaussian で表現している。

forward process も同様に gaussian で transition する Markov chain で定義している。
$ β_t $ は拡散度合いを示すパラメタであり、variational bound と reparametrizatoin trick を使って学習するものとしたり、シンプルに手で与えるパラメタとしたりする。

どちらも gaussian で定義すると言ってえいやと与えているが、ここで気になるのは reverse process（最終的にはこれで画像生成をする）の表現力が十分なものであるか？ということである。この表現力が十分でないと良い画像生成ができないわけだが、ここでは何も考えずに gaussian で transition するものとして与えているだけである。

この点に関して上で挙げた論文で既に言及されている。
上で挙げた論文で参照している https://digitalassets.lib.berkeley.edu/math/ucb/text/math_s1_article-21.pdf によって、process を連続極限を取って $ β_t $ が小さいとした近似において、どちらの process も同じ関数系を取ることが示されている。
（この論文はちゃんと読んでないが、一旦は信じることにする）

ステップが細かく（$ T $ が十分大きく）て $ β_t $ が小さければ、という前提の下に成り立っている議論であることを頭に入れておく必要がある。

[yoheikikuta]

これを学習するにはよく使う negative log likelihood の variational bound を使えばよい。

ここで一つ注目すべき性質は、forward process q において、解析的に取り扱いがしやすい gaussian 関数を用いているため、任意の timestep $ t $ において $ x_t $ が closed form で sampling できるという点である。
これは以下のように計算できることからわかる。この性質は後で学習のときに任意の timestep を取り出して objective function を計算できるというところと関わってくる。

[yoheikikuta]

導出した $ L $ を以下のように KL divergence で書き直すとより分散を減らすことができる。
ここの計算は論文の appendix A. で丁寧に書いてあるのでここでは省略。

これは KL divergence を使って $ p_θ (x_t-1 | x_t) $ と forward process posteriors $ q $ を直接比較することを可能としている。
forward process $ q $ は $ t-1 → t $ と timestep を刻む方向では tractable ではあるけど、この式の 2 項目の $ q $ はその形になっていない。これは以下のように Bayes' の定理を使って少し計算すると tractable な形に書けることが分かる。

（当然）これも gaussian となるので、objective function たる $ L $ は全て gaussian 間の比較となり、分散が大きくなるモンテカルロではなくて closed form で Rao-Blackwellized 的に計算できるようになる。
この Rao-Blackwellized というのは聞いたことないけど、particle filter とかの文脈でよく出てくるもののようだ。いまは解析的に表現できていることが理解できていればいいので、とりあえず深追いしないでおく。

ここまでが diffusion model の基本的な話で、この論文の新しい発見は含まれてない。
この部分が理解したかったことなので、計算は少し真面目に追って理解できたので満足。gaussian は色々解析的に計算できるので気持ちがいいね。

[yoheikikuta]

diffusion model は latent variable model の一クラスとなるが、実装の自由度は大きい。

• forward process q の分散 β_t
• reverse process p_θ のモデルアーキテクチャと gaussian distribution parametrization

どういう方針で進めるかを objective function (5) 式の各項を詳細に見ていくことで理解する。

forward process q と L_T term

これは簡単のため $ β_t $ が学習可能ということを無視して constant として手で与える。
これによって $ q $ は学習可能なパラメタを持たないことになるので、$ L_T $ term は学習においては constant となる。

[yoheikikuta]

reverse process p_θ と L_{1:T-1} term

$ p_θ (x_t-1 | x_t) = N(x_t-1 ; μ_θ (x_t, t), \sum_θ (x_t, t)) $ におけるパラメタをモデル化する。

まず、分散は学習するものではなく時間依存の定数 $ \sum_θ (x_t, t) = σ^2_t I $ とする。
これは説明はそんなになされてないが、$ x_t → x_t-1 $ でノイズを除去して画像生成をしていくときに、ピクセル毎に割り当てられる平均（ここはちょっと書き方が難しいけど、ピクセル値の平均ではなくて、あくまで $ μ_θ (x_t, t) $ であり、例えばモデルを解してピクセル一部/全体の特徴量を踏まえた上でこのピクセル値はつぎは平均的にはこのピクセル値を取った方が画像生成にはよさそうという値を表現するものになる）の方が共分散よりは重要だよなという気はするので、共分散の方はまず（時間依存）定数でやるという感じなのだろう。

次に本丸の $ μ_θ (x_t, t) $ を見ていく。
まずは先ほどの $ \sum_θ (x_t, t) = σ^2_t I $ を使って、gaussian の KL divergence であることに注意すると、以下のように書ける。

$ C $ は $ θ $ に寄らない定数。
これを使うと、$ μ_θ $ の最も straightforward な parametrization は forward process $ q $ の posterior mean を予測するようなモデルというものになる。

これからもう一歩踏み込んでピクセル値の言葉で表現するために、(4) 式により、$ ε \sim N(0, I) $ としたときに $ x_t(x_t, ε) = \sqrt{ \bar{α} } x_0 + \sqrt{ (1 - \bar{α}) } ε → \sqrt{ \bar{α} } x_0 = x_t(x_t, ε) - \sqrt{ (1 - \bar{α}) } ε $ と書けることに注意して、簡単な計算を経て以下の表式を得る。

今考えているのは $ L_{t-1} $ であり、その際に $ x_t $ は given であることを考えると、結局のところモデルで表現したいのは $ ε_θ $ であり、平均 $ μ_θ $ は以下のような parametrization が取れることになる。

この parametrization を用いると、objective function は次のように ε の言葉で書ける。

これは最初の $ μ $ の言葉で書かれていたときよりももっと直接的に意味がわかりやすい。
さまざまな元画像 $ x_0 $ を出発して、ある time step $ t $ において、もともと forward process $ t-1 → t $ において（パラメタ $ β_t $ も含む）ノイズ $ ε $ がピクセル値 $ x_{t-1} $ に対して乗って $ x_t $ になるわけだが、モデルで学習する reverse process $ t → t-1 $ においてそのノイズを学習することでノイズを除去（これは画像生成をしていくことになる）する能力を有するようにしていこう、ということである。

[yoheikikuta]

data scaling, reverse process decoder, and L_0

データは {0,1,...,255} のピクセル値を取るとして、それを [-1, 1] に rescale する。
discrete な log likelihoods を得るために、以下のように last term を設計する。

$ D $ はデータ次元。
このチョイスの理由は以下とのことだが、これはちょっと界隈の知識がないのでパッとは分からない。
一旦はそういうものとしておいてどっかでまた遭遇したらもう少し勉強してみようかな。

Similar to the discretized continuous distributions used in VAE decoders and autoregressive models [34, 52], our choice here ensures that the variational bound is a lossless codelength of discrete data, without need of adding noise to the data or incorporating the Jacobian of the scaling operation into the log likelihood.

[yoheikikuta]

以上で differentiable な objective function が手に入ったが、この論文では以下の簡単化したものが生成画像の質を高めるのにも有効だし実装も楽だと言っている。

$ t $ は 1 から $ T $ までを取るもので uniform distribution からサンプルする。

$ L_T $ に関しては $ β_t $ を const として扱っているので無視して影響なし。
$ L_{t-1} $ に関しては $ t>1 $ に対応していて、これは単に係数を無視した形にしている。
$ L_0 $ に関しては $ t=1 $ に対応していて、これは (13) 式で gaussian を Talyor 展開して一次の項だけを残し、分散を無視し、そして x_0 が離散化された状態で期待値計算をするという点で”近似”になっている。元々の議論はなんだったんだというレベルの簡単化だが。

この論文では以降はこの簡単化した objective function を使用する。
この説明として、係数を無視することで、小さい $ t $ のとき（noise が弱いとき）に down weight されることになるのでそこの学習は相対的に重要性が下がり、大きい $ t $ のとき（noise が強いときで、つまり学習がもっと難しいとき）に学習が集中できるだろうと言っている。
これの方が良い結果になるというのはここまで議論してきたことのある程度の部分を無視することになるのでもっと詳しい説明が必要だけど特にされてないので、これは理論的に何か説明できそうというよりは現状ではほぼ empirical なものという印象。

[yoheikikuta]

ということで、学習とサンプリング（画像生成）のアルゴリズムは以下のように書ける。

学習においては任意の time step $ t $ で objective function が計算できるということから、元画像 $ x_0 $ を与え、time step $ t $ を uniform distribution からサンプリングして、その time step $ t $ で objective function を計算してその微分を計算してパラメタを更新する。

サンプリング（画像生成）においては、初期値となる標準正規分布からノイズをサンプルし、reverse process を辿って各 step で書くピクセルから学習した $ ε_θ $ でノイズを取り除いていく。その際に別途 noise $ z $ を加えているが、これは $ t $ が小さくなるにつれ弱くしていって、最後のステップでは 0 となるものである。

---

ちなみにここの Algorithm 2 は $ ε_θ $ を potential term としたときに Langevin dynamics となるので似ていると言っている。
あと (12) 式は https://arxiv.org/abs/1907.05600 の denoising score matching に似ていると言っている。
これらはまあそうという感じで、論文では似ているので良さそうでしょ！でもこれらとは違う新しいモデルで、確かに実験的にも有効であることを示しまたよ、面白そうでしょ？もっと色々研究してみてね、というメッセージを感じる。

[yoheikikuta]

あとは実験で性能を確かめる。
パラメタは以下のように設定。

• T = 1000
• β_1 = 10^-4, β_T = 0.01 で線形に増やす（これは画像のピクセル値 [-1, 1] と比べて十分に小さくなるようにしている）

これは最初の方に述べて以下を満たすために設定している。十分なのか？というのは実験結果以外からは読み取れないのでそんなもんなのかな、って思うくらいしかできないが。

ステップが細かく（$ T $ が十分大きく）て $ β_t $ が小さければ、という前提の下に成り立っている議論であることを頭に入れておく必要がある。

reverse process を表現するために PixelCNN++ を踏襲したアーキテクチャを使っている。
実験詳細は appendix B に書いてるので割愛。モデルや実験設定について書いてあるけど特別注意すべき点はなさそうだった。

画像生成のクオリティを Inception Score などで測った結果が以下。

Inception Score や Fréchet Inception Distance は画像生成でよく使われる指標で、NLL Test (Train) は negative log likelihood をテスト（学習）画像で計算したもの。
そしてこの論文の手法は unconditional な生成モデルなのでそちらにだけ結果が載っている。

Table 2 の方は色々書いてあるが、結論としては ε を予測するもので簡単化した objective function のものが一番良い結果であるということ。
この結果だけ見てもなんでそれがいいの？みたいな疑問には答えられないし、結果が不安定なところもある。
（この辺りだけみるとこのモデルめっちゃいいから詳しく理解しようという気持ちにまではならないが、diffusion model はその後流行っていく）

In Table 2, we show the sample quality effects of reverse process parameterizations and training objectives (Section 3.2). We find that the baseline option of predicting µ˜ works well only when trained on the true variational bound instead of unweighted mean squared error, a simplified objective akin to Eq. (14). We also see that learning reverse process variances (by incorporating a parameterized diagonal Σθ(xt) into the variational bound) leads to unstable training and poorer sample quality compared to fixed variances. Predicting, as we proposed, performs approximately as well as predicting µ˜ when trained on the variational bound with fixed variances, but much better when trained with our simplified objective

画像生成系は metrics だけ見てもなかなか判断が難しいが、実際生成された画像は以下のような感じ。
ビックリするほどではないけどなかなか良さそう、というレベル。

[yoheikikuta]

その他にも信号処理の問題と捉えて rate-distortion の振る舞いもチェックしている。
以下のように send/recieve プロセスを定義する。特に revieve のプロセスで stochastic な要素を除いている（これは単に解析のしやすさのため）ことに注意。

これで各 time step において rate/distortion をプロットしたものは以下（distortion は root mean squared error で計算）。

distortion は早い段階から減るが、rate が上昇するのは後半になってからである。
このことから、元画像が持ってる情報量の多くは imperceptible distoriton として使われているということが分かる。
このことから diffusion model は優秀な lossy compressor （元画像を noise にするというときに、元画像の知覚的な情報を適切に保持したまま非知覚的な情報は捨て去っている）となる inductive bias を有するということを推測している。

これは面白そうな話ではあるが、如何せん他のモデルとの比較は不十分なので diffusion model でない生成モデルの場合と比べてどうなのかは気になるところではある。

[yoheikikuta]

diffusion model での生成の振る舞いを見るためにいくつかの additional な実験もしている。

まずはどのように denoise されて画像が生成されていくか。
これは reverse process で time step 毎に denoise された画像を追っていけばいいのみ。
まず大域的な情報が作られ、そこから detail ができあがっていく様子が見て取れる。

次に $ x_{0, given} $ から始めて diffusion process のある time step $ t $ での $ x_{t, given} $ を取得し、その $ x_t $ で条件づけて $ p_θ (x_0 | x_{t, given}) $ をサンプルしたもの。
4 枚一組の右下が $ x_{t, given} $ であり、これは右から左に進むにつれて noise が乗せられていっている。
他の3枚は $ x_{t, given} $ に近い画像を元画像集合からサンプルしてくることになり、最初は当然同じ画像だが、noise が乗るにつれて大域的な情報（輪郭全体やデカめのサングラスなど）が保持されている様子が見て取れる。

最後は画像の補間で、まず元画像として $ x_0, x'_0 $ の 2 つを準備し、それを forwar process で noise にした $ x_t, x'_t $ にする。その foward process で $ \bar{x}_t = (1 - λ) x_t + λ x'_t $ として中間点を取り、それを reverse process で元画像空間での画像にする（論文の $ \bar{x}_t $ の表式は間違っているので注意）。
スムースに変化していく様子が見て取れる。

非常に新しい発見があるというわけではないが、モデルを使って色々面白い振る舞いが観察できるのは画像生成系の楽しいところ。

[yoheikikuta]

ということで一通り読み終えた。

この論文の主題ではないけど、そもそも diffusion model の基礎が分かってなかったのでそこをある程度計算を追って理解できたのはよかった。

diffusion model はその後の発展でより面白い結果を出しているものがありそうなので、次はその辺を読んでいこう。