最优传输几何变分法的C++实现
依随深度学习方法的深入发展,人们逐渐意识到最优传输理论所起到的奠基作用。深度学习可以抽象概括为学习流形上的概率分布,因此概率分布之间的变换和衡量概率测度之间的距离是深度学习最为核心的主题。而最优传输理论为解决这些问题提供了理论基础和计算工具。更进一步,最优传输理论为所有概率测度构成的空间定义了黎曼度量和协变微分,从而使得在所有测度构成的空间中进行变分优化成为可能。因此,我们可以预见在深度学习热潮渐落之后,最优传输理论经过大浪淘沙,会凝练成思想精华,全面渗透入工程领域。
但是,描述最优传输映射的偏微分方程具有强烈的非线性,因此精确计算最优传输一直具有挑战性。目前在工程领域,特别是在深度学习领域,人们都是将概率测度离散化,将Kantorovich能量加上熵正则项,得到光滑的近似,逼近误差用一个参数进行控制,例如常见的Sinkhorn算法,Nesterov近似等等。我们经过大量实验,发现这些算法有很多内在缺陷。如果我们希望精确计算连续分布之间的最优传输映射,这些算法需要对源区域和目标区域稠密采样,如此计算复杂度过高,无法用普通硬件实现。同时这些算法给出了近似解,如果我们希望通过减小控制参数来提高逼近精度,因为算法依赖超越运算(指数对数运算),算法稳定性迅速变差,无法收敛。如果我们细致地调节参数以保证算法的稳定性,这时得到的近似解与真解相距甚远。
我们一直希望能够对最优传输映射进行深入细致地理论研究,因此要求能够得到高精度的数值解,同时可以控制误差范围。Sinkhorn算法虽然火热,但是理论上无法真正满足需求。我们发现,目前几乎不存在高精度最优传输算法的软件工具。在基础和应用数学领域,虽然大家数十年如一日地研究蒙日-安培方程,但是并没有比较普及的计算工具;在计算机科学领域,绝大多数的算法都是基于Kantorovich的线性规划及其光滑近似(Sinkhorn)。这些方法空间复杂度很高,逼近误差较大,更为严重的是,这些方法破坏了最优传输问题本身的理论结构,摒弃了内在的几何特性。目前工程领域对于这些算法的热衷,很大程度上是用方便来取代有效,并且将最优传输视为一个成型的工具,而没有深入理解其内在的理论结构。
因此,在我们的课程上,我们着重讲解了最优传输问题的几何求解方法,这一方法与蒙日-安培方程的弱解理论相一致,与Minkowski和Alexandrov的凸几何理论相吻合,从几何角度来理解诠释最优传输理论,给人以强烈的几何直觉和深刻细致的洞察。关于最优传输映射,我们依然有太多的未知,我们希望几何方法提供的计算工具能够帮助大家更加深入地研究这一理论,对于蒙日-安培方程具有更加深刻的理解,从而进一步推动工程应用的发展。
理论回顾
给定欧氏空间的概率测度及其支集和,满足总测度相等,密度函数和,。一个映射是保持测度的,如果对一切,都有
记为。给定传输代价函数,蒙日问题是求解所有保测度映射中传输代价最小者,
如果传输代价是欧氏距离的平方,和紧致,为凸区域,那么Brenier定理断言存在一个凸函数,即所谓的Brenier势能函数,其梯度映射给出了最优传输映射,,并且这种最优传输映射是惟一的。Brenier势能函数满足蒙日-安培方程(Monge-Ampere)
并且满足边界条件。由此,求解最优传输映射等价于求解蒙日-安培方程。
我们在中离散采样,得到采样点集,用Dirac测度之和来逼近目标测度,由此来考虑半连续的最优传输问题。由Brenier定理,最优传输映射等于Brenier势能函数的梯度映射,这时Brenier势能函数为分片线性函数,
这里,每个采样点对应一个支撑平面,高度为未知量,的图是这些支撑平面的上包络(upper envelope),其勒让德对偶的图是点的凸包(convex hull)。
图1. Brenier势能函数与勒让德对偶。
如(1)图所示,Brenier势能函数的投影得到上的一个Power图(power diagram),
Brenier势能函数的Legendre对偶的投影,得到的Weighted Delaunay三角剖分 。通过调节每个对应支撑平面的高度,我们可以得到这样一个Power图,每个Power胞腔的体积等于,
丘成桐先生,笔者和罗锋【6】曾经证明,所求高度向量是下面能量的最大化子,
在可容许高度空间
这一能量是严格凹的,最大化子在可容许空间的内部。由定义,我们有能量的梯度分量等于,能量Hessian矩阵的分量等于Power Voronoi边长与相应的Weighted Delaunay边长之比。由此,我们可以用牛顿法来优化凸能量从而得到唯一解。
计算几何算法
图2. 凸包算法的结果。
我们看到,计算所需要的主要概念和算法都来自经典的计算几何。最为主要的是Legendre对偶,这需要计算的凸包。我们可以采用Edelsbrunner的经典增量凸包算法。这也等价于计算Weighted Delaunay三角剖分,我们可以用Lawson的边翻转(edge flip)算法来实现。如果图(2)所示。
图3. 上包络算法的结果。
得到凸包之后,我们计算其相应的上包络,即计算的勒让德对偶,凸包上每一个顶点对偶于一个平面,每条边对偶于两个顶点对偶平面的交线,每个面对偶于三个顶点对偶平面的交点。同时我们添加一个无穷远平面,如此得到平面族的上包络,如图(3)所示。这里,边界支撑平面被无穷远平面相截。
图4. Sutherland-Hodgman算法。
上包络在上的投影得到Power图,我们需要计算每个Power胞腔与的交集,即用多边形来剪裁。平面多边形裁剪有多种算法,比较高效的是Bentley-Ottman的扫描线算法(sweepline),比较简单的是Sutherland-Hodgman的凸多边形剪裁算法。如图(4)所示,我们用凸的红色多边形来剪裁蓝色多边形,每一步我们用红色多边形的一条边来将平面一分为二,保留内侧去除外侧。如此重复,每次切掉一角,直至穷尽红色多边形的所有边。
图5. Brenier势能函数,边界裁剪之后的结果。
图(5)显示了裁剪之后的Brenier势能函数,和相应的Power图。然后,我们计算每个Power胞腔的面积,每条Power Diagram边的长度,构造能量的梯度和Hessian矩阵,计算更新的高度向量,
然后我们进行试探性的更新:选定一个步长,。我们需要判断新的高度向量是否落在可容许空间内部。我们先用新的高度向量计算凸包,如果存在某个不落在凸包表面,则意味着新的高度向量非法。或者,所有的顶点都落在凸包表面上,但是某个Power胞腔落在之外,则新的高度向量非法。若新的高度向量非可容许,则我们将步长减半,重新试探。这种方法我们称之为阻尼算法。几年前,笔者与丘成桐先生和汪徐家教授分别探讨过求解蒙日-安培方程的计算问题,他们都一针见血地指出:求解过程的关键是保证中间解的严格凸性,优化过程一直在可容许解空间之中。所以,这一步是整个算法的关键。
图6. 算法的收敛过程。
图(6)显示了一个计算实例,由于我们采用了牛顿法来优化凸能量,算法收敛很快。一般情形,如人脸曲面,我们用黎曼映照共形映射到平面圆盘,将共形因子作为概率密度函数,为均匀分布。如果在人脸上均匀采样个点,步长调整为,则步迭代就可以达到相对误差(每个Power胞腔的目标面积与现有面积之差处于目标面积),绝对平方误差。计算的稳定性和难度取决于区域的严格凸性,密度函数和采样点的分布。
难点分析
几何方法求解最优传输映射,理论是最为严密的,精度也是最高的,但是实现起来有很大难度,主要的难点在于下列几个方面。
拓扑结构的空间复杂度高。为了表示Weighted Delaunay三角剖分,我们需要用复杂的数据结构来表达单纯复形(simplical complex),特别是拓扑结构在优化过程中需要动态变化,这需要复杂的指针技巧,编程非常容易出错,从而造成全局崩溃。如果我们计算高维的最优传输映射,拓扑结构更加复杂,空间复杂度是单纯复形维度的指数级函数。一种应对策略是用Dart数据结构,这种结构用群论语言来表达各种拓扑操作(Euler 操作),例如群作用的轨道来表示各个维度的单纯形,因此我们可以用代数规则来保证操作的正确性,减小思维负担。
在优化过程中,我们需要动态变换Weighted Delaunay三角剖分。由盖尔芳德的理论和我们近期的研究成果Secondary Diagram【5】,所有可能的Weighted Delaunay三角剖分的个数是数量级,显然这个复杂度关于采样点的个数是指数级的,这从一个侧面反映了蒙日-安培方程的复杂性。但是,目前的算法在每次优化过程中会遇到的Weighted Delaunay三角剖分的个数是,这里是所在欧氏空间的维数,因此是多项式级别的。在每次更新高度向量之后,如果我们能够动态修改当前的三角剖分使之成为Weighted Delaunay,那么这将会极大地降低计算复杂度。在二维情形,任意两个几何三角剖分都可以从一个通过边翻转操作(edge flip)变换到另外一个。在高维情形,不同的三角剖分之间可能无法从一个通过类似翻转变换成另外一个。这方面的组合理论依然不完备,这为高维的最优传输计算带来理论上的困难。如果每一步我们都直接用凸包算法,则可以保证正确性,但是效率会降低。
另外一个内在困难在于基本几何计算的精度问题。很多计算几何算法的设计非常精巧,为了提高效率算法要求每一步都在几何上正确的,并且对于退化情形,非常依赖于底层计算的精确性。如果有一步出现错误,则整个算法崩溃。例如,Bentley-Ottman的扫描线算法,要求对于所有线段严格排序,如果中途出现小错,则后继计算会出现更多错误。这里,两条线段交点的坐标是线段端点坐标的有理多项式,多项式的次数越高,需要的计算精度越高。在实际计算中,底层几何运算非常容易超越机器精度,从而在接近退化情形给出错误的判断,由此在实际应用中Bentley-Ottman类似的算法非常脆弱。一种解决方法是用有理数域来进行计算,用软件长整数的比值来完成所有底层计算,这会极大地降低计算速度。另一种方法是用多个自适应精度计算方法。总而言之,理论上非常简单的计算几何算法,在实际中需要大量的技巧才能变得实用。对于球面最优传输映射,我们无可避免地用到超越运算,计算精度和稳定性更加关键。
最优传输几何算法最为关键的是保证中间解一定是严格凸的,高度向量一直在可容许高度空间内部。我们前面提出的阻尼算法给出了一个可靠的解法。如果我们在搜索过程中,非常接近可容许空间的边界,这时搜索路径会变得比较曲折。如何选取搜索方向和步长,如何检测可容许空间的边界,这些都是非常有意义的探索方向。
软件使用和下载
目前,最优传输映射成为学术热点,但是缺乏实实在在的高精度计算工具,抽象的基础理论很难被学生们直接理解。为了缓解人人都在谈论最优传输,却又只能纸上谈兵的尴尬局面,我们将基于几何变分的最优传输映射算法设计成课程的一次作业,并且将源码公布,以起到抛砖引玉的作用。算法用通用C++语言实现,可以在Windows,Linux平台上编译运行,图形界面基于OpenGL,线性方程组求解基于Eigen库,底层几何计算基于Shewchuk的自适应算术库。程序主要有两个类,一个用于计算Power Delaunay三角剖分,包括Lawson边对换算法和Sutherland-Hodgman算法,另一个类计算最优传输映射,包含阻尼算法和牛顿法。程序输入是三角网格,目标曲率定义在顶点上。下面的视频给出了程序的命令行,热键和运算流程:
演示程序、源代码、测试数据都可以从【1】下载,或者扫描下方二维码:
解压演示程序后,双击bat文件可以自动演示。程序的源码用CMake可以生成项目文件,再用Visual Studio等编译器进行编译。我们也提供了测试数据,格式为通用的obj文件。
因为我们提供了源代码,大家可以根据自己的学术研究目的自行修改。我们提供了强大的图形界面,优化过程中的每一个步骤都可以停顿下来进行可视化,主要的数据结构,中间计算结果也都可以被3D渲染。例如,凸包、上包络、Weighted Delaunay三角剖分,Power Diagram,裁剪前后的Power胞腔,学生们都可以通过视觉来仔细观察,从而深入理解算法的精髓。
算法的理论基础来自于丘成桐先生和郑绍远教授关于蒙日-安培方程的理论,我们提出的几何变分原理【6】,同时得到斯杭博士、罗锋教授的建议和帮助,雷娜教授进行具体指导。程序比较短小精悍,但是由于其内在的难度,开发周期却比较长。开发团队主要成员包括齐鑫、陈伟、安东生、郭洋、赵彤、温成峰、曾薇等博士、博士后。这一算法已经应用于计算机视觉领域中的曲面配准【2】,医学图像领域中的大脑研究【3】以及深度学习中的几何GAN【4】等等。
未来,我们会进一步公布和开发更加通用计算工具,例如球面,高维和一般密度函数的最优传输映射,同时也欢迎更多的同学加入到开发团队之中。希望大家发现算法漏洞,提出改进方法,更加欢迎深入的学术探讨,宝贵的建议和批评!
【1】https://www3.cs.stonybrook.edu/~gu/software/OT/index.html
【2】Zhengyu Su, Yalin Wang, Rui Shi, Wei Zeng, Jian Sun, Feng Luo and Xianfeng Gu, "Optimal Mass Transport for Shape Matching and Comparison," IEEE Transactions on Pattern Analysis and Machine Intelligence (TPAMI), 37(11), 2246-2259, 2015.
【3】Saad Nadeem, Zhengyu Su, Wei Zeng, Arie Kaufman and Xianfeng Gu, "Spherical Parameterization Balancing Angle and Area Distortions", IEEE Transaction on Visualization and Computer Graphics, 23(6), Pages:1663-1676, 2016.
【4】Na Lei, Dongsheng An, Yang Guo, Kehua Su, Shixia Liu, Zhongxuan Luo, Shing-Tung Yau and Xianfeng Gu, "Geometric Understanding of Deep Learning", Journal of Engineering, 2020.
【5】Na Lei,Zhongxuan Luo, Wei Chen, Hang Si and Xianfeng Gu, "Secondary Polytope and Secondary Diagram", 2019.
【6】Xianfeng Gu, Feng Luo, Jian Sun and Shing-Tung Yau, "Variational Principles for Minkowski Type Problems, Discrete Optimal Transport, and Discrete Monge-Ampere Equations", AJM 20(2), 383-398, 2016.
请长按下方二维码,选择 “识别图中二维码”,即可关注。
【老顾谈几何】邀请国内国际著名纯粹数学家,应用数学家,理论物理学家和计算机科学家,讲授现代拓扑和几何的理论,算法和应用。