# Lecture 2: Flow Matching (Chapter 2) A velocity field assigns one vector to every state and every time. Follow it from a starting point and you trace 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. To generate a sample we integrate the learned field with Euler steps. `flow_matching_mnist.py` is the lecture's main script, the one Section 2.20 of the notes walks through line by line. Underneath it we have three more files, which build 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. | `velocity_fields.py` integrates a field we already know, 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, and only the trigonometric interpolation keeps identity covariance at every time. The trigonometric schedule also 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 among the scripts here that run alone takes long enough that you will leave it going. `--steps 20` trains for twenty optimizer steps and then samples, so you can check quickly that the file runs. ## What is simplified - We write the marginal velocity `v_t(x) = E[U_t | X_t = x]` in closed form in `flow_matching_2d.py`. We can do that 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. - We measure the equivalence 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, and its covariance sits away from the identity. - `velocity_fields.py` differentiates the densities with autograd where the chapter differentiates by hand, so every derivative in that file comes out of the code. It also runs in double precision, because the Euler ladder it prints is arithmetic quoted to six decimal places. - We leave out the Gaussian translation of Section 2.17. Its coupling is deterministic, so the conditional velocity is the constant 4 at every state and time and already equals the marginal velocity. - We leave the cubic Hermite interpolant of Section 2.18 out of `interpolants.py`. It rests on endpoint slopes, and the four-Gaussian data are positions alone. - 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.