Stable Diffusionにおけるサンプラーの役割を理解する

目的

Stable Diffusionには様々なサンプラーが用意されており、チューニングすべきパラメータとしてよく認知されていると思います。ユーザ観点からすれば、サンプラーの選択によって計算時間と絵の品質が変わることさえ知っておけば良いのですが、そもそもサンプラーとは何をやっているのか知っておきたい人向けにまとめてみます。できるだけ数式の世界には踏み込まず、理解のイメージを膨らませることだけを目的とします。Stable diffusionは拡散モデルの一種に過ぎないので、以下では一般性のため「拡散モデル」と呼びましょう。

拡散モデルのおさらい

拡散モデルの導入では以下のような解説がされることが一般的であり、非専門家にも直感で仕組みが伝わるかと思います。

"学習用画像に微小ノイズを加え、そのノイズを除去するようにニューラルネットワークをトレーニングする。このステップを何度も繰り返すことで画像は完全なランダムノイズのようになるが、そこから元の画像に戻す術をニューラルネットワークが学んでいる。これを利用し、ランダムノイズから始めてニューラルネットワークで徐々にノイズを除去していけば画像を生成することができる。"

ただ、この文脈からはサンプラーの必要性が全く読み取れません。このような非専門家向けの拡散モデルの説明から入ってきた人たち(私を含め)が抱えるモヤモヤを少しでも解消してみたいと思います。

サンプラーの役割

サンプラーがなぜ必要なのか？ その答えは、ニューラルネットワークが予測するのは加えられたノイズの平均値であって、実際に加えられたノイズの値そのものを予測していないからです。サンプラーは、ニューラルネットワークが予測したノイズの平均値をもとに、具体的なノイズのサンプルを生成します。

画像生成過程を数式で書くと、$x_0$を完全なランダムノイズとして、1ステップ後の画像$x_1$は
$x_1 = x_0 - n_0$
ここで$n_0$は除去すべきノイズを表します。
$n_0$の平均値を$\mu_0$としますと、$\mu_0$はニューラルネットワークが算出してくれますが$n_0$そのものを算出するわけではありません。サンプラーが平均値$\mu_0$のノイズサンプルを発生させるのです。

なぜランダム性が必要なのか？

ここで率直な疑問が湧くのではないでしょうか？。サンプラーでランダム性を加えることなんてせず単に$n_0 = \mu_0$とすればいいのでは、と。

ランダム性が必要な理由はいくつかありますが直感的にもわかるものを挙げると、まず第一に多様性が失われます。仮にノイズ除去過程にランダム性がないと、出来上がり画像は初期ノイズだけに依存します。初期ノイズ自体にランダム性はあるものの、生成過程が確定的動作であれば生成物がワンパターンになりがちなのはイメージできます。

第二に生成画像のリアリティが損なわれる点。例えば「空」の画像を生成したとします。仮に$n_0 = \mu_0$とすると空全体が同じ値になりがちで、空一面真っ青なリアリティのない画像が作られることでしょう。一方$n_0$をサンプラーで決めると、画素の値にメリハリがつき、ある部分は青空、別の部分は雲、のように現実の空が描かれていくことになります。

様々なサンプラーがある理由

ノイズは正規分布に従う乱数です。なので平均$\mu_0$と分散$\sigma_0$が決まれば発生させることができます。
平均は上述の通りニューラルネットワークが与える一方で、分散をどのように決めるかは自由度があります。これが様々なサンプラーが存在する理由です。典型的には、初期のステップほど大きな分散とし徐々に分散を下げていきますが、サンプラーによってメリハリの付け具体が変わってきます。

様々サンプラーがある理由を微分方程式でより深く理解する

拡散モデルについて深く勉強していると(確率的)微分方程式なるワードが登場してさぞ面喰らうことかと思います。これは拡散過程が微分方程式で表現できるためです。ランダムノイズから始まる画像の列$x_0, x_1, ...$は離散系列ですが、連続変化だと思って時間$t$に対する関数$x(t)$だと考えてみましょう。
そしてこの関数$x(t)$の時間変化率すなわち微分を書いてみると、導出過程は省略しますが(というか自分も理解していない)、
$\frac{dx(t)}{dt} = F(x(t), \mu(t), \sigma(t))$
という形になります。はい、微分方程式ですね。

そもそも拡散モデルの理論的背景にはこの微分方程式があって、数値的に離散化して解いているものがStable diffusionなど拡散モデルに基づくアプリケーション、という訳なのです。

微分方程式の数値的解法には最も初歩的なEuler法を始め様々なものがあります。Stable diffusionに用意されているサンプラー"Euler"はまさにこの微分方程式をEuler法で愚直に解こうとしているのですね。また、近似によって高速化する研究もされており、その成果が様々なサンプラーの実装となってStable diffusionに搭載されていています。