Delete lecture_6

lecture_6/README.md

@@ -1,150 +0,0 @@
-# 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. This lecture learns 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. Differentiating the cocycle in the end time
-gives the Lagrangian identity, in the start time the Eulerian identity, and leaving it finite
-gives the composition identity. Each one becomes a squared residual, each is paired with a flow
-matching anchor on the diagonal, and the rest of the lecture carries 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 takes the shared flags of `cis6270.runner`. `--steps 20 --quiet` finishes in a few
-seconds. In `flow_map.py`, `stochastic_map.py` and `posterior_map.py` nothing is trained, so
-`--steps` sets the Runge-Kutta substep count, the Monte Carlo draw count and the inner substep
-count instead, 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 -6.10e-08 | -0.0600 and -4.6171e-08, see below |
-| 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 |
-
-One of those disagrees with the notes. Section 6.27.4 quotes a saturated squared-error gradient of
-`-6.10e-08`. Its own expression, 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) ,
-
-evaluated at `p = (0.06, 0.12, 0.76, 0.06)` and the student moved to `q_1 = 1e-06` with the other
-three renormalized, gives `-4.6171e-08`. The first term on its own is `-6.00e-08`, so the quoted
-number is that term with the second one dropped. `fmlm.py` prints the full expression, takes both
-gradients by automatic differentiation, and reproduces the unsaturated pair of Section 6.27.4
-exactly. The point the chapter draws from the number is unaffected, since `-0.0600` against
-`-4.6171e-08` is still six orders of magnitude.
-
-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
-carries 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, which is the one-step map overtaking 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, which is what Section 6.28.4
-  predicts, since the bound is minimized by a prediction that does not move with 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.
-
-Timings at the default flags, on one core. 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, which is the point
-  of Section 6.20.
-- `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.
-- The decoding clock of `fmlm.py` is computed by quadrature on the one-position integral, which is
-  exact for the Gaussian interpolant and says nothing about correlations between positions.
-- `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.
-- `stochastic_map.py` carries three Legendre coefficients and fuses them with the Chen relations
-  written out for `n = 0, 1, 2`. A fourth coefficient needs the constants `c_{3,m}`, which the
-  chapter does not list.
-- 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, so the line in each file is the one that acts when a reader runs the
-  script directly.

lecture_6/categorical_map.py

@@ -1,244 +0,0 @@
-"""Categorical Flow Maps, the Lagrangian identity on the simplex and the bound that replaces it.
-
-A softmax head makes the two-time prediction a distribution at every pair of times, and the
-Lagrangian identity rewritten in those coordinates reads pi_{s,t} = pi_{t,t}(X_{s,t}) minus a
-coefficient times the end-time derivative of the prediction. The derivative has no sign
-constraint, so the residual splits into a teacher gap, which is a difference of two distributions,
-and a drift, which is the scaled derivative. A parallelogram bound and Pinsker's inequality turn
-the squared residual into four KL divergences plus two squared drifts, and at the exact map the
-two pieces are antiparallel, so the bound is strictly positive where the residual is zero.
-
-Chapter 6, Section 6.28 of the notes. Trained on the four-token phrases of
-``cis6270.data.phrases``. Prints the chapter's worked cancellation on three tokens, where the gap
-is the exact negative of the drift and the bound reads 0.199 around a residual of zero, and then
-two models trained on the residual and on its bound.
-
-    python lecture_6/categorical_map.py                 # two models, 1000 steps each
-    python lecture_6/categorical_map.py --steps 20      # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-import torch
-
-from cis6270.data import phrases
-from cis6270.metrics import token_entropy, unique_fraction
-from cis6270.nets import TwoTimeMLP, count_parameters
-from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
-
-
-torch.set_num_threads(1)  # These networks are small, and splitting one matrix product costs more
-# than it saves. ``tests/conftest.py`` sets the same thing for the test session, so the line
-# here is the one that acts when a reader runs this script directly.
-
-BOOK_PRIOR = (0.5, 0.3, 0.2)  # The three data probabilities of the chapter's worked example.
-BOOK_SOURCE = (-0.5739, -0.8344, -0.0584)  # The source draw I_0 the chapter starts from.
-BOOK_STATES = 3000  # States the chapter averages the two terms of the bound over.
-
-
-# %% 02. The exact model on three tokens
-#
-# With a Gaussian source and the linear interpolant, the posterior over the data token is a
-# softmax whose only state dependence is the coordinate of that token. The velocity follows from
-# D_t(x) = x + (1-t) b(x,t), and the flow integrates to twelve digits with Runge-Kutta.
-def exact_denoiser(x, t, prior):
-    """pi_{t,t}(x), the posterior over the data token given the state. [B, K]."""
-    return torch.softmax(prior.log() + t * x / (1.0 - t) ** 2, dim=-1)
-
-
-def exact_velocity(x, t, prior):
-    """b(x,t) = (pi_{t,t}(x) - x)/(1 - t) for the linear interpolant. [B, K]."""
-    return (exact_denoiser(x, t, prior) - x) / (1.0 - t)
-
-
-def exact_flow(x, s, t, prior, n: int = 2000):
-    """X_{s,t}(x) by fourth-order Runge-Kutta on the exact velocity."""
-    h = (t - s) / n
-    for k in range(n):
-        u = s + k * h
-        k1 = exact_velocity(x, u, prior)
-        k2 = exact_velocity(x + 0.5 * h * k1, u + 0.5 * h, prior)
-        k3 = exact_velocity(x + 0.5 * h * k2, u + 0.5 * h, prior)
-        k4 = exact_velocity(x + h * k3, u + h, prior)
-        x = x + (h / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)
-    return x
-
-
-def exact_two_time(x, s, t, prior):
-    """pi_{s,t}(x) read off the convex parameterization, [(1-s) X_{s,t}(x) - (1-t) x]/(t-s)."""
-    return ((1.0 - s) * exact_flow(x, s, t, prior) - (1.0 - t) * x) / (t - s)
-
-
-# %% 03. The coefficient, the gap and the drift
-#
-# gamma_{s,t} = (t-s)(1-t)/(1-s) vanishes on the diagonal and at t = 1 and never exceeds 1/4 on
-# the square. The residual scaled by the remaining time splits as (1-t) r = a + b, with a the gap
-# between the student's prediction and the prediction at the transported state.
-def gamma_coefficient(s, t):
-    """gamma_{s,t} = (t-s)(1-t)/(1-s)."""
-    return (t - s) * (1.0 - t) / (1.0 - s).clamp(min=1e-6)
-
-
-def end_time_derivative(x, s, t, prior, n: int = 2000):
-    """d pi_{s,t}/dt from the integral representation, (1-s)/((t-s)(1-t)) times the bracket."""
-    transported = exact_flow(x, s, t, prior, n)
-    prediction = ((1.0 - s) * transported - (1.0 - t) * x) / (t - s)
-    at_output = exact_denoiser(transported, t, prior)
-    return (1.0 - s) / ((t - s) * (1.0 - t)) * (at_output - prediction), prediction, at_output
-
-
-# %% 04. The worked cancellation of Section 6.28
-#
-# The derivative above comes from the closed form of Section 6.26, which already carries the
-# cancellation inside it, so the central difference of pi_{s,t} in t below is the entry that checks
-# the two pieces against each other without assuming the identity.
-def worked_cancellation(s: float = 0.2, t: float = 0.6):
-    """The three-token example, with every number integrated or differentiated here."""
-    prior = torch.tensor(BOOK_PRIOR, dtype=torch.float64)
-    source = torch.tensor(BOOK_SOURCE, dtype=torch.float64)
-    x = (1.0 - s) * source + s * torch.eye(3, dtype=torch.float64)[0]  # Data token is the first.
-    derivative, prediction, at_output = end_time_derivative(x, s, t, prior)
-    gamma = (t - s) * (1.0 - t) / (1.0 - s)
-    gap = prediction - at_output  # a_{s,t}, a difference of two simplex points.
-    drift = gamma * derivative  # b_{s,t}, the scaled derivative of the prediction.
-    divergence = (at_output * (at_output / prediction).log()).sum()  # KL(teacher || student).
-    difference = (exact_two_time(x, s, t + 1e-5, prior)
-                  - exact_two_time(x, s, t - 1e-5, prior)) / 2e-5  # A central check.
-    return {
-        "state x": [round(v, 4) for v in x.tolist()],
-        "X_s,t(x)": [round(v, 4) for v in exact_flow(x, s, t, prior).tolist()],
-        "pi_s,t(x)": [round(v, 4) for v in prediction.tolist()],
-        "pi_t,t(X_s,t(x))": [round(v, 4) for v in at_output.tolist()],
-        "gamma_s,t": gamma,
-        "d pi / dt": [round(v, 4) for v in derivative.tolist()],
-        "central difference": [round(v, 4) for v in difference.tolist()],
-        "gap a": [round(v, 3) for v in gap.tolist()],
-        "drift b": [round(v, 3) for v in drift.tolist()],
-        "residual |a + b|^2": float(((gap + drift) ** 2).sum()),
-        "4 KL": 4.0 * float(divergence),
-        "2 |b|^2": 2.0 * float((drift ** 2).sum()),
-        "the bound": 4.0 * float(divergence) + 2.0 * float((drift ** 2).sum()),
-        "teacher entropy": -float((at_output * at_output.log()).sum()),
-        "cross-entropy": -float((at_output * prediction.log()).sum()),
-    }
-
-
-def averaged_bound(count: int = BOOK_STATES, s: float = 0.2, t: float = 0.6, seed: int = 0,
-                   n: int = 400):
-    """The same two terms averaged over states drawn from the interpolant at the same times."""
-    prior = torch.tensor(BOOK_PRIOR, dtype=torch.float64)
-    generator = torch.Generator().manual_seed(seed)
-    which = torch.multinomial(prior.expand(count, 3), 1, generator=generator).squeeze(1)
-    x1 = torch.eye(3, dtype=torch.float64)[which]  # [N, 3] the lifted data token.
-    x0 = torch.randn(count, 3, generator=generator, dtype=torch.float64)  # [N, 3] the source.
-    x = (1.0 - s) * x0 + s * x1  # [N, 3] the interpolant at the start time.
-    derivative, prediction, at_output = end_time_derivative(x, s, t, prior, n)
-    drift = ((t - s) * (1.0 - t) / (1.0 - s)) * derivative
-    divergence = (at_output * (at_output / prediction).log()).sum(-1).mean()
-    return {"averaged 4 KL": 4.0 * float(divergence),
-            "averaged 2 |b|^2": 2.0 * float((drift ** 2).sum(-1).mean()),
-            "averaged residual": float(((prediction - at_output + drift) ** 2).sum(-1).mean())}
-
-
-# %% 05. The student, and the two losses it can be trained on
-def predict(model, x, s, t, length, vocab):
-    """pi_{s,t}(x) at every position, a softmax head over the vocabulary."""
-    return torch.softmax(model(x.flatten(1), s, t).view(-1, length, vocab), dim=-1)
-
-
-def pieces(model, xs, s, t, length, vocab):
-    """The gap a and the drift b of the split, with the teacher detached."""
-    f = lambda tt: predict(model, xs, s, tt, length, vocab)
-    prediction, derivative = torch.func.jvp(f, (t,), (torch.ones_like(t),))  # One forward mode pass.
-    keep, move = (1.0 - t) / (1.0 - s), (t - s) / (1.0 - s)
-    transported = keep[:, None, None] * xs + move[:, None, None] * prediction  # X_{s,t}(x).
-    teacher = predict(model, transported.detach(), t, t, length, vocab).detach()  # pi_{t,t}.
-    drift = gamma_coefficient(s, t)[:, None, None] * derivative
-    return prediction - teacher, drift, teacher, prediction
-
-
-def make_loss(model, data, length, vocab, batch_size, device, bounded: bool):
-    eye = torch.eye(vocab, device=device)
-
-    def loss_fn(step: int) -> torch.Tensor:
-        index = torch.randint(0, len(data), (batch_size,), device=device)
-        x1, ids = eye[data[index]], data[index]  # [B, L, K] and [B, L].
-        x0 = torch.randn_like(x1)  # [B, L, K] the Gaussian source.
-        t_diag = torch.rand(batch_size, device=device)
-        xt = (1 - t_diag[:, None, None]) * x0 + t_diag[:, None, None] * x1
-        anchor = -(predict(model, xt, t_diag, t_diag, length, vocab).clamp_min(1e-9).log()
-                   .gather(-1, ids[..., None])).mean()  # Cross-entropy on the diagonal.
-
-        s = 0.9 * torch.rand(batch_size, device=device)  # [B] start time, kept off t = 1.
-        t = s + (0.95 - s) * torch.rand(batch_size, device=device)  # [B] end time below one.
-        xs = (1 - s[:, None, None]) * x0 + s[:, None, None] * x1
-        gap, drift, teacher, prediction = pieces(model, xs, s, t, length, vocab)
-        if bounded:  # Four KL divergences plus two squared drifts, the bound of Section 6.28.
-            divergence = (teacher * (teacher.clamp_min(1e-9)
-                                     / prediction.clamp_min(1e-9)).log()).sum(-1).mean()
-            return anchor + 4.0 * divergence + 2.0 * (drift ** 2).sum(-1).mean()
-        return anchor + ((gap + drift) ** 2).sum(-1).mean()  # The residual itself.
-
-    return loss_fn
-
-
-# %% 06. Sampling, and what a valid phrase is
-@torch.no_grad()
-def sample(model, count, length, vocab, steps, device):
-    """Walk a uniform grid with the convex form and read the phrase off the last prediction."""
-    x = torch.randn(count, length, vocab, device=device)
-    grid = torch.linspace(0.0, 1.0, steps + 1)
-    for a, b in zip(grid[:-1], grid[1:]):
-        s, t = a.expand(count).to(device), b.expand(count).to(device)
-        prediction = predict(model, x, s, t, length, vocab)
-        keep, move = (1.0 - t) / (1.0 - s), (t - s) / (1.0 - s)
-        x = keep[:, None, None] * x + move[:, None, None] * prediction
-    return x.argmax(-1)  # [N, L].
-
-
-def valid_fraction(drawn, group: int = 4):
-    """A phrase is in the language when position i holds a token from the i-th quarter."""
-    slots = torch.arange(drawn.shape[1], device=drawn.device)  # [L].
-    return float(((drawn // group) == slots[None, :]).all(dim=1).to(torch.get_default_dtype()).mean())
-
-
-# %% 07. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=1000, batch_size=128, samples=256, sample_steps=4)
-    args = parser.parse_args()
-    set_seed(args.seed)
-    device = pick_device(args.device)
-    tokens, vocabulary = phrases(num_phrases=1024, seed=args.seed)
-    data, length, vocab = tokens.to(device), tokens.shape[1], len(vocabulary)
-
-    report = worked_cancellation()
-    report.update(averaged_bound(seed=args.seed))  # The chapter's three thousand states.
-    if not args.quiet:
-        print(f"gap {report['gap a']} and drift {report['drift b']} cancel to "
-              f"{report['residual |a + b|^2']:.2e}, bound {report['the bound']:.3f}")
-
-    samples: list[str] = []
-    for name, bounded in (("the residual", False), ("the bound", True)):
-        set_seed(args.seed)  # Both students start from the same weights and see the same batches.
-        model = TwoTimeMLP(dim=length * vocab, width=args.width).to(device)
-        if not args.quiet:
-            print(f"\ntraining on {name}  parameters={count_parameters(model):,}")
-        losses = train(model, make_loss(model, data, length, vocab, args.batch_size, device,
-                                        bounded), steps=args.steps, lr=args.lr, quiet=args.quiet)
-        for steps in (1, args.sample_steps):
-            torch.manual_seed(args.seed + 1)
-            drawn = sample(model, args.samples, length, vocab, steps, device)
-            report[f"{name}, {steps} step(s): valid"] = valid_fraction(drawn)
-            report[f"{name}, {steps} step(s): distinct"] = unique_fraction(drawn)
-            report[f"{name}, {steps} step(s): entropy"] = token_entropy(drawn, vocab)
-        report[f"{name}: final loss"] = sum(losses[-20:]) / len(losses[-20:])
-        drawn = sample(model, 4, length, vocab, args.sample_steps, device)
-        samples.append(f"{name:<14s} " + " | ".join(
-            " ".join(vocabulary[int(i)] for i in row) for row in drawn[:2]))
-
-    print_report("Categorical Flow Maps, Chapter 6", report, samples)
-    save_run(args.out, config=vars(args), report=report, samples=samples)
-
-
-if __name__ == "__main__":
-    main()

lecture_6/consistency.py

@@ -1,231 +0,0 @@
-"""Consistency models and consistency trajectory models, one time fixed and then both free.
-
-A consistency model fixes the end time at the data end and learns the top edge of the triangle of
-start and end times, with the identity axiom imposed at the corner by the boundary coefficients
-c_skip and c_out. A consistency trajectory model puts the second time back and learns the whole
-triangle, writing the map as a weighted average of the input and a head that predicts a
-destination. Both are distilled here from a teacher velocity available in closed form.
-
-Chapter 6, Section 6.20 of the notes. Trained on the four Gaussian blobs of
-``cis6270.data.gaussian_mixture_2d``. Prints the boundary parameterization on dx/dt = x, where the
-model returns e^{0.5}, e^{0.1} and 1 at three start times, the teacher's error floor of
-8.189e-05, and the destination head of a consistency trajectory model at two pairs of times.
-
-    python lecture_6/consistency.py                 # two models, 1500 steps each
-    python lecture_6/consistency.py --steps 20      # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-import math
-
-import torch
-
-from cis6270.data import gaussian_mixture_2d
-from cis6270.metrics import mean_log_density, mode_report
-from cis6270.nets import TwoTimeMLP, count_parameters
-from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
-
-
-torch.set_num_threads(1)  # These networks are small, and splitting one matrix product costs more
-# than it saves. ``tests/conftest.py`` sets the same thing for the test session, so the line
-# here is the one that acts when a reader runs this script directly.
-
-TEACHER_STEP = 0.02  # h, the length of the teacher's solver step inside the distillation loss.
-
-
-# %% 02. The teacher velocity, in closed form on the empirical data
-#
-# The posterior over the data point given the interpolant is a softmax over the training set, so
-# the exact marginal velocity of the four blobs needs no training. A trained flow matching model
-# would serve in its place, with its own error on top of the solver's.
-def teacher_denoiser(x, t, data):
-    """D_t(x) = E[X_1 | X_t = x], the posterior mean over the training points. [B, 2]."""
-    spread = (1.0 - t).clamp(min=1e-4)[:, None]  # [B, 1] the width of the conditional.
-    squared = ((x[:, None, :] - t[:, None, None] * data[None]) ** 2).sum(-1)  # [B, N].
-    weights = torch.softmax(-squared / (2.0 * spread ** 2), dim=1)  # [B, N] posterior over data.
-    return weights @ data  # [B, 2] the posterior mean.
-
-
-def teacher_velocity(x, t, data):
-    """b(x,t) = (D_t(x) - x)/(1 - t) for the linear interpolant. [B, 2]."""
-    return (teacher_denoiser(x, t, data) - x) / (1.0 - t).clamp(min=1e-4)[:, None]
-
-
-def solver_step(x, t, h, data):
-    """One Euler step of the teacher, which the chapter writes as Phi_h."""
-    return x + h * teacher_velocity(x, t, data)  # [B, 2] the state at time t + h.
-
-
-# %% 03. The boundary parameterization, checked on the scalar flow
-#
-# f_theta(x, s) = c_skip(s) x + c_out(s) F_theta(x, s) with c_skip = 1 and c_out = 1 - s. At the
-# data end the output coefficient is zero and the model returns its input, so no training signal
-# is spent on a condition we already know.
-def scalar_consistency(x_value: float = 1.0):
-    """The velocity head, the output coefficient and the model at three start times."""
-    out = {}
-    for s in (0.5, 0.9, 1.0):
-        remaining = 1.0 - s  # c_out(s), which vanishes at the data end.
-        head = x_value * math.expm1(remaining) / remaining if remaining > 0 else x_value
-        out[f"c_out({s})"] = remaining
-        out[f"v_{s},1(1)"] = head
-        out[f"f({s})"] = x_value + remaining * head  # Equals e^{1-s} on this flow.
-    return out
-
-
-def teacher_floor(x_value: float = 1.0, s: float = 0.5, h: float = 0.01):
-    """The distillation residual at the exact map, which is the teacher's own Euler error."""
-    exact = lambda x, t: x * math.exp(1.0 - t)  # f(x, t) = x e^{1-t} on dx/dt = x.
-    moved = x_value + h * x_value  # One teacher Euler step, which undershoots e^{h}.
-    residual = exact(x_value, s) - exact(moved, s + h)
-    local = math.exp(h) - 1.0 - h  # The local error of that step.
-    return {"f(1, 0.5)": exact(x_value, s), "f(1.01, 0.51)": exact(moved, s + h),
-            "distillation residual": residual, "local Euler error": local,
-            "carried to the data end": local * math.exp(1.0 - s - h)}
-
-
-# %% 04. Consistency trajectory models, with both times free
-#
-# X_theta(x, s, t) = a_{s,t} x + (1 - a_{s,t}) g_theta(x, s, t) with a_{s,t} = (1-t)/(1-s), which
-# equals one on the diagonal, so the map returns its input there for any head whatever.
-def ctm_weight(s, t):
-    """a_{s,t} = (1-t)/(1-s), the share of the input kept by the step."""
-    return (1.0 - t) / (1.0 - s).clamp(min=1e-4)
-
-
-def ctm_map(model, x, s, t):
-    """The weighted average of the input and the destination head."""
-    a = ctm_weight(s, t)[:, None]  # [B, 1].
-    return a * x + (1.0 - a) * model(x, s, t)  # [B, 2].
-
-
-def scalar_ctm(x_value: float = 1.0):
-    """The head solved from the exact map at two pairs of times, and the two-route sampler."""
-    out = {}
-    for s, t in ((0.0, 0.5), (0.0, 1.0)):
-        a = (1.0 - t) / (1.0 - s)
-        reached = x_value * math.exp(t - s)  # The exact map of dx/dt = x.
-        out[f"a_{s},{t}"] = a
-        out[f"g_{s},{t}"] = (reached - a * x_value) / (1.0 - a)  # Solve the weighted average.
-    short = x_value * math.exp(0.5)  # X_{0,0.5}(1), the direct short step.
-    long_then_back = (x_value * math.exp(1.0)) * math.exp(-0.5)  # X_{1,0.5}(X_{0,1}(1)).
-    out["short step"], out["long jump and return"] = short, long_then_back
-    out["two routes apart"] = abs(short - long_then_back)
-    return out
-
-
-# %% 05. The two losses
-#
-# Both compare the model at one time against the model at a state the teacher has advanced by h,
-# with the target detached, which is the frozen copy of the weights the chapter writes as theta
-# minus. The consistency model holds the end time at one; the trajectory model draws it.
-def make_consistency_loss(model, data, batch_size, device):
-    ones = torch.ones(batch_size, device=device)
-
-    def loss_fn(step: int) -> torch.Tensor:
-        index = torch.randint(0, len(data), (batch_size,), device=device)
-        x1, x0 = data[index], torch.randn(batch_size, 2, device=device)
-        s = (1.0 - TEACHER_STEP) * torch.rand(batch_size, device=device)  # [B] with s + h <= 1.
-        xs = (1 - s[:, None]) * x0 + s[:, None] * x1  # [B, 2] the interpolant at the start time.
-        with torch.no_grad():  # The second leg is the teacher's, and no gradient reaches it.
-            advanced = solver_step(xs, s, TEACHER_STEP, data)
-            target = model.flow_map(advanced, s + TEACHER_STEP, ones)  # f_{theta-}(Phi_h(x_s), s+h).
-        return ((model.flow_map(xs, s, ones) - target) ** 2).mean()
-
-    return loss_fn
-
-
-def make_trajectory_loss(model, data, batch_size, device):
-
-    def loss_fn(step: int) -> torch.Tensor:
-        index = torch.randint(0, len(data), (batch_size,), device=device)
-        x1, x0 = data[index], torch.randn(batch_size, 2, device=device)
-        s = (1.0 - TEACHER_STEP) * torch.rand(batch_size, device=device)  # [B] start time.
-        t = s + TEACHER_STEP + (1.0 - s - TEACHER_STEP) * torch.rand(batch_size, device=device)
-        xs = (1 - s[:, None]) * x0 + s[:, None] * x1  # [B, 2].
-        with torch.no_grad():
-            advanced = solver_step(xs, s, TEACHER_STEP, data)
-            target = ctm_map(model, advanced, s + TEACHER_STEP, t)  # The same end time, one h later.
-        return ((ctm_map(model, xs, s, t) - target) ** 2).mean()
-
-    return loss_fn
-
-
-# %% 06. Sampling
-@torch.no_grad()
-def consistency_sample(model, x0):
-    """One call of the edge map, from the source straight to the data end."""
-    zeros = torch.zeros(len(x0), device=x0.device)
-    return model.flow_map(x0, zeros, torch.ones_like(zeros))  # [N, 2].
-
-
-@torch.no_grad()
-def trajectory_sample(model, x0, steps: int):
-    """Cross the triangle in equal pieces, one evaluation of the weighted average per piece."""
-    x = x0
-    grid = torch.linspace(0.0, 1.0, steps + 1, device=x0.device)  # [steps+1] the time grid.
-    for s, t in zip(grid[:-1], grid[1:]):
-        x = ctm_map(model, x, s.expand(len(x)), t.expand(len(x)))
-    return x  # [N, 2].
-
-
-# %% 07. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=1500, batch_size=128, samples=512)
-    args = parser.parse_args()
-    set_seed(args.seed)
-    device = pick_device(args.device)
-    data = gaussian_mixture_2d(num_points=1024, seed=args.seed).to(device)  # [1024, 2].
-
-    report = scalar_consistency()
-    report.update(teacher_floor())
-    report.update(scalar_ctm())
-    if not args.quiet:
-        print(f"boundary parameterization: f(0.5) = {report['f(0.5)']:.6f}, "
-              f"f(0.9) = {report['f(0.9)']:.6f}, f(1.0) = {report['f(1.0)']:.6f}")
-        print(f"teacher floor {report['distillation residual']:.6e} against the carried local "
-              f"error {report['carried to the data end']:.6e}")
-
-    samples: list[str] = []
-    set_seed(args.seed)
-    consistency = TwoTimeMLP(dim=2, width=args.width).to(device)
-    if not args.quiet:
-        print(f"\ntraining the consistency model  parameters={count_parameters(consistency):,}")
-    edge_losses = train(consistency,
-                        make_consistency_loss(consistency, data, args.batch_size, device),
-                        steps=args.steps, lr=args.lr, quiet=args.quiet)
-    report["consistency final loss"] = sum(edge_losses[-20:]) / len(edge_losses[-20:])
-    torch.manual_seed(args.seed + 1)
-    x0 = torch.randn(args.samples, 2, device=device)  # [N, 2] the shared source points.
-    points = consistency_sample(consistency, x0)
-    stats = mode_report(points)
-    report["consistency one-step nearest"] = stats["nearest_distance"]
-    report["consistency one-step log density"] = mean_log_density(points.cpu())
-    samples.append(f"consistency, 1 step: nearest {stats['nearest_distance']:.3f}  "
-                   f"modes {stats['modes_found']}  entropy {stats['mode_entropy']:.3f}")
-
-    set_seed(args.seed)
-    trajectory = TwoTimeMLP(dim=2, width=args.width).to(device)
-    if not args.quiet:
-        print(f"\ntraining the trajectory model  parameters={count_parameters(trajectory):,}")
-    losses = train(trajectory, make_trajectory_loss(trajectory, data, args.batch_size, device),
-                   steps=args.steps, lr=args.lr, quiet=args.quiet)
-    for steps in (1, 2, 4):
-        points = trajectory_sample(trajectory, x0, steps)
-        stats = mode_report(points)
-        report[f"trajectory {steps}-step nearest"] = stats["nearest_distance"]
-        samples.append(f"trajectory, {steps} step(s): nearest {stats['nearest_distance']:.3f}  "
-                       f"modes {stats['modes_found']}  entropy {stats['mode_entropy']:.3f}")
-    report["trajectory one step against four"] = float(
-        (trajectory_sample(trajectory, x0, 1) - trajectory_sample(trajectory, x0, 4)).norm(dim=1)
-        .mean())
-    report["trajectory final loss"] = sum(losses[-20:]) / len(losses[-20:])
-
-    print_report("Consistency and consistency trajectory models, Chapter 6", report, samples)
-    save_run(args.out, model=trajectory, config=vars(args), report=report, losses=losses,
-             samples=samples)
-
-
-if __name__ == "__main__":
-    main()

lecture_6/discrete_map.py

@@ -1,218 +0,0 @@
-"""Discrete Flow Maps, the derivative identities repaired in logit space.
-
-A curve that stays on the simplex has coordinates summing to one, so its velocity has coordinates
-summing to zero, and a target built from that velocity has a negative entry as soon as one
-coordinate of the prediction is smaller than the derivative pulling it down. Written in logits the
-same derivative becomes a positive multiplicative factor, dividing by it keeps every coordinate
-positive, and a softmax normalizes what is left, so the Kullback-Leibler divergence against the
-repaired target is defined at every state.
-
-Chapter 6, Sections 6.26 and 6.29 of the notes. Trained on the synthetic DNA of
-``cis6270.data.synthetic_dna``. Prints the derivative target of Section 6.26, three of whose four
-entries are negative, the three-token repair, where the additive target (-0.284, 0.828, 0.456)
-becomes the distribution (0.0194, 0.6259, 0.3547), and samples from the trained map.
-
-    python lecture_6/discrete_map.py                 # 1500 steps, about half a minute
-    python lecture_6/discrete_map.py --steps 20      # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-import torch
-
-from cis6270.data import DNA_SIZE, decode_dna, synthetic_dna
-from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
-from cis6270.nets import TwoTimeMLP, count_parameters
-from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
-
-
-torch.set_num_threads(1)  # These networks are small, and splitting one matrix product costs more
-# than it saves. ``tests/conftest.py`` sets the same thing for the test session, so the line
-# here is the one that acts when a reader runs this script directly.
-
-FLOOR = 1e-3  # The factor is clipped away from zero before the logarithm, as the chapter notes.
-
-
-# %% 02. Why the derivative target leaves the simplex
-#
-# d delta_{s,t}/dt = (1-s)/((t-s)(1-t)) [D_t(X_{s,t}(x)) - delta_{s,t}(x)]. The bracket is a
-# difference of two simplex points, so its coordinates sum to zero, and the prefactor only scales
-# it up. Nothing about the alphabet or the prior changes the sign pattern.
-def derivative_target(prediction, at_output, s: float, t: float):
-    """The end-time derivative of the two-time denoiser, a direction and not a distribution."""
-    prefactor = (1.0 - s) / ((t - s) * (1.0 - t))  # 6.6667 at s = 0.2, t = 0.8.
-    return prefactor * (at_output - prediction), prefactor
-
-
-def worked_derivative():
-    """The chapter's DNA example at s = 0.2, t = 0.8, whose target has three negative entries."""
-    prediction = torch.tensor([0.10, 0.20, 0.60, 0.10])  # delta_{0.2,0.8} over (A, C, G, T).
-    at_output = torch.tensor([0.05, 0.10, 0.80, 0.05])  # D_{0.8} at the transported state.
-    target, prefactor = derivative_target(prediction, at_output, 0.2, 0.8)
-    return {"derivative prefactor": prefactor,
-            "bracket": [round(v, 4) for v in (at_output - prediction).tolist()],
-            "derivative target": [round(v, 4) for v in target.tolist()],
-            "target sum": float(target.sum()),
-            "negative entries": int((target < 0).sum())}
-
-
-# %% 03. Differentiating a softmax
-#
-# The Jacobian of the softmax is diag(psi) - psi psi^T, so contracting it with the logit
-# derivative gives psi times the centered derivative. The centring makes the coordinates of
-# d psi/dt sum to zero and makes the identity insensitive to the shift invariance of the softmax.
-def centered(prediction, logit_derivative):
-    """c_{s,t} = d z/dt - <psi, d z/dt> 1, the quantity the whole identity is built on."""
-    mean = (prediction * logit_derivative).sum(-1, keepdim=True)  # <psi, dz/dt>.
-    return logit_derivative - mean  # [..., K] with <psi, c> = 0.
-
-
-# %% 04. The identity as a multiplicative relation, and the repaired teacher
-#
-# psi_{t,t}(X_{s,t}(x)) = psi_{s,t}(x) * (1 + gamma c_{s,t}(x)) coordinatewise. Solving for the
-# student divides by a strictly positive factor, and in logits the division is a subtraction of
-# its logarithm, so the target is a softmax of an expression the network already produces.
-def gamma_lagrangian(s, t):
-    """gamma_{s,t} = (t-s)(1-t)/(1-s), which never exceeds 1/4 on the square."""
-    return (t - s) * (1.0 - t) / (1.0 - s).clamp(min=1e-6)
-
-
-def gamma_eulerian(s, t):
-    """gamma^E_{s,t} = (1-s)(t-s)/(1-t), larger than the Lagrangian one and growing toward t = 1."""
-    return (1.0 - s) * (t - s) / (1.0 - t).clamp(min=1e-6)
-
-
-def repaired_teacher(teacher_logits, factor):
-    """Softmax(z_{t,t}(X_{s,t}(x)) - log(1 + gamma c)), a distribution by construction."""
-    return torch.softmax(teacher_logits - factor.clamp(min=FLOOR).log(), dim=-1)
-
-
-# %% 05. The repair worked on three tokens
-def worked_repair(s: float = 0.2, t: float = 0.6):
-    """The chapter's three-token repair, with the additive target beside the multiplicative one."""
-    prediction = torch.tensor([0.2, 0.5, 0.3])  # psi_{s,t}(x).
-    logit_derivative = torch.tensor([9.60, -2.28, -2.60])  # d z_{s,t}/dt.
-    gamma = (t - s) * (1.0 - t) / (1.0 - s)
-    c = centered(prediction, logit_derivative)  # Already centered on this example.
-    derivative = prediction * c  # d psi/dt, whose coordinates sum to zero.
-    teacher = torch.tensor([0.10, 0.60, 0.30])  # psi_{t,t}(X_{s,t}(x)) during training.
-    additive = teacher - gamma * derivative  # Read straight off the identity.
-    factor = 1.0 + gamma * c  # The multiplicative factor, positive in every coordinate.
-    quotient = teacher / factor  # The teacher divided coordinatewise.
-    return {
-        "weighted average of dz/dt": float((prediction * logit_derivative).sum()),
-        "centered derivative": [round(v, 4) for v in c.tolist()],
-        "d psi/dt": [round(v, 4) for v in derivative.tolist()],
-        "d psi/dt sum": float(derivative.sum()),
-        "additive target": [round(v, 4) for v in additive.tolist()],
-        "additive target sum": float(additive.sum()),
-        "additive negative entries": int((additive < 0).sum()),
-        "multiplicative factor": [round(v, 4) for v in factor.tolist()],
-        "quotient": [round(v, 4) for v in quotient.tolist()],
-        "quotient mass": float(quotient.sum()),
-        "repaired target": [round(v, 4) for v in (quotient / quotient.sum()).tolist()],
-        "log of the factor": [round(v, 4) for v in factor.log().tolist()],
-        "gamma_s,t": gamma,
-        "gamma^E_s,t": (1.0 - s) * (t - s) / (1.0 - t),
-    }
-
-
-# %% 06. The student, and the loss built on the repaired target
-def logits_of(model, x, s, t, length, vocab):
-    """z_{s,t}(x), one vector of length K per position."""
-    return model(x.flatten(1), s, t).view(-1, length, vocab)  # [B, L, K].
-
-
-def make_loss(model, data, length, vocab, batch_size, device, counter):
-    eye = torch.eye(vocab, device=device)
-
-    def loss_fn(step: int) -> torch.Tensor:
-        index = torch.randint(0, len(data), (batch_size,), device=device)
-        ids = data[index]  # [B, L] the data tokens.
-        x1, x0 = eye[ids], torch.randn(batch_size, length, vocab, device=device)
-        t_diag = torch.rand(batch_size, device=device)
-        xt = (1 - t_diag[:, None, None]) * x0 + t_diag[:, None, None] * x1
-        diagonal_logits = logits_of(model, xt, t_diag, t_diag, length, vocab)
-        anchor = torch.nn.functional.cross_entropy(  # KL to the one-hot token, the diagonal anchor.
-            diagonal_logits.reshape(-1, vocab), ids.reshape(-1))
-
-        s = 0.9 * torch.rand(batch_size, device=device)  # [B] start time.
-        t = s + (0.95 - s) * torch.rand(batch_size, device=device)  # [B] end time below one.
-        xs = (1 - s[:, None, None]) * x0 + s[:, None, None] * x1  # [B, L, K].
-        f = lambda tt: logits_of(model, xs, s, tt, length, vocab)
-        z, dz = torch.func.jvp(f, (t,), (torch.ones_like(t),))  # One forward-mode pass in t.
-        prediction = torch.softmax(z, dim=-1)  # psi_{s,t}(x).
-        gamma = gamma_lagrangian(s, t)[:, None, None]
-        factor = 1.0 + gamma * centered(prediction, dz)  # The multiplicative factor.
-        keep, move = (1.0 - t) / (1.0 - s), (t - s) / (1.0 - s)
-        transported = keep[:, None, None] * xs + move[:, None, None] * prediction  # X_{s,t}(x).
-        with torch.no_grad():  # The teacher is the same network read on the diagonal, detached.
-            teacher_logits = logits_of(model, transported.detach(), t, t, length, vocab)
-            target = repaired_teacher(teacher_logits, factor.detach())
-            additive = torch.softmax(teacher_logits, -1) - gamma * prediction * centered(
-                prediction, dz)  # The unrepaired target, measured and not trained on.
-            counter["negative"] += float((additive < 0).to(torch.get_default_dtype()).mean())
-            counter["batches"] += 1.0
-        student = torch.log_softmax(z, dim=-1)  # The student in log space.
-        return anchor + (target * ((target.clamp_min(1e-9)).log() - student)).sum(-1).mean()
-
-    return loss_fn
-
-
-# %% 07. Sampling
-@torch.no_grad()
-def sample(model, count, length, vocab, steps, device):
-    """Walk a uniform grid with the convex form and read the alphabet off the last prediction."""
-    x = torch.randn(count, length, vocab, device=device)
-    grid = torch.linspace(0.0, 1.0, steps + 1)
-    for a, b in zip(grid[:-1], grid[1:]):
-        s, t = a.expand(count).to(device), b.expand(count).to(device)
-        prediction = torch.softmax(logits_of(model, x, s, t, length, vocab), dim=-1)
-        keep, move = (1.0 - t) / (1.0 - s), (t - s) / (1.0 - s)
-        x = keep[:, None, None] * x + move[:, None, None] * prediction
-    return x.argmax(-1)  # [N, L].
-
-
-# %% 08. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=1500, batch_size=128, samples=64, length=16, sample_steps=4)
-    args = parser.parse_args()
-    set_seed(args.seed)
-    device = pick_device(args.device)
-    tokens, _ = synthetic_dna(num_sequences=1024, length=args.length, seed=args.seed)
-    data = tokens.to(device)  # [1024, L].
-
-    report = worked_derivative()
-    report.update(worked_repair())
-    if not args.quiet:
-        print(f"derivative target {report['derivative target']} has "
-              f"{report['negative entries']} negative entries and sums to "
-              f"{report['target sum']:.4f}")
-        print(f"additive {report['additive target']} repairs to {report['repaired target']}")
-
-    counter = {"negative": 0.0, "batches": 0.0}
-    model = TwoTimeMLP(dim=args.length * DNA_SIZE, width=args.width).to(device)
-    if not args.quiet:
-        print(f"\ntraining with the repaired target  parameters={count_parameters(model):,}")
-    losses = train(model, make_loss(model, data, args.length, DNA_SIZE, args.batch_size, device,
-                                    counter), steps=args.steps, lr=args.lr, quiet=args.quiet)
-    report["additive target negative share"] = counter["negative"] / max(counter["batches"], 1.0)
-
-    samples: list[str] = []
-    for steps in (1, args.sample_steps):
-        torch.manual_seed(args.seed + 1)  # The same source noise at every step count.
-        drawn = sample(model, args.samples, args.length, DNA_SIZE, steps, device)
-        report[f"{steps}-step GC fraction"] = gc_fraction(drawn)
-        report[f"{steps}-step motif fraction"] = motif_fraction(drawn)
-        report[f"{steps}-step token entropy"] = token_entropy(drawn, DNA_SIZE)
-        report[f"{steps}-step unique fraction"] = unique_fraction(drawn)
-        samples.append(f"{steps} step(s): {decode_dna(drawn[:1])[0]}  GC {gc_fraction(drawn):.3f}")
-    report["training GC fraction"] = gc_fraction(data)
-    report["final loss"] = sum(losses[-20:]) / len(losses[-20:])
-
-    print_report("Discrete Flow Maps, Chapter 6", report, samples)
-    save_run(args.out, model=model, config=vars(args), report=report, losses=losses, samples=samples)
-
-
-if __name__ == "__main__":
-    main()

lecture_6/expanding.py

@@ -1,217 +0,0 @@
-"""Expanding flow maps, where the state gains coordinates as time runs.
-
-The Jacobian of a map between states of different length is rectangular, so it has no determinant
-and no inverse, and Liouville's formula, the exact likelihood and the inverse map all stop
-applying. Splitting the move into an expansion that fills new coordinates with source noise and a
-transport that acts at the longer length restores a square Jacobian. Each coordinate then carries
-its own clock, rescaled to the interval it occupies, and the composition identity is still
-available, because a finite difference between two maps needs no derivative in time.
-
-Chapter 6, Section 6.32 of the notes. Trained on the synthetic DNA of
-``cis6270.data.synthetic_dna``, grown from half its length to all of it. Prints the local clocks
-of the chapter's figure, which read 0.800, 0.692 and 0.333 at a global time of 0.8, the
-composition weight 0.125 at (0.3, 0.6, 0.9), and samples from the trained map.
-
-    python lecture_6/expanding.py                 # 1500 steps, about half a minute
-    python lecture_6/expanding.py --steps 20      # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-import math
-
-import torch
-
-from cis6270.data import DNA_SIZE, decode_dna, synthetic_dna
-from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
-from cis6270.nets import TwoTimeMLP, count_parameters
-from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
-
-
-torch.set_num_threads(1)  # These networks are small, and splitting one matrix product costs more
-# than it saves. ``tests/conftest.py`` sets the same thing for the test session, so the line
-# here is the one that acts when a reader runs this script directly.
-
-BOOK_BIRTHS = (0.0, 0.0, 0.35, 0.7)  # The insertion times of the chapter's four-tile figure.
-
-
-# %% 02. Local clocks
-#
-# A coordinate inserted at t_ins has 1 - t_ins of global time left, so its own clock runs at
-# 1/(1 - t_ins) and reaches one exactly at the data end. Before its insertion the clock reads
-# zero and the coordinate holds its source value.
-def local_time(t, birth):
-    """t_i = (t - t_ins)/(1 - t_ins), clamped to [0, 1] so an unborn coordinate reads zero."""
-    return ((t - birth) / (1.0 - birth).clamp(min=1e-6)).clamp(0.0, 1.0)
-
-
-def worked_clocks(t: float = 0.8):
-    """The clocks, the speeds and the interpolant shares of the chapter's figure at one time."""
-    births = torch.tensor(BOOK_BIRTHS)
-    clocks = local_time(torch.tensor(t), births)  # [4].
-    speeds = 1.0 / (1.0 - births)  # [4] dt_i/dt.
-    return {
-        "insertion times": [round(v, 2) for v in births.tolist()],
-        "local clocks at t=0.8": [round(v, 3) for v in clocks.tolist()],
-        "clock speeds": [round(v, 3) for v in speeds.tolist()],
-        "source share": [round(v, 3) for v in (1.0 - clocks).tolist()],
-        "last clock from its speed": float(speeds[-1] * (1.0 - t)),
-        "live coordinates at 0.8": int((births <= t).sum()),
-        "live coordinates at 0.6": int((births <= 0.6).sum()),
-    }
-
-
-# %% 03. The composition weight under growth
-#
-# The convexity weight is a ratio of remaining times and makes no reference to the dimension, so
-# the derivation at fixed length goes through once every leg has been lifted to the final length.
-# A coordinate born at zero carries the chapter's omega; a later birth carries its own.
-def omega(s, u, t):
-    """omega_{s,u,t} = (u-s)(1-t)/((t-s)(1-u)), the share of the earlier leg."""
-    return (u - s) * (1.0 - t) / (((t - s) * (1.0 - u)).clamp(min=1e-9))
-
-
-def worked_omega(s: float = 0.3, u: float = 0.6, t: float = 0.9):
-    """The weight at the chapter's triple, for coordinates born at three different times."""
-    out = {"omega at (0.3, 0.6, 0.9)": float(omega(*(torch.tensor([v]) for v in (s, u, t))))}
-    for birth in (0.0, 0.35, 0.7):
-        b = torch.tensor([birth])
-        local = [local_time(torch.tensor([v]), b) for v in (s, u, t)]
-        out[f"omega for a birth at {birth}"] = float(omega(*local))
-    return out
-
-
-# %% 04. A state whose length grows, and where the fresh noise enters
-#
-# The source values of every coordinate that will ever appear are drawn once, at the final
-# length, and a leg that stops before an insertion leaves that coordinate at its source value.
-# One draw then serves all three maps of the composition identity.
-def lifted_state(source, data_onehot, t, births):
-    """The state at global time t, each coordinate interpolated on its own clock. [B, L, K]."""
-    clock = local_time(t[:, None], births[None, :])[:, :, None]  # [B, L, 1].
-    return (1.0 - clock) * source + clock * data_onehot  # [B, L, K].
-
-
-def expansion_log_density(source, births, s: float, t: float):
-    """The source log density of the coordinates inserted inside (s, t], the expansion term."""
-    inserted = (births > s) & (births <= t)  # [L] which coordinates appear on this interval.
-    per_coordinate = (-0.5 * source ** 2 - 0.5 * math.log(2.0 * math.pi)).sum(-1)  # [B, L].
-    return float(per_coordinate[:, inserted].sum(-1).mean()), int(inserted.sum())
-
-
-# %% 05. The model, a two-time denoiser that also reads each coordinate's clock
-def predict(model, x, s, t, births, length, vocab):
-    """psi_{s,t}(x) at every position, with each coordinate's own clock as an extra channel.
-
-    The network has the same width in and out, so every position carries K + 1 channels, the K
-    coordinates of the state and the local clock. The clock channel of the output is dropped.
-    """
-    clock = local_time(s[:, None], births[None, :])[:, :, None]  # [B, L, 1] where each stands.
-    features = torch.cat([x, clock], dim=-1).flatten(1)  # [B, L*(K+1)].
-    logits = model(features, s, t).view(-1, length, vocab + 1)[..., :vocab]  # [B, L, K].
-    return torch.softmax(logits, dim=-1)  # [B, L, K] on the simplex at every position.
-
-
-def step_map(x, prediction, s, t, births):
-    """The convex form, with the weights read on each coordinate's own clock."""
-    s_local = local_time(s[:, None], births[None, :])[:, :, None]  # [B, L, 1].
-    t_local = local_time(t[:, None], births[None, :])[:, :, None]  # [B, L, 1].
-    keep = (1.0 - t_local) / (1.0 - s_local).clamp(min=1e-6)
-    return keep * x + (1.0 - keep) * prediction  # [B, L, K].
-
-
-def make_loss(model, data, births, length, vocab, batch_size, device):
-    eye = torch.eye(vocab, device=device)
-
-    def loss_fn(step: int) -> torch.Tensor:
-        index = torch.randint(0, len(data), (batch_size,), device=device)
-        ids = data[index]  # [B, L].
-        x1 = eye[ids]  # [B, L, K] the data, lifted.
-        source = torch.randn(batch_size, length, vocab, device=device)  # Drawn once, at full length.
-        t_diag = torch.rand(batch_size, device=device)
-        xt = lifted_state(source, x1, t_diag, births)
-        diagonal = -(predict(model, xt, t_diag, t_diag, births, length, vocab).clamp_min(1e-9)
-                     .log().gather(-1, ids[..., None])).mean()  # Cross-entropy on the diagonal.
-
-        s = torch.rand(batch_size, device=device)
-        t = s + (1.0 - s) * torch.rand(batch_size, device=device)
-        u = s + (t - s) * torch.rand(batch_size, device=device)
-        xs = lifted_state(source, x1, s, births)  # The same source draw serves all three legs.
-        with torch.no_grad():
-            first = predict(model, xs, s, u, births, length, vocab)
-            second = predict(model, step_map(xs, first, s, u, births), u, t, births, length, vocab)
-            s_l, u_l, t_l = (local_time(v[:, None], births[None, :]) for v in (s, u, t))
-            weight = omega(s_l, u_l, t_l)[:, :, None]  # [B, L, 1] one weight per coordinate.
-            target = weight * first + (1.0 - weight) * second  # Still on the simplex.
-        student = predict(model, xs, s, t, births, length, vocab).clamp_min(1e-9)
-        return diagonal + (target * (target.clamp_min(1e-9) / student).log()).sum(-1).mean()
-
-    return loss_fn
-
-
-# %% 06. Sampling
-@torch.no_grad()
-def sample(model, count, births, length, vocab, steps, device):
-    """Walk a uniform grid, with every coordinate stepping on its own clock."""
-    x = torch.randn(count, length, vocab, device=device)  # The source values, drawn once.
-    grid = torch.linspace(0.0, 1.0, steps + 1)
-    for a, b in zip(grid[:-1], grid[1:]):
-        s, t = a.expand(count).to(device), b.expand(count).to(device)
-        x = step_map(x, predict(model, x, s, t, births, length, vocab), s, t, births)
-    return x.argmax(-1)  # [N, L].
-
-
-# %% 07. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=1500, batch_size=128, samples=64, length=16, sample_steps=4)
-    args = parser.parse_args()
-    set_seed(args.seed)
-    device = pick_device(args.device)
-    tokens, _ = synthetic_dna(num_sequences=1024, length=args.length, seed=args.seed)
-    data = tokens.to(device)  # [1024, L] the full-length sequences.
-
-    # Half the coordinates exist at the source; the rest are inserted at evenly spaced times.
-    half = args.length // 2
-    births = torch.cat([torch.zeros(half, device=device),
-                        torch.linspace(0.1, 0.8, args.length - half, device=device)])  # [L].
-
-    report = worked_clocks()
-    report.update(worked_omega())
-    report["coordinates at the source"] = int((births == 0).sum())
-    report["coordinates at the data"] = args.length
-    source = torch.randn(256, args.length, DNA_SIZE)
-    added, count = expansion_log_density(source, births.cpu(), 0.0, 0.5)
-    report["coordinates inserted in (0, 0.5]"] = count
-    report["source log density of those"] = added
-    if not args.quiet:
-        print(f"local clocks at t = 0.8: {report['local clocks at t=0.8']} with speeds "
-              f"{report['clock speeds']}")
-        print(f"omega at (0.3, 0.6, 0.9) = {report['omega at (0.3, 0.6, 0.9)']:.3f}")
-
-    model = TwoTimeMLP(dim=args.length * (DNA_SIZE + 1), width=args.width).to(device)
-    if not args.quiet:
-        print(f"\ntraining the expanding map  parameters={count_parameters(model):,}")
-    losses = train(model, make_loss(model, data, births, args.length, DNA_SIZE, args.batch_size,
-                                    device), steps=args.steps, lr=args.lr, quiet=args.quiet)
-
-    samples: list[str] = []
-    early, late = births == 0, births > 0  # Coordinates born at the source and born later.
-    for steps in (1, args.sample_steps):
-        torch.manual_seed(args.seed + 1)
-        drawn = sample(model, args.samples, births, args.length, DNA_SIZE, steps, device)
-        report[f"{steps}-step GC fraction"] = gc_fraction(drawn)
-        report[f"{steps}-step GC, born at the source"] = gc_fraction(drawn[:, early])
-        report[f"{steps}-step GC, born later"] = gc_fraction(drawn[:, late])
-        report[f"{steps}-step motif fraction"] = motif_fraction(drawn)
-        report[f"{steps}-step token entropy"] = token_entropy(drawn, DNA_SIZE)
-        report[f"{steps}-step unique fraction"] = unique_fraction(drawn)
-        samples.append(f"{steps} step(s): {decode_dna(drawn[:1])[0]}  GC {gc_fraction(drawn):.3f}")
-    report["training GC fraction"] = gc_fraction(data)
-    report["final loss"] = sum(losses[-20:]) / len(losses[-20:])
-
-    print_report("Expanding flow maps, Chapter 6", report, samples)
-    save_run(args.out, model=model, config=vars(args), report=report, losses=losses, samples=samples)
-
-
-if __name__ == "__main__":
-    main()

lecture_6/flow_map.py

@@ -1,206 +0,0 @@
-"""The flow map of an ordinary differential equation, with no learning anywhere.
-
-A flow map X_{s,t}(x) returns the state at time t of the trajectory that passes through x at time
-s, so one evaluation replaces the N Euler steps of Chapter 2. Here we integrate a velocity we
-already know, the Gaussian flow of the chapter, and check against it every property the rest of
-the lecture trains for: the three axioms, the inverse, the Jacobian and its differential equation,
-Liouville's formula, the transport of a density, and the generator.
-
-Chapter 6, Sections 6.3 to 6.12 of the notes. Runs on the scalar Gaussian flow and on a batch from
-``cis6270.data.gaussian_mixture_2d``. Prints the four numbers of the chapter's listing, 4.100000
-for the map, the cocycle agreeing with it, 0.700000 for the inverse and a Jacobian determinant of
-3.000000, together with log 3, the density ratio and the log likelihood.
-
-    python lecture_6/flow_map.py                 # 200 Runge-Kutta substeps per crossing
-    python lecture_6/flow_map.py --steps 20      # 20 substeps, a second, and the same report
-"""
-
-# %% 01. Imports and configuration
-import math
-
-import torch
-from torch.func import jacrev  # Differentiates the whole integration in its starting point.
-
-from cis6270.data import GAUSSIAN_STD, gaussian_mixture_2d
-from cis6270.runner import common_parser, print_report, save_run, set_seed
-
-torch.set_num_threads(1)  # The integration is a scalar loop over a small batch, so one thread is
-# the fastest here. ``tests/conftest.py`` sets the same thing for the test session, so the line here
-# is the one that acts when a reader runs this script directly.
-torch.set_default_dtype(torch.float64)  # The axioms are checked below the float32 noise floor.
-
-
-SUBSTEPS = 200  # n, the Runge-Kutta substeps per crossing. No optimizer runs in this file.
-START = 0.7  # The state x at which the chapter evaluates every check.
-
-
-# %% 02. The velocity field of the Gaussian flow
-#
-# Source N(0,1) and data N(2,9), paired by x_1 = 2 + 3 x_0 and joined by a straight line. The
-# velocity is a translation at speed 2 plus an expansion whose strength 2/(1+2t) decays as the
-# cloud widens.
-def b(x, t):  # velocity field
-    return 2.0 + 2.0 * (x - 2.0 * t) / (1.0 + 2.0 * t)  # Section 6.9, b(x,t).
-
-
-def exact_map(x, s, t):
-    """X_{s,t}(x) = 2t + (1+2t)(x - 2s)/(1 + 2s), the closed form the integrator has to reproduce."""
-    return 2.0 * t + (1.0 + 2.0 * t) * (x - 2.0 * s) / (1.0 + 2.0 * s)  # [B] or scalar.
-
-
-# %% 03. The flow map by fourth-order Runge-Kutta
-#
-# Four velocity evaluations per substep, at its start, twice at its midpoint and at its end,
-# combined with the weights 1, 2, 2, 1 over six. The step h is negative when t is below s, so the
-# same function integrates backward with no change.
-def flow(x, s, t, n=SUBSTEPS):  # X_{s,t}(x) by RK4
-    h = (t - s) / n
-    for k in range(n):
-        u = s + k * h
-        k1 = b(x, u)
-        k2 = b(x + 0.5 * h * k1, u + 0.5 * h)
-        k3 = b(x + 0.5 * h * k2, u + 0.5 * h)
-        k4 = b(x + h * k3, u + h)
-        x = x + (h / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)
-    return x
-
-
-# %% 04. The three axioms, checked numerically
-#
-# Identity, cocycle and tangent characterize the map (Proposition 6.1). The first holds by
-# definition, the second by uniqueness of the trajectory, the third by the defining equation.
-def check_axioms(x, n):
-    """Return the residual of each axiom at the state x, each of which the true map drives to zero."""
-    identity = flow(x, 0.5, 0.5, n) - x  # X_{t,t} = id.
-    cocycle = flow(flow(x, 0.0, 0.5, n), 0.5, 1.0, n) - flow(x, 0.0, 1.0, n)  # X_{t,u} o X_{s,t}.
-    h = 1e-6  # One-sided difference quotient at the diagonal.
-    tangent = (flow(x, 0.5, 0.5 + h, n) - x) / h - b(x, 0.5)  # d/dt X at t = s equals b(x,s).
-    return float(identity.abs()), float(cocycle.abs()), float(tangent.abs())
-
-
-# %% 05. The Jacobian and the equation it satisfies
-#
-# dJ/dt = grad b(X_{s,t}(x), t) J, starting from J_{s,s} = I. In one dimension grad b is the single
-# derivative 2/(1+2t), and the solution from zero is J_{0,t} = 1 + 2t.
-def jacobian_by_ode(x, s, t, n):
-    """Integrate the state and its Jacobian together, with one Euler step each on the same grid."""
-    jacobian = torch.ones_like(x)  # J_{s,s} = I, the identity axiom differentiated in x.
-    h = (t - s) / n
-    for k in range(n):
-        u = s + k * h
-        grad_b = 2.0 / (1.0 + 2.0 * u)  # d b / d x at time u, the divergence in one dimension.
-        jacobian = jacobian + h * grad_b * jacobian  # Linear in J, with the coefficient frozen.
-        x = x + h * b(x, u)  # The trajectory the coefficient is read along.
-    return jacobian
-
-
-def liouville_log_det(s, t, n):
-    """Integrate d/dt log det J = div b, which on this flow is the number log 3 over [0, 1]."""
-    h = (t - s) / n
-    grid = s + h * (torch.arange(n) + 0.5)  # Midpoints of the substeps.
-    return float((2.0 / (1.0 + 2.0 * grid)).sum() * h)  # Riemann sum of the divergence.
-
-
-# %% 06. Transporting a density along the flow
-#
-# p_t(X_{s,t}(x)) |det J_{s,t}(x)| = p_s(x). The flow stretches every length by 3 over the unit
-# interval, so it divides every density by 3.
-def normal_density(x, mean, std):
-    """The scalar Gaussian density, used at both ends of the transport statement."""
-    return math.exp(-((x - mean) ** 2) / (2.0 * std ** 2)) / (std * math.sqrt(2.0 * math.pi))
-
-
-def mixture_log_density_2d(points, scale, shift):
-    """Log density of the four-Gaussian mixture after the affine map y = shift + scale * x.
-
-    The pushforward of a Gaussian mixture under an affine map is a Gaussian mixture with the
-    centers and the widths carried through, which gives the transported density in closed form.
-    """
-    from cis6270.data import GAUSSIAN_CENTERS
-
-    centers = shift + scale * GAUSSIAN_CENTERS.to(points)  # [4, 2] transported centers.
-    std = scale * GAUSSIAN_STD  # Each component is widened by the same factor.
-    squared = ((points[:, None, :] - centers) ** 2).sum(dim=-1)  # [N, 4].
-    component = -0.5 * squared / std ** 2 - math.log(2.0 * math.pi * std ** 2)  # [N, 4].
-    return torch.logsumexp(component + math.log(0.25), dim=-1)  # [N] mixture log density.
-
-
-# %% 07. The generator of the flow
-#
-# (L_t f)(x) = b(x,t) . grad f(x), and the rate of change of f along a trajectory is the generator
-# applied to f and read at the moving point.
-def generator(f, grad_f, x, t):
-    """(L_t f)(x) for a scalar test function, which in one dimension is b times the derivative."""
-    return b(x, t) * grad_f(x)  # Section 6.11, the generator.
-
-
-# %% 08. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=SUBSTEPS, samples=512)
-    args = parser.parse_args()  # --steps is the substep count n here, since nothing is trained.
-    set_seed(args.seed)
-    n = max(args.steps, 4)  # Runge-Kutta needs a handful of substeps to be worth its name.
-
-    # %% 09. The chapter's four lines
-    x = torch.tensor(START)  # The single starting point every check below uses.
-    one_crossing = flow(x, 0.0, 1.0, n)  # 2 + 3x = 4.1.
-    two_crossings = flow(flow(x, 0.0, 0.5, n), 0.5, 1.0, n)  # The cocycle.
-    there_and_back = flow(flow(x, 0.0, 1.0, n), 1.0, 0.0, n)  # The inverse.
-    det_jacobian = jacrev(lambda z: flow(z, 0.0, 1.0, n))(x)  # det J = 3.
-    if not args.quiet:
-        print(f"X_0,1({START})      = {float(one_crossing):.6f}   closed form "
-              f"{float(exact_map(x, 0.0, 1.0)):.6f}")
-        print(f"cocycle through 0.5 = {float(two_crossings):.6f}")
-        print(f"inverse X_1,0 o X_0,1 = {float(there_and_back):.6f}")
-        print(f"det J_0,1           = {float(det_jacobian):.6f}")
-
-    identity, cocycle, tangent = check_axioms(x, n)
-
-    # %% 10. The Jacobian two ways, and Liouville's formula
-    jacobian_ode = jacobian_by_ode(x, 0.0, 1.0, 4000)  # One Euler solve of dJ/dt = grad b . J.
-    log_det = liouville_log_det(0.0, 1.0, 4000)  # Integral of the divergence over [0, 1].
-
-    # %% 11. Transport, in one dimension and in two
-    image = float(one_crossing)  # X_{0,1}(0.7) = 4.1.
-    p_source = normal_density(START, 0.0, 1.0)  # p_0(0.7).
-    p_data = normal_density(image, 2.0, 3.0)  # p_1(4.1).
-    log_likelihood = math.log(p_source) - log_det  # Section 6.10, with the divergence integral.
-
-    points = gaussian_mixture_2d(num_points=args.samples, seed=args.seed)  # [N, 2] source cloud.
-    moved = exact_map(points, 0.0, 1.0)  # The same affine flow applied to both coordinates.
-    log_p0 = mixture_log_density_2d(points, 1.0, 0.0)  # [N] density before transport.
-    log_p1 = mixture_log_density_2d(moved, 3.0, 2.0)  # [N] density of the transported mixture.
-    transport_gap = float((log_p1 + 2.0 * math.log(3.0) - log_p0).abs().max())  # det J = 3^2 = 9.
-
-    # %% 12. The generator along one trajectory
-    f, grad_f = (lambda z: z ** 2), (lambda z: 2.0 * z)  # A test function and its derivative.
-    h = 1e-5
-    along = (f(flow(x, 0.0, 0.5 + h, n)) - f(flow(x, 0.0, 0.5 - h, n))) / (2.0 * h)  # d/dt f(X).
-    generator_gap = abs(float(along) - float(generator(f, grad_f, flow(x, 0.0, 0.5, n), 0.5)))
-
-    report = {
-        "X_0,1(0.7)": float(one_crossing),
-        "cocycle through u=0.5": float(two_crossings),
-        "inverse back to the start": float(there_and_back),
-        "det J_0,1 by jacrev": float(det_jacobian),
-        "det J_0,1 by its own ODE": float(jacobian_ode),
-        "log det J_0,1 by Liouville": log_det,
-        "identity axiom residual": identity,
-        "cocycle axiom residual": cocycle,
-        "tangent axiom residual": tangent,
-        "p_0(0.7)": p_source,
-        "p_1(4.1)": p_data,
-        "density ratio p_0/p_1": p_source / p_data,
-        "log p_1(4.1) by transport": log_likelihood,
-        "log p_1(4.1) directly": math.log(p_data),
-        "2-D transport gap": transport_gap,
-        "generator gap": generator_gap,
-        "substeps n": n,
-    }
-    print_report("Flow map of the Gaussian flow, Chapter 6", report)
-    save_run(args.out, config=vars(args), report=report)
-
-
-if __name__ == "__main__":
-    main()

lecture_6/fmlm.py

@@ -1,291 +0,0 @@
-"""Flow Map Language Models, the composition identity in denoiser coordinates.
-
-Each token is lifted to a one-hot vector, the source is Gaussian over that space, and the state
-moves through R^{L x K}, so no atom has to be split. Stepping by the full remaining time 1 - s in
-place of the elapsed time t - s turns the average velocity into a two-time denoiser that lands on
-the simplex, the map becomes a convex combination of staying put and jumping to the prediction,
-and the cocycle becomes a convex combination of two shorter predictions, which a KL divergence
-can measure.
-
-Chapter 6, Sections 6.25 to 6.27 of the notes. Trained on the synthetic DNA of
-``cis6270.data.synthetic_dna``. Prints the convex weights, the composition weight gamma = 0.2 at
-(0.2, 0.5, 0.8), the worked KL of 0.030747 nats, the two gradients that decide the loss, the
-decoding clock of a 30,522-token vocabulary and the four steps it produces, and samples from the
-trained map.
-
-    python lecture_6/fmlm.py                 # 1500 steps, about half a minute
-    python lecture_6/fmlm.py --steps 20      # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-import math
-
-import torch
-
-from cis6270.data import DNA_SIZE, decode_dna, synthetic_dna
-from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
-from cis6270.nets import TwoTimeMLP, count_parameters
-from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
-
-
-torch.set_num_threads(1)  # These networks are small, and splitting one matrix product costs more
-# than it saves. ``tests/conftest.py`` sets the same thing for the test session, so the line
-# here is the one that acts when a reader runs this script directly.
-
-BOOK_VOCAB = 30522  # The vocabulary the chapter puts on the clock, a standard tokenizer.
-ENDPOINT_FRACTION = 0.25  # Share of each batch drawn with t = 1, the slice the sampler reads.
-
-
-# %% 02. The convex form and the composition weight
-#
-# X_{s,t}(x) = (1-t)/(1-s) x + (t-s)/(1-s) delta_{s,t}(x). Both weights are non-negative for
-# s <= t <= 1 and they sum to one, so each step is a convex combination of staying put and jumping
-# to the predicted endpoint. At t = s the pair is (1, 0) and the identity axiom holds.
-def convex_weights(s, t):
-    """The pair ((1-t)/(1-s), (t-s)/(1-s)) that the step mixes the state and the prediction with."""
-    remaining = (1.0 - s).clamp(min=1e-6)
-    return (1.0 - t) / remaining, (t - s) / remaining
-
-
-def composition_weight(s, u, t):
-    """gamma = (1-t)(u-s)/((1-u)(t-s)), the share of the long prediction made over the first leg."""
-    return (1.0 - t) * (u - s) / (((1.0 - u) * (t - s)).clamp(min=1e-6))
-
-
-# %% 03. The cocycle in denoiser coordinates, checked on the scalar flow
-#
-# delta_{s,t}(x) = gamma delta_{s,u}(x) + (1-gamma) delta_{u,t}(X_{s,u}(x)). On dx/dt = x the
-# two-time denoiser is x[1 + (1-s)(e^{t-s}-1)/(t-s)], so every term is available in closed form.
-def scalar_denoiser(x, s, t):
-    """delta_{s,t}(x) on the flow dx/dt = x."""
-    return x * (1.0 + (1.0 - s) * math.expm1(t - s) / (t - s))
-
-
-def scalar_cocycle(x: float = 1.0, s: float = 0.2, u: float = 0.5, t: float = 0.8):
-    """Every term of the denoiser cocycle at the chapter's triple of times."""
-    gamma = (1.0 - t) * (u - s) / ((1.0 - u) * (t - s))
-    first = scalar_denoiser(x, s, u)  # The prediction over the first leg.
-    reached = x * math.exp(u - s)  # X_{s,u}(x), the state the first leg arrives at.
-    second = scalar_denoiser(reached, u, t)  # The prediction over the second leg.
-    long = scalar_denoiser(x, s, t)  # The prediction over the whole interval.
-    keep, move = (1.0 - t) / (1.0 - s), (t - s) / (1.0 - s)
-    return {"gamma": gamma, "delta_s,u": first, "delta_u,t": second, "delta_s,t": long,
-            "two legs combined": gamma * first + (1.0 - gamma) * second,
-            "convex weights": (round(keep, 4), round(move, 4)),
-            "convex form": keep * x + move * long, "X_s,t(x)": x * math.exp(t - s)}
-
-
-# %% 04. One term of the loss, on the DNA alphabet
-#
-# The composed target is a convex combination of two simplex points, so it is a distribution and
-# the divergence against it needs no projection. Its gradient pushes mass onto the token with more
-# weight in the two-leg decomposition than in the long prediction.
-def worked_kl():
-    """The chapter's four-token example, with the target composed from two student predictions."""
-    first = torch.tensor([0.10, 0.20, 0.60, 0.10])  # delta_{0.2,0.5} over (A, C, G, T).
-    second = torch.tensor([0.05, 0.10, 0.80, 0.05])  # delta_{0.5,0.8} at the state the leg reaches.
-    gamma = (1.0 - 0.8) * (0.5 - 0.2) / ((1.0 - 0.5) * (0.8 - 0.2))
-    target = gamma * first + (1.0 - gamma) * second  # [4] the composed target.
-    student = torch.tensor([0.10, 0.15, 0.65, 0.10])  # The student's long prediction.
-    terms = target * (target / student).log()  # [4] the four terms of the divergence.
-    return {"worked gamma": gamma, "composed target": [round(v, 4) for v in target.tolist()],
-            "target mass": float(target.sum()), "worked KL": float(terms.sum()),
-            "the four terms": [round(v, 6) for v in terms.tolist()]}
-
-
-# %% 05. The two gradients a softmax head sees
-#
-# With a softmax head the KL gradient in the logits is q - p, so the size of the correction in
-# coordinate i is the error in coordinate i. The squared error passes through the softmax Jacobian
-# q_j(delta_ij - q_i) and carries a factor q_i in every term, so a token the model has learned to
-# rule out gets almost no correction. Both gradients are taken by automatic differentiation here,
-# which checks the two expressions of Section 6.27 as well as evaluating them.
-def logit_gradients(target, student):
-    """The gradient of KL(p || q) and of half the squared error in the logits of q. Two [K]."""
-    for_kl = student.log().clone().requires_grad_(True)  # Logits whose softmax is the student.
-    divergence = (target * (target / torch.softmax(for_kl, -1)).log()).sum()
-    kl_gradient, = torch.autograd.grad(divergence, for_kl)  # [K] equal to q - p.
-    for_squared = student.log().clone().requires_grad_(True)  # A second graph on the same logits.
-    squared = 0.5 * ((torch.softmax(for_squared, -1) - target) ** 2).sum()
-    squared_gradient, = torch.autograd.grad(squared, for_squared)  # [K] with a q_i throughout.
-    return kl_gradient, squared_gradient
-
-
-def gradient_comparison(floor: float = 1e-6):
-    """The two gradients at the chapter's student, and again once that student has saturated."""
-    target = torch.tensor([0.06, 0.12, 0.76, 0.06])  # p, the composed target of Section 6.27.
-    student = torch.tensor([0.10, 0.15, 0.65, 0.10])  # q, the student's long prediction.
-    kl, squared = logit_gradients(target, student)
-    ratio = (squared / kl).abs()  # Between a quarter and a half, coordinate by coordinate.
-    rest = student[1:] * (1.0 - floor) / student[1:].sum()  # The other three keep their shares.
-    saturated = torch.cat([torch.full_like(student[:1], floor), rest])  # [K], summing to one.
-    kl_low, squared_low = logit_gradients(target, saturated)
-    return {
-        "KL gradient": [round(v, 5) for v in kl.tolist()],
-        "squared-error gradient": [round(v, 5) for v in squared.tolist()],
-        "squared over KL": [round(v, 4) for v in ratio.tolist()],
-        "saturated student": [round(v, 6) for v in saturated.tolist()],
-        "saturated KL gradient": float(kl_low[0]),
-        # The surviving squared-error gradient is of order 1e-08, which four printed decimals
-        # cannot show, so the report carries it in units of 1e-08.
-        "saturated squared-error gradient, in 1e-08": 1e8 * float(squared_low[0]),
-        "orders of magnitude between them": math.log10(
-            abs(float(kl_low[0])) / abs(float(squared_low[0]))),
-    }
-
-
-# %% 06. The decoding clock
-#
-# A position decodes correctly when its data coordinate is the largest of the |V| coordinates,
-# a shifted standard normal against the maximum of the rest. The reparameterized time tau is the
-# fraction of the initial decoding error already removed, and it is the clock the sampler uses.
-def decoding_error(t: float, vocab_size: int, nodes: int = 4001) -> float:
-    """P_e(t), the probability that one position decodes to the wrong token at time t."""
-    if t >= 1.0:
-        return 0.0
-    shift = t / (1.0 - t)  # a = t/(1-t), the gap the data coordinate opens.
-    z = torch.linspace(-12.0, 14.0, nodes, dtype=torch.float64)  # Quadrature grid over the draw.
-    density = torch.exp(-0.5 * z * z) / math.sqrt(2.0 * math.pi)  # phi(z).
-    cumulative = 0.5 * torch.erfc(-(z + shift) / math.sqrt(2.0))  # Phi(z + a).
-    log_all = (vocab_size - 1) * cumulative.clamp(min=1e-300).log()  # All the others fall short.
-    return 1.0 - float(torch.trapz(density * log_all.exp(), z))
-
-
-def clock(vocab_size: int):
-    """tau(t) = (P_e(0) - P_e(t))/P_e(0), as a callable, with the source error rate beside it."""
-    start = decoding_error(0.0, vocab_size)
-    return (lambda t: (start - decoding_error(t, vocab_size)) / start), start
-
-
-def even_in_tau(vocab_size: int, pieces: int = 4):
-    """Solve tau(t) = k/pieces by bisection and return the grid with its two endpoints."""
-    tau, _ = clock(vocab_size)
-    grid = [0.0]
-    for k in range(1, pieces):
-        low, high = 0.0, 1.0 - 1e-6
-        for _ in range(60):  # Bisection; tau is increasing in t.
-            middle = 0.5 * (low + high)
-            low, high = (middle, high) if tau(middle) < k / pieces else (low, middle)
-        grid.append(0.5 * (low + high))
-    return grid + [1.0]
-
-
-# %% 07. The model, a two-time denoiser with a softmax head
-def predict(model, x, s, t, length, vocab):
-    """delta_{s,t}(x) at every position, a point of the simplex produced by a softmax head."""
-    logits = model(x.flatten(1), s, t).view(-1, length, vocab)  # [B, L, K] the network output.
-    return torch.softmax(logits, dim=-1)  # [B, L, K] on the simplex at every position.
-
-
-def step_map(x, prediction, s, t):
-    """The convex form, which is the flow map written in denoiser coordinates."""
-    keep, move = convex_weights(s, t)
-    return keep[:, None, None] * x + move[:, None, None] * prediction  # [B, L, K].
-
-
-def make_loss(model, data, length, vocab, batch_size, device):
-    eye = torch.eye(vocab, device=device)  # The one-hot lift of the alphabet.
-
-    def loss_fn(step: int) -> torch.Tensor:
-        index = torch.randint(0, len(data), (batch_size,), device=device)
-        x1 = eye[data[index]]  # [B, L, K] the data, lifted.
-        x0 = torch.randn_like(x1)  # [B, L, K] the Gaussian source.
-        t_diag = torch.rand(batch_size, device=device)  # [B] a time for the diagonal term.
-        xt = (1 - t_diag[:, None, None]) * x0 + t_diag[:, None, None] * x1  # [B, L, K].
-        diagonal = -(predict(model, xt, t_diag, t_diag, length, vocab).clamp_min(1e-9).log()
-                     .gather(-1, data[index][..., None])).mean()  # Cross-entropy to the token.
-
-        s = torch.rand(batch_size, device=device)  # [B] start time.
-        t = s + (1.0 - s) * torch.rand(batch_size, device=device)  # [B] end time.
-        # The slice t = 1 is the one-step case the sampler reads, and a continuous draw never
-        # lands on it, so a fixed share of the batch is set there on purpose (Section 6.26).
-        endpoint = torch.rand(batch_size, device=device) < ENDPOINT_FRACTION  # [B] which rows.
-        t = torch.where(endpoint, torch.ones_like(t), t)
-        u = s + (t - s) * torch.rand(batch_size, device=device)  # [B] the split point.
-        xs = (1 - s[:, None, None]) * x0 + s[:, None, None] * x1  # [B, L, K] the interpolant.
-        with torch.no_grad():  # Both legs and the intermediate state are frozen targets.
-            first = predict(model, xs, s, u, length, vocab)
-            second = predict(model, step_map(xs, first, s, u), u, t, length, vocab)
-            gamma = composition_weight(s, u, t)[:, None, None]
-            target = gamma * first + (1.0 - gamma) * second  # [B, L, K], still on the simplex.
-        student = predict(model, xs, s, t, length, vocab).clamp_min(1e-9)
-        off_diagonal = (target * (target.clamp_min(1e-9) / student).log()).sum(-1).mean()
-        return diagonal + off_diagonal
-
-    return loss_fn
-
-
-# %% 08. Sampling on a grid of times
-@torch.no_grad()
-def sample(model, count, length, vocab, grid, device):
-    """Walk the grid with the convex form, and read the alphabet off the last prediction."""
-    x = torch.randn(count, length, vocab, device=device)  # [N, L, K] the Gaussian source.
-    for a, b in zip(grid[:-1], grid[1:]):
-        s = torch.full((count,), a, device=device)
-        t = torch.full((count,), b, device=device)
-        x = step_map(x, predict(model, x, s, t, length, vocab), s, t)
-    return x.argmax(-1)  # [N, L] token ids, read by an argmax at every position.
-
-
-# %% 09. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=1500, batch_size=128, samples=64, length=16)
-    args = parser.parse_args()
-    set_seed(args.seed)
-    device = pick_device(args.device)
-    tokens, _ = synthetic_dna(num_sequences=1024, length=args.length, seed=args.seed)
-    data = tokens.to(device)  # [1024, L] the training sequences.
-
-    report = scalar_cocycle()
-    report.update(worked_kl())
-    report.update(gradient_comparison())
-    if not args.quiet:
-        print(f"at the saturated student the KL gradient is "
-              f"{report['saturated KL gradient']:.4f} and the squared-error gradient "
-              f"{1e-8 * report['saturated squared-error gradient, in 1e-08']:.3e}")
-    book_tau, book_start = clock(BOOK_VOCAB)
-    for t in (0.60, 0.75, 0.80, 0.85):
-        report[f"tau({t:.2f}) at |V|=30522"] = book_tau(t)
-    book_grid = even_in_tau(BOOK_VOCAB)
-    report["four steps on the clock"] = [round(t, 6) for t in book_grid]
-    rounded = [round(t, 3) for t in book_grid]  # The chapter quotes the grid to three decimals.
-    report["convex weights on that grid"] = [
-        (round((1 - b) / (1 - a), 4), round((b - a) / (1 - a), 4))
-        for a, b in zip(rounded[:-1], rounded[1:])]
-    report["gamma on the last three"] = float(composition_weight(
-        *(torch.tensor([v]) for v in rounded[1:4])))
-    report["gamma over the long first leg"] = float(composition_weight(
-        *(torch.tensor([v]) for v in (0.0, rounded[1], rounded[2]))))
-    report["tau on a uniform grid"] = [round(book_tau(t), 5) for t in (0.25, 0.5, 0.75)]
-    if not args.quiet:
-        print(f"clock at |V|=30522: P_e(0) = {book_start:.6f}, four steps at "
-              f"{[round(t, 3) for t in book_grid]}")
-
-    model = TwoTimeMLP(dim=args.length * DNA_SIZE, width=args.width).to(device)
-    if not args.quiet:
-        print(f"\ntraining the two-time denoiser  parameters={count_parameters(model):,}")
-    losses = train(model, make_loss(model, data, args.length, DNA_SIZE, args.batch_size, device),
-                   steps=args.steps, lr=args.lr, quiet=args.quiet)
-
-    dna_grid = even_in_tau(DNA_SIZE)  # The same clock, read at the four-letter alphabet.
-    report["DNA clock grid"] = [round(t, 4) for t in dna_grid]
-    samples: list[str] = []
-    for name, grid in (("one step", [0.0, 1.0]), ("four, uniform", [0.0, 0.25, 0.5, 0.75, 1.0]),
-                       ("four, on the clock", dna_grid)):
-        torch.manual_seed(args.seed + 1)  # The same source noise for every schedule.
-        drawn = sample(model, args.samples, args.length, DNA_SIZE, grid, device)
-        report[f"{name}: GC fraction"] = gc_fraction(drawn)
-        report[f"{name}: motif fraction"] = motif_fraction(drawn)
-        report[f"{name}: token entropy"] = token_entropy(drawn, DNA_SIZE)
-        report[f"{name}: unique fraction"] = unique_fraction(drawn)
-        samples.append(f"{name:<20s} {decode_dna(drawn[:1])[0]}  GC {gc_fraction(drawn):.3f}")
-    report["training GC fraction"] = gc_fraction(data)
-    report["final loss"] = sum(losses[-20:]) / len(losses[-20:])
-
-    print_report("Flow Map Language Models, Chapter 6", report, samples)
-    save_run(args.out, model=model, config=vars(args), report=report, losses=losses, samples=samples)
-
-
-if __name__ == "__main__":
-    main()

lecture_6/meanflow.py

@@ -1,175 +0,0 @@
-"""MeanFlow, the Eulerian identity written in the average-velocity parameterization.
-
-Substituting X_{s,t}(x) = x + (t-s) v_{s,t}(x) into the Eulerian identity leaves one equation,
-v_{s,t}(x) = b(x,s) + (t-s) d/ds v_{s,t}(x), whose correction term is a total derivative along the
-trajectory. One forward-mode pass along the tangent (b, 1, 0) in the arguments (x, s, t) computes
-that derivative without forming a Jacobian, and the target uses the realized velocity x_1 - x_0 of
-the interpolant in place of b, so a step costs one forward pass and one derivative.
-
-Chapter 6, Sections 6.13, 6.15 and 6.18 of the notes. Trained on the four Gaussian blobs of
-``cis6270.data.gaussian_mixture_2d``. Prints the identity worked by hand on dx/dt = x, where the
-average velocity is 1.718x, the instantaneous velocity is x and the correction is 0.718x, and then
-one-step, two-step and four-step samples from the trained map.
-
-    python lecture_6/meanflow.py                 # 2000 steps, about twenty seconds
-    python lecture_6/meanflow.py --steps 20      # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-import math
-
-import torch
-
-from cis6270.data import gaussian_mixture_2d
-from cis6270.metrics import mean_log_density, mode_report
-from cis6270.nets import TwoTimeMLP, count_parameters
-from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
-
-
-torch.set_num_threads(1)  # These networks are small, and splitting one matrix product costs more
-# than it saves. ``tests/conftest.py`` sets the same thing for the test session, so the line
-# here is the one that acts when a reader runs this script directly.
-
-DIAGONAL_FRACTION = 0.25  # Share of each batch drawn with t = s, where the identity is flow matching.
-
-
-# %% 02. The average velocity of the scalar flow, in closed form
-#
-# The exact map of dx/dt = x multiplies the state by e^{t-s}, so the displacement over the elapsed
-# time is x (e^{t-s} - 1)/(t-s). On the diagonal the limit is x, which is b(x,t).
-def scalar_average_velocity(x, s, t):
-    """v_{s,t}(x) = x (e^{t-s} - 1)/(t - s) on the flow dx/dt = x."""
-    elapsed = t - s
-    return x * torch.expm1(elapsed) / elapsed  # expm1 keeps the limit accurate near the diagonal.
-
-
-# %% 03. The identity, checked on the scalar flow by hand
-#
-# v_{s,t}(x) = b(x,s) + (t-s) d/ds v_{s,t}(x), with d/ds = partial_s + b(x,s) . grad_x. At s = 0,
-# t = 1 the partial derivative is -x, the spatial gradient is e - 1, and the transport term is
-# therefore (e-1)x, so the correction is x(e-2) = 0.718x.
-def scalar_check(x_value: float = 1.0):
-    """Return every term of the identity at s = 0, t = 1, each one computed and not quoted."""
-    x = torch.tensor(x_value)
-    s, t = torch.tensor(0.0), torch.tensor(1.0)
-    average = scalar_average_velocity(x, s, t)  # v_{0,1}(x) = x(e-1).
-    instant = x  # b(x, 0) = x on this flow.
-    partial_s = torch.func.jvp(  # Derivative in the start time alone, with x and t held fixed.
-        lambda ss: scalar_average_velocity(x, ss, t), (s,), (torch.ones_like(s),))[1]
-    gradient_x = torch.func.jvp(  # Derivative in the state alone, the other half of d/ds.
-        lambda xx: scalar_average_velocity(xx, s, t), (x,), (torch.ones_like(x),))[1]
-    transport = instant * gradient_x  # b(x,s) . grad_x v, the second piece of the total derivative.
-    total = partial_s + transport  # d/ds v_{s,t}(x) = x(e - 2).
-    return {
-        "v_0,1(1)": float(average),
-        "b(1, 0)": float(instant),
-        "partial_s v": float(partial_s),
-        "grad_x v_0,1": float(gradient_x),
-        "d/ds v_0,1": float(total),
-        "identity gap": float((average - (instant + (t - s) * total)).abs()),
-        "one map evaluation": float(x + (t - s) * average),
-        "one Euler step": float(x + (t - s) * instant),
-        "e": math.e,
-    }
-
-
-# %% 04. The Gaussian flow, where the correction vanishes
-#
-# A linear interpolant between paired Gaussians produces straight trajectories, and on a straight
-# trajectory the average velocity equals the initial velocity, so the total derivative is zero.
-def gaussian_check(x_value: float = 0.7):
-    """v_{0,t}(x) = 2 + 2x for every end time, and b(x,0) is the same number."""
-    x = torch.tensor(x_value)
-    average = torch.stack([  # Three end times, all giving the same average velocity.
-        (2.0 * t + (1.0 + 2.0 * t) * x - x) / t for t in (0.25, 0.5, 1.0)])
-    instant = 2.0 + 2.0 * (x - 0.0) / 1.0  # b(x, 0) on the Gaussian flow.
-    return {"v_0,t(0.7) spread": float(average.max() - average.min()),
-            "v_0,1(0.7)": float(average[-1]), "b(0.7, 0)": float(instant)}
-
-
-# %% 05. The MeanFlow target, by one forward-mode pass
-#
-# The tangent (x_1 - x_0, 1, 0) in the arguments (x, s, t) contracts the full Jacobian of v into
-# the total derivative along the trajectory: the realized velocity in the state slot, one in the
-# start-time slot, and zero in the end-time slot because the end time is held fixed.
-def mean_flow_target(model, xs, s, t, realized):
-    """b(x,s) + (t-s) d/ds v_{s,t}(x), detached, which is the right-hand side of the identity."""
-    f = lambda xx, ss: model(xx, ss, t)
-    _, total = torch.func.jvp(f, (xs, s), (realized, torch.ones_like(s)))  # [B, d] d/ds v.
-    return (realized + (t - s)[:, None] * total).detach()  # [B, d] the target, held fixed.
-
-
-def make_loss(model, data, batch_size, device):
-    """Draw a coupling, draw times from the triangle, and regress v onto its own identity."""
-
-    def loss_fn(step: int) -> torch.Tensor:
-        index = torch.randint(0, len(data), (batch_size,), device=device)
-        x1 = data[index]  # [B, 2] the data draw.
-        x0 = torch.randn_like(x1)  # [B, 2] the source draw.
-        s = torch.rand(batch_size, device=device)  # [B] start time.
-        t = s + (1.0 - s) * torch.rand(batch_size, device=device)  # [B] end time, s <= t <= 1.
-        diagonal = torch.rand(batch_size, device=device) < DIAGONAL_FRACTION  # [B] which rows anchor.
-        t = torch.where(diagonal, s, t)  # On those rows the correction term has a factor of zero.
-        xs = (1 - s[:, None]) * x0 + s[:, None] * x1  # [B, 2] the interpolant I_s.
-        realized = x1 - x0  # [B, 2] the derivative of the interpolant, available in closed form.
-        target = mean_flow_target(model, xs, s, t, realized)
-        return ((model(xs, s, t) - target) ** 2).mean()
-
-    return loss_fn
-
-
-# %% 06. Sampling, at any step count the reader chooses
-@torch.no_grad()
-def sample(model, x0, steps: int):
-    """Cross [0, 1] in equal pieces, one evaluation of the map per piece."""
-    x = x0
-    grid = torch.linspace(0.0, 1.0, steps + 1, device=x0.device)  # [steps+1] the time grid.
-    for s, t in zip(grid[:-1], grid[1:]):
-        x = model.flow_map(x, s.expand(len(x)), t.expand(len(x)))  # X_{s,t}(x).
-    return x  # [N, 2] generated points.
-
-
-# %% 07. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=2000, batch_size=128, samples=512)
-    args = parser.parse_args()
-    set_seed(args.seed)
-    device = pick_device(args.device)
-    data = gaussian_mixture_2d(num_points=2048, seed=args.seed).to(device)  # [2048, 2].
-
-    report = scalar_check()
-    report.update(gaussian_check())
-    if not args.quiet:
-        print(f"on dx/dt = x at s=0, t=1: v = {report['v_0,1(1)']:.6f} = b {report['b(1, 0)']:.6f}"
-              f" + correction {report['d/ds v_0,1']:.6f}")
-        print(f"one map evaluation {report['one map evaluation']:.6f} against one Euler step "
-              f"{report['one Euler step']:.6f}, exact e = {math.e:.6f}")
-
-    model = TwoTimeMLP(dim=2, width=args.width).to(device)
-    if not args.quiet:
-        print(f"\ntraining MeanFlow  parameters={count_parameters(model):,}")
-    losses = train(model, make_loss(model, data, args.batch_size, device),
-                   steps=args.steps, lr=args.lr, quiet=args.quiet)
-
-    torch.manual_seed(args.seed + 1)
-    x0 = torch.randn(args.samples, 2, device=device)  # [N, 2] the source points every run shares.
-    samples: list[str] = []
-    for steps in (1, 2, 4):
-        points = sample(model, x0, steps)
-        stats = mode_report(points)
-        report[f"{steps}-step nearest distance"] = stats["nearest_distance"]
-        report[f"{steps}-step mode entropy"] = stats["mode_entropy"]
-        report[f"{steps}-step log density"] = mean_log_density(points.cpu())
-        samples.append(f"{steps} step(s): nearest {stats['nearest_distance']:.3f}  "
-                       f"modes {stats['modes_found']}  entropy {stats['mode_entropy']:.3f}")
-    one, four = sample(model, x0, 1), sample(model, x0, 4)
-    report["one step against four"] = float((one - four).norm(dim=1).mean())
-    report["final loss"] = sum(losses[-20:]) / len(losses[-20:])
-
-    print_report("MeanFlow, Chapter 6", report, samples)
-    save_run(args.out, model=model, config=vars(args), report=report, losses=losses, samples=samples)
-
-
-if __name__ == "__main__":
-    main()

lecture_6/posterior_map.py

@@ -1,221 +0,0 @@
-"""Diamond Maps and Meta Flow Maps, the posterior over endpoints that guidance needs.
-
-Steering a flow toward a reward needs the value function, the log expected exponential reward over
-the terminal states reachable from the current one, and that average runs over a posterior. A
-deterministic map returns one endpoint, so the only estimate it supports is the reward at that
-endpoint, which is low by the Jensen gap. Giving the map a fresh noise argument and indexing it by
-the state it is conditioned on turns it into a family of inner flow maps, one per conditioning
-pair, whose draws are posterior samples.
-
-Chapter 6, Sections 6.33 and 6.34 of the notes. Runs on a two-point reward problem and on the four
-Gaussian blobs of ``cis6270.data.gaussian_mixture_2d``. Prints the Jensen gap of 0.434, the bias
-of the K-sample value at K = 1, 2, 4 and 8, the softmax-weighted value gradient of -0.4483, the
-importance weight that pulls a renoised proposal from 0.420 back to 0.0131, and the stop-gradient
-fixed point 0.625 of Meta Flow Maps.
-
-    python lecture_6/posterior_map.py                 # 16 inner samples per state
-    python lecture_6/posterior_map.py --steps 20      # fewer inner substeps, for a first look
-"""
-
-# %% 01. Imports and configuration
-import math
-
-import torch
-
-from cis6270.data import gaussian_mixture_2d
-from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed
-
-torch.set_num_threads(1)  # The work here is small, and one thread keeps the timing stable.
-# ``tests/conftest.py`` sets the same thing for the test session, so the line here is the one
-# that acts when a reader runs this script directly.
-torch.set_default_dtype(torch.float64)  # Exact posteriors, compared against sample averages.
-
-
-INNER_STEPS = 48  # Substeps of the inner flow, which carries noise to a posterior sample.
-
-
-# %% 02. The tilt, the value function and the Jensen gap
-#
-# The tilted target is the data times the exponential reward, and the same reweighting acts on the
-# posterior over endpoints at every intermediate state. A deterministic map returns the posterior
-# mean, so its estimate of the value is the reward there, and the exponential is convex, so that
-# estimate is always the low one.
-def two_point(reward_gap: float = 2.0, x: float = 0.7, t: float = 0.8):
-    """The chapter's two-point example, with z in {0, 2}, r(z) = z and a linear schedule."""
-    weights = torch.tensor([1.0, math.exp(reward_gap)])  # e^{r(0)} and e^{r(2)}.
-    tilted = weights / weights.sum()  # The tilted target over the two points.
-    far = math.exp(-((x - t * 2.0) ** 2) / (2.0 * (1.0 - t) ** 2))  # Unnormalized, at z = 2.
-    near = math.exp(-(x ** 2) / (2.0 * (1.0 - t) ** 2))  # Unnormalized, at z = 0.
-    posterior = torch.tensor([near, far]) / (near + far)  # The posterior over endpoints at x.
-    value = math.log(float((posterior * weights).sum()))  # V_t(x) at that state.
-    tilted_posterior = posterior * weights / float((posterior * weights).sum())
-    exact = math.log(0.5 + 0.5 * math.exp(reward_gap))  # V over the data itself.
-    return {
-        "tilted target": [round(v, 4) for v in tilted.tolist()],
-        "posterior at x=0.7": [round(v, 4) for v in posterior.tolist()],
-        "tilted posterior": [round(v, 4) for v in tilted_posterior.tolist()],
-        "V_t(0.7)": value,
-        "V over the data": exact,
-        "single-endpoint estimate": 1.0 * reward_gap / 2.0,  # r at the posterior mean z = 1.
-        "Jensen gap": exact - reward_gap / 2.0,
-    }
-
-
-def value_bias(reward_gap: float = 2.0):
-    """E[V_K] over K independent posterior draws, which approaches the exact value like 1/K."""
-    exact = math.log(0.5 + 0.5 * math.exp(reward_gap))
-    out = {}
-    for k in (1, 2, 4, 8):
-        expectation = sum(
-            math.comb(k, j) * 2.0 ** (-k)
-            * math.log(((k - j) + j * math.exp(reward_gap)) / k) for j in range(k + 1))
-        out[f"E[V_{k}]"] = expectation
-        out[f"bias at K={k}"] = exact - expectation
-    return out
-
-
-# %% 03. The value gradient from K samples
-#
-# Differentiating the logarithm of an average of exponentials produces a softmax over the rewards,
-# so a sample whose reward is far above the others dominates the direction and a set with equal
-# rewards gives their plain average.
-def value_gradient(rewards, gradients):
-    """The softmax weights over the rewards, and the weighted average of the reward gradients."""
-    weights = torch.softmax(rewards, dim=0)  # [K].
-    return weights, float((weights * gradients).sum()), float(gradients.mean())
-
-
-# %% 04. The inner clock
-#
-# For a fixed conditioning pair the map is an ordinary two-time flow map on an inner time axis
-# running from the noise draw at 0 to a posterior sample at 1. Its velocity is the flow matching
-# velocity of the posterior, which is available in closed form over the training points.
-def posterior_weights(x, t, data):
-    """p_{1|t}(z_i | x), a softmax over the training points under the linear interpolant. [B, N]."""
-    spread = max(1.0 - t, 1e-6)
-    squared = ((x[:, None, :] - t * data[None]) ** 2).sum(-1)  # [B, N].
-    return torch.softmax(-squared / (2.0 * spread ** 2), dim=1)  # [B, N].
-
-
-def inner_map(noise, weights, data, substeps: int):
-    """Y_{0,1}(eps | x, t), the inner flow from a Gaussian draw to a posterior sample. [B, K, 2]."""
-    y = noise  # [B, K, 2] one inner state per sample.
-    step = 1.0 / substeps
-    for i in range(substeps):
-        a = i * step  # The inner time.
-        squared = ((y[:, :, None, :] - a * data[None, None]) ** 2).sum(-1)  # [B, K, N].
-        inner = torch.softmax(torch.log(weights[:, None, :].clamp_min(1e-30))
-                              - squared / (2.0 * max(1.0 - a, 1e-6) ** 2), dim=-1)  # [B, K, N].
-        denoiser = inner @ data  # [B, K, 2] the inner posterior mean.
-        y = y + step * (denoiser - y) / max(1.0 - a, 1e-6)  # Euler on the inner velocity.
-    return y
-
-
-# %% 05. Posterior samples from a map we already have
-#
-# We can renoise the current state back to an earlier time with fresh noise and push it through
-# the existing map, which discards information and widens the law of the endpoint. The correction
-# is an importance weight, written as a shift of the reward by the log likelihood of the current
-# state given that endpoint.
-def renoise_example(x: float = 0.7, t: float = 0.8, earlier: float = 0.6):
-    """The chapter's renoised proposal, and the importance weight that recovers the posterior."""
-    scale = earlier / t  # alpha_{t'}/alpha_t on the linear schedule.
-    variance = (1.0 - earlier) ** 2 - scale ** 2 * (1.0 - t) ** 2  # The variance added back.
-    mean = scale * x  # The center of the proposal.
-    boundary = (earlier * 2.0 / 2.0 - mean) / math.sqrt(variance)  # Where the map switches mode.
-    proposal_far = 0.5 * math.erfc(boundary / math.sqrt(2.0))  # Share of draws reaching z = 2.
-    penalty_far = (x - t * 2.0) ** 2 / (2.0 * (1.0 - t) ** 2)  # Log likelihood of x given z = 2.
-    penalty_near = x ** 2 / (2.0 * (1.0 - t) ** 2)  # The same for z = 0.
-    far = proposal_far * math.exp(-penalty_far)
-    near = (1.0 - proposal_far) * math.exp(-penalty_near)
-    return {"renoise scale": scale, "renoise variance": variance,
-            "renoise standard deviation": math.sqrt(variance),
-            "proposal weight on the far mode": proposal_far,
-            "penalty at z=2": penalty_far, "penalty at z=0": penalty_near,
-            "corrected weight on the far mode": far / (far + near)}
-
-
-# %% 06. Meta Flow Maps, the stop-gradient form
-#
-# The controlled velocity needs a ratio of two expectations, and regressing onto a per-sample
-# ratio is biased at every sample size. Multiplying through by the denominator leaves a condition
-# with one stop-gradient whose fixed point is the ratio of the expectations.
-def stop_gradient_fixed_point(weights, gradients, diffusion_squared: float = 2.0):
-    """The fixed point of the loss, the per-sample ratio it is compared against, and the gap."""
-    exact = 0.5 * diffusion_squared * float(gradients.mean()) / float(weights.mean())
-    naive = float((gradients / weights).mean()) * 0.5 * diffusion_squared
-    covariance = float(gradients.mean()) - float((gradients / weights).mean() * weights.mean())
-    return {"E[w]": float(weights.mean()), "E[G]": float(gradients.mean()),
-            "stop-gradient fixed point": exact, "per-sample ratio average": naive,
-            "gap": naive - exact, "covariance of G/w and w": covariance,
-            "gap from the covariance": -covariance / float(weights.mean())}
-
-
-# %% 07. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=INNER_STEPS, samples=16)
-    args = parser.parse_args()  # --steps is the inner substep count here, since nothing trains.
-    set_seed(args.seed)
-    pick_device(args.device)  # Every number here is double precision, which not every
-    # accelerator carries, so the work stays on the CPU whatever --device names.
-    substeps = max(args.steps, 8)
-
-    report = two_point()
-    report.update(value_bias())
-    weights, weighted, plain = value_gradient(torch.tensor([0.2, 1.5, 0.9]),
-                                              torch.tensor([0.4, -1.2, 0.5]))
-    report["softmax weights"] = [round(v, 4) for v in weights.tolist()]
-    report["weighted value gradient"] = weighted
-    report["plain average gradient"] = plain
-    report.update(renoise_example())
-    report.update(stop_gradient_fixed_point(torch.tensor([1.0, 2.0, 5.0]),
-                                            torch.tensor([0.4, 1.0, 3.6])))
-    if not args.quiet:
-        print(f"Jensen gap {report['Jensen gap']:.3f} between the exact value "
-              f"{report['V over the data']:.4f} and the single endpoint "
-              f"{report['single-endpoint estimate']:.3f}")
-        print(f"renoised proposal {report['proposal weight on the far mode']:.3f} corrected to "
-              f"{report['corrected weight on the far mode']:.4f} against the posterior "
-              f"{report['posterior at x=0.7'][1]:.4f}")
-
-    # %% 08. The same quantities on the four Gaussian blobs
-    data = gaussian_mixture_2d(num_points=256, seed=args.seed).double()  # [N, 2] the endpoints.
-    outer_t = 0.6  # The outer time the sampler is conditioned at.
-    torch.manual_seed(args.seed + 1)
-    index = torch.randint(0, len(data), (32,))
-    states = ((1.0 - outer_t) * torch.randn(32, 2) + outer_t * data[index])  # [32, 2] contexts.
-    states.requires_grad_(True)
-    exact_weights = posterior_weights(states, outer_t, data)  # [32, N] the exact posterior.
-    reward = data[:, 0]  # r(z) = the first coordinate, a reward with a clear direction.
-    exact_value = torch.log((exact_weights * reward.exp()[None, :]).sum(1))  # [32] V_t(x).
-    exact_grad = torch.autograd.grad(exact_value.sum(), states, retain_graph=True)[0]  # [32, 2].
-
-    noise = torch.randn(32, args.samples, 2)  # [32, K, 2] inner draws, one per posterior sample.
-    drawn = inner_map(noise, exact_weights, data, substeps)  # [32, K, 2] posterior samples.
-    sampled_reward = drawn[:, :, 0]  # [32, K] the reward at each sample.
-    estimate = torch.logsumexp(sampled_reward, dim=1) - math.log(args.samples)  # [32] V hat.
-    sample_grad = torch.autograd.grad(estimate.sum(), states, retain_graph=True)[0]  # [32, 2].
-    deterministic = (exact_weights @ reward)  # [32] the reward at the posterior mean.
-
-    exact_value, estimate = exact_value.detach(), estimate.detach()  # Past the two gradients.
-    deterministic, drawn = deterministic.detach(), drawn.detach()
-    report["exact V, averaged"] = float(exact_value.mean())
-    report[f"V from {args.samples} samples"] = float(estimate.mean())
-    report["single endpoint, averaged"] = float(deterministic.mean())
-    report["Jensen gap on the blobs"] = float((exact_value - deterministic).mean())
-    report["posterior mean error"] = float(
-        (drawn.mean(1) - exact_weights.detach() @ data).norm(dim=1).mean())
-    report["value gradient error"] = float((sample_grad - exact_grad).norm(dim=1).mean())
-    report["exact value gradient norm"] = float(exact_grad.norm(dim=1).mean())
-    states = states.detach()
-    samples = [f"state ({float(x[0]):+.2f}, {float(x[1]):+.2f})  exact V {float(v):.3f}  "
-               f"{args.samples} samples {float(e):.3f}  one endpoint {float(d):.3f}"
-               for x, v, e, d in zip(states[:4], exact_value[:4], estimate[:4], deterministic[:4])]
-
-    print_report("Diamond Maps and Meta Flow Maps, Chapter 6", report, samples)
-    save_run(args.out, config=vars(args), report=report, samples=samples)
-
-
-if __name__ == "__main__":
-    main()

lecture_6/residuals.py

@@ -1,204 +0,0 @@
-"""The three residuals of a flow map, and the diagonal anchor added to each one.
-
-The cocycle differentiated in the end time is the Lagrangian identity, differentiated in the start
-time it is the Eulerian identity, and left finite it is the composition identity. Each becomes a
-loss by squaring the amount the model violates it. None of the three mentions the data, so each is
-satisfied by a model that is self-consistent and wrong, and a flow matching loss on the diagonal
-supplies the one piece of information about the target.
-
-Chapter 6, Sections 6.15 to 6.17 and 6.23 of the notes. Trained on the four Gaussian blobs
-of ``cis6270.data.gaussian_mixture_2d``. Prints the stay-still residuals of Section 6.17, which are
-1, 1 and 0, and then four trained models side by side in one-step sampling.
-
-    python lecture_6/residuals.py                 # four models, 600 steps each
-    python lecture_6/residuals.py --steps 20      # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-import torch
-
-from cis6270.data import GAUSSIAN_STD, gaussian_mixture_2d
-from cis6270.metrics import mode_report
-from cis6270.nets import TwoTimeMLP, count_parameters
-from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
-
-
-torch.set_num_threads(1)  # These networks are small, and splitting one matrix product costs more
-# than it saves. ``tests/conftest.py`` sets the same thing for the test session, so the line
-# here is the one that acts when a reader runs this script directly.
-
-ON_MODE = 3.0 * GAUSSIAN_STD  # A sample counts as on a mode within three data standard deviations.
-ODE_STEPS = 128  # Substeps used to integrate the model's own diagonal velocity accurately.
-
-
-# %% 02. The model, whose identity axiom is free
-#
-# The network produces the average velocity v_{s,t}(x), and the map is the state plus the elapsed
-# time times that velocity. At t = s the second term vanishes whatever the weights are, so no term
-# of the loss is spent on the identity axiom.
-class Map(torch.nn.Module):
-    """X_{s,t}(x) = x + (t - s) v_{s,t}(x), with v the two-time network of ``cis6270.nets``."""
-
-    def __init__(self, dim: int = 2, width: int = 128):
-        super().__init__()
-        self.net = TwoTimeMLP(dim=dim, width=width)  # Reads x, s and t; returns a vector.
-
-    def v(self, x, s, t):  # the network
-        return self.net(x, s, t)  # [B, d] the average velocity over [s, t].
-
-    def X(self, x, s, t):  # identity axiom, free
-        return self.net.flow_map(x, s, t)  # [B, d] = x + (t - s) v_{s,t}(x).
-
-
-# %% 03. The diagonal anchor, which is the tangent axiom
-#
-# I_t = (1-t) x_0 + t x_1 and its derivative x_1 - x_0 are available in closed form, so the anchor
-# costs one network call and no simulation. Its minimizer is E[I_dot | I_t], the true velocity.
-def anchor(m, x0, x1, t):  # tangent axiom
-    xt = (1 - t[:, None]) * x0 + t[:, None] * x1  # [B, d] the interpolant at time t.
-    return ((m.v(xt, t, t) - (x1 - x0)) ** 2).mean()
-
-
-# %% 04. The Lagrangian residual
-#
-# d/dt X = b at the output. The closure holds the state and the start time fixed, and a tangent of
-# one in t returns the map and its end-time derivative from a single forward-mode pass. The target
-# is the model's own diagonal velocity, detached, so the diagonal trains only through the anchor.
-def lagrangian(m, xs, s, t):  # d/dt X = b at the output
-    f = lambda tt: m.X(xs, s, tt)
-    X, dXdt = torch.func.jvp(f, (t,), (torch.ones_like(t),))
-    return ((dXdt - m.v(X, t, t).detach()) ** 2).mean()
-
-
-# %% 05. The Eulerian residual
-#
-# d/ds X + b . grad X = 0. Both the state and the start time are open in the closure, and the
-# tangent (v, 1) makes the one returned directional derivative the whole left-hand side. The
-# identity is already written as zero, so the residual needs no target at all.
-def eulerian(m, xs, s, t):  # d/ds X + b . grad X = 0
-    vs = m.v(xs, s, s).detach()
-    f = lambda xx, ss: m.X(xx, ss, t)
-    _, dX = torch.func.jvp(f, (xs, s), (vs, torch.ones_like(s)))
-    return (dX ** 2).mean()
-
-
-# %% 06. The composition residual
-#
-# Two short maps against one long one, with no derivative of any kind. The detach covers the whole
-# composed target, so the two legs are frozen and only the long map is corrected.
-def composition(m, xs, s, u, t):  # two short maps, one long one
-    tgt = m.X(m.X(xs, s, u), u, t).detach()
-    return ((m.X(xs, s, t) - tgt) ** 2).mean()
-
-
-# %% 07. The stay-still solution each residual has to rule out
-#
-# The average velocity equals the true b when the two times coincide and zero everywhere else, so
-# the map is the identity for every pair of times and the sampler generates the source. The
-# identity axiom, the cocycle and the anchor all hold on it, and the two differential residuals
-# separate it from the truth by a full unit.
-class StayStill(torch.nn.Module):
-    """v_{s,t}(x) = b(x,t) on the diagonal and 0 off it, on the scalar flow dx/dt = x."""
-
-    def v(self, x, s, t):
-        on_diagonal = (s == t)[:, None]  # [B, 1] true where the two times meet.
-        return torch.where(on_diagonal, x, torch.zeros_like(x))  # b(x,t) = x on this flow.
-
-    def X(self, x, s, t):
-        return x  # Off the diagonal (t-s) times zero, on it zero times b, so the map is the identity.
-
-
-def stay_still_residuals():
-    """The three residuals of Section 6.17 at x = 1, s = 0, t = 1, which are 1, 1 and 0."""
-    m = StayStill()
-    x = torch.ones(1, 1)  # [1, 1] the state x = 1.
-    s, t, u = torch.zeros(1), torch.ones(1), torch.full((1,), 0.5)
-    diagonal_gap = float(((m.v(x, t, t) - x) ** 2).mean())  # The anchor, at its minimum.
-    return {
-        "stay-still Lagrangian": float(lagrangian(m, x, s, t)),
-        "stay-still Eulerian": float(eulerian(m, x, s, t)),
-        "stay-still composition": float(composition(m, x, s, u, t)),
-        "stay-still diagonal gap": diagonal_gap,
-    }
-
-
-# %% 08. One training step, the same shape whichever residual we add
-def make_loss(model, data, residual, batch_size, device):
-    """Draw a coupling, draw times from the triangle s <= t, and add the residual to the anchor."""
-
-    def loss_fn(step: int) -> torch.Tensor:
-        index = torch.randint(0, len(data), (batch_size,), device=device)
-        x1 = data[index]  # [B, 2] the data draw.
-        x0 = torch.randn_like(x1)  # [B, 2] the independent source draw.
-        s = torch.rand(batch_size, device=device)  # [B] start time.
-        t = s + (1.0 - s) * torch.rand(batch_size, device=device)  # [B] end time, s <= t <= 1.
-        u = s + (t - s) * torch.rand(batch_size, device=device)  # [B] intermediate time.
-        total = anchor(model, x0, x1, torch.rand(batch_size, device=device))  # The data term.
-        if residual is None:
-            return total  # The diagonal alone, which is flow matching read in one step.
-        xs = (1 - s[:, None]) * x0 + s[:, None] * x1  # [B, 2] the interpolant at the start time.
-        if residual is composition:
-            return total + residual(model, xs, s, u, t)
-        return total + residual(model, xs, s, t)
-
-    return loss_fn
-
-
-# %% 09. One-step sampling, and the model's own ordinary differential equation
-@torch.no_grad()
-def evaluate(model, x0):
-    """One-step output, the endpoint of the model's own diagonal velocity, and the mode statistics."""
-    from cis6270.data import GAUSSIAN_CENTERS
-
-    zeros = torch.zeros(len(x0), device=x0.device)
-    one_step = model.X(x0, zeros, torch.ones_like(zeros))  # [N, 2] a single evaluation of the map.
-    x, h = x0.clone(), 1.0 / ODE_STEPS  # Integrate dx/dt = v_{t,t}(x) with the same weights.
-    for k in range(ODE_STEPS):
-        t = torch.full((len(x0),), k * h, device=x0.device)
-        x = x + h * model.v(x, t, t)  # Euler on the diagonal velocity the anchor trains.
-    distance = float((one_step - x).norm(dim=1).mean())  # Column two of Table 6.2.
-    nearest = torch.cdist(one_step, GAUSSIAN_CENTERS.to(one_step)).min(dim=1).values  # [N].
-    on_mode = float((nearest < ON_MODE).to(torch.get_default_dtype()).mean())  # Column three of Table 6.2.
-    return one_step, distance, on_mode, mode_report(one_step)
-
-
-# %% 10. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(samples=512)
-    args = parser.parse_args()
-    set_seed(args.seed)
-    device = pick_device(args.device)
-    data = gaussian_mixture_2d(num_points=2048, seed=args.seed).to(device)  # [2048, 2] the target.
-
-    report = stay_still_residuals()  # Section 6.17, before any training.
-    if not args.quiet:
-        print(f"stay-still at x=1, s=0, t=1: Lagrangian {report['stay-still Lagrangian']:.3f}  "
-              f"Eulerian {report['stay-still Eulerian']:.3f}  "
-              f"composition {report['stay-still composition']:.3f}")
-
-    recipes = [("the diagonal alone", None), ("+ Lagrangian", lagrangian),
-               ("+ Eulerian", eulerian), ("+ composition", composition)]
-    samples: list[str] = []
-    for name, residual in recipes:
-        set_seed(args.seed)  # Same starting weights, same batches and same noise for all four.
-        model = Map(dim=2, width=args.width).to(device)
-        if not args.quiet:
-            print(f"\ntraining {name}  parameters={count_parameters(model):,}")
-        losses = train(model, make_loss(model, data, residual, args.batch_size, device),
-                       steps=args.steps, lr=args.lr, quiet=args.quiet)
-        torch.manual_seed(args.seed + 1)  # The same source points for every model.
-        x0 = torch.randn(args.samples, 2, device=device)
-        _, distance, on_mode, stats = evaluate(model, x0)
-        report[f"{name}: distance to its own ODE"] = distance
-        report[f"{name}: fraction on a mode"] = on_mode
-        report[f"{name}: final loss"] = sum(losses[-20:]) / len(losses[-20:])
-        samples.append(f"{name:<20s} distance {distance:.3f}  on a mode {on_mode:.2f}  "
-                       f"modes found {stats['modes_found']}")
-
-    print_report("Three residuals and the diagonal anchor, Chapter 6", report, samples)
-    save_run(args.out, config=vars(args), report=report, samples=samples)
-
-
-if __name__ == "__main__":
-    main()

lecture_6/shortcut.py

@@ -1,184 +0,0 @@
-"""Shortcut models, the composition identity indexed by step size.
-
-The second input of the network is the length of the step and not the time at which it ends, and
-the identity is enforced only between a step and the step twice as long. One doubling constraint
-suffices because the ladder chains: with a finest step of 1/128 and seven doublings, the one-step
-map at level seven rests on seven relations and one regression against the data.
-
-Chapter 6, Section 6.20 of the notes. Trained on the four Gaussian blobs of
-``cis6270.data.gaussian_mixture_2d``. Prints the doubling worked on dx/dt = x, where two steps of
-0.25 average to 1.297443 and the direct step of 0.5 returns the same number, then the dyadic ladder
-and the trained model's doubling residual at every rung.
-
-    python lecture_6/shortcut.py                 # 2000 steps, about twenty seconds
-    python lecture_6/shortcut.py --steps 20      # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-import math
-
-import torch
-
-from cis6270.data import gaussian_mixture_2d
-from cis6270.metrics import mean_log_density, mode_report
-from cis6270.nets import TwoTimeMLP, count_parameters
-from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
-
-
-torch.set_num_threads(1)  # These networks are small, and splitting one matrix product costs more
-# than it saves. ``tests/conftest.py`` sets the same thing for the test session, so the line
-# here is the one that acts when a reader runs this script directly.
-
-FINEST = 1.0 / 128.0  # d_0, the finest step of the ladder.
-LEVELS = 7  # log2(128) doublings, so the step sizes are d_k = 2^k d_0 for k = 0 to 7.
-ZERO_FRACTION = 0.75  # Three quarters of each batch is drawn at step size zero.
-
-
-# %% 02. The network in step-size coordinates
-#
-# The substitution s_theta(x, t, d) = v_{t, t+d}(x) turns the published network into the two-time
-# average velocity of this chapter, so the same weights serve both readings.
-def shortcut(model, x, t, d):
-    """s_theta(x, t, d), the average velocity over the step of length d starting at time t."""
-    return model(x, t, t + d)  # [B, 2] = v_{t, t+d}(x).
-
-
-def take_step(model, x, t, d):
-    """One step of length d, which is the flow map X_{t, t+d} in these coordinates."""
-    return x + d[:, None] * shortcut(model, x, t, d)  # [B, 2].
-
-
-# %% 03. The doubling identity, worked on the scalar flow
-#
-# s(x, t, 2d) = half s(x, t, d) + half s(x', t+d, d) with x' = x + d s(x, t, d). On dx/dt = x the
-# average velocity is x (e^d - 1)/d, and the identity holds at every step length at once.
-def scalar_doubling(x_value: float = 1.0, d: float = 0.25):
-    """Every term of the doubling identity on dx/dt = x, each one computed here."""
-    first = x_value * math.expm1(d) / d  # s(1, 0, 0.25).
-    reached = x_value + d * first  # x', which is e^{0.25} to six figures.
-    second = reached * math.expm1(d) / d  # s(x', 0.25, 0.25), faster because the state has grown.
-    doubled = x_value * math.expm1(2.0 * d) / (2.0 * d)  # s(1, 0, 0.5), computed directly.
-    return {
-        "s(1, 0, d)": first,
-        "state after the first step": reached,
-        "s(x', d, d)": second,
-        "average of the two": 0.5 * (first + second),
-        "s(1, 0, 2d) direct": doubled,
-        "doubling gap": abs(0.5 * (first + second) - doubled),
-        "state after the doubled step": x_value + 2.0 * d * doubled,
-        "exp(2d)": math.exp(2.0 * d),
-    }
-
-
-# %% 04. The dyadic ladder
-#
-# Eight step sizes from 1/128 up to 1, each training pair tying level k to level k+1. A state at
-# level zero is already trained against the diagonal by the anchor, so the correctness of the
-# one-step map rests on a chain of seven relations and one regression.
-def ladder():
-    """The step sizes d_k = 2^k d_0, from the finest to the whole interval."""
-    return [FINEST * 2.0 ** k for k in range(LEVELS + 1)]  # [1/128, ..., 1.0].
-
-
-# %% 05. The loss
-#
-# At step size zero the target is the flow matching target of Chapter 2, and above it the batch
-# draws a level of the ladder, forms the two short steps and regresses the doubled step onto their
-# average, with the target detached. The chapter puts three quarters of the batch on the zero step
-# and the rest on the doubling. Both terms are evaluated on the whole batch here and weighted in
-# that proportion, which has the same expectation at one network call more per step.
-def make_loss(model, data, batch_size, device):
-    sizes = torch.tensor(ladder()[:LEVELS], device=device)  # Levels whose doubling still fits in [0,1].
-
-    def loss_fn(step: int) -> torch.Tensor:
-        index = torch.randint(0, len(data), (batch_size,), device=device)
-        x1 = data[index]  # [B, 2] the data draw.
-        x0 = torch.randn_like(x1)  # [B, 2] the source draw.
-        t = torch.rand(batch_size, device=device)  # [B] a time for the zero-step rows.
-        xt = (1 - t[:, None]) * x0 + t[:, None] * x1  # [B, 2] the interpolant.
-        zero = torch.zeros(batch_size, device=device)
-        anchor = ((shortcut(model, xt, t, zero) - (x1 - x0)) ** 2).mean()  # Step size zero.
-
-        d = sizes[torch.randint(0, LEVELS, (batch_size,), device=device)]  # [B] one rung each.
-        start = (1.0 - 2.0 * d) * torch.rand(batch_size, device=device)  # [B] t with t + 2d <= 1.
-        xs = (1 - start[:, None]) * x0 + start[:, None] * x1  # [B, 2] the interpolant at t.
-        with torch.no_grad():  # The two legs are frozen; only the doubled step is corrected.
-            first = shortcut(model, xs, start, d)
-            reached = xs + d[:, None] * first
-            second = shortcut(model, reached, start + d, d)
-            target = 0.5 * (first + second)
-        doubled = ((shortcut(model, xs, start, 2.0 * d) - target) ** 2).mean()
-        return ZERO_FRACTION * anchor + (1.0 - ZERO_FRACTION) * doubled
-
-    return loss_fn
-
-
-# %% 06. Sampling, and the residual at every rung
-@torch.no_grad()
-def sample(model, x0, steps: int):
-    """Cross [0, 1] in equal steps of length 1/steps, one network evaluation per step."""
-    x, d = x0, torch.full((len(x0),), 1.0 / steps, device=x0.device)
-    for k in range(steps):
-        t = torch.full((len(x0),), k / steps, device=x0.device)
-        x = take_step(model, x, t, d)
-    return x  # [N, 2].
-
-
-@torch.no_grad()
-def rung_residuals(model, x0):
-    """The squared doubling gap at each level of the ladder, measured on source points."""
-    out = {}
-    for k, d in enumerate(ladder()[:LEVELS]):
-        step = torch.full((len(x0),), d, device=x0.device)
-        start = torch.zeros(len(x0), device=x0.device)
-        first = shortcut(model, x0, start, step)
-        second = shortcut(model, x0 + step[:, None] * first, start + step, step)
-        gap = shortcut(model, x0, start, 2.0 * step) - 0.5 * (first + second)
-        out[f"rung {k}: d = 1/{int(round(1.0 / d))}"] = float((gap ** 2).mean())
-    return out
-
-
-# %% 07. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=2000, batch_size=128, samples=512)
-    args = parser.parse_args()
-    set_seed(args.seed)
-    device = pick_device(args.device)
-    data = gaussian_mixture_2d(num_points=2048, seed=args.seed).to(device)  # [2048, 2].
-
-    report = scalar_doubling()
-    report["ladder levels"] = LEVELS
-    report["finest step d_0"] = FINEST
-    if not args.quiet:
-        print(f"on dx/dt = x: two steps of 0.25 average to {report['average of the two']:.6f}, "
-              f"one step of 0.5 returns {report['s(1, 0, 2d) direct']:.6f}")
-        print(f"ladder {[round(d, 5) for d in ladder()]}")
-
-    model = TwoTimeMLP(dim=2, width=args.width).to(device)
-    if not args.quiet:
-        print(f"\ntraining the shortcut model  parameters={count_parameters(model):,}")
-    losses = train(model, make_loss(model, data, args.batch_size, device),
-                   steps=args.steps, lr=args.lr, quiet=args.quiet)
-
-    torch.manual_seed(args.seed + 1)
-    x0 = torch.randn(args.samples, 2, device=device)  # [N, 2] the shared source points.
-    report.update(rung_residuals(model, x0))
-    samples: list[str] = []
-    for steps in (1, 2, 4, 8):
-        points = sample(model, x0, steps)
-        stats = mode_report(points)
-        report[f"{steps}-step nearest distance"] = stats["nearest_distance"]
-        report[f"{steps}-step log density"] = mean_log_density(points.cpu())
-        samples.append(f"{steps} step(s): nearest {stats['nearest_distance']:.3f}  "
-                       f"modes {stats['modes_found']}  entropy {stats['mode_entropy']:.3f}")
-    report["one step against eight"] = float(
-        (sample(model, x0, 1) - sample(model, x0, 8)).norm(dim=1).mean())
-    report["final loss"] = sum(losses[-20:]) / len(losses[-20:])
-
-    print_report("Shortcut models, Chapter 6", report, samples)
-    save_run(args.out, model=model, config=vars(args), report=report, losses=losses, samples=samples)
-
-
-if __name__ == "__main__":
-    main()

lecture_6/stochastic_map.py

@@ -1,218 +0,0 @@
-"""Strong stochastic flow maps, where the driving path is an argument of the map.
-
-A map that matches only the transition kernel draws fresh noise on every leg, so one full step and
-two half steps follow different trajectories and their composition residual has a floor at twice
-the variance of the interval. A strong map takes the path itself, projected onto the first few
-shifted Legendre polynomials, and the Chen relations build the coefficients of a long interval
-from the coefficients of its halves, so all three legs of the cocycle are driven by one
-realization.
-
-Chapter 6, Section 6.31 of the notes. Runs on the linear equation dX = X dt + sigma dW
-and on a cloud of starting states from ``cis6270.data.gaussian_mixture_2d``. Prints the weak
-map's floor of 2K_1 = 6.389, the Legendre variances 1, 1/3 and 1/5, the Chen fusion that turns
-four half-interval draws into (0.50, -0.20), and the pathwise error against the coefficient count.
-
-    python lecture_6/stochastic_map.py                 # 400,000 Monte Carlo draws, a few seconds
-    python lecture_6/stochastic_map.py --steps 20      # the floor of 1,000 draws, for a first look
-"""
-
-# %% 01. Imports and configuration
-import math
-
-import torch
-
-from cis6270.data import gaussian_mixture_2d
-from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed
-
-torch.set_num_threads(1)  # The work here is elementwise, and one thread keeps the timing stable.
-# ``tests/conftest.py`` sets the same thing for the test session, so the line here is the one
-# that acts when a reader runs this script directly.
-torch.set_default_dtype(torch.float64)  # The checks below sit well under the float32 noise floor.
-
-
-SIGMA = 1.0  # The noise level of dX = X dt + sigma dW.
-ORDER = 3  # Coefficients carried per interval, the number used in practice.
-FINEST = 32  # Subintervals the path is drawn on before the Chen relations fuse it upward.
-
-
-# %% 02. The transition kernel, and the weak map built from it
-#
-# Starting at x at time s, the state at t is e^{t-s} x plus a centered Gaussian whose variance is
-# K_{t-s}. A weak map reproduces that law on every interval with one fresh draw.
-def kernel_variance(d):
-    """K_d = sigma^2 (e^{2d} - 1)/2, the variance the equation accumulates over a length d."""
-    return 0.5 * SIGMA ** 2 * (math.exp(2.0 * d) - 1.0)
-
-
-def weak_step(x, d, noise):
-    """The weak map, which takes a state and one fresh standard normal draw."""
-    return math.exp(d) * x + math.sqrt(kernel_variance(d)) * noise
-
-
-# %% 03. The cost of matching only the transition kernel
-#
-# The variances compose correctly, so nothing is wrong at the level of distributions. At the level
-# of samples the three draws are independent, so the squared gap between one full step and two
-# half steps averages to twice the full-interval variance however well the map is trained.
-def weak_floor(x: float = 0.3, u: float = 0.5, t: float = 1.0, draws: int = 400_000, seed: int = 0):
-    """The worked comparison of Section 6.31, and a Monte Carlo average of the floor."""
-    full = weak_step(x, t, 0.8)  # One step with the chapter's third draw.
-    half = weak_step(weak_step(x, u, 0.5), t - u, -1.2)  # Two legs with the first two draws.
-    carried = math.exp(2.0 * (t - u)) * kernel_variance(u) + kernel_variance(t - u)
-    matching = (math.exp(t - u) * math.sqrt(kernel_variance(u)) * 0.5
-                - math.sqrt(kernel_variance(t - u)) * 1.2) / math.sqrt(kernel_variance(t))
-    generator = torch.Generator().manual_seed(seed)
-    e1, e2, e3 = (torch.randn(draws, generator=generator) for _ in range(3))
-    gap = (math.sqrt(kernel_variance(t)) * e3
-           - math.exp(t - u) * math.sqrt(kernel_variance(u)) * e1
-           - math.sqrt(kernel_variance(t - u)) * e2)  # [draws] the difference of the two routes.
-    return {
-        "K_1": kernel_variance(t), "sqrt K_1": math.sqrt(kernel_variance(t)),
-        "sqrt K_0.5": math.sqrt(kernel_variance(u)),
-        "one full step": full, "two half steps": half, "gap": full - half,
-        "squared gap": (full - half) ** 2,
-        "carried variance equals K_1": carried, "floor 2 K_1": 2.0 * kernel_variance(t),
-        "the draw that would close it": matching,
-        "Monte Carlo floor": float((gap ** 2).mean()),
-    }
-
-
-# %% 04. The path as an input, through the shifted Legendre polynomials
-#
-# I^{(n)}_{s,t} integrates the n-th shifted Legendre polynomial of the normalized time against the
-# increment. The zeroth coefficient is the plain increment, the first separates a path that did
-# its travelling early from one that did it late, and distinct coefficients are independent.
-def legendre(n: int, r):
-    """The n-th shifted Legendre polynomial on [0, 1]."""
-    return [torch.ones_like(r), 2.0 * r - 1.0, 6.0 * r * r - 6.0 * r + 1.0][n]
-
-
-def draw_coefficients(shape, length: float, generator):
-    """I^{(n)} ~ N(0, (t-s)/(2n+1)), drawn directly, with no path simulated."""
-    scale = torch.tensor([math.sqrt(length / (2 * n + 1)) for n in range(ORDER)])  # [ORDER].
-    return torch.randn(*shape, ORDER, generator=generator) * scale  # [..., ORDER].
-
-
-def kernel_coefficients(length: float, nodes: int = 8001):
-    """a_n(h) in e^{h(1-r)} = sum_n a_n P_n(r), which turns the Ito integral into a finite sum."""
-    r = torch.linspace(0.0, 1.0, nodes)
-    weight = torch.exp(length * (1.0 - r))  # The kernel e^{t-u} written in normalized time.
-    return torch.stack([(2 * n + 1) * torch.trapz(weight * legendre(n, r), r)
-                        for n in range(ORDER)])  # [ORDER].
-
-
-def strong_step(x, length: float, coefficients):
-    """The solution map over one interval, driven by the coefficients of the path on it."""
-    kernel = kernel_coefficients(length).to(x)  # [ORDER].
-    return math.exp(length) * x + SIGMA * (coefficients * kernel).sum(-1)  # Same shape as x.
-
-
-# %% 05. The Chen relations
-#
-# A shifted Legendre polynomial of the long interval expands in the shifted Legendre basis of each
-# half, so the coefficients fuse by a fixed linear rule with signs alternating in the parity of
-# n + m. The first rule says the increments add, and the minus sign in the second is the oddness
-# of P_1 about the midpoint.
-def chen(first, second):
-    """Fuse the coefficients of two equal halves into the coefficients of their union."""
-    return torch.stack([
-        first[..., 0] + second[..., 0],  # The increments add.
-        -0.5 * (first[..., 0] - second[..., 0]) + 0.5 * (first[..., 1] + second[..., 1]),
-        0.25 * (first[..., 2] + second[..., 2]) - 0.75 * (first[..., 1] - second[..., 1]),
-    ], dim=-1)
-
-
-def fuse_to(pieces, target: int):
-    """Halve the number of intervals by Chen fusion until ``target`` intervals remain."""
-    while pieces.shape[-2] > target:
-        pieces = chen(pieces[..., 0::2, :], pieces[..., 1::2, :])  # Fuse neighbouring pairs.
-    return pieces  # [..., target, ORDER].
-
-
-def worked_chen():
-    """The chapter's four half-interval draws, fused into the coefficients of the unit interval."""
-    halves = torch.tensor([[0.80, 0.20, 0.0], [-0.30, 0.50, 0.0]])  # [2, ORDER], the two legs.
-    fused = chen(halves[0], halves[1])
-    return {"I0 over [0,1]": float(fused[0]), "I1 over [0,1]": float(fused[1]),
-            "variance of I1 at h=1": 1.0 / 3.0,
-            "fused variance at h=1": 0.25 * 0.5 + 0.25 * 0.5 + 0.25 / 6.0 + 0.25 / 6.0}
-
-
-# %% 06. The pathwise error against the number of coefficients
-#
-# A fine simulation gives both the exact Ito integral and the coefficients of the same path, so
-# the truncation can be measured on the realization it was taken from.
-def truncation_error(paths: int, grid: int, seed: int):
-    """Mean absolute pathwise error of a one-step map reading 1, 2 and 3 coefficients."""
-    generator = torch.Generator().manual_seed(seed)
-    step = 1.0 / grid
-    time = torch.arange(grid) * step  # [grid] left endpoints of the substeps.
-    increments = torch.randn(paths, grid, generator=generator) * math.sqrt(step)  # [P, grid].
-    exact = (torch.exp(1.0 - time)[None, :] * increments).sum(1)  # [P] the Ito integral.
-    measured = torch.stack([(legendre(n, time)[None, :] * increments).sum(1)
-                            for n in range(ORDER)], dim=1)  # [P, ORDER].
-    kernel = kernel_coefficients(1.0)  # [ORDER].
-    out = {f"coefficient variance I{n}": float(measured[:, n].var()) for n in range(ORDER)}
-    for keep in range(1, ORDER + 1):
-        approximate = (measured[:, :keep] * kernel[:keep]).sum(1)  # [P].
-        out[f"pathwise error with {keep}"] = float((exact - approximate).abs().mean())
-    return out
-
-
-# %% 07. One path, several step counts
-#
-# A strong map crosses one trajectory of one realization in any number of hops, so its terminal
-# state barely moves with the step count. A weak map redraws at every change of the step count,
-# so its terminal state moves by the floor of Section 6.31.
-def cross(x, coefficients, steps: int, weak: bool, generator=None):
-    """Cross [0, 1] in ``steps`` equal hops, either along the fused path or with fresh draws."""
-    pieces = fuse_to(coefficients, steps)  # [..., steps, ORDER] on this grid.
-    length = 1.0 / steps
-    for k in range(steps):
-        if weak:
-            noise = torch.randn(x.shape, generator=generator)  # A fresh draw for every leg.
-            x = weak_step(x, length, noise)
-        else:
-            x = strong_step(x, length, pieces[..., k, :])
-    return x
-
-
-# %% 08. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=400_000, samples=512)
-    args = parser.parse_args()  # --steps is the Monte Carlo draw count here, since nothing trains.
-    set_seed(args.seed)
-    pick_device(args.device)  # Every number here is double precision, which not every
-    # accelerator carries, so the work stays on the CPU whatever --device names.
-
-    report = weak_floor(draws=max(args.steps, 1000), seed=args.seed)
-    report.update(worked_chen())
-    report.update(truncation_error(paths=1024, grid=4096, seed=args.seed))
-    if not args.quiet:
-        print(f"weak floor 2K_1 = {report['floor 2 K_1']:.3f}, Monte Carlo "
-              f"{report['Monte Carlo floor']:.3f}")
-        print(f"pathwise error {report['pathwise error with 1']:.4f} with the increment alone "
-              f"against {report['pathwise error with 3']:.6f} with three coefficients")
-
-    # %% 09. The same cloud at four step counts
-    points = gaussian_mixture_2d(num_points=args.samples, seed=args.seed).double()  # [N, 2].
-    generator = torch.Generator().manual_seed(args.seed + 1)
-    path = draw_coefficients((args.samples, 2, FINEST), 1.0 / FINEST, generator)  # [N, 2, F, 3].
-    reference = cross(points, path, FINEST, weak=False)  # The finest crossing, as the yardstick.
-    samples: list[str] = []
-    for steps in (2, 4, 8, 16):  # Short of FINEST, which would be the yardstick against itself.
-        strong = cross(points, path, steps, weak=False)  # One realization, several hop counts.
-        weak_generator = torch.Generator().manual_seed(args.seed + 2)
-        weak = cross(points, path, steps, weak=True, generator=weak_generator)
-        report[f"{steps} hops: strong spread"] = float((strong - reference).norm(dim=1).mean())
-        report[f"{steps} hops: weak spread"] = float((weak - reference).norm(dim=1).mean())
-        samples.append(f"{steps:2d} hops  strong {float((strong - reference).norm(dim=1).mean()):.6f}"
-                       f"  weak {float((weak - reference).norm(dim=1).mean()):.3f}")
-
-    print_report("Strong stochastic flow maps, Chapter 6", report, samples)
-    save_run(args.out, config=vars(args), report=report, samples=samples)
-
-
-if __name__ == "__main__":
-    main()