Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow

Xingchao Liu; Chengyue Gong; Qiang Liu

Source-linked AI summary

Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow

Xingchao Liu, Chengyue Gong, Qiang Liu

arXiv:2209.03003v1Sep 7, 2022cs.LG

TL;DR

Unpaired distribution transport lacks the paired data required by standard supervised learning. Rectified flow causalizes straight interpolation paths with an ODE, producing deterministic couplings with no larger convex transport costs and enabling one-step simulation.

Problem
Unsupervised distribution transport lacks paired input/output data for directly applying standard supervised learning.
Method
Rectified flow fits an ODE to causalize prescribed straight interpolation paths between two distributions.
Results
Rectification converts arbitrary distribution couplings into deterministic couplings with no larger transport cost for every convex cost function.
Takeaways & Limitations
Straight rectified flows can support fast inference because they can be simulated with one Euler step.
Takeaways & Limitations
Exact recovery of finite empirical target points requires known conditional densities and can completely overfit, making the procedure impractical.

Abstract

from arXiv · show

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 non-increasing 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, image-to-image translation, and domain adaptation. 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.

1 Introduction

The introduction frames generative modeling and domain transfer as transport-map problems between empirical distributions, then presents rectified flow as a unified ODE approach that favors straight trajectories for theoretical and computational reasons. It emphasizes simple least-squares training, reduced inference cost, convex transport-cost guarantees, and strong image-generation results with few Euler steps.

Introduction: Rectified flow uses a simple, scalable unconstrained least-squares objective instead of GAN minimax training, intractable MLE, or diffusion hyper-parameter selection.The introduction motivates this design by citing GAN instability and mode collapse, alongside high design complexity in denoising diffusion methods.
Introduction: Rectified flow unifies these tasks with an ODE that follows straight-line paths between endpoint distributions as closely as possible.Straight paths are shortest between endpoints and can be simulated exactly without time discretization.
Introduction: Rectification yields deterministic couplings with jointly non-increasing transport costs for all convex cost functions and progressively straighter flows under recursive application.These properties support accurate simulation with coarse time discretization during inference.
Introduction: FID 4.85 and recall 0.51 are reported on CIFAR10, while one reflow step makes image-generation trajectories nearly straight and effective with one Euler step.The passage describes this as a state-of-the-art CIFAR10 result and a substantial improvement over standard denoising diffusion methods.

2 Method

Rectified flow causalizes linear interpolation by fitting an ODE drift to endpoint displacements, while preserving intermediate marginals and producing valid endpoint couplings. Recursive rectification straightens trajectories and weakly reduces every convex transport cost, enabling accurate coarse-step simulation.

Core construction: Rectified flow fits the drift to X1 − X0, causalizing the linear interpolation into an ODE that does not require the future endpoint.The original interpolation update is anticipating because it directly uses X1.
Core construction: The resulting flow preserves Law(Zt) = Law(Xt) for every t ∈[0, 1], so (Z0, Z1) remains a coupling of π0 and π1.This preserves the input process’s marginals while changing its joint trajectory distribution.
Theoretical properties: For every convex cost c, E[c(Z1 − Z0)] ≤ E[c(X1 − X0)], regardless of the input coupling.Rectification therefore descends simultaneously across the collection of convex transport costs rather than optimizing one selected cost.
Recursive rectification: Recursive rectification makes flow paths increasingly straight, decreasing the straightness measure S(Z) toward zero and reducing numerical simulation error.Perfectly straight paths can be simulated exactly with one Euler step.
Recursive rectification: One reflow step can already produce nearly straight flows with good single-Euler-step performance, whereas too many steps may accumulate estimation error in vX.The method balances trajectory straightening against errors introduced during drift estimation.
Generalized rectified flows: Generalized rectified flows replace linear interpolation with any time-differentiable connecting curve, preserving marginals but losing guaranteed convex-cost reduction and straightening.Probability flows and DDIM can be viewed approximately as special cases of this generalized framework.

3 Theoretical Analysis

The theoretical analysis establishes marginal preservation, convex transport-cost reduction, and progressive straightening under rectification. It also characterizes straight and c-optimal couplings and identifies PF-ODEs as nonlinear rectified flows.

3.1 Marginal preservation: Nonlinear rectified flows preserve the marginal law Law(Z_t) = Law(X_t) for every t, regardless of interpolation straightness.This holds when X is rectifiable and Z is its induced rectified flow.
3.2 Convex transport costs: With canonical linear interpolation, rectification yields a coupling with no larger expected cost for any convex function c.The reduction follows from the shared marginals of X_t and Z_t and Jensen’s inequality.
3.3 Reflow: A coupling is straight exactly when it is a Rectify(·) fixed point, its rectified flow equals the linear interpolation, and interpolation paths do not intersect.Straight couplings are desirable because their flows can be simulated exactly with one numerical-solver step.
3.3 Reflow: Recursive rectification makes flows increasingly straight, with the deviation from straightness converging at rate O(1/K).The result assumes each recursively generated coupling is rectifiable.
3.4 Straight and c-optimal couplings: For strictly convex costs, c-optimality implies straightness, but in dimensions d ≥2 straightness does not generally guarantee optimality for a specific cost.In one dimension, a straight coupling is deterministic and monotonic; when it exists, it is unique and jointly optimal for all applicable convex costs.
3.5 PF-ODEs: PF-ODEs are instances of nonlinear rectified flows with interpolation X_t = α_tX_1 + β_tξ.The equivalence follows under the stated assumptions in Proposition 3.11.

4 Related Works and Discussion

This section positions rectified flow as a simple ODE-based alternative to likelihood-based, diffusion-derived, and adversarial generative methods. It argues that pre-specified straight paths enable efficient one-step inference while avoiding several computational and theoretical drawbacks of competing approaches.

Related generative models: Reflow plus distillation offers a promising route to one-step models without GAN minimax instability or likelihood-based intractability.The approach is presented as an alternative to classical GAN, VAE, and normalizing-flow methods.
Learning ODEs: MLE-trained neural ODEs require repeated ODE simulation and backpropagation through time, making large-scale training computationally expensive and potentially unstable.MLE also matches only the final outcome distribution, leaving infinitely many possible intermediate trajectories.
Learning ODEs: Rectified flow avoids these MLE drawbacks by pre-specifying an interpolation curve, or “road,” for the ODE to follow rather than inferring paths implicitly.The framework permits any interpolation when the neural drift is a universal approximator.
ODEs versus SDEs: By removing diffusion noise, rectified flow shows that straight paths can guide ODE learning and may inform the design of bridge processes for SDEs.The paper argues that diffusion mechanisms should be decoupled from the other merits of denoising diffusion models.
Optimal versus straight transport: Reflow seeks straight rather than cost-optimal couplings, because all straight couplings produce straight rectified flows that can be simulated with one Euler step.Optimal couplings are computationally challenging to find, and transport cost need not align with individual-task performance.

5 Experiments

Experiments show that reflow progressively straightens rectified flows and enables strong image generation, translation, and domain adaptation with few or even one Euler step. The method achieves state-of-the-art results while using simple training and avoiding adversarial or cycle-consistency optimization for translation.

Experimental procedure: Repeatedly fitting rectified flows on simulated endpoint pairs produces k-rectified flows, which can then be distilled into a one-step model.The procedure begins with independent draws from π_0 and π_1, repeats rectification k times, and fits z_1 = z_0 + v(z_0, 0) for distillation.
Image generation: On CIFAR-10, a fully solved 1-rectified flow achieves the lowest FID (2.58) and highest recall (0.57) among ODE-based methods.It uses the DDPM++ architecture and the RK45 solver, while requiring fewer generation steps than VE, VP, and sub-VP ODEs.
Image generation: Reflow improves FID and recall in the small-step regime, with each reflow producing a noticeable improvement and distillation further boosting one-step generation.The improvement is reported for N⪅80, whereas performance worsens in the large-step regime because of accumulated estimation error.
Image generation: The distilled 2-rectified flow reaches an FID of 4.85 versus 8.91 for TDPM, while 2-rectified flow (0.50) and 3-rectified flow (0.51) exceed StyleGAN2+ADA recall (0.49).These results indicate improved one-step generation quality and diversity relative to the cited U-net generative and GAN baselines.
Straightness and high-resolution generation: Reflow straightens trajectories, and 2-rectified flow produces almost straight paths; 1-rectified flow yields recognizable images at t ≈0.1 versus t ≈0.6 for sub-VP ODE.A 1-rectified flow also generates high-quality 256 × 256 images across LSUN Bedroom, LSUN Church, CelebA HQ, and AFHQ Cat with one or few Euler steps.
Translation and domain adaptation: Rectified flows produce high-quality, diverse image translations and achieve state-of-the-art domain-adaptation performance on DomainNet and OfficeHome.Translation requires neither adversarial optimization nor cycle-consistency regularization, and 2-rectified flow gives good results with a single Euler step (N = 1).

A Additional Experiment Details

The experiments specify training configurations for CIFAR10 generation, image-to-image translation, and domain adaptation, and provide qualitative evidence that rectification improves few-step generation, interpolation, and reconstruction. In particular, 2-rectified flow produces clear images and recognizable latent-space results with very few solver steps.

CIFAR10 Generation: CIFAR10 generation uses 32 × 32 images, a DDPM++ network, EMA smoothing at ratio 0.999999, Adam with learning rate 2e−4, and dropout 0.15.The model configuration follows DDPM++ while applying exponential moving average during training.
CIFAR10 Generation: Reflow generates 4 million (z0, z1) pairs and fine-tunes each rectified-flow model for 300,000 steps before few-step distillation.For k-step generators, fine-tuning samples t from the discrete grid {0, 1/k, · · ·, (k−1)/k}; for k = 1, the L2 loss is replaced with LPIPS similarity.
Few-Step Generation: After one rectification, 2-rectified flow generates clear images in 1, 2, or 3 steps, whereas 1-rectified flow remains blurry at those step counts.The comparison includes VE, VP, and sub-VP ODEs.
Image-to-Image Translation: For image-to-image translation, the experiments use AFHQ, MetFace, and CelebA-HQ with DDPM++ U-Nets, AdamW, batch size 4, 1,000 epochs, weight decay 0.1, and dropout 0.1.AFHQ contains 15,000 512 × 512 animal-face images, MetFace contains 1,336 1024 × 1024 art-derived human-face images, and CelebA-HQ contains 30,000 1024 × 1024 images.
Latent Space and Reconstruction: Latent-space visualizations show smooth interpolation for 1- and 2-rectified flow, while 2-rectified flow remains visually recognizable with one-step sampling and supports meaningful one-step reconstruction.Sample trajectories also indicate that the extrapolation of 2-rectified flow is nearly independent of t, consistent with an almost straight trajectory.