Flow Straight and Fast: A Simple and Unified Approach to Generative Modeling, Domain Transfer, and Optimal Transport

Abstract: We consider the problem of learning a transport mapping between two distributions that are only observed through unpaired data points. This problem provides a unified framework for a variety of fundamental tasks in machine learning: generative modeling is about transforming a Gaussian (or other elementary) random variable to realistic data points; domain transfer concerns with transferring data points from one domain to another; optimal transport (OT) solves the more challenging problem of finding a "best" transport map that minimizes certain transport cost. Unfortunately, despite the unified view, there lacks an algorithm that can solve the transport mapping problem efficiently in all settings. The existing algorithms need to be developed case by case, and tend to be complicated or computationally expensive. In this talk, I will show you that the problem can be addressed with a pretty simple algorithm.

This algorithm, called rectified flow, learns an ordinary differential equation (ODE) model to transfer between the two distributions by following straight paths as much as possible. The algorithm only requires solving a sequence of nonlinear least squares optimization problems, which guarantees to yield monotonically non-increasing couplings w.r.t. all convex transport costs. The straight paths are special and preferred because they are the shortest paths between two points, and can be simulated exactly without time discretization, yielding computationally efficient models. In practice, the ODE models learned by our method generate can high-quality images with a single discretization step, which is a significant speedup over existing diffusion generative models. Moreover, with a proper modification, our method can be used to solve the optimal transport problems on high-dimensional continuous distributions, a challenging problem for which no well accepted efficient algorithms exist.