この記事では拡散モデルの生成過程で使われるサンプラーの一種であるUniPCについて解説する。
これまでと同様に、SMLDの前提・表記を用いる。

関連記事

• DDPM: Denoising Diffusion Probabilistic Models
• DDPMの関連技術
• Stable Diffusionのモデル構造
• 拡散モデルの確率微分方程式
• 拡散モデルのサンプラー (1) - 離散過程・古典的数値計算法・LMS
• 拡散モデルのサンプラー (2) - DPM-Solver
• 拡散モデルのサンプラー (3) - UniPC (本記事)
• 拡散モデルのサンプラー - まとめ・Timestep Scheduler

目次

• UniPC
  • 前提
  • UniP
  • UniC
  • データ予測モデルのUniPC
• 参考

UniPC

UniPCはDPM-Solverの考え方を拡張したアルゴリズム。

前提

以下の定義・表記を用いる。

• $s=t_{i+1}<t=t_i<t_{i-1}<\cdots <t_{i-p+1}$
• $r_j:=t_{i-j}(j\in \{1,2,\cdots ,p-1\})$
• $r_p:=t_{i+1}$
  
  • したがって、$s=r_p<t<r_1<\cdots <r_{p-1}$
• $h:={\lambda }_s-{\lambda }_t={\lambda }_{t_{i+1}}-{\lambda }_{t_i}$
• $k_j:=\frac{{\lambda }_{r_j}-{\lambda }_t}{h}(j\in \{1,2,\cdots ,p\})$
  
  • したがって、$k_p=1$

本投稿における表記

UniP

DPM-Solverによると、拡散モデルのODEは

$$\begin{array}{rl}d{y}_t& =\epsilon (y_t)d{\varsigma }_t\\ & =-e^{{\lambda }_t}\hat{\epsilon }({\lambda }_t)d{\lambda }_t\end{array}$$

であり、

$$\begin{array}{rl}{y}_s& =y_t-{\int }_{{\lambda }_t}^{{\lambda }_s}{e}^{-\lambda }\hat{\epsilon }(\lambda )d\lambda \\ & =y_t-{\int }_{{\lambda }_t}^{{\lambda }_s}{e}^{-\lambda }\sum _{n=0}^{p-1}\frac{(\lambda -{\lambda }_t{)}^n}{n!}{\hat{\epsilon }}^{(n)}({\lambda }_t)d\lambda +\mathcal{O}(h^{p+1})\\ & =y_t-{\varsigma }_s\sum _{n=0}^{p-1}{h}^{n+1}{\varphi }_{n+1}(h){\hat{\epsilon }}^{(n)}({\lambda }_t)+\mathcal{O}(h^{p+1})\\ & =y_t-{\varsigma }_s(e^h-1)\hat{\epsilon }(y_t)-{\varsigma }_s\sum _{n=1}^{p-1}{h}^{n+1}{\varphi }_{n+1}(h){\hat{\epsilon }}^{(n)}({\lambda }_t)+\mathcal{O}(h^{p+1})\\ & =y_t+({\varsigma }_s-{\varsigma }_t)\hat{\epsilon }(y_t)-{\varsigma }_s\sum _{n=1}^{p-1}{h}^{n+1}{\varphi }_{n+1}(h){\hat{\epsilon }}^{(n)}({\lambda }_t)+\mathcal{O}(h^{p+1})\end{array}$$

が成り立つことを既に解説した。ただし、${\varphi }_n(h)=\sum _{m=0}^{\mathrm{\infty }}\frac{h^m}{(n+m)!}$

既に計算済みの$\{\epsilon (y_{r_j}){\}}_{j=1}^{p-1}$と$\epsilon (y_t)$を用いて、

$$\begin{array}{r}\sum _{n=1}^{p-1}{h}^{n+1}{\varphi }_{n+1}(h){\hat{\epsilon }}^{(n)}({\lambda }_t)=\sum _{j=1}^{p-1}{u}_j(h)(\epsilon (y_{r_j})-\epsilon (y_t))+\mathcal{O}(h^{p+1})\end{array}$$

と表すことを試みる。この条件を満たす関数の集合$\{u_j{\}}_{j=1}^{p-1}$を探す必要がある。

式を変形すると、

$$\begin{array}{rl}\sum _{n=1}^{p-1}{h}^{n+1}{\varphi }_{n+1}(h){\hat{\epsilon }}^{(n)}({\lambda }_t)& =\sum _{j=1}^{p-1}{u}_j(h)(\epsilon (y_{r_j})-\epsilon (y_t))+\mathcal{O}(h^{p+1})\\ & =\sum _{j=1}^{p-1}{u}_j(h)(\hat{\epsilon }({\lambda }_t+k_{j}h)-\hat{\epsilon }({\lambda }_t))+\mathcal{O}(h^{p+1})\\ & =\sum _{j=1}^{p-1}{u}_j(h)\sum _{n=1}^{p-1}\frac{(k_{j}h{)}^n}{n!}{\widehat{ϵ}}^{(n)}({\lambda }_t)+\mathcal{O}(h^{p+1})\\ & =\sum _{n=1}^{p-1}\frac{h^n}{n!}(\sum _{j=1}^{p-1}{k}_j^n{u}_j(h)){\widehat{ϵ}}^{(n)}({\lambda }_t)+\mathcal{O}(h^{p+1})\end{array}$$

なので、各$n\in \{1,2,\cdots ,p-1\}$に対して、

$$\begin{array}{r}{h}^{n+1}{\varphi }_{n+1}(h)=\frac{h^n}{n!}\sum _{j=1}^{p-1}{k}_j^n{u}_j(h)+\mathcal{O}(h^{p+1})\end{array}$$

$$\begin{array}{r}n!h{\varphi }_{n+1}(h)=\sum _{j=1}^{p-1}{k}_j^n{u}_j(h)+\mathcal{O}(h^{p+1-n})\end{array}$$

となれば良い。これはつまり、

$$\begin{array}{r}h\left(\begin{array}{c}1!{\varphi }_2(h)\\ 2!{\varphi }_3(h)\\ ⋮\\ (p-1)!{\varphi }_p(h)\end{array}\right)=\left(\begin{array}{cccc}1& 1& \cdots & 1\\ k_1& k_2& \cdots & k_{p-1}\\ ⋮& ⋮& \ddots & ⋮\\ k_1^{p-2}& k_2^{p-2}& \cdots & k_{p-1}^{p-2}\end{array}\right)\left(\begin{array}{c}{k}_1{u}_1(h)\\ k_2{u}_2(h)\\ ⋮\\ k_{p-1}{u}_{p-1}(h)\end{array}\right)+\left(\begin{array}{c}\mathcal{O}(h^p)\\ \mathcal{O}(h^{p-1})\\ ⋮\\ \mathcal{O}(h^2)\end{array}\right)\end{array}$$

となることを意味する。ここで、

$$\begin{array}{r}{V}_{p-1}:=\left(\begin{array}{cccc}1& 1& \cdots & 1\\ k_1& k_2& \cdots & k_{p-1}\\ ⋮& ⋮& \ddots & ⋮\\ k_1^{p-2}& k_2^{p-2}& \cdots & k_{p-1}^{p-2}\end{array}\right)\end{array}$$

はVandermonde行列と呼ばれる行列である。

式を解くと、

$$\begin{array}{r}\left(\begin{array}{c}{k}_1{u}_1(h)\\ k_2{u}_2(h)\\ ⋮\\ k_{p-1}{u}_{p-1}(h)\end{array}\right)=h{V}_{p-1}^{-1}(\left(\begin{array}{c}1!{\varphi }_2(h)\\ 2!{\varphi }_3(h)\\ ⋮\\ (p-1)!{\varphi }_p(h)\end{array}\right)+\left(\begin{array}{c}\mathcal{O}(h^{p-1})\\ \mathcal{O}(h^{p-2})\\ ⋮\\ \mathcal{O}(h)\end{array}\right))\end{array}$$

となる。
この式から、各$u_j(h)$は$u_j(h)=\mathcal{O}(h)$となることがわかる。

各成分の$\mathcal{O}(h^{p-n})$の部分はそれぞれそのオーダー以上の任意の関数に置き換えることができるが、特に全て$0$とした場合を考え、

$$\begin{array}{r}\left(\begin{array}{c}{k}_1{u}_1(h)\\ k_2{u}_2(h)\\ ⋮\\ k_{p-1}{u}_{p-1}(h)\end{array}\right)=h{V}_{p-1}^{-1}\left(\begin{array}{c}1!{\varphi }_2(h)\\ 2!{\varphi }_3(h)\\ ⋮\\ (p-1)!{\varphi }_p(h)\end{array}\right)\end{array}$$

とすれば良い。

このようにして求められたアルゴリズムをUniP(Unified Predictor)と言う。収束精度は$p$となる。

UniP

例 (p=1)

$p=1$の場合、総和の項がなくなりEuler法と完全に一致する。

$$\begin{array}{r}{\overline{y}}_{t_s}={\overline{y}}_t+({\varsigma }_s-{\varsigma }_t)\epsilon ({\overline{y}}_t)\end{array}$$

例 (p=2)

$p=2$の場合、行列やベクトルは全て1次元の値になる。

$$\begin{array}{rl}{v}_1(h)& :=h\cdot 1\cdot 1!{\varphi }_2(h)\\ & =\frac{e^h-h-1}{h}\end{array}$$

となるので、UniPのアルゴリズムは次のようになる。

$$\begin{array}{rl}{\overline{y}}_{t_s}& ={\overline{y}}_t+({\varsigma }_s-{\varsigma }_t)\epsilon ({\overline{y}}_t)-{\varsigma }_{t_s}\frac{e^h-h-1}{k_{1}h}(\epsilon ({\overline{y}}_{r_1})-\epsilon ({\overline{y}}_t))\\ & ={\overline{y}}_t+({\varsigma }_s-{\varsigma }_t)(\epsilon ({\overline{y}}_t)+\frac{1}{2k_1}\frac{2(e^h-h-1)}{h(e^h-1)}(\epsilon ({\overline{y}}_{r_1})-\epsilon ({\overline{y}}_t)))\end{array}$$

DPM-Solver-2の導出過程で$e^h-h-1=\frac{h(e^h-1)}{2}+\mathcal{O}(h^3)$と近似する場面があるが、その近似を実行しなかった場合に上の式と一致する。

例 (p=3)

$p=3$の場合、

$$\begin{array}{rl}\left(\begin{array}{c}{v}_1(h)\\ v_2(h)\end{array}\right)& :=h{\left(\begin{array}{cc}1& 1\\ k_1& k_2\end{array}\right)}^{-1}\left(\begin{array}{c}1!{\varphi }_2(h)\\ 2!{\varphi }_3(h)\end{array}\right)\\ & =\frac{h}{k_2-k_1}\left(\begin{array}{cc}{k}_2& -1\\ -k_1& 1\end{array}\right)\left(\begin{array}{c}{\varphi }_2(h)\\ 2{\varphi }_3(h)\end{array}\right)\\ & =\frac{h}{k_2-k_1}\left(\begin{array}{c}{k}_2{\varphi }_2(h)-2{\varphi }_3(h)\\ -k_1{\varphi }_2(h)+2{\varphi }_3(h)\end{array}\right)\end{array}$$

となるので、UniPのアルゴリズムは次のようになる。

$$\begin{array}{rl}{\overline{y}}_{t_s}& ={\overline{y}}_t+({\varsigma }_s-{\varsigma }_t)\epsilon ({\overline{y}}_t)-{\varsigma }_{t_s}h\frac{k_2{\varphi }_2(h)-2{\varphi }_3(h)}{k_1(k_2-k_1)}(ϵ({\overline{y}}_{t_{i-1}})-\epsilon ({\overline{y}}_t))-{\varsigma }_{t_s}h\frac{-k_1{\varphi }_2(h)+2{\varphi }_3(h)}{k_2(k_2-k_1)}(ϵ({\overline{y}}_{t_{i-2}})-\epsilon ({\overline{y}}_t))\end{array}$$

UniC

UniPCではUniPに加えて、UniC(Unified Corrector)というアルゴリズムを組み合わせて利用する。
UniCでは、UniPで一旦予測した$y_s$を利用して再び$y_s$を予測(補正)する。UniPの出力と区別するために、UniCの出力を$y_s^c$と表すことにする。

UniCのアルゴリズムの考え方はUniPと全く同じである。UniPでは$p$個の値

$$\begin{array}{rl}& \{\epsilon (y_t),\epsilon (y_{r_1}),\epsilon (y_{r_2}),\cdots ,\epsilon (y_{r_{p-1}})\}\\ =& \{\epsilon (y_t),\epsilon (y_{t_{i-1}}),\epsilon (y_{t_{i-2}}),\cdots ,\epsilon (y_{t_{i-p+1}})\}\end{array}$$

を利用していたが、UniCでは$p+1$個の値

$$\begin{array}{rl}& \{\epsilon (y_t),\epsilon (y_{r_1}),\epsilon (y_{r_2}),\cdots ,\epsilon (y_{r_{p-1}}),\epsilon (y_{r_p})\}\\ =& \{\epsilon (y_t),\epsilon (y_{t_{i-1}}),\epsilon (y_{t_{i-2}}),\cdots ,\epsilon (y_{t_{i-p+1}}),\epsilon (y_{t_{i+1}})\}\end{array}$$

を利用して$p$個の関数$v_j(h)$を求める。

総和の個数が1つ増えたことで、最終的な収束次数は$p+1$になる。

UniC

UniPとUniCを交互に実行することで$y_t^c$を生成していくアルゴリズムがUniPCである。

一方、$y_t$の予測には必ずしもUniPを使う必要はなく、DPM-Solverなど他のアルゴリズムとUniCを組み合わせることもできる。論文ではUniPCだけでなく、DDIM + UniCやDPM-Solver++(2M) + UniCなどについても実験が行われ、性能が比較されている。

データ予測モデルのUniPC

データ予測モデルのUniPも同様に導出できる。
DPM-Solver++によると、

$$\begin{array}{rl}{y}_s& =\frac{{\varsigma }_s}{{\varsigma }_t}{y}_t+{\varsigma }_s{\int }_{{\lambda }_s}^{{\lambda }_t}{e}^{\lambda }{\hat{x}}_{\theta }(\lambda )d\lambda \\ & =\frac{{\varsigma }_s}{{\varsigma }_t}{y}_t+\sum _{n=0}^{p-1}{h}^{n+1}{\psi }_{n+1}(h){\hat{x}}_{\theta }^{(n)}({\lambda }_t)+\mathcal{O}(h^{p+1})\\ & =\frac{{\varsigma }_s}{{\varsigma }_t}{y}_t+(1-\frac{{\varsigma }_s}{{\varsigma }_t}){\hat{x}}_{\theta }(y_t)+\sum _{n=1}^{p-1}{h}^{n+1}{\psi }_{n+1}(h){\hat{x}}_{\theta }^{(n)}({\lambda }_t)+\mathcal{O}(h^{p+1})\end{array}$$

が成り立つ。ただし、${\psi }_n(h)$は次のように定義される関数。

$$\begin{array}{rl}{\psi }_n(h)& ={\varphi }_n(-h)\\ & =\sum _{m=0}^{\mathrm{\infty }}\frac{(-h{)}^m}{(n+m)!}\end{array}$$

既に計算済みの$\{x_{\theta }(y_{r_j}){\}}_{j=1}^{p-1}$と$x_{\theta }(y_t)$を用いて、

$$\begin{array}{r}\sum _{n=1}^{p-1}{h}^{n+1}{\psi }_{n+1}(h){\widehat{x_{\theta }}}^{(n)}({\lambda }_t)=\sum _{j=1}^{p-1}{u}_j(h)(x_{\theta }(y_{r_j})-x_{\theta }(y_t))+\mathcal{O}(h^{p+1})\end{array}$$

と表すことを試みる。

ノイズ予測モデルの場合と記号が置き換わっただけなので、以降は同様に求めることができる。UniCの導出も同様。

参考

• [2302.04867] UniPC: A Unified Predictor-Corrector Framework for Fast Sampling of Diffusion Models
• UniPC/example/score_sde_pytorch/uni_pc.py at 25acd34b153fd82a4b1b9a945a520133419a02e1 · wl-zhao/UniPC · GitHub
• diffusers/src/diffusers/schedulers/scheduling_unipc_multistep.py at v0.22.3 · huggingface/diffusers · GitHub