File size: 6,556 Bytes
0ba9d09 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | # 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.
|