The Challenge of Distribution Transport

In machine learning, we often encounter the problem of “distribution transport.” This involves learning how to transform data from one distribution (let's call it \(\pi_0\)) to another (\(\pi_1\)). This challenge is central to many tasks, including:

• Generative Modeling: Creating new data samples that resemble a target distribution (e.g., generating realistic images from random noise).
• Domain Transfer: Translating data from one domain to another (e.g., converting horse images to zebra images).

Traditionally, methods like Generative Adversarial Networks (GANs) and Diffusion Models have been popular. While powerful, they can sometimes face issues like training instability, mode collapse, or high computational costs during inference, especially for continuous-time models that require many steps to simulate.

Rectified Flow offers a fresh perspective by leveraging Ordinary Differential Equation (ODE) models to tackle this problem in a unified and often more efficient way.

What is Rectified Flow? The Idea of “Straightening” Paths

The core idea behind Rectified Flow is beautifully simple: learn an ODE that transports data points from \(\pi_0\) to \(\pi_1\) by following paths that are as straight as possible.

Think of it like this: if you have a starting point and an ending point, the shortest and most direct way to get from one to the other is a straight line. Rectified Flow aims to make the “flow” of data points between distributions mimic these straight-line movements.

Why are straight paths so desirable?

• Computational Efficiency: Perfectly straight paths can be simulated exactly with just a single step, making inference much faster compared to multi-step ODE/SDE solvers. This bridges the gap between computationally intensive continuous-time models and faster one-step models.
• Theoretical Soundness: Straight paths are the shortest paths, which has nice implications for transport costs (more on this later!).

The Rectified Flow defines an ODE on time \(t \in [0,1]\):

where \(Z_t\) represents a data point at time \(t\), and \(v(Z_t,t)\) is a “velocity field” that dictates the direction and speed of movement. Our goal is to learn this \(v\).

Learning the Velocity Field: A Least Squares Approach

The magic of Rectified Flow lies in how it learns this velocity field \(v\). It's surprisingly straightforward!

First, let's consider a simple linear interpolation between a source point \(X_0 \sim \pi_0\) and a target point \(X_1 \sim \pi_1\):

where \(t \in [0,1]\). The derivative of this linear interpolation is simply \((X_1-X_0)\), which represents the straight-line direction from \(X_0\) to \(X_1\).

However, this linear interpolation is “non-causal” because to know \(X_t\), you need to know the final point \(X_1\). The Rectified Flow aims to “causalize” this by learning a velocity field \(v(X_t,t)\) that closely approximates \((X_1-X_0)\). This is achieved by solving a simple least squares regression problem:

In practice, this means we train a neural network to represent \(v\). We feed the network the interpolated point \(X_t\) and the time \(t\), and it tries to predict the “straight” direction \((X_1-X_0)\).

Here's an illustration of the concept:

The Rectification Process: Training and Sampling

The overall procedure for Rectified Flow is outlined in the paper's Algorithm 1:

1. Training

Inputs: Draws from a coupling \((X_0,X_1)\) of your source distribution \(\pi_0\) and target distribution \(\pi_1\). Often, we just draw \(X_0 \sim \pi_0\) and \(X_1 \sim \pi_1\) independently.

Objective: Train your velocity model \(v_\theta\) by minimizing the least squares objective mentioned above. This typically involves sampling \(t\) uniformly from \([0,1]\) during training.

2. Sampling (Inference)

Once the velocity model \(v_{\hat{\theta}}\) is trained, you can use it to generate samples:

1. Draw a starting point \(Z_0 \sim \pi_0\).
2. Solve the ODE \(dZ_t=v_{\hat{\theta}}(Z_t,t)dt\) from \(t=0\) to \(t=1\).
3. The final point \(Z_1\) will then follow the target distribution \(\pi_1\).

For backward sampling (transferring \(\pi_1\) to \(\pi_0\)), you simply solve \(d\tilde{X}_t=-v(\tilde{X}_t,t)dt\) starting from \(\tilde{X}_0 \sim \pi_1\).

Key Properties and Why They Matter

Rectified Flow comes with several attractive properties that make it a powerful tool:

1. Flows Avoid Crossing

A fundamental property of well-defined ODEs is that their trajectories cannot cross each other. While the linear interpolation \(X_t\) can have crossing paths, the rectified flow \(Z_t\) “rewires” these trajectories to avoid intersections while still preserving the overall marginal distributions. This means that the flow maintains a consistent mapping between points from \(\pi_0\) to \(\pi_1\).

2. Reducing Transport Costs

The rectification process is guaranteed to reduce or maintain the transport cost for any convex cost function. In simpler terms, it finds a more “efficient” way to move mass from \(\pi_0\) to \(\pi_1\). For instance, for the common L2 cost, the expected length of the paths \(\|Z_1-Z_0\|\) will be less than or equal to \(\|X_1-X_0\|\).

3. Reflow: Making Paths Even Straighter

This is where things get really interesting! The paper introduces an optional but highly effective procedure called Reflow. This involves recursively applying the rectification process:

1. You start with your initial data coupling \((X_0,X_1)\).
2. You learn the first rectified flow, \(Z^1=\operatorname{RectFlow}((X_0,X_1))\).
3. Then, you use the end points of this first flow, \((Z_0^1,Z_1^1)\), as the new input for the next rectification step to learn \(Z^2=\operatorname{RectFlow}((Z_0^1,Z_1^1))\).

This process can be repeated for \(k\) steps to obtain a \(k\)-rectified flow, \(Z^k\).

The magical effect of this reflow procedure is that it makes the paths of the rectified flows increasingly straight. The straightness of a flow \(Z=\{Z_t\}\) can be measured by:

A value of \(S(Z)=0\) indicates perfect straightness. The paper theoretically proves that applying rectification recursively decreases \(S(Z)\) towards zero.

This straightening is incredibly beneficial: a nearly straight flow can be simulated accurately with a very small number of numerical integration steps, even just a single Euler step! This drastically reduces inference time.

4. Distillation (Optional)

After obtaining a highly straight \(k\)-rectified flow, you can further speed up inference by distilling the flow into a one-step neural network that directly maps \(Z_0^k\) to \(Z_1^k\). This is useful for fine-tuning the model for even faster one-step generation.

My Implementation Insights

For my implementation, I focused on a simple 2D toy example to visualize the concepts.

Dataset:

I used two distinct distributions:

• \(\pi_0\): Two Gaussian blobs (source distribution)
• \(\pi_1\): Two “square” shapes (target distribution), which are intentionally placed to encourage curved paths in naive interpolation, making the straightening effect more evident.

The goal was to see how the Rectified Flow learns to smoothly transform points from the Gaussian blobs to the squares.

The training loop iteratively updates the VectorField network based on the least squares objective. After training, I visualized the trajectories generated by simulating the learned ODE.

Here's what I observed, consistent with the paper's findings:

These visualizations clearly demonstrate the “straightening” effect that Rectified Flow achieves through its training and reflow procedure.

Concluding Thoughts

Rectified Flow presents a really unique way to tackle distribution transport. We can apply these techniques to image and audio data as well. Recent papers like Diff2Flow and Mean Flows hint at a vibrant future for this area. I'm excited to explore its potential further and see where it leads!

Resources

Here’s where I learned it from, for you to explore:

• Original Paper
• Blog by Dr. Scott Hawley
• MIT Computer Science Class 6.S184: Generative AI with Stochastic Differential Equations
• Flow Matching | Explanation + PyTorch Implementation