もし生物情報科学専攻の大学院生が "StableDiffusion" を理解しようとしたら 9 ~Stable Diffusion②~

前回

cake-by-the-river.hatenablog.jp

前回の続きです。今回は、image2imageに軽く触れたのち、主にサンプリングに関する技術をまとめます。この辺りは今までの話とは別の流れで、これまた難しめではあるので、可能な限り難しい理論には踏み込まないような説明を試みようと思います。

image2image, inpainting

image2imageは画像を入力として画像を生成する技術ですが、これは今までの知識だけからすぐわかります。つまり、画像を入力として、拡散過程をある程度実行し、得られたノイジーな潜在表現を逆拡散して戻せばよいです。この際、どこまで戻るかを示すパラメータとして、AUTOMATIC1111（Stable Diffusionの実行環境の一つ）では、Denoising Strength が用いられます。Denoising Strength は [0, 1] の範囲で指定し、逆拡散にかかる総ステップ数とこの数値の積で実際に戻るステップ数を決めます。デフォルトは 0.75 ですが、これは、実際のステップ数に対し 75 % まで拡散するという意味になります。

Denoising Strength の数値によってどれほど画像が変化しうるのかについて、以下で少し検証してみました。Promptが何もない状態と、"pink hair"というPromptを入れた場合について、異なるDenoising Strengthでの結果を表示してみます。今回用いたモデルは、AbyssOrangeMix3と呼ばれるものになります。

Denoising Strengthの強さと、プロンプトの影響。ステップ数の初期値は20、サンプラーはEuler aを使用（サンプラーについては後述）。

Promptがない場合、Denoising Strengthが0.4あたりまでは元画像の復元が割と出来ていますが、それ以上では人物像が崩れていきます。とくに、0.8を超えるものでは、元画像の構図すらも崩れていることが分かります。一方、Promptでピンクの髪という指定を行った場合、0.4までのDenoising Strengthでは髪色を大幅に変えることはできておらず（それどころか、目の色がピンクになるという副作用が生じていることに注意）、0.6以降で上手く反映していっていることが分かります。実際、経験的にDenoising Strengthが0.4~0.8あたりの数値がPromptの反映と元画像の構図の維持の両立が可能だと考えられているようです。

image2imageのほかに、Inpaintingという機能も使うことが出来ます。これも非常に単純で、画像の一部にマスクを施し、その部分についてimage2imageを行い、元画像にはめ込むという処理になります。

サンプリング

DDPMに残された最後の問題として、そのままではサンプリングに多大な時間（とリソース）が必要になることが挙げられます。拡散過程は、十分な時間 $T$ をかけることで逆拡散過程にもガウス分布を使うことの正当性が保たれるため、標準で $T=1000$ という長い時間が必要になります。今回はこのサンプリングを工夫していくことで、 $T=20$ のような現実的な時間ステップ数での画像生成を行う方法を見ていきます。

確率微分方程式

NCSNの解説の際、ランジュバン・モンテカルロ法は以下の確率過程に等しい操作を行っている、という説明をしたことを覚えているでしょうか？

$d\mathbf{x}_t = -\nabla_{\mathbf{x}} E(\mathbf{x}_t) dt + \sqrt{2} d\mathbf{w}_t$

これは $\mathbf{x}_t$ の微小な変化 $d\mathbf{x}_t$ がエネルギーの勾配とウィーナー過程（ブラウン運動） $\mathbf{w}_t$ の微小な変化によって書けるという意味ですが、今回はこのような微小な変化の方程式をもう少ししっかりと見ていくことにします。

arxiv.org

確率微分方程式（SDE）とは、主に以下の形式で書ける微分方程式の総称です。

$dx_t = f(x,t)dt + g(x,t)dw_t$

$dw_t$ はウィーナー過程の微小変化量であり、 $f(x,t)$ はドリフト、 $g(x,t)$ は拡散係数と呼ばれています。ただし、今回は簡単のために拡散係数が $x$ には依存しない $g(t)$ で書けるものと考えます。

ウィーナー過程 $w_t$ は、初期状態 $w_0$ が0で、

$w_t - w_s \sim \mathcal{N}(0, t-s)$

に従うようなほぼ連続な確率過程（時間変化する確率変数）です。上の性質から、微小量 $dw_t$ は微小な時間 $dt$ において、平均 $0$ 分散 $dt$ のガウス分布に従う数値を返すのだと考えることが出来ます。ウィーナー過程の凄いところは、この微小な変化の足し合わせ（積分）もまた、再生性からガウス分布になるところであり、これは私たちが一般に想像する拡散の過程の直感と一致します。したがって、ブラウン運動を表すモデルとして最も一般的に用いられています。

$\int_0^t dw_s ds \sim \mathcal{N}(0, \sqrt{t})$

さて、私たちはガウス分布に従うモデルとしてNCSNやDDPMを導入してきました。これらは、分散発散型と分散保存型のモデルとして、シグナルノイズ比SNRのもとに統一的な枠組みで語ることが出来ることを以前紹介しました。実は、それぞれの時間の連続極限をとることで、SDEへと変換できることが知られています。ここでは、DDPMの例を見てみましょう。

まずは、DDPMの式を改めて提示します。

$x_i = \sqrt{1-\beta_i} x_{i-1} + \sqrt{\beta_i} z_{i-1},\ z_{i-1} \sim \mathcal{N}(0, \mathbf{I})$

DDPM自体は $t=1,...,T$ に対し、 $i=t$ として上の式を行っているわけです。しかし、今回は時刻の幅 $t=0 \sim T$ をそのままに、この $i=1,2,...,N$ について $N \to \infty$ の極限をとることにします。これにより、有限時間の範囲内かつ連続的な時間 $t$ での定式化を行うことが出来るわけです。

$N \to \infty$ の極限において、時間の幅 $\Delta t = 1 / N$ は最終的に微小量 $dt$ となります。この $dt$ に関する微分方程式へと式変形するために、 $\beta_i$ を

$\hat{\beta}_i = N \beta_i$

として変数変換を行います。

$\displaystyle \begin{eqnarray} x_i &=& \sqrt{1-\hat{\beta}_i / N} x_{i-1} + \sqrt{\hat{\beta}_i / N} z_{i-1} \\ &=& \sqrt{1-\hat{\beta}_i \Delta t} x_{i-1} + \sqrt{\hat{\beta}_i \Delta t} z_{i-1} \end{eqnarray}$

$N \to \infty$ の極限にて、各変数 $x_i, \hat{\beta}_i$ は $i$ の変数から $t$ の変数へと書き換えられ、それぞれ $x(t), \beta(t)$ と表すことができます。これを踏まえて $\Delta t$ のテイラー展開を行います（ $\Delta t \to 0$ を考えるため一次までの近似とする）。

$\displaystyle \begin{eqnarray} x(t+\Delta t) &=& \sqrt{1-\beta(t+\Delta t)\Delta t} x(t) + \sqrt{\beta(t+\Delta t)\Delta t} z(t) \\ &=& x(t) - \frac{1}{2}\beta(t+\Delta t)\Delta t x(t) + \sqrt{\beta(t+\Delta t)\Delta t} z(t) + O(\Delta t^2) \\ &\approx& x(t) - \frac{1}{2} \beta(t) \Delta t x(t) + \sqrt{\beta(t)} \sqrt{\Delta t} z(t)\end{eqnarray}$

最後に $N \to \infty$ すなわち $\Delta t \to 0$ の極限をとってあげると、ウィーナー過程の性質から、無事SDEとしての定式化が得られます。

$\displaystyle dx = \lim_{\Delta t \to 0} x(t+\Delta t) - x(t) = - \frac{1}{2} \beta(t) x(t)\ dt + \sqrt{\beta(t)} dW_t$

さて、逆拡散過程も各時刻での遷移がガウス分布となるマルコフ過程で書けるはずでしたが、SDEとしてのDDPMでも、逆拡散過程がSDEの形で書けます。

$\displaystyle dx_\tau = [f(x, \tau) - g(\tau)^2 \nabla_{x} \mathrm{log}\ p_\tau(x)] d\tau + g(\tau) d\tilde{w}_\tau$

ここで、 $\tau$ は反転した時間を表し、 $\tilde{w}_\tau$ は $w_t$ とは異なるウィーナー過程を表します。この導出には式変形を多用する必要があるため、詳細は割愛しますが（知りたい方は、上の論文の参考文献（Anderson,1982）か、拡散モデル本を参照してください）、制御工学の知見からその発想の根底を考えることが出来ます。これについて末尾にて追記するので、式のお気持ちが気になる方は、そちらをご覧下さい。

逆拡散過程のSDEは、拡散過程のSDEにスコア $\mathrm{log}\ p_\tau(x)$ の項が追加されているだけの形式です。したがって、スコアさえ計算できれば、計算することが可能です。NCSN（および等価なDDPM）は、スコアそのものを推定する明示的スコアマッチングと等価な目的関数をもっていたため、デノイジングの過程でスコア計算が可能であり、SDEの計算も出来ます。

確率フローODE

このように逆拡散過程のSDEは計算可能であり、汎用のソルバーを用いることも出来ますが、より効率的なサンプリングを可能にすることが出来ます。それは、SDEで動きが決まる確率変数 $x$ について、その変数が従う分布 $p_t(x)$ （時刻ごとに異なる）を完全に再現できる常微分方程式（ODE）、すなわち確率的ではなく決定的に時間変化する系を用意することです。このSDEと等価なODEは、確率フローODEと名付けられています。

確率フローODEを導出するためには、Fokker-Planck方程式と呼ばれる非平衡統計力学の知識を避けることができません。ここでは、出来るだけ簡単にその流れを説明してみます。

（詳しい導出）
mathlog.info

SDEは、 $x$ の微小変化に関する方程式でしたが、分布 $p_t(x)$ に対する微分方程式の形を考えたのが Fokker-Planck方程式になります。Fokker-Planck方程式は、（波動方程式や拡散方程式のように）分布 $p_t(x)$ の時間微分 $\frac{\partial p_t(x)}{\partial t}$ を、位置 $x$ に関してテイラー展開した形を考えます。

$\displaystyle \frac{\partial p_t(x)}{\partial t} = \left(-\frac{\partial}{\partial x}a(x,t) + \frac{1}{2} \frac{\partial^2}{\partial x^2} b(x,t) \right)p_t(x)$

右辺の $a(x,t), b(x,t)$ はそれぞれ、1階微分 $\frac{\partial}{\partial x}$ と2階微分 $\frac{\partial^2}{\partial x^2}$ の係数に相当するといえます。

そして、SDEがWiener過程、すなわち時間に従うガウス分布の形で示されていたことを踏まえると、分布の"速度" $\frac{\partial p_t(x)}{\partial t}$ の係数も、SDEの右辺に由来するガウス分布の係数（平均・分散）が反映されそうなことが想像できます。

実際、SDEを $dx/dt$ の形で考えると、平均 $0$ , 分散 $1$ のガウス分布 $dw/dt$ を、 $g(t)$ 倍拡大したのち $f(x,t)$ だけ移動させた、といった雑な解釈が可能であり（ランジュバン・モンテカルロ法の動きも同様）、導出されるFokker-Planck方程式も、

$\displaystyle \frac{\partial p_t(x)}{\partial t} = \left(-\frac{\partial}{\partial x}f(x,t) + \frac{1}{2} \frac{\partial^2}{\partial x^2} g(t)^2 \right)p_t(x)$

のように対応して導出ができます（細かいことは無視しています）。

さて、Fokker-Planck方程式の $a(x,t), b(x,t)$ はそれぞれ、SDEの $f(x,t), g(x,t)$ に対応する係数部分でした。もしSDEから確率的な項を取り除こうと考えるならば、 $g(x,t) = 0$ となるような式変形を求める必要があります。したがって、Fokker-Planck方程式を式変形させ、特に第二項の2階微分 $\frac{\partial^2}{\partial x^2}$ を上手く消去することが出来れば、決定的な微分方程式 = ODE を求めることが出来そうです。

今回は拡散係数 $g(t)$ が $x$ に依存しないものとして計算しており、Fokker-Planck方程式の二階微分の項は単に

$\displaystyle \frac{\partial^2}{\partial x^2}p_t(x)$

の計算のみで済みます。そしてこれは、対数の微分 $(\mathrm{log}\ p(x))' = p(x)'/p(x)$ を逆手に取ることで、1階微分の項に押し込めることが出来ます。

$\displaystyle \begin{eqnarray} \frac{\partial^2}{\partial x^2}p_t(x) &=& \frac{\partial}{\partial x} \left(\frac{\partial}{\partial x}p_t(x)\right) \\ &=& \frac{\partial}{\partial x} \left(\frac{\partial}{\partial x}p_t(x) \times \frac{p_t(x)}{p_t(x)}\right) \\ &=& \frac{\partial}{\partial x} \left(\frac{\partial}{\partial x}(\mathrm{log}\ p_t(x))\ p_t(x)\right)\end{eqnarray}$

$\displaystyle \therefore \displaystyle \frac{\partial p_t(x)}{\partial t} = \left(-\frac{\partial}{\partial x}[f(x,t) + \frac{1}{2}g(t)^2\frac{\partial}{\partial x}\mathrm{log}\ p_t(x)] \right) p_t(x)$

SDEとFokker-Planck方程式との対応関係を考えると、これは

$dx = [ f(x,t) - \frac{1}{2}g(t)^2\nabla_x \mathrm{log}\ p_t(x) ] dt + 0 dw$

というSDEに対応し、確率的な項がなくなり、ODEとして扱うことが出来ました。

拡散過程・逆拡散過程のSDEと確率フローODE。確率過程でありギザギザした動きをするSDEに対し、確率フローODE（白色）は滑らかで決定的な動きをしているところに注目。論文より引用。

k-diffusion

さて、確率フローODEを頑張って導出したわけですが、これをどのように使えばよいのでしょうか？

SDEは、その確率的な動きが微小な時刻 $dt$ のうちに組み込まれている関係で、ステップ数を減らしにくくなっています。一方、確率フローODEは決定的な動きであり、その動きを離散化したジャンプにより近似することで、実際にかかるデノイジングステップ数を大幅に削減することが可能です。とくに、この近似法を色々と工夫することで、ステップ数の削減と実際に得られる画像の特徴に違いが生じます。こうしてステップ数を削減しながら画像を生成（サンプリング）する方法を、サンプラーと呼んでいます。

Stable Diffusionの実行環境であるAUTOMATIC1111では、サンプラーとして様々なものが挙げられています。とくに、よく使われているものとして、

Euler a
DPM++ 2M Karras
DDIM

などがあります。今回はこの辺りのアルゴリズムの基本である Karras et al. (2022) の論文について触れ、一部サンプラーの仕組み（k-diffusion）に迫ろうと思います。

arxiv.org

なお、

DDIM
UniPC

については後の節で扱うこととします。

Karras et al. (2022) の論文は、以下の流れでサンプリングアルゴリズムを構築します。

入力分布が徐々にガウス分布へと拡散する過程を表現できる、拡散方程式を探し出す
拡散方程式を、Fokker-Planck方程式とSDEの対応関係を用いて、新しいSDEに変換する
新しいSDEを基に、Predictor-CorrectorライクなSDEソルバーをサンプラーとする

以下、それぞれのステップについて解説していきます。

まずは今後の記法のために、DDPMなど拡散モデルの式を、シグナル・ノイズの係数の関係式に変換します。DDPMは $\alpha_t$ を用いて、初期画像から現在の画像を次の式で表現することができます。

$\displaystyle x_t = \sqrt{\bar{\alpha_t}} x_0 + \sqrt{1 - \bar{\alpha_t}} z,\ \ \bar{\alpha_t} = \prod_{s=1}^t \alpha_s,\ \ z \sim \mathcal{N}(0, \mathbb{I})$

これは $q(x_t|x_0)$ の式（ガウス分布の再生性に依る）からすぐにわかります（DDPM記事を参照）。この式を

$\displaystyle \sigma(t) = \sqrt{\frac{1 - \bar{\alpha_t}}{\bar{\alpha_t}}}$

により書き換えます。

$\displaystyle x_t = \sqrt{\bar{\alpha_t}} \left( x_0 + \sigma(t) z \right) = \frac{1}{\sqrt{1 + \sigma(t)^2}} \left( x_0 + \sigma(t) z \right)$

続いて、変換した式が拡散方程式（分布が拡散していくときに従う方程式）

$\frac{\partial q(x,t)}{\partial t} = \kappa(t) \frac{\partial^2}{\partial x^2} q(x,t)$

に従うものと仮定し、係数 $\kappa(t)$ を求めてみましょう。今、NCSNの例などを思い出してみると、元の画像が為す分布 $p(x_0) = p_{data}(x_0)$ は、時間が経つにつれて徐々に拡散していき、最終的に単純なガウス分布へと変化していく状況を考えていました。この部分は、具体的には各時点でデータ分布 $p_{data}$ にノイズ $\mathcal{N}(0, \sigma(t)^2 \mathbb{I})$ が畳み込まれた分布

$q(x, t) = p_{data}(x)*\mathcal{N}(0, \sigma(t)^2 \mathbb{I})$

で表せるのでした（NCSNを参照）。このような畳み込まれた分布の時間微分を考える際には、Fourier変換を用いるのが定石となっています。Fourier変換は分布の畳み込みを分布の積の形に変換することが出来るため、計算が簡単になります。今、位置 $x$ を運動量 $\nu$ の空間にFourier変換した後の分布は

$\hat{q}(\nu, t) = \hat{p}_{data}(\nu) \mathrm{exp}(-\frac{1}{2} |\nu|^2 \sigma(t)^2)$

と表されます（証明は割愛）。この分布の時間微分は、合成関数の微分により

$\displaystyle \begin{eqnarray} \frac{\partial \hat{q}(\nu,t)}{\partial t} &=& - \dot{\sigma}(t) \sigma(t) |\nu|^2 \hat{p}_{data}(\nu) \mathrm{exp}(-\frac{1}{2} |\nu|^2 \sigma(t)^2) \\ &=& - \dot{\sigma}(t) \sigma(t) |\nu|^2 \hat{q}(\nu, t) \end{eqnarray}$

と書ける一方、拡散方程式のFourier変換は

$\frac{\partial \hat{q}(\nu,t)}{\partial t} = - \kappa(t) |\nu|^2 \hat{q}(\nu,t)$

と書けることが知られているため、この二つを比較することで、係数 $\kappa(t)$ が

$\kappa(t) = \dot{\sigma}(t) \sigma(t)$

となることが分かりました。

次に、今求めた拡散方程式が、確率微分方程式の形式に従うとし、そのFokker-Planck方程式と一致するようなときを考えてみましょう。Fokker-Planck方程式は

$\displaystyle \frac{\partial q(x,t)}{\partial t} = \left(-\frac{\partial}{\partial x}f(x,t) + \frac{1}{2} \frac{\partial^2}{\partial x^2} g(t)^2 \right)q(x,t)$

と書けました。今、これが拡散方程式と同じであるならば、

$\displaystyle \left(-\frac{\partial}{\partial x}f(x,t) + \frac{1}{2} \frac{\partial^2}{\partial x^2} g(t)^2 \right)q(x,t) = \dot{\sigma}(t) \sigma(t) \frac{\partial^2}{\partial x^2} q(x,t)$

$\displaystyle \frac{\partial}{\partial x}(f(x,t)q(x,t)) = \left( \frac{1}{2} g(t)^2 - \dot{\sigma}(t) \sigma(t) \right) \frac{\partial^2}{\partial x^2} q(x,t)$

$\displaystyle \begin{eqnarray} f(x,t) &=& \left( \frac{1}{2} g(t)^2 - \dot{\sigma}(t) \sigma(t) \right) \frac{\partial}{\partial x} q(x,t) \frac{1}{q(x,t)} \\ &=& \left( \frac{1}{2} g(t)^2 - \dot{\sigma}(t) \sigma(t) \right) \nabla_x \mathrm{log}\ q(x,t) \end{eqnarray}$

このように、 $f$ の条件を導くことが出来ました。したがって、任意の $g(t)$ で、対応するSDEが

$\displaystyle dx = \left( \frac{1}{2} g(t)^2 - \dot{\sigma}(t) \sigma(t) \right) \nabla_x \mathrm{log}\ q(x,t) dt + g(t) dw$

として書けます。同様の計算から、逆過程のSDEも、類似した形で書けます。

$\displaystyle dx = \left( - \frac{1}{2} g(t)^2 - \dot{\sigma}(t) \sigma(t) \right) \nabla_x \mathrm{log}\ p(x,t) dt + g(t) dw$

この式について、少し解釈を述べておきましょう。例えば $g(t) = 0$ を考えた場合、対応する微分方程式は

$dx = - \dot{\sigma}(t) \sigma(t) \nabla_x \mathrm{log}\ q(x,t) dt$

となりますが、これは実はNCSNの確率フローODEと一致することが示せます。すなわち、上記のSDEは、拡散モデルの確率フローODEに確率的な項をいくらか加えたものだと考えられます。一方、確率的な項は、スコア $\nabla_x \mathrm{log}\ p(x,t)$ によって減衰する項と、新たに生まれる項 $g(t) dw$ とに分けられており、新しいノイズと古いノイズを入れ替える関係にある（その係数が $g(t)$ ）と考えられます。

こうして新たに出来たSDEは元の確率フローODEと比べて何が有用なのでしょうか？

Karras 達は、一定確率で古いノイズを新しいノイズで置き換える部分が、初期のサンプリングプロセスで生じた誤差（ODEを近似した場合）を緩和するのに役立つと主張しています。一方、新しいノイズが強すぎてもまた誤差を助長するため、適切な $g(t)$ を選ぶことで、確率フローODEより良いサンプリングを可能に出来ると考えたそうです。

さて、最後に実際のサンプリングアルゴリズムを見てみましょう。k-diffusion では、今考えた「ノイズの追加と減衰」による確率フローODEの改善を、各種ODEの近似解法（Euler法, Heun法, DPM法など）に適用するアルゴリズムを提案しています。ここではEuler法に適用した場合の概略を示します。

Euler法は非常にシンプルな解法で、時刻 $t$ でのODEの進行方向（傾き） $d = dx / dt$ を用いて次の時刻での解を $x(t_{next}) = x(t) + d \times dt$ として求めるものです。一方、提案したアルゴリズムでは、この時の傾きを、あえてノイズを更に加えた結果から推定するようにします。

ガウスノイズ $\epsilon$ を生成し、少しノイズレベル $\sigma(t)$ が強くなった $\hat{\sigma}(t)$ に対応するようノイズを入力に加える $\hat{x}(t) = x(t) + \sqrt{\hat{\sigma}(t)^2 - \sigma(t)^2} \epsilon$
加えたノイズをデノイジングするよう拡散モデル（U-Net）に入力し、デノイジング後の結果 $D_\theta (\hat{x}(t), \hat{\sigma}(t))$ を求める
デノイジングの方向を表すベクトル $d = (x(t) - D_\theta (\hat{x}(t), \hat{\sigma}(t))) / \hat{\sigma}(t)$ を用いて、次の時刻 $t_{next}$ でのデノイジング結果を（ステップ幅 $d \sigma(t)$ のEuler法近似で）予測する $x(t_{next}) = \hat{x}(t) + d\sigma(t) d = \hat{x}(t) + (\sigma(t_{next}) - \hat{\sigma}(t)) d$

これは、まずノイズを加えてそのデノイジング（修正, Correction）を行った後、その結果を拡張するように次ステップの値を予測（Prediction）する、という形式を取っており、Corrector-Predictor型のアルゴリズムだと言われています。このようなアルゴリズムは、単にその時刻で学習したデノイジングのスコアを用いるよりも、サンプリングをより安定にすることが出来ると言えます。

確率フローODEとEuler法の例。右端がガウス分布、左端がデータ分布とし、 $T=1000$ のステップ数のDDPMと対応する確率フローODE、および20ステップのEuler法を記している。なお、Euler法の実装はかなり怪しい（雰囲気でお願いします）

確率フローODEとEuler法の例。右端がガウス分布、左端がデータ分布とし、 $T=1000$ のステップ数のDDPMと対応する確率フローODE、および20ステップのEuler法を記している。なお、Euler法の実装はかなり怪しい（雰囲気でお願いします）

上記はEuler法を用いた場合でしたが、同様の枠組みで異なるODE解法を用いることが出来、それらは k-diffusion にて実装されています。

github.com

例えば、Euler法にancestral samplingという手法を適用した Euler a サンプラーや、これらのODEを精度良く解くために開発された DPM-Solver およびその改良版（DPM-Solver-2, DPM-Solver++ など）を適用したものが利用されています。どのような方法であっても、その根本は、確率フローODEに対し確率性を少し加え安定性を確保したものである、ということが重要だと言えます。以下、少し各手法について触れます。

Euler aに代表されるancestral sampling（祖先サンプリング）という手法は、簡単に言えば各時刻での結果が前時刻の結果に依存して決定されるように生成するものです。Euler aの場合は、ステップ幅 $d\sigma(t)$ をあえてやや小さくし、足りない分をガウスノイズに変えています。これにより、各ステップで固有のガウスノイズの結果が出力に必要になり、ancestral samplingの形式をとることが出来ます。このような処理はEuler法で不安定になる予測結果を安定化させる効果があるようですが、かえって画像を収束させるのが難しくなるようです。

DPM-Solverシリーズは、比較的最近開発されたODEソルバーである指数積分に基づいたもので、中点法のように2時刻の結果を利用する高次のオーダーのソルバーです。DPM-Solver++はDPM-Solverの改良版で、条件付きの画像生成で不安定になりやすいという特徴を改善したものとなり、主に２次のオーダーのものが利用されます（現在時刻 $t_i$ と次の時刻 $t_{i+1}$ の間の数値のみを利用する single-step の方法 (DPM++ 2S) と、過去 $t_{i-1}$ の結果も利用する multi-step の方法 (DPM++ 2M) の2種）。また、確率フローODEからSDEを再構成できたように、ノイズ項を加えたSDEバージョンや、ancestral samplingを採用することもできます。

arxiv.org

まとめると、Euler法が最も単純な（1次の）ODEソルバーであり、より高次の近似はHeun法やDPM-Solver++で行うことでステップ数をさらに削減できます。更に、SDEやancestral samplingを採用した方法は安定性と収束性のトレードオフを取ることが出来ます。これで難解なサンプラー名もおよそ内容の推測が可能になります。例えば、"DPM++ 2S a Karras"というサンプラーは、「DPM-Solver++ の 2nd-order single-step ソルバーを、ancestral sampling と Karras et al. (2022) 由来の $\sigma$ スケジュールで行う」という意味になります。

これら各手法の経験的な違いの解説は巷に無限にあるので、そちらを参考にしてください。

DDIM, UniPC

単なるODE解法には基づかないサンプリングメソッドとして、DDIMやUniPCなどが利用されています。ここでは軽くその内容に触れます。

arxiv.org

DDPMでは各時刻の拡散ステップをマルコフ過程としていましたが、逆拡散過程のマルコフ過程を保ちつつ、拡散過程は非マルコフ過程にしたもの（拡散とは異なるものとなるため、推論過程と呼ばれている）が DDIM となります。

DDIMは、その目的関数がDDPMと定数倍の差しかないことが示せるため、DDPM型のモデル（Stable Diffusionなど）をそのままに、サンプリング部分だけ変えることが出来ます。また、推論 $q(x_{t-1} | q_x, q_0)$ の分散 $\sigma(t)$ を調整することで、ステップ数を削減することが出来、それでも精度が低下しないことが経験的に知られていました。

DDIMが高速化できる理由については、データが低次元の多様体に収まる場合（多様体仮説）、データ空間に近いDDPM拡散過程でのデノイジング関数が、一定の方向になりやすく、その部分のステップを近似的に飛ばしても復元されやすいことに由来するようです。

拡散モデルでDDIMが少ないステップ数で高精度に生成できる理由は謎だったが、多様体仮説においてノイズを予測する式は時刻不変となり、DDIMがそのSDEの離散化とみなせるため。gDDIMはこれを非等方性ノイズ（例BDM, CLD）にも拡張。ステップ数が数十回でも高品質に生成可能 https://t.co/aa3RdignTa
— Daisuke Okanohara / 岡野原大輔 (@hillbig) 2022年11月7日

arxiv.org

多様体仮説とDDIM。各ピクセルの数値のデノイジングが時刻に依らず一定値を取る傾向が右図から確認できる。gDIMMの論文より引用。

一方、UniPCは2023年に登場したサンプラーであり、DPM-Solverなど汎用なサンプラーに対し、過去のデノイジング結果を複数組み合わせることで任意の次元のソルバーを作ることが出来ます。UniPCは、Predictor-Corrector型のアルゴリズムに基づき、UniP（Unified Predictor）とUniC（Unified Corrector）と呼ばれる二つの解析的に同じ形式をした計算器を用います。p次のオーダーのソルバーをp+1次のソルバーにできる UniC-p、およびソルバーに対応する UniP-p を採用することで、任意の次元のソルバーが可能となるわけです。

UniCおよびUniPの更新式は、p=1の場合 DDIM に、p=2の場合 DPM-Solver に等しくなることが示せ、これらのソルバーの一般化であると考えることが出来ます。3次以上のUniPCにすることで、ステップ数を更に下げつつ、精度を下げずに生成することが可能になると考えられます（AUTOMATIC1111では3次のUniPCが実装されています）。

長くなりましたが、今回はimage2imageおよびサンプラーのアルゴリズムについて一通り解説しました。これにてStable Diffusionの基本機能に関しては殆ど解説できたものと考えられます。次回からは、拡張機能（追加学習や高度な条件付けなど）について足を進めていこうと思います。

逆過程のSDEと制御工学

そのうち書きます。