Lecture 2: Flow Matching (Chapter 2)

A velocity field assigns one vector to every state and every time. Following it from a starting point traces a trajectory, the trajectories from every starting point form the flow map, and the continuity equation ties the motion of the density to the same field. Training such a field by maximum likelihood costs an ODE solve and a divergence trace per gradient step, so flow matching replaces it with a regression: pair a Gaussian sample with a data sample, join the pair by a straight line, and regress a network onto the derivative of that line. Least squares recovers the average velocity of all the paths through a state, which is the field that transports the prior onto the data, and generation is Euler integration of the learned field.

flow_matching_mnist.py is the lecture's main script, the one Section 2.20 of the notes walks through line by line. The other three files are the mathematics underneath it: the fields and the solver before any training, the full method on a target whose density is known exactly, and the family of interpolants the straight line belongs to.

File	Command	What it does
`flow_matching_mnist.py`	`python lecture_2/flow_matching_mnist.py`	Trains the small U-Net on MNIST and writes a 4-by-4 grid of digits and a trajectory strip. `--dataset toy` runs offline.
`velocity_fields.py`	`python lecture_2/velocity_fields.py`	No training. Trajectories, the flow map, the Euler refinement ladder, the pushforward of a Gaussian, the flux balance and the continuity equation.
`flow_matching_2d.py`	`python lecture_2/flow_matching_2d.py`	Trains a velocity MLP on the four-Gaussian target and samples by Euler. Measures the conditional flow-matching equivalence against the exact marginal velocity.
`interpolants.py`	`python lecture_2/interpolants.py`	Trains the same MLP three times, on the linear, the t-squared and the trigonometric schedule, and compares the samples at a fine and a coarse step budget.

Nothing is trained in velocity_fields.py, so its --steps sets the largest rung of the Euler refinement ladder. In the other three files --steps counts optimizer steps, and in flow_matching_mnist.py it is the budget for the whole run, which --epochs would otherwise set.

The numbers a default run produces

velocity_fields.py, in about two seconds, reproduces the chapter's arithmetic exactly.

euler steps   1  landing 2.000000  error 0.718282  error x steps 0.7183
euler steps   2  landing 2.250000  error 0.468282  error x steps 0.9366
euler steps   4  landing 2.441406  error 0.276876  error x steps 1.1075
flux_balance  -0.3354      enclosed_mass_rate  -0.3354
continuity_time_derivative  0.1123      continuity_minus_flux_divergence  0.1123
density_ratio_at_origin  2.0000         log_density_change  -0.6931

flow_matching_2d.py, in about six seconds on one core, with 83,202 parameters.

crossing_conditional_loss  [10.0, 9.0, 10.0]     crossing_marginal_loss  [1.0, 0.0, 1.0]
crossing_constant_gap      [9.0, 9.0, 9.0]       crossing_target_variance  9.0000
loss_cfm  10.3243   loss_fm  5.1196   loss_gap  5.2506   loss_cross_term  -0.0459
gradient_cosine  0.9989                          final_loss  5.3169
modes_found  4      mode_entropy  1.3657         nearest_mode_distance  0.4505
mean_log_density  -2.2648                        data_mean_log_density  -2.1402

interpolants.py, in about thirteen seconds, trains three fields from the same initialization.

                 endpoint_error  covariance_deviation  mean_log_density  coarse_log_density
linear                   0.0000                0.4999           -2.5276             -3.0309
quadratic                0.0000                0.4939           -2.5748             -3.6419
trigonometric            0.0000                0.0179           -2.6621             -2.9566

All three schedules reach the four modes. The trigonometric one is the only schedule whose interpolation keeps identity covariance at every time, and it is the one that loses least when the sampler drops from 64 Euler steps to 8, while the t-squared schedule, whose velocity grows linearly in t, loses most.

flow_matching_mnist.py downloads MNIST on its first run. The offline path, --dataset toy, trains the same 471,265-parameter U-Net on 512 bars and blobs in about two minutes on one core.

epoch=20 squared_norm_loss=977.133
final loss 977.133  sample range [-5.15, 3.50]

Twenty epochs over 512 images is 80 optimizer steps, and the generated range is still wider than the data's [-1, 1]. With --epochs 100 the loss reaches 387.136 and the range [-3.25, 3.04], in about nine minutes. The MNIST default is 20 epochs over 60,000 images, and it is the one script here that does not finish while you watch it. --steps 20 trains for twenty optimizer steps and then samples, which is the quick check that the file runs.

What is simplified

The marginal velocity v_t(x) = E[U_t | X_t = x] is written in closed form in flow_matching_2d.py. That is available because the target is a Gaussian mixture and the prior is Gaussian; Section 2.14 of the notes is about the case where it is not, which is why the equivalence it proves is worth proving.
The equivalence is measured on one large batch, so the two losses agree up to that batch's own sampling error. The file prints the cross term that accounts for the residual exactly.
The identity-covariance check uses two independent standard Gaussian endpoints, which is the hypothesis the chapter states. The four-Gaussian target is centered but not white.
velocity_fields.py differentiates the densities with autograd where the chapter differentiates by hand, so no derivative in that file is transcribed from the notes. It also runs in double precision, because the Euler ladder it prints is arithmetic quoted to six decimal places.
The Gaussian translation of Section 2.17 is left out. Its coupling is deterministic, so the conditional velocity is the constant 4 at every state and time and already equals the marginal velocity.
The cubic Hermite interpolant of Section 2.18 is left out of interpolants.py. It needs endpoint slopes, which the four-Gaussian data do not carry.
The three two-dimensional scripts run on the CPU. Each state there has two coordinates, so a GPU would be slower. flow_matching_mnist.py picks CUDA, then Apple MPS, then the CPU.