久章智能JZLOG
转载 · Repost · WECHAT ARTICLE

C++ implementation of optimal transport geometric variation method

With the in-depth development of deep learning methods, people gradually realize the foundational role of optimal transmission theory. Deep learning can be abstractly summarized as learning probability distributions on manifolds. Therefore, the transformation between probability distributions and measuring the distance between probability measures are the core topics of deep learning. The optimal transmission theory provides the theoretical basis and calculation tools for solving these problems. Furthermore, optimal transmission theory defines the Riemann degree for the space composed of all probability measures.

With the in-depth development of deep learning methods, people gradually realize the foundational role of optimal transmission theory. Deep learning can be abstractly summarized as learning probability distributions on manifolds. Therefore, the transformation between probability distributions and measuring the distance between probability measures are the core topics of deep learning. The optimal transmission theory provides the theoretical basis and calculation tools for solving these problems. Furthermore, optimal transmission theory defines Riemannian metrics and covariant differentials for the space composed of all probability measures, making it possible to perform variational optimization in the space composed of all measures. Therefore, we can predict that after the deep learning craze fades, the optimal transmission theory will be condensed into the essence of thought and fully penetrate into the engineering field.

However, the partial differential equations describing the optimal transmission mapping are strongly nonlinear, so it has been challenging to accurately calculate the optimal transmission. Currently in the field of engineering, especially in the field of deep learning, people discretize probability measures, add the Kantorovich energy to the entropy regular term, and obtain a smooth approximation. The approximation error is controlled by a parameter, such as the common Sinkhorn algorithm, Nesterov approximation, etc. After extensive experiments, we found that these algorithms have many inherent flaws. If we want to accurately calculate the optimal transmission mapping between continuous distributions, these algorithms require dense sampling of the source and target regions, which is too computationally complex to be implemented with ordinary hardware. At the same time, these algorithms give approximate solutions. If we hope to improve the approximation accuracy by reducing the control parameters, because the algorithm relies on transcendental operations (exponential logarithmic operations), the stability of the algorithm quickly deteriorates and cannot converge. If we carefully adjust the parameters to ensure the stability of the algorithm, the approximate solution obtained is far from the true solution.

We have always hoped to conduct in-depth and detailed theoretical research on optimal transmission mapping, so we require high-precision numerical solutions and control of the error range. Although the Sinkhorn algorithm is popular, it cannot really meet the needs in theory. We found that few software tools currently exist for high-precision optimal transmission algorithms. In the fields of basic and applied mathematics, although everyone has been studying the Monge-Ampere equation for decades, there are no relatively popular calculation tools; in the field of computer science, the vast majority of algorithms are based on Kantorovich's linear programming and its smooth approximation (Sinkhorn). These methods have high space complexity and large approximation errors. What is more serious is that these methods destroy the theoretical structure of the optimal transmission problem itself and abandon the inherent geometric properties. The current enthusiasm for these algorithms in the engineering field largely replaces effectiveness with convenience, and regards optimal transmission as a molded tool without a deep understanding of its intrinsic theoretical structure.

Therefore, in our courses, we focus on the geometric solution method of the optimal transmission problem. This method is consistent with the weak solution theory of the Monge-Ampère equation and the convex geometry theory of Minkowski and Alexandrov. The optimal transmission theory is understood and interpreted from a geometric perspective, giving people strong geometric intuition and profound and detailed insights. We still have too many unknowns about optimal transmission mapping. We hope that the computational tools provided by geometric methods can help everyone study this theory more deeply and gain a deeper understanding of the Monge-Ampere equation, thereby further promoting the development of engineering applications.

theoretical review

Given the probability measure of the Euclidean space and its support sum, it satisfies that the total measure is equal and the density function sum is , . A map is measure-preserving if, for everything,

Record as. Given a transmission cost function, the Monge problem is to solve the minimum transmission cost among all measure-preserving mappings.

If the transmission cost is the square of the Euclidean distance, and the compactness is a convex region, then Brenier's theorem asserts that there is a convex function, the so-called Brenier potential energy function, whose gradient map gives the optimal transmission map, and this optimal transmission map is unique. The Brenier potential energy function satisfies the Monge-Ampere equation (Monge-Ampere)

and satisfy the boundary conditions. Therefore, solving the optimal transmission map is equivalent to solving the Monge-Ampere equation.

We discretely sample in, obtain the sampling point set, and use the sum of Dirac measures to approximate the target measure, thereby considering the semi-continuous optimal transmission problem. According to Brenier's theorem, the optimal transfer map is equal to the gradient map of the Brenier potential energy function. At this time, the Brenier potential energy function is a piecewise linear function,

Here, each sampling point corresponds to a support plane with an unknown height. The graph is the upper envelope of these support planes, and the graph of its Legendre duality is the convex hull of the point.

Figure 1. Brenier potential energy function and Legendre duality.

As shown in (1), the projection of the Brenier potential energy function results in a Power diagram (power diagram),

Projection of the Legendre dual of the Brenier potential energy function, resulting in the Weighted Delaunay triangulation. By adjusting the height of each corresponding support plane, we can obtain such a Power diagram. The volume of each Power cell is equal to,

Mr. Qiu Chengtong, the author and Luo Feng [6] once proved that the required height vector is the maximizer of the following energy,

within the allowable height space

This energy is strictly concave, with the maximizer inside the admissible space. By definition, we have that the gradient component of energy is equal to, and the component of the energy Hessian matrix is ​​equal to the ratio of the Power Voronoi side length to the corresponding Weighted Delaunay side length. From this, we can use Newton's method to optimize the convex energy to obtain the unique solution.

computational geometry algorithms

Figure 2. Results of the convex hull algorithm.

We see that the main concepts and algorithms required for calculations come from classical computational geometry. The main one is the Legendre duality, which requires computing the convex hull. We can use Edelsbrunner's classic incremental convex hull algorithm. This is also equivalent to computing a Weighted Delaunay triangulation, which we can do using Lawson's edge flip algorithm. If shown in Figure (2).

Figure 3. Results of the upper envelope algorithm.

After obtaining the convex hull, we calculate its corresponding upper envelope, that is, the calculated Legendre duality. Each vertex on the convex hull is dual to a plane, each edge is dual to the intersection of two vertex dual planes, and each surface is dual to the intersection of three vertex dual planes. At the same time, we add an infinite plane, thus obtaining the upper envelope of the plane family, as shown in Figure (3). Here, the boundary support plane is intercepted by the infinity plane.

Figure 4. Sutherland-Hodgman algorithm.

The projection of the upper envelope onto the top obtains the Power diagram. We need to calculate the intersection of each Power cell and , that is, use polygons to clip it. There are many algorithms for plane polygon clipping. The more efficient one is Bentley-Ottman's scan line algorithm (sweepline), and the simpler one is Sutherland-Hodgman's convex polygon clipping algorithm. As shown in Figure (4), we use a convex red polygon to clip the blue polygon. At each step, we use an edge of the red polygon to divide the plane into two, retaining the inner side and removing the outer side. Repeat this, cutting off one corner each time, until all sides of the red polygon are exhausted.

Figure 5. Brenier potential energy function, results after boundary clipping.

Figure (5) shows the Brenier potential energy function after cropping, and the corresponding Power diagram. Then, we calculate the area of ​​each Power cell, the length of each Power Diagram side, construct the gradient of the energy and the Hessian matrix, and calculate the updated height vector,

Then we perform a tentative update: select a step size, . We need to determine whether the new height vector falls within the admissible space. We first calculate the convex hull with the new height vector. If there is a height vector that does not fall on the surface of the convex hull, it means that the new height vector is illegal. Or, if all the vertices fall on the convex hull surface, but a Power cell falls outside, the new height vector is illegal. If the new height vector is not admissible, we halve the step size and try again. We call this method the damping algorithm. A few years ago, the author discussed the calculation issues of solving the Monge-Ampère equation with Mr. Qiu Chengtong and Professor Wang Xujia respectively. They both pointed out that the key to the solution process is to ensure the strict convexity of the intermediate solution, and the optimization process is always within the admissible solution space. Therefore, this step is the key to the entire algorithm.

Figure 6. Convergence process of the algorithm.

Figure (6) shows a calculation example. Since we use Newton's method to optimize the convex energy, the algorithm converges quickly. In general cases, such as the face surface, we use Riemannian mapping to conformally map it to a planar disk, and use the conformal factor as a probability density function, which is a uniform distribution. If points are uniformly sampled on the face and the step size is adjusted to, then the step iteration can achieve the relative error (the difference between the target area of ​​each Power cell and the existing area is within the target area) and the absolute squared error. The stability and difficulty of the calculation depend on the strict convexity of the region, the density function and the distribution of sampling points.

Difficulty analysis

The geometric method to solve the optimal transmission mapping has the most rigorous theory and the highest accuracy, but it is very difficult to implement. The main difficulties lie in the following aspects.

The topological structure has high space complexity. In order to represent Weighted Delaunay triangulation, we need to use complex data structures to express simple complexes. In particular, the topology needs to change dynamically during the optimization process, which requires complex pointer skills. Programming is very error-prone, causing global collapse. If we compute the optimal transport map in high dimensions, the topology is more complex and the space complexity is an exponential function of the dimensions of the simplicial complex. One coping strategy is to use Dart data structures. This structure uses group theory language to express various topological operations (Euler operations). For example, the orbit of group action represents the simplex in each dimension. Therefore, we can use algebraic rules to ensure the correctness of the operations and reduce the thinking burden.

During the optimization process, we need to dynamically transform the Weighted Delaunay triangulation. According to Gelfond's theory and our recent research results Secondary Diagram [5], the number of all possible Weighted Delaunay triangulations is orders of magnitude. Obviously this complexity is exponential with respect to the number of sampling points, which reflects the complexity of the Monge-Ampere equation from one side. However, the number of Weighted Delaunay triangulations that the current algorithm will encounter in each optimization process is, here is the dimension of the Euclidean space, so it is at the polynomial level. After each update of the height vector, if we can dynamically modify the current triangulation to become Weighted Delaunay, then this will greatly reduce the computational complexity. In 2D, any two geometric triangulations can be transformed from one to the other by an edge flip operation. In high-dimensional situations, different triangulations may not be able to be transformed from one to another through similar translations. The combination theory in this area is still incomplete, which brings theoretical difficulties to the calculation of optimal transmission in high dimensions. If we directly use the convex hull algorithm at each step, the correctness can be guaranteed, but the efficiency will be reduced.

Another inherent difficulty lies in the accuracy of basic geometric calculations. The design of many computational geometry algorithms is very sophisticated. In order to improve efficiency, the algorithm requires each step to be geometrically correct, and for degenerate cases, it is very dependent on the accuracy of the underlying calculation. If one step goes wrong, the entire algorithm collapses. For example, Bentley-Ottman's scan line algorithm requires strict ordering of all line segments. If a small error occurs midway, more errors will occur in subsequent calculations. Here, the coordinates of the intersection of the two line segments are rational polynomials of the coordinates of the endpoints of the line segments. The higher the degree of the polynomial, the higher the calculation accuracy required. In actual calculations, the underlying geometric operations can easily exceed machine accuracy, thus giving wrong judgments in near-degenerate situations. Therefore, algorithms similar to Bentley-Ottman are very fragile in practical applications. One solution is to use the rational number domain for calculations and do all low-level calculations using software long integer ratios, which greatly slows down the calculations. Another approach is to use multiple adaptive precision calculation methods. All in all, computational geometry algorithms that are very simple in theory require a lot of skill to become practical in practice. For optimal transmission mapping on the sphere, we inevitably use transcendental operations, and calculation accuracy and stability are more critical.

The most critical thing about the optimal transmission geometry algorithm is to ensure that the intermediate solution must be strictly convex, and the height vector is always within the admissible height space. The damping algorithm we proposed earlier gives a reliable solution. If we are very close to the boundary of the admissible space during the search process, the search path will become more tortuous. How to choose the search direction and step size, and how to detect the boundaries of the admissible space are all very meaningful exploration directions.

Software usage and download

At present, optimal transmission mapping has become an academic hot spot, but there is a lack of real high-precision computing tools, and the abstract basic theory is difficult for students to directly understand. In order to alleviate the embarrassing situation where everyone is talking about optimal transmission, but can only talk about it on paper, we designed the optimal transmission mapping algorithm based on geometric variation as an assignment in the course, and published the source code to serve as a starting point. The algorithm is implemented in general C++ language and can be compiled and run on Windows and Linux platforms. The graphical interface is based on OpenGL, the linear equation solution is based on the Eigen library, and the underlying geometric calculation is based on Shewchuk's adaptive arithmetic library. The program mainly has two classes, one for calculating Power Delaunay triangulation, including the Lawson edge pair conversion algorithm and the Sutherland-Hodgman algorithm, and the other class for calculating the optimal transmission mapping, including the damping algorithm and Newton's method. The program input is a triangular mesh with target curvatures defined at the vertices. The video below shows the command line, hotkeys and operation flow of the program:

The demo program, source code, and test data can be downloaded from [1], or scan the QR code below:

After unzipping the demo program, double-click the bat file to automatically demonstrate. The source code of the program can be generated with CMake and then compiled with a compiler such as Visual Studio. We also provide test data in a common obj file format.

Because we provide the source code, everyone can modify it according to their own academic research purposes. We provide a powerful graphical interface. Every step in the optimization process can be stopped for visualization. The main data structure and intermediate calculation results can also be rendered in 3D. For example, students can carefully observe the convex hull, upper envelope, Weighted Delaunay triangulation, Power Diagram, and Power cell before and after cropping through vision to deeply understand the essence of the algorithm.

The theoretical basis of the algorithm comes from the theory of the Monge-Ampère equation by Mr. Qiu Chengtong and Professor Zheng Shaoyuan, and the geometric variation principle proposed by us [6]. At the same time, we received suggestions and help from Dr. Sihang, Professor Luo Feng, and specific guidance from Professor Reina. The program is relatively short and concise, but due to its inherent difficulty, the development cycle is relatively long. The main members of the development team include Qi Xin, Chen Wei, Antonsheng An, Guo Yang, Zhao Tong, Wen Chengfeng, Zeng Wei and other doctors and postdoctoral fellows. This algorithm has been applied to surface registration in the field of computer vision [2], brain research in the field of medical images [3], and geometric GAN in deep learning [4], etc.

In the future, we will further announce and develop more general computing tools, such as optimal transmission mapping for spheres, high-dimensional and general density functions, and we also welcome more students to join the development team. I hope everyone will find loopholes in the algorithm and propose ways to improve it. In-depth academic discussions, valuable suggestions and criticisms are more welcome!

【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.

Machine-translated from the Chinese original 最优传输几何变分法的C++实现 (published 2020-11-16). WeChat original ↗