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.