Optimal Transport Distance(最优传输距离)

date
Sep 3, 2022
Last edited time
Mar 27, 2023 08:49 AM
status
Published
slug
Optimal_Transport_distance(最优传输距离)
tags
Algorithm
summary
重新翻翻OTA笔记, 发现有些地方不懂, 于是重新看了点相关资料, 整理成这份笔记
type
Post
origin
https://www.notion.so/lazurite/Optimal-Transport-distance-30a6731457d141b383258ab99771d7ea
Field
Plat

阅读资料

Computational Optimal Transport
Optimal transport (OT) theory can be informally described using the words of the French mathematician Gaspard Monge (1746-1818): A worker with a shovel in hand has to move a large pile of sand lying on a construction site.
Computational Optimal Transport
https://arxiv.org/abs/1803.00567
Computational Optimal Transport
1803.00567.pdf
41435.2KB
最优传输 (Optimal transport) 和 sinkhorn 迭代
本文首先简要介绍离散情况下的最优传输问题，然后介绍文献 [1] 所提出的最优传输问题的快速解法以及对应的对应梯度的计算。 为了理解的方便，本文只介绍离散情况下的最优传输问题。 值得注意的是，有很多关于最优传输的资料会用到不同的关键词， 例如 optimal transport / wasserstein distance / earth mover's distance / sinkhorn distance 等。 离散条件下 optimal transport / wasserstein distance / earth mover's distance 可以认为是等价的， Sinkhorn iteration 是一种最优传输的快速迭代解法，后文中会介绍到。 本文涉及到的数学符号均沿用文献[1] ，并补充了详细的说明。
最优传输 (Optimal transport) 和 sinkhorn 迭代
https://kaizhao.net/blog/optimal-transport
NIPS-2013-sinkhorn-distances-lightspeed-computation-of-optimal-transport-Paper.pdf
360.4KB
阅读资料Earth mover’s distanceEMD 度量概率分布的距离熵正则为什么要增加熵正则几个性质Sinkhorn 迭代求解最优传输PyTorch 实现Reference
Video preview

Earth mover’s distance

假设地面上有 个土堆，第 个土堆有 数量的土； 同时有 个坑，第 个坑可以容纳的 数量的土。 我们用
来分别表示这土堆和坑，并假设 ，也就是 所有的 “土” 刚好够填满所有的 “坑”。 定义 为距离矩阵，其中 表示 “从第 个土堆搬运一份土到第 个坑的成本 (cost)”。 现在要将土从 搬到 ，定义 为从 到 的 “传输方案” (transportation paln)， 其中 表示从 搬运到 的土的数量。 最优传输目的是求得一个最优的传输矩阵 ，使得总体的传输消耗 (cost) 最小。
很显然，传输矩阵 的每一行 / 列 加起来应该等于 ，这也是为什么用字母 r(ow)/c(olumn) 来表示的原因。 写成公式， 满足约束：
即所有满足条件的 transportation plan 的集合为：
其中 表示 维的全 1 列向量。 值得一提的是， 构成高维空间的多面体 (polytope) ， 因为上式的所有约束都是线性的，平面上多个线性约束条件围成的区域是多边形， 拓展到高维空间后线性约束围成的集合就是多面体（polytope）。 这里用集合 来表示所有满足条件的传输，在优化中集合 又被称为 “可行域” (feasible set)。
最优传输 定义为 “使得总体运输成本 最小的运输方案”：
其中 为 Frobenius inner product 。 使得上式最小的传输方案 称为最优传输方案。
那么, 这个最小的传输成本定义为:
简而言之，最优传输就是要在可行域 中找到找到使得式2最小的传输方案 。 很显然这是一个线性规划问题，因为不论是式中的约束条件还是式中的目标函数都是线性的。

EMD 度量概率分布的距离

如果 和 满足 且 ， 那么 和 可以看作两个概率分布， 也可以看作是边际分布分别为 和 的 “联合概率分布”。 和 之间的最优传输距离可视作概率分布之间的差异。 在很多应用中都会将输入 和 归一化成概率（比如分类中的 softmax），最后计算两个概率的最优传输距离。 因此最优传输也可用作度量概率分布之间的距离。
那么我们为什么要计算 optimal transport 呢？
如果为了 measure 概率分布之间的距离，有很多现成的 measurements 可以用， 比如非常简单的 KL 散度:
除了 “无法处理两个分布的支撑集不相交的情况” 以及 “不满足对称性” 等原因之外，一个重要的原因就是这种逐点计算的度量 没有考虑分布内的结构信息。 所谓的结构信息，就是分布内的联系。例如在 KL 散度中， 彼此之间都是独立计算最后加起来的， 而大部分情况下它们并不是独立的。
就以我们常见的分类任务为例，分类任务通常用交叉熵损失来度量模型预测和样本标签之间的距离， 交叉熵损失实际上就是在计算 onehot 化的标签和模型预测之间的 KL 散度。 这种逐点计算的损失函数（不论是交叉熵还是 L2）都无法考虑分布内不同事件的相关性。 例如将 “汽车” 误分类成 “卡车” 显然没有把 “汽车” 误分类成 “斑马” 严重。 但是用 KL 散度来度量的话，这两种错误的损失是一样的。
假设 是两个离散概率分布，满足 ，其中 表示全1列向量， 我们可以用最优传输来计算两者之间的距离。 而概率分布内的结构信息可以通过距离矩阵 “嵌入” 到距离度量中。 还是以分类为例，我们可以让 远小于 。

熵正则

在文献 [1] 中，Cuturi 提出了一种快速的最优传输求解方法。该方法首先引入熵正则，使得原问题的可行域更加平滑， 然后将最优传输问题转为 matrix permutation 问题，最后用 sinkhorn 迭代算法求原问题的近似解。
联合分布 的熵 (entropy) 为:
原问题 式2 增加熵正则约束之后，优化目标变为：
式6 称为 “熵正则的最优传输” (entropy-regularized optimal transport)，对应的最优解表示为 。 从概率和信息论的角度，均匀分布的熵最大，因此熵正则会让最优传输矩阵 更趋近于均匀分布。
从优化的角度， 式2中的原问题的解一定存在于多边形的某一顶点处， 因此 是一个稀疏矩阵，绝大部分元素都是 0。 增加熵正则之后，相当于将原来的可行域 向内收缩成光滑的 ，对应的最优解 不再是稀疏矩阵。
可行域 和 包含所有满足约束的点， 是一个高维空间中的多边形 (polytope)，增加熵正则之后的可行域 是一个边界光滑的区域。 蓝色虚线表示目标函数的等高线，距离矩阵 决定最优解所在的方向。 本图参考了 文献[1]的 Fig.1。
可行域 和 包含所有满足约束的点， 是一个高维空间中的多边形 (polytope)，增加熵正则之后的可行域 是一个边界光滑的区域。 蓝色虚线表示目标函数的等高线，距离矩阵 决定最优解所在的方向。 本图参考了 文献[1]的 Fig.1。
当 时， 趋向于 ，最优解 也趋向于 ； 当 时， 收缩成点 ，最优解为 , 。

为什么要增加熵正则

参考文献 [1] 中作者列出了两个使用熵正则的原因：
  1. 式 中的原线性规划问题的解一定是在可行域 的某个顶点 (vertex) 上的，因此得到的解 是一个稀疏矩阵。 稀疏的解 会使得最终的传输方案十分不均衡， 使用熵正则可以让传输矩阵更加均衡，这部分可以参考 [1] 的第三章相关内容。
  1. 更重要的是，增加了熵正则之后，原问题可以使用 sinkhorn 算法得到近似解，而且计算开销大大降低。 关于这部分，可以参考原文第四章。
总的来说，增加熵正则之后的最优传输问题得到的解更加光滑和均衡，同时可以使用 sinkhorn 迭代算法快速求解。

几个性质

  1. 对于 , 解 是唯一的, 可以写为 的形式, 其中 , .
  1. 将约束写在拉格朗日乘数法的形式得到:
      2. 令 .
      由于 的每一个元素均为正数, 依据 Sinkhorn’s theorem, 存在一个唯一的矩阵 . 那么, 我们可以使用 Sinkhorn 不动点迭代来计算这个矩阵.

Sinkhorn 迭代求解最优传输

假设有向量 ，令 的初始值为
然后用 sinkhorn 迭代算法求解：
上式中 (参考 [1] 中 Lemma2)。 迭代收敛后最优传输矩阵 和对应的最小传输距离 可以由以下公式给出：
对 的导数分别是:
有了 公式10 和 公式11， 最优传输距离就可以作为损失函数来训练神经网络。

PyTorch 实现

这里 (vlkit.optimal_transport.sinkhorn) 有一个 sinkhorn 的 PyTorch 实现，我们可以计算并可视化两个分布之间的最优传输。
以高斯分布为例，首先生成两个 1d 高斯分布作为 source 和 target distribution：
notion image
然后通过 sinkhorn 迭代计算最优传输，并可视化结果。
notion image

Reference

  1. Cuturi, Marco. “Sinkhorn distances: Lightspeed computation of optimal transport.” Advances in neural information processing systems 26 (2013): 2292-2300.
  1. Arjovsky, Martin, Soumith Chintala, and Léon Bottou. “Wasserstein generative adversarial networks.” International conference on machine learning. PMLR, 2017.
  1. 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.
  1. 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.
  1. Xie, Yujia, et al. “Differentiable top-k operator with optimal transport.” arXiv preprint arXiv:2002.06504 (2020).
  1. Ge, Zheng, et al. “OTA: Optimal Transport Assignment for Object Detection.” Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2021.
  1. Frogner, Charlie, et al. “Learning with a Wasserstein loss.” Advances in neural information processing systems (2015).

© Lazurite 2021 - 2024