FFT-OT: A Fast Algorithm for Optimal Transportation整理

FFT-OT: A Fast Algorithm for Optimal Transportation

 

FFT-OT：一种最优传输的快速算法

---

一篇使用快速傅里叶变换解决最优传输问题的文章，去年组会上讲过，今天睡不着把翻译翻出来贴一下，顾老师的网站上提供了C++版本的源码和数据集 。Fast Fourier Transform - Optimal Transport

预备知识：不动点 有限差分法 FFT(傅里叶变换） DCT（离散余弦变换）

最优传输部分只需要理解蒙日安培方程怎么来的，为何难解， Brenier理论和Kantorovich势能即可

摘 要

一个最优传输可以找到最经济的方法来将一个概率测度传输到另一个。它已广泛应用于视觉、深度学习和医学图像。根据Briner理论，计算最优传输映射等价于求解一个蒙日-安培方程。由于具有高度非线性的性质，计算最大规模优传输映射是非常具有挑战性的。

本文基于三个关键思想，提出了一种简单但功能强大的方法，即FFT-OT算法，来解决这一难题。首先，将蒙日-安培方程转换为一个不动点问题；其次，将最优运输映射的倾斜条件重新表示为矩形域上的Neumann边界条件；第三，为提高效率，在每次迭代中应用FFT来求解一个泊松方程。

对三维扫描采集和医学成像重建的表面进行了实验，并与现有的其他方法进行了比较。实验结果表明，所提出的FFT-OT算法简单、通用、可扩展，高效、精度高。

1. 介绍

近年来，最有传输的快速发展及其在视觉、深度学习和医学成像领域的广泛应用。一个最优传输映射(OT map)找到了讲一种概率度量传输到另一种概率度量的最经济方法，将总运输成本视为两者之间Wasserstein distance随机分布。因此，OT map被用于测量概率分布之间的差异，以及配准图像或3D图形。

动机

由于高度非线性，提高OT映射的计算效率成为核心挑战之一。研究人员开发了许多不同方法来解决这个难题。例如，Sinkhorn算法通过增加一个熵正则化项，大大提高了坎托洛维奇势能的优化速度，但牺牲了准确性。凸几何变分算法找到精确的解，但具有复杂的动态几何数据结构和自适应运算精度。流体动力学算法在时空上找到最优流，从而增加了维数，降低了计算效率。因此，现有的方法很难有效地找到大规模OT问题的准确解决方案。为了解决这一挑战，本文提出了一种新的算法：FFT-OT利用基本和强大的工具：快速傅里叶变换，提供高效和准确的大规模最优运输映射。

主要思想

该算法基于三个简单的思想。目标是找到两个概率测度之间的OT映射，基于brenier定理，将其简化为通过求解蒙日-安培方程，找到brenier势能

1. 不动点 通过改进[10,4]中的公式，构造了索波列夫空间中的一个算子使蒙格-安培方程的解u∗为P的不动点，即。当n→∞，u(n)收敛到u∗，u(n)→u∗时，可以通过迭代，来实现不动点。迭代过程的收敛性和求解的凸性有理论上的保证。
2. 倾斜边界条件 边界条件很难直接表述。我们的公式是基于（5）式中的倾斜性边界条件。

 

假设域Ω是一个平面矩形，brenier势能被重写为φ是众所周知的Kantorovich势能，倾斜性条件等价于φ的纽曼边界条件

     3.快速傅里叶变换 

在不动点迭代过程中，算子是为了求解一个泊松方程等式（4）和等式 (6).

         由于该域具有自然的规则网格结构，因此使用非常方便用有限差分法求解泊松方程。更重要的是，由于离散的拉普拉斯矩阵具有固定的结构，因此泊松方程可以用FFT方法求解来减少复杂度。

1. 相关工作

Kantorovich公式

Monge-Kantorovich理论通过线性规划技术解决最优传输问题，卡马卡尔提出了一种用于线性规划的多项式时间算法。Sinkhorn提出了一个熵正则化项来高效的得到近似解。近年来，人们提出了许多基于Sinkhorn算法的算法。这些方法首先在素数Kantorovich公式中加入一个熵正则化项，然后求解其素数问题或对偶问题。热纳等提出了一种熵正则化技术基于随机优化的方法来处理大规模的OT问题。在理论上，这类算法只能给出近似值。

Brenier 公式

Gu等建立了凸几何中布Brenier定理和Alexandrov定理之间的理论联系。Su等人提出了一种求解最优输运问题的几何变分方法用于保面积映射。该方法已推广到具有更复杂拓扑的曲面。这些方法利用功率图计算Brenier势能，这需要复杂的几何数据结构和自适应算法。

流体力学、数值方法

略

对比

与坎托罗维奇公式中的方法相比，本文提出的算法找到了精度较高的解，近似误差为O(h2)，其中h为空间步长长度；相比于利用几何优化计算Brenier势能的方法，本文的算法不需要复杂的数据结构，因此大大简化了设计，改进了设计效率；与最直接的数值方法相比，我们的算法使用FFT来求解泊松方程，因此更简单、更快。

1. 基础概念

蒙日问题

给定两个概率分布f(x)dx和g(y)dy，支集分别为Ω和Ω∗。如果C1映射T：Ω→Ω∗有，则T是保测度的，其中DT是映射的雅可比矩阵，并记为。给定一个损失函数，表示将单位质量从x点传输到y的成本。映射T的传输成本定义为蒙日提出了寻找最优传输映射的问题，即具有最小传输成本的保测度映射。

Brenier 理论

布雷尼尔定理在温和的正则性条件下，给出了蒙格问题的答案：如果传输成本是二次欧几里得距离，存在一个凸函数，其梯度映射是唯一的最优传输映射，u被称作为brenier势能，满足蒙日-安培方程

边界条件是

倾斜性边界条件

假设Ω是凸的，Ω∗的边界是二阶光滑的，并且满足一些其他的条件，边界点的法向量和对应像点的法向量夹角小于等于直角。倾斜边界条件蕴含着，对所有边界上的点，都只会被映射到边界上，角点只会被映射到角点，否则法向量的夹角将大于90°

1. 计算算法

由于蒙日安培方程是高度非线性的，非常难解。本文的思想主要是将OT问题转化为不动点问题，并且用FFT进行求解。

不动点

在二维情况下，蒙日-安培方程可以写成

因此有，

可导出泊松方程

其中。由于Brenier势能是凸的，拉普拉斯是非负的，因此我们在上述公式中取正的平方根。这保证了计算过程Brenier势能的凸性。定义一个算子，

 

蒙日-安培方程的解是的不动点u*，，因此，我们可以通过迭代的方法来寻找不动点。

倾斜性条件

在目前的工作中，我们主要关注矩形平面域

如图所示，将边界分为了γ0到γ3四个部分，分别对应法向量n0到n3，给定一个边界点，x∈γk，它的像点也是一个边界点，T(x) ∈ ∂Ω*，通过引理3.2，我们知道边界点的法向量和对应像点的法向量夹角小于等于直角，故有

因此，T(x)也在γk上，即x和T(x)都在同一个γk上，由于每个角点属于2个相邻的段，因此它的图像是它本身，T(pk) = pk。  令，φ是坎托罗维奇势能，我们得到了Du=Dφ+Id和∆u=∆φ+2。(4)式中的运算符，被转化为(6)式中的。倾斜性条件(5)变为Neumann边界条件，即∂φ/∂n=0.

有限差分法

如图3所示，图像域Ω∗被离散为笛卡尔M×N网格。步长由hx=2/M和hy=2/N给出。每个网格点的坐标由给出，我们在边界上添加“幽灵cell”，如图3(a)所示，以计算Neumann边界条件的高阶导数。布雷尼尔势表示为二维M×N矩阵(uij)，用中心差分法近似这些微分，

 

离散的拉普拉斯矩阵具有规范形式（见补充文档），正则结构允许我们应用FFT方法来加快计算速度。假设给出了一个具有Neumann边界条件∂u/n=0的泊松方程∆u=ρ，一个解存在的必要条件可以推导出如下

如果ρ不满足该条件，我们可以向ρ添加一个偏移量c，这样ρ-c的积分为0。用于求解在Neumann边界条件下离散泊松方程

我们可以使用离散余弦变换(DCT)方法。给定二维数组u（i,j），二维DCT为

逆DCT为

其中m, i = 0, 1, . . . , M - 1   n, j = 0, 1, . . . , N - 1

令，离散泊松方程（9）的解是，可证明

 

通过将设置为零，我们可以规范化解。利用快速傅里叶变换(FFT)方法可以有效地计算DCT和逆DCT。算法总结如下。

DCT Poisson方程求解器

蒙日-安培方程求解器

5.实验

略

结论

        最优传输在计算机视觉中起着重要的作用，但计算方法具有挑战性。本文提出了一种新的FFT-OT算法，以提高其计算效率和精度。关键思想是将最优传输问题转化为固定点问题，将倾角条件表示为诺伊曼边界条件，并利用FFT求解泊松方程来迭代求解。实验结果表明，FFT-OT算法简单、通用、可扩展，精度高、效率高。在未来，FFT-OT算法将被推广到处理高维情况和球形情况，以更广泛地应用于视觉和深度学习。