Summary

Rectified flow addresses the problem of learning a transport map between two distributions by formulating it as an ODE. Given a coupling (X₀, X₁), the method constructs linear interpolations X_t = tX₁ + (1-t)X₀ and trains a neural velocity field v(x,t) by minimizing E[‖(X₁ - X₀) - v(X_t, t)‖²]. The optimal solution v*(x,t) = E[X₁ - X₀ | X_t = x] “causalizes” the non-causal interpolation: while X_t requires both endpoints, the learned ODE dZ_t = v(Z_t, t)dt can be simulated forward from Z₀ ~ π₀ alone, with the marginal-preserving property guaranteeing Z₁ ~ π₁.

The key theoretical results are: (1) rectification provably reduces convex transport costs simultaneously for all convex cost functions via Jensen’s inequality; (2) recursive application (reflow) straightens paths at O(1/K) rate, measured by the straightness functional S(Z). Straight flows are computationally ideal — a single Euler step exactly recovers the output. The paper also shows that probability flow ODEs and DDIM are special cases of a generalized nonlinear rectified flow framework, but with curved paths that cannot be straightened by reflow, making the linear interpolation the preferred choice.

Empirically, 1-rectified flow achieves FID 2.58 on CIFAR-10 with an adaptive ODE solver, while distilled 2-rectified flow achieves FID 4.85 with a single function evaluation.

Key Contributions

  • Simple ODE framework unifying generative modeling and domain transfer via least squares regression
  • Provably non-increasing convex transport costs under rectification (Theorem 3.5)
  • Reflow procedure straightening paths at O(1/K) rate for few-step sampling
  • Unification of PF-ODEs and DDIM as special cases of nonlinear rectified flow
  • State-of-the-art one-step generation (FID 4.85 on CIFAR-10)
  • Non-crossing property of ODE flows as the mechanism for creating deterministic couplings

Methodology

Training minimizes E[‖X₁ - X₀ - v_θ(tX₁ + (1-t)X₀, t)‖²] with t ~ Uniform([0,1]). For reflow: (1) train 1-rectified flow; (2) simulate to generate pairs (Z₀¹, Z₁¹); (3) train 2-rectified flow on new pairs. Distillation with LPIPS loss further refines one-step inference. U-Net architecture from DDPM++.

Key Findings

  • 1-rectified flow: FID 2.58 on CIFAR-10 (RK45), recall 0.57 (substantial diversity improvement)
  • Distilled 2-rectified flow: FID 4.85 with single Euler step
  • Reflow dramatically improves few-step generation despite slightly worsening full-solver results
  • After one reflow, extrapolated terminal values become nearly constant, confirming straightness
  • Same algorithm achieves high-quality unpaired image-to-image translation
  • VP/sub-VP ODEs produce curved paths that cannot be straightened by reflow

Important References

  1. Score-Based Generative Modeling through Stochastic Differential Equations — PF-ODEs shown as special case
  2. Denoising Diffusion Probabilistic Models — Foundational diffusion model
  3. Connecting Brownian and Poisson Random Bridges with Rectified Flows — Establishes bridge-flow connection

Atomic Notes


paper