lecture_6/README.md

Lecture 6: Flow Maps

Chapter 6 of the notes, book/chapters/ch6_flow_maps.tex.

Chapters 2 to 5 learned a local change and applied it again and again, so the number of steps set the cost of a sample. In this lecture we learn the map between two times directly. The map X_{s,t}(x) is fixed by three axioms, identity, cocycle and tangent, and every file writes it the same way, as the state plus the elapsed time times an average velocity the network produces, which makes the identity axiom hold for any weights at all. We differentiate the cocycle in the end time to get the Lagrangian identity, in the start time to get the Eulerian identity, and we leave it finite to get the composition identity. Each one becomes a squared residual, we pair each with a flow matching anchor on the diagonal, and from there we carry the construction into five settings: a point in space, a distribution over tokens, a Brownian path, a dimension that grows, and a conditioning state.

Every file is standalone, runs offline on a laptop CPU, and prints the numbers the chapter works by hand.

The files

File	Command	What it builds
flow_map.py	python lecture_6/flow_map.py	No learning. The chapter's velocity, the map by RK4, the three axioms, the inverse, the Jacobian and its differential equation, Liouville's formula, density transport and the generator
residuals.py	python lecture_6/residuals.py	The Lagrangian, Eulerian and composition residuals and the diagonal anchor, from the chapter's listing, with four models trained on the four Gaussian blobs
meanflow.py	python lecture_6/meanflow.py	MeanFlow: the Eulerian identity in the average-velocity parameterization, with the total derivative from one forward-mode pass
shortcut.py	python lecture_6/shortcut.py	Shortcut models: the composition identity indexed by step size, on a dyadic ladder from 1/128 to 1
consistency.py	python lecture_6/consistency.py	Consistency models on the edge t = 1 with the boundary coefficients, and consistency trajectory models on the whole triangle
fmlm.py	python lecture_6/fmlm.py	Flow Map Language Models: the two-time denoiser, the convex form, the cocycle in denoiser coordinates, the two gradients a softmax head sees, two KL terms and the decoding clock, on DNA
categorical_map.py	python lecture_6/categorical_map.py	Categorical Flow Maps: the Lagrangian identity with a softmax head, the teacher gap and drift, and the bound that replaces the residual, on phrases
discrete_map.py	python lecture_6/discrete_map.py	Discrete Flow Maps: why the derivative targets leave the simplex, and the repair in logit space, on DNA
stochastic_map.py	python lecture_6/stochastic_map.py	Strong stochastic flow maps: the solution map of an SDE, the weak map's floor, the Legendre coefficients and the Chen relations
expanding.py	python lecture_6/expanding.py	Expanding flow maps: a state whose length grows, local clocks, the composition identity under growth, and one draw of the insertion noise
posterior_map.py	python lecture_6/posterior_map.py	Diamond Maps and Meta Flow Maps: the inner clock, posterior samples, the value function from K samples, its gradient and the importance weight

Every file accepts the shared flags of cis6270.runner, and --steps 20 --quiet finishes in a few seconds. In flow_map.py, stochastic_map.py and posterior_map.py every quantity comes from the chapter's own formulas, so there --steps sets the Runge-Kutta substep count, the Monte Carlo draw count and the inner substep count, which each file's docstring says.

What a default run prints

The chapter's worked numbers come out exactly, and the test suite asserts them.

Quantity	Chapter	Printed
X_{0,1}(0.7), the cocycle through 0.5, the inverse, det J_{0,1}	4.1, 4.1, 0.7, 3	4.100000, 4.100000, 0.700000, 3.000000
log det J_{0,1} by Liouville, and the density ratio	1.0986, 3.000	1.0986, 3.0000
Stay-still residuals at x = 1, s = 0, t = 1	1, 1, 0	1.0000, 1.0000, 0.0000
MeanFlow on dx/dt = x: 1.718x = x + 0.718x	1.718, 1, 0.718	1.7183, 1.0000, 0.7183
Shortcut doubling: two steps of 0.25 against one of 0.5	1.297443	1.297443
Consistency distillation residual at h = 0.01	8.189e-05	8.189e-05
FMLM composed target, its four KL terms and their sum	(0.06, 0.12, 0.76, 0.06), 0.030747	the same, 0.030747
The two gradients in the logits at that target and student	(0.04, 0.03, -0.11, 0.04) and (0.0099, 0.01335, -0.03315, 0.0099)	the same
The same two once the student saturates to q_1 = 1e-06	-0.0600 and -4.62e-08	-0.0600 and -4.6171e-08
The decoding clock at 0.60, 0.75, 0.80, 0.85	0.006, 0.141, 0.456, 0.931	0.0057, 0.1411, 0.4561, 0.9313
Four steps on the clock	0.773490, 0.804481, 0.828095	the same
Gamma on the last three times, and over the long first leg	0.494620, 0.830144	the same
Categorical gap and drift, and the bound around a zero residual	antiparallel, 0.199	1e-34 residual, 0.1989
The same two terms averaged over three thousand states	0.393, 0.095	0.3955, 0.0963
Discrete derivative target, three entries negative	(-0.3333, -0.6667, 1.3333, -0.3333)	the same
The logit repair of (-0.284, 0.828, 0.456)	(0.0194, 0.6259, 0.3547)	the same
The weak map's floor 2K_1	6.389	6.3886 by Monte Carlo
Chen fusion of four half-interval draws	0.50, -0.20	0.50, -0.20
Local clocks at t = 0.8, and omega at (0.3, 0.6, 0.9)	0.800, 0.692, 0.333 and 0.125	the same
Jensen gap, and E[V_K] at K = 1, 2, 4, 8	0.434 and 1.000, 1.217, 1.340, 1.393	the same
Meta Flow Maps fixed point against the per-sample ratio	0.625 and 0.540	the same

For the saturated pair in that table, fmlm.py works the full expression of Section 6.27.4, the gradient of half the squared error through the softmax Jacobian,

d/dz_i (1/2) sum_j (q_j - p_j)^2  =  q_i (q_i - p_i)  -  q_i sum_j q_j (q_j - p_j) ,

at p = (0.06, 0.12, 0.76, 0.06) with the student moved to q_1 = 1e-06 and the other three renormalized. We also differentiate the two losses automatically and compare, and the two routes agree to four figures on every coordinate. As you can see from the pair, -0.0600 against -4.6171e-08 is the six orders of magnitude the chapter draws its point from.

We take the trained numbers from one default run on a laptop CPU, with seed 0.

Run	Headline
residuals.py, 600 steps each	distance to its own ODE 0.900 for the diagonal alone, 0.481, 0.376 and 0.541 with the Lagrangian, Eulerian and composition residuals; fraction on a mode 0.44 against 0.74, 0.66 and 0.70
meanflow.py, 2000 steps	nearest mode 0.627 at one step, 0.522 at two, 0.505 at four, all four modes found
shortcut.py, 2000 steps	nearest mode 0.605 at one step and 0.508 at eight; the doubling residual rises from below 1e-04 at d = 1/128 to 0.034 at d = 1/2
consistency.py, 1500 steps each	nearest mode 0.513 for the consistency model at one step, 0.515, 0.420 and 0.470 for the trajectory model at one, two and four
fmlm.py, 1500 steps	GC 0.586 at one step against 0.581 in the training data, motif 0.63, token entropy 1.371 of a possible 1.386, every sequence distinct
categorical_map.py, 1000 steps each	every phrase in the language at both step counts; trained on the residual, 0.56 of one-step phrases distinct at entropy 2.75; trained on the bound, 0.33 at 2.54, and 0.12 at 2.05 by four steps
discrete_map.py, 1500 steps	GC 0.609 at one step, motif 0.81, entropy 1.362, every sequence distinct; the unrepaired additive target is negative on 0.05% of its entries
stochastic_map.py	the strong map's terminal cloud sits 8.8e-04 from the finest crossing at two hops, 1.1e-04 at four, 1.4e-05 at eight and 2e-06 at sixteen, the weak map's 3.06 at every hop count
expanding.py, 1500 steps	GC 0.617 at one step, 0.621 on coordinates born at the source and 0.613 on coordinates born later; motif 0.75
posterior_map.py	on the blobs, the exact value -0.1127, the 16-sample estimate -0.1088, the single endpoint -0.1628

The DNA files draw 1024 training sequences of the default length 16, which have GC fraction 0.581 and motif fraction 0.749. They read 64 samples, as the phrase file does, so a reported fraction has about 0.06 of sampling noise and its second digit moves between seeds.

Longer runs, where they change the picture:

• residuals.py --steps 2400 moves the four distances to 1.033, 0.171, 0.247 and 0.511 and the four mode fractions to 0.27, 0.82, 0.71 and 0.79, against Table 6.2's 0.896, 0.292, 0.242 and 0.325 at 0.12, 0.84, 0.82 and 0.84. The gap between the anchor alone and the anchor plus any residual widens with the budget, and the three residuals stay inside each other's spread.
• meanflow.py --steps 8000 brings the one-step nearest mode from 0.627 to 0.437, below its own two-step and four-step figures, so the one-step map overtakes the grid.
• categorical_map.py --steps 3000 widens the gap between the two students, to 0.60 of one-step phrases distinct for the residual against 0.27 for the bound, and 0.54 against 0.09 at four steps. The bounded student loses variety as its budget grows, as Section 6.28.4 predicts, since the bound is minimized by a prediction that is constant in the end time.
• The DNA runs are at their plateau by their default budgets. At --steps 6000 the one-step motif fraction of fmlm.py reads 0.61 and of discrete_map.py 0.69, inside the sampling noise of the default run.

At the default flags on one core, we get three or four seconds for flow_map.py, stochastic_map.py and posterior_map.py, which train nothing, and twelve to twenty-seven seconds for the eight that do, including residuals.py with its four models and categorical_map.py and consistency.py with their two.

What is simplified

• The four blob models of residuals.py train for 600 steps on a four-layer network, so the distances are larger than Table 6.2 of the chapter, which trains longer. The ordering the chapter draws, the anchor alone against the anchor plus any residual, comes out at this budget.
• shortcut.py evaluates the zero-step term and the doubling term on the whole batch and weights them three to one, in place of the chapter's split of the batch in that proportion. The two have the same expectation and the weighted form costs one network call more per step.
• consistency.py distils from the exact marginal velocity of the training set, which is a softmax over the data points, in place of a separately trained flow matching teacher. The teacher's discretization error is then the only error the student inherits, as Section 6.20 sets out.
• posterior_map.py builds the inner map from the exact posterior over the training points, in place of training a conditional family, so the inner velocity is a closed form and the value, its gradient and the Jensen gap are all available to compare against.
• We compute the decoding clock of fmlm.py by quadrature on the one-position integral, which is exact for the Gaussian interpolant and treats each position on its own, so the clock stops where correlations between positions begin.
• categorical_map.py and discrete_map.py keep the end time below one, because the drift term of each loss divides by the remaining time. fmlm.py has no such term, so it draws a quarter of each batch at t = 1, the slice Section 6.26 asks for by name.
• In stochastic_map.py we carry three Legendre coefficients and fuse them with the Chen relations written out for n = 0, 1, 2. We stop at three, because the chapter lists the constants c_{3,m} only up to that order.
• Every network is cis6270.nets.TwoTimeMLP. The discrete files flatten the one-hot lift into its input and read its output as per-position logits, so one architecture covers both halves of the lecture and any difference in the samples comes from the loss.
• Each file sets torch.set_num_threads(1). These networks are small enough that splitting one matrix product across cores costs more than it saves. tests/conftest.py sets the same thing for the test session, and the line in each file acts when a reader runs the script directly.