Linear-Time Transport with Rectified Flows

Abstract

Matching probability distributions allows to compare or interpolate them, or model their manifold. Optimal transport is a tool that solves this matching problem. However, despite the development of numerous exact and approximate algorithms, these approaches remain too slow for large datasets due to the inherent challenge of optimizing transport plans. Taking intuitions from recent advances in rectified flows we propose an algorithm that, while not resulting in optimal transport plans, produces transport plans from uniform densities to densities stored on grids that resemble the optimal ones in practice. Our algorithm has linear-time complexity with respect to the problem size and is embarrassingly parallel. It is also trivial to implement, essentially computing three summed-area tables and advecting particles with velocities easily computed from these tables using simple arithmetic. This already allows for applications such as stippling and area-preserving mesh parameterization. Combined with linearized transport ideas, we further extend our approach to match two non-uniform distributions. This allows for wider applications such as shape interpolation or barycenters, matching the quality of more complex optimal or approximate transport solvers while resulting in orders of magnitude speedups. We illustrate our applications in 2D and 3D.

Our transport algorithm computes shape interpolations and image stipplings in linear time, with results approaching those of optimal transport. In 2D, the interpolation is performed on a 256x256 grid with 65,536 samples and timings correspond to computing 5 interpolation steps; stippling is performed on a 1024x1024 grid with 8,192 samples. We compare our results to a GPU-optimized linearized transport with Sinkhorn divergences [Feydy et al. 2019b], the Instant Transport of Nader and Guennebaud [2018], the high-quality stippling of BNOT [de Goes et al. 2012], and that obtained using Instant Transport (timings for Instant OT stippling include 23.5s of precomputations that only depend on grid resolution). In our 3D example, our process takes 815s to compute transport maps on a 2563 resolution with 16.8 millions of particles, and each interpolation step then requires 3s (splatting particles and mesh reconstruction).

We compute barycenters for four voxelized 3d meshes on $256^3$ grids (advecting $256^3$ particles in 1936.31s, splatting them on voxels, and marching cubes reconstruction –3.6s per interpolation step–).