FFT-OT: Fast Fourier Transform method for optimal transmission mapping

Optimal Transportation has increasingly become one of the theoretical foundations of deep learning. In the statistical learning theoretical framework first proposed by Mumford and Zhu Songchun, each "concept" can be regarded as a "probability distribution" in the data space, and its support is a certain low-dimensional manifold. The core of deep learning is to learn probability distribution. Deep learning algorithms essentially perform variations in the space (Wasserstein Space) composed of all probability distributions on the manifold, optimizing specific energy functions under conditions that meet specific constraints, such as the maximum entropy principle. The optimal transmission theory defines the Riemannian metric in the Wasserstein space, thereby obtaining the covariant differential, so that the variational method can be implemented. From this perspective, optimal transmission theory explains data-driven statistical learning methods from first principles.

But in practice, optimal transmission mapping boils down to solving the Monge-Ampere Equation. Since the Monge-Ampere equation has strong nonlinearity, the computational complexity is high. Efficiently and accurately computing optimal transfer mappings is currently a hot topic in the field of deep learning. So far, most researchers are committed to improving algorithms based on Kantorovich theory, such as adding entropy regularization terms, or Nestorov smoothing methods, etc. The author and his collaborators have been improving the algorithm based on Brenier theory, essentially because this method has rich geometric connotations. Recently, we found that the classic fast Fourier transform method can be used to speed up the solution of optimal transmission mapping. More details can be found in [1].

Recently, the author's team compiled and summarized the classroom lecture notes and research reports on optimal transmission theory and algorithms in the past few years and published it in a book [2]. The new book will be released starting in mid-October 2021. I hope it can be helpful to the majority of young students, and welcome your criticisms and corrections!

Figure 1. Optimal transmission mapping calculated with FFT-OT. The Maitreya Buddha surface (upper left frame) is conformally mapped to a planar square (lower left frame), and the surface elements of the surface are pushed forward onto the square, thus defining a measure on the plane. Our algorithm computes the optimal transfer mapping from the bin measure to the Lebesgue measure (lower right frame), and the corresponding potential energy function is shown in the upper right frame.

Introduction to Optimal Transport Mapping

Let us recall the classic optimal transmission problem. Given a probability distribution sum, with a density function sum such that the total mass is equal:

Brenier's theorem states that if , then there is a convex function (Brenier potential energy function) such that the optimal transmission map is equal to the gradient of Brenier potential energy: . Substituting the Brenier potential energy function into the Jacobi equation, we get the classic Monge-Ampere equation:

Furthermore, if the boundary sum of the support set is smooth, the optimal transmission map satisfies the so-called obliqueness condition (Obliqueness condition), that is, the external normal vector at the region boundary point satisfies:

Core idea

fixed pointIn the two-dimensional case, the Brenier potential energy function is obtained from the Monge-Ampere equation,

Therefore we have

From this we get Poisson's equation

Here is the Laplace operator. We define the operator of Sobolev space:

Then the solution of the Monge-Ampere equation (Brenier potential energy function) is the fixed point of the operator.

tilt boundary conditionWe consider that and are the same plane rectangular area. According to the second type of boundary conditions and tilt conditions, for any non-corner boundary point, its image point must also be a non-corner boundary point, and at the same time, it satisfies

This means that and are both on the same side of the rectangle. From this, we can get that the image of the corner point is equal to itself.

We rewrite the Brenier potential energy function,

Then it is the fixed point of the operator,

We start by setting to , and then iterate, and we can prove that the iterative process converges to the fixed point.

At this time, the slope boundary condition of the Brenier potential energy function is converted into the Neumann boundary condition

finite difference methodAt each step in the iterative process, we solve the Poisson equation under Neumann boundary value conditions. We discretize the region into a two-dimensional lattice, and the differential operator is approximated by the difference operator:

Here and are the horizontal and vertical step sizes. In this way, Poisson's equation becomes a difference equation

Fast Fourier Transform

The Poisson equation with Neumann boundary conditions can be solved by discrete cosine transform,

Inverse transformation to

Here, other situations. Let,, can be directly obtained

That is, in the frequency domain, the solution to the Poisson equation can be written directly. The discrete cosine transform can be implemented using the fast Fourier transform, thereby further improving computational efficiency.

Summary

The core idea of FFT-OT is as follows: the Monge-Ampère equation is a strongly nonlinear degenerate elliptic partial differential equation. We can locally approximate it with a linear elliptic partial differential equation and solve it using an iterative method. Linear elliptic partial differential equations can be converted into algebraic equations in the frequency domain using Fourier transform and have analytical solutions. Fast Fourier transform can be supported by hardware, such as CUDA or FPGA, to further improve computing efficiency.

In principle, the FFT-OT method is applicable to optimal transmission mapping in any dimension, but currently high-dimensional FFT-OT needs further development. We look forward to the future adoption of this algorithm by the industry, leading to the development of dedicated libraries and chips to further promote the development of statistical deep learning.

(This document is only for pure idealists: Senior Brother Dr. Shi Wanxiong.)

1. N. Lei and X. Gu, FFT-OT: AFast Algorithm for Optimal Tansportation, ICCV 2021.

2. N. Lei and X. Gu, Optimal transmission theory and calculation, Higher Education Press, October 2021.
   

Machine-translated from the Chinese original FFT-OT: 最优传输映射的快速傅里叶变换方法 (published 2021-10-15). WeChat original ↗