xiangze's sparse blog

機械学習、ベイズ統計、コンピュータビジョンと関連する数学について

拡散モデルについて思ったこと、統計モデリング等との関係

久しぶりの投稿になります

その間世の中では様々なことが起こりましたがStable diffusion,novelAI, DALL-e, midjourneyなど画像生成AIの流行もその一つです。
画像生成AIの性能の中核とも言える拡散モデルに関して他の機械学習手法との関連、そして多くのアイデアの源泉となった非平衡統計物理学，統計モデリング、進化生物学などとの関係に関して思ったことを書きます。
ためになる論文、本のリンクも貼ります。

目次

拡散モデルの基本的な定義は他のブログや論文に譲る(ためになるリンク参照)としてここでは以下のような分野との関連を書きます。

prompton.io

score matchingとの関係

score matchingは未知の確率分布p(x)に対してそのscore関数 \nabla \log p(x)を推定するというタスクです。分母(分配関数)によらない値をとるという点でマルコフ連鎖モンテカルロ法(MCMC)を用いた(事後)分布関数の推定と共通します。これが所与のデータから計算可能である( 正確には近似との誤差を最小化できる)というのがEstimation of Non-Normalized Statistical Models by Score MatchingのTheorem 1の主張です(証明はAppendixにあります短いです)。この論文にあるようにTweedieの公式などを使ったやや複雑な計算を経るとscore関数を推定するというタスクと等価になります。
このことはA Connection Between Score Matching and Denoisingで指摘されており(「ベイズ深層学習」にもscore matchingの説明がありますp.124 )、また
VAEが提唱される以前からAutoencoderとscore matchingの関係は指摘されていました。(参考 Denoising Autoencoderとその一般化)


他の生成モデルとの関係

拡散モデルは変分オートエンコーダー(VAE)と深いつながりがある、というかVAEを階層化したものが拡散モデルです。
VAEとの違いは拡散モデルは各ステップで次元が変わらない(次元圧縮しない)という点があります。
レビュー論文の6章ではGAN、自己回帰モデル、正規化流など他の深層生成モデルとの関係が述べられています。
「ベイズ深層学習」でも自己回帰モデル、正規化流が詳しく説明されています(6.1)。

非平衡統計物理学との関係

拡散モデルの生成過程ではランジュバン動力学によるサンプリングを使っています。ランジュバン動力学は物理学においては熱雑音と外部ポテンシャルのある環境での粒子の動きのモデルです。拡散モデルの起源となった論文の一つDeep Unsupervised Learning using Nonequilibrium Thermodynamicsでは拡散過程は温度を変えながら最小の自由エネルギーを探索するシミュレーテッドアニーリング(疑似焼きなまし法)に相当するとしています。
ただそこで同様に言及されているJarzynski等式

\< \exp(-\beta W) \> = \exp(-\beta \delta F)
非平衡過程に関する平均=2つの平衡状態の熱力学量
(βは逆温度、Wは系を変化させるために与えた仕事、Fは自由エネルギー、<>は経路に関するアンサンブル平均)

との関連性についてはイマイチよくわかりませんでした()。
統計物理との関連で言えば温度を下げる過程が準静的(ざっくりゆうとゆっくり)でないと過程が成り立たない気がしますが、実務的には計算ステップをはしょる高速化がバンバン進んでいて素朴な物理モデルからは離れつつあるのだと思います。

prompton.io

追記:詳しくは局所詳細釣り合い、ゆらぎの定理、Jarzynski等式と拡散モデルの関係

統計モデリングの手法との共通点、相違点

事後分布の推定手段としてMCMCなどを用いたサンプリングと変分ベイズによる近似計算を使うことができるという点で拡散モデルを含む機械学習の生成モデルとベイズ統計モデリングの計算方法は共通しています。そのためpyroやtensorflow probabilityといった同じフレームワーク で実装することができます。
しかしその目的は異なります。機械学習が入力に対して予測した識別結果(生成モデルの場合は望む出力)を可能な限り高い精度で出力するのが目的なのに対し、統計モデリングは手持ちのデータの性質を尤もらしく説明するの目的です。そのため機械学習ではcross validation のためのデータをとっておいて汎化誤差、過学習の程度を見積もるのに対し、統計モデリングでは手持ちのデータで全てを使ってそれを説明するモデルを探索します。
よくある誤解としては機械学習は統計分析を自動化するもの、ディープラーニングは既存の統計モデルを複雑、大規模化したものというのがあります。

個人的にはここらへんの共通点、相違点については理解が曖昧です。以下のブログの意見を指針としています。

また実装としては統計モデリングで使う変分ベイズではモデルと近似モデル(pyroでいうModelとGuide)で変数の依存関係が同じ方向を向いているのに対し、VAEや拡散モデルではEncoderとDecoderが逆方向を向いているため、ベイズ則を使って向きを揃えてKLダイバージェンスを使ったELBOを導出するという風になっています(
Edwardを試してみた（2）変分オートエンコーダー(VAE) - xiangze's sparse blog
)。ここで拡散が正規分布(ガウシアンノイズ)に従えば向きを変えたものも正規分布になるという性質が効いています(正規分布以外を使ったモデルとしてPoisson Flow Generative Modelsというのがあるらしいです)。
Stan,PyMC等に実装されているHMC,NUTS等を拡散モデルにいれてみるというのもありかもしれません(「ベイズ深層学習」ではランジュバン動力学はHMCと比べるとランダムウォーク的挙動が多くなり分布の探索効率が下がるとしています(p.91)が、拡散モデルを含む生成モデルでは分布ではなく最適値を得られればいいのでこの限りではないかもしれません) 。

情報幾何との関係

拡散モデルの学習は真の分布と近似のFisherダイバージェンスの2乗誤差を最小化するというタスクに落とし込めます。Fisherダイバージェンスは情報幾何的にはリーマン計量に相当します。Fisher行列とKL Divergenceの関係とEMアルゴリズム、変分ベイズ推定について少し - xiangze's sparse blog
情報幾何の枠組みではEMアルゴリズムの計算における反復は指数型分布と混合ガウス分布それぞれに対応する測地線を経て観測データの作る多様体とモデルの多様体との間の射影の繰り返しとして理解することが出来、これをemアルゴリズムとしています。
「情報幾何の新展開」のやばさ - xiangze's sparse blog
そして指数型分布と混合ガウス分布の計量は双対の関係にありました。このように機械学習の最適化は測地線という幾何学的な概念で捉えることができます。

温度(分散)を変化させていく拡散モデルのアニーリング過程が何らかのダイバージェンスで記述できるのかというと 確率分布とは限らない密度関数πi、ρ=log(π )、凸関数fと \Psi_f[\rho ] = \int f(\rho\pi(x)))dx, \tau=f’(\rho)に対して定義されるρ-τ divergence
D_f[ \rho(\pi_i ) // \tau_(\pi_j) ] = \Psi_f [\rho_{\pi_i}] + \Psi_{f^*}[\tau_{\pi_j}] -\int \rho_{\pi_i}(x)\tau_{\pi_j}(x)dx
とそれを混合比βi( \sum_{i=1}^N \beta_i =1 で混合した
I_{f,\rho}[p,q] := \min_r \sum_{i=1}^N \beta_i D_f [ \rho_(\pi_i) // \rho_(r) ] というのを使うらしいです。Rho-Tau Bregman Information and the Geometry of Annealing Pathsのp.11に定義が載っています。
そしてこれは関数ρに対するquasi arithmetic mean
\pi_{\beta}^{(\rho)} := \rho^{-1} ( \sum_{i=1}^N \beta_i \rho( \pi_i(x))) =argmin_{r(x)} \sum_{i=1}^N \beta_i D_f [ \rho_(\pi_i) //\rho_(r) ]
に対応するらしいです…
これをどう解釈すべきか。

prompton.io

進化生物学との関係

タンパク質、RNAは1次元の記号列(1次構造)と見なすことができますが生体内では折り畳まれて酵素などの生物学的機能を発揮します。折り畳まれた後の状態のものを3次構造、その部分を2次構造といいます。
The structure of genotype-phenotype maps makes fitness landscapes navigable | Nature Ecology & Evolution

この論文ではタンパク質、RNAの2次構造に得点(適応度)を割り当てることでどのような構造に進化していくかを数値実験したものでこのような研究はだいぶ昔からされていました(https://pubmed.ncbi.nlm.nih.gov/12447981/)が上記の論文はオープンで大規模な実験データを使っや数値実験で性質を明らかにしたという点で独自性があるのかもしれません。
このような適応度の1次構造に対する関数グラフはfitiness landscapeと呼ばれWaddingtonという人が提唱しました。下図で言うボールが重力に従って下に落ちていくのが生物の個体発生で、下のゴールに行きやすいように地形を変化させるのが生物の進化です。fitiness landscapeは他の人(Sewall Wright)が最初に提唱したもので、Waddingtonが提唱したのは遺伝的同化(genetic assimilation)あるいは運河化(Canalisation)と呼ばれるメカニズムで下図はもともと多細胞生物の細胞分化の文脈で提唱されたものですが、進化によりlandscapeが決まり、個体発生はそのlandscapeをボールが転がるように決定して低いところに行き、その低さを適応度(fitness)によって進化が起こるという点は共通です。

拡散モデルとの関連としては進化(系統発生)が学習、個体発生が生成のプロセスに対応するというアナロジーを考えることができます。

Waddingtonの提唱したfitiness landscape
適応度の高い点AとBが低次元(1次元)では離れているが、高次元(3次元)では"尾根"を介してつながっている図(https://www.nature.com/articles/s41559-022-01867-zより)

「深層学習の原理に迫る」では深層学習が過学習を起こさず高い汎化性能を示すのを説明する仮説として最適な解は高次元のパラメーター空間で平坦に分布している、別の言い方をすれば最適解、それに近い値同士はつながっているという説が紹介されています(ほか2つは初期値周辺に最適解がある、パラメーターの値が狭い範囲でしか動かないという仮説)。これはこの論文の言うところと共通しています。ただし「深層学習の原理に迫る」での仮説は学習時に変化する重みパラメータに関する話で、拡散モデル、進化では生成、個体発生時の値であるところが異なります。
別の違いとしては生物進化の場合は個体が自然選択を受けてlandscapeが変化していくのに対し、拡散モデルでは学習プロセスは生成プロセスとは分離していて一度学習でパラメーターが定まってしまえばそれ以上学習した確率分布関数(landscape)は変化しないという点があります。
またStable diffusionなどでは単一の学習対象(適応度関数)があるのではなく、prompt textなどで決まる数多くの対象を一度に学習するという点も違いがあります。これらの相違点がお互いの理解にどう影響するのでしょうか…

専用ハードウェアへのヒント？

(Nvidia,AMDの)GPU以外のディープラーニング実行の選択肢としてはGoogleのTPUのような例がありCNN用のアクセラレータは数多くあります。学習はサーバーで行い、カメラなどエッジと呼ばれる部分での識別に専用アクセラレータを使うというアイデアが広く受け入れられています。しかしながらCNNに特化したハードウェアではAttension(注意機構)の流行にはついていけませんでした。用途を絞れば使えるかもしれません。
同様に拡散モデル専用計算機を考えることができます。拡散モデルは各ステップ乱数生成のコストが他の生成モデルに比べて高く、高速化の研究ではステップ数を減らすような方法が提唱されているそうです。CNN,Attentionも同様ですが現状ソフトウェアスタックが揃っているNvidiaのGPUに特化した最適化の研究開発が進んでおり、他のアクセラレータ製作者は最先端の研究についていくことは困難です。OpenCLのような中立的ライブラリもありますが、この状況を変えることは難しいのでしょうか。

手元でAIお絵描きとかをする用途に専用回路は不要かもしれませんが、異常検知向けなら応用はあり得るかもしれません。

prompton.io

なぜ拡散モデルはうまくいくのか

Understanding Diffusion Models: A Unified Perspectiveではscore matchingの

  • データ全体が高次元空間の中の低次元な部分に分布している場合(多様体仮説)にscore functionが定義されない
  • データの分布が多峰であった場合に一部の峰にしか到達できない
  • データが希薄な領域においてscore functionがうまく推定できない

といった困難が時間とともに温度を下げていくことで解決するのがうまくいく理由としてあげられています。最初は高温で激しく動き希薄な空間も捉え、また目標とする分布が多峰の場合には最も高い峰の近くに到達します。しかし個人的には「データが低次元に分布している場合にscore functionが定義されない」というのがおそらく温度を徐々に下げていくことで低次元のみを動くようになるということなのでしょうが直感的には理解できません*1。


(複数の峰に対し拡散過程(シミュレテッドアニーリング)を用いることで到達する様子[2208.11970] Understanding Diffusion Models: A Unified Perspectiveのfigure6より引用)

「深層学習の原理に迫る」ではシミュレーテッドアニーリング、ランジュバンによる学習はあまりうまくいかないとされています(107p)。一方拡散モデルでは生成の過程でこれらの手法を使ったサンプリングを使っています。それが性能の高さに結びついているのか、またText promptなどの情報を何度も掛け合わせるような方法、単一の学習対象だけでなくtextで表現されるような多数の画像、さらにCLIPで未知のpromptへも対応していることなどが性能に結びついているのかどうかもよくわかりません(個人的には)。

コンテンツ生成以外の応用

GANと同様に異常例を水増ししたりできるので異常検知が考えられます(これから調査)。

その他

元論文ではLatent diffusion modelの実装のVQVAEの部分は認知できない解像度を圧縮する機能としていましたが、実際には画風変換の機能を持ちうるそうです。
https://twitter.com/_determina_/status/1590525136394149888?s=20&t=-OXxVxb34v-CgdjmcoSEIw

zenn.dev

prompt engineeringの流行や生成過程で探索のようなことをやっているのと同様に様々な部品から構成されているstable diffusionでの各部品の役割、境界が曖昧になっているのは興味深いです。

個人的にはmidjourneyで最初の生成ではpromptの指示には従っていないものの，その画像を初期値としてバリエーション画像の生成を繰り返すと指示に近くなっていくように見える現象を体験し，複雑なLandscapeを探索しているのではないか．またバリエーション画像を生成していくうちに画像がだんだん変化していくのは”尾根”を探索しているのではないかという感覚がありました．

prompton.io

Next Step(ToDo)

  • 「深層学習の原理に迫る」の読書メモ　3つの仮説の整理と参考文献リスト
  • 機械学習の生成モデルと統計モデリングの類似点と相違点 - xiangze's sparse blogを書く
  • 正規化流、自己回帰モデルとの関連(サーベイに書かれている)を読む
  • "尾根"仮説を検証する数値実験を考える(既にあるかも「深層学習の原理に迫る」)
  • "尾根"周りには特異点はあるのか　「代数幾何と学習理論」、「情報幾何学の新展開」で触れられていた
  • 進化生物学論文の理解、適応度の設定について論文を読んでちゃんと理解する
  • 拡散モデルと進化生物学の相違点のまとめ、もう一度整理する
  • ρ-τ divergenceに関わるdivergenceの整理
  • 拡散モデルを情報幾何的に理解する(emアルゴリズム、変分ベイズのように),その利点を挙げる
  • 乱数生成コスト他の見積もり( 専用ハードウェア用、そんなに高くないかもしれない)attensionの高速化も参照
  • 拡散モデルの異常検知応用まとめ

ためになるリンク

日本語でも多くの人が拡散モデルの仕組みを紹介されています。


高速化まで解説されている！

英語

ためになる本

日本語の本はまだあまりありませんがすばやく同人誌を書いた人たちがいます。

prompton.io

Fitness landscapeに関するリンク

普遍生物学 - 東京大学出版会 7章最初にepigenetic landscapeの図が説明されています。

Fitness_landscape定義など
幹細胞の2つの顔を暴け!!! 未分化性維持と特異的分化との狭間で... - 世界の幹細胞（関連）論文紹介 - 慶應義塾大学 グローバルCOEプログラム 幹細胞医学のための教育研究拠点　下の方に論文解説がある。
統合と自己：細胞の個性と個体の個性 | JT生命誌研究館
細胞運命の拘束とワディントンのランドスケープモデル | Proteintech Group
https://www.sciencedirect.com/science/article/abs/pii/S1567134817301193

情報幾何に関する本、論文、リンク

情報幾何と機械学習
情報幾何と機械学習(pdf)赤穂先生
Entropy | Free Full-Text | An Elementary Introduction to Information Geometry無料で読める総説論文としては最も充実してそう
Information Geometry甘利先生による解説、Wasserstein距離について書かれている
https://www.amazon.co.jp/%E6%95%B0%E7%90%86%E7%A7%91%E5%AD%A6-2020%E5%B9%B4-11-%E6%9C%88%E5%8F%B7-%E9%9B%91%E8%AA%8C/dp/B08KFWM7MG
https://www.amazon.co.jp/%E6%83%85%E5%A0%B1%E5%B9%BE%E4%BD%95%E5%AD%A6%E3%81%AE%E5%9F%BA%E7%A4%8E-%E6%95%B0%E7%90%86%E6%83%85%E5%A0%B1%E7%A7%91%E5%AD%A6%E3%82%B7%E3%83%AA%E3%83%BC%E3%82%BA-%E8%97%A4%E5%8E%9F-%E5%BD%B0%E5%A4%AB/dp/4434208810
論文：非平衡化学反応系の情報幾何学構造最近アツいらしいです、伊藤先生の関連研究も注目
[2112.11008] Inferring nonequilibrium thermodynamics from tilted equilibrium using information-geometric Legendre transform
[2209.00527] Geometric thermodynamics for the Fokker-Planck equation: Stochastic thermodynamic links between information geometry and optimal transport


www.youtube.com
ヨビノリさんが解説してる。

広告

「情報幾何の新展開」のやばさ - xiangze's sparse blogもぜひ読んでください。

prompton.io
https://prompton.ioではAIイラストを描いてほしい、描きたい人を募集しています。

*1:いや、"尾根"と関係あるかも