1781 年,Gaspard Monge 在研究土木工程时提出了一个问题:如何以最小的代价将一堆土从一个地方搬到另一个地方?这个看似质朴的问题,在两百多年后成为概率论、优化和机器学习的交汇点。最优传输理论不仅给出了一种衡量概率分布之间距离的原则性方法,还揭示了概率度量与函数空间之间的深刻对偶关系。

这篇文章从 Monge 的原始问题出发,解释 Kantorovich 松弛如何克服 Monge 问题的困难,推导 Wasserstein 距离的定义与对偶形式,并回答一个在生成模型中至关重要的问题:为什么 Wasserstein 距离比 KL 散度更适合衡量分布之间的差异?

Monge 问题:搬运土堆的最小代价

给定源分布 μ\mu(土堆的当前位置)和目标分布 ν\nu(土堆的目标位置),Monge 问题寻找一个映射 T:XYT: X \to Y,使得搬运总代价最小:

infTXc(x,T(x))dμ(x)s.t.T#μ=ν\inf_{T} \int_X c(x, T(x)) d\mu(x) \quad \text{s.t.} \quad T_\#\mu = \nu

其中 c(x,y)c(x, y) 是将单位质量从 xx 运到 yy 的代价函数(通常取 c(x,y)=xypc(x,y) = \|x - y\|^p),T#μ=νT_\#\mu = \nu 表示 TTμ\mu 推前到 ν\nu——即对于任意可测集 BBν(B)=μ(T1(B))\nu(B) = \mu(T^{-1}(B))

Monge 问题的核心困难在于映射 TT 是确定性的:每个源点 xx 必须映射到唯一的目标点 T(x)T(x)。这意味着源分布中一个点的质量不能被"分裂"到多个目标点。当源分布和目标分布的质量分布不兼容时(例如 μ\mu 是连续分布而 ν\nu 有离散点),满足 T#μ=νT_\#\mu = \nu 的映射 TT 可能根本不存在。

Kantorovich 放松:从确定性映射到随机耦合

1942 年,Leonid Kantorovich 提出了一个关键松弛:将确定性映射 TT 替换为随机耦合(传输计划)γ\gamma。Kantorovich 问题描述为:

Wpp(μ,ν)=infγΠ(μ,ν)X×Yc(x,y)dγ(x,y)W_p^p(\mu, \nu) = \inf_{\gamma \in \Pi(\mu, \nu)} \int_{X \times Y} c(x, y) d\gamma(x, y)

其中 Π(μ,ν)\Pi(\mu, \nu) 是所有以 μ\muν\nu 为边缘分布的联合分布的集合:

Π(μ,ν)={γ0:γ(A×Y)=μ(A),  γ(X×B)=ν(B)}\Pi(\mu, \nu) = \{\gamma \geq 0 : \gamma(A \times Y) = \mu(A), \; \gamma(X \times B) = \nu(B)\}

这个松弛的关键优势是:Π(μ,ν)\Pi(\mu, \nu) 总是非空的(μν\mu \otimes \nu 就是一个平凡的耦合),所以 Kantorovich 问题总是有解。此外,当 Monge 问题的最优映射存在时,Kantorovich 问题的最优传输计划 γ\gamma^* 恰好是 (Id,T)#μ(Id, T^*)_\#\mu,即 Monge 解的图。

直观上,传输计划 γ(x,y)\gamma(x, y) 描述了"从 xx 搬运多少质量到 yy"。与 Monge 的确定性映射不同,一个源点的质量可以被分配到多个目标点,这正对应了"质量分裂"的需求。

1D 最优传输映射
Loading visualization...
上方:源分布(两个高斯峰),下方:目标分布(三个高斯峰),中间为传输计划的彩色映射线,线宽正比于传输质量。拖动滑块调整分布参数。

Wasserstein 距离的定义

Kantorovich 问题的最优值定义了 Wasserstein 距离(也称 Kantorovich-Monge 距离)。对于概率测度空间 Pp(Rd)\mathcal{P}_p(\mathbb{R}^d) 上的两个分布 μ\muν\nupp-Wasserstein 距离定义为:

Wpp(μ,ν)=infγΠ(μ,ν)Rd×Rdxypdγ(x,y)W_p^p(\mu, \nu) = \inf_{\gamma \in \Pi(\mu, \nu)} \int_{\mathbb{R}^d \times \mathbb{R}^d} \|x - y\|^p d\gamma(x, y)

其中 Π(μ,ν)\Pi(\mu, \nu) 是以 μ\muν\nu 为边缘分布的所有联合分布之集合。WpW_p 满足度量公理(非负性、对称性、三角不等式),这使得它是一个真正的距离函数——这一点与 KL 散度形成鲜明对比。

在 1D 情况下,WpW_p 有解析解。设 FμF_\muFνF_\nu 分别为 μ\muν\nu 的累积分布函数,Fμ1F_\mu^{-1}Fν1F_\nu^{-1} 为分位函数,则:

Wpp(μ,ν)=01Fμ1(t)Fν1(t)pdtW_p^p(\mu, \nu) = \int_0^1 |F_\mu^{-1}(t) - F_\nu^{-1}(t)|^p dt

这个结果反映了 1D 最优传输的本质:按分位数对齐。排在第 tt 分位的源点被传输到排在第 tt 分位的目标点,这是单调递增映射,也是 Monge 最优映射。

Kantorovich-Rubinstein 对偶

线性规划的对偶理论可以应用于 Kantorovich 问题。对于 p=1p = 1 的特殊情况,Wasserstein 距离有一个优美的对偶形式——Kantorovich-Rubinstein 对偶:

W1(μ,ν)=supfL1Eμ[f]Eν[f]W_1(\mu, \nu) = \sup_{\|f\|_L \leq 1} \left|\mathbb{E}_\mu[f] - \mathbb{E}_\nu[f]\right|

其中 fL=supxyf(x)f(y)xy\|f\|_L = \sup_{x \neq y} \frac{|f(x) - f(y)|}{\|x - y\|} 是 Lipschitz 半范数。对偶问题的含义是:寻找一个 1-Lipschitz 函数 ff,使得 ff 在两个分布下的期望差异最大。这个最大的差异就是 W1W_1 距离。

推导路线图:从离散情形到连续对偶的逻辑链如下(详细推导见苏剑林《从Wasserstein距离、对偶理论到WGAN》):

  1. 线性规划形式:离散情形下,传输计划 γ\gamma 变为矩阵,Kantorovich 问题写为 minγC,γ\min_{\gamma} \langle C, \gamma \rangle,约束为 γ0\gamma \geq 0γ1=a\gamma \mathbf{1} = \mathbf{a}γT1=b\gamma^T \mathbf{1} = \mathbf{b},其中 CC 是代价矩阵,a,b\mathbf{a}, \mathbf{b} 是边缘分布向量。

  2. 弱对偶(Lagrange 松弛):对等式约束引入 Lagrange 乘子 f,gf, g,构造 Lagrangian L=C,γfT(γ1a)gT(γT1b)L = \langle C, \gamma \rangle - f^T(\gamma \mathbf{1} - \mathbf{a}) - g^T(\gamma^T \mathbf{1} - \mathbf{b})。对偶函数给出原问题的下界——这就是弱对偶。

  3. Farkas 引理的几何直觉:弱对偶的间隙何时为零?Farkas 引理提供了关键:要么 b\mathbf{b} 在锥 {Ax:x0}\{A\mathbf{x}: \mathbf{x} \geq 0\} 内(存在可行解),要么存在分离超平面 y\mathbf{y} 使得 yTA0\mathbf{y}^T A \geq 0yTb<0\mathbf{y}^T \mathbf{b} < 0。这保证了当原问题有可行解时,对偶间隙为零。

  4. 强对偶:Farkas 引理保证了弱对偶间隙为零——原问题的最优值等于对偶问题的最优值。

  5. 连续化:将求和 \sum 替换为积分 \int,向量 f,g\mathbf{f}, \mathbf{g} 替换为函数 f(x),g(y)f(x), g(y),代价矩阵 CijC_{ij} 替换为代价函数 c(x,y)c(x, y),即得到连续形式的 Kantorovich-Rubinstein 对偶。

对偶形式在机器学习中有直接的应用。WGAN(Wasserstein GAN)的判别器就是在学习这个最优 1-Lipschitz 函数 ff,其输出值近似 W1(Pr,Pg)W_1(\mathbb{P}_r, \mathbb{P}_g)。判别器的 Lipschitz 约束通过权重裁剪(weight clipping)或梯度惩罚(gradient penalty)来实现。

对于一般的 pp,对偶形式为:

Wpp(μ,ν)=supf,g{Xfdμ+Ygdν:f(x)+g(y)c(x,y)}W_p^p(\mu, \nu) = \sup_{f, g} \left\{\int_X f d\mu + \int_Y g d\nu : f(x) + g(y) \leq c(x, y)\right\}

其中 ffgg 是一对对偶势函数。这个对偶形式在 Sinkhorn 算法和 entropic OT 中有重要应用。

KL vs W2W_2:为什么 W2W_2 更适合衡量分布距离?

这是最优传输理论在生成模型中最核心的洞察。KL 散度与 Wasserstein 距离在度量分布差异时表现出截然不同的行为,而这种差异直接影响了训练的稳定性和收敛性。

KL 散度的三个缺陷

  1. 不对称性KL(μν)KL(νμ)\text{KL}(\mu \| \nu) \neq \text{KL}(\nu \| \mu),KL 不满足度量公理中的对称性。这使得基于 KL 的优化目标缺乏几何直觉。

  2. 无穷大问题:当 μ\muν\nu 的支撑集没有重叠时,KL(μν)=+\text{KL}(\mu \| \nu) = +\infty。考虑两个一维高斯分布 μ=N(0,1)\mu = \mathcal{N}(0, 1)ν=N(θ,1)\nu = \mathcal{N}(\theta, 1),当 θ\theta00 变到 1010 时,KL 散度单调增长,但永远不会跳到无穷大——因为高斯分布的支撑集是整个 R\mathbb{R}。然而,如果考虑两个定义在不同区间上的均匀分布,例如 μ=U[0,1]\mu = U[0, 1]ν=U[θ,θ+1]\nu = U[\theta, \theta+1],当 θ>1\theta > 1KL(μν)=+\text{KL}(\mu \| \nu) = +\infty

  3. 梯度消失:在实践中更严重的问题是,即使 KL 不是无穷大,当两个分布几乎无重叠时,KL 的梯度也会趋于零或极不稳定。这意味着生成器无法从判别器那里获得有效的梯度信号来改善生成质量。

W2W_2 的优势

W2W_2 克服了上述所有问题:

  • 对称性W2(μ,ν)=W2(ν,μ)W_2(\mu, \nu) = W_2(\nu, \mu)
  • 有限值:即使两个分布的支撑集完全不重叠,W2W_2 仍然是有限值。例如 W2(U[0,1],U[θ,θ+1])=θ+constW_2(U[0,1], U[\theta, \theta+1]) = |\theta| + \text{const}
  • 有意义的梯度W2W_2 随分布之间的距离平滑变化,即使分布没有重叠,也能提供有意义的梯度信号
KL vs W₂ 的关键差异
Loading visualization...
两个高斯分布逐步分离。左:KL 散度值——当重叠减小时急剧增长;右:W₂ 距离值——平滑线性增长。拖动滑块控制分布间距。

苏剑林的核心发现:扩散模型与 W2W_2 的联系

苏剑林在《生成扩散模型漫谈(十六):W距离 ≤ 得分匹配》中指出,扩散模型的得分匹配损失可以写成 W2W_2 距离的一个上界。具体地,对于扩散 ODE 的训练目标:

LDSM=Et,x0,ϵ[sθ(xt,t)xtlogpt0(xtx0)2]\mathcal{L}_{\text{DSM}} = \mathbb{E}_{t, x_0, \epsilon}\left[\|s_\theta(x_t, t) - \nabla_{x_t}\log p_{t|0}(x_t|x_0)\|^2\right]

可以证明这个目标与 W2W_2 距离存在不等式关系。这意味着优化得分匹配损失隐含地在缩小 W2W_2 距离,这从另一个角度解释了扩散模型为什么能产生高质量的生成结果——它在优化一个比 KL 更有几何意义的度量。

证明思路分三步概括:

  1. 构造传输方案:将扩散 ODE 的轨迹视为从 p0p_0pTp_T 的传输方案。ODE dx=f(x,t)dtdx = f(x, t)dt 的解给出了一条从噪声到数据的连续路径,这恰好定义了一个将概率质量从 p0p_0 传输到 pTp_T 的方案。

  2. 柯西不等式放缩:对传输代价 dx/dt2dt\int \|dx/dt\|^2 dt 使用柯西不等式,上界中包含 xlogpt(x)2\|\nabla_x \log p_t(x)\|^2 项——这正是得分函数的范数平方。这一步将传输代价与得分匹配目标联系起来。

  3. 常数变易法:利用 ODE 的解结构,将时间积分转化为训练目标的形式。扩散 ODE 的漂移项 f(x,t)=12xlogpt(x)f(x, t) = \frac{1}{2}\nabla_x \log p_t(x) 与得分函数直接相关,常数变易法将累积的传输代价表达为得分匹配损失的时间积分。

需要指出的是,上述证明在 ODE 情形下是完整的,但 SDE 情形尚未完全闭合——涉及噪声调度 β(t)\beta(t) 的单调递增性和初始项 W2(p~T,pT)W_2(\tilde{p}_T, p_T) 的控制,这些条件在一般情形下不一定满足。

应用

RLHF 中的分布对齐

在强化学习人类反馈(RLHF)中,语言模型需要从初始策略 πref\pi_{\text{ref}} 更新到对齐策略 πθ\pi_\theta,同时不能偏离太远。KL 散度 KL(πθπref)\text{KL}(\pi_\theta \| \pi_{\text{ref}}) 是最常用的约束,但它有一个已知问题:当 πθ\pi_\thetaπref\pi_{\text{ref}} 几乎为零的区域分配概率时,KL 会给出无穷大的惩罚。W2W_2 距离提供了一种替代约束,它对概率质量的"移动距离"敏感而非对概率比敏感,在理论上更适合控制策略更新的幅度。

Flow Matching 与 OT-CFM

Flow Matching 是一种新型的连续归一化流方法。其关键变体 OT-CFM(Optimal Transport Conditional Flow Matching)直接使用最优传输映射来构建从噪声到数据的流路径。与使用独立高斯噪声的标准 Flow Matching 相比,OT-CFM 的路径更短、更直,训练收敛更快。这利用了 W2W_2 最优传输映射的几何性质:最短路径。

WGAN 的训练稳定性

WGAN 用 W1W_1 距离替代 JS 散度作为 GAN 的训练目标。Arjovsky 等人 (2017) 证明,在数据分布和生成分布不重叠的情况下,JS 散度的梯度为零,而 W1W_1 的梯度始终存在。这解释了 WGAN 相比原始 GAN 的训练稳定性优势。WGAN-GP 进一步用梯度惩罚替代权重裁剪,更好地实现了 Lipschitz 约束。

讨论:WGAN 的成功真的来自 W 距离吗?

苏剑林在《WGAN的成功,可能跟Wasserstein距离没啥关系》(kexue.fm/archives/8244) 中提出了一个重要的反面观点:WGAN 的训练稳定性可能主要归功于梯度惩罚(或权重裁剪)对判别器的 Lipschitz 约束的正则化效果,而非 WW 距离本身的数学优越性。

实验证据支持这一判断。即使 WW 距离的近似质量很差——例如判别器容量不足,无法精确估计 W1W_1——WGAN-GP 仍然能够稳定训练并产生高质量的生成结果。这意味着稳定性来自 Lipschitz 约束对判别器梯度的正则化(防止梯度爆炸/消失),而非来自 WW 距离的精确估计。Lipschitz 约束保证了判别器输出不会随输入微小变化而剧烈震荡,从而为生成器提供平滑的梯度信号——这与具体的距离度量无关。

这个反面观点提醒我们:理论优雅性和工程有效性之间可能有鸿沟。WW 距离在数学上优于 JS 散度是确定的,但 WGAN 训练稳定性的工程优势可能更多来自 Lipschitz 正则化这一副作用,而非距离度量本身的优越性。区分这两者有助于在后续研究中做出更准确的归因。

参考文献

  1. 苏剑林. 生成扩散模型漫谈(十六):W距离 ≤ 得分匹配. https://kexue.fm/archives/9467
  2. 苏剑林. 从Wasserstein距离、对偶理论到WGAN. https://kexue.fm/archives/6282
  3. Villani, C. (2009). Optimal Transport: Old and New. Springer.
  4. Arjovsky, M., Chintala, S., & Bottou, L. (2017). Wasserstein Generative Adversarial Networks. ICML.
  5. Gulrajani, I., et al. (2017). Improved Training of Wasserstein GANs. NeurIPS.
  6. Lipman, Y., et al. (2023). Flow Matching for Generative Modeling. ICLR.
  7. Tong, A., et al. (2023). Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport. TMLR.