最优传输问题和Sinkhorn

最优传输问题

假设有$M$堆土，每堆土的大小是$a_m$，有$N$个坑，每个坑的大小是$b_n$，把单位土从土堆$m$运送到坑$n$的代价是$c(m,n)$，如何找到一种运输方法填满坑，并且代价最小，这就是最优传输问题（optimal transport (OT) problem）。

假设有两个概率分布，如何以最小的成本将一种概率分布转换为另一种概率分布，这也是最优传输问题。这个最小的成本可以作为度量两个概率分布的距离，被称为Wasserstein距离，或者推土机距离（Earth Mover’s Distance（EMD））。

在离散的情况下，假设$\mathbf{r},\mathbf{c}\in {\mathbb{R}}_{+}^d$是两个概率向量，也就是元素求和为1。$1_d$是维度为$d$所有元素为1的向量。
运输多面体（transport polytope ）$U(\mathbf{r},\mathbf{c})$被定义为：
U ( r , c ) : = { P ∈ R + d × d ∣ P 1 d = r , P ⊤ 1 d = c } U(\mathbf r,\mathbf c) := \{ \mathbf P \in \mathbb R^{d \times d}_+ | \mathbf P \mathbf 1_d = \mathbf r, \mathbf P^\top \mathbf 1_d = \mathbf c\} U(r,c):={P∈R+d×d​∣P1d​=r,P⊤1d​=c}给定一个费用矩阵$\mathbf{M}\in {\mathbb{R}}^{d\times d}$，$\mathbf{r}$到$\mathbf{c}$的最优传输距离被定义为：
$$d_{\mathbf{M}}(\mathbf{r},\mathbf{c}):=\underset{\mathbf{P}\in U(\mathbf{r},\mathbf{c})}{\min }<\mathbf{P},\mathbf{M}>=\sum _{i,j}{\mathbf{P}}_{ij}{\mathbf{M}}_{ij}$$对于一般的矩阵$\mathbf{M}$，目前提出的最佳算法在最坏情况下的复杂度是 $O(d^3\log d)$。在实践中复杂度也被证明是超立方的。

Sinkhorn距离

直接求解最优传输问题的复杂度非常高。为了解决这个问题，考虑在限定的范围内求解。
定义凸集（convex set）：
U α ( r , c ) : = { P ∈ U ( r , c ) ∣ K L ( P ∣ ∣ r c ⊤ ) ≤ α } = { P ∈ U ( r , c ) ∣ h ( P ≥ h ( r ) + h ( r ) − α } ⊂ U ( r , c ) U_\alpha(\mathbf r,\mathbf c) := \{\mathbf P \in U(\mathbf r,\mathbf c) | KL(\mathbf P || \mathbf r \mathbf c^\top) \leq \alpha\} = \{\mathbf P \in U(\mathbf r,\mathbf c) | h(\mathbf P \geq h(\mathbf r) + h(\mathbf r) - \alpha\} \subset U(\mathbf r,\mathbf c) Uα​(r,c):={P∈U(r,c)∣KL(P∣∣rc⊤)≤α}={P∈U(r,c)∣h(P≥h(r)+h(r)−α}⊂U(r,c)其中$h(\cdot )$是香浓熵（Shannon entropy）:
$$\begin{gathered} h(\mathbf{r})=-\sum _i{\mathbf{r}}_i\log {\mathbf{r}}_i \\ h(\mathbf{P})=-\sum _{i,j}{\mathbf{P}}_{ij}\log {\mathbf{P}}_{ij} \end{gathered}$$Sinkhorn distance被定义为：$$d_{\mathbf{M},\alpha }(\mathbf{r},\mathbf{c}):=\underset{\mathbf{P}\in U_{\alpha }(\mathbf{r},\mathbf{c})}{\min }\sum _{i,j}{\mathbf{P}}_{ij}{\mathbf{M}}_{ij}$$这是熵约束的最优传输问题。

上面的熵约束的最优传输问题可以通过拉格朗日乘数法（Lagrange multiplier）转换为
$$d_{\mathbf{M}}^{\lambda }(\mathbf{r},\mathbf{c})=\underset{\mathbf{P}\in U(\mathbf{r},\mathbf{c})}{\min }\sum _{i,j}{\mathbf{P}}_{ij}{\mathbf{M}}_{ij}-\frac{1}{\lambda }h(\mathbf{P})\text{\hspace{1em}(1)}$$$d_{\mathbf{M}}^{\lambda }(\mathbf{r},\mathbf{c})$被称为dual-Sinkhorn divergence。
通过对偶理论可以知道，对任意$\alpha$，有一个对应的$\lambda \in [0,\infty ]$使得$U_{\alpha }(\mathbf{r},\mathbf{c})=d_{\mathbf{M}}^{\lambda }(\mathbf{r},\mathbf{c})$。
这可以看成为最优传输问题加上熵正则化。
当$\lambda \to 0$时，上面问题的解是${\mathbf{P}}_{ij}={\mathbf{r}}_i{\mathbf{c}}_j$；当$\lambda \to \infty$时，回到了原始的最优输运问题。
香浓熵要求分配更加均匀， 参数$\lambda$权衡了按花费分配和平分。

加上熵正则的最优传输问题变得更好计算了，因为解变得平滑。
Sinkhorn定理被用来寻找熵正则化最优输运问题的解。

Sinkhorn定理

Sinkhorn 定理指出每个所有元素为正的方阵都可以写成某种标准形式。
具体而言，假设$\mathbf{A}$是一个$n\times n$的所有元素为正的方阵，则存在所有元素为正的向量${\mathbf{d}}_1$和${\mathbf{d}}_2$，使得$\text{diag}({\mathbf{d}}_1)\mathbf{A}\text{diag}({\mathbf{d}}_1)$是双随机（doubly stochastic）的。双随机矩阵是非负实数方阵，且每个行和列求和均为1。${\mathbf{d}}_1$和${\mathbf{d}}_2$在常数因子倍上是唯一的。

Sinkhorn算法非常简单，通过迭代的方法，交替地缩放$\mathbf{A}$的所有行和所有列使其和为 1。
$({\mathbf{d}}_1,{\mathbf{d}}_1)\leftarrow (1./\mathbf{A}{\mathbf{d}}_2,1./{\mathbf{A}}^{\top }{\mathbf{d}}_1)$

使用Sinkhorn算法求解熵正则化最优输运问题

可以证明公式(1)具有唯一解，且解具有形式${\mathbf{P}}^{\lambda }=\text{diag}(\mathbf{u})\mathbf{K}\text{diag}(\mathbf{v})$，$\mathbf{u},\mathbf{v}$是所有元素为正的向量，$\mathbf{K}:=e^{-\lambda \mathbf{M}}$。
这可以通过Sinkhorn算法求解。注意这不是原始的Sinkhorn算法，因为${\mathbf{P}}^{\lambda }$的每个行和列的和由$\mathbf{r}$和$\mathbf{c}$确定，而不再是1。
$(\mathbf{u},\mathbf{v})\leftarrow (\mathbf{r}./\mathbf{Kv},\mathbf{c}./{\mathbf{K}}^{\top }\mathbf{u})$

参考资料

Wiki Sinkhorn’s theorem
Notes on Optimal Transport
http://alexhwilliams.info/itsneuronalblog/2020/10/09/optimal-transport/
https://zipjiang.github.io/2020/11/23/sinkhorn’s-theorem-,-sinkhorn-algorithm-and-applications.html