This is a minimal repo demonstrating Flow Matching — a method to train generative models via supervised learning of ODE trajectories.
Flow Matching is a recently proposed method (Lipman et al., 2022) for training generative models by learning a vector field that transforms samples from a simple distribution (e.g., Gaussian noise) into samples from a complex data distribution (e.g., images, shapes, point clouds).
We define a linear interpolation between a pair of points:
x₀is a sample from the prior distribution (e.g., Gaussian noise)x₁is a sample from the target distribution (e.g., real data)xₜis a point interpolated between the twot ∈ [0, 1]is a random time step
We want to learn a time-dependent vector field v(xₜ, t) that matches the true constant velocity from x₀ to x₁, defined as:
This defines a trivial ODE where each point moves at constant velocity along a straight line.
We train a neural network v_θ(xₜ, t) to match this target velocity using an MSE loss:
Effectively, the model learns to predict the direction and magnitude each point should move in at any timestep t along its trajectory from x₀ to x₁.
To generate new samples:
- Sample
x₀ ∼ p₀(e.g., standard normal distribution) - Solve the ODE:
from t = 0 to t = 1, using an ODE solver (e.g., Euler, Runge-Kutta).
The resulting x₁ is a sample from the learned data distribution.
.gif)
This animation shows how points are moved from the prior (noise) distribution to the data distribution by integrating the learned flow field.
This repo is based on the ideas from:
Y. Lipman, G. Batzolis, and S. Achlioptas.
"Flow Matching for Generative Modeling", NeurIPS 2022.
[arXiv:2210.02747]
This repo is a fun and lightweight implementation for understanding the mechanics of flow matching. It is not optimized for performance or scale, but focuses on clarity and visual intuition.