4-最优传输

最优传输（Optimal Transport, OT）

摘要：最优传输（Optimal Transport, OT）是概率论与几何分析的交叉领域，为度量概率分布之间的距离提供了坚实的数学基础。近年来，OT 理论在深度学习与生成式模型中发挥了核心作用——从 WGAN 的梯度惩罚、扩散模型的 flow matching，到对比学习中的 Sinkhorn Divergence，都可见其影响。本笔记面向深度学习与生成式模型科研人员，系统阐述 OT 的数学基础、计算方法及其与生成模型的深刻联系。

---

1. 最优传输的基本问题

1.1 Monge 问题

定义 1.1（Monge 问题）。给定两个概率空间 和 ，以及代价函数 ，Monge 问题寻求一个确定性映射 ，使得

其中约束 表示 将 的质量传输到 ，即对任意可测集 有 。

Monge 问题的物理直觉是：有一堆”泥土”（由 描述）分布在 上，需要将其搬运到目标堆（由 描述）分布在 上，每单位泥土从 运到 的成本是 ，求总成本最小的搬运方案。

关键性质：Monge 问题不是凸的——映射空间上的约束 使得可行集非凸。因此直接优化 (1.1) 通常非常困难。当 且 时，若 绝对连续且满足某些正则性条件，唯一解由 Brenier 势函数 的梯度给出：。

---

1.2 Kantorovich 松弛

定义 1.2（Kantorovich 问题）。为克服 Monge 问题的非凸性，Kantorovich (1942) 引入联合分布（耦合）来代替确定性映射：

其中 表示所有以 和 为边缘分布的联合概率分布：

即 在 上的边际为 ，在 上的边际为 。

Kantorovich 形式的物理直觉是：不规定”每粒泥土具体去哪里”，只要求整体分布匹配 和 。这使得问题变成关于联合分布 的线性规划，是凸的。

定理 1.1（Monge-Kantorovich 对偶）。在适当的紧性条件下，最优传输问题有如下对偶形式：

证明梗概：将 (1.2) 视为线性规划。其对偶变量是边际约束的拉格朗日乘子，即对偶问题中寻找函数对 使得约束 下的线性泛函最大化。根据线性规划的对偶理论，强对偶性成立（当边际分布是紧概率空间时）。

---

1.3 两者的等价性与适用场景

关系：

• 如果 Monge 问题存在最优映射 ，则由 诱导的传输计划 （即 ）是 Kantorovich 问题的最优解，且两目标值相等。
• 反之不一定成立：当 是严格凸函数（如 ）时，若 绝对连续，则最优传输映射存在且唯一（Brenier 定理）；但对于一般 ，最优解可能是非确定性的，此时 Kantorovich 松弛严格更优。

选择原则：

场景	推荐形式	原因
二次代价 （欧式空间）	Monge（Brenier 映射）	有显式结构，映射唯一
离散分布 / 图结构	Kantorovich	线性规划易于求解
计算梯度（神经网络）	Kantorovich（熵正则化）	对偶形式可微

---

2. Wasserstein 距离

2.1 -Wasserstein 距离的定义

定义 2.1（-Wasserstein 距离）。给定 上的两个概率分布 ，以及代价函数 （通常为度量），-Wasserstein 距离定义为

当 是度量时， 确实构成 （-阶矩有限的概率分布空间）上的一个度量。

---

2.2 距离的特例

当 （欧氏距离的平方）时，得到 Wasserstein-2 距离：

在生成模型中最为常用，因为：

1. 当 绝对连续时，存在唯一的最优传输映射 ，使得 。
2. 由 Kantorovich-Rubinstein 对偶， 距离（）有简化的对偶形式（见 5.1 节），但 的几何结构更丰富。

---

2.3 Wasserstein 距离的性质

定理 2.1（度量结构）。 满足：

• 非负性：，且
• 对称性：
• 三角不等式：
• 弱收敛拓扑： 收敛等价于弱收敛 + -阶矩收敛（在紧空间上）

性质 2.1（相对于 距离的优势）。与 KL 散度、JS 散度不同，Wasserstein 距离不需要分布的支撑集有重叠就能有效度量距离。这使其特别适合处理生成分布 与目标分布 支撑不重叠的情形（如 GAN 中常见的模式坍塌问题）。

---

3. Sinkhorn 算法

3.1 熵正则化最优传输

直接求解 Kantorovich 问题 (1.2) 的计算复杂度为 （离散情形下， 为支撑点数），不适用于大规模数据。

Cuturi (2013) 引入熵正则化：

其中 是代价矩阵， 是联合分布的熵， 是正则化参数。

当 时，(3.1) 的解收敛到原始 Kantorovich 问题 (1.2) 的解；当 时，解趋近于独立乘积 。

---

3.2 Sinkhorn 迭代

将 约束为 中的双随机矩阵（即 ，），用拉格朗日形式写出 (3.1) 的KKT条件，可导出对偶变量 和 ，使得最优耦合具有 Sinkhorn 形式：

其中 称为 Gibbs 核。

定理 3.1（Sinkhorn 迭代）。固定点迭代给出如下交替投影：

其中 ， 是边际分布向量。

矩阵-向量形式：

每轮交替更新称为一次 Sinkhorn 迭代（sinkhorn normalization）。

---

3.3 矩阵形式的 Sinkhorn 循环

设 为代价矩阵，。Sinkhorn 算法维护两个缩放向量 和 ，迭代直至收敛：

输入: 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)

最优传输计划为 ，对应的目标值：

---

3.4 收敛性分析

定理 3.2（Sinkhorn 收敛速度）。设代价矩阵 的 entries 有界，，则 Sinkhorn 迭代线性收敛到唯一最优解：

其中 是最优耦合矩阵（双随机）。

实际注意事项：

• 当 过小时， 极易溢出（underflow），需要 log-space 稳定化。
• 迭代次数 通常设为 量级以达到 -近似。
• 当支撑点数量 很大时（> 10^5），全矩阵 无法存储，可使用 kernel approximation 或 sparse Sinkhorn。

---

4. 无熵约束的 Sinkhorn（稳定版）

4.1 log-space 稳定计算

当 很小或 很大时，直接计算 会下溢（round to zero）。解决方案是始终在 log-space 操作，存储 ，然后使用 log-sum-exp 技巧来避免数值不稳定。

对 (3.4) 取对数，设 ，，则：

对每一行 ，计算 ，然后通过 log-sum-exp 得到 ：

其中 ，可通过减去 来稳定计算：

---

4.2 完整 log-space Sinkhorn

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 域进行，防止 溢出
2. log-sum-exp 通过 (4.3) 稳定计算
3. 收敛判定使用对偶变量差的 范数

---

5. Wasserstein 距离的梯度

5.1 Sinkhorn 的对偶形式

熵正则化问题的对偶（由 (1.3) 扩展而来）为：

其中 是对偶变量。

设 ，，则 (5.1) 的稳定性条件给出 Sinkhorn 迭代的固定点 。

---

5.2 Sinkhorn 梯度在神经网络中的应用

设损失函数 ，其中 是神经网络参数化 生成的分布（如生成器输出的隐空间分布）， 是真实数据分布。

由于 Sinkhorn 算法给出的是近似最优耦合 ，我们可以计算 对 的梯度。

链式法则：

而 ，对 求导涉及神经网络输出的 ：

其中 ，所以

这允许通过标准反向传播计算梯度，尽管需要存储 （或在对数空间中存储 ）。

---

5.3 梯度正则化问题

由于 是近似解，梯度估计存在偏差（bias）：

• 当 较大时，近似误差大，但方差小
• 当 较小时，近似误差小，但 的条件数变差，梯度方差增大

解决方案：

1. Shamir et al. (2023) 的”unbiased Sinkhorn”：通过重参数化技巧消除近似偏差
2. Neural Sinkhorn（Genevay et al., 2016）：用神经网络学习最优传输映射而非迭代求解
3. 截断梯度（truncated gradient）：设定 的下界阈值以避免极端值

---

6. 与生成模型的关系

6.1 WGAN：Wasserstein-1 距离代替 JS 散度

Arjovsky et al. (2017) 提出使用 距离替代 JS 散度来度量生成分布 与真实分布 的距离。

Kantorovich-Rubinstein 对偶给出 的简化形式：

即 是 1-Lipschitz 函数空间上的最大均值差异（MMD）。

关键洞察：对判别器 施加 Lipschitz 约束（而非 Jensen-Shannon 的 log-likelihood 约束），可以在分布支撑不重叠时仍提供有意义的梯度——这解决了 standard GAN 在模式坍塌时梯度消失的问题。

WGAN 的目标函数：

---

6.2 梯度惩罚（Gradient Penalty）的推导与作用

Gulrajani et al. (2017) 发现直接约束 在实践中难以优化，因此提出梯度惩罚（WGAN-GP）：

其中 ，，，。

推导： 的最优判别器 满足 almost everywhere（在 和 的支撑上）。因此梯度惩罚强制学到的 逼近这一性质。

正则化解释：在 的对偶问题中，最优解对应约束 （次梯度），梯度惩罚是对该约束的恰当正则化，使得问题稳定可解。

---

6.3 扩散模型中的 OT：Optimal Transport in Flow Matching

Flow Matching（Lipman et al., 2022; Pooladian et al., 2023）将扩散模型的生成过程建模为路径测度（path measure），在连续时间下从噪声分布 流向数据分布 。

最优传输 flow matching 使用 OT 映射来设计更好的条件路径：

其中 是从 到 的最优传输映射。

关键优势：

• OT-FM 避免了标准 FM 的边缘不匹配问题
• 训练目标简化为回归一个确定的向量场（而非随机路径）
• 在图像生成任务中，使用 OT-FM 可以显著减少推理步骤（从 1000 步降至约 50 步仍保持高质量）

则 OT flow 给出确定性向量场：

这正是 Jordan-Kinder-Pearson (JKP) 方程 在 时的极限形式。

---

6.4 对比学习中的 OT：Sinkhorn Divergence

正面样本对 和负面样本对 在特征空间中应该满足：与正面样本的距离应小于与负面样本的距离。

Sinkhorn Divergence（Cuturi et al., 2019; Mensch et al., 2022）定义为：

其中 是 Sinkhorn 近似的熵正则化传输代价。

几何直觉：

• 给出 到 的有向传输代价
• 减去 使 具有对称性且满足
• 当 时，；当 时， 的某种变形

在对比学习中，可以将 作为损失的一部分：

这鼓励模型将同一样本的正负面映射到特征空间中 OT-最优的位置。

---

7. 理论延伸

7.1 Gromov-Wasserstein：结构匹配（非度量空间）

当两个分布所在的空间不兼容直接比较时（例如， 上的代价函数无法定义为有意义的距离），Gromov-Wasserstein (GW) 距离提供了一种基于结构相似性的比较框架。

定义 7.1（Gromov-Wasserstein 距离）。给定两个度量空间 和 ，-Gromov-Wasserstein 距离定义为：

物理直觉：不要求 与 之间的代价可直接比较，而是比较成对距离结构——如果两个空间中的点对距离结构相似，则它们在 GW 意义下接近。

应用场景：

• 图匹配（graph matching）：节点之间无法直接对应，但图的结构（邻接矩阵特征值、路径长度）需要匹配
• 单细胞 RNA 测序：不同批次细胞的基因表达谱直接比较无意义，但细胞间相似性结构需要保持
• 3D 形状匹配：点云之间的欧氏距离不能直接用于匹配不同拓扑的形状

与 OT 的关系：当存在isometric embedding 使得 时，。因此 GW 是 OT 在结构匹配方向的推广。

---

7.2 无参数双样本检验

OT 距离可用于双样本检验（Two-Sample Testing），即判断两个样本集是否来自同一分布。

基于 Wasserstein 距离的检验统计量：

其中 是从两个样本集估计的经验分布。

定理 7.1（渐近分布）（Ramdas et al., 2017）。在 null hypothesis 下，当 且 时：

其中 i.i.d.， 是某个依赖 的特征值序列。

检验流程：

1. 计算 （使用 Sinkhorn 近似以加速）
2. 通过置换检验（permutation test）估计 -value：在 null 下交换样本标签，重计算 ，估计经验分布
3. 若 ，拒绝

这提供了比 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.

---