Variational & Flow Models advanced 7 min read 12 flashcards

Rectified Flow and Trajectory Straightening

Why the generative ODE learned from independent noise-data pairs is curved, how reflow iteratively straightens it, and what each round of straightening costs in fidelity.

If the trajectory from noise to image were a straight line, one Euler step would land exactly on the image and sampling would cost a single forward pass. Nothing in the training objective forces it to be straight, and in general it is not. Rectified flow (Liu et al., 2023, Flow Straight and Fast, arXiv:2209.03003) is a procedure for making it straighter, and it is the mechanism behind several few-step image generators.

Why the field curves

Flow matching trains on straight conditional paths \(x_t = (1-t)x_0 + t x_1\) with \(x_0\) and \(x_1\) drawn independently. Those individual segments are perfectly straight. The problem is that they cross: two different noise-data pairs can pass through the same point \(x_t\) at the same time with different velocities. The network, unable to distinguish them, learns the average velocity at that point, and the resulting marginal field bends around every crossing region.

A curved trajectory cannot be integrated accurately in one step. This is the entire reason a naively trained flow-matching model still needs 20 to 50 sampling steps despite being trained on straight lines.

Reflow

The fix is to remove the crossings by changing the coupling. Take the trained model, generate a large set of pairs by actually running the ODE, so that each noise sample is paired with the image the model deterministically produces from it, and retrain flow matching on those pairs instead of on independent ones. This new coupling is a transport map, not a random pairing, so its straight segments cross far less.

The key property is that reflow preserves the marginal distributions while reducing transport cost. The set of images produced does not change in distribution; only the paths to them get straighter. Repeating the procedure straightens further, and each round is called a rectification.

\[\text{1-rectified flow} \xrightarrow{\text{generate pairs, retrain}} \text{2-rectified flow} \xrightarrow{\ \cdots\ } \text{k-rectified flow}\]

After one or two rounds the trajectory is straight enough that a small number of Euler steps, sometimes one, produces a usable image. A final distillation step, training the model to match the multi-step output in a single step, closes the remaining gap.

The cost of each round

Reflow is not free and it is not lossless. Each round trains on data generated by the previous model, so any error in that model becomes ground truth for the next. Fine detail degrades, diversity narrows slightly, and the characteristic artefacts of the parent model are inherited and amplified. Two rounds is a common practical limit for this reason; beyond that the fidelity loss usually outweighs the step-count gain.

The compute cost is also real. Generating a reflow dataset means running full multi-step sampling over hundreds of thousands or millions of noise samples, which is comparable to a substantial fraction of the original training budget.

When it breaks

Straightness and quality are not the same axis. A perfectly straight field can transport noise to a distribution that is slightly wrong. Measuring straightness, usually as the deviation of the velocity along a trajectory from its average, tells you how few steps you can get away with, not how good the samples are. Both need to be tracked.

Reflow amplifies whatever the teacher got wrong. Because the training pairs come from the model itself, systematic biases such as oversaturation from high guidance, or a mode the model over-represents, are baked in as targets. Generating the reflow dataset at the guidance scale you intend to deploy at, rather than at a higher one, matters more than it sounds.

One-step models lose the ability to trade compute for quality. A multi-step model gives users a dial: more steps, better image. A distilled one-step model has no such dial, and improving quality means retraining. For interactive products that is often the right trade; for anything where a user may want a high-effort render, it removes an option that costs nothing to keep.

Straightening interacts with conditioning. The optimal transport map from noise to images depends on the condition, so a field straightened for one prompt distribution is not straightened for another. Models rectified on a narrow prompt distribution generalise worse in step count than in quality, which is a subtle failure: the images stay acceptable while the required step count quietly rises.

Check yourself

12 flashcards for this concept

Click a card to reveal the answer.

Drill the whole track