0. Main Takeaway
本文主要是对Notes on Optimal Transport 的翻译及整理, 同时重构了原文提供的代码, 整理到了一个Jupyter Notebook中并增加了注释, 方便理解 (见我的github).
本文看完可以了解到
- Sec 1.: 什么是最优输运; Wasserstein metric; Sinkhorn distance; Sinkhorn-Knopp算法
- Sec 2.: 最优输运的实际运用(含代码), 包括训练集和测试集分布不一致时如何插值; 风格迁移.
- Sec 3.: 提供了整理后的代码.
1. 从一次聚会讲起
It’s Party Time! 假设实验室要举办一场聚会, 一共有n=8个人, 有m=5种小吃. 我们用$\mathbf{r} = (3,3,3,4,2,2,2,1)^{\top}\in \mathbb{R}^n$来代表每个人的食量, 用$\mathbf{c}=(4,2,6,4,4)^{\top}\in \mathbb{R}^m$来代表每个小吃的份数, 如下图:
同时, 为了让每个人都吃到自己最喜欢的小吃, 我们调查了每个人对每种小吃的喜爱程度, 如下表, $+2$表示最喜欢, $-2$表示最讨厌.
| merveilleux | eclair | chocolate mousse | bavarois | carrot cake | |
|---|---|---|---|---|---|
| Bernard | 2 | 2 | 1 | 0 | 0 |
| Jan | 0 | -2 | -2 | -2 | 2 |
| Willem | 1 | 2 | 2 | 2 | -1 |
| Hilde | 2 | 1 | 0 | 1 | -1 |
| Steffie | 0.5 | 2 | 2 | 1 | 0 |
| Marlies | 0 | 1 | 1 | 1 | -1 |
| Tim | -2 | 2 | 2 | 1 | 1 |
| Wouter | 2 | 1 | 2 | 1 | -1 |
接下来问题来了, 我们要如何分配我们的小吃, 从而保证在所有人都吃饱, 且所有小吃都被分完的前提下, 每个人都尽可能吃到自己最喜欢吃的小吃呢?
1.1. 最优输运 (Optimal Transport)
回顾前文, 我们用$\mathbf{r} = (3,3,3,4,2,2,2,1)^{\top}\in \mathbb{R}^n$来代表每个人的食量, 用$\mathbf{c}=(4,2,6,4,4)^{\top}\in \mathbb{R}^m$来代表每个小吃的份数(我们假设份数是可被连续分割的). 则我们的目标是求解一个维度为$n\times m$的矩阵$\mathbf{P}$, 其中$\mathbf{P}_{ij}$代表第$i$个人应该被分到几份第$j$个小吃.
我们用$U(\mathbf{r}, \mathbf{c})$用来表示所有可能的解空间 (即$\mathbf{P}$可能的取值空间):
$$ U(\mathbf{r}, \mathbf{c}) = \{ \mathbf{P} \in \mathbb{R}^{n\times m}_{\geq 0}\mid \mathbf{P}\mathbf{1}_m = \mathbf{r}, \mathbf{P}^{\top} \mathbf{1}_n =\mathbf{c}\}. $$
出于惯例, 我们把每个人对小吃的喜好矩阵取反后得到厌恶/代价矩阵 (cost matrix), 并记为$\mathbf{M}\in \mathbb{R}^{n\times m}$, 则我们的优化任务为:
$$ d_\mathbf{M}(\mathbf{r}, \mathbf{c}) = \min\limits_{\mathbf{P}\in U(\mathbf{r}, \mathbf{c})}\sum\limits_{ij}\mathbf{P}_{ij}\mathbf{M}_{ij} $$
这便被称之为$\mathbf{r}$和$\mathbf{c}$之间的最优输运问题. 注意目标函数和限制函数都是线性的, 所以这个最优输运问题可以被线性规划有效求解 (例如单纯形法, 对偶单纯形法等).其中该问题的最优解$d_{\mathbf{M}}(\mathbf{r}, \mathbf{c})$也被称之为Wasserstein distance, 当$\mathbf{r}$, $\mathbf{c}$都是概率分布时, Wasserstein distance便是这两个分布之间的距离的一种度量 (常见的度量还有KL-Divergence, TV-Divergence等).
Optimal Transport有时候也被称为earth mover distance, 因为当你把$\mathbf{r}$, $\mathbf{c}$想象成两座小沙丘时, $\mathbf{P}_{ij}^{*}$其实描述了你要怎么搬运沙子, 从而以最小的代价让$\mathbf{r}$变成$\mathbf{c}$. 举例
| 0.2 | 0.6 | 0.2 | |
|---|---|---|---|
| 0.6 | 0 | 0.6 | 0 |
| 0.4 | 0.2 | 0 | 0.2 |
假设$\mathbf{r} = (0.6, 0.4)^{\top}$, $\mathbf{c} = (0.2, 0.6, 0.2)^{\top}$, 我们先不考虑代价矩阵$\mathbf{M}$. 上表就描述了一种搬运方式: 即把$\mathbf{r}$的第一堆沙子全搬到$\mathbf{c}$的第二堆, 把$\mathbf{r}$的第二堆沙子平均搬到$\mathbf{c}$的第一堆和第三堆.
Sinkhorn distance和Sinkhorn-Knopp算法
对于一般的最优输运问题, 我们可以使用线性规划求解, 但是当$n$和$m$较大时, 求解会十分耗时, 于是我们不妨考虑一个"更难"的加了正则项问题:
$$ d^{\lambda}_\mathbf{M}(\mathbf{r}, \mathbf{c}) = \min\limits_{\mathbf{P}\in U(\mathbf{r}, \mathbf{c})}\sum\limits_{ij}\mathbf{P}_{ij}\mathbf{M}_{ij} - \frac{1}{\lambda} h(\mathbf{P}), \lambda > 0. $$
这里$h(\mathbf{P})$是矩阵$\mathbf{P}$的信息熵矩阵 (可以先normalize $\mathbf{P}$再求):
$$ h(\mathbf{P}) = -\sum\limits_{ij} \mathbf{P}_{ij}\log \mathbf{P}_{ij} $$
这里引入的正则项$h(\mathbf{P})$越大, 表示分配矩阵$\mathbf{P}$越均匀: 即每个人会越均匀的分配到每个小吃, 而$\lambda$作为超参, 用来平衡求解最优输运和让分配矩阵$\mathbf{P}$越均匀.
若我们记$\mathbf{P}^{*} _{\lambda}$为上述问题对应的最优解, 则$\sum\limits_{ij}(\mathbf{P}^{*}_{\lambda})_{ij}\mathbf{M}_{ij}$被称为Sinkhorn distance. (注意: 正则项$h(\mathbf{P})$的值是不被计入Sinkhorn distance的.) 在一些问题中, 使用Sinkhorn distance得到的结果会比Wasserstein distance要更好, 因为引入$h(\mathbf{P})$相当于引入了一个先验: 即在不考虑代价矩阵的情况下, 我们希望每个人分到的越均匀越好.
虽然一般来说, 加入正则项后会让问题的求解更复杂, 但是对于Sinkhorn distance则是一个例外, 事实上, 我们有一个非常高效且简洁的算法用来求解Sinkhorn distance, 即Sinkhorn-Knopp算法.
该算法的正确性我们这里不做展开, 感兴趣的同学可以自行查阅《the sinkhorn-knopp algorithm: convergence and applications》.
Sinkhorn-Knopp的算法基于一个发现, $d^{\lambda}_\mathbf{M}(\mathbf{r}, \mathbf{c})$的最优解$\mathbf{P}^{*} _{\lambda}$一定满足如下形式:
$$ \left(\mathbf{P}_{\lambda}^{\star}\right)_{i j}=\alpha_{i} \beta_{j} e^{-\lambda \mathbf{M}_{i j}} $$
其中$\alpha_i$, $\beta_j$是待求解的常数, 从而保证$\mathbf{P}_{\lambda}^{\star}$横向求和是$\mathbf{r}$, 纵向求和是$\mathbf{c}$. 因此我们采用一种迭代的思想去求解, 即
Input: $\mathbf{M}, \mathbf{r}, \mathbf{c}, \lambda$
Initialize: $(\mathbf{P}_\lambda)_{ij} = e^{-\lambda \mathbf{M}_{ij}}$
Loop:
- Scale每一行, 从而横向求和是$\mathbf{r}$.
- Scale每一列, 从而纵向求和是$\mathbf{c}$.
Until convergence.
下面给出对应的Python实现:
def compute\_optimal\_transport(M, r, c, lam, epsilon=1e-5):
"""
Computes the optimal transport matrix and Slinkhorn distance using the
Sinkhorn-Knopp algorithm
Inputs:
- M : cost matrix (n x m)
- r : vector of marginals (n, )
- c : vector of marginals (m, )
- lam : strength of the entropic regularization
- epsilon : convergence parameter
Output:
- P : optimal transport matrix (n x m)
- dist : Sinkhorn distance
"""
n, m = M.shape
P = np.exp(- lam * M)
# Avoiding poor math condition
P /= P.sum()
u = np.zeros(n)
# Normalize this matrix so that P.sum(1) == r, P.sum(0) == c
while np.max(np.abs(u - P.sum(1))) > epsilon:
# Shape (n, )
u = P.sum(1)
P *= (r / u).reshape((-1, 1))
P *= (c / P.sum(0)).reshape((1, -1))
return P, np.sum(P * M)
借助Sinkhorn-Knopp算法, 我们可以求解上述加入了entropy正则项的派对问题, 结果如下:
可以注意到, 当$\lambda$的值越小, 我们加入的惩罚项$\frac{1}{\lambda} h(\mathbf{P})$会使得最终解更偏向homogeneous的分布, 即每个人分到的小吃会更加均匀, 这也意味着每个人可能会被更多地分到自己不喜欢吃的小吃.
事实上, 最优输运有着一个很优雅的几何解释:
上图中, $U(\mathbf{r}, \mathbf{c})$包含了所有可行解, 而代价矩阵$\mathbf{M}$则给定了一个方向, 从而指明了什么样的解$\mathbf{P}$是更优的. 在没有正则项$\frac{1}{\lambda} h(\mathbf{P})$的情况下, 最优输运问题的解会在$U(\mathbf{r}, \mathbf{c})$的一个极点取到 (别忘了最优输运本质上是一个线性规划问题). 而在加入正则项$\frac{1}{\lambda} h(\mathbf{P})$后, 我们希望得到的是一个更smooth的解, 例如图中红色椭圆与黄线相切的位置. 这也就解释了为什么加入正则项后求解会更加容易, 因为我们可以不用再管$U(\mathbf{r}, \mathbf{c})$复杂的边界了.
在极端情况下, 当$\lambda \rightarrow \infty$, $\mathbf{P}_\lambda^{*}$会无限接近$\mathbf{P}^{*}$ (直到出现数值不稳定性). 而当$\lambda \rightarrow 0$的时候, 我们只会考虑最大化entropy项$h(\mathbf{P})$, 从而$\mathbf{P}^{*}_{0} =\mathbf{r}^{\top}\mathbf{c}$. (这是满足约束条件的可行解中entropy最大的解.)
2. 最优输运的应用实例(含代码)
在了解了最优输运的原理后, 我们给出几个实例: 包括分布匹配, 风格迁移. 代码即注释见github.
2.1. 分布匹配
分布匹配所关心的问题是, 当我们有两个不同的分布$p_1(X)$和$p_2(X)$, 我们要如何修改分布$p_1$使得其和$p_2$尽可能接近. 这听起来可能有些抽象, 我们不妨考虑一个很具体的例子: 假设我们要做一个目标检测的任务, 但是我们训练集全都是白天的图像, 而测试集全都是黑夜的图像, 那我们要如何调整训练集, 从而使得我们的模型最终能泛化到黑夜上呢?
Note:假设我们有数据集$\left \{ x_i\right \}_{i=1,\cdots, N}$, 则它其实就定义了一个分布, 一般称为经验分布, 该分布的概率密度函数为:
$$ p(X=x) = \frac{1}{N}\sum_i^N \mathbf{1}_{x_i = x} $$
所以两个数据集之间的差异就可以被转化为两个分布间的差异问题.
Domain Adaptation
假设我们有两个数据集, 分别为$X_1$ (蓝色) 和 $X_2$ (橘色), 其中$X_1$有$n$个点, $X_2$有$m$个点. 我们不妨就认为$X_1$是白天的图像, $X_2$是夜晚的图像.
注意这里$\mathbf{r}$和$\mathbf{c}$都是均匀分布 ( $p(X=x) = \frac{1}{N}\sum\limits_i^N \mathbf{1}_{x_i = x}$, 我们假设图片之间两两不同 ). 对于代价矩阵$\mathbf{M}$, $\mathbf{M}_{ij}$代表了把图$i$变成图$j$的代价, 这里我们就取Euclidean Distance, 即$\lVert x_i -x_j \rVert$.
接下来, 我们取$\lambda =100$, 便可以用上文的Sinkhorn-Knopp算法计算最优转移矩阵$\mathbf{P}^{*}_\lambda$, 我们将其可视化如下:
图中$X_1$中的每一个点都被soft mapping到了$X_2$中, 即$(\mathbf{P}^{*}_\lambda)_{ij}$可以被理解为是$X_1$中第$i$个样本和$X_2$中第$j$个样本之间的相似关系, 则接下来我们通过插值:
$$ x’_i = \alpha x_i + (1-\alpha)\times\sum\limits_{j=1}^{m}(\mathbf{P}^{*}_\lambda)_{ij}x_j $$
便可以得到将$X_1$的分布向$X_2$靠近后的结果, 其中$\alpha$控制了靠近的程度. 注意上式中的$\mathbf{P}^{*}_\lambda$需要逐行做归一化, 即需要保证每一行求和为$1$. 插值的结果如下:
因此Optimal Transport为这种Domain Adaptation问题提供了一种通用的解决方法:
- 将分布不同的数据集$X_1$和数据集$X_2$转化为一个最优输运问题.
- 通过最小化$X_1$和$X_2$之间的Wasserstein或Sinkhorn distance 求得输运矩阵 $\mathbf{P}^{*}$.
- 借助输运矩阵$\mathbf{P}^{*}$通过插值得到$X_1’$, 并在$X_1’$上训练模型, 从而能够泛化到$X_2$上.
Optimal Transport还可以用来求解半监督问题, 例如在半监督分类问题中, 我们有少量标注数据, 和大量无标注数据:
我们同样可以利用Optimal Transport, 计算最优输运矩阵$\mathbf{P}^{*}$, 从而将无标注样本点soft mapping到有标注样本点上, 之后借助这个mapping去计算soft label或hard label (one-hot).
Color Transfering
另一个使用Optimal Transfer来解决的直观例子是颜色迁移/风格, 不过这和近些年比较火的Deep Transfer还是不太一样的, Optimal Transfer相对来说更加朴素但是简单. 颜色迁移的任务很简单, 左边是两张给定的原图$X_1$和$X_2$, 希望转换为右边的$X_1’$和$X_2’$, 从而使得$X_1’$和$X_2$相似; $X_2’$和$X_1$相似.
我们首先需要意识到, 图片由像素构成, 一个像素不过是一个三维向量 (R,G,B), 因此一张图片就可以被理解为一个数据集 (一个分布). 在理解了这一点后, 使用Optimal Transport来实现颜色迁移就相对简单了. 我们就只描述$X_2\rightarrow X_2’$的过程.
- 首先计算从$X_2$到$X_1$的Optimal Transport问题, 我们选择Sinkhorn distance, 通过Sinkhorn-Knopp算法求得$\mathbf{P}^{*}_\lambda$.
- 通过$\mathbf{P}^{*}_\lambda$我们插值得到$X_2’$, 换言之, 对于$X_2$中的每一个样本点$x_j$, 我们有其对应的样本点$x_j’$.
- 训练一个模型$f: X_2\mapsto X_2’$, 并取$f(X_2)$作为最终结果.
为什么这里我们需要最后第三部训练模型$f$, 而不是直接取第二部得到的$X_2’$作为最终结果呢, 主要原因有两点:
- 出于训练效率原因, 实际实现时我们并不是取了$X_2$的所有像素点做训练, 而是只取了一部分, 因此最终我们需要借助模型$f$的泛化能力来将未被选中的像素点映射到$X_2’$中.
- 最终模型我们选取的是KNeighborsRegressor, 它在对$x_j$做预测时, 并不是只考虑$x_j’$. 还会考虑$x_j$的$K$个neighbor的取值, 因此我们最终得到的$f(X_2)$相对于$X_2’$来说色彩会更加平滑.
2.2. 衡量分布间的距离
此外, Optimal Transport还可以被用来作为衡量两个分布之间的距离的工具. 当然单纯衡量分布之间距离的方式有很多, 例如KL-Divergence, TV-Divergence. 但是Optimal Transport的好处在于, 我们可以通过代价矩阵$\mathbf{M}$来引入先验知识.
例如如果想比较两道菜之间的差距, 这可能是十分困难的. 例如红烧牛肉和清蒸鲈鱼之间的差异是多少呢? 但是Optimal Transport提供了一个直觉上的思路.
首先我们将两道菜打散成菜谱, 例如$\text{红烧牛肉} = \begin{bmatrix} 牛肉\\酱油\\\vdots\\八角\end{bmatrix}$, $\text{清蒸鲈鱼}=\begin{bmatrix} 鲈鱼\\\vdots\\耗油\end{bmatrix}$, 相对来说度量菜谱里的每一个原材料之间的距离更为简单, 我们便可以得到代价矩阵$\mathbf{M}$. 之后如果我们还知道每道菜里原材料之间的相对比例, 我们还可以将其嵌入到$\mathbf{r}$和$\mathbf{c}$中. 因此, 在嵌入了如此多的知识之后, 我们有理由相信通过Optimal Transport求得的距离是更加合理的.