| # 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. |
|
|
