Rectified Flows & Reflow

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. Why Straight Paths Matter
5. Rectified Flow Definition
6. The Reflow Procedure
7. One-Step Distillation
8. Analysis of Straightness
9. Connection to Consistency Models
10. Practical Results
11. Common Pitfalls
12. Summary
13. Exercises

Learning Objectives

1. Explain why straight trajectories enable few-step generation.
2. Derive the rectified flow objective and its geometric interpretation.
3. Describe the reflow procedure for iterative trajectory straightening.
4. Connect rectified flows to consistency distillation.
5. Analyze the quality-steps Pareto frontier.

Notation

• $(\mathbf{X}_0, \mathbf{X}_1) \sim \gamma$ — coupling between noise and data
• $\mathbf{X}_t = (1-t)\mathbf{X}_0 + t\mathbf{X}_1$ — linear interpolation
• $v_\theta$ — learned velocity field
• $\text{Straightness}(v)$ — trajectory straightness metric

Core Intuition

A flow that follows straight lines from noise to data can be integrated in ONE Euler step: $\mathbf{x}_1 = \mathbf{x}_0 + v_\theta(\mathbf{x}_0, 0)$. But with random noise-data pairing, trajectories cross and the learned velocity must be curved to avoid collisions. Rectified flow iteratively improves the coupling $(\mathbf{X}_0, \mathbf{X}_1)$ so that trajectories become progressively straighter, approaching one-step generation.

Why Straight Paths Matter

Euler discretization error: For an ODE $dx/dt = v(x, t)$ with step size $h$:

For straight paths: $v(x, t) = \text{const}$ along each trajectory → $v'' = 0$ → zero discretization error. One Euler step is exact.

For curved paths: Need many steps to follow the curvature. A path with curvature $\kappa$ requires $O(\kappa / \epsilon)$ steps for error $\epsilon$.

Goal: Make the velocity field as constant as possible along trajectories.

Rectified Flow Definition

Given coupling $\gamma$ over $(\mathbf{X}_0, \mathbf{X}_1)$:

Interpretation: Train the velocity field to predict the direction from noise to data along the linear interpolation. The target $\mathbf{X}_1 - \mathbf{X}_0$ is the straight-line velocity.

Initial coupling (Rectified Flow 1): $\gamma_0 = p_0 \otimes p_1$ (independent). Random pairing → crossing paths → curved learned velocity.

The Reflow Procedure

Key idea: Use the trained model to create a BETTER coupling, then retrain.

Algorithm:

1. RF-1: Train $v_{\theta_1}$ on random coupling $\gamma_0 = p_0 \otimes p_1$.
2. Generate pairs: For each $\mathbf{x}_0 \sim p_0$, solve ODE to get $\hat{\mathbf{x}}_1 = \text{ODE}(v_{\theta_1}, \mathbf{x}_0, 0 \to 1)$.
3. New coupling: $\gamma_1 = \{(\mathbf{x}_0, \hat{\mathbf{x}}_1)\}$ — now the pairs are connected by the flow.
4. RF-2: Train $v_{\theta_2}$ on coupling $\gamma_1$.
5. Repeat: $\gamma_2 \to \text{RF-3}$, etc.

Why it helps: After RF-1, the coupling $\gamma_1$ has non-crossing paths (they were generated by an ODE). Training on non-crossing paths yields straighter velocity fields.

Convergence: 2-3 iterations typically sufficient. Diminishing returns after that.

One-Step Distillation

After reflow, the flow is nearly straight. Distill to exactly one step:

where $\hat{\mathbf{x}}_1$ is generated by the multi-step teacher.

Alternative: self-consistency loss. Enforce that the one-step prediction equals the multi-step result:

Analysis of Straightness

Straightness metric: For a trajectory $\{\mathbf{x}_t\}_{t=0}^1$:

• Straightness = 1: perfectly straight (constant-speed straight line).
• Straightness less than 1: curved path (longer than straight-line distance).

After each reflow:

• RF-1: Straightness $\approx 0.7$–$0.8$ (significant curvature).
• RF-2: Straightness $\approx 0.9$–$0.95$.
• RF-3: Straightness $\approx 0.95$–$0.99$.

Connection to Consistency Models

Consistency models (Song et al., 2023): Train a model $f_\theta(\mathbf{x}_t, t)$ that maps ANY point on a trajectory to the endpoint:

Connection: A perfectly straight rectified flow satisfies:

So the one-step prediction from ANY time $t$ equals $\mathbf{x}_1$. Straight rectified flows ARE consistency models.

Difference: Consistency models enforce this property directly; rectified flows achieve it by straightening trajectories.

Practical Results

Image generation (CIFAR-10, ImageNet):

• 1-step rectified flow (after reflow): FID $\approx$ 4-6.
• 2-step: FID $\approx$ 2-3.
• 50-step diffusion: FID $\approx$ 2.

Used in production:

• Stable Diffusion 3: Flow matching + reflow-style distillation.
• SDXL Turbo: Adversarial distillation of flow to 1-4 steps.

Speed: 1-step generation is 25-50x faster than 50-step diffusion. Real-time generation on consumer GPUs.

Common Pitfalls

Pitfall 1. Expecting RF-1 to produce one-step quality. The first rectified flow still has significant curvature; reflow is essential for few-step generation.

Pitfall 2. Generating too few coupling pairs for reflow. Need diverse pairs (100K+) covering the full distribution. Undercoverage leads to mode collapse.

Pitfall 3. Using high-order ODE solvers during reflow generation. This is correct but expensive. The whole point of reflow is to eventually avoid ODE solving — but the reflow process itself requires it.

Summary

• Straight paths → zero Euler discretization error → one-step generation.
• Rectified flow: Linear interpolation objective with noise-data coupling.
• Reflow: Iteratively improve coupling using trained model; 2-3 iterations for near-straight paths.
• Straightness converges: 0.7 → 0.9 → 0.95+ across reflow iterations.
• Equivalent to consistency models when perfectly straight.
• Enables real-time generation (1-4 steps) with near-diffusion quality.

Exercises

Exercise 1. Prove that non-crossing paths must exist for an ODE-generated coupling (hint: uniqueness of ODE solutions).

Exercise 2. Compute the Euler discretization error for a circular trajectory vs a straight trajectory with the same endpoints.

Exercise 3. For a 2D Gaussian mixture: visualize the trajectories of RF-1 vs RF-2 and measure straightness.

Exercise 4. Derive the relationship between the number of Euler steps and FID for a flow with straightness $S$ (estimate error scaling).

Exercise 5. Compare the training cost of 3 rounds of reflow vs consistency distillation for achieving 1-step generation.