最优传输与 Wasserstein 距离

以将质量从一个分布搬运到另一个分布的最小代价来度量两个概率分布之间的距离。

最优传输用把一个分布的质量重排为另一个分布所需的最小代价,来度量两个概率分布相距多远。先把要解决的问题摆清楚:给定两个概率分布 μ\mu(源)与 ν\nu(目标),以及一个把”从 xx 搬一单位质量到 yy“的代价记为 c(x,y)c(x,y) 的代价函数,我们要找一种搬运方式,既把 μ\mu 全部清空、把 ν\nu 恰好填满,又使总代价最小。这个最小总代价,就是 μ\muν\nu 之间的”传输距离”。

在一维且代价取平方距离的特例中,最优解有一个尤为简洁的形式:把源点与目标点各自排序后按序配对,所得的单调耦合即为最优。这一情形既给出了闭式解,也直观呈现出”质量沿最短总路程移动”的核心思想。Monge 形式寻找一个确定性映射 TT,以最小代价将 μ\mu 推前为 ν\nu

源点 μ目标点 ν
源点 μ目标点 ν最优匹配
总传输代价: 13.79

一维平方代价下,最优耦合就是把源点与目标点各自排序后按序配对(单调匹配)。任意交叉的配对都会严格增大总代价,因此连线从不交叉;拖动偏移时质量整体平移,总代价随之平滑变化。

infT#μ=νc(x,T(x))dμ(x),\inf_{T_\# \mu = \nu}\int c\big(x, T(x)\big)\,d\mu(x),

逐符号读:TT 是一个把每个源点 xx 唯一送往某个目标点 T(x)T(x) 的函数;约束 T#μ=νT_\#\mu=\nu(“推前”)要求被 TT 搬过去的质量恰好堆成 ν\nu,即对任意区域 AA,落入 AA 的源质量 μ(T1(A))\mu\big(T^{-1}(A)\big) 等于 ν(A)\nu(A);积分把每点搬运代价 c(x,T(x))c\big(x,T(x)\big) 按源密度 μ\mu 加权累加,inf\inf 取遍所有满足约束的 TT 求下确界。难点在于这样的映射未必存在:当一处源点的质量必须分给多个目标点时(“质量分裂”),没有任何单值函数 TT 能做到。不过对平方代价、且源连续时,由 Brenier 定理可知它终究存在,且是某个凸函数的梯度。

Kantorovich 松弛绕开了这个存在性障碍。它不再要求每个 xx 只能去一个地方,而用一个传输方案取代映射,即以 μ\muν\nu 为边缘的耦合 π\pi

infπΠ(μ,ν)c(x,y)dπ(x,y).\inf_{\pi \in \Pi(\mu,\nu)} \int c(x,y)\,d\pi(x,y).

耦合 π(x,y)\pi(x,y)(x,y)(x,y) 上的联合分布,记录”从 xx 搬多少到 yy”;Π(μ,ν)\Pi(\mu,\nu) 是所有边缘分别为 μ,ν\mu,\nu 的耦合之集 —— 边缘约束 π(x,)dy=μ\int\pi(x,\cdot)\,dy=\muπ(,y)dx=ν\int\pi(\cdot,y)\,dx=\nu 正是”清空 μ\mu、填满 ν\nu“的形式化。每个 Monge 映射都对应一个集中在曲线 y=T(x)y=T(x) 上的耦合,所以 Kantorovich 问题的解只会更优、绝不会更差。关键好处是:耦合总是存在(乘积测度 μν\mu\otimes\nu 就是一个),问题随之化为关于 π\pi 的线性规划 —— 目标 cdπ\int c\,d\piπ\pi 线性,边缘约束也线性。离散情形下 π\pi 是一个非负矩阵,其行和、列和被钉死为 μ,ν\mu,\nu;可行域是一个多胞形,最优解落在某个顶点上(对应一个稀疏的、近乎确定性的指派)。(本站约定:π\pi 始终表示耦合,γ\gamma 留作熵正则里的”温度”。)

每个这样的问题都有一个对偶。线性规划的对偶把”选哪张方案”换成”给传输定价”:引入一个势函数 f(x)f(x)g(y)g(y),在约束 f(x)+g(y)c(x,y)f(x)+g(y)\le c(x,y) 下最大化 fdμ+gdν\int f\,d\mu+\int g\,d\nu。直觉上 f,gf,g 是一套”运费表”,约束保证这套定价不会比直接搬运更贵,而对偶最优恰好等于原问题最优(强对偶)。对 W1W_1p=1p=1)而言,最优 g=fg=-fff 须为 1-Lipschitz,对偶塌缩成对单个 1-Lipschitz 势函数取上确界 —— 详见 Kantorovich 对偶;正是这个形式把 W1W_1 变成了 Wasserstein-GAN 的目标(判别器即那个被约束为 1-Lipschitz 的势函数)。

当代价取距离的 pp 次幂 c(x,y)=xypc(x,y)=\|x-y\|^p 时,最优值定义了 Wasserstein-pp 距离

Wp(μ,ν)=(infπΠ(μ,ν)xypdπ)1/p.W_p(\mu,\nu) = \left(\inf_{\pi\in\Pi(\mu,\nu)}\int \|x-y\|^p\,d\pi\right)^{1/p}.

外层的 1/p1/p 次方是为了让 WpW_p 满足三角不等式、成为一个真正的度量。最常用的是 p=2p=2(平方代价)与 p=1p=1(线性代价、对应 WGAN)。与 KL 散度不同,即便两个分布支撑不相交,WpW_p 仍然有限且有意义,并且随着一个分布的平移而平滑变化 —— 正是这一性质使其作为训练损失颇具吸引力。

直觉

耦合 π(x,y)\pi(x,y) 是一张运输时刻表:它规定有多少质量从 xx 搬到 yy。在所有能正确清空 μ\mu、填满 ν\nu 的时刻表中,最优传输选取代价最小者。KL 问的是”一个分布在另一个分布下有多意外”,而 Wasserstein 问的是”质量在物理上需要搬多远”。

深入

加入一项相对熵罚即得熵正则最优传输(Sinkhorn)。以参考耦合 κ\kappa 为基准、温度 γ>0\gamma>0 为权重:

minπΠ(μ,ν)cdπ  +  γKL(πκ).\min_{\pi\in\Pi(\mu,\nu)} \int c\,d\pi \;+\; \gamma\,\mathrm{KL}(\pi\Vert\kappa).

κ=μν\kappa=\mu\otimes\nu(独立乘积测度)即得标准 Sinkhorn 罚项。此时最优耦合取吉布斯–玻尔兹曼形式 π(x,y)κ(x,y)ec(x,y)/γ\pi^\star(x,y)\propto\kappa(x,y)\,e^{-c(x,y)/\gamma},处处严格为正。熵罚阻止传输方案塌缩到稀疏、近乎确定性的支撑上,因而所得传输是平滑的,Sinkhorn 迭代则通过交替缩放高效求解。完整推导与 Cryo-ET 中的对应见 熵正则最优传输

直觉:搬土

把概率质量想成一堆堆的土。μ\mu 是当前各处土堆的高度,ν\nu 是希望达到的目标轮廓。搬一铲土走一段距离要花力气,代价随距离增长;总代价就是所有土在搬运中走过的路程之和(按平方距离计权)。最优传输问的就是:怎么搬,总力气最省。“Earth-mover’s distance”(推土机距离)正由此得名。

一维平方代价下答案出奇简单:把两边的土堆各自从小到大排好,再按名次一一对应。考虑四份等重的土,源点位于 {0,1,4,5}\{0,1,4,5\},目标位于 {2,3,6,7}\{2,3,6,7\}。排序后按序配对得 0 ⁣ ⁣2,  1 ⁣ ⁣3,  4 ⁣ ⁣6,  5 ⁣ ⁣70\!\to\!2,\;1\!\to\!3,\;4\!\to\!6,\;5\!\to\!7,每份各走 22,平均平方代价

14(22+22+22+22)=4.\tfrac14\big(2^2+2^2+2^2+2^2\big)=4 .

任何”交叉”的方案都更糟:若把 0 ⁣ ⁣30\!\to\!31 ⁣ ⁣21\!\to\!2 交叉配对,这两份的代价变成 32+12=103^2+1^2=10,已超过原先的 22+22=82^2+2^2=8。单调(不交叉)耦合之所以最优,根源在于平方代价的次模性 —— 交叉总能通过”解开”而降低代价。这就是为何一维情形有闭式解,也是高维直觉的起点。需要强调的是这条捷径只在一维成立:到了高维就再没有”排序”可言,必须真的去解上面那个线性规划,或转而用下文的熵正则 / 切片近似。

计算的代价:为什么一维便宜、高维昂贵

把上面三种形式按”求解难度”排一排,恰好画出整条最优传输工具链的动机。一维有闭式解:排序后按名次配对,代价由排序主导,O(nlogn)O(n\log n)。但精确的 Kantorovich 线性规划在 nn 个源点对 nn 个目标点时,未知数有 n2n^2 个、约束 2n2n 条,最优解虽稀疏(仅约 2n12n-1 个非零),网络单纯形 / 拍卖算法的代价仍达约 O(n3logn)O(n^3\log n) —— 当 nn 是一整个数据集时根本跑不动。这正是后两条出路存在的理由:熵正则把线性规划软化成强凸问题,用 Sinkhorn 迭代以几次矩阵–向量乘法逼近解;而本页末尾的切片思路则干脆把昂贵的高维传输换成大量廉价的一维传输。深度学习里几乎从不直接解精确 Kantorovich 线性规划,用的都是这两类可微的廉价近似。

与 KL 的对比:支撑不相交时为何仍可用

设两个分布支撑完全不相交:μ\mux=0x=0 处的点质量,ν\nux=δx=\delta 处的点质量。此时 KL(μν)=\mathrm{KL}(\mu\Vert\nu)=\infty,无论 δ\delta 多小都不变 —— KL 看不见”差一点”还是”差很多”,因而梯度无处可用(详见 熵与 KL 散度)。Wasserstein 则给出 W2(μ,ν)=δW_2(\mu,\nu)=\delta:随两堆土的间距平滑变化,处处可微。

这一差别在训练早期尤为致命。生成分布 qxq_x 起初与数据分布 pyp_y 几乎没有重叠 —— 正是支撑不相交的情形。基于 KL / 密度比的对抗损失(如判别器)此时要么饱和、要么给出爆炸或消失的梯度,训练不稳、易模式坍缩;而最优传输代价始终随两分布的几何接近度平滑下降,提供始终有效的下降方向。这正是 CryoGEN-II 用 OT 损失取代 CryoGEN-I 式对抗 min–max 后训练更稳定的根本原因。

Wasserstein 空间的几何

正因为 W2W_2 随质量移动而平滑变化,它把分布的空间变成了一种几何 —— 我们可以沿最短路径在 μ\muν\nu 之间插值,而非把一个渐隐成另一个。朴素的线性(混合)插值 (1t)μ+tν(1-t)\,\mu + t\,\nu 让两堆土原地不动、只重新计权:μ\mu 的质量在原处淡出、ν\nu 的质量在别处淡入,于是这条路径要穿过一团双峰的模糊,两端哪个都不像。McCann 的位移插值则让质量沿直线搬运。设 TT 为从 μ\muν\nu 的最优 Monge 映射,用「已走完比例 tt」的部分映射把 μ\mu 推前:

μt=((1t)id+tT)#μ.\mu_t = \big((1-t)\,\mathrm{id} + t\,T\big)_{\#}\,\mu .

逐符号读:id\mathrm{id} 是恒等映射(xxx\mapsto x,即”留在原地”),TT 是”走到终点”,二者的凸组合 (1t)id+tT(1-t)\,\mathrm{id}+t\,T 就是”走了 tt 比例的路”这个部分映射;下标 #{}_\# 表示用它把 μ\mu 推前,得到中间时刻 t[0,1]t\in[0,1] 的分布 μt\mu_t,两端 μ0=μ\mu_0=\muμ1=ν\mu_1=\nu。每份质量以匀速从 xx 滑向 T(x)T(x),于是单座土堆始终是单座土堆 —— 它平移着穿过,而非先溶解再重组。这条路径正是 W2W_2测地线:两个分布之间最短的曲线,其长度以恒定速率累积,W2(μ,μt)=tW2(μ,ν)W_2(\mu,\mu_t) = t\,W_2(\mu,\nu)

源 μ目标 ν插值 μₜ

拖动 t 从 0 到 1。线性混合是 (1−t)μ + tν:质量在左侧淡出、在右侧重现,中途穿过双峰模糊,两端哪个都不像。位移插值则让单座土堆的均值从 μ 滑到 ν,始终单峰 —— 这正是 Wasserstein 空间里的测地线(最短路径),把「对画面取平均」换成了「对位置取平均」。

直觉

线性插值是把质量瞬移:它把一堆调暗、把另一堆调亮,于是中点那一帧同时出现两座半高的土堆。位移插值则是让质量走过去:中点那一帧,单座土堆恰好停在源与目标的正中间。前者对画面取平均,后者对位置取平均。

在 Cryo-ET 中:把谁搬到谁

在 Cryo-ET 里,CryoGEN-II 要搬的两个分布是具体的。目标是真实观测分布 pyp_y(大量真实拍到的断层图)。是模型生成的干净重构分布 qxq_x —— 网络对各观测输出的重构 xx 在重构空间 X\mathcal{X} 中的聚合分布。两者活在不同空间,无法直接比 xxyy;代价先把重构用与 CryoGEN 同一个缺失楔形退化算子 TM\mathcal{T}_M 打回观测空间,再取平方距离

c(y,TM(x))=yTM(x)2,c\big(y,\mathcal{T}_M(x)\big)=\big\lVert y-\mathcal{T}_M(x)\big\rVert^2 ,

于是 CryoGEN-II 最小化的传输代价为

Wc(py,qx)=infπΠ(py,qx)E(y,x)πyTM(x)2.\mathcal{W}_c(p_y,q_x)=\inf_{\pi\in\Pi(p_y,q_x)}\mathbb{E}_{(y,x)\sim\pi}\big\lVert y-\mathcal{T}_M(x)\big\rVert^2 .

逐符号读:TM\mathcal{T}_M 是缺失楔形退化算子(把一个干净体 xx 变成它在显微镜下本该呈现的、缺了一楔形傅里叶信息的样子);yTM(x)2\lVert y-\mathcal{T}_M(x)\rVert^2 度量”退化后的生成体”与”真实观测”差多远;πΠ(py,qx)\pi\in\Pi(p_y,q_x) 取遍以真实分布 pyp_y 与生成分布 qxq_x 为边缘的耦合;inf\inf 与期望合起来就是这两整个分布之间的最优传输代价。最小化它,要求”生成体退化后”的整体分布与”真实拍到”的整体分布在聚合意义上对齐 —— 即聚合重构 q(xy)p(y)dy\int q(x\mid y)\,p(y)\,dy 贴近干净结构的先验 p(x)p(x)。全程不需要任何真值 xx:由于 p(x)p(x) 不可得,监督完全来自真实观测 pyp_y(即 CryoGEN 的 PYP_Y proxy)。它只是一次确定性优化,每个观测仍只得到一个重构 —— 把它升级为一族重构是 CryoWGEN 的事。

直觉

同一退化算子 TM\mathcal{T}_M,两种用法。CryoGEN-I逐张 MAP:对每个 yy 单独求一个使 TM(x)y2\lVert\mathcal{T}_M(x)-y\rVert^2 加能量先验最小的 xx^*,答案逐图独立。CryoGEN-II 做全局分布匹配:不再逐张对齐,而是要求重构的聚合分布 qxq_xTM\mathcal{T}_M 退化后整体搬到 pyp_y 上,代价由上式的 Wc\mathcal{W}_c 度量。前者保真到每一张图,后者保真到整个数据集的统计 —— 后者更不易产生与真实结构分布相悖的伪影。

未正则的最优传输是 CryoGEN-II 全局分布匹配的基础。在此代价上再加一项熵罚,则进一步给出 CryoWGEN / EVIA 的玻尔兹曼后验,把单一重构升级为一族重构。

从距离到自编码器

把 Wasserstein 距离从「一个待计算的数」变成「一个可训练的目标」,正是 Wasserstein 自编码器(WAE)的思路:用编码器把数据压进隐空间、用解码器还原,再要求隐空间里的聚合后验与先验之间的传输代价最小(Tolstikhin 等, 2018)。这里的”聚合后验”指把所有数据点经编码器投到隐空间后的整体分布;要求它整体对齐到先验,而非像 VAE 那样把每个点的后验逐一拉向先验 —— 这个从”逐点”到”聚合”的差别,正是 WAE 相对 VAE 的关键转变,也与上一节 CryoGEN-II 的”聚合分布匹配”是同一思想。这条「最优传输 → 自编码器」的线,正是 WAEEVIA,乃至 CryoGEN-IICryoWGEN 共享的骨架。它的熵正则版本(见 Sinkhorn)对应 Sinkhorn 自编码器(Patrini 等, 2020);在高维下,若用随机投影把昂贵的高维传输化为大量廉价的一维传输,则得到切片 Wasserstein 自编码器(Kolouri 等, 2019)。

最后这一点,用的正是本页开头那张演示所展示的闭式解。把两个分布投影到某个随机方向 θ\theta 上 —— θ#μ\theta_{\#}\muθ#ν\theta_{\#}\nu 此刻都成了一维 —— 再对所有方向把所得的一维距离取平均,便是切片 Wasserstein 距离:

SW22(μ,ν)=Eθ[W22(θ#μ,  θ#ν)].\mathrm{SW}_2^2(\mu,\nu) = \mathbb{E}_{\theta}\big[\,W_2^2(\theta_{\#}\mu,\;\theta_{\#}\nu)\,\big].

逐符号读:θ\theta 是单位球面上一个随机方向,θ#μ\theta_\#\mu 是把 μ\mu 沿 θ\theta 投影(即把每个高维点 xx 换成标量 θ,x\langle\theta,x\rangle)所得的一维分布;W22(θ#μ,θ#ν)W_2^2(\theta_\#\mu,\theta_\#\nu) 是这对一维分布的(平方)Wasserstein 距离;外层期望 Eθ\mathbb{E}_\theta 对所有方向取平均。每一个切片都正是本页顶端那个一维问题 —— 把投影后的源样本与目标样本各自排序、按名次配对、再把平方间距求和 —— 代价为 O(nlogn)O(n\log n),完全由排序主导,既无需线性规划,也无需 Sinkhorn 迭代。对少数几个随机 θ\theta 取平均即得一个廉价且可微的估计;又因每个切片是度量、度量的平均仍是度量,SW2\mathrm{SW}_2 自身便是一个真正的距离(Bonneel 等, 2015)。

Wasserstein 距离支撑着 Wasserstein GAN 以及 Wasserstein 自编码器,并为 CryoWGENCryoGEN 所用的分布匹配目标提供基础。

延伸阅读

← 最优传输