Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow
Abstract
We present rectified flow, a surprisingly 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 special and preferred because they are the shortest paths between two points, and can be simulated exactly without time discretization and hence yield computationally efficient models.
We show that the procedure of learning a rectified flow from data, called rectification, turns an arbitrary coupling of π 0 and π 1 to a new deterministic coupling with provably nonincreasing convex transport costs.
In addition, recursively applying rectification allows us to 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 .
The full paper can be found at https://arxiv.org/abs/2209.03003 .
Introduction
Compared with supervised learning, the shared difficulty of various forms of unsupervised learning is the lack of paired input/output data with which standard regression or classification tasks can be invoked.
The gist of most unsupervised methods is to find, in one way or another, meaningful correspondences between points from two distributions.
For example, generative models such as generative adversarial networks (GAN) and variational autoencoders (VAE)
[e.g., 4, 8, 2] seek to map data points to latent codes following a simple elementary (Gaussian) distribution with which the data can be generated and manipulated.
Representation learning rests on the idea that if a sufficiently smooth function can map a structured data distribution to an elementary distribution, it can (likely) be endowed with certain semantically meaningful interpretation and useful for various downstream learning tasks.
On the other hand, domain transfer methods find mappings to transfer points from two different data distributions, both observed empirically, for the purpose of image-to-image translation, style transfer, and domain adaption [e.g., 21, 3, 17, 13].
All these tasks can be framed unifiedly as finding a transport map between two distributions:
The Transport Mapping Problem Given empirical observations of two distributions X 0 ∼ π 0 , X 1 ∼ π 1 on R d , find a transport map T : R d → R d (hopefully nice or optimal in certain sense), such that
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 , π 1 .
We present a simple approach to this transport mapping problem.
Our method, called rectified flow , learns an ODE model that transports distribution π 0 to π 1 by following straight line paths as much as possible.
The straight paths are preferred both theoretically because it is the shortest path between two end points, and computationally because it can be exactly simulated without time discretization.
Hence, flows with straight paths are simultaneously continuous-time (or infinite-step) models, and one-step models (as the case of GANs and VAEs).
Algorithmically, the rectified flow is trained with a simple and scalable unconstrained least squares optimization procedure, which avoids the instability issues of GANs, the intractable likelihood of MLE methods, and the subtle hyper-parameter decisions of denoising diffusion models.
The procedure of obtaining the rectified flow from the training data has the attractive theoretical property of 1) yielding a coupling with non-increasing transport cost jointly for all convex cost c , and 2) making the paths of flow increasingly straight and hence incurring lower error with numerical solvers.
Therefore, with a reflow procedure that iteratively trains new rectified flows with the data simulated from the previously obtained rectified flow, we obtain nearly straight flows that yield good results even with the coarsest time discretization, i.e., one Euler step.
Our method is purely ODE-based, and is both conceptually simpler and practically faster in inference time than the SDE-based approaches of [5, 16, 15].
Empirically, rectified flow can yield high-quality results for image generation when simulated with a very few number of Euler steps (see Figure 1, top row).
Moreover, with just one step of reflow, the flow becomes nearly straight and hence yield good results with a single Euler discretization step (Figure 1, the second row).
This substantially improves over the standard denoising diffusion methods.
Quantitatively, we claim a state-of-the-art result of FID (4.85) and recall (0.51) on CIFAR for one-step fast diffusion/flow models [1, 11, 18, 20, 10].
The same algorithm also achieves superb result on domain transfer tasks such as image-to-image translation (the bottom two rows of Figure 1). [Figure]
Method
Rectified flow Given empirical observations of X 0 ∼ π 0 , X 1 ∼ π 1 , the rectified flow induced from ( X 0 , X 1 ) is an ordinary differentiable model (ODE) on time t ∈ [0 , 1] ,
[Equation]
which converts Z 0 from π 0 to a Z 1 following π 1 .
The drift force v : R d → R d is set to drive the flow to follow the direction ( X 1 − X 0 ) of the linear path pointing from X 0 to X 1 as much as possible, by solving a simple least squares regression problem:
[Equation]
where X t is the linear interpolation of X 0 and X 1 .
Naviely, X t follows the ODE of d X t = ( X 1 − X 0 )d t, which is non-causal (or anticipating) as the update of X t requires the information of the final point X 1 .
By fitting the drift v with X 1 − X 0 , the rectified flow causalizes the paths of linear interpolation X t , yielding an ODE flow that can be simulated without seeing the future.
In practice, we parameterize v with a neural network or other nonlinear models and solve (1) with any off-the-shelf stochastic optimizer, such as SGD, with empirical draws of ( X 0 , X 1 ) .
See Algorithm 1.
Flows avoid crossing A key to understanding the method is the non-crossing property of flows: the different paths following a well defined ODE d Z t = v ( Z t , t )d t , whose solution exists and is unique, cannot cross each other at any time t ∈ [0 , 1) .
Specifically, there exists no location z ∈ R d and time t ∈ [0 , 1) , such that two paths go across z at time t along different directions, because otherwise the solution of the ODE would be non-unique.
On the other hand, the paths of the interpolation process X t may intersect with each other (Figure 2a), which makes it non-causal.
Hence, as shown in Figure 2b, the rectified flow rewires the individual trajectories passing through the intersection points to avoid crossing, while tracing out the same density map as the linear interpolation paths due to the optimization of (1).
Rectified coupling reduces transport costs If (1) is solved exactly, the pair ( Z 0 , Z 1 ) of the rectified flow has the following two key properties:
) ( Z 0 , Z 1 ) is guaranteed to be a valid coupling of π 0 , π 1 , that is, Z 1 follows π 1 if Z
∼
π Z 0 , Z 1 ) guarantees to yield no larger transport cost than the data pair ( X 0 , X 1 ) simultaneously for all convex cost functions c , that is,
[Equation]
The data pair ( X 0 , X 1 ) can be an arbitrary coupling of π 0 , π 1 , typically independent (i.e., ( X 0 , X 1 ) [-∼π 
-]∼ π 0 × π 1 ) as dictated by the lack of meaningfully paired observations in practical problems.
In comparison, the rectified coupling ( Z 0 , Z 1 ) has a deterministic dependency as it is constructed from an ODE model.
Hence, the rectification procedure above converts an arbitrary coupling into a deterministic coupling with lower or equal convex transport costs.
Straight line flows yield fast inference Following Algorithm 1, denote by Z =
RectFlow (( X 0 , X 1 ))
the rectified flow induced from ( X 0 , X 1 ) .
Applying this operator [-recursivelyyields a sequence of rectified flows Z k +
 = RectFlow (( Z k 0 , Z k 1 )) with ( Z 0 0 , Z 0 1 ) = ( X 0 , X 1 ) ,where Z k is the k -th rectified flow, or simply k -rectified flow , induced from ( X 0 , X 1 ) .-]recursively
Procedure : Z = RectFlow (( X 0 , X 1 )) :
Inputs : Draws from a coupling ( X 0 , X 1 ) ; velocity model v θ : R d → R d with parameter θ.
Training : ˆ θ = arg min θ
E (cid:104) ∥
X 1 − X 0 − v ( tX 1 + (1 − t ) X 0 , t ) ∥ 2 (cid:105) , with t ∼ Uniform ([0 , 1]) .
Sampling : Draw ( Z 0 , Z 1 ) from d Z t = v ˆ θ ( Z t , t )d t with Z 0 ∼ π 0 (or backwardly Z 1 ∼ π 1 ).
Return : Z = { Z t : t ∈ [0 , 1] } .
Reflow (optional): Z k +1 =
RectFlow (( Z k 0 , Z k 1 )) , starting from ( Z 00 , Z 01 ) =
( X 0 , X 1 ) .
Distill (optional): Learn a neural network ˆ T to distill the k -rectified flow, such that Z k 1 ≈ ˆ T ( Z k 0 ) .
{+yields a sequence of rectified flows Z k +
+}
=
RectFlow (( Z k 0 , Z k 1 )) with ( Z 00 , Z 01 ) =
( X 0 , X 1 ) , where Z k is the k -th rectified flow, or simply k -rectified flow , induced from ( X 0 , X 1 ) .
This reflow procedure not only decreases transport cost, but also has the important effect of straightening paths of rectified flows, that is, making the paths of the flow more straight. This is highly attractive computationally as flows with nearly straight paths incur small time-discretization error in numerical simulation. Indeed, perfectly straight paths can be simulated exactly with a single Euler step and is effectively a one-step model, addressing the very bottleneck of slow inference of ODE/SDE models.  Experiments
Unconditional Image Generation on CIFAR-10 Among one-step generative models, the distilled 2-rectified flow achieves an FID of 4 . 85 , beating the best known one-step generative model with U-net architecture, 8 . {+
+}
(TDPM, Table 1 (b)).
The recalls of both 2-rectified flow ( 0 . 50 ) and 3-rectified flow ( 0 . 51 ) outperform the best known results of GANs ( 0 . 49 from StyleGAN2+ADA) showing an advantage in diversity.
Image-to-Image Translation As shown in Figure 3, rectified flow can effectively transfer images between different domains in a unsupervised manner, without the burdensome cycle consistency regularization in CycleGAN [21].
[Figure]
[Table]
[Figure]
[Figure]
[Figure] [Figure]
[Figure] [Figure]