FFT-OT: 最优传输映射的快速傅里叶变换方法
最优传输（Optimal Transportation）日益成为深度学习的理论基础之一。在由芒福德（MumFord）、朱松纯学派率先提出的统计学习理论框架中，每个“概念”可以被视为数据空间中的一个“概率分布”，其支集为某种低维流形，深度学习核心是学习概率分布。深度学习算法本质上是在流形上所有概率分布构成的空间（Wasserstein Space）中进行变分（variation），在满足特定限制的条件下优化特定能量函数，例如最大熵原则（maximum entropy principle）。最优传输理论在Wasserstein空间中定义了黎曼度量（Riemannian metric），从而得到协变微分，这样变分法才可以实行。从这个角度而言，最优传输理论用第一性原理解释了数据驱动的统计学习方法。
但是在实践中，最优传输映射归结为求解蒙日-安培方程（Monge-Ampere Equation）。由于蒙日-安培方程具有强烈的非线性，计算复杂度较高。高效而精确地计算最优传输映射目前成为深度学习领域的一个热门话题。迄今为止，绝大多数的研究者都致力于改进基于Kantorovich理论的算法，例如添加熵正则项，或者Nestorov光滑化方法等等。而笔者和合作者们一直改进基于Brenier理论的算法，本质上是因为这种方法具有丰富的几何内涵。近期，我们发现经典的快速傅里叶变换方法可以被用来加速求解最优传输映射，更多细节可参阅【1】。
最近笔者团队将过去几年中有关最优传输理论和算法的课堂讲义与研究报告加以整理总结，成书出版【2】。新书将于2021年10月中旬开始发行。希望能够对于广大的青年学生有所帮助，同时欢迎大家的批评指正！
图1. 用FFT-OT计算的最优传输映射。弥勒佛曲面（左上帧）共形地映射到平面正方形上（左下帧），曲面的面元被推前（push forward）到正方形上，从而定义了一个平面上的测度。我们的算法计算了从面元测度到勒贝格测度之间的最优传输映射（右下帧），相应的势能函数显示在右上帧中。
最优传输映射简介
让我们回忆一下经典的最优传输问题。给定概率分布和，具有密度函数和，满足总质量相等：
假设映射是光滑，满足下面的Jacobi方程这里是映射的Jacobi矩阵，则我们说映射是“保测度”的，记为。给定代价函数，表示将单位质量从传输到所需的代价，则映射的总传输代价为蒙日（Monge）提出了经典的最优传输问题：在所有保测度的映射中，使得总传输代价最小者被称为是最优传输映射（Optimal Transportation Map）。Brenier定理指出，如果，则存在凸函数（Brenier势能函数），使得最优传输映射等于Brenier势能的梯度：。将Brenier势能函数代入Jacobi方程，我们得到经典的蒙日-安培方程：
具有所谓的第二类边界条件：更进一步，如果支撑集的边界和是光滑的，则最优传输映射满足所谓的倾斜条件(Obliqueness condition)，即区域边界点处的外法向量满足：
这里, 。核心思路
不动点 在二维情形，由蒙日-安培方程得到Brenier势能函数，
因此我们有
由此我们得到泊松方程
这里Laplace算子。我们定义索伯列夫空间的算子：
则蒙日-安培方程的解（Brenier势能函数）是算子的不动点。
倾斜边界条件 我们考虑和是同一个平面长方形区域，由第二类边界条件和倾斜条件, 对于任意非角点的边界点， 其像点也一定是非角点的边界点，同时满足
这意味着和都在长方形的同一条边上。由此，我们可以得到角点的像等于自身。
我们改写Brenier势能函数，
那么是算子的不动点，
我们开始设置为，然后迭代，可以证明迭代过程收敛到不动点。
这时，Brenier势能函数的倾斜边界条件转换成的Neumann边界条件
有限差分法 迭代过程中的每一步，我们都是求解Neumann边值条件下的Poisson方程。我们将区域离散化成二维点阵，微分算子用差分算子近似：
这里和是水平和铅直方向的步长。如此，Poisson方程变成差分方程
快速傅里叶变换
具有Neumann边值条件的Poisson方程, 可以用离散余弦变换来解，
逆变换为
这里，其他情形。令，, 可以直接得到
即在频域，Poisson方程的解可以直接写出。离散余弦变换可以用快速傅里叶变换实现，从而进一步提高计算效率。
小结
FFT-OT的核心想法如下：蒙日-安培方程是强烈非线性的退化椭圆型偏微分方程，局部上我们可以用线性椭圆型偏微分方程来逼近，从而用迭代方法求解。线性椭圆型偏微分方程可以用傅里叶变换转换成频域的代数方程，具有解析解。快速傅里叶变换可以用硬件支持，例如CUDA或者FPGA，进一步提高计算效率。
原则上，FFT-OT方法对于任意维度的最优传输映射都适用，但是目前高维的FFT-OT有待进一步研发。我们期待未来这种算法被工业界采用，从而开发出专用程序库和芯片，进一步推动统计深度学习的发展。
（ 仅以此文献给纯粹的理想主义者：师兄施皖雄博士。）
N. Lei and X. Gu, FFT-OT: AFast Algorithm for Optimal Tansportation, ICCV 2021.
N. Lei and X. Gu, 最优传输理论与计算，高等教育出版社，2021年10月。
请长按下方二维码，选择 “识别图中二维码”，即可关注。
【老顾谈几何】邀请国内国际著名纯粹数学家，应用数学家，理论物理学家和计算机科学家，讲授现代拓扑和几何的理论，算法和应用。