FLOW STRAIGHT AND FAST: LEARNING TO GENER-ATE AND TRANSFER DATA WITH RECTIFIED FLOW

Abstract

We present rectified flow, a simple approach to learning (neural) ordinary differential equation (ODE) models to transport between two empirically observed distributions π 0 and π 1 , hence providing a unified solution to generative modeling and domain transfer, among various other tasks involving distribution transport. The idea of rectified flow is to learn the ODE to follow the straight paths connecting the points drawn from π 0 and π 1 as much as possible. This is achieved by solving a straightforward nonlinear least squares optimization problem, which can be easily scaled to large models without introducing extra parameters beyond standard supervised learning. The straight paths are the shortest paths between two points, and can be simulated exactly without time discretization and hence yield computationally efficient models. We show that, by learning a rectified flow from data, we effectively turn an arbitrary coupling of π 0 and π 1 to a new deterministic coupling with provably non-increasing convex transport costs. In addition, with a "reflow" procedure that iteratively learns a new rectified flow from the data bootstrapped from the previous one, we obtain a sequence of flows with increasingly straight paths, which can be simulated accurately with coarse time discretization in the inference phase. In empirical studies, we show that rectified flow performs superbly on image generation and image-to-image translation. In particular, on image generation and translation, our method yields nearly straight flows that give high quality results even with a single Euler discretization step. Code is available at https://github.com/gnobitab/RectifiedFlow.

1. INTRODUCTION

Compared with supervised learning, the shared difficulty of various forms of unsupervised learning is the lack of paired input/output data that makes standard regression or classification tasks possible. The crux of many unsupervised methods is to find meaningful correspondences between points from two distributions. For example, generative models such as generative adversarial networks (GAN) and variational autoencoders (VAE) (e.g., Goodfellow et al., 2014; Kingma & Welling, 2013; Dinh et al., 2016) seek to map data points to latent codes following a simple elementary (e.g., Gaussian) distribution with which the data can be generated and manipulated. On the other hand, domain transfer methods find mappings to transfer points between two different data distributions, both observed empirically, for the purpose of image-to-image translation, style transfer, and domain adaption (e.g., Zhu et al., 2017; Flamary et al., 2016; Trigila & Tabak, 2016; Peyré et al., 2019) . These tasks can be framed unifiedly as finding a transport map between two distributions: Learning Transport Mapping Given empirical observations of two distributions π 0 , π 1 on R d , find a transport map T : R d → R d , which, in the infinite data limit, gives Z 1 := T (Z 0 ) ∼ π 1 when Z 0 ∼ π 0 , that is, (Z 0 , Z 1 ) is a coupling (a.k.a transport plan) of π 0 and π 1 . We should note that the answers of this problem are not unique because there are often infinitely many transport maps between two distributions. Optimal transport (OT) (e.g., Villani, 2021; Ambrosio et al., 2021; Figalli & Glaudo, 2021; Peyré et al., 2019) addresses the more challenging

