Rectified Flow

Why Does This Exist?

Flow Matching solved a lot of the messy math of standard diffusion models by defining straight-line paths between pure noise and real images. However, when you train a neural network to learn those paths, a problem arises during inference. Because you randomly pair Noise Sample A with Image B during training, the network ends up learning a vector field where paths constantly cross each other.

When paths cross in a vector field, the math forces the neural network to average the intersecting vectors. This turns the theoretical "straight line" into a curved, slightly blurry trajectory in practice. Because the path is curved, you still have to use an ODE solver taking multiple small steps (e.g., 20 steps) to generate an image without veering off course.

Rectified Flow was invented to fix this crossing-path problem. By running a process called "Reflow," it iteratively straightens the learned vector field. The ultimate goal is perfectly straight, non-intersecting lines. If the path from noise to data is a perfectly straight line, you only need to evaluate the neural network once to generate a pristine, high-resolution image.

Think of It Like This

How It Actually Works

Rectified Flow is an iterative algorithm that relies on the model's own predictions to train a better version of itself.

1. Base Flow (Iteration 0)

You start by training a standard Flow Matching model. You pair random noise $z_0$ with random images $x_1$ and train the network to predict the straight line between them. Because the pairings are random, the learned vector field $v_0$ is highly curved and tangled.

2. The Reflow Process (Iteration 1)

To untangle the paths, we use the Base Flow model to generate a new training dataset. We sample a batch of pure noise $z_0$. We run the Base Flow model using an ODE solver (taking many steps) to generate the corresponding images $\hat{x}_1$.

Now, we have perfectly coupled pairs: $z_0$ and $\hat{x}_1$. We know exactly which noise particle naturally wants to become which image according to the model. We train a new Flow Matching model (Iteration 1), but this time we explicitly pair $z_0$ with $\hat{x}_1$. Because we are pairing points that already naturally connect, the paths don't cross nearly as much. The learned vector field $v_1$ is significantly straighter.

3. Iteration N (The Straight Line)

You can repeat the Reflow process indefinitely. Generate pairs using the Iteration 1 model, train an Iteration 2 model on those pairs, and so on. With each iteration, the vector field becomes straighter, closely approximating Optimal Transport (the mathematically shortest, non-intersecting path between two distributions).

4. Distillation to 1-Step

Once the vector field is perfectly straight, the ODE trajectory has no curvature. If there is no curvature, Euler's method has zero truncation error. You can evaluate the neural network exactly once at $t=0$, point in the predicted direction, take a single massive step of size $dt=1.0$, and arrive perfectly at the high-quality image $x_1$.

Show Me the Code

This code outlines the Reflow data generation and training loop concept.

1import torch2
3def generate_reflow_dataset(base_model, num_samples=1000):4    """5    Generates coupled (Noise, Image) pairs using the base model to untangle paths.6    """7    dataset = []8    for _ in range(num_samples):9        # 1. Sample pure noise10        z_0 = torch.randn(1, 3, 64, 64)11        12        # 2. Use an ODE solver (e.g., Euler with 50 steps) to follow the curved path13        # This gives us the exact image that this specific noise maps to14        x_1_generated = solve_ode(base_model, z_0, steps=50)15        16        # 3. Store the perfectly coupled pair17        dataset.append((z_0, x_1_generated))18        19    return dataset20
21def train_rectified_flow(reflow_dataset, new_model):22    """23    Trains the next iteration model on the coupled pairs.24    """25    for z_0, x_1 in reflow_dataset:26        # Sample random time t27        t = torch.rand(1)28        29        # Interpolate along the perfectly straight line between the coupled pair30        x_t = t * x_1 + (1 - t) * z_031        32        # The target direction is just x_1 - z_033        target_direction = x_1 - z_034        35        # Train new_model to predict this straight-line direction36        predicted_direction = new_model(x_t, t)37        loss = torch.nn.functional.mse_loss(predicted_direction, target_direction)38        loss.backward()39        40    return new_model

Watch Out For

The Quick Version

• Standard Flow Matching learns vector fields with crossing paths, resulting in curved trajectories that require multiple steps to traverse.
• Rectified Flow solves this by using the model to generate coupled (Noise, Image) pairs, and retraining a new model on those specific pairs.
• This "Reflow" process straightens the vector field by preventing paths from crossing.
• After 1 or 2 Reflow iterations, the trajectories become perfectly straight lines.
• Because the paths are straight, numerical ODE solvers can jump from noise to a pristine image in a single, massive step, enabling real-time generation.

What to Read Next

• Read Consistency Models to see an alternative approach to 1-step generation that forces the model to map every point on the trajectory directly to the origin.
• Read Flow Matching if you need a refresher on how the base continuous vector fields are defined.