Flow Matching

Setup

This project uses uv.

1. Clone the repository:

git clone https://github.com/lmartinez2001/flow-matching
cd flow-matching

2. Install dependencies:

uv sync

Usage

Training the Model

Train a flow matching model on the two moons dataset:

uv run train.py

This will:

• Generate a synthetic two moons dataset
• Train an MLP to learn the vector field
• Save the trained model to out/model.pth
• Save training losses to out/losses.npy

Training parameters can be modified in train.py:

• lr: Learning rate (default: 1e-4)
• epochs: Number of training epochs (default: 1000)
• batch_size: Batch size (default: 4096)
• n_samples: Number of training samples (default: 8192)

Generating Visualizations

Create animated visualizations of the learned flow:

uv run visualize.py

This will generate:

• Vector field visualizations showing the learned flow at different time steps
• Flow evolution animations showing how samples are transported from noise to data
• GIF saved to out/animations/flow_matching.gif

References

• Paper: Flow Matching for Generative Modeling - Lipman et al., 2022
• Video: Flow Matching Explained