👺

拡散モデルにおける熱力学的不確定性

2023/06/21に公開

拡散モデル

統計物理

非平衡統計

idea

に続いて

の3章とAppendix Bに基づいて非平衡熱力学系の熱力学的不確定性とその拡散モデルにおける対応を考えていきます。
実際には論文をかなり参考にしなければいけませんでした。

熱力学的不確定性(TUR)とは

変数xの関数である統計量Θ(x)とエントロピー生成 $\Delta S$ に対して
$\frac{<\Delta \Theta(x)^2>}{<\Theta(x)>^2} \geq \frac{2}{\Delta S}$
(<>はx平均) となる関係で非平衡な状態でも成り立つことから注目されています。
以下ではUncertainty relations in stochastic processes: An information inequality approachに基づき拡散過程に似たLangevin系の場合の導出を紹介します。

導出

パラメーター $\theta$ と位置座標x(一般には多次元),時間tのLangevin系
$\dot{x}=A_\theta(x,t)+C(x,t)\xi(t)$
とxの関数である統計量Θ(x)に対するCramer-Raoの不等式
$\frac{Var_\theta [\Theta(\Gamma)]}{(\partial_\theta<\Theta(\Gamma)>)^2} \geq \frac{1}{I(\theta)}$
( $I(\theta)$ はFisher情報行列, $\Gamma={x(t)}_{t=0}^T$ は経路の変数)を出発点とします。

ここで定常状態(stable state,ss)(分布関数 $P^{ss}$ )に対して
$A_{\theta,i}:=(\theta+1)A_i(x)-\frac{\theta}{P^{ss}}\sum_j \partial_{x_j} B_{ij}(x)P^{ss}$
と天下り的に定義します。またxに依存した量 $\Lambda(x)$ の流れ
$\Theta_{cur}(x):=\int_0^T\Lambda(x)^T\cdot \dot{x}dt$
を考えると定常状態のカレント(流れ) $J^{ss}(x)$ を用いて
$=T\int dx \Lambda(x)^T J^{ss}_{\theta}(x)$ (定常性)
$=T\int dx \Lambda(x)^T (\theta+1)J^{ss}(x)$ (θの微小性)
$=(\theta+1)j$
となります。ここで平均カレントjを
$j:=<\Theta_{cur}(x)>_{\theta=0} =T\int dx \Lambda(x)^T J^{ss}(x)$
と定義しています。

経路 $\Gamma$ に対する確率分布 $P(\Gamma|x^0)$ は経路積分を用いて
$P(\Gamma|x^0)=N\exp[-\int_0^T dt A_{\theta}(x_t,t)]$
$A_{\theta}(x_t,t):=\frac{1}{4}(\dot{x}-A_\theta)^T B^{-1}(\dot{x}-A_\theta)^T )$

と書かれます。これからFisher情報行列
$I(\theta):=<(\partial_\theta \log P_\theta(\Gamma))^2 >=-<\partial_\theta^2 \log P_\theta(\Gamma)>$
は
$\log P(\Gamma|x^0)=\log N+\log P_\theta(x^0) -\int_0^T dt\frac{1}{4}(\dot{x}-A_\theta)^T B^{-1}(\dot{x}-A_\theta)^T)$

$\partial_\theta^2 \log P(\Gamma|x^0) = \partial_\theta^2 \log P(x^0) -\frac{1}{2}\int dt(\partial_\theta A_\theta)^TB^{-1}(\partial_\theta A_\theta) +\frac{1}{2}\int dt (\dot{x} - A_\theta)^T \bullet B^{-1}(\partial_\theta^2 A_\theta)$
(・は伊藤積分)であることから

$I(\theta)=-<\partial_\theta^2 \log P_(x^0)>_\theta +\frac{1}{2}<\int_0^T dt (\partial_\theta A_\theta)^TB^{-1}(\partial_\theta A_\theta)>_\theta$
となり、

$I(0)=\frac{1}{2}<\int_0^T dt(\frac{J^{ss}(x)^T}{P^{ss}(x)}) B(x)^{-1}(\frac{J^{ss}(x)}{P^{ss}(x)})>$
$=\frac{T}{2}\int dx \frac{J^{ss}(x)^T B(x)^{-1}J^{ss}(x) }{ P^{ss}(x) }$ (定常性から)
と書けます。

ここで系全体のエントロピー生成 $S_{tot}$ を
$\Delta S_{tot}:=T\int dx\frac{J^{ss}(x)^T B(x)^{-1}J^{ss}(x) }{ P^{ss}(x) }$
と定義すると
熱力学的不確定性(TUR)の不等式
$\frac{Var_{\theta=0}[\Theta_{cur}(\Gamma)]}{j^2} \geq \frac{1}{\Delta S_{tot}}$
が得られます(平均カレントは $j=\partial_\theta<\Theta_{cur}(\Gamma)>_\theta$ と書けることから)。
この $A_{\theta,i}$ と $S_{tot}$ の定義はFluctuation-response inequality out of equilibriumで対応する物理的な力がない摂動としてvirtual perturbationと呼ばれている考えに基づいています。

ただこのエントロピーの定義は唐突で解釈が難しいのでここで等式が成り立つ場合を考える。するとコーシー・シュワルツの不等式が等式になる条件から分布関数は
$\partial_\theta \log P_\theta(\Gamma) =\mu(\theta)[\Theta(\Gamma)-\phi(\theta)]$

という条件を満たされなければいけないことがわかります(定数 $\mu(\theta)$ はスケーリング係数の役割を果たす)。この場合には(詳細は論文)
$\Theta_{tot}:=\int^T_0 dt \frac{1}{P^{ss}B(x)}\circ\dot{x}$
に対してエントロピー生成は
$\dot{s}_{tot}:=A(x)^TD\circ\dot{x}-\frac{d}{dt}\log P^{ss}(x)=\frac{J^{ss}(x)D^{-1}}{P^{ss}(x)}\circ\dot{x}$
と書け、第1項は熱力学的( $\frac{\dot{q}}{D}$ )、第2項は情報理論的解釈ができ
のである条件ではエントロピー生成として解釈できると言えます。

一方Fluctuation-response inequality out of equilibriumではまずパラメーターaと観測量hに関するキュムラントの母関数の式
$K^a_{\Delta \Theta}(h):=\log<\exp^{h(\Theta-<\Theta>_a)}>^a=K^a_{\Theta}(h)-h<\Theta>_a$
からJensenの不等式を使ってKLダイバージェンスとの間の不等式を導き、そしてある経路 $\Gamma={x}$ に対する確率分布 $P(\Gamma)$ とその逆経路の分布 $P^\dagger(\Gamma^\dagger)$ 間のKLダイバージェンス $D_{KL}=\int d\Gamma P(\Gamma)\log\frac{P(\Gamma)}{P^\dagger(\Gamma^\dagger))}$ がエントロピー生成に相当するという論理で

$\sqrt{2}<\Theta>\leq\sqrt{\Delta S<\Delta \Theta^2>}\inf_{a>0}[\frac{1}{2a}+\frac{1}{2a\Delta S}K_{\Delta r}(-\sqrt{\frac{2\Delta S}{<\Delta \Theta^2>}a})]$

というより一般的な式を導出しています。

「非平衡統計力学」では特定の物理系(マルコフジャンプ過程、Langevin系)の場合についてCramer-Rao不等式を用いて証明しています。この他「数理科学2020年 11 月号　情報幾何の探求」では時間tをパラメーターとみなしてFocker-Plank方程式を変形させた同一の式
$\frac{d<D>_t}{dt}=\int_{-\infty}^{\infty}dx(D(x)-<D(x)>_t) p(x,t)\frac{\partial \log p(x,t)}{\partial t}$
$=\int dx \frac{\partial D(x)}{\partial x}v(x,t)p(x,t)$
$(v:=F(x;t)-kT\partial_x \log p(x;t)$ 　Fは外力)
からコーシー・シュワルツの不等式
$(\int dxf(x)g(x) )^2\leq(\int dx f(x))^2(\int dx g(x))^2$
を使って1つ目の式からCramer-Raoの不等式、2つ目の式から熱力学不確定性を導出しています。

解釈

ある場合、等式が成り立つ場合にエントロピー生成に物理的意味づけができることは上に書いたとおりです。
またFluctuation-response inequality out of equilibriumではある経路に対する分布関数と逆過程の経路の分布関数間のKLダイバージェンスがエントロピーとして解釈できることから不確定性関係の高次の項を足す形が導き出されています(式(30),(31))
機械学習、AIの学習、生成過程においては原理的にはすべての変数は計算機の中にあるので見ることができます(GPUからCPUに値を移さないほうがずっと高速に計算できますが)。なので右辺から左辺あるいは逆に推定することはあまり意味がないかもしれません。しかしエントロピーは分布の形状に依存し、KLダイバージェンスは変形前後の確率分布間の差異です。変数x(画像や画像を入力としたVAEの出力である隠れ変数)の移動の分散がエントロピーによって抑えられるというのは生成サンプリングと分布の間の関係であり自明でないようにも思えます。
そのため物理系と同様に意義があるかもしれません(わかりやすいマンガ )。
一方で拡散モデルでは時間と一緒に関数形(パラメーター)が変化するので式変形は異なり、非平衡定常過程ではありません。これが関係の導出に影響を与えそうです。

TURはいつ満たされるか

TURが成り立たない場合としてUncertainty relations in stochastic processesではモデルが非線形な場合、Fluctuation-response inequality out of equilibriumでは離散マルコフ連鎖(MCMCの詳細釣り合い条件は満たさない)が挙げられています。ただし前者ではリミットサイクルの力学系の数値計算でもかなりの範囲で近似的に成り立ち、後者ではTURの一般化した式が厳密解を満たすことが図から見てとれます。拡散過程、逆拡散過程は非平衡非定常であり、リミットサイクルも長時間平均が取れます^[1]。あるいは長時間平均を取らないと意味がないのか？

補足非定常な場合のTUR

「非平衡統計力学」によると統計量 $\Theta$ に対して　
$\sigma\geq \frac{(\tau\Theta(\tau) -v\partial_v<\hat{\Theta}>)^2}{<\Delta\hat{\Theta}^2>}$
という形で書かれるそうです。導出は

Thermodynamic Uncertainty Relation for Time-Dependent Driving
Current fluctuations and transport efficiency for general Langevin systems
などに書かれているそうです。後者の議論はキュムラント母関数から始まるのが共通しています。

展望、やるべきこと

パラメーターと時間がともに変化する場合の式の導出
数値実験 diffusersを使いたい、どうにかして非線形性の影響、TURの破れを見る。
DDIM, PNDM,DPMSolverなど様々な拡散モデルの生成方法の違いあるいは等価性の説明
速度限界の導出
学習過程への適用無限次元ランジュバンを用いる？
「非線形統計力学」4章ではマックスウェルの悪魔についての説明、つまり情報論的エントロピーと熱力学的エントロピーの統一的扱いについてさらに詳しく説明されている。「ディープラーニングと物理学」の最初で触れられているように”マックスウェルの悪魔が絵を描く”状況担ってしまったのでその様子を非平衡熱力学の言葉で理解したい。
参考にした本、論文ではTURの導出とCramer-Raoとの関係が全て微妙に違うように見えます。「非平衡統計力学」によると大偏差原理を使った方法と情報理論的な方法があるらしいです。統一的な理解がしたいです。

参考書、論文

Uncertainty relations in stochastic processes: An information inequality approach
Multidimensional thermodynamic uncertainty relations　多次元ランジュバン系でのTURが述べられている。
Fluctuation-response inequality out of equilibrium
数理科学2020年 11 月号　情報幾何の探求

Cramer-Rao不等式の証明に関しては普通の統計学の本

や情報幾何学の本

を参照しました。後者には大偏差原理と情報幾何の関係も少し書かれています。

さらなる情報幾何の話

拡散モデルについては

さらなる参考文献

脚注

論文ではリミットサイクルの周期未満のデータで不確定性に一致する結果が得られています(FIG2 b ↩︎