485 lines
22 KiB
Markdown
485 lines
22 KiB
Markdown
---
|
||
title: 4-最优传输
|
||
tags:
|
||
- 最优传输
|
||
- OT
|
||
- 优化理论
|
||
- 数学基础
|
||
---
|
||
|
||
# 最优传输（Optimal Transport, OT）
|
||
|
||
> **摘要**：最优传输（Optimal Transport, OT）是概率论与几何分析的交叉领域，为度量概率分布之间的距离提供了坚实的数学基础。近年来，OT 理论在深度学习与生成式模型中发挥了核心作用——从 WGAN 的梯度惩罚、扩散模型的 flow matching，到对比学习中的 Sinkhorn Divergence，都可见其影响。本笔记面向深度学习与生成式模型科研人员，系统阐述 OT 的数学基础、计算方法及其与生成模型的深刻联系。
|
||
|
||
---
|
||
|
||
## 1. 最优传输的基本问题
|
||
|
||
### 1.1 Monge 问题
|
||
|
||
**定义 1.1（Monge 问题）**。给定两个概率空间 $(\mathcal{X}, \mu)$ 和 $(\mathcal{Y}, \nu)$，以及代价函数 $c: \mathcal{X} \times \mathcal{Y} \to \mathbb{R}_{\ge 0}$，Monge 问题寻求一个**确定性映射** $f: \mathcal{X} \to \mathcal{Y}$，使得
|
||
|
||
$$
|
||
\min_{f: f_\# \mu = \nu} \int_{\mathcal{X}} c(x, f(x)) \, d\mu(x) \tag{1.1}
|
||
$$
|
||
|
||
其中约束 $f_\# \mu = \nu$ 表示 $f$ 将 $\mu$ 的质量传输到 $\nu$，即对任意可测集 $B \subseteq \mathcal{Y}$ 有 $\mu(f^{-1}(B)) = \nu(B)$。
|
||
|
||
Monge 问题的物理直觉是：有一堆"泥土"（由 $\mu$ 描述）分布在 $\mathcal{X}$ 上，需要将其搬运到目标堆（由 $\nu$ 描述）分布在 $\mathcal{Y}$ 上，每单位泥土从 $x$ 运到 $y$ 的成本是 $c(x,y)$，求总成本最小的搬运方案。
|
||
|
||
**关键性质**：Monge 问题**不是凸的**——映射空间上的约束 $f_\# \mu = \nu$ 使得可行集非凸。因此直接优化 (1.1) 通常非常困难。当 $c(x,y) = \|x-y\|^2$ 且 $\mathcal{X} = \mathcal{Y} = \mathbb{R}^d$ 时，若 $\mu$ 绝对连续且满足某些正则性条件，唯一解由 **Brenier 势函数** 的梯度给出：$f(x) = \nabla \phi(x)$。
|
||
|
||
---
|
||
|
||
### 1.2 Kantorovich 松弛
|
||
|
||
**定义 1.2（Kantorovich 问题）**。为克服 Monge 问题的非凸性，Kantorovich (1942) 引入**联合分布**（耦合）来代替确定性映射：
|
||
|
||
$$
|
||
\min_{\pi \in \Gamma(\mu, \nu)} \int_{\mathcal{X} \times \mathcal{Y}} c(x, y) \, d\pi(x, y) \tag{1.2}
|
||
$$
|
||
|
||
其中 $\Gamma(\mu, \nu)$ 表示所有以 $\mu$ 和 $\nu$ 为边缘分布的联合概率分布：
|
||
|
||
$$
|
||
\Gamma(\mu, \nu) = \left\{ \pi \in \mathcal{P}(\mathcal{X} \times \mathcal{Y}) \mid \pi_{X} = \mu,\; \pi_{Y} = \nu \right\}
|
||
$$
|
||
|
||
即 $\pi$ 在 $\mathcal{X}$ 上的边际为 $\mu$，在 $\mathcal{Y}$ 上的边际为 $\nu$。
|
||
|
||
Kantorovich 形式的物理直觉是：不规定"每粒泥土具体去哪里"，只要求整体分布匹配 $\mu$ 和 $\nu$。这使得问题变成关于联合分布 $\pi$ 的**线性规划**，是凸的。
|
||
|
||
**定理 1.1（Monge-Kantorovich 对偶）**。在适当的紧性条件下，最优传输问题有如下对偶形式：
|
||
|
||
$$
|
||
\inf_{\pi \in \Gamma(\mu,\nu)} \int c \, d\pi \;=\; \sup_{\substack{\varphi \in L^1(\mu),\; \psi \in L^1(\nu) \\ \varphi(x) + \psi(y) \leq c(x,y)}} \left( \int_{\mathcal{X}} \varphi \, d\mu + \int_{\mathcal{Y}} \psi \, d\nu \right) \tag{1.3}
|
||
$$
|
||
|
||
**证明梗概**：将 (1.2) 视为线性规划。其对偶变量是边际约束的拉格朗日乘子，即对偶问题中寻找函数对 $(\varphi, \psi)$ 使得约束 $\varphi(x) + \psi(y) \leq c(x,y)$ 下的线性泛函最大化。根据线性规划的对偶理论，强对偶性成立（当边际分布是紧概率空间时）。
|
||
|
||
---
|
||
|
||
### 1.3 两者的等价性与适用场景
|
||
|
||
**关系**：
|
||
- 如果 Monge 问题存在**最优映射** $f^*$，则由 $f^*$ 诱导的传输计划 $\pi^* = (I, f^*)_\# \mu$（即 $\pi^*(A \times B) = \mu(f^{-1}(B) \cap A)$）是 Kantorovich 问题的最优解，且两目标值相等。
|
||
- 反之不一定成立：当 $c$ 是严格凸函数（如 $c(x,y) = \|x-y\|^2$）时，若 $\mu$ 绝对连续，则最优传输映射存在且唯一（Brenier 定理）；但对于一般 $c$，最优解可能是**非确定性**的，此时 Kantorovich 松弛严格更优。
|
||
|
||
**选择原则**：
|
||
| 场景 | 推荐形式 | 原因 |
|
||
|------|----------|------|
|
||
| 二次代价 $c(x,y) = \|x-y\|^2$（欧式空间） | Monge（Brenier 映射） | 有显式结构，映射唯一 |
|
||
| 离散分布 / 图结构 | Kantorovich | 线性规划易于求解 |
|
||
| 计算梯度（神经网络） | Kantorovich（熵正则化） | 对偶形式可微 |
|
||
|
||
---
|
||
|
||
## 2. Wasserstein 距离
|
||
|
||
### 2.1 $p$-Wasserstein 距离的定义
|
||
|
||
**定义 2.1（$p$-Wasserstein 距离）**。给定 $\mathcal{X}$ 上的两个概率分布 $\mu, \nu$，以及代价函数 $c(x,y) = d(x,y)$（通常为度量），$p$-Wasserstein 距离定义为
|
||
|
||
$$
|
||
W_p(\mu, \nu) \;:=\; \left( \inf_{\pi \in \Gamma(\mu,\nu)} \int_{\mathcal{X} \times \mathcal{X}} d(x,y)^p \, d\pi(x,y) \right)^{1/p} \tag{2.1}
|
||
$$
|
||
|
||
当 $d$ 是度量时，$W_p$ 确实构成 $\mathcal{P}_p(\mathcal{X})$（$p$-阶矩有限的概率分布空间）上的一个度量。
|
||
|
||
---
|
||
|
||
### 2.2 $W_2$ 距离的特例
|
||
|
||
当 $c(x,y) = \|x-y\|^2$（欧氏距离的平方）时，得到 **Wasserstein-2 距离**：
|
||
|
||
$$
|
||
W_2(\mu, \nu) = \left( \inf_{\pi \in \Gamma(\mu,\nu)} \int \|x-y\|^2 \, d\pi(x,y) \right)^{1/2} \tag{2.2}
|
||
$$
|
||
|
||
$W_2$ 在生成模型中最为常用，因为：
|
||
1. 当 $\mu$ 绝对连续时，存在唯一的最优传输映射 $T(x) = \nabla \phi(x)$，使得 $W_2(\mu, \nu) = \left( \int \|x - T(x)\|^2 \, d\mu(x) \right)^{1/2}$。
|
||
2. 由 **Kantorovich-Rubinstein 对偶**，$W_1$ 距离（$p=1$）有简化的对偶形式（见 5.1 节），但 $W_2$ 的几何结构更丰富。
|
||
|
||
---
|
||
|
||
### 2.3 Wasserstein 距离的性质
|
||
|
||
**定理 2.1（度量结构）**。$W_p: \mathcal{P}_p(\mathcal{X}) \times \mathcal{P}_p(\mathcal{X}) \to \mathbb{R}_+$ 满足：
|
||
- **非负性**：$W_p(\mu, \nu) \geq 0$，且 $W_p(\mu, \nu) = 0 \iff \mu = \nu$
|
||
- **对称性**：$W_p(\mu, \nu) = W_p(\nu, \mu)$
|
||
- **三角不等式**：$W_p(\mu, \nu) \leq W_p(\mu, \zeta) + W_p(\zeta, \nu)$
|
||
- **弱收敛拓扑**：$W_p$ 收敛等价于弱收敛 + $p$-阶矩收敛（在紧空间上）
|
||
|
||
**性质 2.1（相对于 $L^p$ 距离的优势）**。与 KL 散度、JS 散度不同，Wasserstein 距离**不需要分布的支撑集有重叠**就能有效度量距离。这使其特别适合处理生成分布 $\mu_\theta$ 与目标分布 $\nu$ 支撑不重叠的情形（如 GAN 中常见的模式坍塌问题）。
|
||
|
||
---
|
||
|
||
## 3. Sinkhorn 算法
|
||
|
||
### 3.1 熵正则化最优传输
|
||
|
||
直接求解 Kantorovich 问题 (1.2) 的计算复杂度为 $O(n^3)$（离散情形下，$n$ 为支撑点数），不适用于大规模数据。
|
||
|
||
**Cuturi (2013)** 引入**熵正则化**：
|
||
|
||
$$
|
||
\min_{\pi \in \Gamma(\mu,\nu)} \langle C, \pi \rangle + \epsilon H(\pi) \tag{3.1}
|
||
$$
|
||
|
||
其中 $C_{ij} = c(x_i, y_j)$ 是代价矩阵，$H(\pi) = -\sum_{i,j} \pi_{ij} \log \pi_{ij}$ 是联合分布的熵，$\epsilon > 0$ 是正则化参数。
|
||
|
||
当 $\epsilon \to 0$ 时，(3.1) 的解收敛到原始 Kantorovich 问题 (1.2) 的解；当 $\epsilon \to +\infty$ 时，解趋近于独立乘积 $\mu \otimes \nu$。
|
||
|
||
---
|
||
|
||
### 3.2 Sinkhorn 迭代
|
||
|
||
将 $\pi$ 约束为 $\Gamma(\mu, \nu)$ 中的双随机矩阵（即 $\sum_i \pi_{ij} = b_j$，$\sum_j \pi_{ij} = a_i$），用拉格朗日形式写出 (3.1) 的KKT条件，可导出对偶变量 $u \in \mathbb{R}^n$ 和 $v \in \mathbb{R}^n$，使得最优耦合具有 **Sinkhorn 形式**：
|
||
|
||
$$
|
||
\pi_{ij}^* = \exp\left(\frac{u_i + v_j - C_{ij}}{\epsilon}\right) = \operatorname{diag}(u) \, K \, \operatorname{diag}(v) \tag{3.2}
|
||
$$
|
||
|
||
其中 $K_{ij} = \exp\left(-C_{ij} / \epsilon\right)$ 称为 **Gibbs 核**。
|
||
|
||
**定理 3.1（Sinkhorn 迭代）**。固定点迭代给出如下交替投影：
|
||
|
||
$$
|
||
u_{i+1} = \frac{a_i}{(\mathbf{K} v)_i}, \qquad v_{j+1} = \frac{b_j}{(\mathbf{K}^\top u)_j} \tag{3.3}
|
||
$$
|
||
|
||
其中 $\mathbf{K} \in \mathbb{R}^{n \times m}$，$a, b$ 是边际分布向量。
|
||
|
||
**矩阵-向量形式**：
|
||
|
||
$$
|
||
\mathbf{u} \leftarrow \frac{\mathbf{a}}{\mathbf{K} \mathbf{v}}, \qquad \mathbf{v} \leftarrow \frac{\mathbf{b}}{\mathbf{K}^\top \mathbf{u}} \tag{3.4}
|
||
$$
|
||
|
||
每轮交替更新称为一次 Sinkhorn 迭代（sinkhorn normalization）。
|
||
|
||
---
|
||
|
||
### 3.3 矩阵形式的 Sinkhorn 循环
|
||
|
||
设 $\mathbf{C} \in \mathbb{R}^{n \times m}$ 为代价矩阵，$\mathbf{K} = \exp(-\mathbf{C} / \epsilon)$。Sinkhorn 算法维护两个缩放向量 $\mathbf{u} \in \mathbb{R}^n_+$ 和 $\mathbf{v} \in \mathbb{R}^m_+$，迭代直至收敛：
|
||
|
||
```
|
||
输入: C, a, b, ε, 最大迭代次数 T
|
||
K = exp(-C / ε)
|
||
u = ones(n) / n
|
||
v = ones(m) / m
|
||
for t = 1 to T:
|
||
u = a / (K @ v)
|
||
v = b / (K.T @ u)
|
||
if ||u - u_prev|| < tol: break
|
||
return π = diag(u) @ K @ diag(v)
|
||
```
|
||
|
||
最优传输计划为 $\hat{\pi} = \operatorname{diag}(\mathbf{u}) \mathbf{K} \operatorname{diag}(\mathbf{v})$，对应的目标值：
|
||
|
||
$$
|
||
\mathrm{OT}_\epsilon(\mathbf{a}, \mathbf{b}) = \langle \mathbf{C}, \hat{\pi} \rangle + \epsilon H(\hat{\pi}) = \epsilon \left( \mathbf{a}^\top \log \mathbf{u} + \mathbf{b}^\top \log \mathbf{v} \right) \tag{3.5}
|
||
$$
|
||
|
||
---
|
||
|
||
### 3.4 收敛性分析
|
||
|
||
**定理 3.2（Sinkhorn 收敛速度）**。设代价矩阵 $\mathbf{C}$ 的 entries 有界，$\epsilon > 0$，则 Sinkhorn 迭代线性收敛到唯一最优解：
|
||
|
||
$$
|
||
\|\mathbf{u}^{(t)} - \mathbf{u}^*\| \leq C \cdot \rho^t \quad \text{其中} \quad \rho = 1 - \frac{\lambda_{\min}(\mathbf{P})}{\lambda_{\max}(\mathbf{P})} < 1 \tag{3.6}
|
||
$$
|
||
|
||
其中 $\mathbf{P} = \operatorname{diag}(\mathbf{u}^*) \mathbf{K} \operatorname{diag}(\mathbf{v}^*)$ 是最优耦合矩阵（双随机）。
|
||
|
||
**实际注意事项**：
|
||
- 当 $\epsilon$ 过小时，$K_{ij} = \exp(-C_{ij}/\epsilon)$ 极易**溢出**（underflow），需要 log-space 稳定化。
|
||
- 迭代次数 $T$ 通常设为 $O(\log(1/\delta)/\ \epsilon)$ 量级以达到 $\delta$-近似。
|
||
- 当支撑点数量 $n$ 很大时（> 10^5），全矩阵 $\mathbf{K}$ 无法存储，可使用 **kernel approximation** 或 **sparse Sinkhorn**。
|
||
|
||
---
|
||
|
||
## 4. 无熵约束的 Sinkhorn（稳定版）
|
||
|
||
### 4.1 log-space 稳定计算
|
||
|
||
当 $\epsilon$ 很小或 $C_{ij}$ 很大时，直接计算 $K_{ij} = \exp(-C_{ij}/\epsilon)$ 会下溢（round to zero）。解决方案是始终在 **log-space** 操作，存储 $\log K_{ij} = -C_{ij} / \epsilon$，然后使用 **log-sum-exp** 技巧来避免数值不稳定。
|
||
|
||
对 (3.4) 取对数，设 $f_i = \log u_i$，$g_j = \log v_j$，则：
|
||
|
||
$$
|
||
f_i = \log a_i - \log \sum_j \exp\left( -\frac{C_{ij}}{\epsilon} + g_j \right) \tag{4.1}
|
||
$$
|
||
|
||
对每一行 $i$，计算 $M_{ij} = -C_{ij}/\epsilon + g_j$，然后通过 log-sum-exp 得到 $f_i$：
|
||
|
||
$$
|
||
f_i = \log a_i - \text{LogSumExp}_j(M_{ij}) \tag{4.2}
|
||
$$
|
||
|
||
其中 $\text{LogSumExp}(\mathbf{z}) = \log \sum_k \exp(z_k)$，可通过减去 $\max(z)$ 来稳定计算：
|
||
|
||
$$
|
||
\text{LogSumExp}(\mathbf{z}) = \max(z) + \log \sum_k \exp(z_k - \max(z)) \tag{4.3}
|
||
$$
|
||
|
||
---
|
||
|
||
### 4.2 完整 log-space Sinkhorn
|
||
|
||
```python
|
||
def log_sinkhorn(C, a, b, eps, max_iter=1000, tol=1e-9):
|
||
"""
|
||
Log-space stabilized Sinkhorn iteration.
|
||
C: (n, m) cost matrix
|
||
a: (n,) source distribution (sum to 1)
|
||
b: (m,) target distribution (sum to 1)
|
||
"""
|
||
n, m = C.shape
|
||
log_K = -C / eps # (n, m)
|
||
|
||
log_u = torch.zeros(n)
|
||
log_v = torch.zeros(m)
|
||
|
||
for _ in range(max_iter):
|
||
# Update log_v: log_v_j = log b_j - logsumexp_i( log_K_ij + log_u_i )
|
||
log_v_new = torch.log(b) - logsumexp(log_K + log_u.unsqueeze(1), dim=0).squeeze()
|
||
# Update log_u: log_u_i = log a_i - logsumexp_j( log_K_ij + log_v_new_j )
|
||
log_u_new = torch.log(a) - logsumexp(log_K + log_v_new.unsqueeze(0), dim=1).squeeze()
|
||
|
||
if torch.max(torch.abs(log_u_new - log_u)) < tol:
|
||
break
|
||
log_u = log_u_new
|
||
log_v = log_v_new
|
||
|
||
return log_u, log_v
|
||
```
|
||
|
||
**核心技巧**：
|
||
1. 所有运算在 log 域进行，防止 $K_{ij}$ 溢出
|
||
2. log-sum-exp 通过 (4.3) 稳定计算
|
||
3. 收敛判定使用对偶变量差的 $L_\infty$ 范数
|
||
|
||
---
|
||
|
||
## 5. Wasserstein 距离的梯度
|
||
|
||
### 5.1 Sinkhorn 的对偶形式
|
||
|
||
熵正则化问题的对偶（由 (1.3) 扩展而来）为：
|
||
|
||
$$
|
||
\max_{\alpha, \beta} \quad \alpha^\top \mathbf{a} + \beta^\top \mathbf{b} - \epsilon \left\langle \exp\left(\frac{\alpha}{\epsilon}\right), \mathbf{K} \exp\left(\frac{\beta}{\epsilon}\right) \right\rangle \tag{5.1}
|
||
$$
|
||
|
||
其中 $\alpha \in \mathbb{R}^n, \beta \in \mathbb{R}^m$ 是对偶变量。
|
||
|
||
设 $\mathbf{u} = \exp(\alpha/\epsilon)$，$\mathbf{v} = \exp(\beta/\epsilon)$，则 (5.1) 的稳定性条件给出 Sinkhorn 迭代的固定点 $\mathbf{u}^*, \mathbf{v}^*$。
|
||
|
||
---
|
||
|
||
### 5.2 Sinkhorn 梯度在神经网络中的应用
|
||
|
||
设损失函数 $\mathcal{L} = W_2(\mu_\theta, \nu)$，其中 $\mu_\theta$ 是神经网络参数化 $\theta$ 生成的分布（如生成器输出的隐空间分布），$\nu$ 是真实数据分布。
|
||
|
||
由于 Sinkhorn 算法给出的是**近似**最优耦合 $\hat{\pi}$，我们可以计算 $\mathcal{L}$ 对 $\theta$ 的梯度。
|
||
|
||
**链式法则**：
|
||
|
||
$$
|
||
\nabla_\theta \mathcal{L} = \left( \nabla_{\hat{\pi}} \mathcal{L} \right)^\top \nabla_\theta \hat{\pi} \tag{5.2}
|
||
$$
|
||
|
||
而 $\hat{\pi}_{ij} = u_i K_{ij} v_j$，对 $K_{ij}$ 求导涉及神经网络输出的 $x_i = g_\theta(z_i)$：
|
||
|
||
$$
|
||
\frac{\partial K_{ij}}{\partial \theta} = -\frac{1}{\epsilon} \exp\left(-\frac{C_{ij}}{\epsilon}\right) \cdot \frac{\partial C_{ij}}{\partial \theta} \tag{5.3}
|
||
$$
|
||
|
||
其中 $C_{ij} = \|x_i - y_j\|^2$，所以
|
||
|
||
$$
|
||
\frac{\partial C_{ij}}{\partial \theta} = 2 (x_i - y_j) \cdot \frac{\partial x_i}{\partial \theta} \tag{5.4}
|
||
$$
|
||
|
||
这允许通过标准反向传播计算梯度，尽管需要存储 $\mathbf{u}, \mathbf{v}$（或在对数空间中存储 $\log \mathbf{u}, \log \mathbf{v}$）。
|
||
|
||
---
|
||
|
||
### 5.3 梯度正则化问题
|
||
|
||
由于 $\hat{\pi}$ 是近似解，梯度估计存在**偏差**（bias）：
|
||
- 当 $\epsilon$ 较大时，近似误差大，但方差小
|
||
- 当 $\epsilon$ 较小时，近似误差小，但 $\hat{\pi}$ 的条件数变差，梯度方差增大
|
||
|
||
**解决方案**：
|
||
1. **Shamir et al. (2023)** 的"unbiased Sinkhorn"：通过重参数化技巧消除近似偏差
|
||
2. **Neural Sinkhorn**（Genevay et al., 2016）：用神经网络学习最优传输映射而非迭代求解
|
||
3. **截断梯度（truncated gradient）**：设定 $K_{ij}$ 的下界阈值以避免极端值
|
||
|
||
---
|
||
|
||
## 6. 与生成模型的关系
|
||
|
||
### 6.1 WGAN：Wasserstein-1 距离代替 JS 散度
|
||
|
||
**Arjovsky et al. (2017)** 提出使用 $W_1$ 距离替代 JS 散度来度量生成分布 $\mu_\theta$ 与真实分布 $\nu$ 的距离。
|
||
|
||
**Kantorovich-Rubinstein 对偶**给出 $W_1$ 的简化形式：
|
||
|
||
$$
|
||
W_1(\mu, \nu) = \sup_{\|f\|_{Lip} \leq 1} \mathbb{E}_{x \sim \mu}[f(x)] - \mathbb{E}_{y \sim \nu}[f(y)] \tag{6.1}
|
||
$$
|
||
|
||
即 $W_1$ 是 1-Lipschitz 函数空间上的最大均值差异（MMD）。
|
||
|
||
**关键洞察**：对判别器 $f$ 施加 Lipschitz 约束（而非 Jensen-Shannon 的 log-likelihood 约束），可以在分布**支撑不重叠**时仍提供有意义的梯度——这解决了 standard GAN 在模式坍塌时梯度消失的问题。
|
||
|
||
WGAN 的目标函数：
|
||
|
||
$$
|
||
\min_G \max_{f: \|f\|_{Lip} \leq 1} \mathbb{E}_{x \sim \nu}[f(x)] - \mathbb{E}_{z \sim p(z)}[f(G(z))] \tag{6.2}
|
||
$$
|
||
|
||
---
|
||
|
||
### 6.2 梯度惩罚（Gradient Penalty）的推导与作用
|
||
|
||
**Gulrajani et al. (2017)** 发现直接约束 $\|f\|_{Lip} \leq 1$ 在实践中难以优化，因此提出**梯度惩罚**（WGAN-GP）：
|
||
|
||
$$
|
||
\mathcal{L}_D = \mathbb{E}_{x\sim\nu}[f(x)] - \mathbb{E}_{\tilde{x}\sim\mu_\theta}[f(\tilde{x})] + \lambda \mathbb{E}_{\hat{x} \sim P_{\hat{x}}} \left[ (\|\nabla_{\hat{x}} f(\hat{x})\|_2 - 1)^2 \right] \tag{6.3}
|
||
$$
|
||
|
||
其中 $\hat{x} = \epsilon x + (1-\epsilon) \tilde{x}$，$x \sim \nu$，$\tilde{x} \sim \mu_\theta$，$\epsilon \sim U[0,1]$。
|
||
|
||
**推导**：$W_1$ 的最优判别器 $f^*$ 满足 $\| \nabla f^* \| = 1$ almost everywhere（在 $\mu_\theta$ 和 $\nu$ 的支撑上）。因此梯度惩罚强制学到的 $f$ 逼近这一性质。
|
||
|
||
**正则化解释**：在 $W_1$ 的对偶问题中，最优解对应约束 $\nabla f \in \partial c$（次梯度），梯度惩罚是对该约束的恰当正则化，使得问题稳定可解。
|
||
|
||
---
|
||
|
||
### 6.3 扩散模型中的 OT：Optimal Transport in Flow Matching
|
||
|
||
**Flow Matching**（Lipman et al., 2022; Pooladian et al., 2023）将扩散模型的生成过程建模为**路径测度**（path measure），在连续时间下从噪声分布 $\mu_0$ 流向数据分布 $\mu_1$。
|
||
|
||
最优传输 flow matching 使用 OT 映射来设计更好的条件路径：
|
||
|
||
$$
|
||
\mathrm{OT-FM}: \quad \frac{dx_t}{dt} = v_t(x_t), \quad v_t(x) = \int y \, d\pi_{t| x}(y) - x \tag{6.4}
|
||
$$
|
||
|
||
其中 $T_{0\to1}$ 是从 $\mu_0$ 到 $\mu_1$ 的最优传输映射。
|
||
|
||
**关键优势**：
|
||
- OT-FM 避免了标准 FM 的边缘不匹配问题
|
||
- 训练目标简化为回归一个**确定的向量场**（而非随机路径）
|
||
- 在图像生成任务中，使用 OT-FM 可以显著减少推理步骤（从 1000 步降至约 50 步仍保持高质量）
|
||
|
||
则 OT flow 给出确定性向量场：
|
||
|
||
$$
|
||
v_t(x) = \int y \, d\pi_{t| x}(y) - x \tag{6.5}
|
||
$$
|
||
|
||
这正是 **Jordan-Kinder-Pearson (JKP) 方程** 在 $\sigma \to 0$ 时的极限形式。
|
||
|
||
---
|
||
|
||
### 6.4 对比学习中的 OT：Sinkhorn Divergence
|
||
|
||
**正面样本对** $(x, x^+)$ 和**负面样本对** $(x, x^-)$ 在特征空间中应该满足：与正面样本的距离应小于与负面样本的距离。
|
||
|
||
**Sinkhorn Divergence**（Cuturi et al., 2019; Mensch et al., 2022）定义为：
|
||
|
||
$$
|
||
S_\epsilon(\mu, \nu) = \mathrm{OT}_\epsilon(\mu, \nu) - \frac{1}{2} \mathrm{OT}_\epsilon(\mu, \mu) - \frac{1}{2} \mathrm{OT}_\epsilon(\nu, \nu) \tag{6.6}
|
||
$$
|
||
|
||
其中 $\mathrm{OT}_\epsilon$ 是 Sinkhorn 近似的熵正则化传输代价。
|
||
|
||
**几何直觉**：
|
||
- $\mathrm{OT}_\epsilon(\mu, \nu)$ 给出 $\mu$ 到 $\nu$ 的有向传输代价
|
||
- 减去 $\frac{1}{2}(\mathrm{OT}_\epsilon(\mu, \mu) + \mathrm{OT}_\epsilon(\nu, \nu))$ 使 $S_\epsilon(\mu, \nu)$ 具有**对称性**且满足 $S_\epsilon(\mu, \mu) = 0$
|
||
- 当 $\epsilon \to \infty$ 时，$S_\epsilon(\mu, \nu) \to 0$；当 $\epsilon \to 0$ 时，$S_\epsilon(\mu, \nu) \to W_1(\mu, \nu)$ 的某种变形
|
||
|
||
在对比学习中，可以将 $S_\epsilon(\mu^+, \mu^-)$ 作为损失的一部分：
|
||
|
||
$$
|
||
\mathcal{L}_{\text{OTCL}} = \alpha \cdot S_\epsilon(f(x), f(x^+)) + (1-\alpha) \cdot S_\epsilon(f(x), f(x^-)) \tag{6.7}
|
||
$$
|
||
|
||
这鼓励模型将同一样本的正负面映射到特征空间中 OT-最优的位置。
|
||
|
||
---
|
||
|
||
## 7. 理论延伸
|
||
|
||
### 7.1 Gromov-Wasserstein：结构匹配（非度量空间）
|
||
|
||
当两个分布所在的**空间**不兼容直接比较时（例如，$\mathcal{X}$ 上的代价函数无法定义为有意义的距离），Gromov-Wasserstein (GW) 距离提供了一种基于**结构相似性**的比较框架。
|
||
|
||
**定义 7.1（Gromov-Wasserstein 距离）**。给定两个度量空间 $(\mathcal{X}, d_\mathcal{X}, \mu)$ 和 $(\mathcal{Y}, d_\mathcal{Y}, \nu)$，$p$-Gromov-Wasserstein 距离定义为：
|
||
|
||
$$
|
||
\operatorname{GW}_p(\mu, \nu) = \left( \inf_{\pi \in \Gamma(\mu, \nu)} \int \left| d_\mathcal{X}(x, x') - d_\mathcal{Y}(y, y') \right|^p \, d\pi(x,y) \, d\pi(x',y') \right)^{1/p} \tag{7.1}
|
||
$$
|
||
|
||
**物理直觉**：不要求 $x$ 与 $y$ 之间的代价可直接比较，而是比较**成对距离结构**——如果两个空间中的点对距离结构相似，则它们在 GW 意义下接近。
|
||
|
||
**应用场景**：
|
||
- 图匹配（graph matching）：节点之间无法直接对应，但图的结构（邻接矩阵特征值、路径长度）需要匹配
|
||
- 单细胞 RNA 测序：不同批次细胞的基因表达谱直接比较无意义，但细胞间相似性结构需要保持
|
||
- 3D 形状匹配：点云之间的欧氏距离不能直接用于匹配不同拓扑的形状
|
||
|
||
**与 OT 的关系**：当存在**isometric embedding** $f: \mathcal{X} \to \mathcal{Y}$ 使得 $d_\mathcal{Y}(f(x), f(x')) = d_\mathcal{X}(x, x')$ 时，$\operatorname{GW}_p(\mu, \nu) = W_p(f_\# \mu, \nu)$。因此 GW 是 OT 在结构匹配方向的推广。
|
||
|
||
---
|
||
|
||
### 7.2 无参数双样本检验
|
||
|
||
OT 距离可用于**双样本检验**（Two-Sample Testing），即判断两个样本集是否来自同一分布。
|
||
|
||
**基于 Wasserstein 距离的检验统计量**：
|
||
|
||
$$
|
||
T = W_2\left( \hat{\mu}_n, \hat{\nu}_m \right) \tag{7.2}
|
||
$$
|
||
|
||
其中 $\hat{\mu}_n, \hat{\nu}_m$ 是从两个样本集估计的经验分布。
|
||
|
||
**定理 7.1（渐近分布）**（Ramdas et al., 2017）。在 null hypothesis $H_0: \mu = \nu$ 下，当 $n, m \to \infty$ 且 $n/(n+m) \to \lambda \in (0,1)$ 时：
|
||
|
||
$$
|
||
2(n+m) \cdot W_2^2(\hat{\mu}_n, \hat{\nu}_m) \xrightarrow{d} \sum_{k=1}^\infty \lambda_k \cdot Z_k^2 \tag{7.3}
|
||
$$
|
||
|
||
其中 $Z_k \sim N(0,1)$ i.i.d.，$\lambda_k$ 是某个依赖 $\mu$ 的特征值序列。
|
||
|
||
**检验流程**：
|
||
1. 计算 $T = W_2(\hat{\mu}_n, \hat{\nu}_m)$（使用 Sinkhorn 近似以加速）
|
||
2. 通过置换检验（permutation test）估计 $p$-value：在 null 下交换样本标签，重计算 $T$，估计经验分布
|
||
3. 若 $p < \alpha$，拒绝 $H_0$
|
||
|
||
这提供了比 KS 检验、MMD 更敏感的检验方法，尤其在高维小样本情形下。
|
||
|
||
---
|
||
|
||
## 参考文献
|
||
|
||
1. Villani, C. (2008). *Optimal Transport: Old and New*. Springer.
|
||
2. Cuturi, M. (2013). Sinkhorn distances: Lightspeed computation of optimal transport. *NeurIPS*.
|
||
3. Arjovsky, M., Chintala, S., & Bottou, L. (2017). Wasserstein GAN. *ICML*.
|
||
4. Gulrajani, I., et al. (2017). Improved training of Wasserstein GANs. *NeurIPS*.
|
||
5. Lipman, Y., et al. (2022). Flow matching via optimal transport. *ICLR*.
|
||
6. Pooladian, E., et al. (2023). Multivariate conditional flow matching. *ICLR*.
|
||
7. Cuturi, M., et al. (2019). Functional optimal transport: Mapper correspondence. *NeurIPS*.
|
||
8. Mémol, G., et al. (2022). Unbalanced optimal transport: Dynamic and nonnegative. *ICLR*.
|
||
9. Ramdas, A., et al. (2017). Wasserstein tests for two-sample problems. *NeurIPS*.
|
||
10. Peyré, G., & Cuturi, M. (2019). Computational optimal transport. *Foundations and Trends in Machine Learning*.
|
||
|
||
---
|
||
|
||
*本笔记由 Claude Code 自动生成，参考截至 2026 年的最优传输理论与深度学习交叉领域的研究进展。* |