This repository contains a Jupyter notebook implementing and explaining the Rectified Flow Matching algorithm for generating samples from a target distribution. The tutorial demonstrates the transformation from a simple Gaussian distribution to a complex star-shaped distribution.
Rectified Flow Matching (RFM) is a powerful technique for learning transformations between probability distributions. It builds on the idea of learning an explicit velocity field that transports samples from source to target distributions along smooth paths, as introduced in Flow Straight and Fast.
Recent developments such as Flow Matching in Latent Space and Conditional Generation and Downstream Tasks explore how these methods can be extended to more complex generative tasks. This tutorial focuses on the core concepts of RFM and the ReFlow technique.
This tutorial provides:
- Detailed implementation of RFM in PyTorch
- Visual demonstrations of the transformation process
- Implementation of the ReFlow training technique for faster sampling
For a simplified explanation of RFM concepts and real-world applications in generative AI, see SIMPLE.md.
This tutorial demonstrates the core components of RFM and ReFlow. The goal is to show how to learn a velocity field that transports samples between distributions, and how to efficiently perform sampling using both time-dependent and time-independent approaches.
The following components are covered in the notebook:
-
Distribution transformation
- Transform samples from a standard Gaussian to a star-shaped mixture of Gaussians
- Visualize the transformation process step by step
-
Model architecture
- Neural network implementation of the velocity field
- Sinusoidal position embeddings for time encoding
- Batch normalization and residual connections
-
Training methods
- Standard RFM training
- ReFlow training for single-step generation
- Visualization of both approaches
- Open the
explain_reflow.ipynbnotebook in Jupyter - Run all cells sequentially
- Observe the transformation process through generated visualizations
