Add Lecture 6 flow maps with tested training, sampling, and course navigation

Add 16 complete Lecture 6 method examples, mathematical tests, paper/code provenance, slide mappings, and actual seeded verification records. Update the course README and preserve earlier lecture files.

README.md

@@ -34,6 +34,7 @@ and implementation notes; additional directories will accompany later lectures.
 | 3 | Flow matching, diffusion, and guidance for ESM-2 residue embeddings | [Guide](lecture_3/README.md) · [Flow matching](lecture_3/esm2_flow_guidance.py) · [Diffusion](lecture_3/esm2_diffusion_guidance.py) |
 | 4 | Discrete diffusion, masked and uniform corruption, block generation, and guidance | [Guide](lecture_4/README.md) · [Training and generation](lecture_4/run.py) · [Slide code map](lecture_4/SLIDE_CODE_MAP.md) |
 | 5 | Discrete flow matching, Dirichlet and Fisher paths, Gumbel-Softmax, rectification, and multi-objective generation | [Guide](lecture_5/README.md) · [Training and generation](lecture_5/run.py) · [Slide code map](lecture_5/SLIDE_CODE_MAP.md) |
+| 6 | Continuous, latent, categorical, posterior, expanding, and strong stochastic flow maps | [Guide](lecture_6/README.md) · [Training and generation](lecture_6/run.py) · [Slide code map](lecture_6/SLIDE_CODE_MAP.md) |
 
 ## Installation
 
@@ -41,7 +42,7 @@ Use Python 3.11, or another compatible Python version at least 3.10, in a new
 virtual environment. The shared `requirements.txt` pins PyTorch 2.9.1,
 TorchVision 0.24.1, and Transformers 4.57.6. Lecture 2 uses PyTorch and
 TorchVision; Lecture 3 also uses Transformers. Lectures 4 and 5 use PyTorch,
-NumPy, and SciPy. Their lecture folders also provide minimal requirements.
+NumPy, and SciPy. Lecture 6 uses PyTorch and NumPy. These lecture folders also provide minimal requirements.
 
 ```bash
 git clone https://huggingface.co/ChatterjeeLab/CIS6270
@@ -144,6 +145,29 @@ Both folders include numerical examples, mathematical tests, and saved results
 from seeded CPU runs. The guides explain finite endpoint approximations and
 classroom simplifications for each method.
 
+## Lecture 6 - Flow Maps
+
+Learn finite-time motion after the local flows from Lecture 5. The new folder
+contains 16 complete examples covering flow-map matching and self-distillation,
+consistency, Shortcut, MeanFlow, learned latent representations, Flow Map
+Language Models, Categorical and Discrete Flow Maps, Diamond Maps, Meta Flow
+Maps, Expanding Flow Maps, and Strong Stochastic Flow Maps.
+
+```bash
+python lecture_6/run.py --method self-distill --out lecture_6/outputs/self-distill
+python lecture_6/run.py --mode sample --out lecture_6/outputs/self-distill --sample-steps 1
+python lecture_6/run_all.py --quick
+```
+
+The [Lecture 6 guide](lecture_6/README.md) includes all method commands,
+training objectives, data formats, and numerical assumptions. The
+[slide code map](lecture_6/SLIDE_CODE_MAP.md) links implementations to the
+[Flow Maps presentation](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit).
+[Source notes](lecture_6/SOURCES.md) identify the exact papers and inspected
+author-code revisions, including differences between paper pseudocode and
+released implementations. [Verified examples](lecture_6/verified_examples/README.md)
+contain actual training logs, generated samples, and checkpoint-reload checks.
+
 ## Repository organization
 
 | Location | Contents |
@@ -153,6 +177,7 @@ classroom simplifications for each method.
 | [`lecture_3/`](https://huggingface.co/ChatterjeeLab/CIS6270/tree/main/lecture_3) | ESM-2 flow and diffusion guidance scripts, sequence data, guide, and mathematical notes |
 | [`lecture_4/`](https://huggingface.co/ChatterjeeLab/CIS6270/tree/main/lecture_4) | Seven discrete diffusion and guidance examples, synthetic DNA, slide code map, and verified outputs |
 | [`lecture_5/`](https://huggingface.co/ChatterjeeLab/CIS6270/tree/main/lecture_5) | Eight discrete and simplex flow examples, synthetic DNA, slide code map, and verified outputs |
+| [`lecture_6/`](https://huggingface.co/ChatterjeeLab/CIS6270/tree/main/lecture_6) | Sixteen flow-map implementations, text data, mathematical notes, slide links, and verified outputs |
 | [`tests/`](https://huggingface.co/ChatterjeeLab/CIS6270/tree/main/tests) | Offline checks for the Lecture 3 examples |
 
 Installation instructions and the lecture index are maintained at the
@@ -165,6 +190,7 @@ references accompany the corresponding code.
 python -m unittest discover -s tests -v
 python -m unittest discover -s lecture_4/tests -v
 python -m unittest discover -s lecture_5/tests -v
+python -m unittest discover -s lecture_6/tests -v
 ```
 
 The Lecture 3 unit tests cover property annotations, scalarization weights, reward
@@ -172,7 +198,10 @@ gradients, DDPM schedule indexing, and constrained decoding.
 
 Lecture 4 tests check reverse KL losses and guidance calculations. Lecture 5
 tests check the master equation, Fisher geometry, Gumbel path derivatives, and
-MH detailed balance. Run a short end-to-end check of every method from its
+MH detailed balance. Lecture 6 checks map identities and JVPs, categorical
+teachers, GLASS conditioning, reward gradients, insertion clocks, and Brownian
+composition. Its end-to-end runner also verifies checkpoint reloads.
+Run a short end-to-end check of every method from its
 lecture folder with `python run_all.py --quick`.
 
 ## License

lecture_6/.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+.venv/
+outputs/
+*.pt

lecture_6/MATHEMATICS.md

@@ -0,0 +1,225 @@
+# From the lecture equations to the running code
+
+The [slide code map](SLIDE_CODE_MAP.md) supplies the corresponding slide links.
+Unless stated otherwise, time zero denotes Gaussian noise and time one denotes
+data. All expectations below are population quantities; the training code uses
+finite minibatches and learned approximations.
+
+## We first learn the motion at one time
+
+For a noise sample $X_0$ and a data sample $X_1$, construct
+
+$$I_t=(1-t)X_0+tX_1,\qquad \dot I_t=X_1-X_0.$$
+
+Here $I_t$ is the sampled interpolant and $t$ is its progress from noise to data.
+Regressing the displacement against the noisy state learns
+
+$$b_t(x)=\mathbb E[X_1-X_0\mid I_t=x].$$
+
+`interpolate` constructs the pair; `diagonal_loss` performs the regression.
+The conditional-mean field transports the same one-time marginals as the
+interpolant. Individual interpolant lines generally differ from its ODE paths.
+
+## Now predict the destination over an interval
+
+$$F_{s,t}(x)=x+(t-s)v_{s,t}(x).$$
+
+$s$ is the start, $t$ is the arrival time, and $v_{s,t}$ is the average velocity
+along that ODE trajectory. `finite_map` uses this residual parameterization,
+which gives $F_{s,s}(x)=x$ exactly. The diagonal average $v_{s,s}$ must still be
+trained to equal $b_s$.
+
+For $\dot x=x$, the exact map is $e^{t-s}x$. From $x=1$ over a unit interval,
+the average velocity is $e-1\approx1.71828$. The residual update reaches
+$2.71828$, while one Euler step using the initial velocity reaches only $2$.
+`numerical_examples.py` trains this scalar map and reports its actual error.
+
+## The ODE supplies three finite-interval identities
+
+$$\partial_t F_{s,t}(x)=b_t(F_{s,t}(x)),$$
+
+$$\partial_sF_{s,t}(x)+J_xF_{s,t}(x)b_s(x)=0,$$
+
+$$F_{s,t}(x)=F_{u,t}(F_{s,u}(x)).$$
+
+The first changes the destination time. The second advances the start along the
+same path, leaving the destination fixed. The third splits the interval at an
+intermediate time $u$. `lagrangian_residual` and `eulerian_residual` use JVPs;
+`semigroup_loss` evaluates the direct and split paths.
+
+The identity map satisfies composition for every interval but has zero velocity.
+This is why composition needs the diagonal anchor. A detached target also
+changes the optimization: only the student branch receives its gradient.
+
+## Shortcut and MeanFlow choose particular targets
+
+Shortcut learns a velocity for a longer interval from two shorter intervals:
+
+$$v_{s,s+2d}(x)\approx\operatorname{sg}\left[
+\tfrac12 v_{s,s+d}(x)+\tfrac12 v_{s+d,s+2d}(x+d\,v_{s,s+d}(x))\right].$$
+
+$d$ is the half-step length. The second prediction receives the moved state.
+`shortcut_loss` samples dyadic lengths and builds that detached target.
+
+MeanFlow uses the opposite clock. Its $z_t=(1-t)X_{\rm data}+tX_{\rm noise}$
+has data at zero and noise at one. For $r\leq t$,
+
+$$z_r=z_t-(t-r)u(z_t,r,t),$$
+
+$$u=v-(t-r)(\partial_tu+J_zu\,v).$$
+
+$v=X_{\rm noise}-X_{\rm data}$ is the conditional training velocity.
+`meanflow_loss` evaluates the JVP in direction $(v,0,1)$ for inputs $(z,r,t)$,
+then detaches the complete right-hand target. The sampler retains the minus sign.
+
+## A latent representation changes the coordinates
+
+For encoder $E$, decoder $D$, and data distribution $p$, fit a map to $E_\#p$.
+Generation is $D(F_{0,1}(Z_0))$. The code first trains an autoencoder and freezes
+it; a stored mean and standard deviation normalize its latent coordinates.
+`train_latent` saves the encoder, decoder, normalization, and map together.
+
+For an $L_D$-Lipschitz decoder, the distributional error is bounded by latent
+transport error multiplied by $L_D$, plus reconstruction error. The code
+reports reconstruction MSE separately; it does not assert a measured global
+Lipschitz constant or equate that MSE with a Wasserstein distance.
+
+## Categorical maps constrain their predictions
+
+For one-hot endpoints, the posterior mean denoiser $D_s(x)$ lies on the
+probability simplex. The noisy state $x$ may lie outside it. The finite map is
+
+$$F_{s,t}(x)=\frac{1-t}{1-s}x+\frac{t-s}{1-s}\psi_{s,t}(x),$$
+
+where $\psi_{s,t}$ is the two-time softmax prediction. `categorical_map` implements
+this expression. At $s=.25,t=.75$, state $(-.2,.6,1.1)$ and prediction
+$(.1,.7,.2)$ produce $(0,2/3,.5)$; the intermediate state is not normalized.
+
+Composition induces
+
+$$q=\gamma\psi_{s,u}(x)+(1-\gamma)\psi_{u,t}(F_{s,u}(x)),\qquad
+\gamma=\frac{(1-t)(u-s)}{(1-u)(t-s)}.$$
+
+At $(s,u,t)=(0,.5,.75)$, the weight is $1/3$. Combining predictions $(.8,.2)$
+and $(.2,.8)$ gives $(.4,.6)$. `composition_target` builds this probability-valued
+target; a detached KL objective has logit gradient $p_{\rm student}-q$.
+
+The decoding clock is $\tau(t)=1-\frac{V}{V-1}P_e(t)$, with vocabulary size $V$
+and single-token corruption error $P_e$. `DecodingClock` numerically inverts it.
+Endpoint enforcement prevents floating-point accuracy saturation from stopping
+the sampler before physical time one.
+
+## Categorical and discrete consistency constrain different residuals
+
+Categorical ECLD combines endpoint agreement with the time derivative of the
+finite denoiser. With $\eta=(t-s)/(1-s)$, the scaled Lagrangian residual is
+
+$$r=\psi_{s,t}-\psi_{t,t}(F_{s,t})+(1-t)\eta\partial_t\psi_{s,t}.$$
+
+The endpoint term can vanish while the derivative term remains nonzero. The
+implementation includes both and follows the released code's finite weights.
+
+For Discrete Flow Maps, write $\psi=\operatorname{softmax}(z)$ and center the
+arrival-time logit derivative:
+
+$$\delta_k=\partial_tz_k-\sum_j\psi_j\partial_tz_j,\qquad
+c=\frac{(t-s)(1-t)}{1-s}.$$
+
+The Lagrangian teacher is
+
+$$T_{\rm LSD}=\operatorname{softmax}\left[z_{t,t}(F_{s,t})-\log(1+c\delta)\right].$$
+
+For the Eulerian teacher, replace the derivative with
+$D_sz=\partial_sz+J_xz\,b_s$, center it in the same way, and use
+
+$$T_{\rm ESD}=\operatorname{softmax}\left[z_{s,s}-
+\log\left(1-\frac{(1-s)(t-s)}{1-t}\delta\right)\right].$$
+
+`corrected_logit_teacher` computes both versions. Log arguments must be positive;
+the code records its stabilization frequency so that violations remain visible.
+
+## Posterior maps need a second noise draw
+
+Meta training constructs
+
+$$I_t=(1-t)X_0+tX_1,\qquad \bar I_s=(1-s)\bar X_0+sX_1,$$
+
+with independent $X_0,\bar X_0$. The outer observation $(t,I_t)$ stays fixed
+while the inner map advances from $s$ to its target time. At that fixed context,
+the shared $X_1$ has exactly the desired posterior distribution.
+
+Diamond distillation obtains the corresponding conditional velocity through
+GLASS. For two independent observations, the likelihood precisions add:
+
+$$\Lambda=\frac{s^2}{(1-s)^2}+\frac{t^2}{(1-t)^2},\qquad
+S=\Lambda^{-1}\left[\frac{s\bar x}{(1-s)^2}+\frac{tx}{(1-t)^2}\right].$$
+
+The equivalent linear-interpolant time is
+$t^*=\sqrt\Lambda/(1+\sqrt\Lambda)$. Evaluating the original denoiser at
+$(t^*,t^*S)$ supplies the posterior mean conditioned on both observations.
+`glass_denoiser` is checked against an independent Gaussian conditioning formula.
+
+## Posterior samples make reward averages computable
+
+For reward $r$ and conditional samples $Z_k$,
+
+$$\widehat V_t(x)=\log\left(\frac1K\sum_{k=1}^K e^{r(Z_k)}\right),\qquad
+\widehat D_t^r(x)=\sum_k\operatorname{softmax}(r(Z))_k Z_k.$$
+
+`posterior_value` differentiates through the learned conditional samples.
+`guided_samples` uses $(\widehat D_t^r(x)-x)/(1-t)$ as the guided velocity.
+The ratio is biased at finite $K$, even with an exact conditional sampler.
+
+For Meta fine-tuning, let $d=b_{\rm student}-b_{\rm base}$, $w=e^{r(Z)}$, and
+$a=g_t^2/2$. The detached surrogate is
+
+$$\ell=\left\|d+(w-1)\operatorname{sg}(d)-a\operatorname{sg}(\nabla_xw)\right\|^2.$$
+
+Its expected gradient in $d$ is $2\mathbb E[wd-a\nabla_xw]$. Squaring $wd-a\nabla w$
+without that stop-gradient placement introduces an extra weight. The test suite
+checks this distinction directly. For the linear path, $a=(1-t)/t$; the optional
+training demonstration therefore excludes the singular initial time.
+
+## Expansion gives each new token its own clock
+
+A token born at $b_i$ has local time
+
+$$\tau_i(t)=\max\left(0,\frac{t-b_i}{1-b_i}\right),\qquad
+\frac{d\tau_i}{dt}=\frac1{1-b_i}\quad(t>b_i).$$
+
+`local_clock` and `local_map` preserve this information during transport.
+`insert_tokens` keeps token order, inserted noise, and birth metadata aligned.
+The linear insertion CDF gives conditional birth fraction
+$\rho_{s,t}=(t-s)/(1-s)$. Predicted remaining gap means are converted into
+interval means using this factor.
+
+The count divergence is
+
+$$\phi(a,b)=b-a+a\log(a/b),\qquad b>0,$$
+
+where $a$ is a realized count and $b$ is the predicted mean. Its population
+minimum is $b=\mathbb E[a]$. The zero-count branch is implemented with a finite
+autodiff expression. Sampling uses bounded binomial proposals and reports
+counts removed by global budget capping.
+
+## Strong maps must retain the same driving path
+
+For an interval of length $h$, the first two Brownian integrals satisfy
+
+$$I_0\sim\mathcal N(0,h),\qquad I_1\sim\mathcal N(0,h/3),\qquad I_0\perp I_1.$$
+
+For left/right subintervals of lengths $h_L,h_R$, Chen composition is
+
+$$I_0=I_0^L+I_0^R,$$
+
+$$I_1=\frac{h_LI_1^L+h_RI_1^R-h_RI_0^L+h_LI_0^R}{h_L+h_R}.$$
+
+With equal halves, $(I_0^L,I_1^L)=(.2,.04)$ and
+$(I_0^R,I_1^R)=(-.1,-.02)$ combine into $(.1,-.14)$.
+`chen_two` and `aggregate_tree` ensure the direct and split maps refer to that
+same history. Independent coarse noise would compare different sample paths.
+
+The SSFM objective combines a short stochastic matching step with same-noise
+composition. The reported strong error compares the learned result to a fine
+Euler-Maruyama trajectory using the same increments. Terminal variance checks
+the marginal separately; matching variance cannot establish pathwise accuracy.

lecture_6/README.md

@@ -0,0 +1,225 @@
+# CIS 6270 - Lecture 6 - Flow Maps
+
+Course hub: [ChatterjeeLab/CIS6270](https://huggingface.co/ChatterjeeLab/CIS6270).
+Slides: [Lecture 6 - Flow Maps](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit).
+Previous lecture: [Discrete Flow Matching](../lecture_5/README.md).
+
+Lecture 5 learned how a distribution moves locally. Here we learn the motion over
+an entire time interval. The examples progress from continuous flow maps to
+latent representations, categorical text, posterior maps, expanding states,
+and maps driven by a shared Brownian path.
+
+This folder contains **16 complete training and generation examples**, with
+checkpoint loading, data, mathematical checks, and saved execution results.
+The code uses small PyTorch networks and inspectable datasets. Each method's
+objective and sampler are implemented here; the large image and language
+benchmark runs from the papers require their original architectures, datasets,
+and training budgets. [SOURCES.md](SOURCES.md) records the exact papers, inspected
+author-code revisions, and differences from those implementations.
+
+## Train a velocity, then learn a finite map
+
+Use Python 3.11 or later. The minimal requirements retain the course's PyTorch
+2.9.1 baseline. All examples run on CPU; `--device cuda` selects an available GPU.
+The data require no downloads.
+
+From the course repository root:
+
+```bash
+cd lecture_6
+python -m venv .venv
+source .venv/bin/activate
+python -m pip install -r requirements.txt
+
+python run.py --method flow-matching --sample-steps 32 --out outputs/flow-matching
+python run.py --method self-distill --out outputs/self-distill
+python run.py --mode sample --out outputs/self-distill --sample-steps 1
+```
+
+The first command fits the conditional-mean velocity and integrates it with
+Heun's method. The second trains a two-time map using a diagonal velocity loss
+and a detached composition target. The third loads its checkpoint and generates
+in one map evaluation. The baseline uses two velocity evaluations per Heun step;
+finite maps use one evaluation per step.
+
+Training writes `checkpoint.pt`, `config.json`, `losses.json`, `report.json`, and
+`samples.txt`. A teacher or autoencoder stage also writes `teacher_losses.json`.
+Generation-only mode reads model architecture and vocabulary from the checkpoint
+and writes `resampled.txt` and `sample_report.json`, preserving training records.
+The checkpoint contains the inference state, not optimizer state for resuming
+training. `--mode train` saves training records without generating samples.
+
+## Run every method in the same structure as Lectures 4 and 5
+
+```bash
+python run_all.py --quick
+python run_all.py
+python -m unittest discover -s tests -v
+python numerical_examples.py --output outputs/numerical
+```
+
+`--quick` performs 20 optimization steps per method and verifies execution and
+checkpoint reloading. The regular command uses 1,000 steps per training stage.
+Both run every method and require seeded generated samples to agree exactly
+after checkpoint reload. The quick run measures execution correctness; inspect
+the longer run's losses and sample metrics when discussing learned behavior.
+
+| Method | Training | Generation |
+| --- | --- | --- |
+| `flow-matching` | Conditional displacement regression on a linear interpolant | Heun integration of the learned diagonal velocity |
+| `fmm-lagrangian` | Train a velocity teacher, then match the map's arrival-time derivative to the teacher | Residual finite-time maps |
+| `fmm-eulerian` | Train a velocity teacher, then match the start-time directional derivative to zero | Residual finite-time maps |
+| `self-distill` | Diagonal flow matching and EMA two-subinterval composition targets | One or more direct map updates |
+| `consistency` | Endpoint consistency along short Heun-solved teacher intervals, with the diagonal anchor | One endpoint prediction |
+| `shortcut` | Conditional flow matching and two-half-step average-velocity targets on dyadic intervals | One or more shortcuts |
+| `meanflow` | Backward average-velocity target with a directional JVP and detached adaptive weights | Backward updates from noise at time one to data at zero |
+| `latent` | Train a 3D-to-2D autoencoder, freeze it, standardize its encodings, then train a flow map | Latent map followed by the saved decoder |
+| `fmlm` | Diagonal token CE and weighted progressive denoiser consistency | Warped-time affine categorical maps, followed by argmax |
+| `categorical` | Diagonal CE and endpoint consistency plus temporal-derivative energy | The same categorical map parameterization |
+| `discrete-lsd` | Diagonal CE and Lagrangian logit-corrected KL teacher | Categorical finite-time maps |
+| `discrete-esd` | Diagonal CE and Eulerian directional-JVP logit teacher | Categorical finite-time maps |
+| `diamond` | Distill a GLASS posterior velocity built from an analytic Gaussian-mixture denoiser | Conditional posterior samples, differentiable value estimates, and reward steering |
+| `meta` | Independent inner and outer noise with a shared data endpoint, diagonal regression, and conditional composition | Posterior samples and reward-weighted posterior-mean steering |
+| `expanding` | Local-clock token CE, compatible-canvas composition, and remaining/interval gap-count losses | Learned binomial gap insertion, budget capping, and local-clock transport |
+| `ssfm` | OU drift anchoring, small-step stochastic matching, and same-noise composition | Strong maps using consistently aggregated Brownian coefficients |
+
+Each method also has its own entry point:
+
+```bash
+python examples/fmm_lagrangian.py
+python examples/fmm_eulerian.py
+python examples/consistency.py
+python examples/shortcut.py
+python examples/meanflow.py --sample-steps 1
+python examples/latent.py --teacher-steps 2000
+python examples/fmlm.py --sample-steps 1
+python examples/categorical.py
+python examples/discrete_lsd.py
+python examples/discrete_esd.py
+python examples/diamond.py --posterior-steps 4 --particles 64
+python examples/meta.py --finetune-steps 300
+python examples/expanding.py --sample-steps 4
+python examples/ssfm.py --sample-steps 4
+python examples/ssfm.py --ssfm-target paper --lr 0.0001 --train-steps 2000 --out outputs/ssfm-paper
+```
+
+The common options are `--train-steps`, `--teacher-steps`, `--batch-size`,
+`--width`, `--lr`, `--seed`, `--samples`, `--sample-steps`, and `--out`.
+`--teacher-steps` controls the teacher or autoencoder stage where one exists.
+Run `python run.py --help` for the remaining options. The `consistency` example
+always makes one endpoint prediction. SSFM requires a power-of-two sampling
+step count so its finite Brownian tree supports every requested partition.
+
+## Change one component at a time
+
+`lecture_core.py` exposes the functions used in the slide walkthroughs.
+`common.py` contains networks, data handling, optimization, and reference
+distributions. Each method family has a separate implementation file.
+
+| File | Responsibility |
+| --- | --- |
+| `continuous.py` | Diagonal, Lagrangian, Eulerian, semigroup, Shortcut, MeanFlow, and latent training |
+| `categorical.py` | Affine probability-valued maps, inverse decoding clock, ECLD, and logit teachers |
+| `posterior.py` | GLASS fusion, conditional map training, posterior values, importance correction, and reward fine-tuning |
+| `expanding.py` | Birth times, compaction, gap labels, count losses, and expansion/transport sampling |
+| `stochastic.py` | Brownian integrals, Chen composition, OU training, and strong-error evaluation |
+| `numerical_examples.py` | The lecture's numerical examples, including the trained exponential-flow example |
+| `SLIDE_CODE_MAP.md` | Stable links from every code slide and major method to the implementation |
+| `MATHEMATICS.md` | Equations, variable definitions, and the meaning of each loss |
+| `verified_examples/` | Actual seeded training logs, generated samples, and verification reports |
+
+The small sequence MLP receives the entire sequence. Its per-position output
+can therefore depend on every input position. The variable-length network also
+receives padding masks and birth-time metadata. Replacing these backbones with
+transformers does not require changing the displayed map identities.
+
+## Data and evaluation
+
+The continuous examples use four 2D Gaussians with known centers and variance.
+Reports include nearest-center distance, mode fractions, mode entropy, and exact
+target log density. These metrics describe different aspects of the samples;
+high mode entropy alone does not establish sample fidelity.
+
+`data/phrases.txt` contains correlated four-token phrases such as `red circle
+moves left`. `data/variable_text.txt` contains two-, four-, and five-token
+versions. Reports record token entropy, unique fraction, training-support
+fraction, and lengths. These tiny-data statistics are not LM1B/OWT perplexity or
+large-language-model evaluations. The split is a seeded 80/20 split of examples;
+the small grammar intentionally appears in both splits.
+
+Custom fixed-length text can be supplied with `--data my_phrases.txt`.
+The expanding example also accepts variable-length lines, up to 32 tokens.
+Continuous methods accept `--data points.csv`, with two finite numeric columns
+and no header. Diamond, Meta, and SSFM retain their specified reference
+distributions so their analytic diagnostics remain valid. See [data/README.md](data/README.md).
+
+## Numerical and theorem boundaries
+
+- **Continuous maps.** The diagonal is exactly the identity as a map, while its
+  predicted average velocity is trained by flow matching. Interpolant samples
+  are not individual ODE trajectories. Diagonal regression has an irreducible
+  conditional-variance floor, so raw training loss need not approach zero.
+- **Teacher distillation.** The finite-map teachers are learned in a separate
+  first stage. This retains teacher approximation error. The fixed MLP, bounded
+  interval sampling, and loss weights are classroom choices.
+- **Consistency and Shortcut.** Endpoint consistency is a compact
+  flow-matching-clock teaching implementation. Shortcut uses dyadic interval
+  lengths, a continuous start-time sample, EMA targets, and a separate diagonal
+  loss. The authors' large-image code has additional schedules and clipping.
+- **MeanFlow.** Data are at time zero and noise at one. The directional JVP uses
+  the conditional training velocity, and the entire corrected target is
+  detached. Seventy-five percent of training examples use equal times. Uniform
+  time sampling replaces the author's default logit-normal schedule.
+- **Latent maps.** This section demonstrates a learned representation and
+  decoder. It does not claim a separate canonical paper named Latent Flow Maps.
+  Reconstruction error and latent transport error are reported separately.
+- **Categorical maps.** Only the predicted denoiser lies on the simplex during
+  the trajectory. The final step reaches time one exactly; argmax is an explicit
+  decoding approximation for an imperfect model. We use Gaussian quadrature
+  for the decoding clock. Categorical ECLD uses the released code's bounded
+  endpoint-plus-time-energy objective, with an EMA endpoint target. Detached KL
+  and CE have the same student gradient; they log different scalar values.
+- **Discrete logit teachers.** The logarithmic correction requires positive
+  denominators. The implementation clamps them at 0.05 and logs the affected
+  fraction. This changes the teacher away from the exact solution; softmax
+  alone cannot repair an invalid logarithm. Differential objectives stay below
+  terminal time 0.97. No unreported gradient surgery is applied.
+- **Diamond and Meta.** The two noises are independent and share a clean
+  endpoint. Diamond uses an exact, evaluable mixture denoiser as its GLASS
+  teacher; Meta learns directly from samples. Posterior diagnostics compare
+  both means and variances. Finite importance estimates and reward-weighted
+  posterior means have self-normalization bias. The weighted-Diamond companion
+  uses an explicit full-support Gaussian proposal with a known density.
+- **Meta fine-tuning.** `--finetune-steps` runs and saves the detached surrogate
+  from Equation 43. It uses a bounded coordinate reward and the analytic base
+  mixture drift. Training is restricted to outer times 0.05 through 0.85;
+  extending the learned drift outside that interval is extrapolation.
+- **Expanding maps.** The data use a linear birth CDF. The network predicts
+  remaining gap means and interval means through the conditional insertion
+  factor. Binomial proposals and global budget capping approximate a joint count
+  law; a mean alone does not identify that law. A finite sampling jump introduces
+  new coordinates at its start, then transports them. Training compares routes
+  on a shared destination canvas with compatible noise and local clocks.
+  These are explicit finite-step choices, not a proof of exact conditional-law
+  recovery by the fitted network. Empty generated sequences remain visible in
+  the reports; the sampler does not silently replace them.
+- **Strong stochastic maps.** The example is an additive-noise OU SDE with a
+  known diffusion coefficient. It uses the first two Legendre integrals and a
+  learned average drift. Default `official-code` targets detach an EMA split
+  prediction; `--ssfm-target paper` detaches the direct prediction, as in the
+  lecture's pseudocode. Both reuse exactly the same Brownian history. The
+  reference is a fine Euler-Maruyama solve on that history, so its residual
+  discretization error is retained. No claim is made for multiplicative noise
+  or exact pathwise recovery with only two coefficients. The optional paper
+  orientation remained inaccurate in the recorded experiments, including the
+  lower-learning-rate recipe above; see the explicit sensitivity results in
+  `verified_examples/README.md`. Use the default released-code orientation for
+  the successful OU demonstration.
+
+## Read the saved results with the implementation
+
+[verified_examples/README.md](verified_examples/README.md) records the executed
+environment, commands, and metrics. Checkpoints are generated locally and are
+excluded from this folder's Git history, matching Lectures 4 and 5. The saved
+text/JSON records let students inspect a complete run before training their own
+models. The [slide code map](SLIDE_CODE_MAP.md) connects that run to the lecture.

lecture_6/SLIDE_CODE_MAP.md

@@ -0,0 +1,36 @@
+# Lecture 6 slide-to-code map
+
+These links use stable native slide identifiers. Each code slide connects to the same functions used by the training runner.
+
+| Code walkthrough | Slide | Implementation |
+| --- | --- | --- |
+| The map takes the state and both times as inputs | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u038_b0) | `common.finite_map`, `continuous.diagonal_loss`, `continuous.semigroup_loss`, `numerical_examples.train_scalar` |
+| The two losses connect local motion to a longer move | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u039_b0) | `common.finite_map`, `continuous.diagonal_loss`, `continuous.semigroup_loss`, `numerical_examples.train_scalar` |
+| A JVP computes that correction without forming a Jacobian | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u051_b0) | `continuous.meanflow_loss` |
+| Softmax constrains the prediction before the map is applied | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u080_b0) | `categorical.categorical_map`, `categorical.composition_target`, `categorical.probability_kl` |
+| The target combines probabilities from two different states | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u081_b0) | `categorical.categorical_map`, `categorical.composition_target`, `categorical.probability_kl` |
+| We can differentiate through the posterior samples | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u122_b0) | `posterior.posterior_samples`, `posterior.posterior_value` |
+| The stop-gradient placement is essential in this fine-tuning loss | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u139_b0) | `posterior.fine_tune_surrogate`, `posterior.train_reward_drift` |
+| Keep the count, noise, and clock updates together | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u168_b0) | `expanding.bounded_counts`, `expanding.insert_tokens`, `expanding.local_clock` |
+| Sampling and combining two coefficients is only a few lines | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u198_b0) | `stochastic.sample_coefficients`, `stochastic.chen_two`, `stochastic.train_stochastic` |
+| The consistency update compares those two routes | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u199_b0) | `stochastic.sample_coefficients`, `stochastic.chen_two`, `stochastic.train_stochastic` |
+
+## Method sections
+
+| Lecture section | First slide | Code |
+| --- | --- | --- |
+| From velocities to finite motion | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u001_b0) | [continuous.py](continuous.py) |
+| Learning the identities | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u025_b0) | [continuous.py](continuous.py) |
+| Related finite-step formulations | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u041_b0) | [continuous.py](continuous.py) |
+| MeanFlow and the direction of time | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u047_b0) | [continuous.py](continuous.py) |
+| Flow maps in learned latent spaces | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u054_b0) | [posterior.py](posterior.py) |
+| Flow Map Language Models | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u060_b0) | [categorical.py](categorical.py) |
+| Categorical Flow Maps | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u083_b0) | [categorical.py](categorical.py) |
+| Discrete Flow Maps and logit consistency | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u093_b0) | [categorical.py](categorical.py) |
+| Diamond Maps and posterior lookahead | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u103_b0) | [posterior.py](posterior.py) |
+| Meta Flow Maps | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u125_b0) | [posterior.py](posterior.py) |
+| Expanding Flow Maps | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u143_b0) | [expanding.py](expanding.py) |
+| Strong Stochastic Flow Maps | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u173_b0) | [stochastic.py](stochastic.py) |
+| Connecting the formulations | [Open](https://docs.google.com/presentation/d/1wfLAazqYveMoyvdy0y5ILeionp7eUjcPgxu-5QFlWwI/edit#slide=id.fm6_u206_b0) | [stochastic.py](stochastic.py) |
+
+The exponential-flow code slides use the exact scalar drift `b(x)=x`; the generative runners use learned fields on a Gaussian mixture. The scalar training experiment remains in `numerical_examples.py`. The SSFM consistency code slide corresponds to `--ssfm-target paper`; the default follows the released code orientation, described in `SOURCES.md`.

lecture_6/SOURCES.md

@@ -0,0 +1,97 @@
+# Primary papers and author-code checks
+
+The teaching implementations were written for this course. Upstream training files were inspected to check objectives and conventions; their code is not vendored here. Small networks, synthetic data, and explicit numerical choices make the methods runnable on CPU. Inspection date: 2026-09-13.
+
+## Flow Map Matching
+
+[Paper](https://arxiv.org/abs/2406.07507v2). Class methods: `fmm-lagrangian, fmm-eulerian`.
+
+[Inspected author source](https://github.com/nmboffi/flow-maps/blob/2f115a07fa9073553193e4b265dfc303827af2b0/py/common/losses.py), revision `2f115a07fa90`.
+
+Frozen learned velocity teachers; residual-map JVPs; small MLP and bounded time range.
+
+## How to build a consistency model: Learning flow maps via self-distillation
+
+[Paper](https://arxiv.org/abs/2505.18825v2). Class methods: `self-distill`.
+
+[Inspected author source](https://github.com/nmboffi/flow-maps/blob/2f115a07fa9073553193e4b265dfc303827af2b0/py/common/losses.py), revision `2f115a07fa90`.
+
+Diagonal regression plus an EMA progressive target, normalized by a floored interval length. Learned uncertainty weights are omitted.
+
+## One Step Diffusion via Shortcut Models
+
+[Paper](https://arxiv.org/abs/2410.12557v3). Class methods: `shortcut`.
+
+[Inspected author source](https://github.com/kvfrans/shortcut-models/blob/601004348667094e1b71f30942199759412d4432/targets_shortcut.py), revision `601004348667`.
+
+Two half-step average target and dyadic intervals retained. Continuous start-time sampling and small MLP replace image-specific schedules and clipping.
+
+## Mean Flows for One-step Generative Modeling
+
+[Paper](https://arxiv.org/abs/2505.13447). Class methods: `meanflow`.
+
+[Inspected author source](https://github.com/Gsunshine/meanflow/blob/d70cb55d298ee03c53bf6da67bec281082e4e2d9/meanflow.py), revision `d70cb55d298e`.
+
+Backward clock, conditional-velocity JVP, detached target, adaptive loss, and 75% diagonal proportion retained. Uniform time sampling; no class guidance.
+
+## Flow Map Language Models: One-step Language Modeling via Continuous Denoising
+
+[Paper](https://arxiv.org/abs/2602.16813v3). Class methods: `fmlm`.
+
+[Inspected author source](https://github.com/david3684/flm/blob/a1918d5164e5038e37d0b7a4fb2010ce75b863b3/algo.py), revision `a1918d5164e5`.
+
+PSD denoiser target and Gaussian decoding clock retained. Independent clean-data diagonal supervision; EMA targets; quadrature lookup and MLP replace full language architecture.
+
+## Categorical Flow Maps
+
+[Paper](https://arxiv.org/abs/2602.12233v1). Class methods: `categorical`.
+
+[Inspected author source](https://github.com/olsdavis/semicat/blob/558602a0fa722514e4a6012f5c46a8ae178b3068/semicat/models/semicat.py), revision `558602a0fa72`.
+
+Released ECLD endpoint CE plus sum-of-squares time energy. Detached-target KL has identical student gradient. EMA target replaces the current network target.
+
+## Discrete Flow Maps
+
+[Paper](https://arxiv.org/abs/2604.09784v1). Class methods: `discrete-lsd, discrete-esd`.
+
+Derived from the paper. The linked project page has a placeholder Code link, so no author implementation was verified. Positive log denominators are explicitly clamped and counted; no gradient surgery.
+
+## Diamond Maps: Stochastic Flow Maps
+
+[Paper](https://arxiv.org/abs/2602.05993). Class methods: `diamond`.
+
+[Inspected author source](https://github.com/PeterHolderrieth/diamond_maps/blob/d30f65c75a169a2ed624f146b2770aeef882543a/posterior_diamond_maps/py/common/losses.py), revision `d30f65c75a16`.
+
+GLASS conditional velocity and Lagrangian posterior distillation. Exact Gaussian-mixture denoiser replaces a pretrained image teacher. Weighted posterior recovery uses an explicit Gaussian proposal with a known density.
+
+## Meta Flow Maps enable scalable reward alignment
+
+[Paper](https://arxiv.org/abs/2601.14430v2). Class methods: `meta`.
+
+[Inspected author source](https://github.com/adh1s/mfm/blob/53c0f60db695cad88bbace8fb26469614e9e7d7d/src/mfm/losses/losses.py), revision `53c0f60db695`.
+
+Data training with independent inner/outer noises and a shared endpoint; conditional semigroup objective. Optional Equation 43 fine-tuning also checked against src/mfm/losses/finetune.py.
+
+## Expanding Flow Maps
+
+[Paper](https://arxiv.org/abs/2607.21585v1). Class methods: `expanding`.
+
+[Inspected author source](https://github.com/sophtang/ExpandingFlowMaps/blob/4defc1ff168d12526d85b551bea78150b9aec605/README.md), revision `4defc1ff168d`.
+
+Paper Algorithms 2-4 and local-clock equations. The inspected repository contains README and artwork only. Classroom implementation includes a shared lifted canvas, remaining/interval count heads, explicit local clocks, and budget-capped binomial insertions.
+
+## Strong Stochastic Flow Maps
+
+[Paper](https://arxiv.org/abs/2606.01086v1). Class methods: `ssfm`.
+
+[Inspected author source](https://github.com/sammccallum/ssfm/blob/24b563620a8ee61c683d791fe802b56668f98f9b/ssfm/losses.py), revision `24b563620a8e`.
+
+First two Legendre integrals, Chen composition, known diffusion, small-step matching. The released loss detaches an EMA split target; paper Algorithm 1 detaches the direct target. Both orientations are selectable. OU replaces image/molecular experiments; no learned uncertainty weights.
+
+## Supporting lecture demonstrations
+
+`flow-matching` supplies the local velocity baseline. `consistency` demonstrates endpoint consistency on a learned flow teacher using the forward clock and a residual endpoint parameterization; it is not a reproduction of the original EDM consistency training system. `latent` demonstrates the encoder/map/decoder construction and does not attribute it to a separate paper called Latent Flow Maps. `numerical_examples.py` retains the lecture’s scalar ODE experiment and exact arithmetic checks.
+
+## Version and result scope
+
+Paper links retain the lecture versions when supplied. Author repositories may have advanced after those paper versions. The source revision above records exactly what was inspected. The lecture’s paper figures and benchmark values remain in the slide deck; the files under `verified_examples/` contain newly executed course examples, with their own data, compute settings, and metrics. No reported FID, generative perplexity, or biological benchmark is reproduced by these CPU teaching runs.

lecture_6/categorical.py

@@ -0,0 +1,145 @@
+"""Simplex-valued mean denoisers for Flow Map Language Models and discrete maps.
+
+The state is Gaussian-noised one-hot text and can be outside the simplex.
+Only the softmax endpoint prediction is a probability vector.
+"""
+import math
+
+import torch
+from torch.nn import functional as F
+
+from common import SequenceNet, draw_batch, ema_copy, interpolate, optimize, time_like
+
+
+def categorical_map(model, x, s, t):
+    s, t = time_like(s, x), time_like(t, x)
+    probability = model(x, s, t).softmax(-1)
+    h = ((t - s) / (1 - s)).unsqueeze(-1)
+    return x + h * (probability - x), probability
+
+
+def composition_target(model, x, s, u, t):
+    mid, p1 = categorical_map(model, x, s, u)
+    _, p2 = categorical_map(model, mid, u, t)
+    gamma = ((1 - t) * (u - s) / ((1 - u) * (t - s))).unsqueeze(-1)
+    return gamma * p1 + (1 - gamma) * p2
+
+
+def probability_kl(logits, target):
+    return F.kl_div(logits.log_softmax(-1), target, reduction='none').sum(-1).mean()
+
+
+class DecodingClock:
+    """Inverse of the FMLM decoding-accuracy clock for isotropic Gaussian noise.
+
+    P(correct) = E_Z[Phi(Z + t/(1-t))^(V-1)]. Gaussian quadrature gives a
+    deterministic lookup. Endpoint values are exact. No empirical vocabulary
+    frequency estimate is mixed into this corruption-only clock.
+    """
+    def __init__(self, vocab, device='cpu'):
+        import numpy as np
+        nodes, weights = np.polynomial.hermite.hermgauss(64)
+        self.t = torch.linspace(0, .999, 1001, device=device)
+        z = torch.as_tensor(nodes * math.sqrt(2), device=device, dtype=torch.float64)
+        w = torch.as_tensor(weights / math.sqrt(math.pi), device=device, dtype=torch.float64)
+        snr = (self.t / (1 - self.t)).double()
+        cdf = .5 * (1 + torch.erf((z[None] + snr[:, None]) / math.sqrt(2)))
+        accuracy = (cdf.pow(vocab - 1) * w).sum(-1)
+        tau = ((vocab * accuracy - 1) / (vocab - 1)).clamp(0, 1).float()
+        self.tau = torch.cummax(tau, 0).values
+        self.tau[0], self.tau[-1], self.t[-1] = 0., 1., 1.
+
+    def inverse(self, tau):
+        indices = torch.searchsorted(self.tau, tau.contiguous()).clamp(1, len(self.tau)-1)
+        low, high = self.tau[indices-1], self.tau[indices]
+        alpha = (tau-low) / (high-low).clamp_min(1e-7)
+        result = self.t[indices-1] + alpha * (self.t[indices]-self.t[indices-1])
+        # Finite precision can saturate the lookup before t=1. Enforce the
+        # mathematical endpoint so generation actually reaches the simplex.
+        return torch.where(tau>=1,torch.ones_like(result),torch.where(tau<=0,torch.zeros_like(result),result))
+
+
+def corrected_logit_teacher(model, teacher, x, s, t, kind, floor=.05):
+    """Lagrangian/Eulerian logit teachers and explicit domain stabilization.
+
+    At an exact solution the denominators are positive. Unconverged networks
+    can violate that domain. We clamp to floor and report the changed fraction.
+    These clamped off-solution targets are an explicit numerical approximation.
+    """
+    if kind == 'discrete-lsd':
+        logits, dz = torch.func.jvp(lambda end: model(x, s, end),
+                                   (t,), (torch.ones_like(t),))
+        p = logits.softmax(-1)
+        delta = dz - (p * dz).sum(-1, keepdim=True)
+        coefficient = ((t-s) * (1-t) / (1-s)).unsqueeze(-1)
+        denominator = 1 + coefficient * delta
+        with torch.no_grad():
+            y, _ = categorical_map(teacher, x, s, t)
+            base_logits = teacher(y, t, t)
+    else:
+        with torch.no_grad():
+            p0 = teacher(x, s, s).softmax(-1)
+            velocity = (p0 - x) / (1-s).unsqueeze(-1)
+            base_logits = teacher(x, s, s)
+        logits, dz = torch.func.jvp(lambda z, start: model(z, start, t),
+                                   (x, s), (velocity, torch.ones_like(s)))
+        p = logits.softmax(-1)
+        delta = dz - (p * dz).sum(-1, keepdim=True)
+        coefficient = ((1-s)*(t-s)/(1-t)).unsqueeze(-1)
+        denominator = 1 - coefficient * delta
+    target = (base_logits - denominator.clamp_min(floor).log()).softmax(-1).detach()
+    return logits, target, (denominator.detach() < floor).float().mean()
+
+
+def train_categorical(method, ids, vocab, args):
+    model = SequenceNet(ids.shape[1], len(vocab), args.width).to(ids.device)
+    teacher = ema_copy(model)
+    clock = DecodingClock(len(vocab), ids.device)
+
+    def objective(step):
+        batch = draw_batch(ids, args.batch_size)
+        clean = F.one_hot(batch, len(vocab)).float()
+        ordered = clock.inverse(torch.rand(len(batch), 2, device=ids.device)).sort(-1).values
+        # Stay away from singular coefficients in differential teacher targets.
+        s, t = (.97 * ordered).split(1, -1)
+        t = torch.maximum(t, s + 1e-4)
+        x, _ = interpolate(clean, s)
+        diagonal = F.cross_entropy(model(x, s, s).flatten(0, 1), batch.flatten())
+        logits = model(x, s, t)
+        extra = {}
+        if method == 'fmlm':
+            with torch.no_grad():
+                target = composition_target(teacher, x, s, (s+t)/2, t)
+            finite = probability_kl(logits, target)
+        elif method == 'categorical':
+            y, _ = categorical_map(model, x, s, t)
+            with torch.no_grad():
+                endpoint = teacher(y.detach(), t, t).softmax(-1)
+            ec = probability_kl(logits, endpoint)
+            _, dt = torch.func.jvp(lambda end: model(x, s, end).softmax(-1),
+                                   (t,), (torch.ones_like(t),))
+            eta = ((t-s)/(1-s)).unsqueeze(-1)
+            temporal = (eta * dt).square().flatten(1).sum(-1).mean()
+            # Released ECLD implementation: CE + sum-of-squares time energy.
+            # Detached-target KL has the same student gradient as its CE term.
+            finite = ec + temporal
+            extra = {'endpoint_kl': ec, 'temporal_derivative': temporal}
+        else:
+            logits, target, fraction = corrected_logit_teacher(model, teacher, x, s, t, method)
+            finite = probability_kl(logits, target)
+            extra = {'denominator_clamp_fraction': fraction}
+        weight = min(1., (step + 1) / max(1, args.train_steps // 4))
+        return diagonal + weight * finite, {'diagonal_ce': diagonal, 'finite': finite, **extra}
+
+    logs = optimize(model, objective, args.train_steps, args.lr, teacher)
+    return model, {'model': model.state_dict(), 'length': ids.shape[1], 'vocab': vocab}, logs
+
+
+@torch.no_grad()
+def sample_categorical(model, count, steps, device):
+    clock = DecodingClock(model.vocab, device)
+    grid = clock.inverse(torch.linspace(0, 1, steps+1, device=device))
+    x = torch.randn(count, model.length, model.vocab, device=device)
+    for s, t in zip(grid[:-1], grid[1:]):
+        x, _ = categorical_map(model, x, s, t)
+    return x.argmax(-1)

lecture_6/common.py

@@ -0,0 +1,207 @@
+"""Shared networks, data, and optimization for CIS 6270 Lecture 6.
+
+The small MLPs keep time derivatives and all training steps visible. Sequence
+networks receive the whole padded sequence, so predictions can depend on other
+positions. No pretrained checkpoint or dataset download is required.
+"""
+import copy
+import json
+import math
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+def seed_all(seed, threads=1):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.set_num_threads(threads)
+
+
+def mlp(inputs, outputs, width):
+    return nn.Sequential(nn.Linear(inputs, width), nn.SiLU(),
+                         nn.Linear(width, width), nn.SiLU(),
+                         nn.Linear(width, outputs))
+
+
+def time_like(t, x):
+    """One scalar time per batch member, shape [B,1]."""
+    t = torch.as_tensor(t, dtype=x.dtype, device=x.device)
+    if t.numel() == 1:
+        return t.expand(len(x), 1)
+    return t.reshape(len(x), 1)
+
+
+class MapNet(nn.Module):
+    def __init__(self, dim=2, width=64, context_dim=0):
+        super().__init__()
+        self.net = mlp(dim + 2 + context_dim, dim, width)
+
+    def forward(self, x, s, t, context=None):
+        inputs = [x, time_like(s, x), time_like(t, x)]
+        if context is not None:
+            inputs.append(context)
+        return self.net(torch.cat(inputs, -1))
+
+
+class SequenceNet(nn.Module):
+    def __init__(self, length, vocab, width=64):
+        super().__init__()
+        self.length, self.vocab = length, vocab
+        self.net = mlp(length * vocab + 2, length * vocab, width)
+
+    def forward(self, x, s, t):
+        inputs = torch.cat([x.flatten(1), time_like(s, x), time_like(t, x)], -1)
+        return self.net(inputs).reshape(-1, self.length, self.vocab)
+
+
+def finite_map(model, x, s, t, context=None):
+    """F(s,t,x) = x + (t-s) v(s,t,x), including exact F(s,s,x)=x."""
+    s, t = time_like(s, x), time_like(t, x)
+    return x + (t - s) * model(x, s, t, context)
+
+
+def ordered_times(x, ceiling=.98):
+    times = ceiling * torch.rand(len(x), 2, device=x.device, dtype=x.dtype)
+    s, t = times.sort(-1).values.split(1, -1)
+    return s, t, (s + t) / 2
+
+
+def draw_batch(data, count):
+    return data[torch.randint(len(data), (count,), device=data.device)]
+
+
+def interpolate(data, time, noise=None):
+    noise = torch.randn_like(data) if noise is None else noise
+    t = time.reshape(len(data), *([1] * (data.ndim - 1)))
+    return (1 - t) * noise + t * data, data - noise
+
+
+def ema_copy(model):
+    result = copy.deepcopy(model).eval()
+    result.requires_grad_(False)
+    return result
+
+
+@torch.no_grad()
+def update_ema(ema, model, decay=.99):
+    for p, q in zip(ema.parameters(), model.parameters()):
+        p.lerp_(q, 1 - decay)
+
+
+def optimize(model, objective, steps, lr=1e-3, ema=None):
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+    logs = []
+    for step in range(steps):
+        loss, details = objective(step)
+        if not torch.isfinite(loss):
+            raise FloatingPointError(f'Nonfinite loss at step {step}: {details}')
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 10.)
+        if not torch.isfinite(norm):
+            raise FloatingPointError(f'Nonfinite gradient at step {step}')
+        optimizer.step()
+        if ema is not None:
+            update_ema(ema, model)
+        logs.append({'step': step, 'loss': float(loss.detach()),
+                     **{k: float(torch.as_tensor(v).detach()) for k, v in details.items()}})
+    return logs
+
+
+CENTERS = torch.tensor([[-1.5, -1.5], [-1.5, 1.5], [1.5, -1.5], [1.5, 1.5]])
+DATA_STD = .22
+
+
+def mixture_data(n, generator=None):
+    ids = torch.randint(4, (n,), generator=generator)
+    return CENTERS[ids] + DATA_STD * torch.randn(n, 2, generator=generator)
+
+
+def mixture_posterior(x, t):
+    """Exact p(X1 | (1-t)X0+tX1=x) for the four Gaussian mixture.
+
+    Return component probabilities, conditional means, scalar variances.
+    This analytic oracle is used for diagnostics and the GLASS example.
+    """
+    t = time_like(t, x)
+    a = 1 - t
+    variance = a.square() + (t * DATA_STD).square()
+    centers = CENTERS.to(x)
+    delta = x[:, None, :] - t[:, None, :] * centers
+    logits = -delta.square().sum(-1) / (2 * variance)
+    prob = logits.softmax(-1)
+    gain = t * DATA_STD**2 / variance
+    means = centers + gain[:, None, :] * delta
+    posterior_variance = DATA_STD**2 * a.square() / variance
+    return prob, means, posterior_variance
+
+
+def exact_denoiser(x, t):
+    p, means, _ = mixture_posterior(x, t)
+    return (p[..., None] * means).sum(1)
+
+
+def exact_velocity(x, t):
+    t = time_like(t, x)
+    # Direct conditional velocity avoids cancellation at t=1.
+    a = 1 - t
+    variance = a.square() + (t * DATA_STD).square()
+    centers = CENTERS.to(x)
+    delta = x[:, None] - t[:, None] * centers
+    weights = (-delta.square().sum(-1) / (2 * variance)).softmax(-1)
+    component_velocity = centers + ((t * DATA_STD**2 - a) / variance)[:, None] * delta
+    return (weights[..., None] * component_velocity).sum(1)
+
+
+def mixture_metrics(samples):
+    x = samples.detach().cpu()
+    distances = (x[:, None] - CENTERS).square().sum(-1)
+    counts = torch.bincount(distances.argmin(-1), minlength=4).float()
+    p = counts / len(x)
+    logp = torch.logsumexp(-distances / (2 * DATA_STD**2), -1)
+    logp -= math.log(4 * 2 * math.pi * DATA_STD**2)
+    return {'finite_samples': bool(torch.isfinite(x).all()),
+            'mean_distance_to_center': float(distances.min(-1).values.sqrt().mean()),
+            'mode_fractions': p.tolist(), 'mode_entropy': float(-(p * p.clamp_min(1e-12).log()).sum()),
+            'mean_target_log_density': float(logp.mean())}
+
+
+def load_text(path, variable=False):
+    lines = [x.strip().split() for x in Path(path).read_text().splitlines() if x.strip()]
+    if len(lines) < 2:
+        raise ValueError('Text data need at least two nonempty lines.')
+    vocab = sorted(set(word for line in lines for word in line))
+    lengths = torch.tensor([len(x) for x in lines])
+    if not variable and len(set(lengths.tolist())) != 1:
+        raise ValueError('Fixed-length methods require equal tokens per line; use expanding otherwise.')
+    if lengths.max() > 32:
+        raise ValueError('Teaching MLP supports up to 32 positions.')
+    ids = {word: i for i, word in enumerate(vocab)}
+    encoded = torch.zeros(len(lines), int(lengths.max()), dtype=torch.long)
+    for i, line in enumerate(lines):
+        encoded[i, :len(line)] = torch.tensor([ids[word] for word in line])
+    return encoded, lengths, vocab
+
+
+def text_metrics(ids, vocab, lengths=None, reference=None):
+    lengths = [ids.shape[1]] * len(ids) if lengths is None else lengths.tolist()
+    texts = [' '.join(vocab[j] for j in row[:length]) for row, length in zip(ids.tolist(), lengths)]
+    flat = [j for row, length in zip(ids.tolist(), lengths) for j in row[:length]]
+    counts = torch.bincount(torch.tensor(flat, dtype=torch.long), minlength=len(vocab)).float()
+    p = counts / counts.sum().clamp_min(1)
+    metrics = {'unique_fraction': len(set(texts)) / len(texts),
+               'token_entropy': float(-(p * p.clamp_min(1e-12).log()).sum()),
+               'mean_length': sum(lengths) / len(lengths), 'empty_fraction': lengths.count(0) / len(lengths)}
+    if reference is not None:
+        metrics['training_support_fraction'] = sum(x in reference for x in texts) / len(texts)
+    return texts, metrics
+
+
+def write_json(path, value):
+    Path(path).write_text(json.dumps(value, indent=2, allow_nan=False) + '\n')

lecture_6/continuous.py

@@ -0,0 +1,183 @@
+"""Flow matching, finite-map identities, Shortcut, MeanFlow, and latent maps.
+
+# %% 1. Learn local motion before learning an interval.
+All times run noise -> data except MeanFlow, whose original backward clock
+r <= t runs data at zero -> noise at one during training.
+"""
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from common import (MapNet, draw_batch, ema_copy, finite_map, interpolate, mlp,
+                    optimize, ordered_times, time_like)
+
+
+def diagonal_loss(model, data):
+    t = torch.rand(len(data), 1, device=data.device)
+    x, displacement = interpolate(data, t)
+    return F.mse_loss(model(x, t, t), displacement)
+
+
+def lagrangian_residual(model, velocity, x, s, t):
+    # d_t F(s,t,x) = b(t,F(s,t,x)). JVP differentiates only arrival time.
+    y, dt = torch.func.jvp(lambda end: finite_map(model, x, s, end),
+                         (t,), (torch.ones_like(t),))
+    return dt - velocity(y.detach(), t).detach()
+
+
+def eulerian_residual(model, velocity, x, s, t):
+    # (d_s + b_s dot grad_x)F = 0. No full Jacobian is materialized.
+    direction = velocity(x, s).detach()
+    _, residual = torch.func.jvp(lambda z, start: finite_map(model, z, start, t),
+                               (x, s), (direction, torch.ones_like(s)))
+    return residual
+
+
+def semigroup_loss(model, teacher, x, s, t, context=None):
+    u = (s + t) / 2
+    with torch.no_grad():
+        mid = finite_map(teacher, x, s, u, context)
+        target = finite_map(teacher, mid, u, t, context)
+    prediction = finite_map(model, x, s, t, context)
+    return ((prediction - target) / (t - s).clamp_min(.05)).square().mean()
+
+
+def meanflow_loss(model, data):
+    # %% 2. Average backward velocity satisfies u = v - (t-r) D_t u.
+    noise = torch.randn_like(data)
+    r, t, _ = ordered_times(data, 1.)
+    # Include diagonal examples with nonzero probability, as in the paper.
+    r = torch.where(torch.rand_like(r) < .75, t, r)
+    z = (1 - t) * data + t * noise
+    v = noise - data
+    average, derivative = torch.func.jvp(model, (z, r, t),
+        (v, torch.zeros_like(r), torch.ones_like(t)))
+    target = (v - (t - r) * derivative).detach()
+    residual = (average - target).square().mean(-1)
+    # Detached adaptive weighting controls large self-distillation residuals.
+    weight = (residual.detach() + .01).pow(-.5)
+    return (weight * residual).mean(), {'mse': residual.mean()}
+
+
+def shortcut_loss(model, teacher, data):
+    t = torch.rand(len(data), 1, device=data.device)
+    x, v = interpolate(data, t)
+    diag = F.mse_loss(model(x, t, t), v)
+    # d is the half-step. Learn the 2d shortcut from two d shortcuts.
+    powers = torch.randint(1, 6, (len(data), 1), device=data.device)
+    d = 2. ** (-powers)
+    start = torch.rand_like(d) * (1 - 2 * d)
+    x, _ = interpolate(data, start)
+    with torch.no_grad():
+        first = teacher(x, start, start + d)
+        second = teacher(x + d * first, start + d, start + 2 * d)
+        target = .5 * (first + second)
+    finite = F.mse_loss(model(x, start, start + 2 * d), target)
+    return diag + finite, {'diagonal': diag, 'shortcut': finite}
+
+
+def train_continuous(method, data, args):
+    model = MapNet(data.shape[1], args.width).to(data)
+    teacher_logs = []
+    teacher = None
+    if method in {'fmm-lagrangian', 'fmm-eulerian', 'consistency'}:
+        teacher = MapNet(data.shape[1], args.width).to(data)
+        teacher_logs = optimize(teacher, lambda _: (diagonal_loss(teacher, draw_batch(data, args.batch_size)), {}),
+                                args.teacher_steps, args.lr)
+        teacher.requires_grad_(False)
+        model.load_state_dict(teacher.state_dict())
+    ema = ema_copy(model)
+
+    def objective(step):
+        batch = draw_batch(data, args.batch_size)
+        if method == 'meanflow':
+            return meanflow_loss(model, batch)
+        if method == 'shortcut':
+            return shortcut_loss(model, ema, batch)
+        diag = diagonal_loss(model, batch)
+        if method == 'flow-matching':
+            return diag, {'diagonal': diag}
+        s, t, _ = ordered_times(batch)
+        x, _ = interpolate(batch, s)
+        if method == 'fmm-lagrangian':
+            finite = lagrangian_residual(model, lambda z, a: teacher(z, a, a), x, s, t).square().mean()
+        elif method == 'fmm-eulerian':
+            finite = eulerian_residual(model, lambda z, a: teacher(z, a, a), x, s, t).square().mean()
+        elif method == 'consistency':
+            # Endpoint consistency distillation on teacher-solved short intervals.
+            # At t=1 the residual endpoint parametrization is exactly identity.
+            t = (s + .1).clamp_max(1.)
+            with torch.no_grad():
+                y = integrate_velocity(lambda z, a: teacher(z, a, a), x, s, t, 4)
+                target = finite_map(ema, y, t, torch.ones_like(t))
+            finite = F.mse_loss(finite_map(model, x, s, torch.ones_like(s)), target)
+        else:
+            finite = semigroup_loss(model, ema, x, s, t)
+        weight = min(1., (step + 1) / max(1, args.train_steps // 5))
+        return diag + weight * finite, {'diagonal': diag, 'finite': finite}
+
+    logs = optimize(model, objective, args.train_steps, args.lr, ema)
+    # Use the trained student; EMA supplies fixed bootstrap targets during training.
+    state = {'model': model.state_dict(), 'dim': data.shape[1]}
+    if teacher is not None:
+        state['teacher'] = teacher.state_dict()
+    return model, state, logs, teacher_logs
+
+
+@torch.no_grad()
+def integrate_velocity(velocity, x, s, t, steps):
+    """Heun integration of a velocity, also supports batch-specific time bounds."""
+    s, t = time_like(s, x), time_like(t, x)
+    h = (t - s) / steps
+    for i in range(steps):
+        a = s + i * h
+        k1 = velocity(x, a)
+        k2 = velocity(x + h * k1, a + h)
+        x = x + h * .5 * (k1 + k2)
+    return x
+
+
+@torch.no_grad()
+def sample_continuous(model, method, noise, steps):
+    if method == 'flow-matching':
+        return integrate_velocity(lambda z, a: model(z, a, a), noise, 0., 1., steps)
+    if method == 'consistency':
+        return finite_map(model, noise, 0., 1.)
+    x = noise
+    for i in range(steps):
+        if method == 'meanflow':
+            t, r = 1 - i / steps, 1 - (i + 1) / steps
+            x = x - (t - r) * model(x, r, t)
+        else:
+            x = finite_map(model, x, i / steps, (i + 1) / steps)
+    return x
+
+
+class Autoencoder(nn.Module):
+    def __init__(self, width):
+        super().__init__()
+        self.encoder = mlp(3, 2, width)
+        self.decoder = mlp(2, 3, width)
+
+
+def embed_surface(x):
+    return torch.cat([x, .3 * (x[:, :1].square() - x[:, 1:].square())], -1)
+
+
+def train_latent(data, args):
+    # %% 3. First fit the representation, then freeze it while fitting its flow.
+    surface = embed_surface(data)
+    ae = Autoencoder(args.width).to(data)
+    def objective(_):
+        batch = draw_batch(surface, args.batch_size)
+        reconstruction = ae.decoder(ae.encoder(batch))
+        return F.mse_loss(reconstruction, batch), {}
+    ae_logs = optimize(ae, objective, args.teacher_steps, args.lr)
+    ae.requires_grad_(False)
+    with torch.no_grad():
+        z = ae.encoder(surface)
+        mean, std = z.mean(0), z.std(0).clamp_min(.05)
+        z = (z - mean) / std
+    model, state, logs, _ = train_continuous('self-distill', z, args)
+    state.update({'autoencoder': ae.state_dict(), 'latent_mean': mean, 'latent_std': std})
+    return model, ae, state, logs, ae_logs

lecture_6/data/README.md

@@ -0,0 +1,34 @@
+# Data for Lecture 6
+
+`phrases.txt` and `variable_text.txt` were generated deterministically with Python
+`random.Random(6270)`. Each contains 640 examples. Their vocabulary describes
+colored shapes and directions; no external dataset is needed.
+
+| Color | Shape | Direction |
+| --- | --- | --- |
+| red | circle | left |
+| blue | square | right |
+| green | triangle | up |
+| gold | star | down |
+
+The fixed-length corpus has four tokens per line. The variable-length corpus
+uses `COLOR SHAPE`, `COLOR SHAPE moves DIRECTION`, or the latter followed by
+`slowly` or `twice`. The color, shape, and direction remain correlated. This
+makes invalid independent combinations easy to inspect.
+
+Training uses a seeded shuffle and 80/20 example split. The grammar has a small
+support, so examples repeat and both partitions contain the same grammar.
+Training-support fraction measures membership in this small support; it is
+not a test of open-ended language generalization.
+
+For a custom corpus, supply a UTF-8 text file with one whitespace-tokenized
+sequence per line. Fixed-length methods require equal token counts. The
+expanding method accepts variable lengths up to 32. Vocabulary and maximum
+length are saved in the checkpoint, so sampling does not need the original file.
+
+The continuous dataset is sampled from four equally weighted Gaussians centered
+at `(−1.5,−1.5)`, `(−1.5,1.5)`, `(1.5,−1.5)`, and `(1.5,1.5)`, with coordinate
+standard deviation 0.22. The latent example embeds these points on
+`(x, y, 0.3(x²−y²))` before learning its autoencoder. SSFM uses analytic
+Ornstein-Uhlenbeck marginals with initial variance one, drift `−x`, and diffusion
+coefficient 0.7.

lecture_6/data/phrases.txt

@@ -0,0 +1,640 @@
+green triangle moves up
+gold star moves down
+green triangle moves up
+gold star moves down
+blue square moves right
+gold star moves down
+green triangle moves up
+red circle moves left
+blue square moves right
+green triangle moves up
+gold star moves down
+gold star moves down
+red circle moves left
+blue square moves right
+blue square moves right
+red circle moves left
+gold star moves down
+blue square moves right
+red circle moves left
+gold star moves down
+green triangle moves up
+red circle moves left
+red circle moves left
+green triangle moves up
+green triangle moves up
+green triangle moves up
+red circle moves left
+red circle moves left
+green triangle moves up
+green triangle moves up
+red circle moves left
+gold star moves down
+red circle moves left
+red circle moves left
+gold star moves down
+gold star moves down
+red circle moves left
+red circle moves left
+red circle moves left
+blue square moves right
+red circle moves left
+blue square moves right
+gold star moves down
+blue square moves right
+red circle moves left
+gold star moves down
+blue square moves right
+blue square moves right
+green triangle moves up
+red circle moves left
+green triangle moves up
+blue square moves right
+red circle moves left
+green triangle moves up
+red circle moves left
+blue square moves right
+green triangle moves up
+green triangle moves up
+green triangle moves up
+red circle moves left
+red circle moves left
+red circle moves left
+gold star moves down
+gold star moves down
+red circle moves left
+gold star moves down
+blue square moves right
+blue square moves right
+red circle moves left
+red circle moves left
+blue square moves right
+green triangle moves up
+blue square moves right
+blue square moves right
+gold star moves down
+green triangle moves up
+red circle moves left
+green triangle moves up
+gold star moves down
+gold star moves down
+gold star moves down
+gold star moves down
+blue square moves right
+gold star moves down
+red circle moves left
+red circle moves left
+green triangle moves up
+gold star moves down
+blue square moves right
+gold star moves down
+red circle moves left
+red circle moves left
+red circle moves left
+green triangle moves up
+blue square moves right
+red circle moves left
+blue square moves right
+gold star moves down
+gold star moves down
+red circle moves left
+green triangle moves up
+green triangle moves up
+gold star moves down
+red circle moves left
+green triangle moves up
+green triangle moves up
+blue square moves right
+green triangle moves up
+gold star moves down
+green triangle moves up
+green triangle moves up
+blue square moves right
+red circle moves left
+blue square moves right
+gold star moves down
+blue square moves right
+blue square moves right
+blue square moves right
+red circle moves left
+blue square moves right
+blue square moves right
+blue square moves right
+gold star moves down
+red circle moves left
+green triangle moves up
+gold star moves down
+red circle moves left
+green triangle moves up
+red circle moves left
+green triangle moves up
+gold star moves down
+blue square moves right
+red circle moves left
+green triangle moves up
+red circle moves left
+blue square moves right
+green triangle moves up
+gold star moves down
+gold star moves down
+gold star moves down
+green triangle moves up
+gold star moves down
+blue square moves right
+blue square moves right
+blue square moves right
+red circle moves left
+gold star moves down
+green triangle moves up
+green triangle moves up
+red circle moves left
+blue square moves right
+gold star moves down
+gold star moves down
+gold star moves down
+gold star moves down
+gold star moves down
+gold star moves down
+red circle moves left
+red circle moves left
+blue square moves right
+blue square moves right
+blue square moves right
+gold star moves down
+blue square moves right
+red circle moves left
+blue square moves right
+gold star moves down
+green triangle moves up
+blue square moves right
+green triangle moves up
+gold star moves down
+blue square moves right
+green triangle moves up
+red circle moves left
+blue square moves right
+green triangle moves up
+gold star moves down
+gold star moves down
+green triangle moves up
+blue square moves right
+red circle moves left
+gold star moves down
+gold star moves down
+red circle moves left
+red circle moves left
+green triangle moves up
+blue square moves right
+gold star moves down
+gold star moves down
+green triangle moves up
+red circle moves left
+green triangle moves up
+green triangle moves up
+red circle moves left
+red circle moves left
+gold star moves down
+red circle moves left
+red circle moves left
+green triangle moves up
+blue square moves right
+green triangle moves up
+blue square moves right
+red circle moves left
+gold star moves down
+blue square moves right
+green triangle moves up
+red circle moves left
+red circle moves left
+gold star moves down
+green triangle moves up
+blue square moves right
+red circle moves left
+gold star moves down
+green triangle moves up
+gold star moves down
+red circle moves left
+blue square moves right
+red circle moves left
+blue square moves right
+green triangle moves up
+blue square moves right
+green triangle moves up
+green triangle moves up
+gold star moves down
+blue square moves right
+blue square moves right
+blue square moves right
+green triangle moves up
+blue square moves right
+blue square moves right
+blue square moves right
+green triangle moves up
+blue square moves right
+green triangle moves up
+red circle moves left
+blue square moves right
+green triangle moves up
+red circle moves left
+blue square moves right
+gold star moves down
+gold star moves down
+red circle moves left
+red circle moves left
+blue square moves right
+red circle moves left
+red circle moves left
+gold star moves down
+gold star moves down
+green triangle moves up
+red circle moves left
+gold star moves down
+red circle moves left
+red circle moves left
+gold star moves down
+gold star moves down
+gold star moves down
+red circle moves left
+green triangle moves up
+gold star moves down
+green triangle moves up
+blue square moves right
+gold star moves down
+green triangle moves up
+red circle moves left
+green triangle moves up
+red circle moves left
+green triangle moves up
+green triangle moves up
+blue square moves right
+red circle moves left
+blue square moves right
+gold star moves down
+gold star moves down
+blue square moves right
+blue square moves right
+red circle moves left
+red circle moves left
+blue square moves right
+green triangle moves up
+green triangle moves up
+blue square moves right
+red circle moves left
+green triangle moves up
+blue square moves right
+red circle moves left
+green triangle moves up
+gold star moves down
+red circle moves left
+red circle moves left
+blue square moves right
+gold star moves down
+red circle moves left
+red circle moves left
+blue square moves right
+red circle moves left
+gold star moves down
+red circle moves left
+blue square moves right
+red circle moves left
+gold star moves down
+green triangle moves up
+blue square moves right
+gold star moves down
+gold star moves down
+gold star moves down
+green triangle moves up
+gold star moves down
+red circle moves left
+gold star moves down
+blue square moves right
+red circle moves left
+green triangle moves up
+green triangle moves up
+gold star moves down
+green triangle moves up
+gold star moves down
+green triangle moves up
+red circle moves left
+blue square moves right
+red circle moves left
+red circle moves left
+blue square moves right
+green triangle moves up
+red circle moves left
+red circle moves left
+gold star moves down
+green triangle moves up
+gold star moves down
+green triangle moves up
+green triangle moves up
+gold star moves down
+red circle moves left
+red circle moves left
+green triangle moves up
+red circle moves left
+green triangle moves up
+blue square moves right
+blue square moves right
+green triangle moves up
+red circle moves left
+green triangle moves up
+green triangle moves up
+blue square moves right
+green triangle moves up
+gold star moves down
+blue square moves right
+green triangle moves up
+green triangle moves up
+blue square moves right
+red circle moves left
+gold star moves down
+red circle moves left
+red circle moves left
+gold star moves down
+blue square moves right
+green triangle moves up
+red circle moves left
+green triangle moves up
+green triangle moves up
+green triangle moves up
+gold star moves down
+green triangle moves up
+gold star moves down
+blue square moves right
+blue square moves right
+green triangle moves up
+green triangle moves up
+blue square moves right
+red circle moves left
+red circle moves left
+blue square moves right
+blue square moves right
+red circle moves left
+red circle moves left
+red circle moves left
+red circle moves left
+blue square moves right
+green triangle moves up
+blue square moves right
+gold star moves down
+blue square moves right
+blue square moves right
+blue square moves right
+green triangle moves up
+red circle moves left
+blue square moves right
+gold star moves down
+green triangle moves up
+gold star moves down
+gold star moves down
+blue square moves right
+gold star moves down
+blue square moves right
+gold star moves down
+gold star moves down
+gold star moves down
+blue square moves right
+gold star moves down
+green triangle moves up
+gold star moves down
+blue square moves right
+gold star moves down
+blue square moves right
+red circle moves left
+gold star moves down
+green triangle moves up
+green triangle moves up
+green triangle moves up
+red circle moves left
+green triangle moves up
+blue square moves right
+blue square moves right
+green triangle moves up
+gold star moves down
+green triangle moves up
+red circle moves left
+gold star moves down
+gold star moves down
+blue square moves right
+red circle moves left
+red circle moves left
+blue square moves right
+blue square moves right
+gold star moves down
+gold star moves down
+blue square moves right
+blue square moves right
+gold star moves down
+blue square moves right
+gold star moves down
+gold star moves down
+red circle moves left
+red circle moves left
+blue square moves right
+blue square moves right
+green triangle moves up
+green triangle moves up
+gold star moves down
+green triangle moves up
+green triangle moves up
+green triangle moves up
+blue square moves right
+green triangle moves up
+gold star moves down
+red circle moves left
+red circle moves left
+gold star moves down
+gold star moves down
+gold star moves down
+blue square moves right
+blue square moves right
+green triangle moves up
+gold star moves down
+blue square moves right
+green triangle moves up
+gold star moves down
+gold star moves down
+blue square moves right
+blue square moves right
+blue square moves right
+gold star moves down
+red circle moves left
+blue square moves right
+green triangle moves up
+green triangle moves up
+green triangle moves up
+blue square moves right
+gold star moves down
+blue square moves right
+gold star moves down
+red circle moves left
+blue square moves right
+green triangle moves up
+green triangle moves up
+green triangle moves up
+green triangle moves up
+gold star moves down
+blue square moves right
+red circle moves left
+green triangle moves up
+blue square moves right
+blue square moves right
+green triangle moves up
+green triangle moves up
+red circle moves left
+blue square moves right
+green triangle moves up
+red circle moves left
+blue square moves right
+gold star moves down
+gold star moves down
+blue square moves right
+blue square moves right
+red circle moves left
+gold star moves down
+blue square moves right
+green triangle moves up
+blue square moves right
+green triangle moves up
+gold star moves down
+green triangle moves up
+green triangle moves up
+green triangle moves up
+blue square moves right
+gold star moves down
+green triangle moves up
+gold star moves down
+red circle moves left
+blue square moves right
+gold star moves down
+gold star moves down
+red circle moves left
+gold star moves down
+green triangle moves up
+gold star moves down
+red circle moves left
+red circle moves left
+blue square moves right
+gold star moves down
+green triangle moves up
+blue square moves right
+red circle moves left
+green triangle moves up
+red circle moves left
+green triangle moves up
+blue square moves right
+red circle moves left
+green triangle moves up
+gold star moves down
+green triangle moves up
+green triangle moves up
+gold star moves down
+green triangle moves up
+blue square moves right
+green triangle moves up
+red circle moves left
+green triangle moves up
+green triangle moves up
+gold star moves down
+red circle moves left
+green triangle moves up
+green triangle moves up
+gold star moves down
+green triangle moves up
+blue square moves right
+green triangle moves up
+green triangle moves up
+green triangle moves up
+blue square moves right
+red circle moves left
+blue square moves right
+red circle moves left
+green triangle moves up
+green triangle moves up
+red circle moves left
+green triangle moves up
+blue square moves right
+red circle moves left
+red circle moves left
+green triangle moves up
+gold star moves down
+blue square moves right
+blue square moves right
+red circle moves left
+green triangle moves up
+red circle moves left
+green triangle moves up
+red circle moves left
+red circle moves left
+blue square moves right
+red circle moves left
+red circle moves left
+red circle moves left
+gold star moves down
+red circle moves left
+green triangle moves up
+green triangle moves up
+green triangle moves up
+red circle moves left
+gold star moves down
+gold star moves down
+gold star moves down
+red circle moves left
+gold star moves down
+green triangle moves up
+green triangle moves up
+blue square moves right
+green triangle moves up
+gold star moves down
+gold star moves down
+green triangle moves up
+gold star moves down
+red circle moves left
+gold star moves down
+green triangle moves up
+blue square moves right
+red circle moves left
+green triangle moves up
+green triangle moves up
+red circle moves left
+green triangle moves up
+green triangle moves up
+red circle moves left
+red circle moves left
+gold star moves down
+blue square moves right
+red circle moves left
+gold star moves down
+gold star moves down
+blue square moves right
+red circle moves left
+blue square moves right
+blue square moves right
+blue square moves right
+red circle moves left
+blue square moves right
+blue square moves right
+green triangle moves up
+blue square moves right
+green triangle moves up
+blue square moves right
+gold star moves down
+gold star moves down
+blue square moves right
+red circle moves left
+red circle moves left
+blue square moves right
+red circle moves left
+green triangle moves up
+red circle moves left
+red circle moves left
+red circle moves left
+gold star moves down
+green triangle moves up
+gold star moves down
+red circle moves left
+blue square moves right
+red circle moves left
+green triangle moves up
+green triangle moves up

lecture_6/data/variable_text.txt

@@ -0,0 +1,640 @@
+green triangle moves up
+gold star moves down twice
+green triangle moves up
+gold star
+blue square moves right
+gold star
+green triangle
+red circle moves left twice
+blue square
+green triangle moves up slowly
+gold star moves down twice
+gold star moves down slowly
+red circle moves left
+blue square moves right slowly
+blue square moves right twice
+red circle moves left
+gold star moves down slowly
+blue square moves right twice
+red circle
+gold star moves down
+green triangle
+red circle
+red circle moves left twice
+green triangle moves up twice
+green triangle moves up
+green triangle moves up slowly
+red circle
+red circle moves left slowly
+green triangle
+green triangle
+red circle moves left twice
+gold star moves down twice
+red circle moves left twice
+red circle
+gold star moves down
+gold star moves down
+red circle
+red circle moves left slowly
+red circle moves left
+blue square moves right twice
+red circle moves left slowly
+blue square
+gold star
+blue square moves right twice
+red circle moves left
+gold star moves down
+blue square moves right
+blue square moves right
+green triangle
+red circle
+green triangle moves up slowly
+blue square moves right
+red circle
+green triangle moves up slowly
+red circle moves left
+blue square moves right slowly
+green triangle
+green triangle moves up slowly
+green triangle moves up slowly
+red circle moves left slowly
+red circle moves left slowly
+red circle moves left slowly
+gold star moves down
+gold star
+red circle moves left
+gold star moves down
+blue square moves right slowly
+blue square moves right
+red circle moves left
+red circle
+blue square moves right twice
+green triangle moves up slowly
+blue square moves right slowly
+blue square
+gold star moves down
+green triangle moves up slowly
+red circle moves left
+green triangle
+gold star
+gold star moves down
+gold star moves down slowly
+gold star moves down twice
+blue square moves right
+gold star moves down
+red circle moves left
+red circle
+green triangle
+gold star moves down
+blue square moves right slowly
+gold star
+red circle moves left
+red circle moves left slowly
+red circle
+green triangle moves up slowly
+blue square moves right
+red circle moves left twice
+blue square moves right twice
+gold star moves down
+gold star moves down twice
+red circle moves left twice
+green triangle moves up slowly
+green triangle
+gold star moves down twice
+red circle moves left
+green triangle moves up slowly
+green triangle moves up
+blue square
+green triangle
+gold star moves down
+green triangle moves up twice
+green triangle moves up slowly
+blue square
+red circle moves left slowly
+blue square moves right twice
+gold star
+blue square moves right
+blue square
+blue square
+red circle
+blue square moves right
+blue square moves right
+blue square moves right slowly
+gold star moves down
+red circle
+green triangle moves up slowly
+gold star moves down
+red circle moves left
+green triangle moves up
+red circle moves left
+green triangle moves up
+gold star
+blue square moves right
+red circle
+green triangle moves up
+red circle moves left twice
+blue square
+green triangle
+gold star
+gold star moves down twice
+gold star
+green triangle moves up
+gold star
+blue square moves right slowly
+blue square
+blue square
+red circle
+gold star moves down slowly
+green triangle moves up
+green triangle
+red circle moves left
+blue square moves right slowly
+gold star moves down
+gold star
+gold star moves down twice
+gold star moves down slowly
+gold star moves down slowly
+gold star moves down
+red circle moves left
+red circle
+blue square moves right
+blue square moves right slowly
+blue square
+gold star
+blue square
+red circle
+blue square moves right
+gold star moves down slowly
+green triangle moves up
+blue square moves right
+green triangle moves up
+gold star
+blue square moves right slowly
+green triangle moves up twice
+red circle moves left slowly
+blue square moves right slowly
+green triangle moves up twice
+gold star moves down
+gold star
+green triangle moves up slowly
+blue square
+red circle moves left slowly
+gold star
+gold star moves down
+red circle moves left
+red circle moves left
+green triangle moves up
+blue square
+gold star moves down
+gold star moves down
+green triangle moves up twice
+red circle moves left
+green triangle
+green triangle
+red circle moves left slowly
+red circle moves left
+gold star moves down slowly
+red circle moves left
+red circle moves left twice
+green triangle moves up slowly
+blue square moves right
+green triangle moves up slowly
+blue square
+red circle moves left twice
+gold star
+blue square
+green triangle
+red circle moves left slowly
+red circle
+gold star moves down slowly
+green triangle
+blue square moves right
+red circle
+gold star
+green triangle moves up twice
+gold star moves down twice
+red circle
+blue square moves right slowly
+red circle moves left slowly
+blue square
+green triangle
+blue square moves right
+green triangle moves up twice
+green triangle moves up slowly
+gold star moves down
+blue square
+blue square moves right
+blue square moves right
+green triangle
+blue square moves right slowly
+blue square moves right slowly
+blue square
+green triangle moves up slowly
+blue square moves right
+green triangle moves up
+red circle moves left slowly
+blue square moves right slowly
+green triangle
+red circle
+blue square moves right
+gold star moves down twice
+gold star moves down
+red circle moves left
+red circle moves left slowly
+blue square moves right slowly
+red circle
+red circle moves left
+gold star
+gold star moves down slowly
+green triangle moves up
+red circle moves left
+gold star moves down twice
+red circle moves left
+red circle moves left slowly
+gold star moves down
+gold star moves down slowly
+gold star
+red circle moves left twice
+green triangle
+gold star moves down slowly
+green triangle
+blue square moves right
+gold star moves down slowly
+green triangle moves up twice
+red circle
+green triangle
+red circle moves left twice
+green triangle
+green triangle moves up slowly
+blue square moves right
+red circle
+blue square
+gold star moves down
+gold star
+blue square
+blue square moves right slowly
+red circle moves left
+red circle moves left slowly
+blue square moves right
+green triangle moves up
+green triangle moves up
+blue square moves right
+red circle moves left twice
+green triangle moves up
+blue square moves right twice
+red circle
+green triangle
+gold star moves down twice
+red circle
+red circle moves left
+blue square moves right
+gold star moves down twice
+red circle moves left slowly
+red circle moves left slowly
+blue square moves right
+red circle moves left
+gold star moves down
+red circle moves left slowly
+blue square
+red circle moves left twice
+gold star moves down
+green triangle
+blue square
+gold star moves down
+gold star moves down
+gold star
+green triangle moves up
+gold star moves down
+red circle moves left
+gold star moves down
+blue square
+red circle
+green triangle
+green triangle moves up
+gold star moves down
+green triangle moves up slowly
+gold star moves down
+green triangle moves up twice
+red circle
+blue square
+red circle moves left twice
+red circle moves left slowly
+blue square moves right
+green triangle moves up
+red circle moves left
+red circle
+gold star moves down
+green triangle moves up slowly
+gold star
+green triangle moves up
+green triangle moves up
+gold star moves down slowly
+red circle
+red circle moves left
+green triangle moves up twice
+red circle
+green triangle moves up
+blue square moves right
+blue square moves right slowly
+green triangle moves up slowly
+red circle moves left
+green triangle moves up twice
+green triangle moves up slowly
+blue square moves right slowly
+green triangle moves up slowly
+gold star moves down twice
+blue square
+green triangle moves up
+green triangle moves up slowly
+blue square moves right twice
+red circle moves left slowly
+gold star moves down
+red circle moves left twice
+red circle moves left slowly
+gold star moves down twice
+blue square moves right
+green triangle moves up
+red circle moves left
+green triangle moves up
+green triangle
+green triangle
+gold star
+green triangle moves up
+gold star moves down
+blue square moves right twice
+blue square moves right slowly
+green triangle
+green triangle moves up slowly
+blue square moves right twice
+red circle
+red circle moves left slowly
+blue square moves right twice
+blue square
+red circle
+red circle moves left
+red circle moves left
+red circle moves left
+blue square moves right twice
+green triangle moves up slowly
+blue square moves right slowly
+gold star moves down slowly
+blue square moves right
+blue square
+blue square moves right slowly
+green triangle moves up
+red circle moves left
+blue square moves right slowly
+gold star
+green triangle moves up
+gold star moves down slowly
+gold star moves down
+blue square
+gold star moves down
+blue square moves right slowly
+gold star moves down twice
+gold star moves down twice
+gold star moves down slowly
+blue square
+gold star moves down
+green triangle
+gold star moves down twice
+blue square moves right slowly
+gold star moves down
+blue square moves right
+red circle moves left twice
+gold star moves down
+green triangle
+green triangle
+green triangle
+red circle
+green triangle
+blue square
+blue square moves right
+green triangle moves up slowly
+gold star moves down
+green triangle moves up
+red circle
+gold star
+gold star moves down
+blue square moves right twice
+red circle moves left
+red circle
+blue square moves right
+blue square moves right twice
+gold star moves down
+gold star moves down twice
+blue square moves right
+blue square moves right twice
+gold star
+blue square
+gold star moves down slowly
+gold star moves down slowly
+red circle moves left
+red circle moves left
+blue square moves right
+blue square moves right
+green triangle moves up twice
+green triangle moves up
+gold star moves down
+green triangle moves up
+green triangle moves up
+green triangle
+blue square moves right slowly
+green triangle moves up slowly
+gold star moves down
+red circle
+red circle moves left twice
+gold star moves down twice
+gold star
+gold star
+blue square moves right twice
+blue square moves right twice
+green triangle moves up
+gold star moves down
+blue square moves right twice
+green triangle moves up slowly
+gold star moves down slowly
+gold star moves down
+blue square moves right twice
+blue square moves right
+blue square moves right slowly
+gold star moves down twice
+red circle moves left twice
+blue square
+green triangle moves up twice
+green triangle
+green triangle
+blue square moves right
+gold star moves down twice
+blue square moves right
+gold star
+red circle moves left twice
+blue square moves right
+green triangle moves up
+green triangle
+green triangle moves up
+green triangle moves up
+gold star
+blue square moves right
+red circle
+green triangle moves up
+blue square moves right twice
+blue square moves right slowly
+green triangle moves up
+green triangle
+red circle
+blue square
+green triangle moves up
+red circle moves left
+blue square
+gold star
+gold star
+blue square moves right
+blue square moves right slowly
+red circle
+gold star
+blue square
+green triangle
+blue square moves right
+green triangle moves up
+gold star moves down twice
+green triangle moves up
+green triangle moves up slowly
+green triangle
+blue square moves right
+gold star moves down slowly
+green triangle
+gold star
+red circle moves left
+blue square moves right
+gold star moves down
+gold star
+red circle moves left twice
+gold star moves down twice
+green triangle moves up slowly
+gold star moves down
+red circle moves left
+red circle moves left slowly
+blue square moves right twice
+gold star moves down
+green triangle moves up slowly
+blue square
+red circle moves left twice
+green triangle moves up twice
+red circle moves left
+green triangle
+blue square
+red circle
+green triangle
+gold star moves down twice
+green triangle moves up twice
+green triangle
+gold star moves down
+green triangle moves up
+blue square moves right twice
+green triangle
+red circle
+green triangle moves up twice
+green triangle moves up
+gold star moves down
+red circle moves left
+green triangle moves up twice
+green triangle
+gold star
+green triangle moves up twice
+blue square moves right slowly
+green triangle
+green triangle moves up
+green triangle moves up
+blue square
+red circle moves left
+blue square moves right twice
+red circle
+green triangle
+green triangle moves up
+red circle
+green triangle moves up twice
+blue square moves right slowly
+red circle moves left slowly
+red circle moves left twice
+green triangle moves up slowly
+gold star moves down slowly
+blue square moves right
+blue square moves right twice
+red circle
+green triangle moves up slowly
+red circle
+green triangle
+red circle moves left slowly
+red circle
+blue square moves right slowly
+red circle moves left slowly
+red circle
+red circle moves left twice
+gold star
+red circle moves left twice
+green triangle moves up twice
+green triangle moves up
+green triangle moves up
+red circle
+gold star moves down slowly
+gold star
+gold star
+red circle
+gold star moves down twice
+green triangle moves up
+green triangle
+blue square
+green triangle moves up slowly
+gold star moves down
+gold star
+green triangle
+gold star
+red circle moves left
+gold star moves down
+green triangle
+blue square moves right slowly
+red circle moves left
+green triangle moves up
+green triangle moves up slowly
+red circle moves left
+green triangle moves up twice
+green triangle moves up
+red circle moves left twice
+red circle moves left
+gold star moves down twice
+blue square moves right
+red circle moves left twice
+gold star
+gold star
+blue square
+red circle
+blue square
+blue square
+blue square moves right twice
+red circle moves left
+blue square moves right slowly
+blue square moves right slowly
+green triangle
+blue square moves right slowly
+green triangle moves up twice
+blue square
+gold star moves down
+gold star moves down twice
+blue square moves right
+red circle
+red circle moves left twice
+blue square moves right twice
+red circle moves left
+green triangle moves up slowly
+red circle
+red circle moves left
+red circle moves left twice
+gold star moves down slowly
+green triangle moves up twice
+gold star moves down
+red circle
+blue square
+red circle
+green triangle
+green triangle

lecture_6/examples/categorical.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Run the complete categorical training and sampling example."""
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from run import main
+if __name__ == '__main__':
+    main(['--method', 'categorical', *sys.argv[1:]])

lecture_6/examples/consistency.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Run the complete consistency training and sampling example."""
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from run import main
+if __name__ == '__main__':
+    main(['--method', 'consistency', *sys.argv[1:]])

lecture_6/examples/diamond.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Run the complete diamond training and sampling example."""
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from run import main
+if __name__ == '__main__':
+    main(['--method', 'diamond', *sys.argv[1:]])

lecture_6/examples/discrete_esd.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Run the complete discrete-esd training and sampling example."""
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from run import main
+if __name__ == '__main__':
+    main(['--method', 'discrete-esd', *sys.argv[1:]])

lecture_6/examples/discrete_lsd.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Run the complete discrete-lsd training and sampling example."""
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from run import main
+if __name__ == '__main__':
+    main(['--method', 'discrete-lsd', *sys.argv[1:]])

lecture_6/examples/expanding.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Run the complete expanding training and sampling example."""
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from run import main
+if __name__ == '__main__':
+    main(['--method', 'expanding', *sys.argv[1:]])

lecture_6/examples/flow_matching.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Run the complete flow-matching training and sampling example."""
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from run import main
+if __name__ == '__main__':
+    main(['--method', 'flow-matching', *sys.argv[1:]])

lecture_6/examples/fmlm.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Run the complete fmlm training and sampling example."""
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from run import main
+if __name__ == '__main__':
+    main(['--method', 'fmlm', *sys.argv[1:]])

lecture_6/examples/fmm_eulerian.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Run the complete fmm-eulerian training and sampling example."""
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from run import main
+if __name__ == '__main__':
+    main(['--method', 'fmm-eulerian', *sys.argv[1:]])

lecture_6/examples/fmm_lagrangian.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Run the complete fmm-lagrangian training and sampling example."""
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from run import main
+if __name__ == '__main__':
+    main(['--method', 'fmm-lagrangian', *sys.argv[1:]])

lecture_6/examples/latent.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Run the complete latent training and sampling example."""
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from run import main
+if __name__ == '__main__':
+    main(['--method', 'latent', *sys.argv[1:]])

lecture_6/examples/meanflow.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Run the complete meanflow training and sampling example."""
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from run import main
+if __name__ == '__main__':
+    main(['--method', 'meanflow', *sys.argv[1:]])

lecture_6/examples/meta.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Run the complete meta training and sampling example."""
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from run import main
+if __name__ == '__main__':
+    main(['--method', 'meta', *sys.argv[1:]])

lecture_6/examples/self_distill.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Run the complete self-distill training and sampling example."""
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from run import main
+if __name__ == '__main__':
+    main(['--method', 'self-distill', *sys.argv[1:]])

lecture_6/examples/shortcut.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Run the complete shortcut training and sampling example."""
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from run import main
+if __name__ == '__main__':
+    main(['--method', 'shortcut', *sys.argv[1:]])

lecture_6/examples/ssfm.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Run the complete ssfm training and sampling example."""
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from run import main
+if __name__ == '__main__':
+    main(['--method', 'ssfm', *sys.argv[1:]])

lecture_6/expanding.py

@@ -0,0 +1,192 @@
+"""Expanding Flow Maps on variable-length text with learned gap insertions.
+
+Based on paper Algorithms 2-4. The state carries birth times alongside token
+coordinates. Finite transport uses local clocks on a shared destination canvas.
+The insertion head learns both remaining-count and interval-count expectations.
+"""
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from common import ema_copy, mlp, optimize
+
+
+class ExpandingNet(nn.Module):
+    def __init__(self, length, vocab, width):
+        super().__init__()
+        self.length, self.vocab = length, vocab
+        # Every token contributes state, two local clocks, birth, and active mask.
+        self.denoiser = mlp(length * (vocab+4), length*vocab, width)
+        self.insertion = mlp(length*(vocab+2)+2, length+1, width)
+
+    def predict(self, x, local_s, local_t, births, mask):
+        features = torch.cat([x, local_s[..., None], local_t[..., None],
+                              births[..., None], mask[..., None].to(x)], -1)
+        return self.denoiser(features.flatten(1)).reshape(-1, self.length, self.vocab)
+
+    def counts(self, x, births, mask, s, t, diagonal=False):
+        features = torch.cat([x, births[..., None], mask[..., None].to(x)], -1).flatten(1)
+        features = torch.cat([features, s, t], -1)
+        remaining = F.softplus(self.insertion(features)) + 1e-5
+        # Linear birth CDF: rho=(t-s)/(1-s). The diagonal predicts the
+        # remaining count; off-diagonal predictions include the interval factor.
+        return remaining if diagonal else (t-s)/(1-s)*remaining
+
+
+def local_clock(t, births):
+    return ((t-births)/(1-births)).clamp(0, 1)
+
+
+def local_map(model, x, a, b, births, mask):
+    p = model.predict(x, a, b, births, mask).softmax(-1)
+    eta = ((b-a)/(1-a).clamp_min(1e-6))[..., None]
+    y = (x + eta*(p-x)) * mask[..., None]
+    return y, p
+
+
+def gap_counts(active_indices, born_indices, length):
+    """Assign each missing index to a gap while retaining original token order."""
+    result = torch.zeros(len(active_indices)+1, device=active_indices.device)
+    if len(born_indices):
+        gap = torch.searchsorted(active_indices, born_indices)
+        result.scatter_add_(0, gap, torch.ones_like(gap, dtype=result.dtype))
+    return result
+
+
+def compact(full, births, selected, max_length):
+    x = full.new_zeros(max_length, full.shape[-1])
+    b = births.new_zeros(max_length)
+    mask = torch.zeros(max_length, dtype=torch.bool, device=full.device)
+    n = len(selected)
+    x[:n], b[:n], mask[:n] = full[selected], births[selected], True
+    return x, b, mask
+
+
+def count_divergence(target, mean):
+    # Poisson NLL up to target-only constants. Finite at target=0.
+    # Avoid xlogy(0,0)'s undefined intermediate derivative in autodiff.
+    return mean - target + target * (target.clamp_min(1e-12).log()-mean.log())
+
+
+def expansion_batch(ids, lengths, vocab, batch_size):
+    index = torch.randint(len(ids), (batch_size,), device=ids.device)
+    fields = {k: [] for k in ['xs','bs','ms','xe','be','me','ys','ye','missing','interval','gapmask']}
+    pairs = (.96*torch.rand(batch_size, 2, device=ids.device)).sort(-1).values
+    s, t = pairs[:, :1], pairs[:, 1:]
+    t = torch.maximum(t, s+1e-4)
+    L = ids.shape[1]
+    for j, row in enumerate(index):
+        n = int(lengths[row])
+        clean = F.one_hot(ids[row, :n], vocab).float()
+        births = .999*torch.rand(n, device=ids.device)
+        noise = torch.randn_like(clean)
+        clock = local_clock(s[j], births)
+        full = (1-clock[:, None])*noise+clock[:, None]*clean
+        active = torch.where(births <= s[j])[0]
+        later = torch.where(births <= t[j])[0]
+        xs, bs, ms = compact(full, births, active, L)
+        xe, be, me = compact(full, births, later, L)
+        ys, _, _ = compact(clean, births, active, L)
+        ye, _, _ = compact(clean, births, later, L)
+        missing = full.new_zeros(L+1)
+        interval = full.new_zeros(L+1)
+        gm = torch.arange(L+1, device=ids.device) <= len(active)
+        missing[gm] = gap_counts(active, torch.where(births>s[j])[0], n)
+        interval[gm] = gap_counts(active, torch.where((births>s[j]) & (births<=t[j]))[0], n)
+        for k, value in zip(fields, [xs,bs,ms,xe,be,me,ys,ye,missing,interval,gm]):
+            fields[k].append(value)
+    return {k: torch.stack(v) for k,v in fields.items()}, s, t
+
+
+def train_expanding(ids, lengths, vocab, args):
+    model = ExpandingNet(ids.shape[1], len(vocab), args.width).to(ids.device)
+    teacher = ema_copy(model)
+    def objective(step):
+        b, s, t = expansion_batch(ids, lengths, len(vocab), args.batch_size)
+        a = local_clock(s, b['bs'])
+        logits = model.predict(b['xs'], a, a, b['bs'], b['ms'])
+        ce = -(b['ys']*logits.log_softmax(-1)).sum(-1)
+        diagonal = (ce*b['ms']).sum()/b['ms'].sum().clamp_min(1)
+        # Lift both paths to the same destination canvas, with the same noise.
+        la, lb = local_clock(s,b['be']), local_clock(t,b['be'])
+        lu = local_clock((s+t)/2,b['be'])
+        with torch.no_grad():
+            middle, p1 = local_map(teacher,b['xe'],la,lu,b['be'],b['me'])
+            _, p2 = local_map(teacher,middle,lu,lb,b['be'],b['me'])
+            gamma = ((1-lb)*(lu-la)/((1-lu)*(lb-la)).clamp_min(1e-6)).clamp(0,1)
+            target = gamma[...,None]*p1+(1-gamma[...,None])*p2
+        logits = model.predict(b['xe'],la,lb,b['be'],b['me'])
+        kl = F.kl_div(logits.log_softmax(-1),target,reduction='none').sum(-1)
+        finite = (kl*b['me']).sum()/b['me'].sum().clamp_min(1)
+        m = model.counts(b['xs'],b['bs'],b['ms'],s,s,diagonal=True)
+        interval = model.counts(b['xs'],b['bs'],b['ms'],s,t)
+        count = count_divergence(b['missing'],m)+count_divergence(b['interval'],interval)
+        insertion = (count*b['gapmask']).sum()/b['gapmask'].sum()
+        ramp = min(1.,(step+1)/max(1,args.train_steps//4))
+        return diagonal+ramp*finite+insertion, {'diagonal_ce':diagonal,'finite':finite,'insertion':insertion}
+    logs = optimize(model,objective,args.train_steps,args.lr,teacher)
+    return model, {'model':model.state_dict(),'length':ids.shape[1],'vocab':vocab}, logs
+
+
+def bounded_counts(means, budget):
+    """Independent binomial proposals, then left-to-right joint budget capping.
+
+    Before capping, each proposal has its predicted mean after [0,budget]
+    clipping. Conditional means alone do not determine the joint count law.
+    """
+    if budget == 0:
+        return torch.zeros_like(means, dtype=torch.long), 0
+    raw = torch.distributions.Binomial(budget, probs=means.clamp(0,budget)/budget).sample().long()
+    kept = raw.clone()
+    remaining = budget
+    for i in range(len(kept)):
+        kept[i] = min(int(kept[i]),remaining)
+        remaining -= int(kept[i])
+    return kept, int((raw-kept).sum())
+
+
+def insert_tokens(state, births, counts, noise, birth_time):
+    pieces, times, offset = [], [], 0
+    for gap, count in enumerate(counts.tolist()):
+        if count:
+            pieces.append(noise[offset:offset+count])
+            times.append(births.new_full((count,),birth_time))
+            offset += count
+        if gap < len(state):
+            pieces.append(state[gap:gap+1]); times.append(births[gap:gap+1])
+    if offset != len(noise):
+        raise ValueError('Noise rows must match insertion counts.')
+    if not pieces:
+        return state, births
+    return torch.cat(pieces),torch.cat(times)
+
+
+@torch.no_grad()
+def sample_expanding(model,count,steps,device):
+    result = torch.zeros(count,model.length,dtype=torch.long,device=device)
+    lengths = torch.zeros(count,dtype=torch.long,device=device)
+    capped = 0
+    traces = []
+    for j in range(count):
+        x = torch.zeros(0,model.vocab,device=device)
+        births = torch.zeros(0,device=device)
+        trace = [0]
+        for k in range(steps):
+            s,t = k/steps,(k+1)/steps
+            padded,bt,mask = compact(x,births,torch.arange(len(x),device=device),model.length)
+            start,end = torch.tensor([[s]],device=device),torch.tensor([[t]],device=device)
+            means = model.counts(padded[None],bt[None],mask[None],start,end)[0,:len(x)+1]
+            counts,discarded = bounded_counts(means,model.length-len(x));capped+=discarded
+            noise = torch.randn(int(counts.sum()),model.vocab,device=device)
+            # Finite expand-then-transport convention: new coordinates enter at s.
+            x,births = insert_tokens(x,births,counts,noise,s)
+            if len(x):
+                padded,bt,mask=compact(x,births,torch.arange(len(x),device=device),model.length)
+                y,_=local_map(model,padded[None],local_clock(start,bt[None]),
+                              local_clock(end,bt[None]),bt[None],mask[None])
+                x=y[0,:len(x)]
+            trace.append(len(x))
+        lengths[j]=len(x)
+        if len(x):result[j,:len(x)]=x.argmax(-1)
+        traces.append(trace)
+    return result,lengths,{'capped_proposed_insertions':capped,'length_trajectories':traces}

lecture_6/lecture_core.py

@@ -0,0 +1,14 @@
+"""Stable imports connecting the lecture's code walkthroughs to full methods.
+
+Implementation lives in the named modules so training and sampling use the
+same functions demonstrated on the slides. See SLIDE_CODE_MAP.md.
+"""
+from common import finite_map
+from continuous import (diagonal_loss,lagrangian_residual,eulerian_residual,
+                        semigroup_loss,meanflow_loss,shortcut_loss)
+from categorical import (categorical_map,composition_target,corrected_logit_teacher,
+                         probability_kl,DecodingClock)
+from posterior import (glass_denoiser,posterior_samples,posterior_value,
+                       fine_tune_surrogate,weighted_diamond_samples)
+from expanding import (local_clock,local_map,count_divergence,bounded_counts,insert_tokens)
+from stochastic import chen_two,sample_coefficients

lecture_6/numerical_examples.py

@@ -0,0 +1,187 @@
+"""CIS 6270 Lecture 6. Small, reproducible flow-map teaching experiments.
+
+Run: python numerical_examples.py --output outputs/numerical
+Dependencies: Python 3.11+, numpy, torch.
+The toy experiments validate the lecture mathematics. They do not reproduce
+the large-scale training or benchmark claims of the cited papers.
+"""
+import argparse, json, math, pathlib
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+def flow(model, x, s, t):
+    """Residual map with an exactly identity diagonal."""
+    return x + (t-s)*model(torch.cat([x,s,t],dim=-1))
+
+def train_scalar(output, iterations=6000):
+    """Self-distill the flow of dx/dt=x, using diagonal and composition only."""
+    torch.manual_seed(6270)
+    torch.set_num_threads(2)
+    model=nn.Sequential(nn.Linear(3,64),nn.SiLU(),nn.Linear(64,64),nn.SiLU(),nn.Linear(64,1))
+    opt=torch.optim.Adam(model.parameters(),lr=1e-3)
+    history=[]
+    for step in range(iterations):
+        # Broad state coverage includes states reached by the split map.
+        x=6*torch.rand(256,1)-3
+        times=torch.rand(256,2).sort(dim=-1).values
+        s,t=times[:,:1],times[:,1:]
+        u=s+(t-s)*torch.rand_like(s)
+        diagonal=model(torch.cat([x,s,s],dim=-1))
+        diag_loss=(diagonal-x).square().mean()
+        with torch.no_grad():
+            target=flow(model,flow(model,x,s,u),u,t)
+        prediction=flow(model,x,s,t)
+        # Normalize the interval residual to prevent tiny intervals dominating
+        # the count of nearly zero-error examples. Keep a finite floor.
+        cons_loss=((prediction-target)/(t-s).clamp_min(0.1)).square().mean()
+        loss=diag_loss + (0 if step<500 else 1)*cons_loss
+        opt.zero_grad();loss.backward();opt.step()
+        if step%500==0:history.append({'step':step,'diagonal':diag_loss.item(),'composition':cons_loss.item()})
+    with torch.no_grad():
+        x=torch.linspace(-1.5,1.5,501)[:,None]
+        s=torch.zeros_like(x);t=torch.ones_like(x);u=0.5*t
+        pred=flow(model,x,s,t);truth=x*math.e
+        split=flow(model,flow(model,x,s,u),u,t)
+        rmse=(pred-truth).square().mean().sqrt().item()
+        composition_rmse=(pred-split).square().mean().sqrt().item()
+        np.savez(output/'scalar_predictions.npz',x=x.numpy().ravel(),pred=pred.numpy().ravel(),truth=truth.numpy().ravel())
+    torch.save(model.state_dict(),output/'scalar_map_weights.pt')
+    return {'endpoint_rmse':rmse,'composition_rmse':composition_rmse,'iterations':iterations,'history':history}
+
+def categorical_map(net,x,s,t):
+    psi=net(x,s,t).softmax(dim=-1)
+    h=(t-s)/(1-s)
+    return (1-h)*x+h*psi,psi
+
+def check_categorical():
+    x=torch.tensor([-.2,.6,1.1],dtype=torch.float64)
+    psi=torch.tensor([.1,.7,.2],dtype=torch.float64)
+    y=(1/3)*x+(2/3)*psi
+    assert torch.allclose(y,torch.tensor([0,2/3,.5],dtype=torch.float64))
+    target=(1/3)*torch.tensor([.8,.2])+(2/3)*torch.tensor([.2,.8])
+    assert torch.allclose(target,torch.tensor([.4,.6]))
+    logits=torch.tensor([.2,-.1],requires_grad=True)
+    loss=F.kl_div(logits.log_softmax(-1),target,reduction='sum')
+    loss.backward()
+    assert torch.allclose(logits.grad,logits.softmax(-1)-target,atol=1e-7)
+    return {'mapped_state':y.tolist(),'state_sum':y.sum().item(),'target':target.tolist()}
+
+def check_meanflow_jvp():
+    # Exact backward average for dz/dt=z. Stay off diagonal for this check.
+    def exact_average(z,r,t):
+        h=t-r
+        return z*(-torch.expm1(-h))/h
+    z=torch.tensor([[1.7]],dtype=torch.float64)
+    r=torch.tensor([[.2]],dtype=torch.float64)
+    t=torch.tensor([[.8]],dtype=torch.float64)
+    value,derivative=torch.func.jvp(exact_average,(z,r,t),(z,torch.zeros_like(r),torch.ones_like(t)))
+    target=z-(t-r)*derivative
+    assert torch.allclose(value,target,atol=1e-10)
+    return {'average':value.item(),'identity_residual':(value-target).abs().item()}
+
+def posterior_value_demo():
+    # Prior Z~N(0,1), observation x=beta*Z+alpha*eps.
+    # Reward r(Z)=c*Z gives an analytic log moment-generating function.
+    torch.manual_seed(13)
+    alpha,beta,c=.7,.6,.4
+    x=torch.tensor(.3,dtype=torch.float64,requires_grad=True)
+    gain=beta/(alpha**2+beta**2)
+    variance=alpha**2/(alpha**2+beta**2)
+    eps=torch.randn(100000,dtype=torch.float64)
+    z=gain*x+math.sqrt(variance)*eps
+    logw=c*z
+    estimate=torch.logsumexp(logw,0)-math.log(len(eps))
+    gradient=torch.autograd.grad(estimate,x)[0]
+    exact=c*gain*x.detach()+.5*c*c*variance
+    exact_gradient=c*gain
+    assert abs(gradient.item()-exact_gradient)<1e-10
+    assert abs(estimate.item()-exact.item())<.006
+    return {'estimated_value':estimate.item(),'exact_value':exact.item(),'gradient':gradient.item(),'exact_gradient':exact_gradient}
+
+def sample_gap_counts(means,remaining_budget):
+    """Paper-style bounded proposals followed by left-to-right budget capping.
+
+    These proposals match the per-gap means before joint truncation. They are
+    not asserted to identify the full conditional count law from means alone.
+    """
+    counts=[];remaining=int(remaining_budget)
+    for mean in means:
+        if remaining_budget==0:count=0
+        else:
+            prob=float(torch.as_tensor(mean).clamp(0,remaining_budget))/remaining_budget
+            count=int(torch.distributions.Binomial(remaining_budget,probs=prob).sample())
+        count=min(count,remaining);counts.append(count);remaining-=count
+    return torch.tensor(counts,dtype=torch.long)
+
+def insert_by_gap(state,old_birth_times,counts,new_noise,birth_time):
+    """Insert ordered noise rows into the n+1 gaps and preserve clock alignment."""
+    assert len(counts)==len(state)+1
+    pieces=[];clocks=[];offset=0
+    for gap,count in enumerate(counts.tolist()):
+        if count:
+            pieces.append(new_noise[offset:offset+count]);offset+=count
+            clocks.append(torch.full((count,),float(birth_time)))
+        if gap<len(state):
+            pieces.append(state[gap:gap+1]);clocks.append(old_birth_times[gap:gap+1])
+    assert offset==len(new_noise)
+    return torch.cat(pieces),torch.cat(clocks)
+
+def check_expansion():
+    state=torch.tensor([[1.,0.],[0.,1.]])
+    births=torch.tensor([0.,.2]);counts=torch.tensor([1,0,1])
+    noise=torch.tensor([[-.2,.4],[.3,-.1]])
+    expanded,bt=insert_by_gap(state,births,counts,noise,.5)
+    local=(.75-bt)/(1-bt)
+    assert expanded.shape==(4,2)
+    assert torch.allclose(bt,torch.tensor([.5,0.,.2,.5]))
+    assert torch.allclose(local,torch.tensor([.5,.75,.6875,.5]))
+    return {'birth_times':bt.tolist(),'local_times':local.tolist(),'expected_insertions':((.75-.25)/(1-.25))*3}
+
+def chen_two(left,right,hL,hR):
+    h=hL+hR
+    return torch.stack([left[...,0]+right[...,0],
+        (hL*left[...,1]+hR*right[...,1]-hR*left[...,0]+hL*right[...,0])/h],dim=-1)
+
+def check_brownian():
+    left=torch.tensor([.2,.04]);right=torch.tensor([-.1,-.02])
+    coarse=chen_two(left,right,.5,.5)
+    assert torch.allclose(coarse,torch.tensor([.1,-.14]))
+    direct=1+.5+.8*coarse[0]
+    split=(1+.25+.8*left[0])+.25+.8*right[0]
+    assert torch.allclose(direct,split)
+    torch.manual_seed(17)
+    scale=torch.tensor([.5,1/6]).sqrt()
+    L=torch.randn(200000,2)*scale;R=torch.randn(200000,2)*scale
+    C=chen_two(L,R,.5,.5)
+    covariance=torch.cov(C.T)
+    assert torch.allclose(covariance,torch.diag(torch.tensor([1.,1/3])),atol=.012)
+    # Exact polynomial restriction identities, tested over unequal intervals.
+    hL,hR=.3,.7
+    q=torch.linspace(0,1,100,dtype=torch.float64)
+    global_left=2*(hL*q)/(hL+hR)-1
+    local_left=(hL/(hL+hR))*(2*q-1)-hR/(hL+hR)
+    assert torch.allclose(global_left,local_left,atol=1e-12)
+    return {'coarse_coefficients':coarse.tolist(),'same_noise_endpoint':direct.item(),'empirical_covariance':covariance.tolist()}
+
+def check_meta_gradient():
+    d=torch.tensor(.3,requires_grad=True)
+    w=torch.tensor([1.,4.]);grad_w=torch.tensor([.2,.8]);a=.5
+    residual=d+(w-1)*d.detach()-a*grad_w
+    residual.square().mean().backward()
+    expected=2*(w*d.detach()-a*grad_w).mean()
+    assert torch.allclose(d.grad,expected)
+    return {'surrogate_gradient':d.grad.item(),'estimating_equation_gradient':expected.item()}
+
+def main():
+    parser=argparse.ArgumentParser();parser.add_argument('--output',default='results');parser.add_argument('--iterations',type=int,default=6000)
+    args=parser.parse_args();out=pathlib.Path(args.output);out.mkdir(parents=True,exist_ok=True)
+    results={'torch_version':torch.__version__,'categorical':check_categorical(),'meanflow_jvp':check_meanflow_jvp(),
+        'posterior_value':posterior_value_demo(),'expansion':check_expansion(),'brownian':check_brownian(),'meta_gradient':check_meta_gradient()}
+    results['scalar_training']=train_scalar(out,args.iterations)
+    (out/'checks.json').write_text(json.dumps(results,indent=2))
+    print(json.dumps({k:v for k,v in results.items() if k!='scalar_training'},indent=2))
+    print('Scalar endpoint RMSE',results['scalar_training']['endpoint_rmse'])
+    print('Scalar composition RMSE',results['scalar_training']['composition_rmse'])
+if __name__=='__main__':main()

lecture_6/posterior.py

@@ -0,0 +1,179 @@
+"""Posterior Diamond Maps, data-trained Meta Flow Maps, and reward steering.
+
+There are two independent noise draws and a shared clean endpoint. The inner
+time evolves while the outer observation and outer time remain fixed.
+"""
+import math
+
+import torch
+from torch.nn import functional as F
+
+from common import (MapNet, draw_batch, ema_copy, exact_denoiser, exact_velocity,
+                    finite_map, interpolate, mixture_posterior, optimize,
+                    ordered_times, time_like)
+from continuous import semigroup_loss
+
+
+def glass_denoiser(inner, s, outer, t):
+    """Fuse independent Gaussian observations by adding their precisions.
+
+    inner=s*X1+(1-s)*eps, outer=t*X1+(1-t)*eps'. The sufficient
+    statistic S has noise variance 1/precision. Convert it to the original
+    linear schedule t*=sqrt(precision)/(1+sqrt(precision)), then denoise.
+    """
+    precision = (s / (1-s)).square() + (t / (1-t)).square()
+    safe = precision.clamp_min(1e-12)
+    statistic = (s * inner / (1-s).square() + t * outer / (1-t).square()) / safe
+    ratio = safe.sqrt()
+    effective_time = ratio / (1 + ratio)
+    value = exact_denoiser(effective_time * statistic, effective_time)
+    return torch.where(precision > 1e-12, value, torch.zeros_like(value))
+
+
+def glass_velocity(inner, s, outer, t):
+    return (glass_denoiser(inner, s, outer, t) - inner) / (1-s)
+
+
+def posterior_context(outer, t):
+    return torch.cat([outer, time_like(t, outer)], -1)
+
+
+def train_posterior(method, data, args):
+    model = MapNet(2, args.width, context_dim=3).to(data)
+    teacher = ema_copy(model)
+
+    def objective(step):
+        clean = draw_batch(data, args.batch_size)
+        outer_t = .94 * torch.rand(len(clean), 1, device=data.device)
+        outer, _ = interpolate(clean, outer_t)
+        context = posterior_context(outer, outer_t)
+        s, t, _ = ordered_times(clean, .97)
+        # Independent inner noise; reusing the outer noise changes the posterior.
+        inner, displacement = interpolate(clean, s)
+        if method == 'diamond':
+            diagonal_target = glass_velocity(inner, s, outer, outer_t).detach()
+        else:
+            diagonal_target = displacement
+        diagonal = F.mse_loss(model(inner, s, s, context), diagonal_target)
+        if method == 'diamond':
+            y, derivative = torch.func.jvp(lambda end: finite_map(model, inner, s, end, context),
+                                          (t,), (torch.ones_like(t),))
+            target = glass_velocity(y.detach(), t, outer, outer_t).detach()
+            finite = F.mse_loss(derivative, target)
+        else:
+            finite = semigroup_loss(model, teacher, inner, s, t, context)
+        weight = min(1., (step+1) / max(1, args.train_steps//4))
+        return diagonal + weight * finite, {'diagonal': diagonal, 'finite': finite}
+
+    logs = optimize(model, objective, args.train_steps, args.lr, teacher)
+    return model, {'model': model.state_dict(), 'dim': 2}, logs
+
+
+def posterior_samples(model, outer, t, count, steps=1, noise=None):
+    context = posterior_context(outer, t)
+    context = context[:, None].expand(-1, count, -1).reshape(-1, 3)
+    x = torch.randn(len(outer)*count, 2, device=outer.device) if noise is None else noise.reshape(-1, 2)
+    for i in range(steps):
+        x = finite_map(model, x, i/steps, (i+1)/steps, context)
+    return x.reshape(len(outer), count, 2)
+
+
+def reward(x, strength=1.):
+    """Bounded differentiable preference for positive first coordinates."""
+    return strength * torch.tanh(x[..., 0])
+
+
+def posterior_value(model, x, t, particles=32, strength=1., steps=1):
+    samples = posterior_samples(model, x, t, particles, steps)
+    log_weights = reward(samples, strength)
+    value = torch.logsumexp(log_weights, -1) - math.log(particles)
+    weighted_mean = (log_weights.softmax(-1)[..., None] * samples).sum(1)
+    return value, weighted_mean, samples
+
+
+def fine_tune_surrogate(delta, weights, grad_weights, coefficient):
+    """Equation 43 of Meta Flow Maps. Both Monte Carlo terms are detached.
+
+    d/d(delta) E[loss] = 2 E[w*delta - coefficient*grad(w)].
+    Differentiating (w*delta - coefficient*grad(w))^2 instead adds a wrong w.
+    """
+    residual = delta + (weights.detach()-1) * delta.detach() - coefficient * grad_weights.detach()
+    return residual.square().mean()
+
+
+def train_reward_drift(posterior, data, args):
+    posterior.requires_grad_(False)
+    drift = MapNet(2, args.width).to(data)
+    def objective(_):
+        clean = draw_batch(data, args.batch_size)
+        t = .05 + .8 * torch.rand(len(clean), 1, device=data.device)
+        x, _ = interpolate(clean, t)
+        x.requires_grad_(True)
+        endpoint = posterior_samples(posterior, x, t, 1, args.posterior_steps)[:, 0]
+        w = reward(endpoint, args.reward_strength).exp().unsqueeze(-1)
+        grad_w = torch.autograd.grad(w.sum(), x)[0]
+        baseline = exact_velocity(x, t).detach()
+        delta = drift(x.detach(), t, t) - baseline
+        # Linear interpolant's compatible probability-flow correction g^2/2=(1-t)/t.
+        # Use t>=.05 to avoid singular coefficients in this demonstration.
+        coefficient = (1-t) / t
+        loss = fine_tune_surrogate(delta, w, grad_w, coefficient)
+        return loss, {'mean_weight': w.mean()}
+    logs = optimize(drift, objective, args.finetune_steps, args.lr)
+    return drift, logs
+
+
+@torch.no_grad()
+def guided_samples(model, count, steps, particles, strength, posterior_steps=1):
+    """Derivative-free reward-weighted posterior mean in the linear flow drift.
+
+    v*(x,t)=(E_reward[X1|x,t]-x)/(1-t). It equals the compatible
+    probability-flow correction for an exact posterior. Finite learned maps
+    and self-normalized Monte Carlo introduce approximation and ratio bias.
+    """
+    x = torch.randn(count, 2, device=next(model.parameters()).device)
+    h = 1/steps
+    for i in range(steps):
+        t = torch.full((count, 1), i/steps, device=x.device)
+        _, mean, _ = posterior_value(model, x, t, particles, strength, posterior_steps)
+        x = x + h * (mean-x) / (1-t)
+    return x
+
+
+def weighted_diamond_samples(outer, t, particles=128):
+    """Exact tractable demonstration of weighted Diamond proposal correction.
+
+    Proposal q=N(0,4I) has evaluable density and full support. Reweight with
+    p1(z)*p_t(outer|z)/q(z), then resample. This supplies an independent
+    check of posterior recovery without assuming an arbitrary learned map's
+    Jacobian density is known. Finite self-normalized importance sampling is biased.
+    """
+    from common import CENTERS, DATA_STD
+    t = time_like(t, outer)
+    z = 2 * torch.randn(len(outer), particles, 2, device=outer.device)
+    delta = z[:, :, None] - CENTERS.to(z)
+    lp = torch.logsumexp(-delta.square().sum(-1)/(2*DATA_STD**2), -1)
+    lp -= math.log(4 * 2 * math.pi * DATA_STD**2)
+    noise = (outer[:, None]-t[:, None]*z)/(1-t[:, None])
+    likelihood = -.5*noise.square().sum(-1) - 2*(1-t).log()
+    lq = -z.square().sum(-1)/8 - math.log(8*math.pi)
+    weights = (lp + likelihood - lq).softmax(-1)
+    mean = (weights[..., None]*z).sum(1)
+    indices = torch.multinomial(weights, 1)
+    selected = z.gather(1, indices[..., None].expand(-1, -1, 2))[:, 0]
+    return selected, mean, 1/weights.square().sum(-1)
+
+
+@torch.no_grad()
+def posterior_diagnostics(model, steps=1):
+    device = next(model.parameters()).device
+    x = torch.tensor([[0., 0.], [1., -.5], [-1., 1.]], device=device)
+    t = torch.tensor([[.15], [.5], [.8]], device=device)
+    p, means, var = mixture_posterior(x, t)
+    exact_mean = (p[..., None]*means).sum(1)
+    exact_variance = (p[..., None]*(var[:, None]+means.square())).sum(1)-exact_mean.square()
+    samples = posterior_samples(model, x, t, 512, steps)
+    mean, variance = samples.mean(1), samples.var(1)
+    return {'posterior_mean_rmse': float((mean-exact_mean).square().mean().sqrt()),
+            'posterior_variance_rmse': float((variance-exact_variance).square().mean().sqrt()),
+            'posterior_reference_mean': exact_mean.tolist(), 'posterior_estimated_mean': mean.tolist()}

lecture_6/requirements.txt

@@ -0,0 +1,3 @@
+# Match the course's PyTorch baseline; torch.func.jvp is required.
+torch==2.9.1
+numpy>=1.24,<3

lecture_6/run.py

@@ -0,0 +1,245 @@
+#!/usr/bin/env python3
+"""Train, save, reload, and sample the Lecture 6 flow-map methods."""
+import argparse
+import json
+import platform
+import time
+from pathlib import Path
+
+import torch
+from torch.nn import functional as F
+
+from common import (MapNet, SequenceNet, interpolate, load_text, mixture_data,
+                    mixture_metrics, seed_all, text_metrics, write_json)
+from continuous import (Autoencoder, embed_surface, integrate_velocity, sample_continuous,
+                        train_continuous, train_latent)
+from categorical import sample_categorical, train_categorical
+from posterior import (guided_samples, posterior_diagnostics, posterior_samples,
+                       posterior_value, reward, train_posterior, train_reward_drift,
+                       weighted_diamond_samples)
+from expanding import ExpandingNet, sample_expanding, train_expanding
+from stochastic import StrongMap, evaluate_stochastic, train_stochastic
+
+ROOT = Path(__file__).resolve().parent
+CONTINUOUS = ['flow-matching','fmm-lagrangian','fmm-eulerian','self-distill',
+              'consistency','shortcut','meanflow','latent']
+CATEGORICAL = ['fmlm','categorical','discrete-lsd','discrete-esd']
+METHODS = CONTINUOUS+CATEGORICAL+['diamond','meta','expanding','ssfm']
+
+
+def parser():
+    p=argparse.ArgumentParser(description=__doc__)
+    p.add_argument('--method',choices=METHODS)
+    p.add_argument('--mode',choices=['train-sample','train','sample'],default='train-sample')
+    p.add_argument('--train-steps',type=int,default=1000)
+    p.add_argument('--teacher-steps',type=int,default=1000)
+    p.add_argument('--finetune-steps',type=int,default=0,help='Optional Meta reward-drift fine-tuning')
+    p.add_argument('--sample-steps',type=int,default=8)
+    p.add_argument('--posterior-steps',type=int,default=4)
+    p.add_argument('--particles',type=int,default=32)
+    p.add_argument('--reward-strength',type=float,default=1.)
+    p.add_argument('--ssfm-target',choices=['official-code','paper'],default='official-code')
+    p.add_argument('--batch-size',type=int,default=64)
+    p.add_argument('--samples',type=int,default=128)
+    p.add_argument('--width',type=int,default=64)
+    p.add_argument('--lr',type=float,default=1e-3)
+    p.add_argument('--seed',type=int,default=6270)
+    p.add_argument('--threads',type=int,default=1)
+    p.add_argument('--device',choices=['cpu','cuda'],default='cpu')
+    p.add_argument('--data',help='Whitespace-separated text, one sequence per line; or continuous CSV with two numeric columns')
+    p.add_argument('--out',help='Run directory; defaults to lecture_6/outputs/METHOD')
+    return p
+
+
+def restore_model(checkpoint,device):
+    method=checkpoint['method'];width=checkpoint['width']
+    if method in CATEGORICAL:
+        model=SequenceNet(checkpoint['length'],len(checkpoint['vocab']),width)
+    elif method=='expanding':
+        model=ExpandingNet(checkpoint['length'],len(checkpoint['vocab']),width)
+    elif method=='ssfm':
+        model=StrongMap(width,checkpoint['sigma'])
+    else:
+        model=MapNet(checkpoint['dim'],width,3 if method in ['diamond','meta'] else 0)
+    model.load_state_dict(checkpoint['model'])
+    model=model.to(device).eval()
+    return model
+
+
+def generate(model,state,args):
+    """All metrics use fresh, seeded samples; no training examples stand in for output."""
+    seed_all(args.seed+100,args.threads)
+    method=state['method'];device=args.device
+    report={'method':method,'sample_steps':args.sample_steps,'samples':args.samples,
+            'sampling_seed':args.seed+100}
+    if method in CATEGORICAL:
+        ids=sample_categorical(model,args.samples,args.sample_steps,device).cpu()
+        texts,metrics=text_metrics(ids,state['vocab'],reference=set(state['reference_text']))
+        report.update(metrics)
+        return texts,report
+    if method=='expanding':
+        ids,lengths,extra=sample_expanding(model,args.samples,args.sample_steps,device)
+        texts,metrics=text_metrics(ids.cpu(),state['vocab'],lengths.cpu(),set(state['reference_text']))
+        report.update(metrics);report.update(extra)
+        return texts,report
+    if method=='ssfm':
+        x,extra=evaluate_stochastic(model,args.samples,args.sample_steps,device)
+        report.update(extra)
+    elif method in ['diamond','meta']:
+        with torch.no_grad():
+            outer=torch.randn(args.samples,2,device=device)
+            t=torch.zeros(args.samples,1,device=device)
+            x=posterior_samples(model,outer,t,1,args.posterior_steps)[:,0]
+            guided=guided_samples(model,args.samples,args.sample_steps,args.particles,
+                                  args.reward_strength,args.posterior_steps)
+            report.update(posterior_diagnostics(model,args.posterior_steps))
+            report['unconditional_mean_reward']=float(reward(x,args.reward_strength).mean())
+            report['guided_mean_reward']=float(reward(guided,args.reward_strength).mean())
+            report['guided_distribution']=mixture_metrics(guided)
+            report['posterior_steps']=args.posterior_steps
+            report['particles']=args.particles
+            observation=torch.tensor([[.5,-.3]],device=device)
+            _,estimate,ess=weighted_diamond_samples(observation,.5,4096)
+            report['importance_posterior_mean']=estimate.tolist()
+            report['importance_effective_sample_size']=ess.tolist()
+        # Verify differentiability through the learned posterior's context.
+        observation=observation.detach().requires_grad_(True)
+        value,_,_=posterior_value(model,observation,.5,args.particles,args.reward_strength,args.posterior_steps)
+        gradient=torch.autograd.grad(value.sum(),observation)[0]
+        report['posterior_value_context_gradient']=gradient.tolist()
+        report['finite_value_gradient']=bool(torch.isfinite(gradient).all())
+        report['guided_samples']=guided.detach().cpu().tolist()
+        if 'reward_drift' in state:
+            drift=MapNet(2,state['width']).to(device)
+            drift.load_state_dict(state['reward_drift'])
+            # A complete sampler for the fitted drift; boundary-time queries
+            # extrapolate beyond the [.05,.85] fine-tuning interval.
+            aligned=integrate_velocity(lambda z,a:drift(z,a,a),outer,0.,1.,args.sample_steps)
+            report['finetuned_mean_reward']=float(reward(aligned,args.reward_strength).mean())
+            report['finetuned_finite_samples']=bool(torch.isfinite(aligned).all())
+            report['finetuned_samples']=aligned.detach().cpu().tolist()
+            report['finetuned_sampling_note']='Full-interval Heun integration; boundary times extrapolate beyond fine-tuning support.'
+    else:
+        noise=torch.randn(args.samples,state['dim'],device=device)
+        x=sample_continuous(model,method,noise,args.sample_steps)
+        if method=='latent':
+            ae=Autoencoder(state['width']).to(device)
+            ae.load_state_dict(state['autoencoder'])
+            with torch.no_grad():
+                x=ae.decoder(x*state['latent_std'].to(device)+state['latent_mean'].to(device))
+                heldout=state['heldout_data'].to(device)
+                reconstruction=ae.decoder(ae.encoder(embed_surface(heldout)))
+                report['heldout_reconstruction_mse']=float((reconstruction-embed_surface(heldout)).square().mean())
+                report['surface_residual_rmse']=float((x[:,2:]-.3*(x[:,:1].square()-x[:,1:2].square())).square().mean().sqrt())
+        if method not in ['consistency','meanflow','flow-matching']:
+            from common import finite_map
+            with torch.no_grad():
+                direct=finite_map(model,noise,0.,1.)
+                split=finite_map(model,finite_map(model,noise,0.,.5),.5,1.)
+                report['composition_rmse']=float((direct-split).square().mean().sqrt())
+    report['finite_samples']=bool(torch.isfinite(x).all())
+    if not report['finite_samples']:raise FloatingPointError('Sampling produced nonfinite coordinates.')
+    if x.shape[1]>=2 and state['data_kind']=='four-gaussian-mixture':
+        report.update(mixture_metrics(x[:,:2]))
+    rows=[' '.join(f'{value:.7f}' for value in row) for row in x.detach().cpu().tolist()]
+    return rows,report
+
+
+def main(argv=None):
+    p=parser();args=p.parse_args(argv)
+    if min(args.train_steps,args.teacher_steps,args.sample_steps,args.posterior_steps,
+           args.particles,args.batch_size,args.samples,args.width,args.threads)<1 or args.finetune_steps<0:
+        p.error('Step counts, sizes, and width must be positive; finetune steps must be nonnegative.')
+    if args.lr<=0:p.error('Learning rate must be positive.')
+    if args.device=='cuda' and not torch.cuda.is_available():p.error('CUDA is unavailable; select cpu.')
+    if args.mode=='sample' and args.out is None:p.error('Sample mode requires --out for its checkpoint.')
+    if args.mode!='sample' and args.method is None:args.method='flow-matching'
+    out=Path(args.out) if args.out else ROOT/'outputs'/args.method
+    out.mkdir(parents=True,exist_ok=True)
+    seed_all(args.seed,args.threads)
+    start=time.perf_counter()
+    if args.mode=='sample':
+        state=torch.load(out/'checkpoint.pt',map_location=args.device,weights_only=True)
+        if args.method is not None and args.method!=state['method']:
+            p.error('Requested method differs from the saved checkpoint.')
+        args.method=state['method']
+        if args.method=='ssfm' and args.sample_steps & (args.sample_steps-1):
+            p.error('SSFM sample steps must be a power of two.')
+        model=restore_model(state,args.device)
+        rows,report=generate(model,state,args)
+        (out/'resampled.txt').write_text('\n'.join(rows)+'\n')
+        write_json(out/'sample_report.json',report)
+        print(json.dumps({'method':args.method,'mode':'sample','out':str(out),'seconds':time.perf_counter()-start}))
+        return report
+    if args.method=='ssfm' and args.sample_steps & (args.sample_steps-1):
+        p.error('SSFM sample steps must be a power of two.')
+    if args.finetune_steps and args.method!='meta':p.error('--finetune-steps applies to meta.')
+    if args.data and args.method in ['diamond','meta','ssfm']:
+        p.error('Analytic posterior/OU examples use their specified reference distributions.')
+    stage_logs=[]
+    if args.method in CATEGORICAL+['expanding']:
+        default='variable_text.txt' if args.method=='expanding' else 'phrases.txt'
+        data_path=Path(args.data) if args.data else ROOT/'data'/default
+        ids,lengths,vocab=load_text(data_path,args.method=='expanding')
+        perm=torch.randperm(len(ids));ids,lengths=ids[perm],lengths[perm]
+        split=max(1,int(.8*len(ids)))
+        training=ids[:split].to(args.device)
+        if args.method=='expanding':
+            model,state,logs=train_expanding(training,lengths[:split].to(args.device),vocab,args)
+        else:
+            model,state,logs=train_categorical(args.method,training,vocab,args)
+        reference,_=text_metrics(ids[:split],vocab,lengths[:split])
+        state.update({'reference_text':reference,'data_kind':'synthetic-text' if args.data is None else 'custom-text'})
+        with torch.no_grad():
+            if args.method in CATEGORICAL:
+                heldout=ids[split:].to(args.device)
+                clean=F.one_hot(heldout,len(vocab)).float();t=torch.full((len(clean),1),.5,device=args.device)
+                x,_=interpolate(clean,t)
+                state['validation_ce']=float(F.cross_entropy(model(x,t,t).flatten(0,1),heldout.flatten()))
+    elif args.method=='ssfm':
+        model,state,logs=train_stochastic(args,args.device)
+        state['data_kind']='ornstein-uhlenbeck'
+    else:
+        if args.data:
+            import numpy as np
+            data=torch.tensor(np.loadtxt(args.data,delimiter=','),dtype=torch.float32)
+            if data.ndim!=2 or data.shape[1]!=2 or len(data)<8 or not torch.isfinite(data).all():
+                p.error('Continuous CSV needs at least eight finite rows and exactly two columns, without a header.')
+            data=data[torch.randperm(len(data))];kind='custom-continuous'
+        else:
+            data=mixture_data(4096);kind='four-gaussian-mixture'
+        split=int(.8*len(data));train=data[:split].to(args.device)
+        if args.method=='latent':
+            model,_,state,logs,stage_logs=train_latent(train,args)
+        elif args.method in ['diamond','meta']:
+            model,state,logs=train_posterior(args.method,train,args)
+            if args.finetune_steps:
+                drift,finetune_logs=train_reward_drift(model,train,args)
+                state['reward_drift']=drift.state_dict()
+                write_json(out/'finetune_losses.json',finetune_logs)
+        else:
+            model,state,logs,stage_logs=train_continuous(args.method,train,args)
+        state.update({'data_kind':kind,'heldout_data':data[split:]})
+    state.update({'format_version':1,'method':args.method,'width':args.width,'seed':args.seed})
+    torch.save(state,out/'checkpoint.pt')
+    write_json(out/'config.json',vars(args))
+    write_json(out/'losses.json',logs)
+    if stage_logs:write_json(out/'teacher_losses.json',stage_logs)
+    report={'method':args.method,'data_kind':state['data_kind'],
+            'train_loss_first_20_mean':sum(x['loss'] for x in logs[:20])/len(logs[:20]),
+            'train_loss_last_20_mean':sum(x['loss'] for x in logs[-20:])/len(logs[-20:]),
+            'training_steps':args.train_steps,'python':platform.python_version(),'torch':str(torch.__version__)}
+    if 'validation_ce' in state:report['heldout_diagonal_ce_at_half_time']=state['validation_ce']
+    if args.mode=='train-sample':
+        model.eval()
+        rows,metrics=generate(model,state,args)
+        (out/'samples.txt').write_text('\n'.join(rows)+'\n')
+        report.update(metrics)
+    report['elapsed_seconds']=time.perf_counter()-start
+    write_json(out/'report.json',report)
+    print(json.dumps({'method':args.method,'out':str(out),'loss':report['train_loss_last_20_mean'],
+                      'seconds':report['elapsed_seconds']}),flush=True)
+    return report
+
+
+if __name__=='__main__':main()

lecture_6/run_all.py

@@ -0,0 +1,32 @@
+#!/usr/bin/env python3
+"""Run every method, then verify generation from each saved checkpoint."""
+import argparse
+from pathlib import Path
+
+from run import METHODS, ROOT, main
+
+
+def run_all():
+    parser=argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--quick',action='store_true',help='20-step execution check; not a quality benchmark')
+    parser.add_argument('--out',default=str(ROOT/'outputs'))
+    parser.add_argument('--train-steps',type=int,default=1000)
+    args=parser.parse_args()
+    for method in METHODS:
+        out=Path(args.out)/method
+        command=['--method',method,'--out',str(out)]
+        if args.quick:
+            command+=['--train-steps','20','--teacher-steps','20','--batch-size','8',
+                      '--samples','8','--width','32','--sample-steps','4','--posterior-steps','2','--particles','4']
+        else:command+=['--train-steps',str(args.train_steps),'--teacher-steps',str(args.train_steps)]
+        main(command)
+        sample_args=['--mode','sample','--out',str(out)]
+        if args.quick:
+            sample_args+=['--samples','8','--sample-steps','4','--posterior-steps','2','--particles','4']
+        main(sample_args)
+        if (out/'samples.txt').read_bytes()!=(out/'resampled.txt').read_bytes():
+            raise AssertionError(f'{method}: seeded generation changed after checkpoint reload')
+    print(f'All {len(METHODS)} methods trained, sampled, and reproduced samples after checkpoint reload.')
+
+
+if __name__=='__main__':run_all()

lecture_6/source_manifest.json

@@ -0,0 +1,132 @@
+[
+  {
+    "paper_title": "Flow Map Matching",
+    "paper_url": "https://arxiv.org/abs/2406.07507v2",
+    "methods": [
+      "fmm-lagrangian",
+      "fmm-eulerian"
+    ],
+    "scope": "Frozen learned velocity teachers; residual-map JVPs; small MLP and bounded time range.",
+    "repository": "nmboffi/flow-maps",
+    "commit": "2f115a07fa9073553193e4b265dfc303827af2b0",
+    "inspected_file": "py/common/losses.py",
+    "inspected_file_sha256": "55d5a6cd2b7602afc93f848bda687aa4a287bfb7d1a4095d5071c7df59334d59"
+  },
+  {
+    "paper_title": "How to build a consistency model: Learning flow maps via self-distillation",
+    "paper_url": "https://arxiv.org/abs/2505.18825v2",
+    "methods": [
+      "self-distill"
+    ],
+    "scope": "Diagonal regression plus an EMA progressive target, normalized by a floored interval length. Learned uncertainty weights are omitted.",
+    "repository": "nmboffi/flow-maps",
+    "commit": "2f115a07fa9073553193e4b265dfc303827af2b0",
+    "inspected_file": "py/common/losses.py",
+    "inspected_file_sha256": "55d5a6cd2b7602afc93f848bda687aa4a287bfb7d1a4095d5071c7df59334d59"
+  },
+  {
+    "paper_title": "One Step Diffusion via Shortcut Models",
+    "paper_url": "https://arxiv.org/abs/2410.12557v3",
+    "methods": [
+      "shortcut"
+    ],
+    "scope": "Two half-step average target and dyadic intervals retained. Continuous start-time sampling and small MLP replace image-specific schedules and clipping.",
+    "repository": "kvfrans/shortcut-models",
+    "commit": "601004348667094e1b71f30942199759412d4432",
+    "inspected_file": "targets_shortcut.py",
+    "inspected_file_sha256": "adb2ac83febc1de012a7cdb713b9117d53616473ccd38f1104fd2475d8b1abc3"
+  },
+  {
+    "paper_title": "Mean Flows for One-step Generative Modeling",
+    "paper_url": "https://arxiv.org/abs/2505.13447",
+    "methods": [
+      "meanflow"
+    ],
+    "scope": "Backward clock, conditional-velocity JVP, detached target, adaptive loss, and 75% diagonal proportion retained. Uniform time sampling; no class guidance.",
+    "repository": "Gsunshine/meanflow",
+    "commit": "d70cb55d298ee03c53bf6da67bec281082e4e2d9",
+    "inspected_file": "meanflow.py",
+    "inspected_file_sha256": "b8dcee1033e6ba01de0019d14fad846a399d356d67e549ee311a5e19d2a7c3a4"
+  },
+  {
+    "paper_title": "Flow Map Language Models: One-step Language Modeling via Continuous Denoising",
+    "paper_url": "https://arxiv.org/abs/2602.16813v3",
+    "methods": [
+      "fmlm"
+    ],
+    "scope": "PSD denoiser target and Gaussian decoding clock retained. Independent clean-data diagonal supervision; EMA targets; quadrature lookup and MLP replace full language architecture.",
+    "repository": "david3684/flm",
+    "commit": "a1918d5164e5038e37d0b7a4fb2010ce75b863b3",
+    "inspected_file": "algo.py",
+    "inspected_file_sha256": "65206adc72013bb9b09b68fd847239b3193dd65fd59669bbdee42e061583d76a"
+  },
+  {
+    "paper_title": "Categorical Flow Maps",
+    "paper_url": "https://arxiv.org/abs/2602.12233v1",
+    "methods": [
+      "categorical"
+    ],
+    "scope": "Released ECLD endpoint CE plus sum-of-squares time energy. Detached-target KL has identical student gradient. EMA target replaces the current network target.",
+    "repository": "olsdavis/semicat",
+    "commit": "558602a0fa722514e4a6012f5c46a8ae178b3068",
+    "inspected_file": "semicat/models/semicat.py",
+    "inspected_file_sha256": "82f4bf2cab2fcfa78ccfc59d9094e6c3132f5c2933c1cd7e61ff0b644ac71425"
+  },
+  {
+    "paper_title": "Discrete Flow Maps",
+    "paper_url": "https://arxiv.org/abs/2604.09784v1",
+    "methods": [
+      "discrete-lsd",
+      "discrete-esd"
+    ],
+    "scope": "Derived from the paper. The linked project page has a placeholder Code link, so no author implementation was verified. Positive log denominators are explicitly clamped and counted; no gradient surgery."
+  },
+  {
+    "paper_title": "Diamond Maps: Stochastic Flow Maps",
+    "paper_url": "https://arxiv.org/abs/2602.05993",
+    "methods": [
+      "diamond"
+    ],
+    "scope": "GLASS conditional velocity and Lagrangian posterior distillation. Exact Gaussian-mixture denoiser replaces a pretrained image teacher. Weighted posterior recovery uses an explicit Gaussian proposal with a known density.",
+    "repository": "PeterHolderrieth/diamond_maps",
+    "commit": "d30f65c75a169a2ed624f146b2770aeef882543a",
+    "inspected_file": "posterior_diamond_maps/py/common/losses.py",
+    "inspected_file_sha256": "acb26f1e717e43a648661c1641cadd4a573ed2672b53210ad1ddb207c6a6ac1e"
+  },
+  {
+    "paper_title": "Meta Flow Maps enable scalable reward alignment",
+    "paper_url": "https://arxiv.org/abs/2601.14430v2",
+    "methods": [
+      "meta"
+    ],
+    "scope": "Data training with independent inner/outer noises and a shared endpoint; conditional semigroup objective. Optional Equation 43 fine-tuning also checked against src/mfm/losses/finetune.py.",
+    "repository": "adh1s/mfm",
+    "commit": "53c0f60db695cad88bbace8fb26469614e9e7d7d",
+    "inspected_file": "src/mfm/losses/losses.py",
+    "inspected_file_sha256": "52b99238199f5d6ac691c15566463d52725a6f7ebcdef1e252b2d28da5eb1df2"
+  },
+  {
+    "paper_title": "Expanding Flow Maps",
+    "paper_url": "https://arxiv.org/abs/2607.21585v1",
+    "methods": [
+      "expanding"
+    ],
+    "scope": "Paper Algorithms 2-4 and local-clock equations. The inspected repository contains README and artwork only. Classroom implementation includes a shared lifted canvas, remaining/interval count heads, explicit local clocks, and budget-capped binomial insertions.",
+    "repository": "sophtang/ExpandingFlowMaps",
+    "commit": "4defc1ff168d12526d85b551bea78150b9aec605",
+    "inspected_file": "README.md",
+    "inspected_file_sha256": "d2a55e221a909f7282609cc661e49c07a241570dcd1dab7b55c6df207f821f4a"
+  },
+  {
+    "paper_title": "Strong Stochastic Flow Maps",
+    "paper_url": "https://arxiv.org/abs/2606.01086v1",
+    "methods": [
+      "ssfm"
+    ],
+    "scope": "First two Legendre integrals, Chen composition, known diffusion, small-step matching. The released loss detaches an EMA split target; paper Algorithm 1 detaches the direct target. Both orientations are selectable. OU replaces image/molecular experiments; no learned uncertainty weights.",
+    "repository": "sammccallum/ssfm",
+    "commit": "24b563620a8ee61c683d791fe802b56668f98f9b",
+    "inspected_file": "ssfm/losses.py",
+    "inspected_file_sha256": "3c4414ffb29c906272831e3642827131582c4741fb2457cb71d122564da29233"
+  }
+]

lecture_6/stochastic.py

@@ -0,0 +1,110 @@
+"""Strong stochastic flow maps for dX=-X dt + sigma dW.
+
+The first two shifted-Legendre Brownian integrals have variances h and h/3.
+Chen composition makes coarse and fine evaluations use the same path.
+"""
+import math
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from common import ema_copy, mlp, optimize, time_like
+
+
+def sample_coefficients(h, count=None):
+    h = torch.as_tensor(h)
+    if h.ndim == 0:
+        if count is None: raise ValueError('count is required for scalar h')
+        h = h.expand(count,1)
+    scale = torch.cat([h.sqrt(),(h/3).sqrt()],-1)
+    return torch.randn_like(scale)*scale
+
+
+def chen_two(left,right,h_left,h_right):
+    h=h_left+h_right
+    return torch.stack([left[...,0]+right[...,0],
+        (h_left*left[...,1]+h_right*right[...,1]-h_right*left[...,0]+h_left*right[...,0])/h],-1)
+
+
+class StrongMap(nn.Module):
+    def __init__(self,width=64,sigma=.7):
+        super().__init__()
+        self.net=mlp(5,1,width)
+        self.sigma=sigma
+
+    def average_drift(self,x,s,t,coefficients):
+        return self.net(torch.cat([x,time_like(s,x),time_like(t,x),coefficients],-1))
+
+    def forward(self,x,s,t,coefficients):
+        s,t=time_like(s,x),time_like(t,x)
+        return x+(t-s)*self.average_drift(x,s,t,coefficients)+self.sigma*coefficients[:,:1]
+
+
+def train_stochastic(args,device):
+    model=StrongMap(args.width).to(device)
+    teacher=ema_copy(model)
+    def objective(step):
+        B=args.batch_size
+        s=.9*torch.rand(B,1,device=device)
+        # Exact one-time OU marginals with X0~N(0,1), not simulated trajectories.
+        variance=torch.exp(-2*s)+model.sigma**2/2*(1-torch.exp(-2*s))
+        x=variance.sqrt()*torch.randn(B,1,device=device)
+        h=.001+.019*torch.rand_like(s)
+        c=sample_coefficients(h)
+        target=x-h*x+model.sigma*c[:,:1]
+        matching=((model(x,s,s+h,c)-target).square()/h).mean()
+        diagonal=F.mse_loss(model.average_drift(x,s,s,torch.zeros_like(c)),-x)
+        t=s+.025+(1-s-.025)*torch.rand_like(s)
+        u=(s+t)/2
+        L,R=sample_coefficients(u-s),sample_coefficients(t-u)
+        C=chen_two(L,R,(u-s).squeeze(-1),(t-u).squeeze(-1))
+        if args.ssfm_target=='paper':
+            # Paper Algorithm 1: EMA direct target; differentiate both splits.
+            with torch.no_grad():direct=teacher(x,s,t,C)
+            split=model(model(x,s,u,L),u,t,R)
+        else:
+            # Released code: EMA split target; differentiate the direct map.
+            with torch.no_grad():split=teacher(teacher(x,s,u,L),u,t,R)
+            direct=model(x,s,t,C)
+        consistency=((split-direct).square()/(t-s)).mean()
+        ramp=min(1.,(step+1)/max(1,args.train_steps//4))
+        return diagonal+matching+ramp*consistency,{'drift':diagonal,'small_step':matching,'strong_consistency':consistency}
+    logs=optimize(model,objective,args.train_steps,args.lr,teacher)
+    return model,{'model':model.state_dict(),'sigma':model.sigma},logs
+
+
+def aggregate_tree(coefficients):
+    """Coefficients are [B,leaves,2], with uniform leaf lengths."""
+    c=coefficients
+    h=1/c.shape[1]
+    while c.shape[1]>1:
+        c=chen_two(c[:,::2],c[:,1::2],h,h)
+        h*=2
+    return c[:,0]
+
+
+@torch.no_grad()
+def evaluate_stochastic(model,count,steps,device):
+    if steps & (steps-1):raise ValueError('SSFM sample steps must be a power of two.')
+    leaves=max(512,steps)
+    c=sample_coefficients(torch.tensor(1/leaves,device=device),count*leaves).reshape(count,leaves,2)
+    x0=torch.randn(count,1,device=device)
+    reference=x0.clone()
+    for i in range(leaves):reference=reference-reference/leaves+model.sigma*c[:,i,:1]
+    all_coeff=aggregate_tree(c)
+    direct=model(x0,0.,1.,all_coeff)
+    x=x0.clone()
+    block=leaves//steps
+    for i in range(steps):
+        # aggregate_tree assumes total duration one; equal-half Chen weights are
+        # scale invariant, so its output is valid for these smaller blocks too.
+        ci=aggregate_tree(c[:,i*block:(i+1)*block])
+        x=model(x,i/steps,(i+1)/steps,ci)
+    exact_variance=math.exp(-2)+model.sigma**2/2*(1-math.exp(-2))
+    return x,{'same_noise_reference_rmse':float((x-reference).square().mean().sqrt()),
+              'direct_vs_split_rmse':float((direct-x).square().mean().sqrt()),
+              'one_step_reference_rmse':float((direct-reference).square().mean().sqrt()),
+              'reference_em_steps':leaves,'sample_variance':float(x.var()),
+              'exact_terminal_variance':exact_variance,
+              'reference_note':'Same Brownian increments; EM reference retains finite discretization error.'}

lecture_6/tests/test_mathematics.py

@@ -0,0 +1,189 @@
+"""Independent mathematical checks, including failure cases from the lecture."""
+import math
+import sys
+import unittest
+from pathlib import Path
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+sys.path.insert(0,str(Path(__file__).resolve().parents[1]))
+from common import (MapNet, SequenceNet, finite_map, exact_denoiser, exact_velocity,
+                    mixture_posterior, seed_all, DATA_STD, CENTERS)
+from continuous import lagrangian_residual, eulerian_residual, semigroup_loss
+from categorical import (DecodingClock, categorical_map, composition_target,
+                         corrected_logit_teacher, probability_kl)
+from posterior import glass_denoiser, fine_tune_surrogate, weighted_diamond_samples
+from expanding import (ExpandingNet, local_clock, local_map, gap_counts,
+                       bounded_counts, insert_tokens, count_divergence)
+from stochastic import sample_coefficients, chen_two, aggregate_tree
+
+
+class ExactExponential(nn.Module):
+    def forward(self,x,s,t,context=None):
+        h=t-s
+        # Stable at h=0 while keeping the exact off-diagonal expression.
+        safe=h.clamp_min(1e-12)
+        ratio=torch.where(h.abs()<1e-8,1+h/2+h.square()/6,torch.expm1(h)/safe)
+        return x*ratio
+
+
+class ConstantLogits(nn.Module):
+    def forward(self,x,s,t):
+        return torch.zeros_like(x)+torch.tensor([.2,-.1],dtype=x.dtype,device=x.device)
+
+
+class MathematicsTests(unittest.TestCase):
+    def setUp(self):seed_all(6270)
+
+    def test_exact_identity_and_semigroup(self):
+        model=ExactExponential();x=torch.tensor([[1.2],[-.4]],dtype=torch.float64)
+        s=torch.zeros_like(x);u=s+.3;t=s+.9
+        torch.testing.assert_close(finite_map(model,x,s,s),x)
+        expected=x*math.exp(.9)
+        torch.testing.assert_close(finite_map(model,x,s,t),expected)
+        torch.testing.assert_close(finite_map(model,finite_map(model,x,s,u),u,t),expected)
+
+    def test_lagrangian_and_eulerian_jvp(self):
+        x=torch.tensor([[1.7]],dtype=torch.float64);s=x*0+.2;t=x*0+.8
+        for residual in [lagrangian_residual,eulerian_residual]:
+            value=residual(ExactExponential(),lambda z,a:z,x,s,t)
+            self.assertLess(float(value.abs().max()),1e-10)
+
+    def test_jvp_losses_reach_parameters(self):
+        for residual in [lagrangian_residual,eulerian_residual]:
+            m=MapNet(2,16);x=torch.randn(8,2);s=torch.rand(8,1)*.3;t=s+.5
+            loss=residual(m,lambda z,a:z,x,s,t).square().mean();loss.backward()
+            self.assertTrue(all(p.grad is not None and torch.isfinite(p.grad).all() for p in m.parameters()))
+            self.assertGreater(sum(float(p.grad.abs().sum()) for p in m.parameters()),0)
+
+    def test_composition_alone_does_not_identify_motion(self):
+        m=MapNet(2,8)
+        for p in m.parameters():p.data.zero_()
+        x=torch.randn(8,2);s=torch.zeros(8,1);t=s+1
+        self.assertEqual(float(semigroup_loss(m,m,x,s,t)),0.)
+        self.assertGreater(float((m(x,s,s)-x).square().mean()),0.)
+
+    def test_meanflow_backward_identity(self):
+        def average(z,r,t):return z*(-torch.expm1(-(t-r)))/(t-r)
+        z=torch.tensor([[1.7]],dtype=torch.float64);r=z*0+.2;t=z*0+.8
+        u,du=torch.func.jvp(average,(z,r,t),(z,torch.zeros_like(r),torch.ones_like(t)))
+        torch.testing.assert_close(u,z-(t-r)*du,atol=1e-11,rtol=1e-11)
+
+    def test_analytic_mixture_velocity(self):
+        x=torch.randn(20,2,dtype=torch.float64);t=torch.rand(20,1,dtype=torch.float64)*.9
+        torch.testing.assert_close(exact_velocity(x,t),(exact_denoiser(x,t)-x)/(1-t),atol=1e-10,rtol=1e-10)
+
+    def test_categorical_numerical_step(self):
+        class Net(nn.Module):
+            def forward(self,x,s,t):return torch.tensor([.1,.7,.2]).log().expand_as(x)
+        x=torch.tensor([[[-.2,.6,1.1]]]);y,p=categorical_map(Net(),x,.25,.75)
+        torch.testing.assert_close(y,torch.tensor([[[0.,2/3,.5]]]),atol=1e-7,rtol=1e-6)
+        self.assertGreater(float(y.sum()),1.1)
+        torch.testing.assert_close(p.sum(-1),torch.ones(1,1))
+
+    def test_weighted_probability_composition_matches_maps(self):
+        net=SequenceNet(3,4,12);x=torch.randn(6,3,4)
+        s=torch.zeros(6,1);u=s+.5;t=s+.75
+        q=composition_target(net,x,s,u,t)
+        mid,_=categorical_map(net,x,s,u);split,_=categorical_map(net,mid,u,t)
+        direct=(1-t[...,None])*x+t[...,None]*q
+        torch.testing.assert_close(direct,split)
+        torch.testing.assert_close(q.sum(-1),torch.ones(6,3))
+
+    def test_detached_kl_gradient(self):
+        logits=torch.tensor([.2,-.1],requires_grad=True);p=torch.tensor([.4,.6])
+        probability_kl(logits,p).backward()
+        torch.testing.assert_close(logits.grad,logits.softmax(-1)-p)
+
+    def test_discrete_teachers_with_zero_derivative(self):
+        x=torch.randn(5,3,2);s=torch.zeros(5,1)+.2;t=s+.5;net=ConstantLogits()
+        for kind in ['discrete-lsd','discrete-esd']:
+            logits,target,clipped=corrected_logit_teacher(net,net,x,s,t,kind)
+            torch.testing.assert_close(target,logits.softmax(-1))
+            self.assertEqual(float(clipped),0.)
+
+    def test_endpoint_match_does_not_remove_temporal_residual(self):
+        dt=torch.tensor([.2,-.2]);residual=.5*.5*dt
+        self.assertAlmostEqual(float(residual.square().sum()),.005,places=7)
+
+    def test_decoding_clock_inverse(self):
+        clock=DecodingClock(13);tau=torch.linspace(0,1,101)
+        time=clock.inverse(tau)
+        self.assertTrue((time[1:]>=time[:-1]).all())
+        self.assertEqual(float(time[0]),0.);self.assertEqual(float(time[-1]),1.)
+
+    def test_glass_matches_two_observation_bayes_rule(self):
+        inner=torch.tensor([[.3,-.2]],dtype=torch.float64);outer=inner+.4
+        s=torch.tensor([[.35]],dtype=torch.float64);t=s+.2
+        # Independent derivation by conditioning each Gaussian component twice.
+        variance=1/(1/DATA_STD**2+(s/(1-s))**2+(t/(1-t))**2)
+        means=variance[:,None]*(CENTERS.double()/DATA_STD**2+s[:,None]*inner[:,None]/(1-s[:,None])**2+t[:,None]*outer[:,None]/(1-t[:,None])**2)
+        covariance=torch.tensor([[float((1-s)**2+s**2*DATA_STD**2),float(s*t*DATA_STD**2)],
+                                 [float(s*t*DATA_STD**2),float((1-t)**2+t**2*DATA_STD**2)]],dtype=torch.float64)
+        residual=torch.stack([inner[:,None]-s[:,None]*CENTERS,outer[:,None]-t[:,None]*CENTERS],-1)
+        logits=-.5*torch.einsum('bkdi,ij,bkdj->bk',residual,torch.linalg.inv(covariance),residual)
+        expected=(logits.softmax(-1)[...,None]*means).sum(1)
+        torch.testing.assert_close(glass_denoiser(inner,s,outer,t),expected,atol=1e-9,rtol=1e-9)
+
+    def test_glass_no_inner_observation(self):
+        outer=torch.randn(5,2);t=torch.full((5,1),.4)
+        torch.testing.assert_close(glass_denoiser(torch.randn(5,2),torch.zeros_like(t),outer,t),exact_denoiser(outer,t))
+
+    def test_meta_surrogate_gradient(self):
+        delta=torch.tensor([[.3]],requires_grad=True)
+        w=torch.tensor([[1.],[4.]]);gw=torch.tensor([[.2],[.8]])
+        loss=fine_tune_surrogate(delta,w,gw,.5);loss.backward()
+        torch.testing.assert_close(delta.grad,2*(w*delta.detach()-.5*gw).mean().reshape(1,1))
+
+    def test_importance_weights_recover_posterior_mean(self):
+        x=torch.tensor([[.5,-.3]])
+        _,mean,ess=weighted_diamond_samples(x,.5,50000)
+        self.assertLess(float((mean-exact_denoiser(x,.5)).abs().max()),.2)
+        self.assertGreater(float(ess),100)
+
+    def test_insertion_preserves_order_and_clocks(self):
+        x=torch.tensor([[1.,0.],[0.,1.]]);b=torch.tensor([0.,.2]);counts=torch.tensor([1,0,1])
+        noise=torch.tensor([[-.2,.4],[.3,-.1]])
+        y,bt=insert_tokens(x,b,counts,noise,.5)
+        torch.testing.assert_close(y[1:3],x)
+        torch.testing.assert_close(bt,torch.tensor([.5,0.,.2,.5]))
+        torch.testing.assert_close(local_clock(.75,bt),torch.tensor([.5,.75,.6875,.5]))
+
+    def test_gap_labels_and_global_budget(self):
+        torch.testing.assert_close(gap_counts(torch.tensor([1,4]),torch.tensor([0,2,3,5]),6),torch.tensor([1.,2.,1.]))
+        counts,_=bounded_counts(torch.tensor([100.,100.,100.]),4)
+        self.assertLessEqual(int(counts.sum()),4)
+        self.assertTrue((counts>=0).all())
+
+    def test_count_divergence_minimizes_conditional_mean(self):
+        prediction=torch.tensor(2.,requires_grad=True)
+        count_divergence(torch.tensor([0.,1.,5.]),prediction).mean().backward()
+        self.assertAlmostEqual(float(prediction.grad),0.,places=6)
+
+    def test_insertion_interval_zero_and_terminal(self):
+        net=ExpandingNet(4,3,8);x=torch.zeros(1,4,3);b=torch.zeros(1,4);m=b.bool();s=torch.tensor([[.3]])
+        self.assertEqual(float(net.counts(x,b,m,s,s).sum()),0.)
+        t=torch.ones_like(s)
+        self.assertTrue(torch.isfinite(net.counts(x,b,m,s,t)).all())
+
+    def test_chen_arithmetic_and_covariance(self):
+        left=torch.tensor([.2,.04]);right=torch.tensor([-.1,-.02])
+        torch.testing.assert_close(chen_two(left,right,.5,.5),torch.tensor([.1,-.14]))
+        L=sample_coefficients(torch.tensor(.3),100000);R=sample_coefficients(torch.tensor(.7),100000)
+        C=chen_two(L,R,.3,.7)
+        torch.testing.assert_close(torch.cov(C.T),torch.diag(torch.tensor([1.,1/3])),atol=.015,rtol=0)
+
+    def test_same_noise_constant_sde_and_wrong_noise_counterexample(self):
+        direct=1+.5+.8*(.2-.1)
+        split=(1+.25+.8*.2)+.25+.8*(-.1)
+        self.assertAlmostEqual(direct,split);self.assertAlmostEqual(direct,1.58)
+        self.assertNotAlmostEqual(direct,1+.5+.8*(-.3))
+
+    def test_chen_tree_matches_increment_sum(self):
+        c=torch.randn(6,16,2)
+        torch.testing.assert_close(aggregate_tree(c)[:,0],c[:,:,0].sum(1))
+
+
+if __name__=='__main__':unittest.main()

lecture_6/verified_examples/README.md

@@ -0,0 +1,76 @@
+# Verified Lecture 6 examples
+
+These are actual outputs from seeded CPU runs on 2026-09-13. All 16 default
+methods completed 1,000 training steps, sampled, and reproduced identical
+seeded samples after checkpoint reload. Teacher/autoencoder stages also used
+1,000 steps where applicable. The complete suite passed on PyTorch 2.9.1+cpu
+and PyTorch 2.14.0+cpu with Python 3.12.14. These checked-in outputs use 2.9.1.
+All 23 independent mathematical tests passed in both environments.
+
+Reproduce the default runs from `lecture_6/`:
+
+```bash
+python run_all.py --out outputs/verified-torch291 --train-steps 1000
+python -m unittest discover -s tests -v
+python numerical_examples.py --output outputs/numerical
+```
+
+The baseline reports use 128 samples, eight sampling steps, width 64, batch size
+64, and seed 6270. Posterior methods use four inner-map steps and 32 reward
+particles. Training losses are not comparable across different objectives.
+
+| Method | First 20 loss mean | Last 20 loss mean | Selected diagnostic |
+| --- | ---: | ---: | --- |
+| [categorical](categorical/report.json) | 2.47260 | 0.53130 | grammar support 1.000; mean length 4.00 |
+| [consistency](consistency/report.json) | 1.97463 | 1.97996 | mean center distance 0.753; mode entropy 1.386 |
+| [diamond](diamond/report.json) | 2.30034 | 1.39093 | posterior mean RMSE 0.285; variance RMSE 0.314 |
+| [discrete-esd](discrete-esd/report.json) | 2.47265 | 0.55223 | grammar support 1.000; mean length 4.00 |
+| [discrete-lsd](discrete-lsd/report.json) | 2.47260 | 0.52785 | grammar support 0.992; mean length 4.00 |
+| [expanding](expanding/report.json) | 4.25715 | 2.91940 | grammar support 0.617; mean length 4.09 |
+| [flow-matching](flow-matching/report.json) | 3.24047 | 2.02429 | mean center distance 0.575; mode entropy 1.383 |
+| [fmlm](fmlm/report.json) | 2.47259 | 0.52103 | grammar support 1.000; mean length 4.00 |
+| [fmm-eulerian](fmm-eulerian/report.json) | 1.99308 | 2.03712 | mean center distance 0.474; mode entropy 1.386 |
+| [fmm-lagrangian](fmm-lagrangian/report.json) | 1.98769 | 2.05232 | mean center distance 0.461; mode entropy 1.384 |
+| [latent](latent/report.json) | 1.95472 | 1.47142 | mean center distance 1.024; mode entropy 1.383 |
+| [meanflow](meanflow/report.json) | 1.67703 | 1.57479 | mean center distance 0.480; mode entropy 1.382 |
+| [meta](meta/report.json) | 3.04584 | 1.50320 | posterior mean RMSE 0.277; variance RMSE 0.487 |
+| [self-distill](self-distill/report.json) | 3.19830 | 2.18310 | mean center distance 0.729; mode entropy 1.381 |
+| [shortcut](shortcut/report.json) | 3.22197 | 2.18371 | mean center distance 0.738; mode entropy 1.385 |
+| [ssfm](ssfm/report.json) | 0.40952 | 0.00012 | same-noise RMSE 0.00483 |
+
+The fixed-length text examples recover almost all phrases in the four-family
+grammar at eight steps. The expanding example still produces invalid
+combinations; its support fraction is retained. Posterior means and variances
+both have measurable error. Continuous samples cover the modes but retain
+finite distributional error. These observations describe the actual small runs,
+and are not evidence of reproducing paper-scale quality.
+
+## Additional verified recipes
+
+```bash
+python examples/meta.py --finetune-steps 300 --out outputs/meta-finetune-verified
+python run.py --mode sample --out outputs/meta-finetune-verified
+python examples/ssfm.py --ssfm-target paper --lr 0.0001 --train-steps 2000 --out outputs/ssfm-paper-stable
+python run.py --mode sample --out outputs/ssfm-paper-stable
+```
+
+`meta-finetune/` records the complete optional reward-drift stage and its
+sampling metrics. `ssfm-paper/` records the paper's gradient orientation with
+its separate lower-learning-rate recipe. Both reproduced seeded samples after
+reloading. The paper-orientation recipe is optimization-sensitive: an initial
+1,000-step run at learning rate 0.001 diverged in accuracy despite finite
+numbers. Its same-noise RMSE was 51.94. This is why the default implementation
+uses the released code's EMA split target. The lower-learning-rate paper run
+also remained inaccurate, with same-noise RMSE 3.85 and one-step RMSE 22.40.
+It is retained as an optimization-sensitivity comparison, not a successful
+strong-solution fit. These failures concern this small-model training recipe;
+they do not establish a general failure of the paper's objective.
+
+`numerical_checks.json` retains the lecture's scalar-flow, posterior-gradient,
+probability-update, insertion-clock, and Brownian-composition checks. The scalar
+map's one-step endpoint RMSE is 0.011255 after 6,000 steps.
+
+The JSON files preserve each executed configuration and the full loss history.
+`checkpoint.pt` files are regenerated by the commands and excluded from Git,
+matching the prior lecture folders. `training_reload_log.txt` and
+`mathematics_test_log.txt` record the completed verification.

lecture_6/verified_examples/categorical/config.json

@@ -0,0 +1,21 @@
+{
+  "method": "categorical",
+  "mode": "train-sample",
+  "train_steps": 1000,
+  "teacher_steps": 1000,
+  "finetune_steps": 0,
+  "sample_steps": 8,
+  "posterior_steps": 4,
+  "particles": 32,
+  "reward_strength": 1.0,
+  "ssfm_target": "official-code",
+  "batch_size": 64,
+  "samples": 128,
+  "width": 64,
+  "lr": 0.001,
+  "seed": 6270,
+  "threads": 1,
+  "device": "cpu",
+  "data": null,
+  "out": "outputs/verified-torch291/categorical"
+}

lecture_6/verified_examples/categorical/report.json

@@ -0,0 +1,19 @@
+{
+  "method": "categorical",
+  "data_kind": "synthetic-text",
+  "train_loss_first_20_mean": 2.472597396373749,
+  "train_loss_last_20_mean": 0.5313047975301742,
+  "training_steps": 1000,
+  "python": "3.12.14",
+  "torch": "2.9.1+cpu",
+  "heldout_diagonal_ce_at_half_time": 0.45874327421188354,
+  "sample_steps": 8,
+  "samples": 128,
+  "sampling_seed": 6370,
+  "unique_fraction": 0.03125,
+  "token_entropy": 2.415271520614624,
+  "mean_length": 4.0,
+  "empty_fraction": 0.0,
+  "training_support_fraction": 1.0,
+  "elapsed_seconds": 5.423348005999287
+}

lecture_6/verified_examples/categorical/samples.txt

@@ -0,0 +1,128 @@
+gold star moves down
+blue square moves right
+gold star moves down
+green triangle moves up
+green triangle moves up
+gold star moves down
+blue square moves right
+green triangle moves up
+gold star moves down
+blue square moves right
+blue square moves right
+red circle moves left
+gold star moves down
+gold star moves down
+blue square moves right
+gold star moves down
+red circle moves left
+green triangle moves up
+blue square moves right
+gold star moves down
+red circle moves left
+red circle moves left
+green triangle moves up
+green triangle moves up
+red circle moves left
+blue square moves right
+green triangle moves up
+green triangle moves up
+gold star moves down
+gold star moves down
+blue square moves right
+blue square moves right
+green triangle moves up
+blue square moves right
+gold star moves down
+gold star moves down
+gold star moves down
+green triangle moves up
+gold star moves down
+green triangle moves up
+blue square moves right
+blue square moves right
+green triangle moves up
+gold star moves down
+gold star moves down
+gold star moves down
+gold star moves down
+green triangle moves up
+red circle moves left
+blue square moves right
+blue square moves right
+red circle moves left
+blue square moves right
+gold star moves down
+blue square moves right
+green triangle moves up
+red circle moves left
+blue square moves right
+green triangle moves up
+gold star moves down
+gold star moves down
+green triangle moves up
+blue square moves right
+green triangle moves up
+blue square moves right
+gold star moves down
+red circle moves left
+red circle moves left
+green triangle moves up
+green triangle moves up
+gold star moves down
+blue square moves right
+green triangle moves up
+blue square moves right
+blue square moves right
+blue square moves right
+gold star moves down
+green triangle moves up
+blue square moves right
+green triangle moves up
+red circle moves left
+gold star moves down
+blue square moves right
+gold star moves down
+red circle moves left
+blue square moves right
+gold star moves down
+gold star moves down
+red circle moves left
+blue square moves right
+red circle moves left
+green triangle moves up
+green triangle moves up
+gold star moves down
+gold star moves down
+green triangle moves up
+red circle moves left
+blue square moves right
+blue square moves right
+red circle moves left
+blue square moves right
+red circle moves left
+red circle moves left
+red circle moves left
+gold star moves down
+green triangle moves up
+blue square moves right
+gold star moves down
+green triangle moves up
+red circle moves left
+gold star moves down
+red circle moves left
+green triangle moves up
+red circle moves left
+red circle moves left
+blue square moves right
+gold star moves down
+green triangle moves up
+blue square moves right
+red circle moves left
+gold star moves down
+gold star moves down
+gold star moves down
+gold star moves down
+red circle moves left
+green triangle moves up
+gold star moves down
+gold star moves down

lecture_6/verified_examples/consistency/config.json

@@ -0,0 +1,21 @@
+{
+  "method": "consistency",
+  "mode": "train-sample",
+  "train_steps": 1000,
+  "teacher_steps": 1000,
+  "finetune_steps": 0,
+  "sample_steps": 8,
+  "posterior_steps": 4,
+  "particles": 32,
+  "reward_strength": 1.0,
+  "ssfm_target": "official-code",
+  "batch_size": 64,
+  "samples": 128,
+  "width": 64,
+  "lr": 0.001,
+  "seed": 6270,
+  "threads": 1,
+  "device": "cpu",
+  "data": null,
+  "out": "outputs/verified-torch291/consistency"
+}

lecture_6/verified_examples/consistency/report.json

@@ -0,0 +1,23 @@
+{
+  "method": "consistency",
+  "data_kind": "four-gaussian-mixture",
+  "train_loss_first_20_mean": 1.9746342718601226,
+  "train_loss_last_20_mean": 1.979964154958725,
+  "training_steps": 1000,
+  "python": "3.12.14",
+  "torch": "2.9.1+cpu",
+  "sample_steps": 8,
+  "samples": 128,
+  "sampling_seed": 6370,
+  "finite_samples": true,
+  "mean_distance_to_center": 0.7534297704696655,
+  "mode_fractions": [
+    0.265625,
+    0.25,
+    0.2421875,
+    0.2421875
+  ],
+  "mode_entropy": 1.3855692148208618,
+  "mean_target_log_density": -7.795166492462158,
+  "elapsed_seconds": 3.6416576160045224
+}

lecture_6/verified_examples/consistency/samples.txt

@@ -0,0 +1,128 @@
+-1.3468293 1.9925021
+-1.9144447 1.9564545
+-1.8007749 -2.0338287
+-1.9767884 0.2636437
+0.4644759 -1.6100454
+-1.8300068 2.0076437
+1.8318725 -1.9341205
+-2.0396407 -1.7563149
+-2.0363870 -1.9806333
+1.0769628 1.7889115
+0.1340955 1.3456073
+1.7798740 -0.5705347
+2.0481019 0.6637622
+0.5697179 -0.0764643
+-1.6694670 -2.1333323
+1.8878694 -2.1138554
+1.3903328 1.7709353
+1.3510761 -1.8566105
+0.3830533 0.7458028
+1.8123627 -2.1073432
+-1.8500829 1.9188887
+-1.9960229 -0.9180380
+-2.0363288 -1.9042180
+-1.7602494 -1.8083441
+0.1228525 -2.1484063
+-1.9425657 0.6813170
+0.9867452 0.2309941
+-0.5291872 -2.0708153
+-1.9266198 1.5953822
+-0.5012224 -1.9294865
+-1.2689531 -1.5546137
+-1.5741237 -1.8497179
+-0.2626170 -1.6006305
+-1.9871480 -2.1482060
+2.0151453 1.6644613
+-0.1489909 -2.0024414
+2.0573957 1.4528581
+-0.2505112 1.9968579
+1.7871954 0.6475998
+1.8512402 -0.2385958
+0.1899848 1.9513735
+-1.9870937 -0.6491414
+-0.2206627 -1.5287057
+1.6779035 1.0743937
+-2.0382123 -2.1339107
+-0.7100272 1.9262421
+-1.6801538 -1.8067182
+1.8546274 -2.1622596
+1.5299554 1.2159443
+-1.9548271 1.3845929
+0.4502607 2.0001235
+-0.0262713 -0.3656653
+0.9616735 -1.0734054
+0.8062696 1.8239052
+-0.4521321 1.8881189
+-1.4971223 1.4280133
+1.2114744 -1.1458678
+-0.5043197 1.1818823
+-2.0388222 -1.5035791
+-1.3591039 0.9025277
+-2.0305381 -0.3237583
+1.2204367 1.9129956
+0.3672223 1.9162316
+2.0437796 0.6819464
+1.7619469 -1.9899399
+1.7444496 -1.7915529
+0.0078025 -1.0781103
+1.8828144 -2.0743697
+-1.9811914 -2.1016579
+-1.9607775 -0.8501123
+1.8939016 1.3304422
+1.7166512 0.3390130
+-0.0189792 -0.1771857
+-0.6227967 1.3938310
+1.9738953 1.8059912
+0.0466411 -2.0397491
+1.7655673 -1.9929256
+1.8075031 -1.8748856
+1.9316881 1.7332070
+1.6287178 -1.4561557
+-1.1077086 1.9380562
+1.8720345 1.2524481
+1.5174096 1.3109862
+-1.0241457 0.1447840
+-1.4152889 -0.7302222
+-1.7965597 -0.9746199
+1.1624088 1.3963469
+-0.4442409 1.4740853
+-1.5870914 1.7438629
+-0.2185191 -0.1900957
+1.6071042 1.2579901
+-1.4373560 1.5747277
+1.4849627 -1.7022858
+1.4948187 1.8095267
+-1.9391402 0.7429604
+-1.0735885 1.5553175
+-0.5654426 1.1886865
+1.4643161 -2.1554084
+-2.0169582 0.1357105
+1.1027896 -2.1490393
+-1.9470890 -2.0728121
+-1.9837592 0.9692687
+-0.8680381 2.0148981
+-1.6903710 -1.3934293
+-2.0227857 -1.9842577
+-1.8039379 1.8491473
+0.8419592 -2.0827856
+1.3736886 -2.1311660
+-1.6144994 -1.8624021
+-1.2275145 1.6786437
+-0.7562458 -2.1603642
+1.5444539 1.1505903
+0.6380618 0.8644766
+1.3242871 -2.1602378
+1.8642752 -1.8713832
+-1.7034441 -1.0856594
+1.6279490 1.2668343
+1.7643193 1.6216198
+1.0740453 -0.7010343
+-1.8715028 -2.0824487
+1.4197307 -1.6388113
+-1.6448362 0.3884219
+-1.9149169 1.7842357
+1.8479285 1.3923659
+1.4299746 -2.1582410
+-1.3168693 1.5584388
+-1.4611561 1.2482233
+1.8228297 -2.0618267

lecture_6/verified_examples/consistency/teacher_losses.json

@@ -0,0 +1,4002 @@
+[
+  {
+    "step": 0,
+    "loss": 2.901901960372925
+  },
+  {
+    "step": 1,
+    "loss": 4.030956745147705
+  },
+  {
+    "step": 2,
+    "loss": 3.7361669540405273
+  },
+  {
+    "step": 3,
+    "loss": 3.1736671924591064
+  },
+  {
+    "step": 4,
+    "loss": 3.2819948196411133
+  },
+  {
+    "step": 5,
+    "loss": 2.6870734691619873
+  },
+  {
+    "step": 6,
+    "loss": 3.2287180423736572
+  },
+  {
+    "step": 7,
+    "loss": 3.540705919265747
+  },
+  {
+    "step": 8,
+    "loss": 2.320319652557373
+  },
+  {
+    "step": 9,
+    "loss": 3.2631494998931885
+  },
+  {
+    "step": 10,
+    "loss": 3.2428276538848877
+  },
+  {
+    "step": 11,
+    "loss": 3.1067261695861816
+  },
+  {
+    "step": 12,
+    "loss": 2.9945061206817627
+  },
+  {
+    "step": 13,
+    "loss": 3.429009199142456
+  },
+  {
+    "step": 14,
+    "loss": 3.202871084213257
+  },
+  {
+    "step": 15,
+    "loss": 3.343949317932129
+  },
+  {
+    "step": 16,
+    "loss": 3.338503837585449
+  },
+  {
+    "step": 17,
+    "loss": 2.27468204498291
+  },
+  {
+    "step": 18,
+    "loss": 3.2059783935546875
+  },
+  {
+    "step": 19,
+    "loss": 3.2105021476745605
+  },
+  {
+    "step": 20,
+    "loss": 3.607123374938965
+  },
+  {
+    "step": 21,
+    "loss": 2.5147407054901123
+  },
+  {
+    "step": 22,
+    "loss": 2.4429540634155273
+  },
+  {
+    "step": 23,
+    "loss": 2.693626880645752
+  },
+  {
+    "step": 24,
+    "loss": 2.922039031982422
+  },
+  {
+    "step": 25,
+    "loss": 3.412623405456543
+  },
+  {
+    "step": 26,
+    "loss": 3.049421548843384
+  },
+  {
+    "step": 27,
+    "loss": 3.4507718086242676
+  },
+  {
+    "step": 28,
+    "loss": 3.2085723876953125
+  },
+  {
+    "step": 29,
+    "loss": 3.431762933731079
+  },
+  {
+    "step": 30,
+    "loss": 3.186777114868164
+  },
+  {
+    "step": 31,
+    "loss": 2.358400821685791
+  },
+  {
+    "step": 32,
+    "loss": 2.739171028137207
+  },
+  {
+    "step": 33,
+    "loss": 3.3288445472717285
+  },
+  {
+    "step": 34,
+    "loss": 2.9850246906280518
+  },
+  {
+    "step": 35,
+    "loss": 2.979214906692505
+  },
+  {
+    "step": 36,
+    "loss": 3.4282126426696777
+  },
+  {
+    "step": 37,
+    "loss": 2.3601138591766357
+  },
+  {
+    "step": 38,
+    "loss": 2.754340410232544
+  },
+  {
+    "step": 39,
+    "loss": 2.772657871246338
+  },
+  {
+    "step": 40,
+    "loss": 3.816230535507202
+  },
+  {
+    "step": 41,
+    "loss": 2.872300624847412
+  },
+  {
+    "step": 42,
+    "loss": 3.0081868171691895
+  },
+  {
+    "step": 43,
+    "loss": 2.6166443824768066
+  },
+  {
+    "step": 44,
+    "loss": 2.4548749923706055
+  },
+  {
+    "step": 45,
+    "loss": 2.7366158962249756
+  },
+  {
+    "step": 46,
+    "loss": 2.7148921489715576
+  },
+  {
+    "step": 47,
+    "loss": 2.800400972366333
+  },
+  {
+    "step": 48,
+    "loss": 2.9447460174560547
+  },
+  {
+    "step": 49,
+    "loss": 2.708866834640503
+  },
+  {
+    "step": 50,
+    "loss": 2.883293628692627
+  },
+  {
+    "step": 51,
+    "loss": 2.8611910343170166
+  },
+  {
+    "step": 52,
+    "loss": 2.477092742919922
+  },
+  {
+    "step": 53,
+    "loss": 2.551604986190796
+  },
+  {
+    "step": 54,
+    "loss": 2.338512897491455
+  },
+  {
+    "step": 55,
+    "loss": 3.476226329803467
+  },
+  {
+    "step": 56,
+    "loss": 2.717756509780884
+  },
+  {
+    "step": 57,
+    "loss": 2.5666332244873047
+  },
+  {
+    "step": 58,
+    "loss": 2.3942244052886963
+  },
+  {
+    "step": 59,
+    "loss": 2.3049826622009277
+  },
+  {
+    "step": 60,
+    "loss": 2.817511796951294
+  },
+  {
+    "step": 61,
+    "loss": 2.9381887912750244
+  },
+  {
+    "step": 62,
+    "loss": 3.0970544815063477
+  },
+  {
+    "step": 63,
+    "loss": 2.7542073726654053
+  },
+  {
+    "step": 64,
+    "loss": 2.95668625831604
+  },
+  {
+    "step": 65,
+    "loss": 3.3345565795898438
+  },
+  {
+    "step": 66,
+    "loss": 3.0203516483306885
+  },
+  {
+    "step": 67,
+    "loss": 3.0535619258880615
+  },
+  {
+    "step": 68,
+    "loss": 2.2485365867614746
+  },
+  {
+    "step": 69,
+    "loss": 3.313269853591919
+  },
+  {
+    "step": 70,
+    "loss": 3.151857614517212
+  },
+  {
+    "step": 71,
+    "loss": 3.3004636764526367
+  },
+  {
+    "step": 72,
+    "loss": 3.13859224319458
+  },
+  {
+    "step": 73,
+    "loss": 2.831691026687622
+  },
+  {
+    "step": 74,
+    "loss": 2.824505090713501
+  },
+  {
+    "step": 75,
+    "loss": 3.2724480628967285
+  },
+  {
+    "step": 76,
+    "loss": 3.3112194538116455
+  },
+  {
+    "step": 77,
+    "loss": 3.5525429248809814
+  },
+  {
+    "step": 78,
+    "loss": 3.02317214012146
+  },
+  {
+    "step": 79,
+    "loss": 2.3117425441741943
+  },
+  {
+    "step": 80,
+    "loss": 2.52921986579895
+  },
+  {
+    "step": 81,
+    "loss": 3.346832752227783
+  },
+  {
+    "step": 82,
+    "loss": 2.557271957397461
+  },
+  {
+    "step": 83,
+    "loss": 2.8063182830810547
+  },
+  {
+    "step": 84,
+    "loss": 3.1693220138549805
+  },
+  {
+    "step": 85,
+    "loss": 2.1511852741241455
+  },
+  {
+    "step": 86,
+    "loss": 2.3931725025177
+  },
+  {
+    "step": 87,
+    "loss": 2.7104268074035645
+  },
+  {
+    "step": 88,
+    "loss": 3.144993305206299
+  },
+  {
+    "step": 89,
+    "loss": 2.5371882915496826
+  },
+  {
+    "step": 90,
+    "loss": 2.635395050048828
+  },
+  {
+    "step": 91,
+    "loss": 2.9128007888793945
+  },
+  {
+    "step": 92,
+    "loss": 2.091165781021118
+  },
+  {
+    "step": 93,
+    "loss": 2.678255081176758
+  },
+  {
+    "step": 94,
+    "loss": 3.1458537578582764
+  },
+  {
+    "step": 95,
+    "loss": 3.0256667137145996
+  },
+  {
+    "step": 96,
+    "loss": 3.504319906234741
+  },
+  {
+    "step": 97,
+    "loss": 2.4840328693389893
+  },
+  {
+    "step": 98,
+    "loss": 3.156357765197754
+  },
+  {
+    "step": 99,
+    "loss": 2.409656047821045
+  },
+  {
+    "step": 100,
+    "loss": 2.607604503631592
+  },
+  {
+    "step": 101,
+    "loss": 3.1481847763061523
+  },
+  {
+    "step": 102,
+    "loss": 2.3221395015716553
+  },
+  {
+    "step": 103,
+    "loss": 2.7403645515441895
+  },
+  {
+    "step": 104,
+    "loss": 2.9177207946777344
+  },
+  {
+    "step": 105,
+    "loss": 3.1296420097351074
+  },
+  {
+    "step": 106,
+    "loss": 2.714864492416382
+  },
+  {
+    "step": 107,
+    "loss": 2.2136282920837402
+  },
+  {
+    "step": 108,
+    "loss": 3.0324344635009766
+  },
+  {
+    "step": 109,
+    "loss": 2.8028464317321777
+  },
+  {
+    "step": 110,
+    "loss": 3.500157356262207
+  },
+  {
+    "step": 111,
+    "loss": 3.659149646759033
+  },
+  {
+    "step": 112,
+    "loss": 2.6796770095825195
+  },
+  {
+    "step": 113,
+    "loss": 4.001369476318359
+  },
+  {
+    "step": 114,
+    "loss": 3.12492299079895
+  },
+  {
+    "step": 115,
+    "loss": 2.186939239501953
+  },
+  {
+    "step": 116,
+    "loss": 3.0609922409057617
+  },
+  {
+    "step": 117,
+    "loss": 2.7359426021575928
+  },
+  {
+    "step": 118,
+    "loss": 2.8246054649353027
+  },
+  {
+    "step": 119,
+    "loss": 2.0663673877716064
+  },
+  {
+    "step": 120,
+    "loss": 2.82535719871521
+  },
+  {
+    "step": 121,
+    "loss": 2.7651166915893555
+  },
+  {
+    "step": 122,
+    "loss": 2.3627724647521973
+  },
+  {
+    "step": 123,
+    "loss": 3.1746630668640137
+  },
+  {
+    "step": 124,
+    "loss": 2.2432596683502197
+  },
+  {
+    "step": 125,
+    "loss": 2.632728338241577
+  },
+  {
+    "step": 126,
+    "loss": 3.509660005569458
+  },
+  {
+    "step": 127,
+    "loss": 2.6030969619750977
+  },
+  {
+    "step": 128,
+    "loss": 2.819704532623291
+  },
+  {
+    "step": 129,
+    "loss": 2.7884159088134766
+  },
+  {
+    "step": 130,
+    "loss": 3.066481590270996
+  },
+  {
+    "step": 131,
+    "loss": 2.5375585556030273
+  },
+  {
+    "step": 132,
+    "loss": 2.6820192337036133
+  },
+  {
+    "step": 133,
+    "loss": 2.8677313327789307
+  },
+  {
+    "step": 134,
+    "loss": 3.295163631439209
+  },
+  {
+    "step": 135,
+    "loss": 3.008159637451172
+  },
+  {
+    "step": 136,
+    "loss": 2.6790380477905273
+  },
+  {
+    "step": 137,
+    "loss": 3.2993409633636475
+  },
+  {
+    "step": 138,
+    "loss": 2.3291473388671875
+  },
+  {
+    "step": 139,
+    "loss": 3.2662582397460938
+  },
+  {
+    "step": 140,
+    "loss": 3.1028847694396973
+  },
+  {
+    "step": 141,
+    "loss": 3.5241870880126953
+  },
+  {
+    "step": 142,
+    "loss": 2.6248207092285156
+  },
+  {
+    "step": 143,
+    "loss": 3.289806365966797
+  },
+  {
+    "step": 144,
+    "loss": 2.761228322982788
+  },
+  {
+    "step": 145,
+    "loss": 2.6195461750030518
+  },
+  {
+    "step": 146,
+    "loss": 3.2968337535858154
+  },
+  {
+    "step": 147,
+    "loss": 2.2302370071411133
+  },
+  {
+    "step": 148,
+    "loss": 2.7895140647888184
+  },
+  {
+    "step": 149,
+    "loss": 3.0598552227020264
+  },
+  {
+    "step": 150,
+    "loss": 2.991975784301758
+  },
+  {
+    "step": 151,
+    "loss": 2.758124589920044
+  },
+  {
+    "step": 152,
+    "loss": 2.8041367530822754
+  },
+  {
+    "step": 153,
+    "loss": 2.7885689735412598
+  },
+  {
+    "step": 154,
+    "loss": 2.642705202102661
+  },
+  {
+    "step": 155,
+    "loss": 3.208401679992676
+  },
+  {
+    "step": 156,
+    "loss": 2.890425205230713
+  },
+  {
+    "step": 157,
+    "loss": 3.0730347633361816
+  },
+  {
+    "step": 158,
+    "loss": 2.9297609329223633
+  },
+  {
+    "step": 159,
+    "loss": 3.1628060340881348
+  },
+  {
+    "step": 160,
+    "loss": 2.235241174697876
+  },
+  {
+    "step": 161,
+    "loss": 3.431405544281006
+  },
+  {
+    "step": 162,
+    "loss": 3.1083548069000244
+  },
+  {
+    "step": 163,
+    "loss": 2.592428207397461
+  },
+  {
+    "step": 164,
+    "loss": 2.538461923599243
+  },
+  {
+    "step": 165,
+    "loss": 2.8032314777374268
+  },
+  {
+    "step": 166,
+    "loss": 2.8171889781951904
+  },
+  {
+    "step": 167,
+    "loss": 2.857837677001953
+  },
+  {
+    "step": 168,
+    "loss": 2.95976185798645
+  },
+  {
+    "step": 169,
+    "loss": 2.500030994415283
+  },
+  {
+    "step": 170,
+    "loss": 2.632225513458252
+  },
+  {
+    "step": 171,
+    "loss": 2.696317672729492
+  },
+  {
+    "step": 172,
+    "loss": 2.671130657196045
+  },
+  {
+    "step": 173,
+    "loss": 3.2583823204040527
+  },
+  {
+    "step": 174,
+    "loss": 2.750561237335205
+  },
+  {
+    "step": 175,
+    "loss": 2.7291007041931152
+  },
+  {
+    "step": 176,
+    "loss": 2.5894086360931396
+  },
+  {
+    "step": 177,
+    "loss": 2.551658868789673
+  },
+  {
+    "step": 178,
+    "loss": 2.7688674926757812
+  },
+  {
+    "step": 179,
+    "loss": 2.5460314750671387
+  },
+  {
+    "step": 180,
+    "loss": 2.573371648788452
+  },
+  {
+    "step": 181,
+    "loss": 2.5819363594055176
+  },
+  {
+    "step": 182,
+    "loss": 2.7828993797302246
+  },
+  {
+    "step": 183,
+    "loss": 2.3213307857513428
+  },
+  {
+    "step": 184,
+    "loss": 2.626124382019043
+  },
+  {
+    "step": 185,
+    "loss": 2.0648884773254395
+  },
+  {
+    "step": 186,
+    "loss": 2.6399617195129395
+  },
+  {
+    "step": 187,
+    "loss": 2.4929964542388916
+  },
+  {
+    "step": 188,
+    "loss": 2.738865613937378
+  },
+  {
+    "step": 189,
+    "loss": 2.869166851043701
+  },
+  {
+    "step": 190,
+    "loss": 2.6773924827575684
+  },
+  {
+    "step": 191,
+    "loss": 3.207840919494629
+  },
+  {
+    "step": 192,
+    "loss": 2.593851327896118
+  },
+  {
+    "step": 193,
+    "loss": 3.145773410797119
+  },
+  {
+    "step": 194,
+    "loss": 2.6272592544555664
+  },
+  {
+    "step": 195,
+    "loss": 3.2182958126068115
+  },
+  {
+    "step": 196,
+    "loss": 2.6175713539123535
+  },
+  {
+    "step": 197,
+    "loss": 2.102283000946045
+  },
+  {
+    "step": 198,
+    "loss": 2.3725454807281494
+  },
+  {
+    "step": 199,
+    "loss": 2.310814142227173
+  },
+  {
+    "step": 200,
+    "loss": 2.554727077484131
+  },
+  {
+    "step": 201,
+    "loss": 2.290694236755371
+  },
+  {
+    "step": 202,
+    "loss": 2.0614757537841797
+  },
+  {
+    "step": 203,
+    "loss": 2.4946200847625732
+  },
+  {
+    "step": 204,
+    "loss": 3.1302361488342285
+  },
+  {
+    "step": 205,
+    "loss": 2.4278793334960938
+  },
+  {
+    "step": 206,
+    "loss": 2.7277281284332275
+  },
+  {
+    "step": 207,
+    "loss": 1.6022748947143555
+  },
+  {
+    "step": 208,
+    "loss": 2.2651426792144775
+  },
+  {
+    "step": 209,
+    "loss": 2.446990489959717
+  },
+  {
+    "step": 210,
+    "loss": 2.829681634902954
+  },
+  {
+    "step": 211,
+    "loss": 2.8909616470336914
+  },
+  {
+    "step": 212,
+    "loss": 2.6464128494262695
+  },
+  {
+    "step": 213,
+    "loss": 3.0840983390808105
+  },
+  {
+    "step": 214,
+    "loss": 2.442357063293457
+  },
+  {
+    "step": 215,
+    "loss": 3.077253580093384
+  },
+  {
+    "step": 216,
+    "loss": 3.0163815021514893
+  },
+  {
+    "step": 217,
+    "loss": 3.3182713985443115
+  },
+  {
+    "step": 218,
+    "loss": 2.6422321796417236
+  },
+  {
+    "step": 219,
+    "loss": 2.7757790088653564
+  },
+  {
+    "step": 220,
+    "loss": 2.817136287689209
+  },
+  {
+    "step": 221,
+    "loss": 2.7244391441345215
+  },
+  {
+    "step": 222,
+    "loss": 3.3463165760040283
+  },
+  {
+    "step": 223,
+    "loss": 2.5178160667419434
+  },
+  {
+    "step": 224,
+    "loss": 2.3083720207214355
+  },
+  {
+    "step": 225,
+    "loss": 2.6672210693359375
+  },
+  {
+    "step": 226,
+    "loss": 2.953312635421753
+  },
+  {
+    "step": 227,
+    "loss": 2.245347023010254
+  },
+  {
+    "step": 228,
+    "loss": 2.6127264499664307
+  },
+  {
+    "step": 229,
+    "loss": 2.3829171657562256
+  },
+  {
+    "step": 230,
+    "loss": 1.9589905738830566
+  },
+  {
+    "step": 231,
+    "loss": 2.7865824699401855
+  },
+  {
+    "step": 232,
+    "loss": 2.784212589263916
+  },
+  {
+    "step": 233,
+    "loss": 2.638892412185669
+  },
+  {
+    "step": 234,
+    "loss": 2.2629337310791016
+  },
+  {
+    "step": 235,
+    "loss": 2.3170275688171387
+  },
+  {
+    "step": 236,
+    "loss": 2.3931238651275635
+  },
+  {
+    "step": 237,
+    "loss": 2.5577549934387207
+  },
+  {
+    "step": 238,
+    "loss": 2.67926025390625
+  },
+  {
+    "step": 239,
+    "loss": 2.399914026260376
+  },
+  {
+    "step": 240,
+    "loss": 2.719255208969116
+  },
+  {
+    "step": 241,
+    "loss": 2.645662784576416
+  },
+  {
+    "step": 242,
+    "loss": 2.1666066646575928
+  },
+  {
+    "step": 243,
+    "loss": 2.710801839828491
+  },
+  {
+    "step": 244,
+    "loss": 2.2600297927856445
+  },
+  {
+    "step": 245,
+    "loss": 2.4880356788635254
+  },
+  {
+    "step": 246,
+    "loss": 2.333265542984009
+  },
+  {
+    "step": 247,
+    "loss": 2.3228213787078857
+  },
+  {
+    "step": 248,
+    "loss": 2.98348331451416
+  },
+  {
+    "step": 249,
+    "loss": 2.44199538230896
+  },
+  {
+    "step": 250,
+    "loss": 2.6789755821228027
+  },
+  {
+    "step": 251,
+    "loss": 2.432734489440918
+  },
+  {
+    "step": 252,
+    "loss": 2.3950858116149902
+  },
+  {
+    "step": 253,
+    "loss": 3.024618148803711
+  },
+  {
+    "step": 254,
+    "loss": 2.7950377464294434
+  },
+  {
+    "step": 255,
+    "loss": 2.6282715797424316
+  },
+  {
+    "step": 256,
+    "loss": 2.5368154048919678
+  },
+  {
+    "step": 257,
+    "loss": 2.5300967693328857
+  },
+  {
+    "step": 258,
+    "loss": 2.71651029586792
+  },
+  {
+    "step": 259,
+    "loss": 2.345527172088623
+  },
+  {
+    "step": 260,
+    "loss": 1.991959571838379
+  },
+  {
+    "step": 261,
+    "loss": 2.2183401584625244
+  },
+  {
+    "step": 262,
+    "loss": 2.7822048664093018
+  },
+  {
+    "step": 263,
+    "loss": 2.955565929412842
+  },
+  {
+    "step": 264,
+    "loss": 2.171292781829834
+  },
+  {
+    "step": 265,
+    "loss": 3.007766008377075
+  },
+  {
+    "step": 266,
+    "loss": 2.5466020107269287
+  },
+  {
+    "step": 267,
+    "loss": 2.4173848628997803
+  },
+  {
+    "step": 268,
+    "loss": 2.6756205558776855
+  },
+  {
+    "step": 269,
+    "loss": 2.4434726238250732
+  },
+  {
+    "step": 270,
+    "loss": 2.3825483322143555
+  },
+  {
+    "step": 271,
+    "loss": 2.5813300609588623
+  },
+  {
+    "step": 272,
+    "loss": 2.342559814453125
+  },
+  {
+    "step": 273,
+    "loss": 2.664748191833496
+  },
+  {
+    "step": 274,
+    "loss": 2.8664631843566895
+  },
+  {
+    "step": 275,
+    "loss": 2.8220643997192383
+  },
+  {
+    "step": 276,
+    "loss": 2.6016016006469727
+  },
+  {
+    "step": 277,
+    "loss": 2.292630672454834
+  },
+  {
+    "step": 278,
+    "loss": 2.3435916900634766
+  },
+  {
+    "step": 279,
+    "loss": 2.225806474685669
+  },
+  {
+    "step": 280,
+    "loss": 2.861757755279541
+  },
+  {
+    "step": 281,
+    "loss": 2.4925742149353027
+  },
+  {
+    "step": 282,
+    "loss": 2.3195114135742188
+  },
+  {
+    "step": 283,
+    "loss": 2.723316192626953
+  },
+  {
+    "step": 284,
+    "loss": 2.3318111896514893
+  },
+  {
+    "step": 285,
+    "loss": 2.3740956783294678
+  },
+  {
+    "step": 286,
+    "loss": 2.7161083221435547
+  },
+  {
+    "step": 287,
+    "loss": 1.8623343706130981
+  },
+  {
+    "step": 288,
+    "loss": 3.0774383544921875
+  },
+  {
+    "step": 289,
+    "loss": 2.684474468231201
+  },
+  {
+    "step": 290,
+    "loss": 2.1260976791381836
+  },
+  {
+    "step": 291,
+    "loss": 2.633331775665283
+  },
+  {
+    "step": 292,
+    "loss": 3.048750400543213
+  },
+  {
+    "step": 293,
+    "loss": 2.1517202854156494
+  },
+  {
+    "step": 294,
+    "loss": 2.703066349029541
+  },
+  {
+    "step": 295,
+    "loss": 2.2772347927093506
+  },
+  {
+    "step": 296,
+    "loss": 2.532550573348999
+  },
+  {
+    "step": 297,
+    "loss": 2.5133583545684814
+  },
+  {
+    "step": 298,
+    "loss": 2.2637734413146973
+  },
+  {
+    "step": 299,
+    "loss": 2.8052215576171875
+  },
+  {
+    "step": 300,
+    "loss": 2.3975658416748047
+  },
+  {
+    "step": 301,
+    "loss": 2.142946720123291
+  },
+  {
+    "step": 302,
+    "loss": 2.3383138179779053
+  },
+  {
+    "step": 303,
+    "loss": 2.956615924835205
+  },
+  {
+    "step": 304,
+    "loss": 2.268173933029175
+  },
+  {
+    "step": 305,
+    "loss": 2.2826666831970215
+  },
+  {
+    "step": 306,
+    "loss": 2.504539728164673
+  },
+  {
+    "step": 307,
+    "loss": 2.5966529846191406
+  },
+  {
+    "step": 308,
+    "loss": 1.9715884923934937
+  },
+  {
+    "step": 309,
+    "loss": 2.8570797443389893
+  },
+  {
+    "step": 310,
+    "loss": 2.3709073066711426
+  },
+  {
+    "step": 311,
+    "loss": 2.3691515922546387
+  },
+  {
+    "step": 312,
+    "loss": 2.5264945030212402
+  },
+  {
+    "step": 313,
+    "loss": 2.853039503097534
+  },
+  {
+    "step": 314,
+    "loss": 2.6715211868286133
+  },
+  {
+    "step": 315,
+    "loss": 2.4609460830688477
+  },
+  {
+    "step": 316,
+    "loss": 2.3723304271698
+  },
+  {
+    "step": 317,
+    "loss": 2.623749017715454
+  },
+  {
+    "step": 318,
+    "loss": 2.3588414192199707
+  },
+  {
+    "step": 319,
+    "loss": 2.2384724617004395
+  },
+  {
+    "step": 320,
+    "loss": 2.6265077590942383
+  },
+  {
+    "step": 321,
+    "loss": 2.3127033710479736
+  },
+  {
+    "step": 322,
+    "loss": 2.5432920455932617
+  },
+  {
+    "step": 323,
+    "loss": 2.902012348175049
+  },
+  {
+    "step": 324,
+    "loss": 2.2379114627838135
+  },
+  {
+    "step": 325,
+    "loss": 2.8483102321624756
+  },
+  {
+    "step": 326,
+    "loss": 3.0475165843963623
+  },
+  {
+    "step": 327,
+    "loss": 2.2012202739715576
+  },
+  {
+    "step": 328,
+    "loss": 3.10513973236084
+  },
+  {
+    "step": 329,
+    "loss": 2.7302069664001465
+  },
+  {
+    "step": 330,
+    "loss": 2.05947208404541
+  },
+  {
+    "step": 331,
+    "loss": 2.0735864639282227
+  },
+  {
+    "step": 332,
+    "loss": 2.351263999938965
+  },
+  {
+    "step": 333,
+    "loss": 2.1889426708221436
+  },
+  {
+    "step": 334,
+    "loss": 1.893702745437622
+  },
+  {
+    "step": 335,
+    "loss": 2.886248826980591
+  },
+  {
+    "step": 336,
+    "loss": 2.281409502029419
+  },
+  {
+    "step": 337,
+    "loss": 2.278639793395996
+  },
+  {
+    "step": 338,
+    "loss": 2.7411670684814453
+  },
+  {
+    "step": 339,
+    "loss": 2.242872714996338
+  },
+  {
+    "step": 340,
+    "loss": 2.567128896713257
+  },
+  {
+    "step": 341,
+    "loss": 2.1635279655456543
+  },
+  {
+    "step": 342,
+    "loss": 2.7505345344543457
+  },
+  {
+    "step": 343,
+    "loss": 2.4954588413238525
+  },
+  {
+    "step": 344,
+    "loss": 2.4626471996307373
+  },
+  {
+    "step": 345,
+    "loss": 2.457416296005249
+  },
+  {
+    "step": 346,
+    "loss": 2.4571924209594727
+  },
+  {
+    "step": 347,
+    "loss": 2.3100287914276123
+  },
+  {
+    "step": 348,
+    "loss": 2.4437308311462402
+  },
+  {
+    "step": 349,
+    "loss": 2.0501744747161865
+  },
+  {
+    "step": 350,
+    "loss": 2.1469335556030273
+  },
+  {
+    "step": 351,
+    "loss": 2.464240550994873
+  },
+  {
+    "step": 352,
+    "loss": 1.8512563705444336
+  },
+  {
+    "step": 353,
+    "loss": 2.4845597743988037
+  },
+  {
+    "step": 354,
+    "loss": 2.224407434463501
+  },
+  {
+    "step": 355,
+    "loss": 2.5380733013153076
+  },
+  {
+    "step": 356,
+    "loss": 2.51659893989563
+  },
+  {
+    "step": 357,
+    "loss": 2.0740654468536377
+  },
+  {
+    "step": 358,
+    "loss": 2.069371461868286
+  },
+  {
+    "step": 359,
+    "loss": 2.823765516281128
+  },
+  {
+    "step": 360,
+    "loss": 2.5208475589752197
+  },
+  {
+    "step": 361,
+    "loss": 2.116133213043213
+  },
+  {
+    "step": 362,
+    "loss": 2.556251049041748
+  },
+  {
+    "step": 363,
+    "loss": 2.223443031311035
+  },
+  {
+    "step": 364,
+    "loss": 2.6511716842651367
+  },
+  {
+    "step": 365,
+    "loss": 1.691718339920044
+  },
+  {
+    "step": 366,
+    "loss": 2.411846160888672
+  },
+  {
+    "step": 367,
+    "loss": 2.582014799118042
+  },
+  {
+    "step": 368,
+    "loss": 2.2078142166137695
+  },
+  {
+    "step": 369,
+    "loss": 2.609196662902832
+  },
+  {
+    "step": 370,
+    "loss": 2.444159984588623
+  },
+  {
+    "step": 371,
+    "loss": 2.0085737705230713
+  },
+  {
+    "step": 372,
+    "loss": 2.4077141284942627
+  },
+  {
+    "step": 373,
+    "loss": 2.8086421489715576
+  },
+  {
+    "step": 374,
+    "loss": 2.359715700149536
+  },
+  {
+    "step": 375,
+    "loss": 2.343597888946533
+  },
+  {
+    "step": 376,
+    "loss": 1.948641061782837
+  },
+  {
+    "step": 377,
+    "loss": 2.3641631603240967
+  },
+  {
+    "step": 378,
+    "loss": 2.0386908054351807
+  },
+  {
+    "step": 379,
+    "loss": 2.9259984493255615
+  },
+  {
+    "step": 380,
+    "loss": 2.6643362045288086
+  },
+  {
+    "step": 381,
+    "loss": 2.8584165573120117
+  },
+  {
+    "step": 382,
+    "loss": 2.6431076526641846
+  },
+  {
+    "step": 383,
+    "loss": 2.8262100219726562
+  },
+  {
+    "step": 384,
+    "loss": 2.1087937355041504
+  },
+  {
+    "step": 385,
+    "loss": 2.412994861602783
+  },
+  {
+    "step": 386,
+    "loss": 2.452765703201294
+  },
+  {
+    "step": 387,
+    "loss": 2.2612946033477783
+  },
+  {
+    "step": 388,
+    "loss": 2.9486331939697266
+  },
+  {
+    "step": 389,
+    "loss": 2.2730305194854736
+  },
+  {
+    "step": 390,
+    "loss": 2.526773691177368
+  },
+  {
+    "step": 391,
+    "loss": 2.397615909576416
+  },
+  {
+    "step": 392,
+    "loss": 2.3714022636413574
+  },
+  {
+    "step": 393,
+    "loss": 2.501068353652954
+  },
+  {
+    "step": 394,
+    "loss": 2.101055860519409
+  },
+  {
+    "step": 395,
+    "loss": 1.8332078456878662
+  },
+  {
+    "step": 396,
+    "loss": 2.620009183883667
+  },
+  {
+    "step": 397,
+    "loss": 2.6437087059020996
+  },
+  {
+    "step": 398,
+    "loss": 2.1969316005706787
+  },
+  {
+    "step": 399,
+    "loss": 2.578944444656372
+  },
+  {
+    "step": 400,
+    "loss": 2.8199615478515625
+  },
+  {
+    "step": 401,
+    "loss": 1.9555848836898804
+  },
+  {
+    "step": 402,
+    "loss": 2.3566293716430664
+  },
+  {
+    "step": 403,
+    "loss": 1.9274892807006836
+  },
+  {
+    "step": 404,
+    "loss": 2.3533687591552734
+  },
+  {
+    "step": 405,
+    "loss": 2.3114140033721924
+  },
+  {
+    "step": 406,
+    "loss": 2.3387277126312256
+  },
+  {
+    "step": 407,
+    "loss": 2.3504040241241455
+  },
+  {
+    "step": 408,
+    "loss": 2.694427490234375
+  },
+  {
+    "step": 409,
+    "loss": 2.370358467102051
+  },
+  {
+    "step": 410,
+    "loss": 1.8371087312698364
+  },
+  {
+    "step": 411,
+    "loss": 2.34598970413208
+  },
+  {
+    "step": 412,
+    "loss": 2.001140832901001
+  },
+  {
+    "step": 413,
+    "loss": 2.0652122497558594
+  },
+  {
+    "step": 414,
+    "loss": 2.631736993789673
+  },
+  {
+    "step": 415,
+    "loss": 2.3046376705169678
+  },
+  {
+    "step": 416,
+    "loss": 2.365978240966797
+  },
+  {
+    "step": 417,
+    "loss": 2.047502040863037
+  },
+  {
+    "step": 418,
+    "loss": 2.6210899353027344
+  },
+  {
+    "step": 419,
+    "loss": 2.1412880420684814
+  },
+  {
+    "step": 420,
+    "loss": 2.2439045906066895
+  },
+  {
+    "step": 421,
+    "loss": 2.1576952934265137
+  },
+  {
+    "step": 422,
+    "loss": 2.051438331604004
+  },
+  {
+    "step": 423,
+    "loss": 2.0313544273376465
+  },
+  {
+    "step": 424,
+    "loss": 3.118997573852539
+  },
+  {
+    "step": 425,
+    "loss": 2.471263885498047
+  },
+  {
+    "step": 426,
+    "loss": 1.698156714439392
+  },
+  {
+    "step": 427,
+    "loss": 2.799497365951538
+  },
+  {
+    "step": 428,
+    "loss": 2.1427876949310303
+  },
+  {
+    "step": 429,
+    "loss": 2.820594549179077
+  },
+  {
+    "step": 430,
+    "loss": 1.9575327634811401
+  },
+  {
+    "step": 431,
+    "loss": 2.6089365482330322
+  },
+  {
+    "step": 432,
+    "loss": 2.0574324131011963
+  },
+  {
+    "step": 433,
+    "loss": 2.868596315383911
+  },
+  {
+    "step": 434,
+    "loss": 2.1540873050689697
+  },
+  {
+    "step": 435,
+    "loss": 2.039017677307129
+  },
+  {
+    "step": 436,
+    "loss": 2.4460833072662354
+  },
+  {
+    "step": 437,
+    "loss": 2.5863380432128906
+  },
+  {
+    "step": 438,
+    "loss": 2.5211329460144043
+  },
+  {
+    "step": 439,
+    "loss": 2.0439255237579346
+  },
+  {
+    "step": 440,
+    "loss": 2.083531141281128
+  },
+  {
+    "step": 441,
+    "loss": 2.2036032676696777
+  },
+  {
+    "step": 442,
+    "loss": 2.1617541313171387
+  },
+  {
+    "step": 443,
+    "loss": 2.2706785202026367
+  },
+  {
+    "step": 444,
+    "loss": 2.3687477111816406
+  },
+  {
+    "step": 445,
+    "loss": 2.1419553756713867
+  },
+  {
+    "step": 446,
+    "loss": 2.3326170444488525
+  },
+  {
+    "step": 447,
+    "loss": 2.0808708667755127
+  },
+  {
+    "step": 448,
+    "loss": 2.309194803237915
+  },
+  {
+    "step": 449,
+    "loss": 2.8360342979431152
+  },
+  {
+    "step": 450,
+    "loss": 2.3126049041748047
+  },
+  {
+    "step": 451,
+    "loss": 2.2485318183898926
+  },
+  {
+    "step": 452,
+    "loss": 2.6139726638793945
+  },
+  {
+    "step": 453,
+    "loss": 2.174654483795166
+  },
+  {
+    "step": 454,
+    "loss": 2.4408035278320312
+  },
+  {
+    "step": 455,
+    "loss": 2.3054635524749756
+  },
+  {
+    "step": 456,
+    "loss": 2.500479221343994
+  },
+  {
+    "step": 457,
+    "loss": 2.712820053100586
+  },
+  {
+    "step": 458,
+    "loss": 2.3475780487060547
+  },
+  {
+    "step": 459,
+    "loss": 2.476759910583496
+  },
+  {
+    "step": 460,
+    "loss": 2.163086175918579
+  },
+  {
+    "step": 461,
+    "loss": 2.1306872367858887
+  },
+  {
+    "step": 462,
+    "loss": 2.569213390350342
+  },
+  {
+    "step": 463,
+    "loss": 1.9032127857208252
+  },
+  {
+    "step": 464,
+    "loss": 2.670799493789673
+  },
+  {
+    "step": 465,
+    "loss": 2.116924285888672
+  },
+  {
+    "step": 466,
+    "loss": 2.200054883956909
+  },
+  {
+    "step": 467,
+    "loss": 2.726794719696045
+  },
+  {
+    "step": 468,
+    "loss": 2.4743218421936035
+  },
+  {
+    "step": 469,
+    "loss": 2.1093218326568604
+  },
+  {
+    "step": 470,
+    "loss": 1.994795322418213
+  },
+  {
+    "step": 471,
+    "loss": 2.4643502235412598
+  },
+  {
+    "step": 472,
+    "loss": 2.4895074367523193
+  },
+  {
+    "step": 473,
+    "loss": 2.212019920349121
+  },
+  {
+    "step": 474,
+    "loss": 1.9118317365646362
+  },
+  {
+    "step": 475,
+    "loss": 1.9181571006774902
+  },
+  {
+    "step": 476,
+    "loss": 2.6781809329986572
+  },
+  {
+    "step": 477,
+    "loss": 1.9821468591690063
+  },
+  {
+    "step": 478,
+    "loss": 2.641598701477051
+  },
+  {
+    "step": 479,
+    "loss": 2.110565185546875
+  },
+  {
+    "step": 480,
+    "loss": 2.528163194656372
+  },
+  {
+    "step": 481,
+    "loss": 2.2186732292175293
+  },
+  {
+    "step": 482,
+    "loss": 1.94599187374115
+  },
+  {
+    "step": 483,
+    "loss": 2.3718581199645996
+  },
+  {
+    "step": 484,
+    "loss": 2.222639560699463
+  },
+  {
+    "step": 485,
+    "loss": 2.4093661308288574
+  },
+  {
+    "step": 486,
+    "loss": 2.54209041595459
+  },
+  {
+    "step": 487,
+    "loss": 2.024200201034546
+  },
+  {
+    "step": 488,
+    "loss": 2.5153119564056396
+  },
+  {
+    "step": 489,
+    "loss": 1.658121943473816
+  },
+  {
+    "step": 490,
+    "loss": 2.2563576698303223
+  },
+  {
+    "step": 491,
+    "loss": 2.3844380378723145
+  },
+  {
+    "step": 492,
+    "loss": 2.1449241638183594
+  },
+  {
+    "step": 493,
+    "loss": 2.2108402252197266
+  },
+  {
+    "step": 494,
+    "loss": 2.344649314880371
+  },
+  {
+    "step": 495,
+    "loss": 2.548719644546509
+  },
+  {
+    "step": 496,
+    "loss": 1.7746530771255493
+  },
+  {
+    "step": 497,
+    "loss": 2.309767246246338
+  },
+  {
+    "step": 498,
+    "loss": 2.0731115341186523
+  },
+  {
+    "step": 499,
+    "loss": 2.5398313999176025
+  },
+  {
+    "step": 500,
+    "loss": 2.091824769973755
+  },
+  {
+    "step": 501,
+    "loss": 1.9239593744277954
+  },
+  {
+    "step": 502,
+    "loss": 1.9768168926239014
+  },
+  {
+    "step": 503,
+    "loss": 2.260188341140747
+  },
+  {
+    "step": 504,
+    "loss": 2.553421974182129
+  },
+  {
+    "step": 505,
+    "loss": 2.4699223041534424
+  },
+  {
+    "step": 506,
+    "loss": 2.630557060241699
+  },
+  {
+    "step": 507,
+    "loss": 2.0340492725372314
+  },
+  {
+    "step": 508,
+    "loss": 2.028682231903076
+  },
+  {
+    "step": 509,
+    "loss": 2.45908260345459
+  },
+  {
+    "step": 510,
+    "loss": 2.2142536640167236
+  },
+  {
+    "step": 511,
+    "loss": 2.4662668704986572
+  },
+  {
+    "step": 512,
+    "loss": 2.27579927444458
+  },
+  {
+    "step": 513,
+    "loss": 2.2575185298919678
+  },
+  {
+    "step": 514,
+    "loss": 2.6080620288848877
+  },
+  {
+    "step": 515,
+    "loss": 2.5972297191619873
+  },
+  {
+    "step": 516,
+    "loss": 2.1600301265716553
+  },
+  {
+    "step": 517,
+    "loss": 1.9268064498901367
+  },
+  {
+    "step": 518,
+    "loss": 2.278988838195801
+  },
+  {
+    "step": 519,
+    "loss": 2.4084134101867676
+  },
+  {
+    "step": 520,
+    "loss": 2.7133266925811768
+  },
+  {
+    "step": 521,
+    "loss": 2.123317003250122
+  },
+  {
+    "step": 522,
+    "loss": 2.0050742626190186
+  },
+  {
+    "step": 523,
+    "loss": 2.311730146408081
+  },
+  {
+    "step": 524,
+    "loss": 2.100144624710083
+  },
+  {
+    "step": 525,
+    "loss": 2.182567834854126
+  },
+  {
+    "step": 526,
+    "loss": 2.6850223541259766
+  },
+  {
+    "step": 527,
+    "loss": 2.64558482170105
+  },
+  {
+    "step": 528,
+    "loss": 2.065207004547119
+  },
+  {
+    "step": 529,
+    "loss": 2.57641339302063
+  },
+  {
+    "step": 530,
+    "loss": 2.37033748626709
+  },
+  {
+    "step": 531,
+    "loss": 2.3202273845672607
+  },
+  {
+    "step": 532,
+    "loss": 2.5949063301086426
+  },
+  {
+    "step": 533,
+    "loss": 2.621664047241211
+  },
+  {
+    "step": 534,
+    "loss": 2.4353041648864746
+  },
+  {
+    "step": 535,
+    "loss": 2.0931382179260254
+  },
+  {
+    "step": 536,
+    "loss": 2.1506476402282715
+  },
+  {
+    "step": 537,
+    "loss": 2.0098910331726074
+  },
+  {
+    "step": 538,
+    "loss": 1.8340765237808228
+  },
+  {
+    "step": 539,
+    "loss": 2.3008437156677246
+  },
+  {
+    "step": 540,
+    "loss": 2.4273364543914795
+  },
+  {
+    "step": 541,
+    "loss": 2.4804060459136963
+  },
+  {
+    "step": 542,
+    "loss": 2.25795841217041
+  },
+  {
+    "step": 543,
+    "loss": 2.357149600982666
+  },
+  {
+    "step": 544,
+    "loss": 2.1500396728515625
+  },
+  {
+    "step": 545,
+    "loss": 2.459463596343994
+  },
+  {
+    "step": 546,
+    "loss": 2.2286086082458496
+  },
+  {
+    "step": 547,
+    "loss": 2.4219624996185303
+  },
+  {
+    "step": 548,
+    "loss": 2.686821222305298
+  },
+  {
+    "step": 549,
+    "loss": 2.2649004459381104
+  },
+  {
+    "step": 550,
+    "loss": 2.307349920272827
+  },
+  {
+    "step": 551,
+    "loss": 1.8991062641143799
+  },
+  {
+    "step": 552,
+    "loss": 1.9229987859725952
+  },
+  {
+    "step": 553,
+    "loss": 2.5056543350219727
+  },
+  {
+    "step": 554,
+    "loss": 2.1609060764312744
+  },
+  {
+    "step": 555,
+    "loss": 1.9778512716293335
+  },
+  {
+    "step": 556,
+    "loss": 2.471957206726074
+  },
+  {
+    "step": 557,
+    "loss": 1.8614418506622314
+  },
+  {
+    "step": 558,
+    "loss": 2.335939884185791
+  },
+  {
+    "step": 559,
+    "loss": 1.712602972984314
+  },
+  {
+    "step": 560,
+    "loss": 2.4900341033935547
+  },
+  {
+    "step": 561,
+    "loss": 2.3448946475982666
+  },
+  {
+    "step": 562,
+    "loss": 2.333690643310547
+  },
+  {
+    "step": 563,
+    "loss": 2.3033447265625
+  },
+  {
+    "step": 564,
+    "loss": 1.985489010810852
+  },
+  {
+    "step": 565,
+    "loss": 2.1400845050811768
+  },
+  {
+    "step": 566,
+    "loss": 2.6713786125183105
+  },
+  {
+    "step": 567,
+    "loss": 2.2288951873779297
+  },
+  {
+    "step": 568,
+    "loss": 1.8685835599899292
+  },
+  {
+    "step": 569,
+    "loss": 1.8806357383728027
+  },
+  {
+    "step": 570,
+    "loss": 2.147825241088867
+  },
+  {
+    "step": 571,
+    "loss": 1.928283452987671
+  },
+  {
+    "step": 572,
+    "loss": 1.8230233192443848
+  },
+  {
+    "step": 573,
+    "loss": 1.8360440731048584
+  },
+  {
+    "step": 574,
+    "loss": 2.6959915161132812
+  },
+  {
+    "step": 575,
+    "loss": 1.9693267345428467
+  },
+  {
+    "step": 576,
+    "loss": 2.7600269317626953
+  },
+  {
+    "step": 577,
+    "loss": 2.2720563411712646
+  },
+  {
+    "step": 578,
+    "loss": 2.299776077270508
+  },
+  {
+    "step": 579,
+    "loss": 2.212860584259033
+  },
+  {
+    "step": 580,
+    "loss": 2.3061680793762207
+  },
+  {
+    "step": 581,
+    "loss": 2.2603089809417725
+  },
+  {
+    "step": 582,
+    "loss": 2.0869574546813965
+  },
+  {
+    "step": 583,
+    "loss": 2.23502254486084
+  },
+  {
+    "step": 584,
+    "loss": 2.2090988159179688
+  },
+  {
+    "step": 585,
+    "loss": 2.244767427444458
+  },
+  {
+    "step": 586,
+    "loss": 2.256420135498047
+  },
+  {
+    "step": 587,
+    "loss": 1.9098711013793945
+  },
+  {
+    "step": 588,
+    "loss": 2.3755404949188232
+  },
+  {
+    "step": 589,
+    "loss": 2.6855807304382324
+  },
+  {
+    "step": 590,
+    "loss": 2.460073471069336
+  },
+  {
+    "step": 591,
+    "loss": 1.8837049007415771
+  },
+  {
+    "step": 592,
+    "loss": 2.081373929977417
+  },
+  {
+    "step": 593,
+    "loss": 2.1925439834594727
+  },
+  {
+    "step": 594,
+    "loss": 1.9564194679260254
+  },
+  {
+    "step": 595,
+    "loss": 2.205615758895874
+  },
+  {
+    "step": 596,
+    "loss": 2.478692054748535
+  },
+  {
+    "step": 597,
+    "loss": 2.0080032348632812
+  },
+  {
+    "step": 598,
+    "loss": 2.469820976257324
+  },
+  {
+    "step": 599,
+    "loss": 1.8856621980667114
+  },
+  {
+    "step": 600,
+    "loss": 1.6023671627044678
+  },
+  {
+    "step": 601,
+    "loss": 2.1479015350341797
+  },
+  {
+    "step": 602,
+    "loss": 2.306637763977051
+  },
+  {
+    "step": 603,
+    "loss": 2.2650392055511475
+  },
+  {
+    "step": 604,
+    "loss": 2.517967939376831
+  },
+  {
+    "step": 605,
+    "loss": 2.248115062713623
+  },
+  {
+    "step": 606,
+    "loss": 2.1899242401123047
+  },
+  {
+    "step": 607,
+    "loss": 2.3952863216400146
+  },
+  {
+    "step": 608,
+    "loss": 2.0763449668884277
+  },
+  {
+    "step": 609,
+    "loss": 2.3016607761383057
+  },
+  {
+    "step": 610,
+    "loss": 2.283357858657837
+  },
+  {
+    "step": 611,
+    "loss": 1.9581166505813599
+  },
+  {
+    "step": 612,
+    "loss": 2.4603822231292725
+  },
+  {
+    "step": 613,
+    "loss": 2.20369553565979
+  },
+  {
+    "step": 614,
+    "loss": 2.3975698947906494
+  },
+  {
+    "step": 615,
+    "loss": 2.3169641494750977
+  },
+  {
+    "step": 616,
+    "loss": 2.176382303237915
+  },
+  {
+    "step": 617,
+    "loss": 2.343982219696045
+  },
+  {
+    "step": 618,
+    "loss": 1.8763999938964844
+  },
+  {
+    "step": 619,
+    "loss": 2.083078384399414
+  },
+  {
+    "step": 620,
+    "loss": 1.902529239654541
+  },
+  {
+    "step": 621,
+    "loss": 2.0607450008392334
+  },
+  {
+    "step": 622,
+    "loss": 2.078382730484009
+  },
+  {
+    "step": 623,
+    "loss": 2.4523744583129883
+  },
+  {
+    "step": 624,
+    "loss": 2.1366493701934814
+  },
+  {
+    "step": 625,
+    "loss": 2.0985474586486816
+  },
+  {
+    "step": 626,
+    "loss": 2.5543007850646973
+  },
+  {
+    "step": 627,
+    "loss": 2.5027101039886475
+  },
+  {
+    "step": 628,
+    "loss": 2.8813798427581787
+  },
+  {
+    "step": 629,
+    "loss": 2.3890798091888428
+  },
+  {
+    "step": 630,
+    "loss": 1.9115877151489258
+  },
+  {
+    "step": 631,
+    "loss": 1.8879164457321167
+  },
+  {
+    "step": 632,
+    "loss": 2.415476083755493
+  },
+  {
+    "step": 633,
+    "loss": 2.3512673377990723
+  },
+  {
+    "step": 634,
+    "loss": 2.330165147781372
+  },
+  {
+    "step": 635,
+    "loss": 2.368882179260254
+  },
+  {
+    "step": 636,
+    "loss": 1.8848557472229004
+  },
+  {
+    "step": 637,
+    "loss": 2.5454752445220947
+  },
+  {
+    "step": 638,
+    "loss": 2.4458537101745605
+  },
+  {
+    "step": 639,
+    "loss": 2.3046398162841797
+  },
+  {
+    "step": 640,
+    "loss": 2.293977737426758
+  },
+  {
+    "step": 641,
+    "loss": 2.0221903324127197
+  },
+  {
+    "step": 642,
+    "loss": 1.678128719329834
+  },
+  {
+    "step": 643,
+    "loss": 2.55545711517334
+  },
+  {
+    "step": 644,
+    "loss": 2.0510752201080322
+  },
+  {
+    "step": 645,
+    "loss": 1.8273422718048096
+  },
+  {
+    "step": 646,
+    "loss": 2.4667928218841553
+  },
+  {
+    "step": 647,
+    "loss": 2.062159538269043
+  },
+  {
+    "step": 648,
+    "loss": 1.8708369731903076
+  },
+  {
+    "step": 649,
+    "loss": 1.8367586135864258
+  },
+  {
+    "step": 650,
+    "loss": 2.3304603099823
+  },
+  {
+    "step": 651,
+    "loss": 2.2214672565460205
+  },
+  {
+    "step": 652,
+    "loss": 2.156398057937622
+  },
+  {
+    "step": 653,
+    "loss": 1.802478551864624
+  },
+  {
+    "step": 654,
+    "loss": 2.1194801330566406
+  },
+  {
+    "step": 655,
+    "loss": 1.832618236541748
+  },
+  {
+    "step": 656,
+    "loss": 2.7607004642486572
+  },
+  {
+    "step": 657,
+    "loss": 1.5378355979919434
+  },
+  {
+    "step": 658,
+    "loss": 2.871095895767212
+  },
+  {
+    "step": 659,
+    "loss": 2.1772589683532715
+  },
+  {
+    "step": 660,
+    "loss": 2.378739356994629
+  },
+  {
+    "step": 661,
+    "loss": 1.9543912410736084
+  },
+  {
+    "step": 662,
+    "loss": 1.7253490686416626
+  },
+  {
+    "step": 663,
+    "loss": 2.3326663970947266
+  },
+  {
+    "step": 664,
+    "loss": 2.0331671237945557
+  },
+  {
+    "step": 665,
+    "loss": 2.8661580085754395
+  },
+  {
+    "step": 666,
+    "loss": 2.4486398696899414
+  },
+  {
+    "step": 667,
+    "loss": 2.2856521606445312
+  },
+  {
+    "step": 668,
+    "loss": 2.0959157943725586
+  },
+  {
+    "step": 669,
+    "loss": 1.828498363494873
+  },
+  {
+    "step": 670,
+    "loss": 2.580878257751465
+  },
+  {
+    "step": 671,
+    "loss": 2.111854076385498
+  },
+  {
+    "step": 672,
+    "loss": 2.3081603050231934
+  },
+  {
+    "step": 673,
+    "loss": 2.1535303592681885
+  },
+  {
+    "step": 674,
+    "loss": 2.4697322845458984
+  },
+  {
+    "step": 675,
+    "loss": 1.8599785566329956
+  },
+  {
+    "step": 676,
+    "loss": 2.156494140625
+  },
+  {
+    "step": 677,
+    "loss": 2.1473135948181152
+  },
+  {
+    "step": 678,
+    "loss": 1.5099081993103027
+  },
+  {
+    "step": 679,
+    "loss": 2.241101026535034
+  },
+  {
+    "step": 680,
+    "loss": 2.242471694946289
+  },
+  {
+    "step": 681,
+    "loss": 2.042699098587036
+  },
+  {
+    "step": 682,
+    "loss": 2.1123008728027344
+  },
+  {
+    "step": 683,
+    "loss": 2.017486095428467
+  },
+  {
+    "step": 684,
+    "loss": 2.0131030082702637
+  },
+  {
+    "step": 685,
+    "loss": 1.8260565996170044
+  },
+  {
+    "step": 686,
+    "loss": 1.8099098205566406
+  },
+  {
+    "step": 687,
+    "loss": 1.9493796825408936
+  },
+  {
+    "step": 688,
+    "loss": 1.7958455085754395
+  },
+  {
+    "step": 689,
+    "loss": 2.12551212310791
+  },
+  {
+    "step": 690,
+    "loss": 1.8851398229599
+  },
+  {
+    "step": 691,
+    "loss": 2.169377326965332
+  },
+  {
+    "step": 692,
+    "loss": 2.345599412918091
+  },
+  {
+    "step": 693,
+    "loss": 2.6371800899505615
+  },
+  {
+    "step": 694,
+    "loss": 2.054779291152954
+  },
+  {
+    "step": 695,
+    "loss": 1.8875532150268555
+  },
+  {
+    "step": 696,
+    "loss": 1.899794101715088
+  },
+  {
+    "step": 697,
+    "loss": 1.9862229824066162
+  },
+  {
+    "step": 698,
+    "loss": 2.069486141204834
+  },
+  {
+    "step": 699,
+    "loss": 2.5258560180664062
+  },
+  {
+    "step": 700,
+    "loss": 2.138136863708496
+  },
+  {
+    "step": 701,
+    "loss": 1.7787747383117676
+  },
+  {
+    "step": 702,
+    "loss": 2.3944613933563232
+  },
+  {
+    "step": 703,
+    "loss": 2.2432310581207275
+  },
+  {
+    "step": 704,
+    "loss": 2.23433780670166
+  },
+  {
+    "step": 705,
+    "loss": 2.3525569438934326
+  },
+  {
+    "step": 706,
+    "loss": 1.683660626411438
+  },
+  {
+    "step": 707,
+    "loss": 1.928051233291626
+  },
+  {
+    "step": 708,
+    "loss": 2.214559316635132
+  },
+  {
+    "step": 709,
+    "loss": 2.444596290588379
+  },
+  {
+    "step": 710,
+    "loss": 2.5637311935424805
+  },
+  {
+    "step": 711,
+    "loss": 2.094400405883789
+  },
+  {
+    "step": 712,
+    "loss": 2.1752796173095703
+  },
+  {
+    "step": 713,
+    "loss": 2.1646933555603027
+  },
+  {
+    "step": 714,
+    "loss": 2.3239498138427734
+  },
+  {
+    "step": 715,
+    "loss": 2.528728485107422
+  },
+  {
+    "step": 716,
+    "loss": 1.7567914724349976
+  },
+  {
+    "step": 717,
+    "loss": 1.8295328617095947
+  },
+  {
+    "step": 718,
+    "loss": 2.579841375350952
+  },
+  {
+    "step": 719,
+    "loss": 1.9017343521118164
+  },
+  {
+    "step": 720,
+    "loss": 2.507222890853882
+  },
+  {
+    "step": 721,
+    "loss": 2.127558946609497
+  },
+  {
+    "step": 722,
+    "loss": 2.1579627990722656
+  },
+  {
+    "step": 723,
+    "loss": 1.957413673400879
+  },
+  {
+    "step": 724,
+    "loss": 1.7173194885253906
+  },
+  {
+    "step": 725,
+    "loss": 1.8139698505401611
+  },
+  {
+    "step": 726,
+    "loss": 2.3142149448394775
+  },
+  {
+    "step": 727,
+    "loss": 2.4957833290100098
+  },
+  {
+    "step": 728,
+    "loss": 2.096470832824707
+  },
+  {
+    "step": 729,
+    "loss": 2.1391425132751465
+  },
+  {
+    "step": 730,
+    "loss": 2.0562264919281006
+  },
+  {
+    "step": 731,
+    "loss": 1.704184651374817
+  },
+  {
+    "step": 732,
+    "loss": 2.186826467514038
+  },
+  {
+    "step": 733,
+    "loss": 2.098335027694702
+  },
+  {
+    "step": 734,
+    "loss": 1.9378281831741333
+  },
+  {
+    "step": 735,
+    "loss": 1.9070260524749756
+  },
+  {
+    "step": 736,
+    "loss": 2.5211658477783203
+  },
+  {
+    "step": 737,
+    "loss": 1.8523521423339844
+  },
+  {
+    "step": 738,
+    "loss": 2.3545069694519043
+  },
+  {
+    "step": 739,
+    "loss": 1.715818166732788
+  },
+  {
+    "step": 740,
+    "loss": 1.9143247604370117
+  },
+  {
+    "step": 741,
+    "loss": 1.833640456199646
+  },
+  {
+    "step": 742,
+    "loss": 3.0031652450561523
+  },
+  {
+    "step": 743,
+    "loss": 1.5467760562896729
+  },
+  {
+    "step": 744,
+    "loss": 1.8195242881774902
+  },
+  {
+    "step": 745,
+    "loss": 1.785135269165039
+  },
+  {
+    "step": 746,
+    "loss": 2.0208964347839355
+  },
+  {
+    "step": 747,
+    "loss": 1.799407958984375
+  },
+  {
+    "step": 748,
+    "loss": 2.357728958129883
+  },
+  {
+    "step": 749,
+    "loss": 2.19157338142395
+  },
+  {
+    "step": 750,
+    "loss": 1.7044637203216553
+  },
+  {
+    "step": 751,
+    "loss": 2.1935999393463135
+  },
+  {
+    "step": 752,
+    "loss": 1.8042278289794922
+  },
+  {
+    "step": 753,
+    "loss": 2.340193271636963
+  },
+  {
+    "step": 754,
+    "loss": 1.8830957412719727
+  },
+  {
+    "step": 755,
+    "loss": 2.1636412143707275
+  },
+  {
+    "step": 756,
+    "loss": 2.573277473449707
+  },
+  {
+    "step": 757,
+    "loss": 2.185680866241455
+  },
+  {
+    "step": 758,
+    "loss": 2.005404472351074
+  },
+  {
+    "step": 759,
+    "loss": 2.041130304336548
+  },
+  {
+    "step": 760,
+    "loss": 2.069777250289917
+  },
+  {
+    "step": 761,
+    "loss": 1.942754864692688
+  },
+  {
+    "step": 762,
+    "loss": 2.273555278778076
+  },
+  {
+    "step": 763,
+    "loss": 1.8626675605773926
+  },
+  {
+    "step": 764,
+    "loss": 3.090045690536499
+  },
+  {
+    "step": 765,
+    "loss": 1.6991760730743408
+  },
+  {
+    "step": 766,
+    "loss": 1.9688661098480225
+  },
+  {
+    "step": 767,
+    "loss": 2.0735771656036377
+  },
+  {
+    "step": 768,
+    "loss": 1.9161018133163452
+  },
+  {
+    "step": 769,
+    "loss": 1.935367226600647
+  },
+  {
+    "step": 770,
+    "loss": 2.512490749359131
+  },
+  {
+    "step": 771,
+    "loss": 1.9653754234313965
+  },
+  {
+    "step": 772,
+    "loss": 1.9006037712097168
+  },
+  {
+    "step": 773,
+    "loss": 2.1403775215148926
+  },
+  {
+    "step": 774,
+    "loss": 1.9985692501068115
+  },
+  {
+    "step": 775,
+    "loss": 2.2423832416534424
+  },
+  {
+    "step": 776,
+    "loss": 2.272416114807129
+  },
+  {
+    "step": 777,
+    "loss": 1.6806894540786743
+  },
+  {
+    "step": 778,
+    "loss": 1.6961314678192139
+  },
+  {
+    "step": 779,
+    "loss": 2.070193290710449
+  },
+  {
+    "step": 780,
+    "loss": 2.30169415473938
+  },
+  {
+    "step": 781,
+    "loss": 1.498708963394165
+  },
+  {
+    "step": 782,
+    "loss": 2.672607898712158
+  },
+  {
+    "step": 783,
+    "loss": 2.000545024871826
+  },
+  {
+    "step": 784,
+    "loss": 2.1863222122192383
+  },
+  {
+    "step": 785,
+    "loss": 1.955033302307129
+  },
+  {
+    "step": 786,
+    "loss": 2.0105948448181152
+  },
+  {
+    "step": 787,
+    "loss": 2.3913588523864746
+  },
+  {
+    "step": 788,
+    "loss": 2.0229620933532715
+  },
+  {
+    "step": 789,
+    "loss": 2.0052614212036133
+  },
+  {
+    "step": 790,
+    "loss": 2.718447685241699
+  },
+  {
+    "step": 791,
+    "loss": 2.2307233810424805
+  },
+  {
+    "step": 792,
+    "loss": 2.1480185985565186
+  },
+  {
+    "step": 793,
+    "loss": 2.508347988128662
+  },
+  {
+    "step": 794,
+    "loss": 1.927217960357666
+  },
+  {
+    "step": 795,
+    "loss": 2.075195789337158
+  },
+  {
+    "step": 796,
+    "loss": 2.0269381999969482
+  },
+  {
+    "step": 797,
+    "loss": 1.7403855323791504
+  },
+  {
+    "step": 798,
+    "loss": 2.3028817176818848
+  },
+  {
+    "step": 799,
+    "loss": 2.2278242111206055
+  },
+  {
+    "step": 800,
+    "loss": 2.3212714195251465
+  },
+  {
+    "step": 801,
+    "loss": 2.1449832916259766
+  },
+  {
+    "step": 802,
+    "loss": 1.9621870517730713
+  },
+  {
+    "step": 803,
+    "loss": 2.1465606689453125
+  },
+  {
+    "step": 804,
+    "loss": 2.28023624420166
+  },
+  {
+    "step": 805,
+    "loss": 2.3931326866149902
+  },
+  {
+    "step": 806,
+    "loss": 1.9575012922286987
+  },
+  {
+    "step": 807,
+    "loss": 2.2189478874206543
+  },
+  {
+    "step": 808,
+    "loss": 2.313659906387329
+  },
+  {
+    "step": 809,
+    "loss": 1.785354495048523
+  },
+  {
+    "step": 810,
+    "loss": 2.360384941101074
+  },
+  {
+    "step": 811,
+    "loss": 1.5126988887786865
+  },
+  {
+    "step": 812,
+    "loss": 2.0700507164001465
+  },
+  {
+    "step": 813,
+    "loss": 2.364596128463745
+  },
+  {
+    "step": 814,
+    "loss": 1.7121427059173584
+  },
+  {
+    "step": 815,
+    "loss": 2.4048376083374023
+  },
+  {
+    "step": 816,
+    "loss": 2.1315109729766846
+  },
+  {
+    "step": 817,
+    "loss": 2.0960779190063477
+  },
+  {
+    "step": 818,
+    "loss": 2.1370670795440674
+  },
+  {
+    "step": 819,
+    "loss": 2.2387843132019043
+  },
+  {
+    "step": 820,
+    "loss": 2.6564176082611084
+  },
+  {
+    "step": 821,
+    "loss": 2.3899805545806885
+  },
+  {
+    "step": 822,
+    "loss": 1.8396912813186646
+  },
+  {
+    "step": 823,
+    "loss": 1.849434494972229
+  },
+  {
+    "step": 824,
+    "loss": 2.4848148822784424
+  },
+  {
+    "step": 825,
+    "loss": 2.300788640975952
+  },
+  {
+    "step": 826,
+    "loss": 1.6741901636123657
+  },
+  {
+    "step": 827,
+    "loss": 1.9888312816619873
+  },
+  {
+    "step": 828,
+    "loss": 2.1556782722473145
+  },
+  {
+    "step": 829,
+    "loss": 2.0511889457702637
+  },
+  {
+    "step": 830,
+    "loss": 2.110619068145752
+  },
+  {
+    "step": 831,
+    "loss": 1.8331103324890137
+  },
+  {
+    "step": 832,
+    "loss": 2.0915446281433105
+  },
+  {
+    "step": 833,
+    "loss": 1.3950932025909424
+  },
+  {
+    "step": 834,
+    "loss": 2.5744056701660156
+  },
+  {
+    "step": 835,
+    "loss": 2.351382255554199
+  },
+  {
+    "step": 836,
+    "loss": 2.207592248916626
+  },
+  {
+    "step": 837,
+    "loss": 2.0599913597106934
+  },
+  {
+    "step": 838,
+    "loss": 2.4455690383911133
+  },
+  {
+    "step": 839,
+    "loss": 2.037224769592285
+  },
+  {
+    "step": 840,
+    "loss": 2.1246533393859863
+  },
+  {
+    "step": 841,
+    "loss": 2.539519786834717
+  },
+  {
+    "step": 842,
+    "loss": 2.3089704513549805
+  },
+  {
+    "step": 843,
+    "loss": 1.8031303882598877
+  },
+  {
+    "step": 844,
+    "loss": 2.1725902557373047
+  },
+  {
+    "step": 845,
+    "loss": 1.9818854331970215
+  },
+  {
+    "step": 846,
+    "loss": 2.469529151916504
+  },
+  {
+    "step": 847,
+    "loss": 1.6893788576126099
+  },
+  {
+    "step": 848,
+    "loss": 1.9128904342651367
+  },
+  {
+    "step": 849,
+    "loss": 2.3154168128967285
+  },
+  {
+    "step": 850,
+    "loss": 1.864374041557312
+  },
+  {
+    "step": 851,
+    "loss": 2.139346122741699
+  },
+  {
+    "step": 852,
+    "loss": 1.9899048805236816
+  },
+  {
+    "step": 853,
+    "loss": 1.8425148725509644
+  },
+  {
+    "step": 854,
+    "loss": 2.196641445159912
+  },
+  {
+    "step": 855,
+    "loss": 2.084127187728882
+  },
+  {
+    "step": 856,
+    "loss": 2.589137077331543
+  },
+  {
+    "step": 857,
+    "loss": 1.9049241542816162
+  },
+  {
+    "step": 858,
+    "loss": 2.1088719367980957
+  },
+  {
+    "step": 859,
+    "loss": 1.8666791915893555
+  },
+  {
+    "step": 860,
+    "loss": 2.291273355484009
+  },
+  {
+    "step": 861,
+    "loss": 2.7720913887023926
+  },
+  {
+    "step": 862,
+    "loss": 2.0325875282287598
+  },
+  {
+    "step": 863,
+    "loss": 2.1262521743774414
+  },
+  {
+    "step": 864,
+    "loss": 2.1930947303771973
+  },
+  {
+    "step": 865,
+    "loss": 2.0483975410461426
+  },
+  {
+    "step": 866,
+    "loss": 2.1810200214385986
+  },
+  {
+    "step": 867,
+    "loss": 2.0406017303466797
+  },
+  {
+    "step": 868,
+    "loss": 1.9476172924041748
+  },
+  {
+    "step": 869,
+    "loss": 2.3409388065338135
+  },
+  {
+    "step": 870,
+    "loss": 2.2193973064422607
+  },
+  {
+    "step": 871,
+    "loss": 2.135343313217163
+  },
+  {
+    "step": 872,
+    "loss": 2.0815932750701904
+  },
+  {
+    "step": 873,
+    "loss": 2.2342116832733154
+  },
+  {
+    "step": 874,
+    "loss": 1.839571237564087
+  },
+  {
+    "step": 875,
+    "loss": 2.225438356399536
+  },
+  {
+    "step": 876,
+    "loss": 2.303863525390625
+  },
+  {
+    "step": 877,
+    "loss": 1.9251621961593628
+  },
+  {
+    "step": 878,
+    "loss": 2.1342930793762207
+  },
+  {
+    "step": 879,
+    "loss": 2.1234943866729736
+  },
+  {
+    "step": 880,
+    "loss": 2.137087821960449
+  },
+  {
+    "step": 881,
+    "loss": 1.9579668045043945
+  },
+  {
+    "step": 882,
+    "loss": 2.0309898853302
+  },
+  {
+    "step": 883,
+    "loss": 1.799654245376587
+  },
+  {
+    "step": 884,
+    "loss": 2.4847488403320312
+  },
+  {
+    "step": 885,
+    "loss": 2.6662018299102783
+  },
+  {
+    "step": 886,
+    "loss": 1.9268032312393188
+  },
+  {
+    "step": 887,
+    "loss": 2.043844223022461
+  },
+  {
+    "step": 888,
+    "loss": 2.038583278656006
+  },
+  {
+    "step": 889,
+    "loss": 2.393343925476074
+  },
+  {
+    "step": 890,
+    "loss": 2.155910015106201
+  },
+  {
+    "step": 891,
+    "loss": 2.047548532485962
+  },
+  {
+    "step": 892,
+    "loss": 2.4230449199676514
+  },
+  {
+    "step": 893,
+    "loss": 1.7401704788208008
+  },
+  {
+    "step": 894,
+    "loss": 2.0996835231781006
+  },
+  {
+    "step": 895,
+    "loss": 1.96319580078125
+  },
+  {
+    "step": 896,
+    "loss": 1.9458671808242798
+  },
+  {
+    "step": 897,
+    "loss": 2.0543224811553955
+  },
+  {
+    "step": 898,
+    "loss": 2.5579683780670166
+  },
+  {
+    "step": 899,
+    "loss": 2.102567434310913
+  },
+  {
+    "step": 900,
+    "loss": 1.797571063041687
+  },
+  {
+    "step": 901,
+    "loss": 2.612189292907715
+  },
+  {
+    "step": 902,
+    "loss": 1.8275275230407715
+  },
+  {
+    "step": 903,
+    "loss": 2.3661575317382812
+  },
+  {
+    "step": 904,
+    "loss": 2.412741184234619
+  },
+  {
+    "step": 905,
+    "loss": 1.7417240142822266
+  },
+  {
+    "step": 906,
+    "loss": 1.9783049821853638
+  },
+  {
+    "step": 907,
+    "loss": 3.020042896270752
+  },
+  {
+    "step": 908,
+    "loss": 2.264007568359375
+  },
+  {
+    "step": 909,
+    "loss": 2.6521012783050537
+  },
+  {
+    "step": 910,
+    "loss": 1.8450102806091309
+  },
+  {
+    "step": 911,
+    "loss": 2.0339226722717285
+  },
+  {
+    "step": 912,
+    "loss": 1.9741251468658447
+  },
+  {
+    "step": 913,
+    "loss": 2.638288736343384
+  },
+  {
+    "step": 914,
+    "loss": 2.128150463104248
+  },
+  {
+    "step": 915,
+    "loss": 2.2978391647338867
+  },
+  {
+    "step": 916,
+    "loss": 2.329345226287842
+  },
+  {
+    "step": 917,
+    "loss": 1.7687416076660156
+  },
+  {
+    "step": 918,
+    "loss": 1.7877784967422485
+  },
+  {
+    "step": 919,
+    "loss": 2.0238168239593506
+  },
+  {
+    "step": 920,
+    "loss": 2.0090293884277344
+  },
+  {
+    "step": 921,
+    "loss": 2.5923585891723633
+  },
+  {
+    "step": 922,
+    "loss": 2.2417752742767334
+  },
+  {
+    "step": 923,
+    "loss": 2.1116175651550293
+  },
+  {
+    "step": 924,
+    "loss": 1.826296091079712
+  },
+  {
+    "step": 925,
+    "loss": 1.888626217842102
+  },
+  {
+    "step": 926,
+    "loss": 1.9428234100341797
+  },
+  {
+    "step": 927,
+    "loss": 1.937356948852539
+  },
+  {
+    "step": 928,
+    "loss": 2.575594902038574
+  },
+  {
+    "step": 929,
+    "loss": 2.1356945037841797
+  },
+  {
+    "step": 930,
+    "loss": 2.0911166667938232
+  },
+  {
+    "step": 931,
+    "loss": 2.0979466438293457
+  },
+  {
+    "step": 932,
+    "loss": 2.5593225955963135
+  },
+  {
+    "step": 933,
+    "loss": 1.8475861549377441
+  },
+  {
+    "step": 934,
+    "loss": 1.9928505420684814
+  },
+  {
+    "step": 935,
+    "loss": 1.6050671339035034
+  },
+  {
+    "step": 936,
+    "loss": 2.3365602493286133
+  },
+  {
+    "step": 937,
+    "loss": 1.6974786520004272
+  },
+  {
+    "step": 938,
+    "loss": 2.5650625228881836
+  },
+  {
+    "step": 939,
+    "loss": 2.400352954864502
+  },
+  {
+    "step": 940,
+    "loss": 2.04494571685791
+  },
+  {
+    "step": 941,
+    "loss": 1.9382586479187012
+  },
+  {
+    "step": 942,
+    "loss": 2.147663116455078
+  },
+  {
+    "step": 943,
+    "loss": 1.8787052631378174
+  },
+  {
+    "step": 944,
+    "loss": 2.428316593170166
+  },
+  {
+    "step": 945,
+    "loss": 2.1108169555664062
+  },
+  {
+    "step": 946,
+    "loss": 1.6542246341705322
+  },
+  {
+    "step": 947,
+    "loss": 2.0597152709960938
+  },
+  {
+    "step": 948,
+    "loss": 2.3472602367401123
+  },
+  {
+    "step": 949,
+    "loss": 1.8338260650634766
+  },
+  {
+    "step": 950,
+    "loss": 1.7510230541229248
+  },
+  {
+    "step": 951,
+    "loss": 1.7876065969467163
+  },
+  {
+    "step": 952,
+    "loss": 2.431952714920044
+  },
+  {
+    "step": 953,
+    "loss": 1.9844751358032227
+  },
+  {
+    "step": 954,
+    "loss": 1.6035586595535278
+  },
+  {
+    "step": 955,
+    "loss": 2.2497174739837646
+  },
+  {
+    "step": 956,
+    "loss": 2.1397483348846436
+  },
+  {
+    "step": 957,
+    "loss": 2.418480396270752
+  },
+  {
+    "step": 958,
+    "loss": 2.0681729316711426
+  },
+  {
+    "step": 959,
+    "loss": 2.1492745876312256
+  },
+  {
+    "step": 960,
+    "loss": 2.292433023452759
+  },
+  {
+    "step": 961,
+    "loss": 1.993412733078003
+  },
+  {
+    "step": 962,
+    "loss": 1.9851655960083008
+  },
+  {
+    "step": 963,
+    "loss": 2.1893622875213623
+  },
+  {
+    "step": 964,
+    "loss": 2.258875608444214
+  },
+  {
+    "step": 965,
+    "loss": 2.1283998489379883
+  },
+  {
+    "step": 966,
+    "loss": 1.8134756088256836
+  },
+  {
+    "step": 967,
+    "loss": 1.8318551778793335
+  },
+  {
+    "step": 968,
+    "loss": 2.4340434074401855
+  },
+  {
+    "step": 969,
+    "loss": 1.9633982181549072
+  },
+  {
+    "step": 970,
+    "loss": 2.085658311843872
+  },
+  {
+    "step": 971,
+    "loss": 1.3723973035812378
+  },
+  {
+    "step": 972,
+    "loss": 3.041292667388916
+  },
+  {
+    "step": 973,
+    "loss": 2.3183510303497314
+  },
+  {
+    "step": 974,
+    "loss": 1.9262609481811523
+  },
+  {
+    "step": 975,
+    "loss": 2.1205694675445557
+  },
+  {
+    "step": 976,
+    "loss": 1.4980692863464355
+  },
+  {
+    "step": 977,
+    "loss": 2.5546534061431885
+  },
+  {
+    "step": 978,
+    "loss": 2.7014734745025635
+  },
+  {
+    "step": 979,
+    "loss": 1.676489233970642
+  },
+  {
+    "step": 980,
+    "loss": 1.5694572925567627
+  },
+  {
+    "step": 981,
+    "loss": 2.120713233947754
+  },
+  {
+    "step": 982,
+    "loss": 2.21758770942688
+  },
+  {
+    "step": 983,
+    "loss": 1.9932773113250732
+  },
+  {
+    "step": 984,
+    "loss": 2.172217607498169
+  },
+  {
+    "step": 985,
+    "loss": 1.980347990989685
+  },
+  {
+    "step": 986,
+    "loss": 2.1737043857574463
+  },
+  {
+    "step": 987,
+    "loss": 2.1837902069091797
+  },
+  {
+    "step": 988,
+    "loss": 1.7776117324829102
+  },
+  {
+    "step": 989,
+    "loss": 2.1192994117736816
+  },
+  {
+    "step": 990,
+    "loss": 1.7178640365600586
+  },
+  {
+    "step": 991,
+    "loss": 2.1105740070343018
+  },
+  {
+    "step": 992,
+    "loss": 1.9493019580841064
+  },
+  {
+    "step": 993,
+    "loss": 2.4504644870758057
+  },
+  {
+    "step": 994,
+    "loss": 2.0112550258636475
+  },
+  {
+    "step": 995,
+    "loss": 2.2686667442321777
+  },
+  {
+    "step": 996,
+    "loss": 2.4370152950286865
+  },
+  {
+    "step": 997,
+    "loss": 1.8057717084884644
+  },
+  {
+    "step": 998,
+    "loss": 1.7339015007019043
+  },
+  {
+    "step": 999,
+    "loss": 1.915513038635254
+  }
+]

lecture_6/verified_examples/diamond/config.json

@@ -0,0 +1,21 @@
+{
+  "method": "diamond",
+  "mode": "train-sample",
+  "train_steps": 1000,
+  "teacher_steps": 1000,
+  "finetune_steps": 0,
+  "sample_steps": 8,
+  "posterior_steps": 4,
+  "particles": 32,
+  "reward_strength": 1.0,
+  "ssfm_target": "official-code",
+  "batch_size": 64,
+  "samples": 128,
+  "width": 64,
+  "lr": 0.001,
+  "seed": 6270,
+  "threads": 1,
+  "device": "cpu",
+  "data": null,
+  "out": "outputs/verified-torch291/diamond"
+}