Wasserstein梯度流基础

1. 引言

Wasserstein梯度流是研究概率分布演化的强大数学工具¹。与传统的 $L^2$ 梯度流不同，Wasserstein梯度流利用**最优传输（Optimal Transport）**的几何结构，为分布动力学提供了深刻的理解。

在深度学习中，Wasserstein梯度流框架已被用于：

• Mean-Field注意力动力学分析²
• 神经网络训练动态³
• 扩散模型理论⁴
• 变分推断⁵

本文档系统介绍Wasserstein梯度流的理论基础及其在深度学习中的应用。

---

2. 最优传输基础

2.1 Monge问题

Monge最优传输问题（1781年）提出：

给定两个概率测度 $\mu ,\nu \in \mathcal{P}({\mathbb{R}}^d)$ 和代价函数 $c(x,y)$，寻找传输映射 $T:{\mathbb{R}}^d\to {\mathbb{R}}^d$ 使得：

$$T_{\#}\mu =\nu \text{且}\int c(x,T(x))d\mu (x)\text{ 最小化}$$

其中 $T_{\#}\mu$ 表示pushforward测度，定义为对于任意Borel集 $A$：

$$(T_{\#}\mu )(A)=\mu (T^{-1}(A))$$

2.2 Kantorovich松弛

Monge问题在一般情况下可能无解（无最优映射）。Kantorovich（1942年）提出了松弛形式：

$$\underset{\pi \in \Gamma (\mu ,\nu )}{\inf }\int c(x,y)d\pi (x,y)$$

其中 $\Gamma (\mu ,\nu )$ 是所有以 $\mu ,\nu$ 为边缘分布的耦合测度集合。

2.3 Wasserstein距离

对于 $p\ge 1$，Wasserstein-$p$距离定义为：

$$W_p(\mu ,\nu )={\left(\underset{\pi \in \Gamma (\mu ,\nu )}{\inf }\int \Vert x-y{\Vert }^{p}d\pi (x,y)\right)}^{1/p}$$

Wasserstein-2距离（最常用）：

$$W_2(\mu ,\nu )={\left(\underset{\pi \in \Gamma (\mu ,\nu )}{\inf }\int \Vert x-y{\Vert }^{2}d\pi (x,y)\right)}^{1/2}$$

2.4 Brenier定理

当 $\mu$ 绝对连续且 $c(x,y)=\frac{1}{2}\Vert x-y{\Vert }^2$ 时，Kantorovich问题的解由Brenier映射给出：

$$T(x)=\nabla \varphi (x)$$

其中 $\varphi :{\mathbb{R}}^d\to \mathbb{R}$ 是凸函数。这建立了最优传输与Monge-Ampère方程的联系。

---

3. Wasserstein空间的几何结构

3.1 测地线

Wasserstein空间 $({\mathcal{P}}_2({\mathbb{R}}^d),W_2)$ 是一个CAT(0)空间，具有以下性质：

对于 ${\mu }_0,{\mu }_1\in {\mathcal{P}}_2({\mathbb{R}}^d)$，它们之间的测地线为：

$${\mu }_t=((1-t)\text{Id}+tT{)}_{\#}{\mu }_0,t\in [0,1]$$

其中 $T$ 是从 ${\mu }_0$ 到 ${\mu }_1$ 的Brenier映射。

3.2 切空间

在点 $\mu$ 处，Wasserstein空间的切空间 $T_{\mu }{\mathcal{P}}_2$ 等同于：

$$T_{\mu }{\mathcal{P}}_2={\overline{\{\nabla \varphi :\varphi \in C_c^{\infty }\}}}^{L^2(\mu )}$$

即 $L^2(\mu )$ 中梯度场的闭包。直观上，切向量对应于速度场，驱动分布沿测地线移动。

3.3 内积

切空间上的内积定义为：

$$⟨v,w{⟩}_{T_{\mu }{\mathcal{P}}_2}=\int ⟨v(x),w(x)⟩d\mu (x)$$

这一定义使得 ${\mathcal{P}}_2$ 成为希尔伯特流形。

---

4. 梯度流理论

4.1 传统梯度流

在欧几里得空间 ${\mathbb{R}}^n$ 中，函数 $\mathcal{E}$ 的梯度流由下式定义：

$$\dot{x}(t)=-\nabla \mathcal{E}(x(t))$$

4.2 Wasserstein梯度流

在Wasserstein空间中，$\mathcal{E}:{\mathcal{P}}_2\to \mathbb{R}\cup \{+\infty \}$ 的Wasserstein梯度流定义为曲线 $\mu (t)$ 满足：

$${\partial }_t{\mu }_t+\nabla \cdot ({\mu }_t{v}_t)=0$$ $$v_t=-{\nabla }_W\mathcal{E}({\mu }_t)$$

其中 ${\nabla }_W$ 表示Wasserstein空间中的梯度。

4.3 能量-耗散等式

Wasserstein梯度流具有优雅的能量-耗散等式：

$$\frac{d}{dt}\mathcal{E}({\mu }_t)=-\int \Vert v_t{\Vert }^{2}d{\mu }_t=-2\mathcal{D}({\mu }_t)$$

其中 $\mathcal{D}({\mu }_t)=\int \Vert {\nabla }_W\mathcal{E}({\mu }_t){\Vert }^{2}d{\mu }_t$ 是耗散泛函。

这表明能量沿轨迹单调递减。

---

5. McKean-Vlasov方程

5.1 方程定义

McKean-Vlasov方程描述了相互作用粒子的Mean-Field极限行为：

$${\partial }_t{\mu }_t+\nabla \cdot ({\mu }_{t}v[{\mu }_t])=0$$

其中速度场依赖于当前分布：

$$v[{\mu }_t](x)=\int F(x,y)d{\mu }_t(y)$$

5.2 与Wasserstein梯度流的联系

当 $F(x,y)=-\nabla \frac{\delta \mathcal{E}}{\delta \mu }(x)$ 时，McKean-Vlasov方程恰好是能量泛函

$$\mathcal{E}(\mu )=\frac{1}{2}\iint K(x,y)d\mu (x)d\mu (y)$$

的Wasserstein梯度流。

5.3 Mean-Field注意力动力学

对于USA注意力动力学²，连续性方程为：

$${\partial }_t{\mu }_t+\nabla \cdot ({\mu }_t{v}_t[{\mu }_t])=0$$ $$v_t(x)=\int e^{\beta ⟨x,y⟩}yd{\mu }_t(y)$$

这正是McKean-Vlasov方程的形式，其对应的能量泛函为：

$${\mathcal{E}}_{\beta }(\mu )=\frac{1}{2}\iint e^{\beta ⟨x,y⟩}d\mu (x)d\mu (y)$$

---

6. JKO格式

6.1 变分离散化

Jordan-Kinderlehrer-Owock（JKO）格式提供了Wasserstein梯度流的变分刻画：

$${\mu }_{k+1}\in \arg \underset{\mu \in {\mathcal{P}}_2}{\min }\left\{\frac{1}{2h}{W}_2^2(\mu ,{\mu }_k)+\mathcal{E}(\mu )\right\}$$

其中 $h>0$ 是时间步长。

6.2 收敛性

当 $h\to 0$ 时，JKO序列 ${\mu }_h(t)={\mu }_{\lfloor t/h\rfloor }$ 收敛到Wasserstein梯度流 ${\mu }_t$。

6.3 应用

JKO格式被广泛用于：

• Fokker-Planck方程：加入噪声项
• 反应-扩散方程：多物种相互作用
• 扩散生成模型：Score-Based模型的理论基础

---

7. 在深度学习中的应用

7.1 神经网络训练的Mean-Field分析

考虑两层神经网络的参数分布 ${\mu }_t$ 随时间的演化。经验风险

$$\mathcal{E}(\mu )={\mathbb{E}}_{(x,y)\sim \rho }[\ell (f(x;\theta ),y)]$$

的Wasserstein梯度流提供了神经网络训练的连续时间模型³。

7.2 扩散模型的概率流 ODE

扩散模型的概率流ODE是Wasserstein梯度流的一种形式：

$$d{X}_t=\nabla \log p_t(X_t)dt$$

这与最优传输理论有深层联系⁴。

7.3 不变度量梯度下降

**Invariant Metric Gradient Descent (IMGD)**使用Wasserstein几何来设计优化算法：

$${\theta }_{t+1}={\exp }_{{\theta }_t}(-\eta \nabla \mathcal{E}({\theta }_t))$$

其中 ${\exp }_{\theta }$ 是Wasserstein空间中的指数映射。

---

8. 数学附录

8.1 核心不等式

Wasserstein不等式：

$$W_2(\mu ,\nu )\ge C\cdot \text{TV}(\mu ,\nu {)}^{1/d}$$

HWI不等式（Hamming-Wasserstein-Information）：

$$H(\mu |\nu )\ge \frac{1}{2}{W}_2^2(\mu ,\nu )/{\sigma }^2$$

8.2 梯度与散度的定义

在Wasserstein空间中，对于泛函 $\mathcal{E}$：

$$⟨{\nabla }_W\mathcal{E}(\mu ),v⟩=\frac{d}{dt}\mathcal{E}({\mu }_t){|}_{t=0}$$

其中 ${\mu }_t$ 由 ${\partial }_t{\mu }_t+\nabla \cdot ({\mu }_{t}v)=0$ 定义。

---

参考文献

---

本文档为Transformer Mean-Field动力学理论提供最优传输与梯度流的数学基础。

Footnotes

1. Ambrosio, L., Gigli, N., Savaré, G. “Gradient Flows in Metric Spaces and in the Space of Probability Measures.” Birkhäuser (2005). ↩

2. Geshkovski, B., et al. “The Mean-Field Dynamics of Transformers.” arXiv:2512.01868v1 (2025). ↩ ↩²

3. Chizat, L., Bach, F. “On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport.” NeurIPS (2018). ↩ ↩²

4. Song, Y., et al. “Score-Based Generative Modeling through Stochastic Differential Equations.” ICLR (2021). ↩ ↩²

5. Liutkus, A., et al. “Sliced-Wasserstein Flow: A Particle-Based Analog of the Wasserstein Gradient Flow.” ICML (2021). ↩