Flow-Based Generative Models — Example Codes

Companion code for the talk From Normalizing Flows to Flow Matching (DMA 2026). Each script in examples/ is a self-contained, single-file PyTorch implementation of one figure from the talk. The goal is pedagogical clarity over performance: small datasets (2D toy distributions), small models (MLPs), short training runs, and one figure per script.

---

What's here

Script	Method	Dataset	Estimator / objective
examples/01_fm_two_moons.py	Flow Matching (OT-CFM)	make_moons	Mini-batch OT pairing on linear interpolant
examples/02_realnvp_two_moons.py	RealNVP (coupling-layer NF)	make_moons	Maximum likelihood (exact)
examples/03_cnf_8gaussians.py	Continuous Normalizing Flow	8 Gaussians on a circle	MLE with exact trace divergence
examples/04_ffjord_8gaussians.py	FFJORD	8 Gaussians on a circle	MLE with Hutchinson stochastic trace
examples/05_fm_compare.py	OT-CFM vs Gaussian-VP-FM	8 Gaussians on a circle	Two FM objectives, side-by-side

Each script saves figures (.pdf + .png), training losses (.npy), and a model checkpoint (.pt) to figures/.

Two flavors of the FM examples: from-scratch and library-based

For the two flow-matching examples, there is also a *_meta.py variant that uses Meta's flow_matching library instead of hand-rolled interpolants and ODE solvers:

From-scratch	Library-based (flow_matching)
examples/01_fm_two_moons.py	examples/01_fm_two_moons_meta.py
examples/05_fm_compare.py	examples/05_fm_compare_meta.py

The pair is meant to be read side-by-side: the from-scratch version makes the math (linear interpolant, target velocity, Heun integrator) explicit; the library version shows the same logic refactored into the standard ProbPath / ModelWrapper / ODESolver abstractions. Same hyperparameters, same figures, different level of abstraction.

License note for the Meta variants. The flow_matching library is released under CC-BY-NC (Creative Commons Attribution-NonCommercial). Importing it from this MIT-licensed repo is fine for research, teaching, and personal use, but commercial use of the *_meta.py scripts is restricted by Meta's license. The from-scratch versions are unrestricted MIT.

---

Setup

These examples target Python 3.10+ and PyTorch 2.x. Two of the examples use zuko for normalizing-flow primitives, and one uses torchdiffeq for adaptive ODE integration.

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

To run the optional *_meta.py variants (Meta's flow_matching library):

pip install -r requirements-meta.txt

A GPU helps (especially for 04_ffjord_8gaussians.py and 05_fm_compare.py) but is not required — every script auto-falls-back to CPU. On CPU, expect each script to take a few minutes (FM two-moons, RealNVP) to ~30 minutes (FFJORD).

---

Running

From the repo root:

python examples/01_fm_two_moons.py
python examples/02_realnvp_two_moons.py
python examples/03_cnf_8gaussians.py
python examples/04_ffjord_8gaussians.py
python examples/05_fm_compare.py

The CNF and FFJORD scripts also accept --render-only to skip training and load the saved checkpoint:

python examples/03_cnf_8gaussians.py --render-only

---

Notation conventions

Across all scripts and the talk:

• $z_0 \sim p_{\text{init}} = \mathcal{N}(\mathbf{0}, I)$ — source / noise sample
• $x_1 \sim p_{\text{data}}$ — data sample
• $t \in [0, 1]$ — time / interpolation variable, with $t = 0$ noise and $t = 1$ data
• $u_\theta(t, x)$ — learned velocity field (the flow-matching network)
• $p_t$ — marginal probability path at time $t$

For RealNVP only (Part 2 of the talk), bold $\mathbf{t}$ denotes the translation vector in the affine coupling, not time.

---

Reduce training time for laptops

Each script's hyperparameters are at the top of the file (SEED, N_STEPS, BATCH, HIDDEN, ...). For a CPU-only test run, halve N_STEPS and BATCH — results will look noisier but the qualitative behavior is preserved.

---

Credits

The implementations build on patterns from:

• Lipman, Chen, Ben-Hamu, Nickel, Le, Flow Matching for Generative Modeling (ICLR 2023) — arXiv:2210.02747
• Grathwohl, Chen, Bettencourt, Sutskever, Duvenaud, FFJORD: Free-Form Continuous Dynamics for Scalable Reversible Generative Models (ICLR 2019) — arXiv:1810.01367
• Dinh, Sohl-Dickstein, Bengio, Density Estimation using Real NVP (ICLR 2017) — arXiv:1605.08803
• Chen, Rubanova, Bettencourt, Duvenaud, Neural Ordinary Differential Equations (NeurIPS 2018) — arXiv:1806.07366
• The flow_matching reference implementation by Meta AI.
• The zuko probabilistic-flow library.

---

License

MIT. See LICENSE.