一维最优传输 [optimal-transport-in-one-dimension]

最优传输一直以"搬土堆"作为直觉上的例子:源分布是一堆土,目标分布规定土最终应当怎么摆,我们要选运输方案使总成本最小。

在一般空间里这是个不折不扣的困难优化问题,但在实数轴上,问题就变得简单了,答案可以用一句话说完:

核心结论 实数轴上的位置天然具有从左到右的全序。对于成本 c(x,y)=|x-y|^pp\ge 1),最优运输遵循相同排名的质量互相匹配的原则:离散等权情形下是排序后逐点配对,一般情形下是匹配相同分位数。

这句话乍一看很合理,但其中蕴含丰富的细节可以思考。本文的主线就是把这句话逐步做实,从直觉和分析两个角度理解这个结论。

1. 记号:位置与质量分开写

实数轴上的离散概率分布写成

\alpha=\sum_{i=1}^n a_i\,\delta_{x_i}, \qquad a_i\ge 0, \qquad \sum_{i=1}^n a_i=1.

这里 \delta_{x_i} 是集中在 x_i 处的单位 Dirac 质量,所以 a_i\delta_{x_i} 读作"在位置 x_i 放置质量 a_i"。位置写在 \delta 的下标里,质量写在系数里,两个角色始终分开:位置负责"在哪",质量负责"有多少"。

若有 n 个样本点 x_1,\dots,x_n,其经验分布写成 \alpha=\frac1n\sum_{i=1}^n\delta_{x_i}。这不是"每个点没有质量",而是每个样本点默认携带相同质量 1/n;若某个位置重复出现,质量自动累加:

\frac14(\delta_{0.1}+\delta_{0.2}+\delta_{0.2}+\delta_{0.9}) = \frac14\delta_{0.1}+\frac12\delta_{0.2}+\frac14\delta_{0.9}.

2. 最简单的情形:等权离散

设两个等权经验分布

\alpha=\frac1n\sum_{i=1}^n\delta_{x_i}, \qquad \beta=\frac1n\sum_{i=1}^n\delta_{y_i},

将两组位置分别排序:x_{(1)}\le\cdots\le x_{(n)}y_{(1)}\le\cdots\le y_{(n)}

这里排序排的是位置;等权时每个点质量都是 \frac1n、一样重,质量不参与排序,只在最后作为系数出现。对于成本 c(x,y)=|x-y|^pp\ge1),一个最优方案就是按排名配对 x_{(i)}\to y_{(i)},于是就有

W_p(\alpha,\beta) = \left(\frac1n\sum_{i=1}^n |x_{(i)}-y_{(i)}|^p\right)^{1/p}.

换句话说,一维等权经验分布之间的 Wasserstein 距离,就是两个排序后坐标向量之间的归一化 \ell^p 距离:

W_p(\alpha,\beta) = n^{-1/p}\,\big\|(x_{(1)},\dots,x_{(n)})-(y_{(1)},\dots,y_{(n)})\big\|_p.

 设 x=(4,0,3)y=(6,2,1)。排序后 x_{(\cdot)}=(0,3,4)y_{(\cdot)}=(1,2,6),最优匹配为 0\to13\to24\to6。当 p=2W_2^2 = \frac13(1^2+1^2+2^2)=2,即 W_2=\sqrt2

这个公式好到令人怀疑:怎么排序向量的 p 范数就够了呢?

排序匹配最优的根本原因是:对凸的距离成本函数来说,交叉运输不会比顺序运输更便宜。先看只有两个点的情形。设 x_1\le x_2y_1\le y_2,比较顺序匹配(x_1\to y_1x_2\to y_2)与交叉匹配(x_1\to y_2x_2\to y_1)。对 p\ge1,总有

|x_1-y_1|^p+|x_2-y_2|^p \le |x_1-y_2|^p+|x_2-y_1|^p.

这就是一维 Wasserstein 成本的 Monge 不等式。直觉上很好接受:交叉意味着两条运输路线有一段路被重复走了两遍,把交叉解开,重复的那段就省了。

证明 令 \varphi(t)=|t|^p,当 p\ge1 时它是凸函数。记 a=x_2-x_1\ge0b=y_2-y_1\ge0z=x_1-y_2,则四个差值分别为 x_1-y_2=zx_1-y_1=z+bx_2-y_2=z+ax_2-y_1=z+a+b,要证的不等式化为

\varphi(z+b)+\varphi(z+a) \le \varphi(z)+\varphi(z+a+b).

对固定的 b\ge0,定义增量 D_b(t)=\varphi(t+b)-\varphi(t)。凸函数的斜率随位置增大而不减,故 D_b 非递减;由 z+a\ge zD_b(z+a)\ge D_b(z),移项即得。\blacksquare

有了这个不等式,可以进一步考虑:对于任何一个含逆序配对的方案,都能找出其中一对交叉(四个点)、按上式交换而不增加成本;反复交换直到没有交叉,剩下的正是排序匹配。所以任何不按顺序来的匹配都不可能严格更优。

两点备注,划定这个结论的边界:

  • 成本函数得是凸函数(当 c(x,y)=|x-y|^p 时,条件 p\ge1 不能丢)。0<p<1|t|^p 是凹函数,排序可能不再最优。例如 x_1=0x_2=1y_1=1y_2=2:顺序匹配成本 1^p+1^p=2,交叉匹配成本 2^p+0,当 0<p<12^p<2,交叉反而更便宜。
  • 最优方案未必唯一。 p>1 且位置严格递增、质量结构不退化时,严格凸性通常给出更强的唯一性;p=1 时排序方案仍最优,但可能与别的方案并列最优;存在重复位置或零距离区间时也可能不唯一。例如 p=1、源的两个点都在目标两个点的同一侧时(如 x=(2,3)y=(0,1)),顺序匹配和交叉匹配的成本同为 4,两种方案并列最优。

3. 放宽一步:质量不相等

现在允许两边质量任意:

\alpha=\sum_{i=1}^n a_i\delta_{x_i}, \qquad \beta=\sum_{j=1}^m b_j\delta_{y_j}, \qquad x_1\le\cdots\le x_n,\quad y_1\le\cdots\le y_m.

乍一看之前基于等权离散的排序方法的前提不在了,排序不能继续用了。但是稍微思考,能把问题稍稍转换成我们已经解决的问题:把一个质量大的点想象成许多个质量很小的等权小点叠在同一个位置。 比如 0.6\delta_0 可以看成 6 个质量 0.1 的小点全站在 0 处。两边都这样打散之后,又变回了等权情形,排序匹配照常进行;所谓"一个点拆给多个目标",无非是这些站在同一位置的小点被分头送往了相同或者不同的地方。

落实成算法就是所谓"双指针":两边各派一个指针从最左的格子出发,每次运走 \min(\text{两边当前剩余质量}),谁的格子清零谁前进,直到走完整根尺子。

 设 \alpha=0.6\delta_0+0.4\delta_3\beta=0.3\delta_1+0.7\delta_2。 双指针过程如下:先运 \min(0.6,0.3)=0.3 的质量 0\to1\beta 的第一格清零、指针前进,\alpha 的第一格还剩 0.3;再运 \min(0.3,0.7)=0.3 的质量 0\to2\alpha 的第一格清零、指针前进,\beta 的第二格还剩 0.4;最后运 \min(0.4,0.4)=0.4 的质量 3\to2,两边同时清零,结束。当 p=1 时总成本为 0.3\times1+0.3\times2+0.4\times1=1.3。可以看到,位置 0 处的质量被"拆"给了两个目标,这正是把它视为一堆小点后小点分头出发的结果。

4. 统一的语言:累积分布与分位数函数

前面的排序匹配适合有限点集,思路已经很好了。但对连续分布、混合分布或任意权重,"打散成等权小点再排序"没法直接照搬:连续分布在每个点上的质量都是零,谈不上一个个"小点";点又有无穷多个,也无法逐一列举出来排序。好在思路本身并没有问题——按质量切小、按位置排队——缺的只是一套对所有分布都通用的语言,把这个思路严格写下来。本节的工具就是累积分布函数和它的广义逆。

先给出两个定义:设 \alpha 是实数轴上的概率测度。

累积分布函数(CDF)F_\alpha(x) = \alpha((-\infty,x]),表示 \alpha 位于 x 左侧(含 x)的累计质量。例如 \alpha=0.3\delta_0+0.7\delta_2 对应

F_\alpha(x)= \begin{cases} 0, & x<0,\\ 0.3, & 0\le x<2,\\ 1, & x\ge 2. \end{cases}

这个概念学过概率论的应该很熟悉。它的一大好处是普适:离散、连续、混合类型的分布可能没有密度函数,但都有累积分布函数,天然适合用来统一描述。

分位数函数(广义逆):

Q_{\alpha}(u) \coloneqq F_{\alpha}^{\leftarrow}(u) \coloneqq\inf\left\{ x\in\mathbb{R} \,\middle|\, F_{\alpha}(x)\ge u \right\}, \qquad u\in(0,1).

Q_\alpha(u)\alpha 的第 u 分位数:Q_\alpha(0.5) 是中位数,Q_\alpha(0.25) 是第一四分位数。对上面的例子,Q_\alpha(u)=00<u\le0.3),Q_\alpha(u)=20.3<u<1)。直观读法:输入一个累计概率 u,输出累计概率第一次达到 u 时所在的位置——一个具体的实数。

有了这两个工具,回头看第 2、3 节"切碎质量、按序搬运"的操作,在这套语言里应该怎么表达。

之所以从累积分布函数下手,除了它人人都有之外,更因为它有一个关键性质:单调不减——沿数轴正方向走,累积质量只增不减(若某段上 F_\alpha 完全不变,说明分布在那一段上根本没有质量)。这意味着"位置的左右顺序"和"质量被累积的先后顺序"天然一致:越靠左的质量,越早被计入累积。

这就提示我们换一根轴来切分质量:不再按位置切,而是按累积质量 u。把总质量摊平成区间 (0,1),切成无穷多份微元 \mathrm{d}u;分位数函数 Q_\alpha(u) 恰好回答"第 u 份微元在数轴上的哪个位置"。每份微元大小相同,这正是第 3 节里等权小点的连续极限版本——只不过小点变成了无穷小的质量微元,"排好序"这件事由 F_\alpha 的单调性自动保证。

匹配也更容易:源分布的第 u 份微元位于 Q_\alpha(u),目标分布的第 u 份微元位于 Q_\beta(u),把前者搬到后者即可。关键的读法是:u\in(0,1) 就是"质量的排名"F_\alpha 把位置翻译成排名,Q_\alpha 把排名翻译回位置,这是一对互逆的查询表。

\alpha,\beta\in\mathcal P_p(\mathbb R)(即有有限 p 阶矩)。对 p\ge1

W_p(\alpha,\beta)^p = \int_0^1 |Q_\alpha(u)-Q_\beta(u)|^p\,\mathrm{d}u.

一维最优传输的本质由此完全显形:把源分布的第 u 分位数送到目标分布的第 u 分位数

5. p=1 的额外礼物:CDF 之间的面积

p=1 时,分位数公式还能换一个坐标轴来算——不积"横着的"分位数差,改积"竖着的"CDF 差:

W_1(\alpha,\beta) = \int_{\mathbb R} |F_\alpha(x)-F_\beta(x)|\,\mathrm{d}x.

(两式相等的几何原因:|Q_\alpha-Q_\beta|u 的积分和 |F_\alpha-F_\beta|x 的积分,量的都是两条单调曲线之间同一块区域的面积,只是一个横着切、一个竖着切。)

它有一个很物理的解释:固定位置 x,在数轴上切一刀。F_\alpha(x)-F_\beta(x) 是切口左边源质量与目标质量的净差额;差额不为零,就至少有 |F_\alpha(x)-F_\beta(x)| 的质量必须穿过这个切口。而一份质量从 a 运到 b 时,恰好穿过 a,b 之间的每一个切口,所以对全部切口积分,正好把"质量 × 运输距离"累计起来——这就是 W_1

6. 总结

一维最优传输之所以简单,是因为实数轴上的位置可以完全排序。对凸成本 c(x,y)=|x-y|^pp\ge1),Monge 不等式保证交叉可以被消除,最优耦合必然单调,于是全部理论收束成一句话:排序离散点,或等价地,匹配相同分位数。三个层次的公式分别是:

  • 离散等权:W_p^p = \frac1n\sum_i|x_{(i)}-y_{(i)}|^p
  • 一般一维分布:W_p^p = \int_0^1|Q_\alpha(u)-Q_\beta(u)|^p\,\mathrm{d}u
  • p=1 特例:W_1 = \int_{\mathbb R}|F_\alpha(x)-F_\beta(x)|\,\mathrm{d}x

其实说的都是同样的事情,终于有点懂了,cool!