🍉

「最適輸送理論，曲率次元条件と熱分布」の勉強ノートと非平衡統計力学、拡散モデルへの適用に向けて

2023/11/15に公開

math

「最適輸送理論，曲率次元条件と熱分布」の勉強ノートと非平衡統計力学、拡散モデルへの適用に向て

最適輸送研究会に影響を受け、最適輸送と非平衡統計力学、機械学習との関係に興味を持ちました。
最適輸送理論に関してはC. Villaniの「Optimal transport, old and new」が1000pageくらいあるのでその内容を確認するのは用意ではないです。一方で最適輸送とリッチフローに関しては日本語での充実した研究ドキュメントpdfがあります。

最適輸送理論，曲率次元条件と熱分布約100page
最適輸送理論とリッチ曲率 ∼ 物を運ぶと曲率が分かる ∼ 約100page
最適輸送理論，Riemann的曲率次元条件と熱分布 77page
最適輸送理論梗概
最後にKLダイバージェンスを使う情報幾何とWasserstein距離との正規分布族に対する比較があります。
距離コストに対する最適輸送問題について

これらでもそれぞれ100ページ近く有ります。そのなかで曲率次元公式という不思議な形の式が登場し、log Sobolev不等式など幾つかの重要な不等式が成り立つ基礎となっています。他にも自分には見慣れない数学的概念が多く相互の関係を把握するのが難しかったです。そこでここでは「最適輸送理論，曲率次元条件と熱分布」の内容をかいつまんだまとめを記します。

最適輸送理論，曲率次元条件と熱分布の概要(6章まで)

各章の概要です。聞き慣れない概念、定理とそれらの間の関係についてまとめます。基本的に関数は有限のL2ノルム(とおそらくsobolevノルム)を持つものとします。定理の証明、数学的厳密性は「最適輸送理論，曲率次元条件と熱分布」を参照してください。

1章　導入

普通のm次元ブラウン運動 $B_0,B_1$ 間の関係として結合法という関係

$E[|B_0(2\alpha t) - B_1(\beta t)|^2]\leq |x_1-x_2|^2+2m(\sqrt{\beta t }-\sqrt{\alpha t})^2]$

を満たし、これに対する最適輸送費用 $T_{d^2}$ は初期値をデルタ関数 $\delta(x_i)$ とする分布関数 $\nu_i$ を用いて

$T_{d^2}(\nu_s,\nu_t)\leq d(x_0,x_1)^2+2N(\sqrt{t}-\sqrt{s})$

が成り立ちます。これはRicc曲率が正かつ次元がN以下という条件と同値であり、これが一般のリーマン多様体を含むどのような枠組みで成り立つか知りたいのが動機らしいです。

2章　Bakry-Emery 理論入門

Bakry-Emeryの曲率次元条件(CD) (BECD)

Focker-Planck方程式 $\dot{f}=Lf:=\Delta fL<\nabla V,\nabla f>$ と対応するLangevin方程式に対して

$\frac{1}{2}L|\nabla f|^2-<\nabla f,\nabla Lf>\geq K|\nabla f|^2+\frac{1}{N}(Lf)^2$

が成り立つのが曲率次元条件(CD)です。これはあくまで条件であってこれが成り立つには以下のWeitzenbock-Bochnerの定理が成り立つ必要が有りますがリーマン多様体では成り立ちます。

Weitzenbock-Bochner公式

$\forall f\in C^\infty(M) \frac{1}{2}L|\nabla|^2-<\nabla f,\nabla Lf>=||Hess(f)||^2+Ric_V^\infty(\nabla f,\nabla f)$

$Ric_V^N:=Ric+HessV-\frac{1}{N-m}\nabla V \cdot\nabla V$

が(重み付き)リーマン多様体では成り立ちます(証明は今野宏著微分幾何学を参考)。

関数に対する $Gamma$ 計算

$\Gamma(f_1,f_2):=\frac{1}{2}(L(f_1f_2)-f_1L(f_2)-f_2L(f_1))$

$\Gamma_2(f_1,f_2):=\frac{1}{2}(L\Gamma(f_1,f_2)-\Gamma(f_1,L(f_2))-\Gamma(f_2L(f_1)))$

を使うとBECDは

$\Gamma(f,f)\geq K\Gamma(f,f)+\frac{1}{N}(L(f)^2)$

と簡単に書けます。

また熱拡散方程式 $\dot{u(x,t)}=Lu(x,t) , u(x,0)=f(x)$ とその作用素の半群 $P_t=e^{tL}$ を用いるとBECDは

任意の $f, t>0$ で $|\nabla P_t f|^2+\frac{1-e^{-2Kt}}{NK}LP_tf|^2\leq e^{-2Kt}P_t(|\nabla f|^2)$

と同値になります。

特殊で重要な式として任意の $f\in C^\infty(M)$ ,測度mに対して

対数sobolev不等式

$\int f\log f dm\leq \frac{1}{2K}\int \frac{|\nabla f|^2}{f}dm$

が導かれます。これはKL divergenceとFisher行列間の関係

$D(p//q)\leq \frac{1}{2K}I(p,q)$

とも書かれます。機械学習への応用上重要です。

3章最適輸送理論の基礎

基本的な記号と概念

$\phi:Y\rightarrow Z$ にとる測度 $\mu \in \mathcal{P}$ の押しだし $\phi_{\sharp}\mu \in \mathcal{P}(Z)$
測度 $\mu_0,\mu_1$ 間のカップリング $\Pi:=[\pi \in \mathcal{P}(X\times X)|(p_i)_\sharp \pi=\mu_i]]$ (2変数関数 $\Pi$ の片方の引数を固定すると $\mu_i$ になる)
測地距離d
$\gamma:[0,T] \rightarrow Y$ が絶対連続とは $\exists g\forall s,t\in[0,1] d(\gamma_x,\gamma_t)\leq \int g dL^1$

Lipschitz連続、定数

$Lip(f):=sub_{y\neq x}\frac{f(y)-f(y)}{d(y,x)}$

最適輸送(OT)

距離 $d,c=d^2$ 測度 $\mu_0\mu_1$ に対する最適輸送写像 $T_c(\mu_0\mu_1)$ は

$T_c(\mu_0\mu_1):=\inf_{\pi\in\Pi}\int c(x,y)d\pi(x,y)$

Kantorovich双対性

測度 $\mu_0\mu_1\in \mathcal{P}$ に対する

$T_c(\mu_0,\mu_1)=\sup[\int gd\mu_1-\int fd\mu_0] =\sup[\int f^cd\mu_1-\int fd\mu_0]$

( $f^c(y):=\inf_x [f(x)+c(x,y)]$ )となること(ここでは(-f,g)の存在までは言っていない)

Kantorovich-Rubinsteinの公式

$T_d(\mu_0,\mu_1)=\sup[\int fd\mu_1-\int fd\mu_0|f\in Lip_b(X)]$

Kantorovichポテンシャル

$g\circ p_1 -f \circ p_0 (:=\int gd\mu_1-\int fd\mu_0) \leq c$
$g\circ p_1 -f \circ p_0= c , \pi a.e$

のとき(-f,g)はKantorovichポテンシャル

Breinerの定理(最適輸送写像の存在)

最適カップリングのとき $\exist T:X\rightarrow X,\pi=(id \times T)$ ,

また $\exist\phi s.t. T=\nabla \phi$

このTを最適輸送写像またはBreiner写像といいます。

Tには一意性があります。

Mongeの問題

$(id \times T)_\mu \subseteq \Pi(\mu_0,\mu_1)$

の形の最適輸送問題

Wasserstein distance

確率分布関数P,Q、あるいは測度 $\mu_0,\mu_1$ 間の距離として

$W_p(P,Q)^p:=\inf_{\pi\in\Pi(\mu,\nu)}[|d|_{L^p(x)}]=\inf_{\pi\in\Pi(\mu,\nu)}[\int |d(x,y)|^p d\pi(x,y)]$

(dはユークリッド距離など点間の距離)これは最適輸送と関連し中心的な意味を持つ距離です。

カップリング $\Pi(x,y)$ が分解でき

$\mu,\nu\in \mathcal{P}(X), \pi \Pi(mu,\nu),\exist \{\nu_x\}\in \mathcal{P}(X)\forall A\subset X x\mapsto \nu_x(A)$ は可測　に対して

$\forall h \int hd\pi=\int_X\int_Xh(y) d\nu_x d\nu_y$

となります。

接着補題

最適輸送の補間

高次元に測地線を持ち上げることができるということらしいです。これによってポテンシャル $\phi$ がハミルトン・ヤコビ方程式を満たすらしいです。

その他重要な概念

重畳原理
双対補間

4章 Otto解析と熱分布

この章では

熱分布が相対エントロピー汎関数の勾配流になる

$\dot{p}_t=\nabla Ent(p_t)$

という重要な、そして字面だけ見ると統計物理学のようなステートメントを証明してます[^1]。

[^1]熱方程式と勾配流(連続方程式, Focker-Planck)を計量の概念を入れて結びつけるのはある意味Dirac方程式に共通します(こちらは時間空間が一次)

Otto解析

上記のために形式的に以下のような接空間 $T$ , 計量g、測度mをもつ多様体を考えます。

$T_\mu\mathcal{P}:=\{ \nabla \phi | \int|\phi|^2d\mu<\infty\}$

$g(Z,Z'):=\int<Z,Z'>d\mu$

連続方程式(Focker-Planck方程式)

$\partial_tp_t=-\nabla_m\cdot (p_t \nabla \phi_t)$

弱形式では $\forall f\in C^\infty_0(R^m)$ に対して

$\frac{d}{dt}\int fd\mu_t=\int<\nabla f,\nabla \phi_t>d\mu_t$

$L^2$ Wasserstern距離に関するBenamou-Breiner公式 (定理4.1)

$W_2(\mu_1,\mu_2)^2=\inf[\int dt \int |\nabla \phi_t|^2 d\mu_t]$

条件(V)(定義4.2)

測度距離空間 $(X,d,m)$ に対して

$\int \exp(-cd(x_0,x_1))mdx<\infty$

相対エントロピー $Ent_m(\mu)$

測度m, $\mu$ に対して相対エントロピー

$Ent_m(\mu)=\int \rho\log(\rho)dm$
$(\mu=\rho m)$

が定義されます。その測度に対する変換は
となります。エントロピーを意味あるものにするために距離の2乗に関する積分が有限であること $\int exp(-c d(x_0,x)^2)m(dx)<\infty$ を要求しています(Def 4.2)。

エントロピーの勾配流方程式(Theorem 4.5)

$\dot{v_t}=\nabla Ent_m(v_t)$

の解が $v_0$ を初期値とする熱方程式

$\dot{v_t}=\mathcal{L}v_t=\Delta p-<\nabla v_t,\nabla v_t>$

の解と一致することが証明されています

Fisher情報量

$F(p,q):=\int (\nabla Ent_m,\nabla Ent_m)=\int_X \frac{|\rho|^2}{\rho}dm(=4\int_X|\nabla\sqrt{\rho}|^2dm)$

Otto解析によるBakry-EmeryのCD $\Rightarrow$ $Ent_m$ の(K,N)凸性

Otto解析,Bakry-Emeryの曲率次元条件が成り立つとすると

$Hess(Ent_m)-\frac{1}{N}(\nabla Ent_m)^2\geq K$

Otto解析による $W_2$ 収縮性

$s^2(\frac{W_2(\mu_s,\nu_t)}{2})\leq e^{-K(s+t)}s(\frac{W_2(\mu_0,\mu_0)}{2})+\frac{N}{2}\frac{1-e^{-K(s+t)}}{K(s+t)}(\sqrt{t}-\sqrt{s})^2$

ここで $s(x)=\frac{\sin(\sqrt{K/N}x)}{(K/N)}$

これは $N=\infty$ の時

$W_2(\mu_t,\nu_t)\leq e^{-2Kt}W_2(\mu_0,\nu_0)^2$

K=0のとき

$W_2(\mu_t,\nu_t)\leq W_2(\mu_0,\nu_0)^2+2N(\sqrt{t}-\sqrt{s})^2$

と単純化されます。

5章最適輸送理論による曲率次元条件(CD)

曲率次元条件(CD)はN大、K小の方向が条件として弱い

Strum/Lott-Villaniの曲率次元条件 ( $CD(K\infty)$ )

$Ent_m(\mu_t)\leq (1-t)Ent_m(\mu_t)+ tEnt_m(\mu_1)-\frac{K}{2}$

$\mu_0,\mu_1$ 間の任意の $W_2$ 測地線でになりたつ場合は強 $CD(K\infty)$ といいます

entropy powerという汎関数

$U_N:=\exp(\frac{1}{N}Ent_m)$

を定義し、その曲率次元条件が以下になります。

エントロピー曲率次元条件( $CD^e(K,N)$ )

$U_N(\mu_t)\geq \frac{s(1-t)W_2(\mu_0,\mu_1)}{sW_2(\mu_0,\mu_1)}U_N(\mu_0)+\frac{s(tW_2(\mu_0,\mu_1))}{sW_2(\mu_0,\mu_1)}U_N(\mu_1)$

以前出てきた微分不等式は $U_N$ の凹性と同値

$Hess(Ent_m)-\frac{1}{N}(\nabla Ent_m)^2 \Leftrightarrow Hess(U_N)\leq -\frac{K}{N}U_N$

であることがわかります。

ここでまた別のエントロピー(Renyi(-Tsallis)エントロピー)

$S_N(\mu):=\int \rho^{1-\frac{1}{N}}dm$

( $\mu=\rho+m$ )を定義すると

Strum/Lott-Villaniの曲率次元条件( $N<\infty$ )

$S_{N'}\leq \int_{X\times X} \Big( (1-t)\frac{s'((1-t)d(x_0,x_1))}{(1-t)s'(d(x_0,x_1))}^{1-\frac{1}{N'}}\rho_0^{-\frac{1}{N'}}(x_0)+t\frac{s'(td(x_0,x_1))}{s'(d(x_0,x_1))}^{1-\frac{1}{N'}}\rho_1^{-\frac{1}{N'}}(x_1) \Big)d\pi(x,y)$

ここで $s'(x)=\frac{\sin(\sqrt{K/N'}x)}{(K/N')}$

reduced CD( $CD^*(K,N)$ )

$S_{N'}\leq \int_{X\times X} \Big( \frac{s'((1-t)d(x_0,x_1))}{s'(d(x_0,x_1))}^{1-\frac{1}{N'}}\rho_0^{-\frac{1}{N'}}(x_0)+\frac{s'(td(x_0,x_1))}{ts'(d(x_0,x_1))}^{1-\frac{1}{N'}}\rho_1^{-\frac{1}{N'}}(x_1) \Big)d\pi(x,y)$

測度収縮性(MCP)

$\forall x\in X$ ,可測な $A\subset X,m(A)$ , $A\subset B_{\sqrt{\frac{N-1}{K\pi}}}(x)$ に対して $\exist \delta_x,\Xi$ (カップリング) s.t

$m\leq(e_t)_\sharp\big( t\frac{s(td(\gamma_0,\gamma_1))}{s(d(\gamma_0,\gamma_1))}^{N-1}m(A)\Xi(d\gamma)$

のとき $(X,d,m)$ はMCPを満たす

曲率次元条件間の関係

$CD(K,N) \Rightarrow CD^*(K,N) \Rightarrow MCP(K,N)$
$K\geq 0$ のとき $CD^*(K,N) \Rightarrow CD((N-1)K/N,N)$ 特に $CD(0,N)\Rightarrow CD^*(0,N)$
$(X,d,m)$ が本質的非分岐 $\Rightarrow (CD^*(K,N) \Leftrightarrow CD^e(K,N))$
強 $CD^*(K,N) \Leftrightarrow$ 強 $CD^e(K,N))$

6章 $L^2$ -Wasserstein 空間上の勾配流

疑問「熱分布に対する $W_2$ 収縮性」を最適輸送の考えから導出できるか、より詳しくは

$(\mathscr{P_2}(X),W_2)$ 上で相対エントロピー $Ent_m$ の勾配流をどう定式化するのか
勾配流はCDe(K,N)条件化で $W_2$ 収縮性をみたすか
その勾配流は熱方程式の解と同一なのか？(以下のJKO schemeを使った導出が基礎にある)

という疑問に分かれます。今まで挙げた曲率次元関係間に幾つかの関係が成り立ちます。

JKO scheme

Fokker-Planck方程式

$\dot{p}=\nabla\cdot(\nabla U(x)p)+\beta^{-1}\Delta p$

の解を(Wasserstein)距離dに対して $d(p^{(k)},p)^2$ を最小化するようなものとみなし反復法でその解を求めようとする方法

エネルギー消散等式(EDE)

曲線 $\{\eta_t\}$ に対して

$-\frac{d}{dt}U(\eta_t)=\frac{1}{2}|\dot{\eta}|(t)^2+\frac{1}{2}|\nabla_U|(\eta_t)^2$

$|\nabla_U|(y):=\bar{\lim}_{z\rightarrow y}\frac{[U(y)-U(z)]_+}{d_Y(y,z)}$

発展変分不等式(K-EVI)

$\frac{1}{2}\frac{d}{dt}d_(\eta_t,z)^2+\frac{K}{2}d_Y(\eta_t,z)\leq U(z)-U(\eta_t)$

を満たす $\{\eta_t\}$ が初期条件 $\eta_0$ とするUのK-EVIの勾配流と呼ばれる

諸関係

EDEの解、K-EVIの解はUの勾配流
Uの勾配流で $Hess U\geq K\Rightarrow$ K-EVIの解
K-EVI ⇒ 収縮性
EVI勾配　= EDE勾配　で一意
$CD(K,\infty) \Leftrightarrow$ K-EVI on (M,d,m)

7章 Riemann 的曲率次元条件と熱分布

ディリクレ形式(エネルギー汎関数)

$E(u,v)=\sum_i \int \frac{\partial u}{\partial x_i}\frac{\partial v}{\partial x_i}$

との関係が述べられています。

8章応用と関連する話題

曲率次元条件が導く色々な不等式を紹介するにとどめます。以下の展望と関係します。

HWI不等式

$CD(K,\infty)$ で

$Ent_m(\mu_0)\leq Ent_m(\mu_1)+W_2(\mu_0,\mu_1)\sqrt{I_m(\mu_0)}-\frac{K}{2}W_2(\mu_0,\mu_1)^2$

log Sobolev不等式

$Ent_m\leq\frac{1}{K}I_m(\mu)$

Talagrand不等式

$W_2(\mu_0,\mu_1)^2\leq \frac{2}{K}Ent_m(\mu)$

大域ポアンカレ不等式

$\int_X|f-\frac{1}{m(X)}\int_X fdm|^2 \leq \frac{1}{K}\int|\nabla f|^2_*dm$

N-log Sobolev不等式

$KN[\exp(\frac{2}{N}Ent_m(\mu))-1]\leq I_m(\mu)$

N-Talagrand不等式

$Ent_m(\mu)\leq N\log\cos(\sqrt{\frac{K}{N}}W_2(\mu,m))$

9章 RCD空間上の解析とに関する幾つかの話題

気になったトピックのみ挙げます

リッチフローへの応用
自己改良
$W_2$ 収縮の拡張
Wエントロピー(Perelmanがポアンカレ予想解決で提唱)
剛性

非平衡統計力学、機械学習、拡散モデルへの展望

非平衡統計力学においては(相対)エントロピー生成 $\sigma$ が任意の関数に対して熱力学的不確定性(TUR)の不等式を満たします。TURはCramer-Raoの不等式の経路版と言えCauchy–Schwarzの不等式から導出されるので上記の不等式とは異なります。上記の不等式をそのまま非平衡統計力学に適用すると異なる関係が出てくると言えるでしょうか。それはGeometrical aspects of entropy production in stochastic thermodynamics based on Wasserstein distanceではL2-Wasserstein距離の満たす三角不等式から熱力学的速度限界がTURとは独立に導出されています。Taragrand不等式に対応しているようにも見えます。TURはエントロピー生成の微分、熱力学速度限界はエントロピー生成そのものの制約です。曲率次元公式はこれより厳しい関係になっているように見えます。
数値計算、実験の値から実効的な次元Nやリッチ曲率Kの上限が求まるかもしれません。拡散モデルに対しても同様のことが言えるはずです。

一般の機械学習の場合は重みパラメーターxはLangevin方程式

$\dot{x}=-\nabla_x U(x)+f(x)-\gamma x+\sqrt{\beta_t}dW_t$

に従って学習されるとモデル化されます。学習データに起因するノイズ部分 $\beta_t$ をうまく取り扱う必要がありそうです。
下記の論文のようにより条件を絞ったTanagrand不等式やlog Sobolev不等式の成立を考えるだけでも意義深いです(Tanagrand不等式が測度限界に対応しているように見えます)。

非平衡統計力学、機械学習と最適輸送の関係

以下に関連論文を挙げます

Geometrical aspects of entropy production in stochastic thermodynamics based on Wasserstein distance
Thermodynamic Unification of Optimal Transport: Thermodynamic Uncertainty Relation, Minimum Dissipation, and Thermodynamic Speed Limits
Score-based Generative Modeling Secretly Minimizes the Wasserstein Distance

KLダイバージェンスに基づいて作られた拡散モデルでもWasserstein距離を最小化できるということを導出しています。log Sobolev不等式やTalagrand不等式と関係していそうですが、論文を見るとどうも自力で導出しているようです。

W-GANでも曲率次元関係相当のことが考えられるはずです。

その他の参考文献

既存の情報幾何(Fisher幾何)に対するWasserstein幾何に関連するものとして

幾何学的話題

https://ymatz.sakura.ne.jp/journal/wp-content/uploads/2021/01/20180225-mathcafe-slide.pdf　Rauchiの比較定理に関して
リッチフローと幾何化予想
Ricci Flow and the Poincare Conjecture

確率不等式

確率不等式
Hoeffdingの不等式
Bernsteinの不等式
Rademacher複雑さとDudley積分
VC次元
局所Rademacher複雑さ
高次元推定問題におけるLasso
N-Talagrand不等式
Minimax最適性
など、本稿とかぶっているのはありません
統計的学習理論概説

Talagrandのcenentration inequalityは非常に汎用性の高い不等式である．実際，局所Rademacher複雑さを用いたfast learning rateの導出に用いられている．

対数 Sobolev 不等式と Talagrand 不等式の改良について

ユークリッド空間上の正規分布に基づくこれらの不等式の改良について考える．とくに，インプットする分散がある意味で大きい場合に最適な形で不等式が改良されることを報告する

defict $\delta$ を不等式の両辺の差として定義する

$\delta_{LSI}(\mu):=\frac{1}{2} I_\gamma(\mu) − Ent_\gamma(\mu)$

$\delta_{Tal}(\mu):=Ent_\gamma(\mu) -\frac{1}{2} W_2(\mu,\gamma)^2$

deficit の非自明な下界を評価する

分散が小さい場合の改良

$0 \lt \beta \le 1$ を満たすとし $\mu \in P^2_{ac}(R^n;\gamma)$ に対し $cov(\mu)\le \beta I_n$ ならば

$\delta_{LSI}(\mu) \ge \delta_{Tal}(\mu)=\frac{n}{2}(\log \beta-1\frac{1}{\beta})$

$\delta_{Tal}(\mu) \ge \frac{n(2(1-\beta)+(1+\beta)\log\beta)}{2(\beta-1)}$

分散が大きい場合の改良

$\mu \in \mathcal{P}^2(R^n), d\mu(x):=f(x)dx$ である確率密度関数fに対し

$\nabla^2 \log f \ge −\frac{1}{\beta}I_n \Rightarrow cov(\mu) \ge \beta I_n$
$\nabla^2 \log f \le −\frac{1}{\beta}I_n \Rightarrow cov(\mu) \le \beta I_n$
$Delta \log f \ge −\frac{n}{\beta}I_n \Rightarrow Tr[cov(\mu)] \ge n\beta$