朗之萬動力學原理簡介
本文的主要內容是基於以下教程：Tutorial on Diffusion Models for Imaging and Vision
https://arxiv.org/abs/2403.18103
此教程寫的非常好，非常推薦大家學習。教程的語言風格也很親切，時不時地蹦出諸如“這是地球人能想出來的公式？”這樣的話，為你枯燥的學習過程增添些許趣味。
朗之萬動力學(Langevin Dynamics)是擴散模型和score matching方法中的取樣過程，是文字生成影像中的一個重要步驟。想要洞悉文生圖的基本原理，朗之萬動力學是繞不開的話題。
給定一個已知的機率分佈 \(p(\mathbf{x})\)，我們的目標是取樣出機率密度更大的哪些樣本。解決這個問題有多種方法，比如生成偽隨機均勻分佈，然後用機率分佈變換的方法；或者用馬爾可夫鏈蒙特卡洛方法(MCMC)。而朗之萬動力學給出的方法是這樣：
隨機選取空間中一個點(這是很簡單的，採用高斯生成與 \(\mathbf{x}\) 同維度的樣本就可以)，然後計算機率密度函式在該點的梯度，然後沿著梯度的方向計算下一個點，直到收斂。為了防止收斂到區域性極大值，我們在每一步迭代計算的時候都會增加高斯噪聲。下面我們給出定義和公式：
朗之萬動力學，是從已知的機率分佈 \(p(\mathbf{x})\) 中進行迭代取樣 \(t=1, \ldots, T\) :
\[\begin{aligned}
\mathbf{x}_{t+1}=\mathbf{x}_t+\tau \nabla_{\mathbf{x}} \log p\left(\mathbf{x}_t\right)+\sqrt{2 \tau} \mathbf{z}, \quad \mathbf{z} \sim \mathcal{N}(0, \mathbf{I}) \quad\quad\quad\quad (1)
\end{aligned}
\\
\]
其中 \(\tau\) 是設定好的迭代步長，用於控制迭代過程，類似SGD中的學習率。初始值 \(\mathbf{x}_0\) 是白噪聲。聰明的讀者已經意識到了，這像極了隨機梯度下降法(SGD)。沒錯！朗之萬動力學就可以理解為梯度下降法！只不過，朗之萬動力學是求機率密度極大值。
朗之萬動力學就是隨機梯度下降！
我們給出一個視覺化結果。假設資料分佈為高斯混合分佈，是兩個高斯分佈混合的情況，那麼，如果每一步不新增噪聲，朗之萬動力學的迭代過程是這個樣子的：
可以看到，收斂到區域性極大值點就不再變了。如果每一步新增噪聲，那麼迭代過程是這個樣子的：
這時候接近區域性極大值點的時候，是逐步震盪收斂的。
下面我們給出公式(1)是怎麼來的，它和物理學中的朗之萬動力學又有什麼關係。大家不要被朗之萬動力學這個名詞嚇到，其實，閱讀此文不需要有大學物理背景，只需要有高中物理知識，同時會一點點微積分就可以了。
一、預備知識
1、牛頓第二定律
力 \(\mathbf{F}\) 等於質量 \(m\) 和 加速度 \(\mathbf{a}(t)\) 的乘積。
\[\begin{aligned}
\mathbf{F} = m \mathbf{a}(t) = m \frac{d\mathbf{v}(t)}{dt}
\end{aligned}
\\
\]
其中加速度 \(\mathbf{a}(t)\) 和 速度 \(\mathbf{v}(t)\) 是時間的函式。\(\mathbf{a}(t)\) 是 \(\mathbf{v}(t)\) 對時間的導數。好，很簡單，沒什麼問題。唯一需要注意的是，我們這裡的向量(vector)用 粗體 表示，標量(scalar)用 細斜體 表示。下面也一樣。
2、力與能量的關係
這個可能部分讀者不太熟悉。我們從動能定理開始。動能定理（kinetic energy theorem）描述的是物體動能的變化量與合外力所做的功的關係，具體內容為：合外力對物體所做的功，等於物體動能的變化量。公式如下：
\[\begin{aligned}
\langle\mathbf{F}, \mathbf{s}\rangle = E_2 - E_1 = \frac{1}{2}m||\mathbf{v}_2||^2 - \frac{1}{2}m||\mathbf{v}_1||^2
\end{aligned}
\\
\]
這裡\(\mathbf{s}\)指的是位移，也就是位置的變化。\(\langle\cdot, \cdot\rangle\) 表示向量內積。\(E\) 代表物體的動能。我們用\(\mathbf{x}\)表示位置，那麼位移就等於位置的變化 \(\mathbf{s} = \Delta \mathbf{x}\)。於是我們重寫上式：
\[\begin{aligned}
\langle\mathbf{F}(\mathbf{x}), \Delta \mathbf{x}\rangle = \Delta E(\mathbf{x})
\end{aligned}
\\
\]
力和能量是隨著物體的位置而變化的，因此是位置 \(x\) 的函式。進一步，我們可以微元化：
\[\begin{aligned}
\langle\mathbf{F}(\mathbf{x}), d \mathbf{x}\rangle &= d E(\mathbf{x}) \\
\mathbf{F}(\mathbf{x}) &= \nabla_{\mathbf{x}} E(\mathbf{x})
\end{aligned}
\\
\]
因此，力是動能對於位移方向的梯度。另外，我們將系統內粒子的能量分為動能和勢能，除了動能之外的所有能量都為勢能，我們用\(\mathbf{U}(\mathbf{x})\)來表示勢能。總能量為\(E_w\)，根據能量守恆定律：\(E_w = U(\mathbf{x}) + E(\mathbf{x})\)，因此：
\[\begin{aligned}
\mathbf{F}(\mathbf{x}) &= \nabla_{\mathbf{x}} E(\mathbf{x}) \\
&= \nabla_{\mathbf{x}} (E_w - U(\mathbf{x})) \\
&= - \nabla_{\mathbf{x}} U(\mathbf{x}) \\
\end{aligned}
\\
\]
我們得出結論，力是粒子勢能的負梯度。
3、玻爾茲曼分佈
在熱力學與統計物理中。每個粒子的分佈的位置是一個隨機變數，粒子的分佈一般用玻爾茲曼分佈來表示：
\[\begin{aligned}
p(\mathbf{x}) &= \frac{1}{Z} \exp\left\{-U(\mathbf{x})\right\} \\
\end{aligned}
\\
\]
其中 \(Z\) 是歸一化係數。我們可以直觀地理解，對於勢能特別大的地方，粒子出現的機率很小；勢能特別小的地方，說明粒子的動能比較大，運動比較劇烈，分佈比較密集。這是符合我們的人類感知的，你看一般人口比較密集的地方都比較喧鬧，如商場和鬧市區；人口稀疏的地方，比較安靜。
二、推導朗之萬動力學
一百多年前，著名物理學家朗之萬給出了布朗運動方程：
\[\begin{aligned}
m \frac{d\mathbf{v}(t)}{dt} &= -\gamma \mathbf{v}(t) + \bm{\eta} \quad\quad\quad \bm{\eta} \sim \mathcal{N}(\mathbf{0}, \sigma^2 \mathbf{I}) \quad\quad\quad (2)\\
\end{aligned}
\\
\]
我們不必糾結這個方程怎麼來的。我們其實可以很輕鬆地讀懂這個方程。這本質上就是牛頓第二定律的擴充，等式左邊就是質量和加速度的乘積，等式右邊是力。其中 \(-\gamma \mathbf{v}(t)\) 是和速度有關的摩擦力，它和物體的運動方向相反，所以加個負號，\(\gamma\) 是動摩擦因數。\(\bm{\eta}\) 是隨機力，服從均值為 \(0\)，協方差矩陣為 \(\sigma^2 \mathbf{I}\) 的高斯分佈。
由牛頓第二定律和力與能量的關係可知：
\[\begin{aligned}
m \frac{d\mathbf{v}(t)}{dt} &= -\nabla_{\mathbf{x}} U(\mathbf{x}) \\
\end{aligned}
\\
\]
代入(2)式，於是有：
\[\begin{aligned}
-\nabla_{\mathbf{x}} U(\mathbf{x}) &= -\gamma \mathbf{v}(t) + \bm{\eta} \\
\mathbf{v}(t) &= \frac{1}{\gamma}\nabla_{\mathbf{x}} U(\mathbf{x}) + \frac{1}{\gamma}\bm{\eta} \\
\end{aligned}
\\
\]
速度是位移隨時間的導數，另外，將隨機力簡化為標準高斯分佈\(\bm{\eta} = \sigma \mathbf{z}\)，有：
\[\begin{aligned}
\frac{d\mathbf{x}(t)}{dt} &= \frac{1}{\gamma}\nabla_{\mathbf{x}} U(\mathbf{x}) + \frac{\sigma}{\gamma}\mathbf{z} \quad \mathbf{z} \sim \mathcal{N}(0, \mathbf{I})\\
\end{aligned}
\\
\]
時間離散化：
\[\begin{aligned}
\frac{\mathbf{x}_{t+1} - \mathbf{x}_{t}}{\Delta t} &= \frac{1}{\gamma}\nabla_{\mathbf{x}} U(\mathbf{x}_{t}) + \frac{\sigma}{\gamma}\mathbf{z} \\
\mathbf{x}_{t+1} - \mathbf{x}_{t} &= \frac{\Delta t}{\gamma}\nabla_{\mathbf{x}} U(\mathbf{x}_{t}) + \frac{\sigma \Delta t}{\gamma}\mathbf{z} \\
\end{aligned}
\\
\]
這裡 \(\Delta t\) 是離散取樣時間間隔，是常數，\(\gamma\) 是動摩擦因數，也是常數。我們用常數 \(\tau = \frac{\Delta t}{\gamma}\)，於是有：
\[\begin{aligned}
\mathbf{x}_{t+1} - \mathbf{x}_{t} &= \tau\nabla_{\mathbf{x}} U(\mathbf{x}_{t}) + \tau \sigma \mathbf{z} \\
\mathbf{x}_{t+1} &= \mathbf{x}_{t} + \tau\nabla_{\mathbf{x}} U(\mathbf{x}_{t}) + \tau \sigma \mathbf{z} \\
\end{aligned}
\\
\]
到這裡我們稍微停一下。根據(2)式所描述的過程，最終粒子會越來越集中，還是越來越分散？對(2)式兩邊取期望，可以發現，粒子的速度是越來越小的，即動能是越來越小的，粒子往勢能增大的方向擴散。根據玻爾茲曼分佈，粒子的機率密度會越來越小。而我們希望的是相反的方向，即取樣出機率密度最大的樣本。因此需要對上式右邊梯度的符號進行修改：
\[\begin{aligned}
\mathbf{x}_{t+1} &= \mathbf{x}_{t} - \tau\nabla_{\mathbf{x}} U(\mathbf{x}_{t}) + \tau \sigma \mathbf{z} \quad\quad\quad (3)\\
\end{aligned}
\\
\]
根據玻爾茲曼分佈：
\[\begin{aligned}
\nabla_{\mathbf{x}} \log p(\mathbf{x}_{t}) &= \nabla_{\mathbf{x}} \left\{ \log \frac{1}{Z} - U(\mathbf{x}_{t})\right\} \\ \\
&= -\nabla_{\mathbf{x}} U(\mathbf{x}_{t})
\end{aligned}
\\
\]
代入(3)式，有：
\[\begin{aligned}
\mathbf{x}_{t+1} &= \mathbf{x}_{t} + \tau \nabla_{\mathbf{x}} \log p(\mathbf{x}_{t}) + \tau \sigma \mathbf{z} \quad\quad\quad (4)\\
\end{aligned}
\\
\]
對比 (1) 式和 (4) 式，是不是感覺幾乎一樣了呢？\(\tau\) 是常數，方差\(\sigma\)也是常數。我們深度學習計算中的常數值幾乎都可以替換或者捨棄，或者可以人工指定。令\(\sigma = \sqrt{2/\tau}\)，我們就得到了 (1) 式：
\[\begin{aligned}
\mathbf{x}_{t+1}=\mathbf{x}_t+\tau \nabla_{\mathbf{x}} \log p\left(\mathbf{x}_t\right)+\sqrt{2 \tau} \mathbf{z}, \quad \mathbf{z} \sim \mathcal{N}(0, \mathbf{I})
\end{aligned}
\\
\]