引言

最优传输理论在机器学习领域受到了越来越多的关注，特别是 Wasserstein-GAN [2] 的出现，引起了大家对最优传输的兴趣。有很多近期的工作将最优传输应用在许多不同的任务中，比如小样本学习 [3]，自监督学习 [4]，信息检索 [5]，目标检测 [6]，多标签分类 [7] 等等。但是目前关于最优传输的中文资料较少，特别是关于最优传输的 Sinkhorn 迭代解法及其 subgradient 的求法，几乎找不到相关的中文材料。

本文首先简要介绍离散情况下的最优传输问题，然后介绍文献 [1] 所提出的最优传输问题的快速解法以及对应的对应梯度的计算。为了理解的方便，本文只介绍离散情况下的最优传输问题。值得注意的是，有很多关于最优传输的资料会用到不同的关键词，例如 optimal transport / wasserstein distance / earth mover's distance / sinkhorn distance 等。离散条件下 optimal transport / wasserstein distance / earth mover's distance 可以认为是等价的，Sinkhorn iteration 是一种最优传输的快速迭代解法，后文中会介绍到。本文涉及到的数学符号均沿用文献 [1]，并补充了详细的说明。

推土机距离：earth mover's distance

Earth mover's distance

假设地面上有 $d_1$ 个土堆，第 $i$ 个土堆有 $r_i$ 数量的土；同时有 $d_2$ 个坑，第 $j$ 个坑可以容纳的 $c_j$ 数量的土。我们用 $$ \begin{split} r\in\mathbb{R}^{d1} \\ c\in\mathbb{R}^{d2} \end{split} $$ 来分别表示这土堆和坑，并假设 $\sum_i r_i = \sum_j c_j$，也就是 所有的 “土” 刚好够填满所有的 “坑”。定义 $\mathbf{M}\in\mathbb{R}^{d_1\times d_2}_+$ 为距离矩阵，其中 $m_{ij}$ 表示 “从第 $i$ 个土堆搬运一份土到第 $j$ 个坑的成本 (cost)”。现在要将土从 $r$ 搬到 $c$，定义 $\mathbf{P}\in\mathbb{R}^{d_1\times d_2}_+$ 为从 $r$ 到 $c$ 的 "传输方案" (transportation paln)，其中 $p_{i,j}$ 表示从 $r_i$ 搬运到 $c_j$ 的土的数量。最优传输目的是求得一个最优的传输矩阵 $P^*$，使得总体的传输消耗 (cost) 最小。

很显然，传输矩阵 $P$ 的每一行/列加起来应该等于 $r$/$c$，这也是为什么用字母 r(ow)/c(olumn) 来表示的原因。写成公式，$P$ 满足约束： $$ \begin{equation} \begin{cases} \sum_i p_{i,j} = c \\ \sum_j p_{i,j} = r \end{cases}\label{eq:t-constrain} \end{equation} $$ 即所有满足条件的 transportation plan $P$ 的集合为： $$ U(r, c) := \big\{ P\in \mathbb{R}_+^{d_1\times d_2} | P \cdot 1_{d2} = r, P^T \cdot 1_{d1} = c \big\} \label{eq:U} $$ 其中 $1_d$ 表示 $d$ 维的全1向量。值得一提的是，$U(r, c)$ 构成高维空间的多面体 (polytope) ，因为式$\ref{eq:t-constrain}$ 所有约束都是线性的，平面上多个线性约束条件围成的区域是多边形，拓展到高维空间后线性约束围成的集合就是多面体（polytope）。这里用集合 $\mathbf{U}$ 来表示所有满足条件的传输，在优化中集合 $\mathbf{U}$ 又被称为“可行域” (feasible set)。

最优传输 $P^*$ 定义为“使得总体运输成本 $\langle P, M\rangle_F$ 最小的运输方案”： $$ \begin{equation} P^* = \inf_{P\in \mathbf{U(r, c)}} \langle P, M\rangle_F \label{eq:ot-origin} \end{equation} $$ 其中 $\langle X, Y\rangle_F = \sum_i X_iY_i$ 为 Frobenius inner product 。使得式$\ref{eq:ot-origin}$ 最小的传输方案 $P^*$ 称 为最优传输方案。

简而言之，最优传输就是要在可行域 $U(r, c)$ 中找到找到使得式$\ref{eq:ot-origin}$ 最小的传输方案 $P^*$。很显然这是一个线性规划问题，因为不论是式$\ref{eq:t-constrain}$中的约束条件还是式$\ref{eq:ot-origin}$中的目标函数 都是线性的。

EMD 度量概率分布的距离

如果 $r$ 和 $c$ 满足 $\sum_i r_i = 1, \ \sum_j c_j = 1 $ 且 $\forall i, j \ \ \ \ r_i \ge 0, c_j \ge 0$，那么 $r$ 和 $c$ 可以看作两个概率分布， $U(r, c)$ 也可以看作是边际分布分别为 $r$ 和 $c$ 的 “联合概率分布”。 $r$ 和 $c$ 之间的最优传输距离可视作概率分布之间的差异。在很多应用中都会将输入 $r$ 和 $c$ 归一化成概率（比如分类中的 softmax），最后计算两个概率的最优传输距离。因此最优传输也可用作度量概率分布之间的距离。

Why optimal transport？

那么我们为什么要计算 optimal transport 呢？如果为了 measure 概率分布之间的距离，有很多现成的 measurements 可以用，比如非常简单的 KL 散度: $$ \text{KL}(p, q) = \sum_i p_i \cdot\log\frac{p_i}{q_i}. $$ 除了 “无法处理两个分布的支撑集不相交的情况”以及“不满足对称性” 等原因之外，一个重要的原因就是这种逐点计算的度量 没有考虑分布内的结构信息。所谓的结构信息，就是分布内的联系。例如在 KL 散度中，$p_i, i=0,1,...$ 彼此之间都是独立计算最后加起来的，而大部分情况下它们并不是独立的。

就以我们常见的分类任务为例，分类任务通常用交叉熵损失来度量模型预测和样本标签之间的距离，交叉熵损失实际上就是在计算 onehot 化的标签和模型预测之间的 KL 散度。这种逐点计算的损失函数（不论是交叉熵还是 L2）都无法考虑分布内不同事件的相关性。例如将“汽车”误分类成“卡车”显然没有把“汽车”误分类成“斑马”严重。但是用 KL 散度来度量的话，这两种错误的损失是一样的。

假设$r \in \mathbb{R}^{d1}_+, c \in \mathbb{R}^{d2}_+$是两个离散概率分布，满足 $r^T\mathbf{1}_{d1} = 1, c^T\mathbf{1}_{d2} = 1$，其中$\mathbf{1}_{d1},\mathbf{1}_{d1}$ 表示全1列向量，我们可以用最优传输来计算两者之间的距离。而概率分布内的结构信息可以通过距离矩阵 $\mathbf{M}$ “嵌入” 到距离度量中。还是以分类为例，我们可以让 $m_{汽车,卡车}$ 远小于 $m_{汽车,斑马}$。

熵正则的最优传输问题

熵正则

在文献[1]中，Cuturi 提出了一种快速的最优传输求解方法。该方法首先引入熵正则，使得原问题的可行域更加平滑，然后将最优传输问题转为 matrix permutation 问题，最后用 sinkhorn 迭代算法求原问题的近似解。

联合分布 $P$ 的熵 (entropy) 为: $$ \begin{equation} H(P) = -\sum_i P_i\log(P_i) \end{equation} $$ 原问题式$\ref{eq:ot-origin}$ 增加熵正则约束之后，优化目标变为： $$ \begin{equation} D_{M, \lambda}(r,c) = \inf_{P\in \mathbf{U}} \langle P, M\rangle - \frac{1}{\lambda}H(P) \label{eq:ot-entropy} \end{equation} $$ 式$\ref{eq:ot-entropy}$ 称为“熵正则的最优传输” (entropy-regularized optimal transport)，对应的最优解表示为 $P^*_{\lambda}$。从概率和信息论的角度，均匀分布的熵最大，因此熵正则会让最优传输矩阵 $P^*_{\lambda}$ 更趋近于均匀分布。

从优化的角度，式$\ref{eq:ot-origin}$ 中的原问题的解一定存在于多边形 $U(r, c)$ 的某一顶点处（这么说似乎不严谨，但是当 $r, c$ 的维度很高时，最优解一定在顶点处取得），因此 $P^*$ 是一个稀疏矩阵，绝大部分元素都是0。增加熵正则之后，相当于将原来的可行域 $U(r, c)$ 向内收缩成光滑的 $U_{\lambda}(r, c)$，对应的最优解 $P^*_{\lambda}$ 不再是稀疏矩阵。

可行域 $U(r, c)$ 和 $U_{\lambda}(r, c)$ 包含所有满足约束的点， $U(r, c)$ 是一个高维空间中的多边形 (polytope)，增加熵正则之后的可行域 $U_{\lambda}(r, c)$ 是一个边界光滑的区域。蓝色虚线表示目标函数的等高线，距离矩阵 $M$ 决定最优解所在的方向。本图参考了文献[1] 的 Fig.1。

当$\lambda \rightarrow \infty$ 时，$U_{\lambda}(r, c)$趋向于 $U_{\lambda}(r, c)$，最优解 $P^*_{\lambda}$ 也趋向于 $P^*$；当$\lambda \rightarrow 0$ 时，$U_{\lambda}(r, c)$ 收缩成点 $rc^T$，最优解为 $P^*_{\lambda}=rc^T$。

为什么要增加熵正则

参考文献 [1] 中作者列出了两个使用熵正则的原因：

式$\ref{eq:ot-origin}$ 中的原线性规划问题的解一定是在可行域 $U(r, c)$ 的某个顶点 (vertex) 上的，因此得到的解 $P^*$ 是一个稀疏矩阵。稀疏的解 $P^*$ 会使得最终的传输方案十分不均衡，使用熵正则可以让传输矩阵更加均衡，这部分可以参考 [1] 的第三章相关内容。
更重要的是，增加了熵正则之后，原问题可以使用 sinkhorn 算法得到近似解，而且计算开销大大降低。关于这部分，可以参考原文第四章。

总的来说，增加熵正则之后的最优传输问题得到的解更加光滑和均衡，同时可以使用 sinkhorn 迭代算法快速求解。

基于 Sinkhorn 迭代的快速解法。

Sinkhorn 迭代求解最优传输

假设有向量 $u \in \mathbb{R}^{d1}, v\in \mathbb{R}^{d2}$，令 $u, v$ 的初始值为 $$ \begin{split} u &= \mathbf{1}_{d1} / d1 \\ v &= \mathbf{1}_{d2} / d2 \end{split} $$ 然后用 sinkhorn 迭代算法求解： $$ \begin{equation} \begin{split} u &\leftarrow r / (K\cdot v) \\ v &\leftarrow c / (K^T\cdot u) \end{split}\label{eq:sinkhorn} \end{equation} $$ 上式中 $K:=e^{-\lambda M}$ (参考[1] 中 Lemma2)。迭代收敛后最优传输矩阵 $P^*$ 和对应的最小传输距离 $C^*$ 可以由以下公式给出： $$ \begin{equation}\begin{split} P^* &= \text{diag}(u)K \text{diag}(v) \\ C^* &= \langle P^*, M\rangle_F \end{split}\label{eq:distance} \end{equation} $$ $C^*$ 对 $r, c$ 的导数分别是: $$ \begin{equation} \begin{split} \frac{\partial C^*}{\partial r} &= \log(u) / \lambda \\ \frac{\partial C^*}{\partial c} &= \log(v) / \lambda \end{split}\label{eq:gradient} \end{equation} $$ 有了公式$\ref{eq:distance}$ 和公式$\ref{eq:gradient}$，最优传输距离就可以作为损失函数来训练神经网络。

PyTorch 实现

这里 (vlkit.optimal_transport.sinkhorn) 有一个 sinkhorn 的 PyTorch 实现，我们可以计算并可视化两个分布之间的最优传输。

以高斯分布为例，首先生成两个 1d 高斯分布作为 source 和 target distribution：

import torch
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import gridspec
from vlkit.optimal_transport import sinkhorn

# generate two gaussians as the source and target
def gaussian(mean=0, std=10, n=100):
    d = (-(torch.arange(n) - mean)**2 / (2 * std**2)).exp()
    d /= d.sum()
    return d

n = 20
d1 = gaussian(mean=12, std=2, n=n)
d2 = gaussian(mean=6, std=4, n=n)

dist = (torch.arange(n).view(1, n) - torch.arange(n).view(n, 1)).abs().float()
dist /= dist.max()

# visualize distr
fig, axes = plt.subplots(1, 2, figsize=(9, 3))
axes[0].bar(torch.arange(n), d1)
axes[0].set_title('Source distribution')
axes[1].bar(torch.arange(n), d2)
axes[1].set_title('Target distribution')
plt.tight_layout()

然后通过 `sinkhorn` 迭代计算最优传输，并可视化结果。

T, u, v = sinkhorn(r=d1.unsqueeze(dim=0), c=d2.unsqueeze(dim=0), reg=1e-2, M=dist.unsqueeze(dim=0))
plt.figure(figsize=(10, 10))
gs = gridspec.GridSpec(3, 3)

ax1 = plt.subplot(gs[0, 1:3])
plt.bar(torch.arange(n), d2, label='Target distribution')

ax2 = plt.subplot(gs[1:, 0])
ax2.barh(torch.arange(n), d1, label='Source distribution')

plt.gca().invert_xaxis()
plt.gca().invert_yaxis()

plt.subplot(gs[1:3, 1:3], sharex=ax1, sharey=ax2)
plt.imshow(T.squeeze(dim=0))
plt.axis('off')

plt.tight_layout()

补充阅读

这里列出了一些关于 optimal transport 的相关阅读材料供大家参考。

如果本文的内容对你撰写学术论文有帮助，希望能考虑引用：

@misc{zhao2020optimal,
title   = {最优传输与 sinkhorn-knop 迭代},
author  = {Kai Zhao},
year    = 2020,
note    = {\url{http://kaizhao.net/blog/optimal-transport}}
}

References

Cuturi, Marco. "Sinkhorn distances: Lightspeed computation of optimal transport." Advances in neural information processing systems 26 (2013): 2292-2300.
Arjovsky, Martin, Soumith Chintala, and Léon Bottou. "Wasserstein generative adversarial networks." International conference on machine learning. PMLR, 2017.
Zhang, Chi, et al. "DeepEMD: Few-Shot Image Classification With Differentiable Earth Mover's Distance and Structured Classifiers." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2020.
Liu, Songtao, Zeming Li, and Jian Sun. "Self-EMD: Self-Supervised Object Detection without ImageNet." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2021.
Xie, Yujia, et al. "Differentiable top-k operator with optimal transport." arXiv preprint arXiv:2002.06504 (2020).
Ge, Zheng, et al. "OTA: Optimal Transport Assignment for Object Detection." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2021.
Frogner, Charlie, et al. "Learning with a Wasserstein loss." Advances in neural information processing systems (2015).

感谢阅读🤗 本文内容谢绝任何形式的转载，如果您想和朋友分享本文内容，请分享本文链接 kaizhao.net/blog/optimal-transport。如果您发现文中的错误，或者有任何疑问，欢迎在下方留言交流 (留言功能基于 disqus，在中国大陆的读者可能需要一些技术手段才能连接🥲)。