拡散モデル基礎： 拡散項をドリフト項に読み替えるのがミソ

### 拡散モデルって何？
 拡散モデルでは、$N$次元空間におけるデータ点の集合(解集合）を生成することを考える。確率的に与えることにすると確率分布$p({\bf x})$からベクトル${\bf x} $をサンプリングする問題である。ただし${\bf x} \in R^{\rm N}$. 
 * 現実的には、$p({\bf x})$もサンプルの集合（以下、解集合とよぶ）$A=${${\bf x}_0,{\bf x}_1,...,{\bf x}_n$}に対して$p({\bf x})=\frac{1}{n}\sum_i\delta({\bf x}-{\bf x}_i)$として与えられている。それゆえ現実的には$A$に類似するものをサンプルする課題であると言える。
 * 条件付きサンプリングにも拡張される（たとえば条件指定した絵を書くなど）。

[拡散モデル、データサイエンスのための確率の知識](https://qiita.com/takaokotani/items/11322dbd3c73a16432ed)

### 分子動力学的な例題を考える
一定の箱の中に複数の原子を閉じ込めておいて安定解を探す分子動力学的な問題を考える。一定の周期性を仮定するとする。原子間力により原子の並びは規定されるが、いろいろな多形が許されうる状況を考える。たとえばゼオライトのようなものを想定し（Siのみを考える）十分大きい箱を考えておくとその中に整合性をもってきちんと収まる周期的なパターン（結晶）というのは多少のずれを許容するなら多数あり得る。

高温にするとランダム分布になり、アニーリングにより高温でのランダムな構造が、温度ゼロでは、秩序をもつ結晶になるということである。

これを数値的にシミュレーションする。物理学的手法では原子間力を仮定して、力に比例した原子の移動を繰り返すことによりアニーリング。一方、拡散モデルもアニーリング方法の一種であるが、原子間力でなく、温度ゼロでの結晶の多形を表す確率分布$p({\bf x})=\frac{1}{n}\sum_i\delta({\bf x}-{\bf x}_i)$を与えておいて、それに戻るようにアニールしていくことになる。最初に与えたどれかの${\bf x}-{\bf x}_i$にもどるだけでは面白くないが、そこの点を工夫すれば類似のものも生成できる。

[拡散モデル、ゼオライト](https://ecalj.sakura.ne.jp/talk/MRSJ2024/mrsj_takano.pdf) まだあんまりうまくできない

### 連続的な拡散モデルからの導出
「ランジュバン方程式と拡散方程式の関係,ドリフト電流＋拡散電流、拡散とランダムウォーク」をだいたい知ってる人には
「連続的な拡散モデル（拡散方程式）を考えてから逆拡散方程式を考える理論」のほうが、わかりやすい。以下それを説明する。
* まず
```math
dx= f(x,t) dt + g(t) dw　　　　\label{eq:1}\tag{1}
```
のドリフト項(速度場）$f(x,t)$つき拡散のランジュバン方程式から出発する。ここで$dw$は正規化ランダムノイズ、$g(t)=\sqrt{2D(t)}$である。$D(t)$が拡散係数である。$x$はベクトルである。

(1)式と等価な拡散方程式（あるいは、移送方程式、フォッカープランク方程式）は、$dp(x,t) \equiv p(x,t+dt)-p(x,t)$として、
```math
d p = \nabla \left(- p f + D \nabla p \right) dt　　　　\label{eq:2}\tag{2}
```
と書ける。($\ref{eq:2}$)式の拡散方程式は確率（総和が１）のダイナミクスの方程式である。($\ref{eq:1}$)式を解いたランダムモーションを多数回おこなって各時刻で考えると、確率$p({\bf x},t)$の分布となっている。
電流＝拡散電流＋ドリフト電流、により電子密度の変化が起こる、という式である。

 * (２)式を変形して
```math
d p = \nabla \left(- p f + D \nabla p  \right) dt\\
= \nabla \left(-p (f - D \nabla\log p ) \right) dt
```
と書く。ここで$f - D \nabla\log p $をドリフト項であると見たててランジュバン方程式とするなら、**拡散項がなくなり、ドリフト項のみの方程式**となり、
```math
dx= (f(x,t) - \frac{1}{2} g^2(t) \nabla\log p(x,t))dt　\label{eq:3}\tag{3}
```
となる。

速度場を与える$\nabla \log p(x,t)$はスコア関数と呼ばれる。もしも、これが既知であれば、($\ref{eq:3}$)式は単なる速度場の方程式であり、時間軸を逆にした方程式を考えることができる。これが拡散モデルの考え方である。ランジュバン方程式の立場における、拡散項をドリフト項$\nabla\log p$（速度場）によるものであると考えるのである。
そもそも初期集合が与えられているなら、拡散方程式によって$p(x,t)$は決定されているはずのものであり、$\nabla\log p$は計算可能なものである。拡散はガウシアンノイズを加える過程なので、有限の$t$でもやはりガウシアンノイズを加える過程になる（ガウシアンの畳み込みはガウシアン）。

実際にどうやるのか？は以下。

### 参考
#### スコア関数の学習の仕方
$\nabla \log p({\bf x},t)$は損失関数$L=||s_\theta({\bf x},t)-\nabla\log p({\bf x},t)||$(二乗和）を最小化するようにモデル関数${\bf s}_\theta({\bf x},t)$のパラメータ$\theta$を最適化することで機械学習させておく。
これは単純なサンプリング和の最小化手法で可能である.最適化のための損失関数$L$は(以下$t$は固定して考えている）
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/482779/e53b4d0c-cea1-f6f4-e65b-6e110a98dcd2.png)
(ここで$C_1,C_2$は$\theta$によらない定数。$p({\bf x},{\bf x}')$は同時確率)
ミソは２つ目の式が$p$の線形項のみである、という点である。なのであえて$p({\bf x},{\bf x}')$という[同時確率](https://qiita.com/takaokotani/items/11322dbd3c73a16432ed#%E6%9D%A1%E4%BB%B6%E4%BB%98%E3%81%8D%E7%A2%BA%E7%8E%87)に置き換えて${\bf x}'$で積分してやるというように変形しても構わない。モデル関数は解析的であり（pytorchに実装されている線形関数、活性化関数などを組み合わせて作る）、$L$の最後の式($C_2$除く）の$p({\bf x},{\bf x})$は以下の簡単な関数である。被積分関数が${\bf x},{\bf x}'$について簡単に定まる関数であるのでサンプリングで評価でき、自動微分による緩和で$\theta$を最適化できる。(参考：岡野原「拡散モデル」p.26（ミソがわかりにくい）)

$p({\bf x},{\bf x}')$: 
順拡散のときの$p_t({\bf x})$を時刻0での$p_0({\bf x})$の分布（サンプルの分布関数）がガウシアンで拡散していくものであるとするので$p_t({\bf x},{\bf x}')\propto \exp (-|{\bf x}-{\bf x}'|^2/(2\sigma(t)^2) p_0({\bf x}')$を用いる。

#### ノイズを付加した逆拡散
(２)式を変形して
```math
d p = \nabla \left(- p f + D \nabla p \right) dt\\
= \nabla \left(-p (f - (D＋D') \nabla\log p )
\right) dt
- \nabla \left(D' \frac{\partial p}{\partial x}\right) dt
```
と書くこともできる。これは、時間反転させた時に$D'$という拡散項を導入した標識。ドリフト項を$(D+D')/D$と大きくしておいた分、ノイズを加えてサンプリングすることになる。$\nabla\log p$は特異な点を多く含むだろうから、ある程度なめらかに動かしたいならこういうのもありかと思われる。

#### モデルのとり方
実際的にはモデルのとり方が問題になる。モデルに局所性がある場合が多い。すなわち近傍構造を見て中心位置を動かす、というようなスコア関数である。このような場合には、最初のサンプル空間でなく、それをも含むもっと大きなサンプル空間へのサンプリングが行われることになる。精度が落ちてしまっていると見ることできるが、多様なものを生成できている、と見ることもできる。


#### 拡散モデルの他の導出法
https://arxiv.org/pdf/2006.11239
が著名論文で拡散モデルをこの導出法にしたがって解説している文献も多い。が、結局は拡散のステップごとに逆拡散ステップを対応させることになり、同様の結論に到達する。離散化したレベルの$L$を考えているので、差分化方程式を導出するときの不定性はない。ただ、離散化したレベルの$L$はいろいろあり得るので方程式もいろいろになる(DDPM,DDIMの違い）。
(学習時の時刻ごとの重みが云々という話もあるがぼくはいまいち理解できてない)

#### 文献
確率フローODE：https://arxiv.org/pdf/2011.13456.pdf
の式(13)（論文https://arxiv.org/pdf/2202.00512.pdfも参照）
がその基礎づけとなる。論文の付録Dに確率フロー方程式の導出があるーGにx依存性がないとして、単純化したバーションを考えるとよい。
それが１式になる。



<!---
$p(x,t=0)$を既知と仮定し（解集合を確率分布として与えておく）、$g(t)$を決めておけば$p(x,t)$は一意的に決定されている量である。$\nabla \log p(x,t)$も一意的に決定されている。

$\nabla \log p(x,t)$は$\epsilon_\theta({\bf x},t)$と等価なものである。 

1の場合は、逆拡散の方法を決定するため、A.ELBOをとる、B.ガウシアンによる逆拡散kernelの形、の2つを仮定することになる。

学習過程においては、アニールで解集合に到達するように各温度ごとに、どのように原子を動かすか？の$\epsilon_\theta({\bf x},t)$関数を学習させる。$\epsilon_\theta({\bf x},t)$は原子に働く力であるが、これを解集合から直接に学習させることになる。運動方程式でなくエネルギーの緩和なので力に比例した形で変位をさせる。
結晶データベースにある結晶構造でトレーニングして、高温からのアニールで結晶が再現できれば成功である(MatterGenなど）。

温度ステップと言ってもいいが、通常は、時間$t$と呼び、$t=T$から$0$までたとえばT=1000ステップで逆拡散させる、という形式を想定するのが通常である。

この関数$\epsilon_\theta({\bf x},t)$は、パラメトライズして（Modelとして）あたえることになる（２次元画像ではUnetなど）。pytorchではモデルを与えるのは容易であり、自動最適化ができる。$\epsilon_\theta({\bf x},t)$は局所性をもつ量であるー自分の周囲の原子をみて自分に働く力を決めるという関数。MatterGenでは原子種をも含めての拡散過程を考えている。

現在,生成AIの主流は、大きく画像生成AIと、トランスフォーマー（特にGPTなどのシーケンシャルトランスフォーマー）に分類される。画像生成AIでは、拡散モデル＋GANs+オートエンコーダーを繰り返し適用して画像生成を行うとのことである。

* ただ、この定式化において雑に離散化すると目的関数が失われうる（離散化した目的関数がない状況）。実際、具体的な方程式導出においては離散化した後でも目的関数が存在している１の定式化に優位性がありうる（必ずしも通常のものでなくても良いがどういう可能性があるかはよくわかってない）。ネーターの定理を考えればいいが目的関数の対称性が保持される場合、有利性がある。離散化したレベルで保存量をもつ方程式であれば（シンプレクティック数値積分法のような）、おかしな発散をおこさずに安定な解を得ることが期待できる。分散保存型拡散なら、離散化していても分散が保存される積分法であるべき。拡散方程式を変分的に書いて離散化したのちに逆拡散方程式を導いたほうが良いようにも思われる。

* 逆拡散過程では$\nabla \log p(x,t)$が与えられていればランダム分布から出発して$p(x,t=0)$を作ることができる。もし$\nabla \log p(x,t)$が具体的に与えられているのならこの計算は原理的に可能である。$\nabla \log p(x,t)$は順方向の拡散過程において計算しておけばよい。任意の$x,t$に対して直ちに計算できるようにするのは困難なので何らかの形でパラメトライズしておき、そのパラメーターをサンプリングにもとづく最適化手法で決定する。
* 拡散モデルでは、「tが大きくなるとg(t)が十分に大きくなるように取る」ので、式（１）の拡散過程により、「元データの確率分布」がおよそ「一様な確率分布」に写像される。
*  拡散モデルの重要な点は、非一様性の高い「元データの確率分布」を「一様な確率分布」に写像してしまう点である。この意味ではVAEも同じである。
* 「元データ」がN枚の画像からなるとしその確率分布として1/Nを割り振る状況を考えると、逆拡散式($\ref{eq:3}$)で指定される写像により、ランダムなノイズが、その「元データ」のどれかに必ず戻ることになる。Nが10億枚であってもそうなるはずである。
* 実際には「元データの確率分布」に戻らない。これは学習させたUnetの作る$\nabla \log p(x,t)$（と離散化誤差）により規定されている。 あくまで「元データの確率分布の学習させたUnetによる近似分布」に戻るだけである。この近似分布に豊富な内容が含まれうる。これはUnetによる学習バイアスが画像合成にとって好ましいものだからである。

* 論文https://arxiv.org/pdf/2202.00512.pdf では、逆拡散を上式を離散化で解く手法と考えて統一的に説明している。離散化には、Stocastic(D'=Dとして離散化）,Euler,Runge-Kutta,DDIMを用いて数値実験している。その上で知識蒸留（サンプリング回数を減らす方法）を考えている。知識蒸留のFig.3は面白いーpdfプリントするだけだとよくわからないが、ディスプレーで見ると、１sample step（完全蒸留）ではやはりかなり崩れてしまっている。Fig.４のようにsampling stepを32までに減らせるということは、まだUnetの32段重ねの計算が必要ということになる。そのうち完全蒸留も可能になるだろう。
* 式($\ref{eq:3}$)をさらに逆に解くことも考えられる。すなわち、ある特定の画像データに対応する「ノイズデータ」を決定できる。

*  基本的な応用は『ノイズ画像==>意味のある画像』という写像を作る話。シンプルなアイデアに基づいているし最近顕著な結果が出てきている。bingで、羽の生えた猫を描いてください、とかお願いしてみるとよい。

--->