最优传输深度实战：从蒙日问题、Kantorovich 松弛与 Wasserstein 距离到 Sinkhorn 熵正则化、WGAN 与最优传输映射的完整工程链路

最优传输（Optimal Transport, OT）研究的是"如何以最小代价把一种分布的质量搬运成另一种分布"。它看似是纯粹的数学——蒙日（Monge）在 1781 年研究土方工程时提出，但今天它已经变成了深度学习、生成模型、域适应、公平机器学习乃至蛋白质结构比对的核心工具。Wasserstein GAN 用它替代 KL/JS 散度让生成训练稳定；领域自适应用它对齐源域与目标域的特征分布；Wasserstein 距离本身是一种几何意义上的"度量"，比 f-散度更能刻画不重叠分布之间的距离。

本文延续本站"深度实战"的风格：从第一性原理出发，不调用任何专用 OT 库，用 NumPy 从零实现 Wasserstein 距离、Monge 映射、Sinkhorn 算法、Wasserstein 重心、Kantorovich–Rubinstein 对偶，并落到颜色迁移这一工程场景，所有代码均可直接运行验证。

一、第一性原理：蒙日问题 vs Kantorovich 松弛

最初的蒙日问题是一个确定性的映射问题：给定源分布 μ 与目标分布 ν，寻找一个映射 T，使得 T 把 μ "推"成 ν（记为 T#μ = ν），并最小化搬运代价：

$$\inf_{T,\;T\#\mu=\nu} \int c(x, T(x))\, d\mu(x)$$

这里 c(x,y) 是单位质量的搬运代价，最常用的是欧氏距离 c(x,y)=‖x−y‖^p。

蒙日问题的致命缺陷是：最优映射未必存在。例如当 μ 是一个 Dirac 单点质量、而 ν 是两个点的混合时，单点无法同时确定性地映射到两个点。1942 年 Kantorovich 用一个更宽松的联合分布（传输计划，coupling） 表述拯救了这个理论：

$$\inf_{\pi \in \Pi(\mu,\nu)} \int c(x,y)\, d\pi(x,y)$$

其中 Π(μ,ν) 是所有边缘分布为 μ 和 ν 的联合分布集合。蒙日映射是"确定性计划"（每行/列只有一个非零元），而 Kantorovich 允许"软计划"（质量可以分裂到多个目标）。这一松弛不仅让问题永远有解，还把它变成了线性规划，从而可以高效求解。

二、一维 Wasserstein 与排序定理

在一维情形，最优传输有闭式解：单调重排（monotone rearrangement）——把源样本排序、目标样本排序，然后按秩一一配对就是最优的。这是"重排不等式"的直接推论。

import numpy as np

def wasserstein_1d(sa, sb, p=1):
    a = np.sort(sa); b = np.sort(sb)
    n = min(len(a), len(b)); a = a[:n]; b = b[:n]
    return np.mean(np.abs(a - b) ** p) ** (1.0 / p)

rng = np.random.default_rng(0)
mu = rng.normal(0, 1, 2000)
nu = rng.normal(3, 2, 2000)
print("W_1 (empirical) =", round(wasserstein_1d(mu, nu, 1), 4))
print("W_2 (empirical) =", round(wasserstein_1d(mu, nu, 2), 4))
# 一维高斯闭式：W_2^2(mu,nu) = (m1-m2)^2 + (s1-s2)^2
w2sq = (0 - 3) ** 2 + (1 - 2) ** 2
print("W_2^2 closed form =", w2sq, " sqrt =", round(w2sq ** 0.5, 4))

运行结果：W₁≈3.02、W₂≈3.18；而两个一维高斯 N(0,1) 与 N(3,2) 的解析 W₂² = (0−3)²+(1−2)² = 10，W₂=3.162，与经验值吻合。注意 W₂ 比均值差 |3−0|=3 略大，因为方差不同（σ=1 vs 2）也会贡献距离——这正是 Wasserstein 相比单纯比较均值更"几何"的地方。

对一维高斯，W₂ 有漂亮的闭式（Bures–Wasserstein）：W₂²(μ,ν) = ‖m_μ−m_ν‖² + Tr(Σ_μ+Σ_ν − 2(Σ_μ^{1/2}Σ_νΣ_μ^{1/2})^{1/2})。一维时退化为 (m₁−m₂)²+(σ₁−σ₂)²，我们会在第七节用大样本经验值验证它。

三、Monge 映射与推前测度：Kantorovich 为何必要

在一维，最优 Monge 映射有显式表达：先做源分布的累积分布函数 F_μ，再做目标分布的 quantile 函数 F_ν⁻¹，两者复合 T = F_ν⁻¹ ∘ F_μ 即最优搬运映射。下面验证它确实把 μ 推成了 ν（看推前样本的均值与方差是否对齐）：

import numpy as np

rng = np.random.default_rng(2)
xa = rng.normal(0, 1, 4000); xb = rng.normal(3, 2, 4000)
sa = np.sort(xa); sb = np.sort(xb)
ranks = np.searchsorted(sa, xa) / (len(sa) - 1)
T = np.quantile(xb, np.clip(ranks, 0, 1))   # Monge 映射 T = F_nu^{-1} o F_mu 的推前样本
print("E[mu]=", round(float(xa.mean()), 3), " E[nu]=", round(float(xb.mean()), 3),
      " E[T#mu]=", round(float(T.mean()), 3))
print("std[mu]=", round(float(xa.std()), 3), " std[nu]=", round(float(xb.std()), 3),
      " std[T#mu]=", round(float(T.std()), 3))

# Monge 会失败的情形：mu=delta_0, nu=0.5 d_-1 + 0.5 d_+1
# 单点无法同时映射到两个目标 -> Kantorovich 让质量分裂
Cd = np.abs(np.array([0.0])[:, None] - np.array([-1.0, 1.0])[None, :])
def sinkhorn_marg(C, a, b, reg=1.0, iters=3000):
    K = np.exp(-C / reg); u = np.ones(len(a)); v = np.ones(len(b))
    for _ in range(iters):
        u = a / (K @ v); v = b / (K.T @ u)
    return np.diag(u) @ K @ np.diag(v)
Pd = sinkhorn_marg(Cd, np.array([1.0]), np.array([0.5, 0.5]), reg=1.0)
print("Kantorovich plan (delta_0 -> 0.5 d_-1 + 0.5 d_+1):")
print("  P =", np.round(Pd, 3), " colsums =", np.round(Pd.sum(axis=0), 3))

推前样本 T#μ 的均值（≈3.10）与标准差（≈1.99）几乎完美复刻 ν——证明 T 确实是最优搬运。随后展示 Kantorovich 松弛的价值：当 μ 是单点 δ₀、ν 是 0.5δ₋₁+0.5δ₊₁ 时，蒙日映射无解，但 Kantorovich 计划给出 P=[[0.5, 0.5]]，让这 1 单位质量分裂到两个目标，列和严格等于 (0.5, 0.5)。这正是"松弛"带来的存在性保证。

四、Sinkhorn 熵正则化：把线性规划变成可微迭代

直接求解 Kantorovich 线性规划在大规模时是 O(n³) 且不可微，无法嵌进神经网络。Cuturi (2013) 的突破是加入熵正则项，把硬约束的传输计划变成可微、可并行的迭代缩放（iterative scaling）：

$$\min_{\pi\in\Pi} \langle C,\pi\rangle + \varepsilon\, \mathrm{KL}(\pi \| \mathbf{1}\mathbf{1}^\top)$$

其最优解形如 π* = diag(u) K diag(v)，其中 K = e^{−C/ε}，u、v 通过对行/列和交替归一化求出。正则项 ε→0 时退回精确 OT；ε 大时偏向均匀耦合（熵最大）。

import numpy as np

def wasserstein_1d(sa, sb, p=1):
    a = np.sort(sa); b = np.sort(sb)
    n = min(len(a), len(b)); a = a[:n]; b = b[:n]
    return np.mean(np.abs(a - b) ** p) ** (1.0 / p)

rng = np.random.default_rng(5)
n = 8
xa = np.sort(rng.normal(0, 1, n)); xb = np.sort(rng.normal(2, 1, n))
C = np.abs(xa[:, None] - xb[None, :])
opt_mean = wasserstein_1d(xa, xb, 1)   # 排序解即最优传输代价（均值）

def sinkhorn(C, reg=0.1, iters=2000):
    K = np.exp(-C / reg)
    a = np.ones(n) / n; b = np.ones(n) / n
    u = np.ones(n); v = np.ones(n)
    for _ in range(iters):
        u = a / (K @ v); v = b / (K.T @ u)
    return np.diag(u) @ K @ np.diag(v)

a = np.ones(n) / n; b = np.ones(n) / n
for reg in [0.5, 0.1, 0.02, 0.005]:
    P = sinkhorn(C, reg)
    rs = P.sum(axis=1); cs = P.sum(axis=0)
    cost = float(np.sum(C * P))
    print(f"reg={reg}: cost={cost:.4f} "
          f"rowsum_dev={float(np.max(np.abs(rs - a))):.2e} "
          f"colsum_dev={float(np.max(np.abs(cs - b))):.2e}")
print("optimal transport cost (mean) =", round(opt_mean, 4))

可以看到：随着 ε 从 0.5 降到 0.005，传输代价从 2.33 平滑收敛到 2.3227，且行和/列和偏差始终维持在 1e-17（即边缘分布约束被严丝合缝地满足）。这正是 Sinkhorn 的工程价值——它用几十次矩阵乘法逼近精确 OT，且整个流程对 GPU 友好、可反向传播，因此能作为神经网络的一层（Differentiable OT Layer）。

工程陷阱：ε 太小会导致 K=e^{−C/ε} 下溢成全零矩阵而数值发散；ε 太大则正则偏置显著、偏离真 OT。实践中 ε 取 0.01~0.1 并在迭代后做代价校正（如 Sinkhorn divergence 减去两个边际自传输）是常见做法。

五、Wasserstein 重心：在分布空间里做"平均"

既然 Wasserstein 给分布之间定义了距离，那"多个分布的平均"就有了几何意义——Wasserstein 重心是使到各分布加权 W₂² 距离之和最小的那一个分布。在一维，重心有闭式：对各个分布的 quantile 函数做加权平均即可。

import numpy as np

rng = np.random.default_rng(3)
qs = np.linspace(0, 1, 2000)
distros = [rng.normal(m, s, 8000) for m, s in [(0, 1), (3, 1), (1, 2)]]
Q = np.mean([np.quantile(d, qs) for d in distros], axis=0)
bary = np.interp(np.linspace(0, 1, 8000), qs, Q)
print("bary mean =", round(float(np.mean(bary)), 3), " (theory (0+3+1)/3 = 1.333)")
print("bary std  =", round(float(np.std(bary)), 3), " (theory (1+1+2)/3 = 1.333)")

三个分布 N(0,1)、N(3,1)、N(1,2) 的重心均值 ≈1.344、标准差 ≈1.326，与解析值 (0+3+1)/3=1.333、(1+1+2)/3=1.333 高度吻合（残差来自 8000 样本的 Monte Carlo 离散化噪声）。重心之所以"反直觉地"把标准差取平均而非按方差混合，是因为 Wasserstein 几何作用在 quantile 函数上、而非概率密度上——这正体现了 OT 与欧氏"均值"的本质不同。该思想已用于：纹理混合、形状插值、以及多任务模型参数的"联邦平均"（Federated Averaging 的分布版）。

六、Kantorovich–Rubinstein 对偶与 WGAN：一个线性规划证明对偶相等

当 p=1 时，Kantorovich 问题有一个优美的对偶形式（Kantorovich–Rubinstein 对偶）：

$$W_1(\mu,\nu) = \sup_{\|f\|_{\mathrm{Lip}}\le 1} \int f\,d\mu - \int f\,d\nu$$

即 Wasserstein 距离等于"所有 1-Lipschitz 函数 f 在 μ 与 ν 上期望差的上确界"。这正是 WGAN 的目标函数：训练一个 1-Lipschitz 的"评判器(critic)" f 去最大化 E[f(X)]−E[f(Y)]，其最优值就是 W₁。下面用线性规划直接求解该对偶，并验证它等于原始的 OT 距离（强对偶），无需任何 OT 专用库：

import numpy as np
from scipy.optimize import linprog

def kr_dual(p_mu, w_mu, p_nu, w_nu):
    pts = np.concatenate([p_mu, p_nu]); m = len(pts)
    obj = np.concatenate([w_mu, -w_nu])
    A, bb = [], []
    for i in range(m):
        for j in range(i + 1, m):
            d = abs(pts[i] - pts[j])
            r = np.zeros(m); r[i] = 1; r[j] = -1; A.append(r); bb.append(d)
            r2 = np.zeros(m); r2[i] = -1; r2[j] = 1; A.append(r2); bb.append(d)
    res = linprog(-obj, A_ub=np.array(A), b_ub=np.array(bb),
                  bounds=[(None, None)] * m, method='highs')
    return -res.fun

# 例1：mu=delta_0, nu=0.5 d_-1 + 0.5 d_+1  -> W_1 = 1.0
ex = kr_dual(np.array([0.0]), np.array([1.0]),
             np.array([-1.0, 1.0]), np.array([0.5, 0.5]))
print("KR dual (mu=delta_0, nu=0.5d_-1+0.5d_+1) =", round(float(ex), 4), " (truth W_1 = 1.0)")

# 例2：mu 在 {-2,-1,0,1}, nu 在 {0,1,2,3}, 均匀  -> W_1 = mean|sorted-sorted| = 2.0
ex2 = kr_dual(np.array([-2, -1, 0, 1.0]), np.ones(4) / 4,
              np.array([0, 1, 2, 3.0]), np.ones(4) / 4)
print("KR dual (4+4 example) =", round(float(ex2), 4), " (truth = 2.0)")

# 真实数据的对偶演示（同一 12 点子集上，对偶 == 原始 W_1）
rng = np.random.default_rng(7)
mu = rng.normal(0, 1, 8000); nu = rng.normal(3, 2, 8000)
sub = rng.choice(len(mu), 12, replace=False)
pm = np.sort(mu[sub]); pn = np.sort(nu[sub])
w1_sub = np.mean(np.abs(np.sort(pm) - np.sort(pn)))   # 排序解
dual_real = kr_dual(pm, np.ones(12) / 12, pn, np.ones(12) / 12)
print("W_1 (primal, sorting, 12+12 subset) =", round(w1_sub, 4))
print("KR dual (same 12+12 subset)         =", round(float(dual_real), 4), " (== W_1 by strong duality)")

两个小例子分别精确返回 1.0 与 2.0（与手算一致）；在 12 点真实数据子集上，对偶值 2.6947 与排序求出的原始 W₁ 2.6947 逐位相等——这就是强对偶的实证：线性规划的原问题与对偶问题给出同一个数。也由此看清 WGAN 的理论根基：它训练的那个 Lipschitz critic，其最优目标值就是这个对偶上确界，也就是真正的 W₁。

工程陷阱：WGAN 原始论文用"权重裁剪(weight clipping)"强制 Lipschitz，但裁剪会限制容量、导致梯度病态。改进方案（WGAN-GP）改为对随机插值点施加梯度范数惩罚 ‖∇f‖≈1，对应到对偶里就是更精确地约束 ‖f‖_Lip≤1。本节的 LP 解法在小规模是精确基准，但 O(n²) 约束规模无法用于百万级样本——这时 Sinkhorn 与 WGAN-GP 才是可扩展之选。

七、工程落地：最优传输域适应与颜色/直方图迁移

OT 最接地气的应用之一是分布对齐：给定源域样本与目标域样本，学习一个映射把源域分布"搬"到目标域。一维下，最优传输映射就是 rank-preserving 的 quantile 重映射——把源像素/特征值按其秩映射到目标的对应分位，从而在不改变结构的前提下匹配整体分布。下面用亮度直方图迁移演示：

import numpy as np

rng = np.random.default_rng(11)
src = rng.normal(0.3, 0.1, 30000)    # 源图像亮度
tgt = rng.normal(0.6, 0.15, 30000)   # 目标图像亮度
src_s = np.sort(src); tgt_s = np.sort(tgt)
ranks = np.searchsorted(src_s, src) / (len(src_s) - 1)
mapped = np.quantile(tgt, np.clip(ranks, 0, 1))   # 最优传输映射
print("src mean =", round(float(src.mean()), 3), " tgt mean =", round(float(tgt.mean()), 3),
      " mapped mean =", round(float(mapped.mean()), 3))
print("src std  =", round(float(src.std()), 3), " tgt std  =", round(float(tgt.std()), 3),
      " mapped std  =", round(float(mapped.std()), 3))

迁移后图像的亮度均值（≈0.599）与标准差（≈0.15）精确复刻了目标分布，而每个像素的相对秩被保留——这正是"风格迁移"中直方图匹配（histogram matching）的最优传输解释。在高维特征空间，这一思想演变为 Optimal Transport for Domain Adaptation (OTDA)：在源/目标特征联合上解 Sinkhorn，用得到的最优计划把带标签的源样本加权投影到目标域做分类。

最后，回到一维高斯的 W₂ 闭式，用大样本经验值验证其精确性作为收尾：

import numpy as np

def wasserstein_1d(sa, sb, p=1):
    a = np.sort(sa); b = np.sort(sb)
    n = min(len(a), len(b)); a = a[:n]; b = b[:n]
    return np.mean(np.abs(a - b) ** p) ** (1.0 / p)

rng = np.random.default_rng(13)
m1, s1, m2, s2 = 0.0, 1.0, 3.0, 2.0
w2sq = (m1 - m2) ** 2 + (s1 - s2) ** 2
print("W_2^2 closed form =", w2sq, " W_2 =", round(w2sq ** 0.5, 4))
emp = wasserstein_1d(rng.normal(m1, s1, 40000), rng.normal(m2, s2, 40000), 2)
print("W_2 empirical     =", round(emp, 4))

闭式 W₂=3.1623，大样本经验值 ≈3.1644，残差来自有限样本离散化，量级在 O(1/√N)，符合统计直觉。

八、与 AI / 深度学习的映射表

九、工程边界与陷阱

1. 计算复杂度：精确 OT 是 O(n³) 线性规划（网络单纯形）；Sinkhorn 是 O(n²·k) 但 ε 小需更多迭代，且对机器精度敏感（易下溢，需 log-sum-exp 稳定版）。
2. 维度诅咒：Wasserstein 在高维的经验估计需要样本量随维度指数增长——实践中常与特征降维、投影、或熵正则结合来缓解。
3. 熵正则偏置：Sinkhorn 距离在有限 ε 下系统性偏小（偏置约 ε·χ²/2），报告指标时应说明 ε 或用 Sinkhorn divergence 校正。
4. Lipschitz 约束：WGAN 的 critic 必须真正 1-Lipschitz，权重裁剪会损害表达力，梯度惩罚（WGAN-GP）或谱归一化更稳健。
5. 质量守恒假设：OT 默认总质量相等（balanced）；真实场景（如异常检测、部分匹配）需用 unbalanced OT（带 KL 边际罚项）。

结语

最优传输把"搬运质量"这个朴素的工程直觉，锻造成了一门有闭式解、有强对偶、可微、且深度嵌入现代机器学习的基础设施。从一维排序的 Wasserstein、到 Sinkhorn 的可微迭代、再到 WGAN 的 Lipschitz 对偶——本文所有结论都由从零实现的 NumPy 代码实证。它的真正魅力在于：当你需要让两个分布"对齐"时，OT 给出的不是近似启发，而是带有几何意义与最优性保证的答案。