（潜在）一貫性モデル/ (Latent) Consistency Models

（潜在）一貫性モデル / (Latent) Consistency Models

23年の年末は拡散モデルをさらに高速に実行できる一貫性モデルが流行っています。簡単にモデルを把握したいので、確認したいと思います。

拡散モデル

まずは簡単に拡散モデルのおさらいをします。

Denoising Diffusion Probabilistic Models
Jonathan Ho, Ajay Jain, Pieter Abbeel

https://arxiv.org/abs/2006.11239

画像：https://arxiv.org/abs/2006.11239

拡散モデルは、あるノイズから少しずつノイズ除去して画像を復元するモデルとなっていて、初期のX_TからX_0へと状態を戻していきます。その戻し方はU-Netと呼ばれるニューラルネットワークを利用するのが一般的ですが、今回はさらにその拡散モデルを拡張した一貫性モデルを見てみたいと思います。

拡散モデルは、初期のノイズ状態が下記のガウス分布から始まり、

p(x_T) = N(x_T;0,I)

途中の過程は、

p_\theta (x_{t-1}|x_t):=N(x_{t-1};\mu_\theta(x_t,t),\sum_\theta(x_t, t))

と定義づけられ、最終的な復元状態はこの過程を経て得られる。逆拡散過程と呼ばれるらしい。

p_{\theta}(x_{0:T}) := p(x_T)\prod_{t=1}^T p_\theta(x_{t-1} | x_t)

式の通り、初期状態から順番に確率的に状態を決めながら画像を復元していくらしい。

一貫性モデル

この一般的に見られる拡散モデルの説明は見た感じ時間ステップで離散的なので、この時間の関数を連続に捉える、連続時間拡散モデルを導入するところからスタートし、理論を拡張するのが一貫性モデルという新しいモデルらしい。この新しいモデルを利用することで、最近の画像生成などで1秒間にたくさんの画像を高速に出力できる理論のベースとなっている。

Consistency Models Yang Song, Prafulla
Dhariwal, Mark Chen, Ilya Sutskever
https://arxiv.org/abs/2303.01469

Score-Based Generative Modeling through Stochastic Differential Equations
Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, Ben Poole
https://arxiv.org/abs/2011.13456

連続時間拡散モデルと確率微分方程式

連続時間拡散モデルは、下記の確率微分方程式の形でかける。

dx_t = \mu(x_t, t) + \sigma(t)dw_t

確率微分方程式（かくりつびぶんほうていしき、英: Stochastic differential equation）とは、1つ以上の項が確率過程である微分方程式であって、その結果、解自身も確率過程となるものである。（wikipediaより）
ウィキペディア

ブラウン運動とウィーナー過程

上記の式は、ブラウン運動の数学的なモデルである、ウィーナー過程の形で書かれている。

ブラウン運動（ブラウンうんどう、英: Brownian motion）とは、液体や気体中に浮遊する微粒子（例：コロイド）が、不規則（ランダム）に運動する現象である。ウィーナー過程は離散型であるランダムウォークの極限となる確率過程。（wikipediaより）
ウィキペディア

下記の形で書かれるのが、

X_t = \mu t + \sigma W_t

ドリフト項\muと無限小分散\sigma^2を持つウィーナー過程と呼ばれる（らしい）
ウィキペディア

フォッカー・プランク方程式と常微分方程式

この確率微分方程式(Stochastic Differential Equation / SDE)は常微分方程式の形に書き換えられる。

常微分方程式（じょうびぶんほうていしき、英: ordinary differential equation, O.D.E.）とは、数学において、未知関数とその導関数からなる等式で定義される方程式である微分方程式の一種で、未知関数が本質的にただ一つの変数を持つものである場合をいう。（wikipediaより）
ウィキペディア

こちらの資料にて、フォッカープランク方程式への導出（？）をされていました。

https://speakerdeck.com/misawann/diffusion-model-part-3-dot-6-probability-flow-ode?slide=3

あと、この本にも同じようなことが書いてありました（上のスライドを見て気づきました）。。。

拡散モデル データ生成技術の数理

下記のような確率微分方程式を、

\frac{dx}{dt} = a(x,t) + b(x,t)R(t)

次のように書き換えられるようです。

\frac{\partial P(x,t)}{\partial t} = -\frac{\partial}{\partial x}\alpha_1 (x,t)P(x,t) + \frac{1}{2}\frac{\partial^2}{\partial x^2} \alpha_2(x,t)P(x,t)

フォッカー・プランク方程式（英: Fokker–Planck equation）とは、統計力学でクラマース・モヤル方程式（英語版）においてn ≥ 3 の項のない次の方程式のことをいう。(wikipediaより)
ウィキペディア

確率フロー常微分方程式

最終的には整理されて、

dx_t = [\mu(x_t,t)-\frac{1}{2}\sigma(t)^2 \nabla log p_t(x_t)]dt

の形に。\nabla log p_t(x)はスコア関数と呼ばれているようです。この式は確率フロー常微分方程式、Probability Flow ODEと呼ばれています。

スコア関数は、尤度関数の自然対数の勾配。(wikipediaより)
wikipedia

経験的確率フロー常微分方程式 （Empirical PF ODE）

上記の式をさらに変形したいが、生成AIで利用されるp_T(x)はガウス分布\pi(x)に近いことから、\mu(x,t) = 0 及び \sigma(t) = \sqrt{2}tを代入し、スコア関数をs_\phi(x,t)とすると、

\frac{dx_t}{dt} = -ts_\phi(x_t,t)

とかける。最終的なノイズの分布は、\pi = N(0,T^2I)で初期化される。

PF ODEを解けば答えが出るので良いが、既存のソルバーや元の逆拡散過程を使って画像を復元するにはステップ数がかかってしまうので、論文では一貫性モデルと呼ばれるモデルを提案している。

一貫性モデル / Consistency Models

上記の確率フローODEを解く際に、今回は一貫性関数という関数を仮定している。これは、ある\{x_t\}_{t\in [\epsilon,T]}に対して、

f:(x_t,t) \mapsto x_\epsilon

となるような関数を決める。この一貫性関数は自己整合性(self-consistency)を持っていて、ある一つのPF ODEの軌跡を取った場合、その軌跡上のどの(x_t,t)をとっても最終的に同じ答えに行き着くというものになっている。そのため、任意のポイントから必ず終状態である(x_0,0)に行き着くような関数となっている。有名（？）な模式図が下のようで、軌跡を選ぶと終点が決まる。

これによって、これまで拡散モデルのようにステップを順番に経て最終の絵が出力されるという状態から、1ステップで高速に絵を出力できるようになったのがLatent Consistency Modelとなっている。

そのうち学習などについても触れてみたいと思います。
以上