Learning Rectified Flow: My Understanding and Implementation

Hey everyone! Today, I want to share my recent deep dive into a fascinating paper titled "Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow". As someone learning about generative models and how they handle different data distributions, this paper offered a surprisingly elegant and effective approach. I've even been playing around with my own implementation, and I'm excited to walk you through the documentation of what I learned.

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 $π_0$) to another ($π_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.

This is where Rectified Flow comes in! It 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 $π_0$ to $π_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 ∈ [0,1]$:

$$ dZ_t/dt=v(Z_t,t) $$

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!