Upload 87 files

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Pranam Chatterjee
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

MIGRATION.md

@@ -0,0 +1,44 @@
+# Replacing the current Hugging Face repository
+
+This tree is a rebuild of `ChatterjeeLab/CIS6270` from the final lecture notes. Everything in it is
+new except the four files listed under "keep" below.
+
+## Keep from the current repository
+
+| File | Why |
+| --- | --- |
+| `lecture_2/flow_unet_mnist.pt` | The published five-epoch checkpoint. `lecture_2/flow_matching_mnist.py` still defines `FlowUNet` with the same state-dict keys, so the file loads unchanged. Move it to `lecture_2/`. |
+| `lecture_2/mnist_selected_examples.png` | The selected generations in the Lecture 2 guide. |
+| `lecture_2/mnist_selected_trajectories.png` | The noise-to-image trajectories in the same guide. |
+| `LICENSE` | Unchanged, MIT. |
+
+The overview figure at the top of the README is hosted on the Hugging Face CDN and is linked, not
+stored, so it carries over on its own.
+
+## Delete
+
+Everything else. The old tree held 67 files per lecture folder of which most were generated:
+
+- `run.py`, `run_all.py`, `lecture_core.py`, `common.py`, `diffusion.py`, `flows.py`, `examples/`.
+  One file per framework replaces the dispatcher, so a method is read in one place.
+- `verified_examples/` in lectures 4 to 7, about 300 JSON files, several over 200 KB. A seeded run
+  is now reproduced by running the script, and the numbers a default run gives are in each
+  lecture's README.
+- `lecture_6/EQUATION_READINGS.md` (138 KB) and `lecture_7/lecture7_ot_sbm_notes.md` (148 KB). The
+  book carries both.
+- `MATHEMATICS.md`, `SOURCES.md`, `SLIDE_CODE_MAP.md`, `source_manifest.json`,
+  `paper_manifest.json`, `lecture7_examples.ipynb`. The slide maps pointed at the decks as they
+  were before the October rewrites.
+- `lecture_4/data/dna_train.tsv` and `lecture_5/data/dna_train.tsv`, which were byte-identical
+  copies of each other. `cis6270.data.synthetic_dna` generates the same data in one place, and
+  `cis6270.data.write_dna_tsv` writes it out in the old format for anyone who wants the file.
+- `lecture_3/esm2_example.csv`, `GUIDANCE_NOTES.md`, and the per-lecture `requirements.txt` files.
+
+## What changed for a reader
+
+- There is a Lecture 1 now, covering the mathematical foundations chapter.
+- `pip install -e .` puts `cis6270` on the path; otherwise run from the repository root with
+  `PYTHONPATH=.`.
+- The shared flags are the same everywhere: `--steps`, `--batch-size`, `--lr`, `--width`,
+  `--length`, `--samples`, `--sample-steps`, `--seed`, `--device`, `--out`, `--quiet`.
+- `python -m pytest tests -q` runs 484 tests, which assert the numbers the notes work by hand.

README.md

@@ -0,0 +1,89 @@
+# CIS 6270 Course Code
+
+![CIS6270_Unified_Generative_Dynamics (1)](https://cdn-uploads.huggingface.co/production/uploads/64cd5b3f0494187a9e8b7c69/myMO3iMdYkquG6D5mhU4S.png)
+
+Code for CIS 6270, Discrete Generative Models, at the University of Pennsylvania. One file per
+framework, written to be read beside the lecture notes: the corruption, the loss and the sampler of
+each method are in that file, in the notes' own symbols, with a comment on almost every line.
+
+Everything runs offline on a laptop CPU in seconds. The one exception is
+`lecture_2/flow_matching_mnist.py`, which downloads MNIST, as its name says; `--dataset toy` runs it
+offline. Lecture 3 goes the other way: its scripts run offline and reach MNIST or ESM-2 only when
+you ask with `--dataset mnist` or `--dataset esm2`.
+
+```bash
+git clone https://huggingface.co/ChatterjeeLab/CIS6270
+cd CIS6270
+python3 -m venv .venv && source .venv/bin/activate
+python -m pip install -r requirements.txt
+
+python lecture_4/mdlm.py          # train masked diffusion on DNA and generate, about 20 seconds
+python -m pytest tests -q         # the notes' worked numbers, checked against the code
+```
+
+On Windows PowerShell, activate with `.venv\Scripts\Activate.ps1`.
+
+## The lectures
+
+| | Lecture | What the code builds | Guide |
+| --- | --- | --- | --- |
+| 1 | Mathematical Foundations | Every worked example of the chapter, printed beside the chapter's own number | [lecture_1](lecture_1/README.md) |
+| 2 | Flow Matching | Velocity fields and the flow map, the conditional path, the objective, MNIST | [lecture_2](lecture_2/README.md) |
+| 3 | Score Matching and Diffusion | The forward SDE, the three score objectives, DDPM, three samplers, guidance, latent diffusion on protein embeddings | [lecture_3](lecture_3/README.md) |
+| 4 | Discrete Diffusion | Jump processes, MDLM, uniform replacement, D3PM, SEDD, block diffusion, guidance, planners, search | [lecture_4](lecture_4/README.md) |
+| 5 | Discrete Flow Matching | The mixture path and its velocity, correctors, kinetic-optimal paths, Edit Flows, Dirichlet, Fisher and Gumbel geometries, guidance, rectification | [lecture_5](lecture_5/README.md) |
+| 6 | Flow Maps | The three axioms, the three residuals, MeanFlow, Shortcut, consistency, the discrete flow maps, stochastic, expanding and posterior maps | [lecture_6](lecture_6/README.md) |
+| 7 | Optimal Transport and Schrödinger Bridges | Kantorovich and duality, Sinkhorn, displacement interpolation, static and dynamic bridges, DSB, DSBM, SF2M, the discrete bridges, branched and entangled | [lecture_7](lecture_7/README.md) |
+
+## How a file is laid out
+
+Every script follows the shape of `lecture_2/flow_matching_mnist.py`: a docstring naming the
+framework and the chapter sections it implements, numbered `# %%` sections in the order the lecture
+develops them, and a `main()` that prints a report.
+
+```
+cis6270/           the scaffolding, read once: data, the four small networks, the training loop,
+                   the sample metrics. None of it is the subject of a lecture.
+lecture_1/ ... 7/  one file per framework, standalone and runnable
+tests/             the chapters' worked numbers, asserted against the code
+```
+
+Nothing in `lecture_N/` imports anything from another `lecture_N/`. To read a method, open its file.
+
+## The data
+
+| Data | Used by | Download |
+| --- | --- | --- |
+| Synthetic DNA, four letters, a planted ACGT motif, a GC label | Lectures 4 to 6 | no |
+| Four 2-D Gaussians, two moons, a checkerboard | Lectures 2, 3, 6, 7 | no |
+| Four-token phrases such as `red circle moves left` | Lecture 6 | no |
+| MNIST | Lectures 2, 3, with `--dataset mnist` | yes |
+| ESM-2 residue embeddings | Lecture 3, with `--dataset esm2` | yes |
+
+The DNA is synthetic and the two guidance objectives are explicit toy properties. They are there so
+that a training run finishes while you watch it, and so that the structure a sampler is supposed to
+recover is one you can check by eye. Nothing here is a benchmark or a biological predictor.
+
+## The shared flags
+
+`--steps`, `--batch-size`, `--lr`, `--width`, `--length`, `--samples`, `--sample-steps`, `--seed`,
+`--device`, `--out`, `--quiet`. Learn them once and they work everywhere. `--out` writes
+`checkpoint.pt`, `config.json`, `report.json`, `losses.json` and `samples.txt`.
+
+The defaults are a short budget, enough to see a method work in well under a minute. A method
+usually looks better with `--steps 1500 --width 128`.
+
+## Tests
+
+```bash
+python -m pytest tests -q
+```
+
+The suite asserts the numbers the notes work by hand: the posterior average of the reverse rates on
+one DNA base, the Sinkhorn iterates at `epsilon = 2`, the cocycle and Jacobian of the Gaussian flow,
+the conditional velocity of the mixture path. A number in the book and a number in the code are the
+same number, and the suite is what keeps them that way.
+
+## License
+
+[MIT](LICENSE). ESM-2 weights come from their original repository under its own terms.

cis6270/__init__.py

@@ -0,0 +1,20 @@
+"""Shared helpers for the CIS 6270 course code.
+
+Nothing in this package is the subject of a lecture. The data, the small networks, the training
+loop and the sample metrics live here so that each lecture file holds only the mathematics of its
+own framework. Read this package once; after that every lecture file stands on its own.
+"""
+
+import warnings
+
+# torch 2.14 warns about `torch.jit.script` from inside its own forward-mode autograd machinery,
+# the first time `torch.func.jvp` runs. Chapter 6 uses that function, and nothing here is
+# deprecated, so the message is filtered and the reader never sees it.
+warnings.filterwarnings(
+    "ignore", message=r"`torch\.jit\.script` is deprecated", category=FutureWarning,
+)
+
+from cis6270 import data, metrics, nets, runner
+
+__all__ = ["data", "metrics", "nets", "runner"]
+__version__ = "1.0.0"

cis6270/data.py

@@ -0,0 +1,317 @@
+"""Datasets for the course.
+
+Everything here runs offline in under a second, so a fresh clone and the test suite need no
+network. MNIST and ESM-2 embeddings download on first use and are reached only through the
+``--dataset`` flag of the scripts that offer them.
+
+The discrete lectures all use the same synthetic DNA, so a loss that improves in Lecture 4 and a
+loss that improves in Lecture 5 are measured on the same sequences.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+
+# %% 01. The DNA alphabet
+#
+# Chapters 4 and 5 run every worked example on DNA, so the vocabulary is four letters. The mask
+# token, when a lecture needs one, is appended as index 4 by that lecture, not here, because masked
+# diffusion has a mask and uniform diffusion does not.
+
+DNA = "ACGT"  # Token 0 is A, 1 is C, 2 is G, 3 is T.
+DNA_SIZE = len(DNA)  # K = 4, the alphabet size the chapters call |V|.
+MOTIF = "ACGT"  # The planted motif, used as the label in the guidance examples.
+
+
+def encode_dna(sequences: list[str]) -> torch.Tensor:
+    """Turn a list of strings over ACGT into a long tensor of shape [N, L]."""
+    index = {letter: i for i, letter in enumerate(DNA)}  # Letter to token id.
+    rows = [[index[letter] for letter in sequence] for sequence in sequences]
+    return torch.tensor(rows, dtype=torch.long)  # [N, L] token ids.
+
+
+def decode_dna(tokens: torch.Tensor) -> list[str]:
+    """Turn a long tensor of shape [N, L] back into strings over ACGT."""
+    return ["".join(DNA[int(i)] for i in row) for row in tokens]  # One string per row.
+
+
+# %% 02. Synthetic DNA with structure a small model can actually learn
+#
+# A uniform random sequence has nothing to learn, so the sampler would look the same before and
+# after training. These sequences carry two pieces of structure instead: a GC-rich background, and
+# the motif ACGT planted at a random position in a fraction of the sequences. The GC fraction then
+# gives a label for the conditional examples, and the motif gives a property to steer toward.
+
+
+def synthetic_dna(
+    num_sequences: int = 512,
+    length: int = 16,
+    motif_fraction: float = 0.75,
+    gc_bias: float = 0.6,
+    seed: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Draw synthetic DNA with a planted motif.
+
+    Returns the tokens as a long tensor of shape [N, L] and a binary label of shape [N] that is
+    one when the sequence has GC fraction at least 0.6. The label is a deterministic function of
+    the sequence, not a measurement of anything biological.
+    """
+    generator = torch.Generator().manual_seed(seed)  # Local generator; global seed untouched.
+
+    # Background letters: C and G each get gc_bias/2, A and T share the rest.
+    background = torch.tensor([
+        (1.0 - gc_bias) / 2.0,  # A
+        gc_bias / 2.0,  # C
+        gc_bias / 2.0,  # G
+        (1.0 - gc_bias) / 2.0,  # T
+    ])
+    tokens = torch.multinomial(
+        background.expand(num_sequences * length, DNA_SIZE),  # One draw per position.
+        num_samples=1,
+        replacement=True,
+        generator=generator,
+    ).view(num_sequences, length)  # [N, L] background sequence.
+
+    # Plant the motif in a fraction of the sequences, at a position drawn uniformly.
+    motif = encode_dna([MOTIF])[0]  # [4] token ids of ACGT.
+    carries_motif = torch.rand(num_sequences, generator=generator) < motif_fraction  # [N] bool.
+    # The motif starts at one of three places, not anywhere, which leaves it a pattern a
+    # small model can learn from a few hundred sequences while still varying its position. A start
+    # is kept only when the whole motif fits, so a short sequence simply has fewer of them.
+    room = length - len(MOTIF)  # Last index at which the motif still fits.
+    choices = torch.tensor([s for s in (2, 6, 10) if s <= room] or [0])
+    starts = choices[torch.randint(0, len(choices), (num_sequences,), generator=generator)]
+    for row in torch.nonzero(carries_motif).flatten().tolist():
+        start = int(starts[row])
+        tokens[row, start:start + len(MOTIF)] = motif  # Overwrite those positions.
+
+    gc = ((tokens == 1) | (tokens == 2)).to(torch.get_default_dtype()).mean(dim=1)  # [N] fraction of C and G.
+    labels = (gc >= 0.6).long()  # [N] the toy label used by the conditional examples.
+    return tokens, labels
+
+
+def dna_splits(
+    num_sequences: int = 512, length: int = 16, seed: int = 0, validation: float = 0.2,
+) -> dict[str, torch.Tensor]:
+    """The same synthetic DNA, split into a training and a validation set."""
+    tokens, labels = synthetic_dna(num_sequences, length, seed=seed)
+    cut = int(num_sequences * (1.0 - validation))  # Index where validation begins.
+    return {
+        "train_tokens": tokens[:cut], "train_labels": labels[:cut],
+        "valid_tokens": tokens[cut:], "valid_labels": labels[cut:],
+    }
+
+
+def write_dna_tsv(path: str | Path, tokens: torch.Tensor, labels: torch.Tensor) -> None:
+    """Write sequences and labels as a two-column TSV, the format ``--data`` reads."""
+    lines = ["sequence\tlabel"]
+    lines += [f"{s}\t{int(y)}" for s, y in zip(decode_dna(tokens), labels)]
+    Path(path).write_text("\n".join(lines) + "\n")
+
+
+def read_dna_tsv(path: str | Path) -> tuple[torch.Tensor, torch.Tensor]:
+    """Read a two-column TSV of sequences and labels."""
+    rows = Path(path).read_text().strip().split("\n")[1:]  # Skip the header line.
+    sequences = [row.split("\t")[0] for row in rows]
+    labels = [int(row.split("\t")[1]) for row in rows]
+    return encode_dna(sequences), torch.tensor(labels, dtype=torch.long)
+
+
+# %% 03. Two-dimensional densities
+#
+# The continuous lectures need a target we can plot and whose exact density we know. Four Gaussians
+# at the corners of a square give both: the modes are far enough apart that a sampler that collapses
+# is obvious, and the log density is a closed form the reports can quote.
+
+GAUSSIAN_CENTERS = torch.tensor([[2.0, 2.0], [2.0, -2.0], [-2.0, 2.0], [-2.0, -2.0]])
+GAUSSIAN_STD = 0.35  # Standard deviation of each component.
+
+
+def gaussian_mixture_2d(num_points: int = 2048, seed: int = 0) -> torch.Tensor:
+    """Draw from four equally weighted Gaussians at the corners of a square. Returns [N, 2]."""
+    generator = torch.Generator().manual_seed(seed)
+    which = torch.randint(0, len(GAUSSIAN_CENTERS), (num_points,), generator=generator)  # [N].
+    noise = torch.randn(num_points, 2, generator=generator) * GAUSSIAN_STD  # [N, 2].
+    return GAUSSIAN_CENTERS[which] + noise  # [N, 2] points.
+
+
+def gaussian_mixture_log_density(points: torch.Tensor) -> torch.Tensor:
+    """Exact log density of the four-Gaussian mixture at each row of ``points`` ([N, 2] -> [N])."""
+    diff = points[:, None, :] - GAUSSIAN_CENTERS.to(points)  # [N, 4, 2] offset to each center.
+    squared = (diff ** 2).sum(dim=-1)  # [N, 4] squared distance to each center.
+    log_component = (
+        -0.5 * squared / GAUSSIAN_STD ** 2  # Gaussian exponent.
+        - torch.log(torch.tensor(2.0 * torch.pi * GAUSSIAN_STD ** 2))  # Normalizer in 2-D.
+    )  # [N, 4].
+    weight = torch.log(torch.tensor(1.0 / len(GAUSSIAN_CENTERS)))  # Equal mixture weights.
+    return torch.logsumexp(log_component + weight, dim=-1)  # [N] mixture log density.
+
+
+def two_moons(num_points: int = 2048, noise: float = 0.08, seed: int = 0) -> torch.Tensor:
+    """Two interleaving half circles, the standard test for a sampler that must bend. [N, 2]."""
+    generator = torch.Generator().manual_seed(seed)
+    half = num_points // 2
+    angle = torch.rand(half, generator=generator) * torch.pi  # [half] angle along each arc.
+    upper = torch.stack([torch.cos(angle), torch.sin(angle)], dim=1)  # Upper half circle.
+    lower = torch.stack([1.0 - torch.cos(angle), 0.5 - torch.sin(angle)], dim=1)  # Lower arc.
+    points = torch.cat([upper, lower], dim=0)  # [2*half, 2].
+    points = points + noise * torch.randn(points.shape, generator=generator)  # Blur the arcs.
+    return (points - points.mean(0)) / points.std(0)  # Standardize to roughly unit scale.
+
+
+def checkerboard(num_points: int = 2048, seed: int = 0) -> torch.Tensor:
+    """A checkerboard of eight squares, a density with disconnected support. [N, 2]."""
+    generator = torch.Generator().manual_seed(seed)
+    x = torch.rand(num_points, generator=generator) * 4.0 - 2.0  # [N] in [-2, 2].
+    y = torch.rand(num_points, generator=generator) - torch.randint(
+        0, 2, (num_points,), generator=generator,
+    ).to(torch.get_default_dtype()) * 2.0  # [N] in [-2, 1] before the shift below.
+    y = y + torch.floor(x) % 2.0  # Shift alternate columns so the squares interleave.
+    return torch.stack([x, y * 2.0], dim=1)  # [N, 2] points.
+
+
+# %% 04. Short phrases, for the discrete examples whose tokens are words
+#
+# Chapter 6 runs its categorical flow maps on four-token phrases. The grammar is small enough that
+# a reader can tell at a glance whether a generated phrase is in the language.
+
+_SUBJECTS = ["red", "blue", "green", "small"]
+_OBJECTS = ["circle", "square", "arrow", "dot"]
+_VERBS = ["moves", "turns", "drifts", "jumps"]
+_DIRECTIONS = ["left", "right", "up", "down"]
+
+
+def phrase_vocabulary() -> list[str]:
+    """The token vocabulary of the phrase language, in a fixed order."""
+    return _SUBJECTS + _OBJECTS + _VERBS + _DIRECTIONS  # 16 tokens, four per slot.
+
+
+def phrases(num_phrases: int = 512, seed: int = 0) -> tuple[torch.Tensor, list[str]]:
+    """Draw four-token phrases such as ``red circle moves left``. Returns [N, 4] and the vocabulary.
+
+    Each slot draws from its own quarter of the vocabulary, so a correct phrase is one whose
+    position i holds a token from group i. That makes the training support easy to check.
+    """
+    generator = torch.Generator().manual_seed(seed)
+    vocabulary = phrase_vocabulary()
+    groups = [_SUBJECTS, _OBJECTS, _VERBS, _DIRECTIONS]
+    columns = []
+    for slot, group in enumerate(groups):
+        offset = slot * len(group)  # Where this slot's tokens start in the vocabulary.
+        picks = torch.randint(0, len(group), (num_phrases,), generator=generator)
+        columns.append(picks + offset)  # [N] token ids inside the full vocabulary.
+    return torch.stack(columns, dim=1), vocabulary  # [N, 4] and the 16 token strings.
+
+
+# %% 05. MNIST and ESM-2, the two datasets that download
+#
+# Both are optional. A script reaches them only when the reader passes --dataset mnist or
+# --dataset esm2, and every such script also runs on the offline data above.
+
+
+def mnist_loader(batch_size: int = 128, root: str | Path = "data", train: bool = True):
+    """MNIST as a DataLoader of images normalized to [-1, 1]. Downloads on first use."""
+    from torch.utils.data import DataLoader  # Imported here so the module loads without torchvision.
+    from torchvision import datasets
+    from torchvision.transforms import v2
+
+    transform = v2.Compose([
+        v2.ToImage(),  # Channel-first image: [1, 28, 28].
+        v2.ToDtype(torch.float32, scale=True),  # Integer pixels 0-255 to floats in [0, 1].
+        v2.Normalize(mean=(0.5,), std=(0.5,)),  # Then [0, 1] to [-1, 1].
+    ])
+    dataset = datasets.MNIST(root=root, train=train, download=True, transform=transform)
+    return DataLoader(dataset, batch_size=batch_size, shuffle=train, num_workers=0)
+
+
+def toy_images(num_images: int = 256, size: int = 28, seed: int = 0) -> torch.Tensor:
+    """Images shaped like MNIST, drawn from a two-mode distribution of bars and blobs.
+
+    These stand in for MNIST wherever a test has to run without a download. They have the same
+    shape and the same [-1, 1] range, so the identical training and sampling code paths run.
+    """
+    generator = torch.Generator().manual_seed(seed)
+    images = torch.full((num_images, 1, size, size), -1.0)  # Background at the low end of [-1, 1].
+    grid = torch.arange(size)
+    for i in range(num_images):
+        if i % 2 == 0:  # A vertical bar.
+            column = int(torch.randint(6, size - 6, (1,), generator=generator))
+            images[i, 0, 4:size - 4, column - 2:column + 2] = 1.0
+        else:  # A round blob.
+            cy = int(torch.randint(8, size - 8, (1,), generator=generator))
+            cx = int(torch.randint(8, size - 8, (1,), generator=generator))
+            dy = (grid[:, None] - cy) ** 2
+            dx = (grid[None, :] - cx) ** 2
+            images[i, 0][(dy + dx) <= 25] = 1.0
+    images = images + 0.05 * torch.randn(images.shape, generator=generator)  # Light pixel noise.
+    return images.clamp(-1.0, 1.0)  # [N, 1, size, size] in the training range.
+
+
+PEPTIDES = [
+    "KWKLFKKIGAVLKVL", "GIGKFLHSAKKFGKA", "FLPIIAKLLSGLL", "ILPWKWPWWPWRR",
+    "KRWWKWWRR", "GLLSKLWKKVFKAFKKFLKK", "ACDEFGHIKLMNPQR", "RRWWRRWRR",
+    "LLGDFFRKSKEKIGKE", "KIAKVALKAL", "VRRFPWWWPFLRR", "GWLKKIGKKIERVGQ",
+]  # Short antimicrobial-peptide-like sequences, used when ESM-2 is unavailable.
+
+
+def esm2_embeddings(
+    sequences: list[str] | None = None,
+    model_name: str = "facebook/esm2_t33_650M_UR50D",
+    max_length: int = 32,
+    device: torch.device | str = "cpu",
+):
+    """Residue embeddings from a frozen ESM-2 encoder. Downloads the checkpoint on first use.
+
+    Returns the embeddings [N, L, D], the attention mask [N, L], the tokenizer and the masked
+    language-model head, which Lecture 3 uses to decode generated latents back to amino acids.
+    """
+    from transformers import AutoModel, AutoModelForMaskedLM, AutoTokenizer
+
+    sequences = sequences or PEPTIDES
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    encoder = AutoModel.from_pretrained(model_name).to(device).eval()
+    head = AutoModelForMaskedLM.from_pretrained(model_name).lm_head.to(device).eval()
+    for parameter in encoder.parameters():
+        parameter.requires_grad = False  # The encoder is frozen; only the denoiser trains.
+
+    batch = tokenizer(
+        sequences, padding="max_length", truncation=True,
+        max_length=max_length, return_tensors="pt",
+    ).to(device)
+    with torch.no_grad():
+        embeddings = encoder(**batch).last_hidden_state  # [N, L, D] residue embeddings.
+    return embeddings, batch["attention_mask"], tokenizer, head
+
+
+def toy_embeddings(
+    num_sequences: int = 128, length: int = 16, width: int = 32, seed: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Embeddings shaped like ESM-2 output, for the offline path of the latent examples.
+
+    Each sequence is a smooth low-frequency curve along the residue axis plus noise, so neighbouring
+    positions are correlated the way real residue embeddings are. Returns [N, L, D] and a mask.
+    """
+    generator = torch.Generator().manual_seed(seed)
+    phase = torch.rand(num_sequences, 1, width, generator=generator) * 2.0 * torch.pi
+    frequency = torch.rand(num_sequences, 1, width, generator=generator) * 0.6 + 0.1
+    position = torch.arange(length).view(1, length, 1).to(torch.get_default_dtype())  # [1, L, 1].
+    signal = torch.sin(frequency * position + phase)  # [N, L, D] smooth along the residue axis.
+    noise = 0.2 * torch.randn(num_sequences, length, width, generator=generator)
+    mask = torch.ones(num_sequences, length)  # No padding in the toy set.
+    return signal + noise, mask
+
+
+def standardize(embeddings: torch.Tensor, mask: torch.Tensor):
+    """Center and scale each embedding coordinate over the valid residues.
+
+    Returns the standardized embeddings and the mean and standard deviation, which the decoder
+    needs in order to undo the transformation before reading amino acids.
+    """
+    valid = mask.bool().unsqueeze(-1)  # [N, L, 1] true at a real residue.
+    count = valid.sum()  # Number of valid residue-coordinate entries.
+    mean = (embeddings * valid).sum(dim=(0, 1)) / count  # [D] per-coordinate mean.
+    variance = ((embeddings - mean) ** 2 * valid).sum(dim=(0, 1)) / count  # [D].
+    std = variance.sqrt().clamp(min=1e-6)  # [D], clamped so a dead coordinate cannot divide by zero.
+    return (embeddings - mean) / std, mean, std

cis6270/metrics.py

@@ -0,0 +1,100 @@
+"""Measurements of a batch of generated samples.
+
+None of these is a benchmark. They are small, exactly computable quantities that tell a reader
+whether a sampler learned the structure planted in the data: the right letter frequencies, the
+planted motif, the right number of modes. A number here is a diagnostic, not a score.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from cis6270.data import GAUSSIAN_CENTERS, MOTIF, decode_dna
+
+
+# %% 01. Measurements on DNA
+
+
+def gc_fraction(tokens: torch.Tensor) -> float:
+    """Fraction of positions holding C or G, averaged over the batch. The training data sit at 0.577."""
+    return float(((tokens == 1) | (tokens == 2)).to(torch.get_default_dtype()).mean())
+
+
+def motif_fraction(tokens: torch.Tensor, motif: str = MOTIF) -> float:
+    """Fraction of sequences that contain the planted motif anywhere."""
+    sequences = decode_dna(tokens)
+    return sum(motif in sequence for sequence in sequences) / max(len(sequences), 1)
+
+
+def unique_fraction(tokens: torch.Tensor) -> float:
+    """Fraction of generated sequences that are distinct, which catches a collapsed sampler.
+
+    The comparison is on the token ids themselves, so this works for any vocabulary, not only the
+    four DNA letters.
+    """
+    rows = [tuple(row.tolist()) for row in tokens]  # One tuple per sequence.
+    return len(set(rows)) / max(len(rows), 1)
+
+
+def token_entropy(tokens: torch.Tensor, vocab_size: int) -> float:
+    """Entropy in nats of the letter frequencies, at most log K. Four letters give log 4 = 1.386."""
+    counts = torch.bincount(tokens.flatten(), minlength=vocab_size).to(torch.get_default_dtype())
+    probabilities = counts / counts.sum().clamp(min=1)
+    nonzero = probabilities[probabilities > 0]
+    return float(-(nonzero * nonzero.log()).sum())
+
+
+def support_fraction(tokens: torch.Tensor, training: torch.Tensor) -> float:
+    """Fraction of generated sequences that also appear in the training set."""
+    seen = {tuple(row.tolist()) for row in training}
+    return sum(tuple(row.tolist()) in seen for row in tokens) / max(len(tokens), 1)
+
+
+# %% 02. Measurements on two-dimensional samples
+
+
+def mode_report(points: torch.Tensor, centers: torch.Tensor | None = None) -> dict:
+    """Assign each point to its nearest mode and report the distance and the coverage.
+
+    ``nearest_distance`` near the data's own standard deviation means the points land on the modes.
+    ``mode_entropy`` near log 4 = 1.386 means all four modes are visited about equally.
+    """
+    centers = GAUSSIAN_CENTERS.to(points) if centers is None else centers.to(points)
+    distance = torch.cdist(points, centers)  # [N, M] distance to every center.
+    nearest, which = distance.min(dim=1)  # [N] distance and index of the closest mode.
+    counts = torch.bincount(which, minlength=len(centers)).to(torch.get_default_dtype())
+    fractions = counts / counts.sum().clamp(min=1)
+    nonzero = fractions[fractions > 0]
+    return {
+        "nearest_distance": float(nearest.mean()),
+        "modes_found": int((counts > 0).sum()),
+        "mode_entropy": float(-(nonzero * nonzero.log()).sum()),
+        "mode_fractions": [round(float(f), 3) for f in fractions],
+    }
+
+
+def mean_log_density(points: torch.Tensor) -> float:
+    """Average exact log density of the four-Gaussian target at the generated points."""
+    from cis6270.data import gaussian_mixture_log_density
+
+    return float(gaussian_mixture_log_density(points).mean())
+
+
+# %% 03. Measurements on a learned distribution over a small state space
+#
+# Several chapters work on a state space small enough to enumerate, where the exact answer is
+# available. Comparing an empirical histogram with that exact answer is the sharpest check the
+# course has.
+
+
+def total_variation(p: torch.Tensor, q: torch.Tensor) -> float:
+    """Half the sum of absolute differences between two distributions over the same states."""
+    return float(0.5 * (p - q).abs().sum())
+
+
+def empirical_distribution(tokens: torch.Tensor, vocab_size: int, length: int) -> torch.Tensor:
+    """Histogram of whole sequences over the K**L states, for a state space small enough to list."""
+    weights = vocab_size ** torch.arange(length - 1, -1, -1, device=tokens.device)  # Place values.
+    indices = (tokens * weights).sum(dim=1)  # [N] index of each sequence in the full state space.
+    counts = torch.bincount(indices, minlength=vocab_size ** length).to(torch.get_default_dtype())
+    return counts / counts.sum().clamp(min=1)  # [K**L] empirical distribution.

cis6270/nets.py

@@ -0,0 +1,231 @@
+"""The small networks every lecture reuses.
+
+Four networks cover the whole course. None of them is the subject of a lecture: a lecture changes
+the loss and the sampler, and keeps the network fixed, so that any difference in the samples comes
+from the mathematics and not from the architecture.
+
+| Network            | Input                      | Output                | Used by |
+| ------------------ | -------------------------- | --------------------- | ------- |
+| `TimeEmbedding`    | a time or step index       | a vector of width D   | all     |
+| `SequenceDenoiser` | tokens [B, L] and a time   | logits [B, L, K]      | 4, 5, 6 |
+| `VelocityMLP`      | a point [B, D] and a time  | a vector [B, D]       | 2, 3, 6, 7 |
+| `TwoTimeMLP`       | a point, a start, an end   | a vector [B, D]       | 6 |
+| `SmallUNet`        | an image [B, 1, 28, 28]    | an image of the same  | 2, 3 |
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+
+
+# %% 01. Embedding a continuous time
+#
+# Every network in the course takes a time as well as a state. A raw scalar in [0, 1] is a poor
+# input to a linear layer, so it is first spread over many frequencies, exactly as in Section 3.28
+# of the notes, and then passed through a small MLP.
+
+
+def sinusoidal_embedding(t: torch.Tensor, width: int) -> torch.Tensor:
+    """Spread a time [B] over ``width`` sine and cosine frequencies. Returns [B, width]."""
+    half = width // 2  # Half the coordinates are sines and half cosines.
+    frequencies = torch.exp(
+        -math.log(10000.0) * torch.arange(half, device=t.device).to(torch.get_default_dtype()) / half
+    )  # [half] geometrically spaced frequencies.
+    angles = t.to(torch.get_default_dtype())[:, None] * frequencies[None, :]  # [B, half].
+    return torch.cat([torch.sin(angles), torch.cos(angles)], dim=1)  # [B, width].
+
+
+class TimeEmbedding(nn.Module):
+    """Sinusoidal features of the time, followed by a two-layer MLP."""
+
+    def __init__(self, width: int = 128, scale: float = 1000.0):
+        super().__init__()
+        self.width = width  # Width of the embedding the rest of the network receives.
+        self.scale = scale  # Times in [0, 1] are multiplied by this before the sines.
+        self.mlp = nn.Sequential(
+            nn.Linear(width, width), nn.SiLU(), nn.Linear(width, width),
+        )
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        """t: [B] -> [B, width]."""
+        return self.mlp(sinusoidal_embedding(t * self.scale, self.width))  # [B, width].
+
+
+# %% 02. A denoiser over token sequences
+#
+# Lectures 4, 5 and 6 all train the same shape of network: it reads a corrupted sequence and a time,
+# and produces a distribution over the vocabulary at every position. The notes call it the denoiser
+# d_theta(. | z, t). Attention lets the prediction at one position depend on every other position,
+# which is what makes the factorized reverse step of Chapter 4 worth taking.
+
+
+class SequenceDenoiser(nn.Module):
+    """A small transformer that maps a token sequence and a time to per-position logits."""
+
+    def __init__(
+        self,
+        vocab_size: int,
+        length: int,
+        width: int = 128,
+        depth: int = 2,
+        heads: int = 4,
+        num_classes: int = 0,
+    ):
+        super().__init__()
+        self.vocab_size = vocab_size  # K, including a mask token when the lecture adds one.
+        self.length = length  # L, the number of positions.
+        self.width = width  # D, the model width.
+        self.token_embedding = nn.Embedding(vocab_size, width)  # One vector per token id.
+        self.position_embedding = nn.Embedding(length, width)  # One vector per position.
+        self.time_embedding = TimeEmbedding(width)  # Shared across positions.
+        # A label embedding, used only by the conditional examples. Index num_classes is the null
+        # label that classifier-free guidance drops to.
+        self.label_embedding = (
+            nn.Embedding(num_classes + 1, width) if num_classes else None
+        )
+        layer = nn.TransformerEncoderLayer(
+            d_model=width, nhead=heads, dim_feedforward=4 * width,
+            batch_first=True, dropout=0.0, norm_first=True,
+        )
+        self.transformer = nn.TransformerEncoder(
+            layer, num_layers=depth, enable_nested_tensor=False,
+        )
+        self.head = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, vocab_size))
+
+    def forward(
+        self, tokens: torch.Tensor, t: torch.Tensor, labels: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """tokens: [B, L], t: [B], labels: [B] or None -> logits [B, L, K]."""
+        positions = torch.arange(tokens.shape[1], device=tokens.device)  # [L].
+        h = (
+            self.token_embedding(tokens)  # [B, L, D] what each position currently holds.
+            + self.position_embedding(positions)[None, :, :]  # [1, L, D] where it sits.
+            + self.time_embedding(t)[:, None, :]  # [B, 1, D] how far along the path we are.
+        )
+        if labels is not None and self.label_embedding is not None:
+            h = h + self.label_embedding(labels)[:, None, :]  # [B, 1, D] the condition.
+        h = self.transformer(h)  # [B, L, D] every position now sees every other.
+        return self.head(h)  # [B, L, K] logits over the vocabulary at each position.
+
+
+# %% 03. A velocity network on a vector state
+#
+# Lectures 2, 3, 6 and 7 learn a field on R^d: a velocity, a score, a noise prediction or a drift.
+# The same MLP serves for all of them, because the only difference is what the loss regresses it
+# onto.
+
+
+class VelocityMLP(nn.Module):
+    """An MLP from a point and a time to a vector of the same width as the point."""
+
+    def __init__(self, dim: int = 2, width: int = 128, depth: int = 3, num_classes: int = 0):
+        super().__init__()
+        self.dim = dim  # d, the width of the state.
+        self.time_embedding = TimeEmbedding(width, scale=1.0)  # Times here already live in [0, 1].
+        self.label_embedding = (
+            nn.Embedding(num_classes + 1, width) if num_classes else None
+        )
+        self.input = nn.Linear(dim, width)  # Lift the state to the hidden width.
+        layers: list[nn.Module] = []
+        for _ in range(depth):
+            layers += [nn.Linear(width, width), nn.SiLU()]
+        self.body = nn.Sequential(*layers)
+        self.output = nn.Linear(width, dim)  # Back down to one vector per coordinate.
+
+    def forward(
+        self, x: torch.Tensor, t: torch.Tensor, labels: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """x: [B, d], t: [B], labels: [B] or None -> [B, d]."""
+        h = self.input(x) + self.time_embedding(t)  # [B, width] state and time together.
+        if labels is not None and self.label_embedding is not None:
+            h = h + self.label_embedding(labels)  # [B, width] the condition.
+        return self.output(self.body(h))  # [B, d].
+
+
+# %% 04. A two-time network, for the flow maps of Lecture 6
+#
+# A flow map carries a state from a start time s to an end time t, so its network takes both. The
+# identity axiom X_{s,s} = id is free if the map is written as the state plus the elapsed time times
+# a network output, which is what `flow_map` below does.
+
+
+class TwoTimeMLP(nn.Module):
+    """An MLP from a point, a start time and an end time to a vector."""
+
+    def __init__(self, dim: int = 2, width: int = 128, depth: int = 3):
+        super().__init__()
+        self.dim = dim
+        self.start_embedding = TimeEmbedding(width, scale=1.0)  # Embedding of s.
+        self.end_embedding = TimeEmbedding(width, scale=1.0)  # Embedding of t.
+        self.input = nn.Linear(dim, width)
+        layers: list[nn.Module] = []
+        for _ in range(depth):
+            layers += [nn.Linear(width, width), nn.SiLU()]
+        self.body = nn.Sequential(*layers)
+        self.output = nn.Linear(width, dim)
+
+    def forward(self, x: torch.Tensor, s: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        """The average velocity v_{s,t}(x). x: [B, d], s: [B], t: [B] -> [B, d]."""
+        h = self.input(x) + self.start_embedding(s) + self.end_embedding(t)  # [B, width].
+        return self.output(self.body(h))  # [B, d].
+
+    def flow_map(self, x: torch.Tensor, s: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        """X_{s,t}(x) = x + (t - s) v_{s,t}(x), so the identity axiom holds by construction."""
+        elapsed = (t - s)[:, None]  # [B, 1] elapsed time, broadcast over coordinates.
+        return x + elapsed * self.forward(x, s, t)  # [B, d].
+
+
+# %% 05. A small U-Net for 28-by-28 images
+#
+# Lectures 2 and 3 generate images. The network is deliberately small, about half a million
+# parameters, so that a few epochs on a laptop CPU already produce recognizable digits.
+
+
+def _conv_block(in_channels: int, out_channels: int) -> nn.Sequential:
+    """Two 3-by-3 convolutions with a SiLU after each, keeping the height and width."""
+    return nn.Sequential(
+        nn.Conv2d(in_channels, out_channels, 3, padding=1), nn.SiLU(),
+        nn.Conv2d(out_channels, out_channels, 3, padding=1), nn.SiLU(),
+    )
+
+
+class SmallUNet(nn.Module):
+    """A three-scale U-Net that maps an image and a time to an image of the same shape.
+
+    The time enters as an extra constant channel, which is the simplest conditioning that works and
+    is what the Lecture 2 script uses.
+    """
+
+    def __init__(self, base_channels: int = 32, in_channels: int = 1):
+        super().__init__()
+        c = base_channels  # Widths are c, 2c, 4c.
+        self.encoder1 = _conv_block(in_channels + 1, c)  # Image plus the time channel.
+        self.encoder2 = _conv_block(c, 2 * c)
+        self.middle = _conv_block(2 * c, 4 * c)
+        self.decoder2 = _conv_block(6 * c, 2 * c)  # Upsampled 4c joined with the 2c skip.
+        self.decoder1 = _conv_block(3 * c, c)  # Upsampled 2c joined with the c skip.
+        self.output = nn.Conv2d(c, in_channels, 1)  # One prediction per pixel.
+        self.pool = nn.MaxPool2d(2)  # Halve the height and width.
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        """x: [B, 1, 28, 28], t: [B] -> [B, 1, 28, 28]."""
+        import torch.nn.functional as F  # Local import keeps the module header short.
+
+        time_channel = t[:, None, None, None].expand(x.shape[0], 1, *x.shape[2:])  # [B, 1, H, W].
+        h = torch.cat([x, time_channel], dim=1)  # [B, 2, H, W].
+        skip1 = self.encoder1(h)  # [B, c, 28, 28].
+        skip2 = self.encoder2(self.pool(skip1))  # [B, 2c, 14, 14].
+        h = self.middle(self.pool(skip2))  # [B, 4c, 7, 7] the compressed representation.
+        h = F.interpolate(h, size=skip2.shape[-2:], mode="nearest")  # 7 to 14.
+        h = self.decoder2(torch.cat([h, skip2], dim=1))  # Rejoin the 14-by-14 detail.
+        h = F.interpolate(h, size=skip1.shape[-2:], mode="nearest")  # 14 to 28.
+        h = self.decoder1(torch.cat([h, skip1], dim=1))  # Rejoin the 28-by-28 detail.
+        return self.output(h)  # [B, 1, 28, 28].
+
+
+def count_parameters(model: nn.Module) -> int:
+    """Number of trainable weights and biases, printed by every script at startup."""
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)

cis6270/runner.py

@@ -0,0 +1,173 @@
+"""Seeding, device selection, the training loop, checkpoints and the shared command line.
+
+Every lecture script ends in the same four lines: parse the arguments, build the model, call
+``train`` with a step function, and print a report. The step function is the only part that differs
+between frameworks, and it is the part the lecture is about.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import random
+import time
+from pathlib import Path
+from typing import Callable
+
+import torch
+
+
+# %% 01. Reproducibility and hardware
+
+
+def set_seed(seed: int) -> None:
+    """Seed Python, NumPy through torch, and torch itself, so a run repeats exactly."""
+    random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def pick_device(requested: str = "auto") -> torch.device:
+    """Choose CUDA, then Apple MPS, then CPU, unless the reader names one."""
+    if requested != "auto":
+        return torch.device(requested)
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+# %% 02. The command line every script shares
+#
+# A reader who has learned the flags once knows them for all fifty scripts. A script adds its own
+# flags on top of these with `parser.add_argument`.
+
+
+def common_parser(description: str) -> argparse.ArgumentParser:
+    """The flags shared by every script in the course."""
+    parser = argparse.ArgumentParser(
+        description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("--steps", type=int, default=600, help="optimizer steps")
+    parser.add_argument("--batch-size", type=int, default=64, help="examples per step")
+    parser.add_argument("--lr", type=float, default=1e-3, help="learning rate")
+    parser.add_argument("--width", type=int, default=128, help="hidden width of the network")
+    parser.add_argument("--length", type=int, default=16, help="sequence length L")
+    parser.add_argument("--samples", type=int, default=16, help="sequences or points to generate")
+    parser.add_argument("--sample-steps", type=int, default=64, help="integration steps")
+    parser.add_argument("--seed", type=int, default=0, help="random seed")
+    parser.add_argument("--device", type=str, default="auto", help="auto, cpu, cuda or mps")
+    parser.add_argument("--out", type=str, default="", help="directory for the run's output")
+    parser.add_argument("--quiet", action="store_true", help="print only the final report")
+    return parser
+
+
+# %% 03. The training loop
+#
+# One loop for the whole course. The caller passes a function that takes the step index and returns
+# a scalar loss with a graph attached; everything else, the optimizer, the logging and the timing,
+# is the same in every lecture.
+
+
+def train(
+    model: torch.nn.Module,
+    loss_fn: Callable[[int], torch.Tensor],
+    steps: int = 600,
+    lr: float = 1e-3,
+    log_every: int = 100,
+    quiet: bool = False,
+    grad_clip: float = 1.0,
+    extra_parameters: list[torch.nn.Parameter] | None = None,
+) -> list[float]:
+    """Run ``steps`` optimizer steps on ``loss_fn`` and return the loss at every step.
+
+    ``loss_fn(step)`` draws its own batch, so a lecture that needs fresh noise at every step gets it
+    without a DataLoader. ``extra_parameters`` covers the rare case of a second trained object, such
+    as the bias force of Chapter 7.
+    """
+    parameters = list(model.parameters()) + list(extra_parameters or [])
+    optimizer = torch.optim.AdamW(parameters, lr=lr, weight_decay=0.0)
+    losses: list[float] = []
+    started = time.time()
+    model.train()
+    for step in range(steps):
+        loss = loss_fn(step)  # The lecture's own objective.
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(parameters, grad_clip)  # Keep a bad batch from derailing.
+        optimizer.step()
+        losses.append(float(loss.detach()))
+        if not quiet and log_every and (step + 1) % log_every == 0:
+            window = losses[-log_every:]  # Average over the last window, which is less jumpy.
+            print(f"  step {step + 1:5d}/{steps}  loss {sum(window) / len(window):.4f}")
+    if not quiet:
+        print(f"  trained {steps} steps in {time.time() - started:.1f}s")
+    model.eval()
+    return losses
+
+
+# %% 04. Saving a run
+#
+# A run writes three files: the weights, the configuration that produced them, and a short report of
+# what came out. No optimizer state, because none of these examples resumes training.
+
+
+def save_run(
+    out: str | Path,
+    model: torch.nn.Module | None = None,
+    config: dict | None = None,
+    report: dict | None = None,
+    losses: list[float] | None = None,
+    samples: list[str] | None = None,
+) -> Path | None:
+    """Write the weights, the configuration, the report, the loss curve and any samples."""
+    if not out:
+        return None
+    directory = Path(out)
+    directory.mkdir(parents=True, exist_ok=True)
+    if model is not None:
+        torch.save({"model": model.state_dict(), "config": config or {}},
+                   directory / "checkpoint.pt")
+    if config is not None:
+        (directory / "config.json").write_text(json.dumps(config, indent=2) + "\n")
+    if report is not None:
+        (directory / "report.json").write_text(json.dumps(report, indent=2) + "\n")
+    if losses is not None:
+        # Thin the curve to at most 200 points so the file stays small and readable.
+        stride = max(1, len(losses) // 200)
+        thinned = [round(v, 5) for v in losses[::stride]]
+        (directory / "losses.json").write_text(json.dumps(thinned) + "\n")
+    if samples is not None:
+        (directory / "samples.txt").write_text("\n".join(samples) + "\n")
+    return directory
+
+
+def load_run(out: str | Path, model: torch.nn.Module) -> dict:
+    """Read the weights back into ``model`` and return the configuration that trained them."""
+    payload = torch.load(Path(out) / "checkpoint.pt", map_location="cpu", weights_only=False)
+    model.load_state_dict(payload["model"])
+    model.eval()
+    return payload.get("config", {})
+
+
+# %% 05. Printing a report
+#
+# Every script ends with the same two-column block, so a reader can compare two methods by reading
+# two reports side by side.
+
+
+def print_report(title: str, report: dict, samples: list[str] | None = None) -> None:
+    """Print a titled block of name and value pairs, and a few samples."""
+    print(f"\n{title}")
+    print("-" * len(title))
+    for key, value in report.items():
+        if isinstance(value, float):
+            print(f"  {key:<34s} {value:.4f}")
+        else:
+            print(f"  {key:<34s} {value}")
+    if samples:
+        print("\n  samples")
+        for sample in samples[:8]:
+            print(f"    {sample}")

lecture_1/README.md

@@ -0,0 +1,83 @@
+# Lecture 1: Introduction and Mathematical Foundations
+
+Chapter 1 of the notes is the toolbox the other six chapters draw on, and it works every piece of
+it on a small example by hand. The six scripts here run those examples. Each one prints the
+chapter's hand-worked value, the value the script computes for it, and the gap between the two, so
+a student reading the chapter can watch its arithmetic come out of the code. Nothing here trains a
+network. The scripts are worked examples, and `tests/test_lecture_1.py` asserts every number they
+print against the chapter that states it.
+
+Run them from the repository root:
+
+```
+cd cis6270
+PYTHONPATH=. python lecture_1/probability.py
+PYTHONPATH=. python -m pytest tests/test_lecture_1.py -q
+```
+
+## The files
+
+| File | Command | What it works through |
+| --- | --- | --- |
+| `probability.py` | `python lecture_1/probability.py` | Sections 1.3 to 1.11. The two-token DNA distribution, the expected adenine count against a Monte Carlo ladder, the worked variance, the three peptides and their covariance matrix, the noising equation, the motif-and-binding table with its marginals and conditionals, Bayes' rule, and the independence check. |
+| `linear_algebra.py` | `python lecture_1/linear_algebra.py` | Sections 1.12, 1.13, 1.15, 1.16, 1.17 and 1.18. The simplex and a mixture, the softmax of `z = (2, 1, 0)`, one step of the two-state chain, the eigenvalues of `A = [[2, 1], [1, 2]]`, powers along an eigendirection, the telescoping `A^n = V Lambda^n V^{-1}`, the stationary distribution and its rate, and `e^{tA}` by both the series and the eigenvalues. |
+| `calculus.py` | `python lecture_1/calculus.py` | Sections 1.22 to 1.30. The Jacobian by automatic differentiation against the chapter's closed form, the determinant as local volume change, change of variables in one and in two dimensions, the log-determinant of a composition, the two inverse branches of the squaring map, the gradient and directional derivative, the divergence of three fields, and the continuity equation checked numerically on `v(x) = x`. |
+| `fields.py` | `python lecture_1/fields.py` | Sections 1.31 and 1.32. The score of the standard Gaussian and of `N(1, 4)`, the normalizing constant differentiated away, the Laplacian of the bowl `x_1^2 + 3 x_2^2` and of the saddle `x_1^2 - x_2^2`, the heat equation at `D = 1/2` on the widening Gaussian `N(0, t)`, and the Fokker-Planck equation with its two reductions, to the continuity equation at `sigma_t = 0` and to the heat equation at `b_t = 0`. The same drift and noise drive an Euler-Maruyama walk, whose sample variance tracks the elapsed time. Chapter 3 starts from both sections. |
+| `information.py` | `python lecture_1/information.py` | Sections 1.33 to 1.35. Information content, entropy, the cross-entropy of `p = (0.8, 0.2)` scored by `q = (0.6, 0.4)`, the split `H(p, q) = H(p) + KL(p \|\| q)`, nonnegativity and asymmetry, and maximum likelihood on ten DNA bases as a KL minimization. |
+| `optimization.py` | `python lecture_1/optimization.py` | Sections 1.36 to 1.38. Gradient descent on `(theta - 3)^2`, an equality constraint through its Lagrangian, normalization as a constraint, the maximum-entropy derivation that produces the exponential family at a GC target of 0.6, the KKT conditions on a bounded quadratic, and the same likelihood problem solved on logits through the softmax. |
+
+## The flags
+
+Every script takes `--steps`, `--seed`, `--out` and `--quiet`, so the flags learned once apply
+across the course. `--steps` drives the iteration count of each script, and every flag a script
+takes changes a line of its report.
+
+| Flag | Where it changes the report |
+| --- | --- |
+| `--steps` | `probability.py`, repeats of the ten-sample Monte Carlo estimate, default 5000. `linear_algebra.py`, iterations of `p_{t+1} = P^T p_t`, default 10. `calculus.py`, Euler steps along the worked flow, default 200. `fields.py`, Euler-Maruyama steps over `[0, 1]`, default 200. `information.py`, repeats of the ten-sample cross-entropy estimate, default 5000. `optimization.py`, gradient-descent steps, default 600. |
+| `--samples` | `probability.py`, the largest Monte Carlo sample size, default 1000. `fields.py`, paths simulated along the SDE, default 20000. `information.py`, the draws behind the sampled cross-entropy and the random distribution pairs, default 20000. `optimization.py`, the competing distributions drawn at the same GC target, default 4000. |
+| `--grid` | `calculus.py` and `fields.py`, the points used in the numerical checks, default 1000. |
+| `--seed` | The Monte Carlo draws, the simulated paths and the random distributions. The same seed gives the same report. |
+| `--out` | A directory for `config.json` and `report.json`, written through `cis6270.runner.save_run`. |
+| `--quiet` | Print the closing report and nothing else. |
+
+## What a default run prints
+
+| Script | Headline numbers |
+| --- | --- |
+| `probability.py` | `E[h(X)] = 1.1000`, Monte Carlo estimate `1.1050` at `N = 1000`, `Var(h(X)) = 0.6900`, spread of 5000 ten-sample averages `0.2616` against the predicted `0.2627`, `rho_XY = 0.9449`, `p(X = 1 \| Y = 1) = 0.8182`, association `0.0960`. |
+| `linear_algebra.py` | `softmax(z)_1 = 0.6652`, `lambda_1 = 3.0000`, `x_10` first entry `118099`, `p*_1 = 0.6000`, second eigenvalue `0.5000`, `e^{tA}_11 = 1.2275` at `t = 0.1`, truncation error `0.0002`. |
+| `calculus.py` | `cos(1) = 0.5403`, `det J_phi = 4.0000` at `x_1 = 2`, `p_Y = 0.1667` inside the image of `A`, `log\|det J\| = 3.1781` for the composition, continuity residual `0.0000`, `log p = -1.0000` along the flow at `t = 1`. |
+| `fields.py` | `s(3) = -3.0000`, `s(5) = -1.0000` under `N(1, 4)`, Laplacian `8.0000` on the bowl and `0.0000` on the saddle, `p_1(0) = 0.3989`, `d p_t(0)/dt = -0.1995`, `D` times the Laplacian `-0.1995`, every residual `0.0000`, sample variance `1.0130` at `t = 1` over 20000 paths. |
+| `information.py` | `H(p) = 0.5004`, `H(p, q) = 0.5919`, `KL(p \|\| q) = 0.0915`, `KL(q \|\| p) = 0.1046`, spread of 5000 ten-sample estimates `0.0502` against the predicted `0.0513`, empirical loss `1.2799` nats, uniform loss `1.3863` nats, `KL(p_hat \|\| uniform) = 0.1064`. |
+| `optimization.py` | `theta = 3.0000` after 600 steps, `lambda = -1.0000` on the equality constraint, `lambda = 0.4055` for a GC target of 0.6, `H(p*) = 1.3662` nats, best drawn competitor `1.3661` nats, `lambda = 4.0000` on the active bound, softmax descent loss `1.2799` nats. |
+
+Every one of these runs in a few seconds on one core. The sampled rows are the ones that move with
+`--seed`. `probability.py --steps 50000` closes the measured spread to `0.2624` and
+`information.py --steps 50000` closes it to `0.0510`, at the cost of a second each.
+
+## What is simplified
+
+- Every script runs in double precision. Chapter 1 quotes five decimal places in places, and the
+  single precision used by the training lectures would not reach them.
+- The Monte Carlo sections redraw their samples, so the sampled values move with `--seed`. The
+  chapter's own ten-sample draw gave 1.2 against the exact 1.1, and that particular draw is not
+  reproduced. The exact value and the `1/sqrt(N)` error scale are.
+- The continuity equation is checked by central differences at interior points of the support of
+  `p_t`, away from the two ends, where a one-sided difference would be needed. The Fokker-Planck
+  residual of `fields.py` uses the same differences, with a second difference for the Laplacian, so
+  it sits near `1e-8` and not at zero.
+- `fields.py` takes the Brownian increment of `dX_t = b_t(X_t) dt + sigma_t dW_t` as a scaled
+  Gaussian draw per step. Chapter 3 defines Brownian motion properly. With a constant drift and a
+  constant noise magnitude the Euler-Maruyama walk has the exact law of the chapter's `N(bt, t)`, so
+  the gap the script prints is Monte Carlo error alone.
+- The branch density `1/(2 sqrt(y))` is integrated numerically over `[0.01, 1]`, because the
+  integrand grows without bound at the origin. Its value over the full `[0, 1]` comes from the
+  antiderivative the chapter gives.
+- The random competitors in the maximum-entropy check and the random pairs in the KL check confirm
+  the chapter's inequalities on a finite sample. They are a check, not the proof, which the chapter
+  gives in full. The competitors cover the two free splits of the GC-constrained family, one inside
+  `{A, T}` and one inside `{C, G}`, drawn independently.
+- Each script has its own three-line `show` helper for the side-by-side printing and its own
+  `WORKED` dictionary of computed values. Neither belongs in `cis6270/`, because no other lecture
+  prints a chapter number beside a computed one.

lecture_1/calculus.py

@@ -0,0 +1,279 @@
+"""Jacobians, change of variables, divergence and the continuity equation.
+
+A map moves points, and the Jacobian is the linear map it looks like near one of them. Its
+determinant is the local volume factor, so the density at an image point is the density at its
+source divided by that factor, the log-determinants add across a composition, and for a
+many-to-one map the density is a sum of one corrected term per inverse branch. In continuous time
+that factor becomes the divergence of a velocity field, and the continuity equation holds it in
+balance with the change in density at every point.
+
+Chapter 1, Sections 1.22 to 1.30 of the notes. Runs on the maps of Sections 1.22.2 and 1.23, the
+squaring map of Section 1.26, and the expanding flow v(x) = x of Section 1.30.2. Every Jacobian
+is taken by automatic differentiation and checked against the chapter's closed form.
+
+    python lecture_1/calculus.py               # every worked example
+    python lecture_1/calculus.py --grid 4000   # finer grids in the numerical checks
+    python lecture_1/calculus.py --steps 2000  # a finer Euler walk along the worked flow
+"""
+
+# %% 01. Imports and configuration
+from __future__ import annotations
+
+import argparse
+import math
+
+import torch
+
+from cis6270.runner import print_report, save_run, set_seed
+
+torch.set_default_dtype(torch.float64)  # Chapter 1 is arithmetic, and five printed digits need it.
+
+
+A_SCALE = torch.tensor([[2.0, 0.0], [0.0, 3.0]])  # The map of Section 1.23, which scales area by 6.
+
+
+WORKED: dict[str, float] = {}  # Every computed value, by label, for tests/test_lecture_1.py.
+
+
+def show(label: str, by_hand: float, computed: float, quiet: bool = False) -> None:
+    """Print the chapter's hand-worked value beside the value this file computes, and keep it."""
+    WORKED[label] = computed  # The test suite asserts these against the chapter's own numbers.
+    if not quiet:  # Under --quiet only the closing report is printed.
+        computed = computed + 0.0  # Adding zero prints a negative zero as a plain zero.
+        print(f"  {label:<36s} chapter {by_hand:>11.5f}   computed {computed:>11.5f}"
+              f"   gap {abs(by_hand - computed):.1e}")
+
+
+def phi_mixed(x: torch.Tensor) -> torch.Tensor:
+    """The map (x_1^2 + x_2, sin x_1 + 3 x_2) of Section 1.22.2, nonlinear in its first component."""
+    return torch.stack([x[0] ** 2 + x[1], torch.sin(x[0]) + 3.0 * x[1]])
+
+
+def phi_square(x: torch.Tensor) -> torch.Tensor:
+    """The map (x_1^2, x_2) of Section 1.23, one-to-one once the domain is restricted to x_1 > 0."""
+    return torch.stack([x[0] ** 2, x[1]])
+
+
+def psi_scale(y: torch.Tensor) -> torch.Tensor:
+    """The linear map y -> A y with A = diag(2, 3), composed after phi_square in Section 1.25."""
+    return A_SCALE @ y
+
+
+def f_bowl(x: torch.Tensor) -> torch.Tensor:
+    """The scalar field f(x_1, x_2) = x_1^2 + 2 x_2^2 of Section 1.27.1."""
+    return x[0] ** 2 + 2.0 * x[1] ** 2
+
+
+def divergence(field, x: torch.Tensor) -> float:
+    """Section 1.29: the divergence is the sum of d v_i / d x_i, which is the trace of J_v."""
+    return float(torch.autograd.functional.jacobian(field, x).trace())
+
+
+# %% 02. The Jacobian matrix and the local linear approximation
+def jacobian(quiet: bool = False) -> dict:
+    """Sections 1.22.2 and 1.22.3: J_phi at (1, 0), and the displacement it predicts."""
+    point = torch.tensor([1.0, 0.0])  # The point at which the chapter evaluates the Jacobian.
+    computed = torch.autograd.functional.jacobian(phi_mixed, point)  # [2, 2], by backpropagation.
+    closed_form = torch.tensor([[2.0 * point[0], 1.0],  # The chapter's entries, 2 x_1 and 1.
+                                [torch.cos(point[0]), 3.0]])  # cos x_1 and 3; [2, 2].
+    delta = torch.tensor([0.01, -0.02])  # The small input displacement of Section 1.22.3.
+    predicted = computed @ delta  # J_phi(x) Delta x, the first-order change in the output; [2].
+    exact = phi_mixed(point + delta) - phi_mixed(point)  # The change the full map produces; [2].
+    show("J_phi(1, 0)_11", 2.0, float(computed[0, 0]), quiet)
+    show("J_phi(1, 0)_21 = cos(1)", 0.5403, float(computed[1, 0]), quiet)
+    show("largest |autograd - closed form|", 0.0, float((computed - closed_form).abs().max()), quiet)
+    show("predicted change, first output", 0.0, float(predicted[0]), quiet)
+    show("predicted change, second output", -0.0546, float(predicted[1]), quiet)
+    show("exact change, first output", 0.0001, float(exact[0]), quiet)
+    show("exact change, second output", -0.05464, float(exact[1]), quiet)
+    return {"jacobian": computed, "predicted": predicted, "exact": exact}
+
+
+# %% 03. The Jacobian determinant as local volume change
+def determinant(quiet: bool = False) -> dict:
+    """Section 1.23: the absolute determinant is the factor by which a small region changes area."""
+    corners = torch.tensor([[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]])  # The unit square.
+    mapped = corners @ A_SCALE.T  # The image rectangle [0, 2] by [0, 3]; [4, 2].
+    area = float((mapped[1, 0] - mapped[0, 0]) * (mapped[3, 1] - mapped[0, 1]))  # Width times height.
+    collapsing = torch.tensor([[1.0, 0.0], [0.0, 0.0]])  # A map that flattens the plane onto a line.
+    stretch = torch.autograd.functional.jacobian(phi_square, torch.tensor([2.0, 1.0]))  # At x_1 = 2.
+    small = torch.tensor([0.25, 1.0])  # A point where the same map contracts, since 2 x_1 is below 1.
+    squeeze = torch.autograd.functional.jacobian(phi_square, small)  # J_phi at x_1 = 0.25; [2, 2].
+    show("det A for diag(2, 3)", 6.0, float(torch.linalg.det(A_SCALE)), quiet)
+    show("area of the image rectangle", 6.0, area, quiet)
+    show("det of the collapsing map", 0.0, float(torch.linalg.det(collapsing)), quiet)
+    show("det J_phi at x_1 = 2", 4.0, float(torch.linalg.det(stretch)), quiet)
+    show("det J_phi at x_1 = 0.25", 0.5, float(torch.linalg.det(squeeze)), quiet)
+    return {"area": area, "stretch": float(torch.linalg.det(stretch))}
+
+
+# %% 04. Change of variables in one and in two dimensions
+def change_of_variables(grid: int, quiet: bool = False) -> dict:
+    """Sections 1.24.1 and 1.24.2: p_Y(phi(x)) = p_X(x) / |det J_phi(x)| on the two worked maps."""
+    inverse_slope = 0.5  # d phi^{-1}/dy for phi^{-1}(y) = (y - 1)/2, the reciprocal of phi' = 2.
+    p_y = 1.0 * abs(inverse_slope)  # p_Y(y) = p_X(phi^{-1}(y)) |d phi^{-1}/dy|, so p_Y = 1/2.
+    preimage = ((1.4 - 1.0) / 2.0, (2.0 - 1.0) / 2.0)  # The inverse image of B = [1.4, 2.0].
+    mass = (preimage[1] - preimage[0]) * 1.0  # P(X in [0.2, 0.5]) under the uniform density.
+    line = 1.0 + (torch.arange(grid) + 0.5) * (2.0 / grid)  # Midpoints of [1, 3]; [grid].
+    normalization = float(torch.full_like(line, p_y).sum() * (2.0 / grid))  # Integral of p_Y there.
+    flat = 1.0 / float(torch.linalg.det(A_SCALE))  # p_Y = 1/6 inside the image of the unit square.
+    step = 1.0 / grid  # Midpoint spacing on the unit interval of the first input coordinate.
+    x_one = (torch.arange(grid) + 0.5) * step  # Midpoints of (0, 1) in that coordinate; [grid].
+    volume_factor = 2.0 * x_one  # |det J_phi| = 2 x_1 for phi = (x_1^2, x_2), at each point; [grid].
+    output_density = 1.0 / (2.0 * (x_one ** 2).sqrt())  # p_Y(y) = 1/(2 sqrt(y_1)) at y_1 = x_1^2.
+    residual = float((output_density * volume_factor - 1.0).abs().max())  # Should recover p_X = 1.
+    transported = float((output_density * volume_factor).sum() * step)  # Total mass by substitution.
+    show("preimage of [1.4, 2.0], lower end", 0.2, preimage[0], quiet)
+    show("preimage of [1.4, 2.0], upper end", 0.5, preimage[1], quiet)
+    show("P(X in the preimage)", 0.3, mass, quiet)
+    show("p_Y on [1, 3] after Y = 2X + 1", 0.5, p_y, quiet)
+    show("integral of p_Y over [1, 3]", 1.0, normalization, quiet)
+    show("p_Y inside the image of A", 1.0 / 6.0, flat, quiet)
+    show("density times area for A", 1.0, flat * 6.0, quiet)
+    show("p_Y(y_1 = 1) for the x_1^2 map", 0.5, 1.0 / (2.0 * math.sqrt(1.0)), quiet)
+    show("largest |p_Y |det J| - p_X|", 0.0, residual, quiet)
+    show("total mass after the x_1^2 map", 1.0, transported, quiet)
+    return {"p_y": p_y, "flat": flat, "residual": residual, "transported": transported}
+
+
+# %% 05. Composed maps and the log-density form
+def composition(quiet: bool = False) -> dict:
+    """Section 1.25: log|det J| adds across a composition, so two factors become one sum."""
+    point = torch.tensor([2.0, 1.0])  # The chapter evaluates the composition at x_1 = 2.
+    inner = torch.autograd.functional.jacobian(phi_square, point)  # J_phi(x), det 2 x_1 = 4; [2, 2].
+    outer = torch.autograd.functional.jacobian(psi_scale, phi_square(point))  # J_psi at phi(x); [2, 2].
+    chained = outer @ inner  # The chain rule of Proposition 1.1, giving diag(8, 3); [2, 2].
+    composed = torch.autograd.functional.jacobian(lambda x: psi_scale(phi_square(x)), point)
+    log_inner = math.log(abs(float(torch.linalg.det(inner))))  # log 4 = 1.3863.
+    log_outer = math.log(abs(float(torch.linalg.det(outer))))  # log 6 = 1.7918.
+    log_density = math.log(1.0) - (log_inner + log_outer)  # log p_Y = log p_X - log|det J|.
+    show("det J_psi", 6.0, float(torch.linalg.det(outer)), quiet)
+    show("chained Jacobian, entry 11", 8.0, float(chained[0, 0]), quiet)
+    show("det of the composition", 24.0, float(torch.linalg.det(composed)), quiet)
+    show("log 4", 1.3863, log_inner, quiet)
+    show("log 6", 1.7918, log_outer, quiet)
+    show("log 4 + log 6 = log 24", 3.1781, log_inner + log_outer, quiet)
+    show("log p_Y from a unit base density", -3.1781, log_density, quiet)
+    return {"log_sum": log_inner + log_outer, "log_density": log_density}
+
+
+# %% 06. Local invertibility and several inverse branches
+def branches(grid: int, quiet: bool = False) -> dict:
+    """Section 1.26: for Y = X^2 on a uniform X, the two preimages each contribute to p_Y(y)."""
+    def branch_density(y: float) -> float:
+        roots = (math.sqrt(y), -math.sqrt(y))  # The two solutions of x^2 = y inside the support.
+        return sum(0.5 / abs(2.0 * root) for root in roots)  # Sum of p_X(x) / |phi'(x)| over roots.
+
+    lower = 0.01  # The integral is cut at 0.01 because the integrand grows without bound at zero.
+    step = (1.0 - lower) / grid  # Midpoint spacing on [0.01, 1].
+    points = lower + (torch.arange(grid) + 0.5) * step  # Midpoints of the integration interval.
+    partial = float((1.0 / (2.0 * points.sqrt())).sum() * step)  # Midpoint integral of 1/(2 sqrt y).
+    show("p_Y(0.25), two branches at ±0.5", 1.0, branch_density(0.25), quiet)
+    show("p_Y(0.04), two branches at ±0.2", 2.5, branch_density(0.04), quiet)
+    show("integral of p_Y over [0.01, 1]", 1.0 - math.sqrt(lower), partial, quiet)
+    return {"at_quarter": branch_density(0.25), "partial": partial}
+
+
+# %% 07. Scalar fields, the gradient and the directional derivative
+def gradient(quiet: bool = False) -> dict:
+    """Sections 1.27.1 and 1.27.2: grad f at (1, 2), its steepest direction and a directional rate."""
+    point = torch.tensor([1.0, 2.0])  # The point at which the chapter evaluates the gradient.
+    grad = torch.autograd.functional.jacobian(f_bowl, point)  # (2 x_1, 4 x_2) = (2, 8); [2].
+    steepest = grad / grad.norm()  # The unit vector of fastest increase; [2].
+    along_x1 = float(grad @ torch.tensor([1.0, 0.0]))  # D_u f for u the first basis direction.
+    along_x2 = float(grad @ torch.tensor([0.0, 1.0]))  # D_u f for u the second basis direction.
+    show("grad f(1, 2), first entry", 2.0, float(grad[0]), quiet)
+    show("grad f(1, 2), second entry", 8.0, float(grad[1]), quiet)
+    show("|grad f(1, 2)| = sqrt(68)", math.sqrt(68.0), float(grad.norm()), quiet)
+    show("steepest direction, first entry", 0.243, float(steepest[0]), quiet)
+    show("steepest direction, second entry", 0.970, float(steepest[1]), quiet)
+    show("D_u f along (1, 0)", 2.0, along_x1, quiet)
+    show("D_u f along (0, 1)", 8.0, along_x2, quiet)
+    return {"gradient": grad, "along_x1": along_x1}
+
+
+# %% 08. Divergence
+def divergences(quiet: bool = False) -> dict:
+    """Section 1.29: the divergence of the chapter's three fields is -3, then 1, then 0."""
+    point = torch.tensor([1.0, 1.0])  # The divergences below are the same at every point.
+    contract = divergence(lambda x: torch.stack([-x[0], -2.0 * x[1]]), point)  # -1 + (-2) = -3.
+    expand = divergence(lambda x: torch.stack([2.0 * x[0], -x[1]]), point)  # 2 - 1 = 1.
+    rotate = divergence(lambda x: torch.stack([-x[1], x[0]]), point)  # 0 + 0 = 0.
+    show("div of (-x_1, -2 x_2)", -3.0, contract, quiet)
+    show("div of (2 x_1, -x_2)", 1.0, expand, quiet)
+    show("div of (-x_2, x_1)", 0.0, rotate, quiet)
+    show("area factor e^{div t} at t = 0.5", 1.649, math.exp(expand * 0.5), quiet)
+    return {"contract": contract, "expand": expand, "rotate": rotate}
+
+
+# %% 09. The continuity equation on the worked flow
+def continuity(grid: int, steps: int, quiet: bool = False) -> dict:
+    """Section 1.30.2: v(x) = x stretches a uniform source, with p_t(x) = e^{-t} on [0, e^t].
+
+    Central differences give the continuity residual at interior points, and one Euler walk
+    accumulates the log-density rate -div v of Section 1.30.1 along a trajectory.
+    """
+    def flux(time: float, position):
+        return math.exp(-time) * position  # j_t(x) = p_t(x) v_t(x), with p_t = e^{-t} and v(x) = x.
+
+    t, h = 0.5, 1e-4  # The time of the residual check, and the central-difference step.
+    x = torch.linspace(0.1, 1.4, grid)  # Interior points of the support [0, e^{0.5}]; [grid].
+    time_rate = (math.exp(-(t + h)) - math.exp(-(t - h))) / (2.0 * h)  # d p_t / d t at fixed x.
+    flux_rate = (flux(t, x + h) - flux(t, x - h)) / (2.0 * h)  # d (p_t v) / d x; [grid].
+    residual = float((time_rate + flux_rate).abs().max())  # The continuity equation, point by point.
+    x_t, log_p, dt = 0.5, 0.0, 1.0 / max(steps, 1)  # One trajectory started inside the source.
+    for _ in range(max(steps, 1)):
+        log_p -= 1.0 * dt  # d/dt log p_t(x_t) = -div v = -1 for the field v(x) = x.
+        x_t = x_t + x_t * dt  # Euler step on dx_t/dt = v(x_t) = x_t.
+    show("largest continuity residual", 0.0, residual, quiet)
+    show("log p along the trajectory at t = 1", -1.0, log_p, quiet)
+    show("p_t(x_t) at t = 1", 0.3679, math.exp(log_p), quiet)
+    show("x_t at t = 1, exact 0.5e", 0.5 * math.e, x_t, quiet)
+    show("density at t = log 2", 0.5, math.exp(-math.log(2.0)), quiet)
+    show("mass at t = log 2", 1.0, 2.0 * math.exp(-math.log(2.0)), quiet)
+    show("mass at t = 1", 1.000, math.e * math.exp(-1.0), quiet)
+    return {"residual": residual, "log_p": log_p, "x_t": x_t}
+
+
+# %% 10. The report
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    parser.add_argument("--grid", type=int, default=1000, help="points in the numerical checks")
+    parser.add_argument("--steps", type=int, default=200, help="Euler steps along the worked flow")
+    parser.add_argument("--seed", type=int, default=0, help="random seed, for a repeatable run")
+    parser.add_argument("--out", type=str, default="", help="directory for the run's output")
+    parser.add_argument("--quiet", action="store_true", help="print only the final report")
+    args = parser.parse_args()
+
+    set_seed(args.seed)  # Seed first, so the same seed gives the same run.
+    if not args.quiet:
+        print("Chapter 1, Sections 1.22 to 1.30, worked example by worked example\n")
+    local = jacobian(args.quiet)
+    volume = determinant(args.quiet)
+    moved = change_of_variables(args.grid, args.quiet)
+    stacked = composition(args.quiet)
+    folded = branches(args.grid, args.quiet)
+    slope = gradient(args.quiet)
+    fields = divergences(args.quiet)
+    flow = continuity(args.grid, args.steps, args.quiet)
+
+    report = {
+        "J_phi(1, 0)_21 = cos(1)": float(local["jacobian"][1, 0]),
+        "det J_phi at x_1 = 2": volume["stretch"],
+        "p_Y on [1, 3] after Y = 2X + 1": moved["p_y"],
+        "p_Y inside the image of A": moved["flat"],
+        "mass after the x_1^2 map": moved["transported"],
+        "log|det J| of the composition": stacked["log_sum"],
+        "p_Y(0.25) summed over branches": folded["at_quarter"],
+        "D_u f along (1, 0)": slope["along_x1"],
+        "div of (2 x_1, -x_2)": fields["expand"],
+        "largest continuity residual": flow["residual"],
+        "log p along the flow at t = 1": flow["log_p"],
+        "x_t at t = 1 after Euler steps": flow["x_t"],
+    }
+    print_report("Calculus, Sections 1.22 to 1.30", report)
+    save_run(args.out, config=vars(args), report=report)
+
+
+if __name__ == "__main__":  # Execute when launched directly.
+    main()  # Work every example of Sections 1.22 to 1.30 and print the report.

lecture_1/fields.py

@@ -0,0 +1,234 @@
+"""Score functions, the Laplacian and the Fokker-Planck equation.
+
+The score of a density is the gradient of its logarithm, a vector field of the same shape as a
+velocity field, and the normalizing constant differentiates away. The Laplacian is the divergence
+of the gradient, and it measures the curvature that drives diffusion, so the heat equation sets the
+rate of change of a density proportional to it. One statement about the density holds both
+mechanisms at once, the Fokker-Planck equation, and it returns the continuity equation when the
+noise magnitude is zero and the heat equation when the drift is.
+
+Chapter 1, Sections 1.31 and 1.32 of the notes. Runs on the Gaussians of Section 1.31, the bowl and
+the saddle of Section 1.32, the widening Gaussian p_t = N(0, t) of Section 1.32.1 and the expanding
+flow v(x) = x of Section 1.30.2. Prints the chapter's number beside the computed one.
+
+    python lecture_1/fields.py               # every worked example
+    python lecture_1/fields.py --steps 2000  # a finer Euler-Maruyama walk
+    python lecture_1/fields.py --quiet       # the closing report only
+"""
+
+# %% 01. Imports and configuration
+from __future__ import annotations
+
+import argparse
+import math
+
+import torch
+
+from cis6270.runner import print_report, save_run, set_seed
+
+torch.set_default_dtype(torch.float64)  # Chapter 1 is arithmetic, and five printed digits need it.
+
+
+SCORE_POINTS = torch.tensor([3.0, -2.0, 0.0])  # The three states scored in Section 1.31.
+WIDE_MEAN, WIDE_VARIANCE = 1.0, 4.0  # The general Gaussian N(mu, sigma^2) of Section 1.31.
+DRIFT = 0.5  # The constant drift b_t(x) = 0.5 of the combined process in Section 1.32.2.
+NOISE = 1.0  # The noise magnitude sigma_t, so that D = sigma_t^2 / 2 = 0.5 as in Section 1.32.1.
+
+
+WORKED: dict[str, float] = {}  # Every computed value, by label, for tests/test_lecture_1.py.
+
+
+def show(label: str, by_hand: float, computed: float, quiet: bool = False) -> None:
+    """Print the chapter's hand-worked value beside the value this file computes, and keep it."""
+    WORKED[label] = computed  # The test suite asserts these against the chapter's own numbers.
+    if not quiet:  # Under --quiet only the closing report is printed.
+        computed = computed + 0.0  # Adding zero prints a negative zero as a plain zero.
+        print(f"  {label:<36s} chapter {by_hand:>11.5f}   computed {computed:>11.5f}"
+              f"   gap {abs(by_hand - computed):.1e}")
+
+
+def log_gaussian(x: torch.Tensor, mu: float = 0.0, variance: float = 1.0) -> torch.Tensor:
+    """log N(x; mu, sigma^2) = -(x - mu)^2 / (2 sigma^2) - (1/2) log(2 pi sigma^2)."""
+    normalizer = 0.5 * math.log(2.0 * math.pi * variance)  # The term with no x in it.
+    return -((x - mu) ** 2) / (2.0 * variance) - normalizer
+
+
+def gaussian(t: float, x: torch.Tensor) -> torch.Tensor:
+    """p_t = N(0, t) in one dimension, the widening Gaussian of Section 1.32.1."""
+    return (-(x ** 2) / (2.0 * t)).exp() / math.sqrt(2.0 * math.pi * t)
+
+
+def score(log_density, x: torch.Tensor) -> torch.Tensor:
+    """Section 1.31: s(x) = grad_x log p(x), taken by automatic differentiation."""
+    point = x.detach().clone().requires_grad_(True)  # A leaf to differentiate against; [N].
+    total = log_density(point).sum()  # The points do not interact, so one pass gives every score.
+    gradient, = torch.autograd.grad(total, point)  # d/dx log p at each point; [N].
+    return gradient
+
+
+def laplacian(field, x: torch.Tensor) -> float:
+    """Section 1.32: the Laplacian is the divergence of the gradient, the trace of the Hessian."""
+    return float(torch.autograd.functional.hessian(field, x).trace())
+
+
+def fokker_planck_residual(p, b, sigma: float, x: torch.Tensor, t: float) -> float:
+    """Section 1.32.2: d p_t / d t + div(b_t p_t) - (1/2) sigma_t^2 Laplacian p_t, by differences.
+
+    Setting sigma to zero leaves the continuity equation of Section 1.30, and setting b to zero
+    leaves the heat equation of Section 1.32.1, so one residual covers all three statements.
+    """
+    h = 1e-4  # The central-difference step, small enough for five digits in double precision.
+    time_rate = (p(t + h, x) - p(t - h, x)) / (2.0 * h)  # d p_t(x) / d t at fixed x; [grid].
+    flux = (b(x + h) * p(t, x + h) - b(x - h) * p(t, x - h)) / (2.0 * h)  # div(b_t p_t); [grid].
+    curvature = (p(t, x + h) - 2.0 * p(t, x) + p(t, x - h)) / h ** 2  # Laplacian p_t; [grid].
+    return float((time_rate + flux - 0.5 * sigma ** 2 * curvature).abs().max())
+
+
+# %% 02. The score of a Gaussian
+def scores(quiet: bool = False) -> dict:
+    """Section 1.31: the score of N(0, 1) is -x, and of N(mu, sigma^2) it is -(x - mu)/sigma^2."""
+    standard = score(log_gaussian, SCORE_POINTS)  # s(x) at x = 3, -2 and 0; [3].
+    closed_form = -SCORE_POINTS  # The chapter's result, with the log(2 pi) term differentiated away.
+    unnormalized = score(lambda z: -(z ** 2) / 2.0, SCORE_POINTS)  # The same density without its Z.
+    wide = score(lambda z: log_gaussian(z, WIDE_MEAN, WIDE_VARIANCE), torch.tensor([5.0]))  # [1].
+    show("s(x) at x = 3", -3.0, float(standard[0]), quiet)
+    show("s(x) at x = -2", 2.0, float(standard[1]), quiet)
+    show("s(x) at the mode x = 0", 0.0, float(standard[2]), quiet)
+    show("largest |s(x) + x|", 0.0, float((standard - closed_form).abs().max()), quiet)
+    show("largest change from dropping Z", 0.0, float((unnormalized - standard).abs().max()), quiet)
+    show("s(5) under N(1, 4)", -1.0, float(wide[0]), quiet)
+    return {"standard": standard, "wide": float(wide[0])}
+
+
+# %% 03. The Laplacian of a scalar field
+def curvatures(quiet: bool = False) -> dict:
+    """Section 1.32: the bowl x_1^2 + 3 x_2^2 has Laplacian 8, and the saddle x_1^2 - x_2^2 has 0."""
+    def bowl(x: torch.Tensor) -> torch.Tensor:
+        return x[0] ** 2 + 3.0 * x[1] ** 2  # The field whose two curvatures are 2 and 6.
+
+    def saddle(x: torch.Tensor) -> torch.Tensor:
+        return x[0] ** 2 - x[1] ** 2  # The field whose two curvatures cancel.
+
+    point = torch.tensor([1.0, 2.0])  # The Laplacians below are the same at every point.
+    second = torch.autograd.functional.hessian(bowl, point)  # The two pure second derivatives; [2, 2].
+    h = 1e-3  # The spacing of the four neighbors used for the average reading of the Laplacian.
+    steps = torch.tensor([[h, 0.0], [-h, 0.0], [0.0, h], [0.0, -h]])  # The four offsets; [4, 2].
+    neighbors = sum(bowl(point + step) for step in steps) - 4.0 * bowl(point)  # Above the center.
+    show("d^2 f / d x_1^2", 2.0, float(second[0, 0]), quiet)
+    show("d^2 f / d x_2^2", 6.0, float(second[1, 1]), quiet)
+    show("Laplacian of x_1^2 + 3 x_2^2", 8.0, laplacian(bowl, point), quiet)
+    show("Laplacian from the four neighbors", 8.0, float(neighbors) / h ** 2, quiet)
+    show("Laplacian of the saddle", 0.0, laplacian(saddle, point), quiet)
+    return {"bowl": laplacian(bowl, point), "saddle": laplacian(saddle, point)}
+
+
+# %% 04. The heat equation
+def heat(grid: int, quiet: bool = False) -> dict:
+    """Section 1.32.1: at D = 1/2 the Gaussian N(0, t) solves d p_t / d t = D Laplacian p_t."""
+    diffusion = 0.5 * NOISE ** 2  # D = sigma_t^2 / 2, the coefficient the chapter takes as 1/2.
+    center = 1.0 / math.sqrt(2.0 * math.pi)  # p_1(0), the height at the middle at t = 1.
+    time_rate = -0.5 / math.sqrt(2.0 * math.pi)  # d p_t(0)/dt = -(1/2)(2 pi)^{-1/2} t^{-3/2} at t = 1.
+    curvature = laplacian(lambda z: gaussian(1.0, z), torch.tensor([0.0]))  # Laplacian p_1 at 0.
+    x = torch.linspace(-3.0, 3.0, grid)  # Points covering the bulk of the mass at t = 1; [grid].
+    residual = fokker_planck_residual(
+        gaussian, torch.zeros_like, NOISE, x, 1.0)  # The heat equation, with no drift.
+    show("p_1(0) = 1 / sqrt(2 pi)", 0.39894, center, quiet)
+    show("d p_t(0) / d t at t = 1", -0.19947, time_rate, quiet)
+    show("Laplacian of p_1 at 0", -0.39894, curvature, quiet)
+    show("-p_1(0) / sigma^2 at sigma^2 = 1", -0.39894, -center / 1.0, quiet)
+    show("D times the Laplacian of p_1 at 0", -0.19947, diffusion * curvature, quiet)
+    show("largest heat-equation residual", 0.0, residual, quiet)
+    return {"center": center, "time_rate": time_rate, "curvature": curvature,
+            "flow": diffusion * curvature, "residual": residual}
+
+
+# %% 05. Drift and diffusion together
+def fokker_planck(grid: int, quiet: bool = False) -> dict:
+    """Section 1.32.2: a drift that translates and a noise that widens, in one equation.
+
+    With a constant drift b and sigma_t = 1 the density N(bt, t) solves the equation, since the
+    drift moves the center and the noise grows the variance at the rate of the heat equation.
+    """
+    def translating(t: float, x: torch.Tensor) -> torch.Tensor:
+        return gaussian(t, x - DRIFT * t)  # N(bt, t), the center carried along by the drift.
+
+    def expanding(t: float, x: torch.Tensor) -> torch.Tensor:
+        return torch.full_like(x, math.exp(-t))  # p_t = e^{-t} on [0, e^t], Section 1.30.2.
+
+    x = torch.linspace(-3.0, 3.0, grid)  # Points across the translating Gaussian; [grid].
+    both = fokker_planck_residual(translating, lambda z: torch.full_like(z, DRIFT), NOISE, x, 1.0)
+    inside = torch.linspace(0.1, 1.4, grid)  # Interior points of the support at t = 0.5; [grid].
+    transport = fokker_planck_residual(expanding, lambda z: z, 0.0, inside, 0.5)  # sigma_t = 0.
+    center = DRIFT * 1.0  # The mean of p_t at t = 1, which the drift alone sets.
+    show("center of p_t at t = 1", 0.5, center, quiet)
+    show("largest residual, drift and noise", 0.0, both, quiet)
+    show("largest residual at sigma_t = 0", 0.0, transport, quiet)
+    return {"both": both, "transport": transport, "center": center}
+
+
+# %% 06. The process behind the density
+def euler_maruyama(samples: int, steps: int, quiet: bool = False) -> dict:
+    """Section 1.32.2: dX_t = b_t(X_t) dt + sigma_t dW_t, simulated, against the density above.
+
+    The variance of p_t = N(bt, t) equals the elapsed time, so the sample mean and the sample
+    variance of the walk track bt and t. Both gaps shrink as --samples grows.
+    """
+    dt = 1.0 / max(steps, 1)  # The time increment of one Euler-Maruyama step.
+    x_t = torch.zeros(max(samples, 1))  # Every path starts at the origin, so p_0 is a point mass.
+    mean_gap, variance_gap = 0.0, 0.0  # The largest departure from bt and from t along the walk.
+    for step in range(max(steps, 1)):
+        noise = torch.randn_like(x_t)  # The Brownian increment dW_t, scaled by sqrt(dt); [samples].
+        x_t = x_t + DRIFT * dt + NOISE * math.sqrt(dt) * noise  # One step of the chapter's SDE.
+        elapsed = (step + 1) * dt  # The time reached after this step.
+        mean_gap = max(mean_gap, abs(float(x_t.mean()) - DRIFT * elapsed))
+        variance_gap = max(variance_gap, abs(float(x_t.var(unbiased=False)) - elapsed))
+    show("sample mean at t = 1, exact 0.5", 0.5, float(x_t.mean()), quiet)
+    show("sample variance at t = 1, exact 1", 1.0, float(x_t.var(unbiased=False)), quiet)
+    show("largest |sample mean - b t|", 0.0, mean_gap, quiet)
+    show("largest |sample variance - t|", 0.0, variance_gap, quiet)
+    return {"mean": float(x_t.mean()), "variance": float(x_t.var(unbiased=False)),
+            "mean_gap": mean_gap, "variance_gap": variance_gap, "steps": max(steps, 1)}
+
+
+# %% 07. The report
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    parser.add_argument("--grid", type=int, default=1000, help="points in the numerical checks")
+    parser.add_argument("--steps", type=int, default=200, help="Euler-Maruyama steps over [0, 1]")
+    parser.add_argument("--samples", type=int, default=20000, help="paths simulated along the SDE")
+    parser.add_argument("--seed", type=int, default=0, help="random seed for the simulated paths")
+    parser.add_argument("--out", type=str, default="", help="directory for the run's output")
+    parser.add_argument("--quiet", action="store_true", help="print only the final report")
+    args = parser.parse_args()
+
+    set_seed(args.seed)  # Seed first, so the same seed gives the same walk.
+    if not args.quiet:
+        print("Chapter 1, Sections 1.31 and 1.32, worked example by worked example\n")
+    uphill = scores(args.quiet)
+    bowl = curvatures(args.quiet)
+    spreading = heat(args.grid, args.quiet)
+    combined = fokker_planck(args.grid, args.quiet)
+    walk = euler_maruyama(args.samples, args.steps, args.quiet)
+
+    report = {
+        "s(x) at x = 3": float(uphill["standard"][0]),
+        "s(5) under N(1, 4)": uphill["wide"],
+        "Laplacian of x_1^2 + 3 x_2^2": bowl["bowl"],
+        "Laplacian of the saddle": bowl["saddle"],
+        "p_1(0)": spreading["center"],
+        "d p_t(0) / d t at t = 1": spreading["time_rate"],
+        "Laplacian of p_1 at 0": spreading["curvature"],
+        "D times that Laplacian": spreading["flow"],
+        "heat-equation residual": spreading["residual"],
+        "Fokker-Planck residual": combined["both"],
+        "residual at sigma_t = 0": combined["transport"],
+        "Euler-Maruyama steps": walk["steps"],
+        "sample variance at t = 1": walk["variance"],
+        "largest |sample variance - t|": walk["variance_gap"],
+    }
+    print_report("Fields, Sections 1.31 and 1.32", report)
+    save_run(args.out, config=vars(args), report=report)
+
+
+if __name__ == "__main__":  # Execute when launched directly.
+    main()  # Work every example of Sections 1.31 and 1.32 and print the report.

lecture_1/information.py

@@ -0,0 +1,219 @@
+"""Entropy, cross-entropy, KL divergence and maximum likelihood, in nats.
+
+Information content measures one observation, entropy averages it under the distribution that
+produced it, cross-entropy averages it under a second distribution supplying the probabilities,
+and the KL divergence is the difference between those two averages. Minimizing a negative
+log-likelihood over a model is therefore minimizing a KL divergence to the data distribution.
+
+Chapter 1, Sections 1.33, 1.34 and 1.35 of the notes. Runs on the pair p = (0.8, 0.2) and
+q = (0.6, 0.4) of Section 1.33.3 and on the ten DNA bases of Section 1.35. Prints the chapter's
+number beside the computed one. Logarithms are natural, so every quantity is in nats.
+
+    python lecture_1/information.py                   # every worked example
+    python lecture_1/information.py --samples 200000  # a tighter Monte Carlo cross-entropy
+    python lecture_1/information.py --quiet           # the closing report only
+"""
+
+# %% 01. Imports and configuration
+from __future__ import annotations
+
+import argparse
+import math
+
+import torch
+
+from cis6270.runner import print_report, save_run, set_seed
+
+torch.set_default_dtype(torch.float64)  # Chapter 1 is arithmetic, and five printed digits need it.
+
+
+P_REF = torch.tensor([0.8, 0.2])  # The reference distribution p of Sections 1.33.3 and 1.34.
+Q_MODEL = torch.tensor([0.6, 0.4])  # The comparison distribution q scored against those samples.
+BASE_COUNTS = torch.tensor([4.0, 3.0, 2.0, 1.0])  # n(A), n(C), n(G), n(T) over ten DNA bases.
+
+
+WORKED: dict[str, float] = {}  # Every computed value, by label, for tests/test_lecture_1.py.
+
+
+def show(label: str, by_hand: float, computed: float, quiet: bool = False) -> None:
+    """Print the chapter's hand-worked value beside the value this file computes, and keep it."""
+    WORKED[label] = computed  # The test suite asserts these against the chapter's own numbers.
+    if not quiet:  # Under --quiet only the closing report is printed.
+        computed = computed + 0.0  # Adding zero prints a negative zero as a plain zero.
+        print(f"  {label:<36s} chapter {by_hand:>11.5f}   computed {computed:>11.5f}"
+              f"   gap {abs(by_hand - computed):.1e}")
+
+
+def entropy(p: torch.Tensor) -> float:
+    """H(p) = -sum_x p(x) log p(x), the expected information in one sample drawn from p."""
+    return float(-(p * p.clamp_min(1e-300).log()).sum())  # Clamp keeps 0 log 0 at zero.
+
+
+def cross_entropy(p: torch.Tensor, q: torch.Tensor) -> float:
+    """H(p, q) = -sum_x p(x) log q(x), the cost of scoring samples from p with the values of q."""
+    return float(-(p * q.log()).sum())  # The weights come from p and the probabilities from q.
+
+
+def divergence(p: torch.Tensor, q: torch.Tensor) -> float:
+    """KL(p || q) = sum_x p(x) log(p(x) / q(x)), the expected log ratio taken under p."""
+    return float((p * (p / q).log()).sum())  # Zero exactly when q matches p at every state.
+
+
+# %% 02. Information content
+def information(quiet: bool = False) -> dict:
+    """Section 1.33.1: I(x) = -log p(x), so a rarer outcome carries more information."""
+    common = -math.log(0.5)  # An outcome of probability one half.
+    rare = -math.log(0.01)  # An outcome fifty times less likely.
+    certain = -math.log(1.0)  # A certain event, which has zero information.
+    show("I(x) at p(x) = 0.5", 0.693, common, quiet)
+    show("I(x) at p(x) = 0.01", 4.605, rare, quiet)
+    show("I(x) at p(x) = 1", 0.0, certain, quiet)
+    show("ratio of the two surprises", 6.6, rare / common, quiet)
+    return {"common": common, "rare": rare}
+
+
+# %% 03. Entropy
+def entropies(quiet: bool = False) -> dict:
+    """Section 1.33.2: entropy is high when the mass is spread and zero when it is all on one state."""
+    skewed = entropy(torch.tensor([0.9, 0.1]))  # 0.09482 + 0.23026 nats.
+    balanced = entropy(torch.tensor([0.5, 0.5]))  # log 2, the largest value on two states.
+    deterministic = entropy(torch.tensor([1.0, 0.0]))  # All mass on one state, so 1 log 1 = 0.
+    uniform_four = entropy(torch.full((4,), 0.25))  # log K at K = 4, proved in Section 1.37.2.
+    show("H((0.9, 0.1))", 0.325, skewed, quiet)
+    show("H((0.5, 0.5)) = log 2", 0.693, balanced, quiet)
+    show("H of a point mass", 0.0, deterministic, quiet)
+    show("H of the uniform on 4 states", 1.386, uniform_four, quiet)
+    return {"skewed": skewed, "balanced": balanced, "uniform_four": uniform_four}
+
+
+# %% 04. Cross-entropy
+def cross_entropies(quiet: bool = False) -> dict:
+    """Section 1.33.3: scoring samples from p with the probabilities of q costs more than H(p)."""
+    cost = cross_entropy(P_REF, Q_MODEL)  # 0.40866 + 0.18326 nats.
+    own = entropy(P_REF)  # 0.17851 + 0.32189 nats, the irreducible part of that cost.
+    same = cross_entropy(P_REF, P_REF)  # H(p, p) recovers the entropy exactly.
+    show("H(p, q)", 0.592, cost, quiet)
+    show("H(p)", 0.500, own, quiet)
+    show("excess cost H(p, q) - H(p)", 0.092, cost - own, quiet)
+    show("H(p, p) - H(p)", 0.0, same - own, quiet)
+    return {"cost": cost, "own": own, "excess": cost - own}
+
+
+# %% 05. Kullback-Leibler divergence
+def kl(samples: int, quiet: bool = False) -> dict:
+    """Section 1.34: the split H(p, q) = H(p) + KL(p || q), nonnegativity, and the asymmetry."""
+    forward = divergence(P_REF, Q_MODEL)  # 0.23015 - 0.13863 nats, the price of using q.
+    reverse = divergence(Q_MODEL, P_REF)  # -0.17261 + 0.27726 nats, the other direction.
+    split = entropy(P_REF) + forward  # The boxed identity, which must return H(p, q).
+    logits = torch.randn(max(samples, 1), 2, 8)  # Random logits for pairs of distributions.
+    pairs = torch.softmax(logits, dim=-1)  # Two distributions over eight states per row; [N, 2, 8].
+    drawn = (pairs[:, 0] * (pairs[:, 0] / pairs[:, 1]).log()).sum(-1)  # KL for every pair; [N].
+    violation = float((-drawn).clamp_min(0.0).max())  # The largest breach of KL >= 0 among them.
+    t = torch.linspace(0.01, 5.0, 500)  # The range on which the inequality log t <= t - 1 is checked.
+    show("KL(p || q)", 0.0915, forward, quiet)
+    show("KL(q || p)", 0.1046, reverse, quiet)
+    show("H(p) + KL(p || q) = H(p, q)", 0.592, split, quiet)
+    show("KL(p || p)", 0.0, divergence(P_REF, P_REF), quiet)
+    show("reverse over forward, percent", 14.0, 100.0 * (reverse / forward - 1.0), quiet)
+    show("largest log t - (t - 1)", 0.0, float((t.log() - (t - 1.0)).max()), quiet)
+    show(f"largest KL breach over {len(drawn)} pairs", 0.0, violation, quiet)
+    return {"forward": forward, "reverse": reverse, "split": split,
+            "smallest": float(drawn.min()), "violation": violation}
+
+
+# %% 06. Cross-entropy by Monte Carlo
+def monte_carlo(samples: int, repeats: int, quiet: bool = False) -> dict:
+    """Section 1.5.3 applied to Section 1.33.3: H(p, q) is an expectation, so an average estimates it.
+
+    We draw states from p and average -log q at the drawn states. The estimate moves with --seed
+    and with --samples, and the exact value it approaches does not. Repeating a ten-sample estimate
+    ``repeats`` times measures the sd that Section 1.5.3 predicts for a ten-sample average.
+    """
+    exact = cross_entropy(P_REF, Q_MODEL)  # The target of the estimator, 0.592 nats.
+    draws = torch.multinomial(P_REF, max(samples, 10), replacement=True)  # [N] states drawn from p.
+    surprise = -Q_MODEL.log()  # The two values -log q(x) the estimator averages over; [2].
+    scored = surprise[draws]  # The summand -log q(x^(n)) at each drawn state; [N].
+    estimate, spread = float(scored.mean()), float(scored.std())  # The average and its per-sample sd.
+    error = spread / math.sqrt(len(draws))  # The sd of the average, which falls like 1/sqrt(N).
+    show("H(p, q) exactly", 0.592, exact, quiet)
+    if not quiet:
+        label = f"H(p, q) from {len(draws)} samples"
+        print(f"  {label:<36s} exact   {exact:>11.5f}   estimate {estimate:>11.5f}"
+              f"   sd {error:.1e}")
+    rounds = max(repeats, 2)  # Two averages is the fewest a standard deviation can be taken over.
+    repeated = torch.multinomial(P_REF, 10 * rounds, replacement=True)  # [10 * rounds] states.
+    averages = surprise[repeated].view(rounds, 10).mean(dim=1)  # One ten-sample estimate per round.
+    exact_spread = float((P_REF @ (surprise - exact).square()).sqrt())  # sd of one scored draw.
+    predicted = exact_spread / math.sqrt(10.0)  # The sd of a ten-sample average, Section 1.5.3.
+    show(f"sd of {rounds} ten-sample estimates", predicted, float(averages.std()), quiet)
+    return {"exact": exact, "estimate": estimate, "error": error, "draws": len(draws),
+            "predicted": predicted, "measured": float(averages.std()), "rounds": rounds}
+
+
+# %% 07. Maximum likelihood as KL minimization
+def maximum_likelihood(quiet: bool = False) -> dict:
+    """Section 1.35: the empirical loss at the frequencies is their entropy, and the gap to any
+    other model is the KL divergence from the frequencies to it."""
+    frequencies = BASE_COUNTS / BASE_COUNTS.sum()  # p_hat = (0.4, 0.3, 0.2, 0.1) over A, C, G, T.
+    uniform = torch.full((4,), 0.25)  # The categorical model that spreads the mass evenly.
+    at_frequencies = cross_entropy(frequencies, frequencies)  # The empirical loss at p_hat.
+    at_uniform = cross_entropy(frequencies, uniform)  # The same data scored by the uniform model.
+    gap = at_uniform - at_frequencies  # The reducible part of the cost, which is a KL divergence.
+    terms = -frequencies * frequencies.log()  # The four terms the chapter adds up; [4].
+    show("term for A, 0.4 log(1/0.4)", 0.36652, float(terms[0]), quiet)
+    show("term for C, 0.3 log(1/0.3)", 0.36119, float(terms[1]), quiet)
+    show("term for G, 0.2 log(1/0.2)", 0.32189, float(terms[2]), quiet)
+    show("term for T, 0.1 log(1/0.1)", 0.23026, float(terms[3]), quiet)
+    show("loss at the frequencies", 1.280, at_frequencies, quiet)
+    show("loss under the uniform model", 1.386, at_uniform, quiet)
+    show("gap between the two losses", 0.1064, gap, quiet)
+    show("KL(p_hat || uniform)", 0.1064, divergence(frequencies, uniform), quiet)
+    return {"frequencies": frequencies, "at_frequencies": at_frequencies,
+            "at_uniform": at_uniform, "gap": gap}
+
+
+# %% 08. The report
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    parser.add_argument("--samples", type=int, default=20000, help="draws in the Monte Carlo checks")
+    parser.add_argument("--steps", type=int, default=5000, help="repeats of the ten-sample estimate")
+    parser.add_argument("--seed", type=int, default=0, help="random seed for the Monte Carlo draws")
+    parser.add_argument("--out", type=str, default="", help="directory for the run's output")
+    parser.add_argument("--quiet", action="store_true", help="print only the final report")
+    args = parser.parse_args()
+
+    set_seed(args.seed)  # Seed first, so the same seed gives the same draws.
+    if not args.quiet:
+        print("Chapter 1, Sections 1.33 to 1.35, worked example by worked example\n")
+    surprise = information(args.quiet)
+    spread = entropies(args.quiet)
+    cost = cross_entropies(args.quiet)
+    gap = kl(args.samples, args.quiet)
+    estimate = monte_carlo(args.samples, args.steps, args.quiet)
+    fitted = maximum_likelihood(args.quiet)
+
+    report = {
+        "I(x) at p(x) = 0.01, nats": surprise["rare"],
+        "H((0.9, 0.1)), nats": spread["skewed"],
+        "H((0.5, 0.5)), nats": spread["balanced"],
+        "H(p), nats": cost["own"],
+        "H(p, q), nats": cost["cost"],
+        "KL(p || q), nats": gap["forward"],
+        "KL(q || p), nats": gap["reverse"],
+        "smallest KL over random pairs": gap["smallest"],
+        "largest KL breach, random pairs": gap["violation"],
+        "H(p, q) from samples": estimate["estimate"],
+        "Monte Carlo samples drawn": estimate["draws"],
+        "ten-sample estimates taken": estimate["rounds"],
+        "sd of those estimates": estimate["measured"],
+        "sd Section 1.5.3 predicts": estimate["predicted"],
+        "empirical loss at p_hat, nats": fitted["at_frequencies"],
+        "loss under the uniform, nats": fitted["at_uniform"],
+        "KL(p_hat || uniform), nats": fitted["gap"],
+    }
+    print_report("Information, Sections 1.33 to 1.35", report)
+    save_run(args.out, config=vars(args), report=report)
+
+
+if __name__ == "__main__":  # Execute when launched directly.
+    main()  # Work every example of Sections 1.33 to 1.35 and print the report.

lecture_1/linear_algebra.py

@@ -0,0 +1,253 @@
+"""The simplex, the softmax, transition matrices and the matrix exponential.
+
+A distribution over K states is a point of the simplex, a softmax maps unrestricted logits onto
+that simplex, and a row-stochastic matrix moves one point of it to the next. Along an
+eigendirection a matrix power is a power of a number, and that is where the stationary
+distribution, the rate at which a chain reaches it, and the matrix exponential of continuous time
+all come from.
+
+Chapter 1, Sections 1.12, 1.13, 1.15, 1.16, 1.17 and 1.18 of the notes. Runs on the three-base
+simplex of Section 1.12.2, the two-state chain of Section 1.15.1 and the symmetric matrix
+A = [[2, 1], [1, 2]] of Section 1.16. Prints the chapter's number beside the computed one.
+
+    python lecture_1/linear_algebra.py             # every worked example
+    python lecture_1/linear_algebra.py --steps 40  # run the chain further toward p*
+    python lecture_1/linear_algebra.py --quiet     # the closing report only
+"""
+
+# %% 01. Imports and configuration
+from __future__ import annotations
+
+import argparse
+import math
+
+import torch
+
+from cis6270.runner import print_report, save_run, set_seed
+
+torch.set_default_dtype(torch.float64)  # Chapter 1 is arithmetic, and five printed digits need it.
+
+
+P_BASES = torch.tensor([0.5, 0.3, 0.2])  # The three-state distribution over A, C, G; Section 1.12.2.
+H_BASES = torch.tensor([1.0, 0.0, 4.0])  # Hydrogen-bond contacts per base, the property h.
+Q_BASES = torch.tensor([0.1, 0.1, 0.8])  # The second point of the simplex, mixed with P_BASES.
+LOGITS = torch.tensor([2.0, 1.0, 0.0])  # The logit vector z of Section 1.13.1.
+P_CHAIN = torch.tensor([[0.8, 0.2],  # Row 1: stay in state one with 0.8, move to state two with 0.2.
+                        [0.3, 0.7]])  # Row 2: move to state one with 0.3, stay with 0.7.
+A_MATRIX = torch.tensor([[2.0, 1.0], [1.0, 2.0]])  # The running matrix of Sections 1.16 and 1.18.
+X_START = torch.tensor([3.0, 1.0])  # The initial vector x_0 = 2 v_1 + v_2 of Section 1.16.2.
+
+
+WORKED: dict[str, float] = {}  # Every computed value, by label, for tests/test_lecture_1.py.
+
+
+def show(label: str, by_hand: float, computed: float, quiet: bool = False) -> None:
+    """Print the chapter's hand-worked value beside the value this file computes, and keep it."""
+    WORKED[label] = computed  # The test suite asserts these against the chapter's own numbers.
+    if not quiet:  # Under --quiet only the closing report is printed.
+        computed = computed + 0.0  # Adding zero prints a negative zero as a plain zero.
+        print(f"  {label:<36s} chapter {by_hand:>11.5f}   computed {computed:>11.5f}"
+              f"   gap {abs(by_hand - computed):.1e}")
+
+
+# %% 02. Probability vectors and the simplex
+def simplex(quiet: bool = False) -> dict:
+    """Sections 1.12.2 and 1.12.3: expectation as a dot product, and a mixture of two points."""
+    expectation = float(P_BASES @ H_BASES)  # E[h(X)] = p^T h = 0.5 + 0 + 0.8 = 1.3.
+    mixture = 0.5 * P_BASES + 0.5 * Q_BASES  # r = lambda p + (1 - lambda) q at lambda = 0.5; [3].
+    corners = 0.3 * torch.tensor([1.0, 0.0]) + 0.7 * torch.tensor([0.0, 1.0])  # Two-state mixture.
+    show("sum_i p_i", 1.0, float(P_BASES.sum()), quiet)
+    show("E[h(X)] = p^T h", 1.3, expectation, quiet)
+    show("mixture r_1 at lambda = 0.5", 0.30, float(mixture[0]), quiet)
+    show("mixture r_2 at lambda = 0.5", 0.20, float(mixture[1]), quiet)
+    show("mixture r_3 at lambda = 0.5", 0.50, float(mixture[2]), quiet)
+    show("sum_i r_i, still a distribution", 1.0, float(mixture.sum()), quiet)
+    show("0.3 e_1 + 0.7 e_2, first entry", 0.3, float(corners[0]), quiet)
+    return {"expectation": expectation, "mixture": mixture}
+
+
+# %% 03. Logits and the softmax
+def softmax(quiet: bool = False) -> dict:
+    """Section 1.13.1: softmax(z)_i = e^{z_i} / sum_j e^{z_j} on the logits z = (2, 1, 0)."""
+    exponentials = LOGITS.exp()  # e^2, e^1, e^0; [3].
+    partition = float(exponentials.sum())  # The normalizing sum, 7.389 + 2.718 + 1.000.
+    probabilities = exponentials / partition  # The point of the simplex the logits name; [3].
+    shifted = torch.softmax(LOGITS + 5.0, dim=0)  # Adding a constant c cancels in the ratio.
+    show("e^{z_1}", 7.389, float(exponentials[0]), quiet)
+    show("e^{z_2}", 2.718, float(exponentials[1]), quiet)
+    show("e^{z_3}", 1.000, float(exponentials[2]), quiet)
+    show("partition function sum_j e^{z_j}", 11.107, partition, quiet)
+    show("softmax(z)_1", 0.665, float(probabilities[0]), quiet)
+    show("softmax(z)_2", 0.245, float(probabilities[1]), quiet)
+    show("softmax(z)_3", 0.090, float(probabilities[2]), quiet)
+    show("sum_i softmax(z)_i", 1.0, float(probabilities.sum()), quiet)
+    show("shift invariance, largest change", 0.0, float((shifted - probabilities).abs().max()), quiet)
+    return {"probabilities": probabilities, "partition": partition}
+
+
+# %% 04. A probability vector under a transition matrix
+def one_transition(quiet: bool = False) -> dict:
+    """Section 1.15.1: p_{t+1} = P^T p_t moves an even split to (0.55, 0.45) in one step."""
+    p_t = torch.tensor([0.5, 0.5])  # The current distribution over the two states; [2].
+    p_next = P_CHAIN.T @ p_t  # Each new entry collects the mass arriving from every state; [2].
+    show("row sums of P, row 1", 1.0, float(P_CHAIN[0].sum()), quiet)
+    show("row sums of P, row 2", 1.0, float(P_CHAIN[1].sum()), quiet)
+    show("p_{t+1}(1)", 0.55, float(p_next[0]), quiet)
+    show("p_{t+1}(2)", 0.45, float(p_next[1]), quiet)
+    show("sum_j p_{t+1}(j)", 1.0, float(p_next.sum()), quiet)
+    return {"p_next": p_next}
+
+
+# %% 05. Eigenvalues and eigenvectors
+def eigen(quiet: bool = False) -> dict:
+    """Section 1.16: the characteristic equation of A, its two eigenvalues and its eigenvectors."""
+    trace = float(A_MATRIX.trace())  # The sum of the eigenvalues, 4.
+    determinant = float(torch.linalg.det(A_MATRIX))  # The product of the eigenvalues, 3.
+    values, vectors = torch.linalg.eigh(A_MATRIX)  # A is symmetric, so the eigenvalues are real.
+    order = values.argsort(descending=True)  # The chapter indexes lambda_1 = 3 before lambda_2 = 1.
+    values, vectors = values[order], vectors[:, order]  # [2] and [2, 2].
+    vectors = vectors / vectors[0]  # Rescale each one to first entry one, which here gives (1, ±1).
+    residual = (A_MATRIX @ vectors - vectors * values).abs().max()  # A v - lambda v at both pairs.
+    show("trace A, lambda_1 + lambda_2", 4.0, trace, quiet)
+    show("det A, lambda_1 lambda_2", 3.0, determinant, quiet)
+    show("lambda_1", 3.0, float(values[0]), quiet)
+    show("lambda_2", 1.0, float(values[1]), quiet)
+    show("v_1 second entry over first", 1.0, float(vectors[1, 0]), quiet)
+    show("v_2 second entry over first", -1.0, float(vectors[1, 1]), quiet)
+    show("largest |A v - lambda v|", 0.0, float(residual), quiet)
+    return {"values": values, "vectors": vectors}
+
+
+# %% 06. Powers along an eigendirection
+def powers(values: torch.Tensor, vectors: torch.Tensor, quiet: bool = False) -> dict:
+    """Section 1.16.2: x_n = c_1 lambda_1^n v_1 + c_2 lambda_2^n v_2, against A^n x_0 taken directly."""
+    coefficients = torch.linalg.solve(vectors, X_START)  # c_1 = 2 and c_2 = 1 for x_0 = (3, 1).
+    eigen_route = vectors @ (coefficients * values ** 2)  # The n = 2 state along the two directions.
+    direct = torch.linalg.matrix_power(A_MATRIX, 2) @ X_START  # A^2 x_0, squaring the matrix first.
+    far = vectors @ (coefficients * values ** 10)  # The same decomposition carried to n = 10.
+    show("c_1", 2.0, float(coefficients[0]), quiet)
+    show("c_2", 1.0, float(coefficients[1]), quiet)
+    show("lambda_1^10", 59049.0, float(values[0]) ** 10, quiet)
+    show("x_2 first entry, eigen route", 19.0, float(eigen_route[0]), quiet)
+    show("x_2 second entry, eigen route", 17.0, float(eigen_route[1]), quiet)
+    show("x_2 first entry, A^2 x_0", 19.0, float(direct[0]), quiet)
+    show("x_10 first entry", 118099.0, float(far[0]), quiet)
+    show("x_10 second entry", 118097.0, float(far[1]), quiet)
+    return {"x_two": eigen_route, "x_ten": far}
+
+
+# %% 07. Diagonalization and the dominant eigenvalue
+def diagonalization(values: torch.Tensor, vectors: torch.Tensor, quiet: bool = False) -> dict:
+    """Sections 1.16.3 and 1.16.4: A = V Lambda V^{-1}, so A^n = V Lambda^n V^{-1} telescopes."""
+    inverse = torch.linalg.inv(vectors)  # V^{-1} = 0.5 [[1, 1], [1, -1]]; [2, 2].
+    rebuilt = vectors @ torch.diag(values) @ inverse  # V Lambda V^{-1} returns A itself; [2, 2].
+    cube = vectors @ torch.diag(values ** 3) @ inverse  # V diag(27, 1) V^{-1}, the n = 3 power.
+    direct_cube = torch.linalg.matrix_power(A_MATRIX, 3)  # A^3 computed by repeated multiplication.
+    dominant = torch.tensor([2.0, 0.5]) ** 10  # The coefficients 2^10 and 0.5^10 of Section 1.16.4.
+    show("V^{-1}_11", 0.5, float(inverse[0, 0]), quiet)
+    show("largest |V Lambda V^{-1} - A|", 0.0, float((rebuilt - A_MATRIX).abs().max()), quiet)
+    show("A^3_11 through the eigenvalues", 14.0, float(cube[0, 0]), quiet)
+    show("A^3_12 through the eigenvalues", 13.0, float(cube[0, 1]), quiet)
+    show("largest |V Lambda^3 V^{-1} - A^3|", 0.0, float((cube - direct_cube).abs().max()), quiet)
+    show("dominant coefficient 2^10", 1024.0, float(dominant[0]), quiet)
+    show("decaying coefficient 0.5^10", 0.000977, float(dominant[1]), quiet)
+    show("ratio of the two, in millions", 1.05, float(dominant[0] / dominant[1]) / 1e6, quiet)
+    return {"inverse": inverse, "cube": cube}
+
+
+# %% 08. Stationary distributions
+def stationary(steps: int, quiet: bool = False) -> dict:
+    """Section 1.17: p* solves (P^T - I) p* = 0 with sum one, and the chain halves its deviation."""
+    drift = (P_CHAIN.T - torch.eye(2))[0]  # The first row of P^T - I, which reads 0.3 p_2 = 0.2 p_1.
+    system = torch.stack([drift, torch.ones(2)])  # That one equation, with normalization; [2, 2].
+    p_star = torch.linalg.solve(system, torch.tensor([0.0, 1.0]))  # The lambda = 1 eigenvector; [2].
+    trace, determinant = float(P_CHAIN.trace()), float(torch.linalg.det(P_CHAIN))  # 1.5 and 0.5.
+    second = (trace - math.sqrt(trace ** 2 - 4.0 * determinant)) / 2.0  # The smaller root, 0.5.
+    p_t, history = torch.tensor([0.5, 0.5]), []  # Start from the even split of Section 1.15.1.
+    for _ in range(max(steps, 1)):  # Iterate p_{t+1} = P^T p_t for as many steps as requested.
+        p_t = P_CHAIN.T @ p_t  # One more application of the transition matrix; [2].
+        history.append(p_t.clone())
+    show("p*_1", 0.6, float(p_star[0]), quiet)
+    show("p*_2", 0.4, float(p_star[1]), quiet)
+    show("largest |P^T p* - p*|", 0.0, float((P_CHAIN.T @ p_star - p_star).abs().max()), quiet)
+    show("trace P", 1.5, trace, quiet)
+    show("det P", 0.5, determinant, quiet)
+    show("second eigenvalue, the rate", 0.5, second, quiet)
+    for index, value in ((0, 0.55), (1, 0.575), (2, 0.5875), (9, 0.59990)):
+        if index < len(history):  # The chapter prints p_1, p_2, p_3 and p_10 of this sequence.
+            show(f"p_{index + 1}(1)", value, float(history[index][0]), quiet)
+    deviation = float((history[-1] - p_star).abs().max())  # |d_n| = 0.5^n |d_0| after n steps.
+    show(f"|d_n| after {len(history)} steps", 0.1 * second ** len(history), deviation, quiet)
+    return {"p_star": p_star, "second": second, "history": history, "deviation": deviation}
+
+
+# %% 09. From discrete steps to continuous time
+def matrix_exponential(values: torch.Tensor, vectors: torch.Tensor, quiet: bool = False) -> dict:
+    """Section 1.18.3: the four-term series and the eigenvalue formula for e^{tA} at t = 0.1."""
+    t = 0.1  # The small time at which the chapter truncates the series.
+    series, power, factorial = torch.eye(2), torch.eye(2), 1.0  # Accumulate I + tA + t^2A^2/2 + ...
+    for k in range(1, 4):  # Terms k = 1, 2, 3, which with the identity make four.
+        power, factorial = power @ A_MATRIX, factorial * k  # A^k and k factorial.
+        series = series + t ** k * power / factorial  # Add t^k A^k / k!; [2, 2].
+    discarded = t ** 4 * (power @ A_MATRIX) / 24.0  # The first omitted term, t^4 A^4 / 4!.
+    scaled = torch.diag((t * values).exp())  # e^{t Lambda}, the diagonal of entry exponentials; [2, 2].
+    eigen_route = vectors @ scaled @ torch.linalg.inv(vectors)  # V e^{t Lambda} V^{-1}; [2, 2].
+    gap = float((series - eigen_route).abs().max())  # How far four terms fall short at t = 0.1.
+    coefficients = torch.linalg.solve(vectors, X_START)  # c_1 = 2, c_2 = 1, as in Section 1.16.2.
+    x_one = vectors @ (coefficients * values.exp())  # x_t = sum_i c_i e^{t lambda_i} v_i at t = 1.
+    show("four-term series, entry 11", 1.22733, float(series[0, 0]), quiet)
+    show("four-term series, entry 12", 0.12217, float(series[0, 1]), quiet)
+    show("eigenvalue route, entry 11", 1.22751, float(eigen_route[0, 0]), quiet)
+    show("eigenvalue route, entry 12", 0.12234, float(eigen_route[0, 1]), quiet)
+    show("largest gap between the routes", 1.8e-4, gap, quiet)
+    show("first discarded term, entry 11", 1.71e-4, float(discarded[0, 0]), quiet)
+    show("first discarded term, entry 12", 1.67e-4, float(discarded[0, 1]), quiet)
+    show("x_1 first entry, 2e^3 + e", 42.889, float(x_one[0]), quiet)
+    show("x_1 second entry, 2e^3 - e", 37.453, float(x_one[1]), quiet)
+    show("decay e^{-2t} at t = 1", 0.1353, math.exp(-2.0), quiet)
+    show("decay e^{-2t} at t = 3", 0.002479, math.exp(-6.0), quiet)
+    show("sqrt(lambda_1), Gaussian long axis", 1.732, math.sqrt(float(values[0])), quiet)
+    return {"series": series, "eigen_route": eigen_route, "gap": gap, "x_one": x_one}
+
+
+# %% 10. The report
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    parser.add_argument("--steps", type=int, default=10, help="iterations of p_{t+1} = P^T p_t")
+    parser.add_argument("--seed", type=int, default=0, help="random seed, for a repeatable run")
+    parser.add_argument("--out", type=str, default="", help="directory for the run's output")
+    parser.add_argument("--quiet", action="store_true", help="print only the final report")
+    args = parser.parse_args()
+
+    set_seed(args.seed)  # Seed first, so the same seed gives the same run.
+    if not args.quiet:
+        print("Chapter 1, Sections 1.12 to 1.18, worked example by worked example\n")
+    triangle = simplex(args.quiet)
+    logits = softmax(args.quiet)
+    step = one_transition(args.quiet)
+    spectrum = eigen(args.quiet)
+    grown = powers(spectrum["values"], spectrum["vectors"], args.quiet)
+    diagonalization(spectrum["values"], spectrum["vectors"], args.quiet)
+    chain = stationary(args.steps, args.quiet)
+    flow = matrix_exponential(spectrum["values"], spectrum["vectors"], args.quiet)
+
+    report = {
+        "E[h(X)] on the three bases": triangle["expectation"],
+        "softmax(z)_1": float(logits["probabilities"][0]),
+        "p_{t+1}(1) after one step": float(step["p_next"][0]),
+        "lambda_1 of A": float(spectrum["values"][0]),
+        "lambda_2 of A": float(spectrum["values"][1]),
+        "x_10 first entry": float(grown["x_ten"][0]),
+        "p*_1, stationary": float(chain["p_star"][0]),
+        "second eigenvalue of P": chain["second"],
+        "chain steps run": len(chain["history"]),
+        "deviation from p* at the end": chain["deviation"],
+        "e^{tA}_11 at t = 0.1": float(flow["eigen_route"][0, 0]),
+        "four-term truncation error": flow["gap"],
+    }
+    print_report("Linear algebra, Sections 1.12 to 1.18", report)
+    save_run(args.out, config=vars(args), report=report)
+
+
+if __name__ == "__main__":  # Execute when launched directly.
+    main()  # Work every example of Sections 1.12 to 1.18 and print the report.

lecture_1/optimization.py

@@ -0,0 +1,223 @@
+"""Unconstrained descent, Lagrange multipliers, maximum entropy and the KKT conditions.
+
+A differentiable minimum has a zero gradient, and a descent step moves against that gradient. An
+equality constraint comes with a multiplier, and at the constrained optimum the two gradients are
+parallel. Applied to entropy under an expectation constraint, that condition leads to an
+exponential family. An inequality comes with a sign requirement and complementary slackness, and
+under the softmax the simplex constraints hold by construction.
+
+Chapter 1, Sections 1.36, 1.37 and 1.38 of the notes. Runs on the quadratic loss of
+Section 1.36, the constrained distance of Section 1.37.1, the GC-content target of
+Section 1.37.3 and the bounded quadratic of Section 1.38.
+
+    python lecture_1/optimization.py              # every worked example
+    python lecture_1/optimization.py --steps 20   # a short descent, far from the minimizer
+    python lecture_1/optimization.py --quiet      # the closing report only
+"""
+
+# %% 01. Imports and configuration
+from __future__ import annotations
+
+import argparse
+import math
+
+import torch
+
+from cis6270.runner import print_report, save_run, set_seed
+
+torch.set_default_dtype(torch.float64)  # Chapter 1 is arithmetic, and five printed digits need it.
+
+
+BASE_COUNTS = torch.tensor([4.0, 3.0, 2.0, 1.0])  # n(A), n(C), n(G), n(T) over ten DNA bases.
+GC_INDICATOR = torch.tensor([0.0, 1.0, 1.0, 0.0])  # h(x) = 1 on the strong bases C and G.
+GC_TARGET = 0.6  # The required expected GC content, the constraint value c of Section 1.37.3.
+
+
+WORKED: dict[str, float] = {}  # Every computed value, by label, for tests/test_lecture_1.py.
+
+
+def show(label: str, by_hand: float, computed: float, quiet: bool = False) -> None:
+    """Print the chapter's hand-worked value beside the value this file computes, and keep it."""
+    WORKED[label] = computed  # The test suite asserts these against the chapter's own numbers.
+    if not quiet:  # Under --quiet only the closing report is printed.
+        computed = computed + 0.0  # Adding zero prints a negative zero as a plain zero.
+        print(f"  {label:<36s} chapter {by_hand:>11.5f}   computed {computed:>11.5f}"
+              f"   gap {abs(by_hand - computed):.1e}")
+
+
+def entropy(p: torch.Tensor) -> float:
+    """H(p) = -sum_i p_i log p_i, the quantity maximized in both constrained problems below."""
+    return float(-(p * p.clamp_min(1e-300).log()).sum())  # Clamp keeps 0 log 0 at zero.
+
+
+# %% 02. Optimization without constraints
+def unconstrained(steps: int, quiet: bool = False) -> dict:
+    """Section 1.36: on L(theta) = (theta - 3)^2 the error contracts by 0.8 at every step."""
+    learning_rate, theta = 0.1, 0.0  # The step size eta and the starting parameter theta_0.
+    history = []  # The iterates, so the first three can be compared with the chapter.
+    for _ in range(max(steps, 1)):
+        theta = theta - learning_rate * 2.0 * (theta - 3.0)  # theta - eta L'(theta) = 0.8 theta + 0.6.
+        history.append(theta)
+    closed_form = 3.0 - 3.0 * 0.8 ** len(history)  # theta_k = 3 - 3(0.8)^k, from the contraction.
+    for index, value in ((0, 0.600), (1, 1.080), (2, 1.464)):
+        if index < len(history):  # The chapter prints theta_1, theta_2 and theta_3.
+            show(f"theta_{index + 1}", value, history[index], quiet)
+    show(f"theta_{len(history)} from the closed form", closed_form, history[-1], quiet)
+    show("gradient at the minimizer theta = 3", 0.0, 2.0 * (3.0 - 3.0), quiet)
+    return {"history": history, "closed_form": closed_form}
+
+
+# %% 03. An equality constraint and its multiplier
+def equality_constraint(quiet: bool = False) -> dict:
+    """Section 1.37.1: minimize x_1^2 + x_2^2 on the line x_1 + x_2 = 1, through the Lagrangian."""
+    system = torch.tensor([[2.0, 0.0, 1.0],  # 2 x_1 + lambda = 0, the first stationarity condition.
+                           [0.0, 2.0, 1.0],  # 2 x_2 + lambda = 0, the second.
+                           [1.0, 1.0, 0.0]])  # x_1 + x_2 = 1, the constraint itself.
+    solution = torch.linalg.solve(system, torch.tensor([0.0, 0.0, 1.0]))  # (x_1, x_2, lambda); [3].
+    optimum = float(solution[:2].square().sum())  # f at the constrained minimizer.
+    endpoint = float(torch.tensor([1.0, 0.0]).square().sum())  # f at either end of the line.
+    grad_f = 2.0 * solution[:2]  # grad f = (2 x_1, 2 x_2) at the solution; [2].
+    parallel = float((grad_f + solution[2] * torch.ones(2)).abs().max())  # grad f + lambda grad g.
+    show("x_1 at the constrained minimum", 0.5, float(solution[0]), quiet)
+    show("x_2 at the constrained minimum", 0.5, float(solution[1]), quiet)
+    show("lambda, the equality multiplier", -1.0, float(solution[2]), quiet)
+    show("f at the minimum", 0.5, optimum, quiet)
+    show("f at the endpoint (1, 0)", 1.0, endpoint, quiet)
+    show("largest |grad f + lambda grad g|", 0.0, parallel, quiet)
+    return {"solution": solution, "optimum": optimum}
+
+
+# %% 04. Normalization as a constraint
+def normalization(quiet: bool = False) -> dict:
+    """Section 1.37.2: entropy under normalization alone is maximized by the uniform distribution,
+    and the empirical negative log-likelihood is minimized by the observed frequencies."""
+    states = 4  # K, the size of the DNA alphabet.
+    uniform = torch.full((states,), 1.0 / states)  # p_i = e^{lambda - 1}, the same at every state.
+    frequencies = BASE_COUNTS / BASE_COUNTS.sum()  # p(x) = n(x) / N, the solution at lambda = 1.
+    loss = float(-(frequencies * frequencies.log()).sum())  # The minimum empirical loss, its entropy.
+    nearby = frequencies + torch.tensor([0.05, -0.05, 0.0, 0.0])  # A different categorical model.
+    worse = float(-(frequencies * nearby.log()).sum())  # Its loss on the same ten bases.
+    show("uniform p_i at K = 4", 0.25, float(uniform[0]), quiet)
+    show("max entropy under normalization", 1.386, entropy(uniform), quiet)
+    show("p_hat(A) = n(A) / N", 0.4, float(frequencies[0]), quiet)
+    show("p_hat(T) = n(T) / N", 0.1, float(frequencies[3]), quiet)
+    show("minimum empirical loss", 1.280, loss, quiet)
+    show("loss drop from moving off p_hat", 0.0, max(0.0, loss - worse), quiet)
+    if not quiet:  # The displaced model is not a chapter number, so it is printed on its own.
+        print(f"  {'loss of a model moved by 0.05':<36s} loss    {worse:>11.5f}"
+              f"   at p_hat {loss:>11.5f}")
+    return {"uniform": uniform, "frequencies": frequencies, "loss": loss, "worse": worse}
+
+
+# %% 05. Maximum entropy under a property constraint
+def maximum_entropy(samples: int, quiet: bool = False) -> dict:
+    """Section 1.37.3: with one multiplier per property the entropy maximizer is p* proportional to
+    exp(lambda h(x)), and the constraint fixes lambda."""
+    strong, weak = float(GC_INDICATOR.sum()), float((1.0 - GC_INDICATOR).sum())  # Two bases each.
+    exp_lambda = (GC_TARGET / (1.0 - GC_TARGET)) * (weak / strong)  # From 2e^L/(2e^L + 2) = 0.6.
+    multiplier = math.log(exp_lambda)  # lambda = log 1.5, the tilt toward the strong bases.
+    unnormalized = (multiplier * GC_INDICATOR).exp()  # exp(lambda h(x)) at the four bases; [4].
+    partition = float(unnormalized.sum())  # Z(lambda) = 2 e^lambda + 2 = 5.
+    p_star = unnormalized / partition  # The maximum-entropy distribution (0.2, 0.3, 0.3, 0.2); [4].
+    uniform = torch.full((4,), 0.25)  # What lambda = 0 returns, with no property requirement.
+    weights = torch.rand(max(samples, 1), 2)  # One split per base class, drawn independently; [N, 2].
+    competitors = torch.stack([  # Distributions that also meet the GC target, in A, C, G, T order.
+        (1.0 - GC_TARGET) * weights[:, 0], GC_TARGET * weights[:, 1],
+        GC_TARGET * (1.0 - weights[:, 1]), (1.0 - GC_TARGET) * (1.0 - weights[:, 0])], dim=1)
+    drawn = -(competitors * competitors.clamp_min(1e-300).log()).sum(1)  # H of every competitor; [N].
+    best = float(drawn.max())  # The largest entropy reached by any distribution meeting the target.
+    show("e^lambda", 1.5, exp_lambda, quiet)
+    show("lambda = log 1.5", 0.4055, multiplier, quiet)
+    show("partition function Z", 5.0, partition, quiet)
+    show("p*(A)", 0.2, float(p_star[0]), quiet)
+    show("p*(C)", 0.3, float(p_star[1]), quiet)
+    show("GC content of p*", 0.6, float(p_star @ GC_INDICATOR), quiet)
+    show("H(p*)", 1.366, entropy(p_star), quiet)
+    show("H of the unconstrained uniform", 1.386, entropy(uniform), quiet)
+    show("entropy cost of the constraint", 0.0201, entropy(uniform) - entropy(p_star), quiet)
+    show(f"H above H(p*) over {len(drawn)} draws", 0.0, max(0.0, best - entropy(p_star)), quiet)
+    return {"multiplier": multiplier, "p_star": p_star, "partition": partition, "best_drawn": best}
+
+
+# %% 06. Inequality constraints and the KKT conditions
+def kkt(quiet: bool = False) -> dict:
+    """Section 1.38: on min (x - 3)^2 with x <= 1 the bound is active, so lambda is positive."""
+    x_star = 1.0  # The unconstrained minimizer 3 breaks the bound, so the solution is at it.
+    multiplier = -2.0 * (x_star - 3.0)  # Stationarity 2(x - 3) + lambda = 0 solved at x = 1.
+    slack = x_star - 1.0  # g(x*) = x* - 1, which must be at or below zero.
+    show("x* at the bound", 1.0, x_star, quiet)
+    show("lambda, the inequality multiplier", 4.0, multiplier, quiet)
+    show("g(x*) = x* - 1", 0.0, slack, quiet)
+    show("complementary slackness lambda g", 0.0, multiplier * slack, quiet)
+    show("stationarity 2(x* - 3) + lambda", 0.0, 2.0 * (x_star - 3.0) + multiplier, quiet)
+    show("f at the unconstrained minimum", 0.0, (3.0 - 3.0) ** 2, quiet)
+    show("f under the bound", 4.0, (x_star - 3.0) ** 2, quiet)
+    show("value drop for epsilon = 0.001", 0.004, -((1.001 - 3.0) ** 2 - 4.0), quiet)
+    return {"multiplier": multiplier, "value": (x_star - 3.0) ** 2}
+
+
+# %% 07. Optimization over the simplex through the softmax
+def over_the_simplex(steps: int, quiet: bool = False) -> dict:
+    """Section 1.38.1: under logits the simplex constraints hold at every step, so the constrained
+    likelihood problem of Section 1.37.2 becomes an unconstrained descent that reaches p_hat."""
+    frequencies = BASE_COUNTS / BASE_COUNTS.sum()  # The minimizer the multiplier route produced.
+    logits = torch.zeros(4, requires_grad=True)  # z in R^4, starting at the uniform distribution.
+    learning_rate, losses = 1.0, []  # Plain gradient descent on z, with the loss after each step.
+    for _ in range(max(steps, 1)):
+        model = torch.softmax(logits, dim=0)  # p_i = e^{z_i} / sum_j e^{z_j}, always on the simplex.
+        loss = -(frequencies * model.log()).sum()  # The empirical negative log-likelihood.
+        gradient, = torch.autograd.grad(loss, logits)  # d L / d z, which is softmax(z) - p_hat.
+        with torch.no_grad():  # The update itself needs no graph.
+            logits -= learning_rate * gradient  # One descent step of Section 1.36, now on logits.
+        losses.append(float(loss.detach()))
+    final = torch.softmax(logits.detach(), dim=0)  # The distribution the descent arrived at; [4].
+    show("sum_i p_i at the end", 1.0, float(final.sum()), quiet)
+    show("entries at or below zero", 0.0, float((final <= 0.0).sum()), quiet)
+    show("p(A) after the descent", 0.4, float(final[0]), quiet)
+    show("p(T) after the descent", 0.1, float(final[3]), quiet)
+    show("loss after the descent", 1.280, losses[-1], quiet)
+    show("largest |p - p_hat|", 0.0, float((final - frequencies).abs().max()), quiet)
+    return {"final": final, "losses": losses, "distance": float((final - frequencies).abs().max())}
+
+
+# %% 08. The report
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    parser.add_argument("--steps", type=int, default=600, help="gradient descent steps")
+    parser.add_argument("--samples", type=int, default=4000, help="competing distributions drawn")
+    parser.add_argument("--seed", type=int, default=0, help="random seed for the drawn competitors")
+    parser.add_argument("--out", type=str, default="", help="directory for the run's output")
+    parser.add_argument("--quiet", action="store_true", help="print only the final report")
+    args = parser.parse_args()
+
+    set_seed(args.seed)  # Seed first, so the same seed gives the same competitors.
+    if not args.quiet:
+        print("Chapter 1, Sections 1.36 to 1.38, worked example by worked example\n")
+    descent = unconstrained(args.steps, args.quiet)
+    constrained = equality_constraint(args.quiet)
+    normalized = normalization(args.quiet)
+    tilted = maximum_entropy(args.samples, args.quiet)
+    bounded = kkt(args.quiet)
+    fitted = over_the_simplex(args.steps, args.quiet)
+
+    report = {
+        "descent steps run": len(descent["history"]),
+        "theta after the descent": descent["history"][-1],
+        "x_1 under x_1 + x_2 = 1": float(constrained["solution"][0]),
+        "lambda, equality constraint": float(constrained["solution"][2]),
+        "minimum empirical loss, nats": normalized["loss"],
+        "lambda for GC content 0.6": tilted["multiplier"],
+        "p*(C) under that constraint": float(tilted["p_star"][1]),
+        "H(p*), nats": entropy(tilted["p_star"]),
+        "best drawn H at the same target": tilted["best_drawn"],
+        "lambda, active inequality": bounded["multiplier"],
+        "f under the bound": bounded["value"],
+        "loss after the softmax descent": fitted["losses"][-1],
+        "largest |p - p_hat|": fitted["distance"],
+    }
+    print_report("Optimization, Sections 1.36 to 1.38", report)
+    save_run(args.out, config=vars(args), report=report)
+
+
+if __name__ == "__main__":  # Execute when launched directly.
+    main()  # Work every example of Sections 1.36 to 1.38 and print the report.

lecture_1/probability.py

@@ -0,0 +1,261 @@
+"""Probability on discrete and continuous states, in the chapter's own numbers.
+
+A distribution allocates one unit of mass over the states, expectation averages a property under
+that allocation, variance and covariance record how the values fluctuate, and conditioning and
+Bayes' rule revise the allocation once something has been observed. The chapter works every
+example below by hand, so each section prints the hand value beside the value this file computes.
+
+Chapter 1, Sections 1.3 to 1.11 of the notes. Runs on the two-token DNA distribution of
+Section 1.3.1, the three peptides of Section 1.7.1 and the motif-and-binding table of
+Section 1.9.1. Prints the chapter's number, the computed number and the gap between them.
+
+    python lecture_1/probability.py                   # every worked example
+    python lecture_1/probability.py --samples 100000  # a longer Monte Carlo ladder
+    python lecture_1/probability.py --quiet           # the closing report only
+"""
+
+# %% 01. Imports and configuration
+from __future__ import annotations
+
+import argparse
+import math
+
+import torch
+
+from cis6270.runner import print_report, save_run, set_seed
+
+torch.set_default_dtype(torch.float64)  # Chapter 1 is arithmetic, and five printed digits need it.
+
+
+STATES = ["AA", "AC", "CA", "CC"]  # The state space X of Section 1.3.1, with L = 2 over {A, C}.
+P_X = torch.tensor([0.4, 0.1, 0.2, 0.3])  # p_X over those four states, in the order of STATES.
+A_COUNT = torch.tensor([2.0, 1.0, 1.0, 0.0])  # h(x), the number of A tokens in each state.
+PEPTIDES = torch.tensor([[1.0, 2.0], [2.0, 3.0], [3.0, 7.0]])  # (hydrophobicity, binding score).
+JOINT = torch.tensor([[0.32, 0.08],  # Row x = 0, motif absent: columns are y = 0 and y = 1.
+                      [0.24, 0.36]])  # Row x = 1, motif present: Table 1.1 of Section 1.9.1.
+
+
+WORKED: dict[str, float] = {}  # Every computed value, by label, for tests/test_lecture_1.py.
+
+
+def show(label: str, by_hand: float, computed: float, quiet: bool = False) -> None:
+    """Print the chapter's hand-worked value beside the value this file computes, and keep it."""
+    WORKED[label] = computed  # The test suite asserts these against the chapter's own numbers.
+    if not quiet:  # Under --quiet only the closing report is printed.
+        computed = computed + 0.0  # Adding zero prints a negative zero as a plain zero.
+        print(f"  {label:<36s} chapter {by_hand:>11.5f}   computed {computed:>11.5f}"
+              f"   gap {abs(by_hand - computed):.1e}")
+
+
+# %% 02. A distribution over length-two sequences
+def two_token_distribution(quiet: bool = False) -> dict:
+    """Section 1.3.1: the four masses sum to one, and the event {x_1 = A} has probability 0.5."""
+    total = float(P_X.sum())  # Normalization, sum over x in X of p_X(x).
+    event = [index for index, state in enumerate(STATES) if state[0] == "A"]  # A = {x : x_1 = A}.
+    first_is_a = float(P_X[event].sum())  # P(X in A), the masses of {AA, AC} added together.
+    show("sum_x p_X(x)", 1.0, total, quiet)
+    show("P(x_1 = A) = p(AA) + p(AC)", 0.5, first_is_a, quiet)
+    return {"normalization": total, "p_first_is_a": first_is_a}
+
+
+# %% 03. Densities for continuous states
+def continuous_densities(quiet: bool = False) -> dict:
+    """Section 1.4: a density is probability per unit length, so its height may exceed one."""
+    uniform_height = 1.0 / 0.1  # The uniform density on [0, 0.1] has the value 1/0.1 there.
+    narrow_mass = uniform_height * 0.02  # P(0 <= X <= 0.02) is the height times the width.
+    gaussian_peak = 1.0 / (math.sqrt(2.0 * math.pi) * 0.1)  # N(0, 0.1^2) at its center.
+    show("uniform[0, 0.1] density height", 10.0, uniform_height, quiet)
+    show("P(0 <= X <= 0.02)", 0.2, narrow_mass, quiet)
+    show("peak of N(0, 0.1^2)", 3.989, gaussian_peak, quiet)
+    return {"uniform_height": uniform_height, "gaussian_peak": gaussian_peak}
+
+
+# %% 04. The expected adenine count
+def expected_adenine(quiet: bool = False) -> dict:
+    """Section 1.5.1: E[h(X)] = sum_x p_X(x) h(x), the probability-weighted average of the count."""
+    expectation = float(P_X @ A_COUNT)  # 0.4(2) + 0.1(1) + 0.2(1) + 0.3(0).
+    across_thousand = 1000.0 * expectation  # The adenines expected in 1000 generated sequences.
+    show("E[h(X)], expected A count", 1.1, expectation, quiet)
+    show("expected A count in 1000 draws", 1100.0, across_thousand, quiet)
+    return {"expectation": expectation}
+
+
+# %% 05. Monte Carlo estimation
+def monte_carlo(samples: int, repeats: int, quiet: bool = False) -> dict:
+    """Section 1.5.3: the sample average estimates E[h(X)] with an error that shrinks like 1/sqrt(N).
+
+    The chapter's own ten-sample draw gave 1.2 against the exact 1.1, a gap well inside the
+    standard deviation 0.263 of a ten-sample average. We redraw here, so the estimate moves with
+    --seed while the exact value and the error scale stay put. Repeating the ten-sample estimate
+    ``repeats`` times turns that predicted 0.263 into a measured spread.
+    """
+    exact = float(P_X @ A_COUNT)  # The target of the estimator, E[h(X)] = 1.1.
+    second_moment = float(P_X @ A_COUNT.square())  # E[h^2] = 0.4(4) + 0.1 + 0.2 = 1.9.
+    variance = second_moment - exact ** 2  # Var(h(X)) = 1.9 - 1.21 = 0.69.
+    show("E[h(X)^2]", 1.9, second_moment, quiet)
+    show("Var(h(X))", 0.69, variance, quiet)
+    ladder, size = [], 10  # The chapter quotes the error at N = 10 and at N = 1000.
+    while size <= max(samples, 10):
+        draws = torch.multinomial(P_X, size, replacement=True)  # [size] state indices from p_X.
+        estimate = float(A_COUNT[draws].mean())  # (1/N) sum_n h(x^(n)), the Monte Carlo average.
+        ladder.append((size, estimate, math.sqrt(variance / size)))
+        size *= 10
+    if not quiet:
+        for size, estimate, error in ladder:
+            print(f"  N = {size:<8d} estimate {estimate:.4f}   exact {exact:.4f}"
+                  f"   sd of the average {error:.4f}")
+    show("sd of a 10-sample average", 0.263, math.sqrt(variance / 10), quiet)
+    show("sd of a 1000-sample average", 0.0263, math.sqrt(variance / 1000), quiet)
+    rounds = max(repeats, 2)  # Two averages is the fewest a standard deviation can be taken over.
+    repeated = torch.multinomial(P_X, 10 * rounds, replacement=True)  # [10 * rounds] state indices.
+    averages = A_COUNT[repeated].view(rounds, 10).mean(dim=1)  # One ten-sample average per round.
+    measured = float(averages.std())  # The spread of those averages, which 0.263 predicts.
+    show(f"sd of {rounds} ten-sample averages", 0.263, measured, quiet)
+    return {"exact": exact, "variance": variance, "ladder": ladder, "measured": measured,
+            "rounds": rounds}
+
+
+# %% 06. A worked variance
+def worked_variance(quiet: bool = False) -> dict:
+    """Section 1.6.2: Var(X) = E[X^2] - (E[X])^2 on the values 0, 1, 2 with masses 0.2, 0.5, 0.3."""
+    values = torch.tensor([0.0, 1.0, 2.0])  # The three values X can take.
+    masses = torch.tensor([0.2, 0.5, 0.3])  # p_X at those three values.
+    first_moment = float(masses @ values)  # E[X] = 0 + 0.5 + 0.6 = 1.1.
+    second_moment = float(masses @ values.square())  # E[X^2] = 0 + 0.5 + 1.2 = 1.7.
+    from_moments = second_moment - first_moment ** 2  # The boxed identity of Section 1.6.1.
+    from_definition = float(masses @ (values - first_moment).square())  # The two-pass route.
+    show("E[X]", 1.1, first_moment, quiet)
+    show("E[X^2]", 1.7, second_moment, quiet)
+    show("Var(X) from the two moments", 0.49, from_moments, quiet)
+    show("Var(X) from the definition", 0.49, from_definition, quiet)
+    show("sigma_X", 0.7, math.sqrt(from_moments), quiet)
+    return {"variance": from_moments, "sigma": math.sqrt(from_moments)}
+
+
+# %% 07. Covariance, correlation and the covariance matrix
+def peptide_covariance(quiet: bool = False) -> dict:
+    """Sections 1.7.1 and 1.7.2: three equally likely peptides, their correlation and their Sigma."""
+    mu = PEPTIDES.mean(dim=0)  # (mu_X, mu_Y) = (2, 4), each coordinate mean; [2].
+    centered = PEPTIDES - mu  # The deviations (-1, -2), (0, -1), (1, 3); [3, 2].
+    sigma_matrix = centered.T @ centered / PEPTIDES.shape[0]  # E[(X-mu)(X-mu)^T]; [2, 2].
+    cov_xy = float(sigma_matrix[0, 1])  # Cov(X, Y) = (2 + 0 + 3)/3 = 5/3.
+    sd_x = math.sqrt(float(sigma_matrix[0, 0]))  # sigma_X = sqrt(2/3).
+    sd_y = math.sqrt(float(sigma_matrix[1, 1]))  # sigma_Y = sqrt(14/3).
+    rho = cov_xy / (sd_x * sd_y)  # Correlation, which divides out the units of both variables.
+    show("mu_X", 2.0, float(mu[0]), quiet)
+    show("mu_Y", 4.0, float(mu[1]), quiet)
+    show("Cov(X, Y)", 5.0 / 3.0, cov_xy, quiet)
+    show("Var(X) = Sigma_11", 2.0 / 3.0, float(sigma_matrix[0, 0]), quiet)
+    show("Var(Y) = Sigma_22", 14.0 / 3.0, float(sigma_matrix[1, 1]), quiet)
+    show("sigma_X across the peptides", 0.8165, sd_x, quiet)
+    show("sigma_Y across the peptides", 2.160, sd_y, quiet)
+    show("rho_XY", 0.9449, rho, quiet)
+    show("Sigma symmetry, Sigma_21 - Sigma_12", 0.0, float(sigma_matrix[1, 0] - cov_xy), quiet)
+    return {"sigma_matrix": sigma_matrix, "rho": rho}
+
+
+# %% 08. Gaussian noise and the noising equation
+def noising_covariance(quiet: bool = False) -> dict:
+    """Section 1.8: Cov(X_t) = alpha_t^2 Cov(X_0) + sigma_t^2 I at alpha_t = 0.6, sigma_t = 0.8."""
+    alpha_t, sigma_t = 0.6, 0.8  # A variance-preserving pair, since 0.36 + 0.64 = 1.
+    identity = torch.eye(2)  # The isotropic noise covariance; [2, 2].
+    isotropic = alpha_t ** 2 * identity + sigma_t ** 2 * identity  # Data covariance I stays I.
+    wide_data = torch.diag(torch.tensor([4.0, 1.0]))  # Cov(X_0) = diag(4, 1); [2, 2].
+    anisotropic = alpha_t ** 2 * wide_data + sigma_t ** 2 * identity  # diag(2.08, 1.00).
+    show("alpha_t^2 + sigma_t^2", 1.0, alpha_t ** 2 + sigma_t ** 2, quiet)
+    show("Cov(X_t)_11 with Cov(X_0) = I", 1.0, float(isotropic[0, 0]), quiet)
+    show("Cov(X_t)_11 with diag(4, 1)", 2.08, float(anisotropic[0, 0]), quiet)
+    show("Cov(X_t)_22 with diag(4, 1)", 1.00, float(anisotropic[1, 1]), quiet)
+    return {"anisotropic": anisotropic}
+
+
+# %% 09. Joint, marginal and conditional distributions
+def joint_marginal_conditional(quiet: bool = False) -> dict:
+    """Sections 1.9.1 to 1.9.3: the motif-and-binding table, its marginals and its conditionals."""
+    total = float(JOINT.sum())  # 0.36 + 0.24 + 0.08 + 0.32 = 1.
+    p_x = JOINT.sum(dim=1)  # p(x) = sum_y p(x, y), motif present or absent; [2].
+    p_y = JOINT.sum(dim=0)  # p(y) = sum_x p(x, y), high or low binding; [2].
+    p_y_given_x = JOINT / p_x[:, None]  # p(y | x) = p(x, y) / p(x), renormalized row by row; [2, 2].
+    show("sum over the four cells", 1.0, total, quiet)
+    show("p(X = 1), motif present", 0.60, float(p_x[1]), quiet)
+    show("p(X = 0), motif absent", 0.40, float(p_x[0]), quiet)
+    show("p(Y = 1), high binding", 0.44, float(p_y[1]), quiet)
+    show("p(Y = 0), low binding", 0.56, float(p_y[0]), quiet)
+    show("p(Y = 1 | X = 1)", 0.6, float(p_y_given_x[1, 1]), quiet)
+    show("p(Y = 0 | X = 1)", 0.4, float(p_y_given_x[1, 0]), quiet)
+    show("p(Y = 1 | X = 0)", 0.2, float(p_y_given_x[0, 1]), quiet)
+    show("p(Y = 0 | X = 0)", 0.8, float(p_y_given_x[0, 0]), quiet)
+    return {"p_x": p_x, "p_y": p_y, "p_y_given_x": p_y_given_x}
+
+
+# %% 10. Bayes' rule and the law of total probability
+def bayes(quiet: bool = False) -> dict:
+    """Section 1.10: the motif probability after observing high binding, with its own denominator."""
+    p_x = JOINT.sum(dim=1)  # The prior p(x) over motif presence; [2].
+    p_y_given_x = JOINT / p_x[:, None]  # The likelihood p(y | x); [2, 2].
+    evidence = float(p_y_given_x[:, 1] @ p_x)  # p(y) = sum_x p(y | x) p(x), summed over the cause.
+    posterior = float(p_y_given_x[1, 1] * p_x[1]) / evidence  # p(X = 1 | Y = 1) = 0.36 / 0.44.
+    show("evidence p(Y = 1)", 0.44, evidence, quiet)
+    show("posterior p(X = 1 | Y = 1)", 0.818, posterior, quiet)
+    show("prior p(X = 1)", 0.60, float(p_x[1]), quiet)
+    show("likelihood ratio 0.6 / 0.2", 3.0, float(p_y_given_x[1, 1] / p_y_given_x[0, 1]), quiet)
+    return {"evidence": evidence, "posterior": posterior}
+
+
+# %% 11. Independence
+def independence(quiet: bool = False) -> dict:
+    """Section 1.11: p(x, y) equals p(x) p(y) under independence, and the table fails that test."""
+    p_x = JOINT.sum(dim=1)  # Marginal over motif presence; [2].
+    p_y = JOINT.sum(dim=0)  # Marginal over binding activity; [2].
+    product = float(p_x[1] * p_y[1])  # 0.60 times 0.44 = 0.264 under a factorized model.
+    joint_entry = float(JOINT[1, 1])  # The table's own entry, 0.36.
+    show("p(X = 1) p(Y = 1)", 0.264, product, quiet)
+    show("p(X = 1, Y = 1)", 0.36, joint_entry, quiet)
+    show("association missed by factorizing", 0.096, joint_entry - product, quiet)
+    return {"product": product, "joint_entry": joint_entry, "gap": joint_entry - product}
+
+
+# %% 12. The report
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    parser.add_argument("--samples", type=int, default=1000, help="largest Monte Carlo sample size")
+    parser.add_argument("--steps", type=int, default=5000, help="repeats of the ten-sample estimate")
+    parser.add_argument("--seed", type=int, default=0, help="random seed for the Monte Carlo draws")
+    parser.add_argument("--out", type=str, default="", help="directory for the run's output")
+    parser.add_argument("--quiet", action="store_true", help="print only the final report")
+    args = parser.parse_args()
+
+    set_seed(args.seed)  # Seed first, so the same seed gives the same Monte Carlo ladder.
+    if not args.quiet:
+        print("Chapter 1, Sections 1.3 to 1.11, worked example by worked example\n")
+    two_token_distribution(args.quiet)
+    continuous_densities(args.quiet)
+    adenine = expected_adenine(args.quiet)
+    estimates = monte_carlo(args.samples, args.steps, args.quiet)
+    spread = worked_variance(args.quiet)
+    peptides = peptide_covariance(args.quiet)
+    noising = noising_covariance(args.quiet)
+    table = joint_marginal_conditional(args.quiet)
+    posterior = bayes(args.quiet)
+    association = independence(args.quiet)
+
+    report = {
+        "E[h(X)], expected A count": adenine["expectation"],
+        "Monte Carlo estimate at N": estimates["ladder"][-1][0],
+        "Monte Carlo estimate": estimates["ladder"][-1][1],
+        "Var(h(X))": estimates["variance"],
+        "ten-sample averages taken": estimates["rounds"],
+        "sd of those averages": estimates["measured"],
+        "Var(X) on 0, 1, 2": spread["variance"],
+        "rho_XY across three peptides": peptides["rho"],
+        "Cov(X_t)_11 from diag(4, 1)": float(noising["anisotropic"][0, 0]),
+        "p(Y = 1), high binding": float(table["p_y"][1]),
+        "p(X = 1 | Y = 1)": posterior["posterior"],
+        "p(X=1,Y=1) - p(X=1)p(Y=1)": association["gap"],
+    }
+    print_report("Probability, Sections 1.3 to 1.11", report)
+    save_run(args.out, config=vars(args), report=report)
+
+
+if __name__ == "__main__":  # Execute when launched directly.
+    main()  # Work every example of Sections 1.3 to 1.11 and print the report.

lecture_2/README.md

@@ -0,0 +1,99 @@
+# Lecture 2: Flow Matching (Chapter 2)
+
+A velocity field assigns one vector to every state and every time. Following it from a starting
+point traces a trajectory, the trajectories from every starting point form the flow map, and the
+continuity equation ties the motion of the density to the same field. Training such a field by
+maximum likelihood costs an ODE solve and a divergence trace per gradient step, so flow matching
+replaces it with a regression: pair a Gaussian sample with a data sample, join the pair by a
+straight line, and regress a network onto the derivative of that line. Least squares recovers the
+average velocity of all the paths through a state, which is the field that transports the prior
+onto the data, and generation is Euler integration of the learned field.
+
+`flow_matching_mnist.py` is the lecture's main script, the one Section 2.20 of the notes walks
+through line by line. The other three files are the mathematics underneath it: the fields and the
+solver before any training, the full method on a target whose density is known exactly, and the
+family of interpolants the straight line belongs to.
+
+| File | Command | What it does |
+| --- | --- | --- |
+| `flow_matching_mnist.py` | `python lecture_2/flow_matching_mnist.py` | Trains the small U-Net on MNIST and writes a 4-by-4 grid of digits and a trajectory strip. `--dataset toy` runs offline. |
+| `velocity_fields.py` | `python lecture_2/velocity_fields.py` | No training. Trajectories, the flow map, the Euler refinement ladder, the pushforward of a Gaussian, the flux balance and the continuity equation. |
+| `flow_matching_2d.py` | `python lecture_2/flow_matching_2d.py` | Trains a velocity MLP on the four-Gaussian target and samples by Euler. Measures the conditional flow-matching equivalence against the exact marginal velocity. |
+| `interpolants.py` | `python lecture_2/interpolants.py` | Trains the same MLP three times, on the linear, the t-squared and the trigonometric schedule, and compares the samples at a fine and a coarse step budget. |
+
+Nothing is trained in `velocity_fields.py`, so its `--steps` sets the largest rung of the Euler
+refinement ladder. In the other three files `--steps` counts optimizer steps, and in
+`flow_matching_mnist.py` it is the budget for the whole run, which `--epochs` would otherwise set.
+
+## The numbers a default run produces
+
+`velocity_fields.py`, in about two seconds, reproduces the chapter's arithmetic exactly.
+
+```
+euler steps   1  landing 2.000000  error 0.718282  error x steps 0.7183
+euler steps   2  landing 2.250000  error 0.468282  error x steps 0.9366
+euler steps   4  landing 2.441406  error 0.276876  error x steps 1.1075
+flux_balance  -0.3354      enclosed_mass_rate  -0.3354
+continuity_time_derivative  0.1123      continuity_minus_flux_divergence  0.1123
+density_ratio_at_origin  2.0000         log_density_change  -0.6931
+```
+
+`flow_matching_2d.py`, in about six seconds on one core, with 83,202 parameters.
+
+```
+crossing_conditional_loss  [10.0, 9.0, 10.0]     crossing_marginal_loss  [1.0, 0.0, 1.0]
+crossing_constant_gap      [9.0, 9.0, 9.0]       crossing_target_variance  9.0000
+loss_cfm  10.3243   loss_fm  5.1196   loss_gap  5.2506   loss_cross_term  -0.0459
+gradient_cosine  0.9989                          final_loss  5.3169
+modes_found  4      mode_entropy  1.3657         nearest_mode_distance  0.4505
+mean_log_density  -2.2648                        data_mean_log_density  -2.1402
+```
+
+`interpolants.py`, in about thirteen seconds, trains three fields from the same initialization.
+
+```
+                 endpoint_error  covariance_deviation  mean_log_density  coarse_log_density
+linear                   0.0000                0.4999           -2.5276             -3.0309
+quadratic                0.0000                0.4939           -2.5748             -3.6419
+trigonometric            0.0000                0.0179           -2.6621             -2.9566
+```
+
+All three schedules reach the four modes. The trigonometric one is the only schedule whose
+interpolation keeps identity covariance at every time, and it is the one that loses least when the
+sampler drops from 64 Euler steps to 8, while the t-squared schedule, whose velocity grows linearly
+in t, loses most.
+
+`flow_matching_mnist.py` downloads MNIST on its first run. The offline path, `--dataset toy`,
+trains the same 471,265-parameter U-Net on 512 bars and blobs in about two minutes on one core.
+
+```
+epoch=20 squared_norm_loss=977.133
+final loss 977.133  sample range [-5.15, 3.50]
+```
+
+Twenty epochs over 512 images is 80 optimizer steps, and the generated range is still wider than
+the data's [-1, 1]. With `--epochs 100` the loss reaches 387.136 and the range [-3.25, 3.04], in
+about nine minutes. The MNIST default is 20 epochs over 60,000 images, and it is the one script
+here that does not finish while you watch it. `--steps 20` trains for twenty optimizer steps and
+then samples, which is the quick check that the file runs.
+
+## What is simplified
+
+- The marginal velocity `v_t(x) = E[U_t | X_t = x]` is written in closed form in
+  `flow_matching_2d.py`. That is available because the target is a Gaussian mixture and the prior is
+  Gaussian; Section 2.14 of the notes is about the case where it is not, which is why the equivalence
+  it proves is worth proving.
+- The equivalence is measured on one large batch, so the two losses agree up to that batch's own
+  sampling error. The file prints the cross term that accounts for the residual exactly.
+- The identity-covariance check uses two independent standard Gaussian endpoints, which is the
+  hypothesis the chapter states. The four-Gaussian target is centered but not white.
+- `velocity_fields.py` differentiates the densities with autograd where the chapter differentiates
+  by hand, so no derivative in that file is transcribed from the notes. It also runs in double
+  precision, because the Euler ladder it prints is arithmetic quoted to six decimal places.
+- The Gaussian translation of Section 2.17 is left out. Its coupling is deterministic, so the
+  conditional velocity is the constant 4 at every state and time and already equals the marginal
+  velocity.
+- The cubic Hermite interpolant of Section 2.18 is left out of `interpolants.py`. It needs endpoint
+  slopes, which the four-Gaussian data do not carry.
+- The three two-dimensional scripts run on the CPU. Each state there has two coordinates, so a GPU
+  would be slower. `flow_matching_mnist.py` picks CUDA, then Apple MPS, then the CPU.

lecture_2/flow_matching_2d.py

@@ -0,0 +1,245 @@
+"""Flow matching on four Gaussians, the whole method end to end.
+
+A coupling pairs a Gaussian prior sample with a data sample, a straight line joins the pair, and the
+derivative of that line is the regression target. Least squares against those per-pair velocities
+recovers the marginal velocity field, and Euler integration of the learned field turns fresh noise
+into data. The four-Gaussian target is small enough that the marginal velocity has a closed form, so
+the conditional flow-matching equivalence can be measured here and not only argued.
+
+Chapter 2, Sections 2.11 to 2.16 and 2.19 of the notes. Trained on
+``cis6270.data.gaussian_mixture_2d``. Prints the worked endpoint pair, the average at the crossing,
+the two losses with their constant gap, the agreement of their gradients, and the mode coverage and
+mean exact log density of the generated points.
+
+    python lecture_2/flow_matching_2d.py                # train and generate, about ten seconds
+    python lecture_2/flow_matching_2d.py --steps 20     # a first look, with no time to learn
+"""
+
+# %% 01. Imports and configuration
+import argparse
+
+import torch
+
+from cis6270.data import GAUSSIAN_CENTERS, GAUSSIAN_STD, gaussian_mixture_2d
+from cis6270.metrics import mean_log_density, mode_report
+from cis6270.nets import VelocityMLP, count_parameters
+from cis6270.runner import print_report, save_run, set_seed, train
+
+DIM = 2  # d, the width of the state.
+TRAIN_STEPS = 1000  # Optimizer steps, each on one fresh batch of endpoint pairs.
+BATCH_SIZE = 256  # Endpoint pairs per step.
+LEARNING_RATE = 3e-3  # Optimizer step size.
+WIDTH = 128  # Hidden width of the velocity network.
+SAMPLE_STEPS = 64  # Euler steps from t = 0 to t = 1, so Delta t = 1/64.
+NUM_SAMPLES = 1024  # Points generated for the mode coverage and the log density.
+DATA_POINTS = 4096  # Size of the training set drawn from the four-Gaussian target.
+EQUIVALENCE_POINTS = 16384  # Pairs used for the Monte Carlo check of Theorem 2.1.
+GRID = torch.tensor([0.0, 0.25, 0.50, 0.75, 1.0])  # [5] the chapter's grid of times.
+
+
+# %% 02. The coupling
+#
+# Section 2.11. A coupling is any joint law over the endpoint pair whose marginals are the prior and
+# the data. Drawing the two ends with no reference to each other gives the independent coupling.
+def sample_coupling(data, batch_size):
+    """Draw (X_0, X_1) ~ pi with pi(x_0, x_1) = p_0(x_0) p_1(x_1). Returns two [B, d] tensors."""
+    index = torch.randint(0, len(data), (batch_size,), device=data.device)  # [B] which data point.
+    x_1 = data[index]  # [B, d] the data endpoint, drawn from p_1.
+    x_0 = torch.randn_like(x_1)  # [B, d] the prior endpoint, drawn from p_0 = N(0, I).
+    return x_0, x_1
+
+
+# %% 03. The conditional path and its velocity
+#
+# Section 2.12. The straight segment at constant speed, and its derivative with the endpoints held
+# fixed. The velocity is called conditional because only someone who knows both endpoints has it.
+def conditional_path(x_0, x_1, t):
+    """X_t = (1 - t) X_0 + t X_1 and U_t = X_1 - X_0. Returns two [B, d] tensors."""
+    x_t = (1.0 - t)[:, None] * x_0 + t[:, None] * x_1  # [B, d] convex combination of the endpoints.
+    u_t = x_1 - x_0  # [B, d] the displacement, the same vector at every time.
+    return x_t, u_t
+
+
+def worked_pair():
+    """The chapter's one-dimensional pair x_0 = -2, x_1 = 3, evaluated on the grid of times."""
+    x_0 = torch.full((len(GRID), 1), -2.0)  # [5, 1] the same prior endpoint at every grid time.
+    x_1 = torch.full((len(GRID), 1), 3.0)  # [5, 1] the same data endpoint at every grid time.
+    x_t, u_t = conditional_path(x_0, x_1, GRID)  # The line x_t = -2 + 5t and its slope.
+    return {
+        "worked_pair_path": [round(float(v), 4) for v in x_t[:, 0]],
+        "worked_pair_velocity": float(u_t[0, 0]),  # The slope 5, so 1.25 units in every quarter.
+    }
+
+
+# %% 04. The marginal velocity
+#
+# Section 2.13. Several conditional paths pass through one state at one time, each with its own
+# U_t, and a deterministic field assigns one vector there, the average v_t(x) = E[U_t | X_t = x].
+# For this target the average is available exactly, and Section 2.14 is about the case where it
+# is not. Conditioning on component k makes (X_1, X_t) jointly Gaussian, so the posterior over
+# components is a softmax of the component log densities at X_t and v_t follows in closed form.
+def marginal_velocity(x_t, t):
+    """v_t(x) = E[U_t | X_t = x] for a standard Gaussian prior and the four-Gaussian target."""
+    centers = GAUSSIAN_CENTERS.to(x_t)  # [M, d] the four means of the mixture.
+    variance = (1.0 - t) ** 2 + (t * GAUSSIAN_STD) ** 2  # [B] variance of X_t given component k.
+    offset = x_t[:, None, :] - t[:, None, None] * centers  # [B, M, d] distance to each shifted mean.
+    log_weight = -0.5 * (offset ** 2).sum(-1) / variance[:, None]  # [B, M] up to a shared constant.
+    weight = torch.softmax(log_weight, dim=1)  # [B, M] posterior over components given X_t.
+    mean_center = (weight[:, :, None] * centers).sum(1)  # [B, d] posterior mean of the center.
+    # Writing E[X_1 | X_t] in terms of that posterior mean and cancelling the factor (1 - t) that
+    # U_t = (X_1 - X_t)/(1 - t) would otherwise divide by leaves one that stays finite at t = 1.
+    return ((1.0 - t)[:, None] * mean_center
+            - (1.0 - t - t * GAUSSIAN_STD ** 2)[:, None] * x_t) / variance[:, None]  # [B, d].
+
+
+def crossing_average():
+    """Two paths of Section 2.13 that meet at x = 2 at t = 0.5, with velocities 4 and -2."""
+    x_0 = torch.tensor([[0.0], [3.0]])  # [2, 1] the two prior endpoints.
+    x_1 = torch.tensor([[4.0], [1.0]])  # [2, 1] the two data endpoints.
+    x_t, u_t = conditional_path(x_0, x_1, torch.full((2,), 0.5))  # Both states land on x = 2.
+    return {
+        "crossing_state": float(x_t.mean()),  # Both paths occupy this state at t = 0.5.
+        "crossing_velocities": [round(float(v), 4) for v in u_t[:, 0]],
+        "crossing_average_velocity": float(u_t.mean()),  # The equally weighted average, 1.
+    }
+
+
+# %% 05. The two objectives
+#
+# Sections 2.14 and 2.15. The flow-matching loss regresses on the marginal velocity, which is not
+# available for real data, and the conditional loss regresses on the sampled U_t, which is.
+def objectives(model, x_0, x_1, t):
+    """L_CFM, L_FM, the gap E||U_t - v_t||^2 and the cross term, on one batch of endpoint pairs."""
+    x_t, u_t = conditional_path(x_0, x_1, t)  # [B, d] the state and its conditional velocity.
+    v_pred = model(x_t, t)  # [B, d] the network sees the state and the time, never the endpoints.
+    v_t = marginal_velocity(x_t, t)  # [B, d] the target of the loss we cannot evaluate in general.
+    loss_cfm = (v_pred - u_t).square().sum(1).mean()  # E||v_theta(X_t, t) - U_t||^2.
+    loss_fm = (v_pred - v_t).square().sum(1).mean()  # E||v_theta(X_t, t) - v_t(X_t)||^2.
+    gap = (u_t - v_t).square().sum(1).mean()  # The conditional variance of the target.
+    # Expanding the square leaves L_CFM = L_FM + cross + gap on any batch. The proof of Theorem 2.1
+    # rests on the cross term having conditional mean zero, so on a large batch it is near zero.
+    cross = 2.0 * ((v_pred - v_t) * (v_t - u_t)).sum(1).mean()
+    return loss_cfm, loss_fm, gap, cross
+
+
+# %% 06. The equivalence, measured
+#
+# Section 2.15. At the crossing the two losses are parabolas in the prediction a, nine units apart
+# at every a, with one derivative between them. The same identity holds in expectation on the
+# two-dimensional problem, up to the Monte Carlo error of the batch that estimates it.
+def crossing_gap(predictions=(0.0, 1.0, 2.0)):
+    """The conditional and marginal losses at x = 2, t = 0.5, where the targets are 4 and -2."""
+    targets = torch.tensor([4.0, -2.0])  # [2] the two conditional velocities meeting there.
+    marginal = targets.mean()  # The conditional mean of the target, which is 1.
+    conditional_values, marginal_values = [], []
+    for a in predictions:  # One candidate prediction at a time.
+        conditional_values.append(round(float((a - targets).square().mean()), 4))
+        marginal_values.append(round(float((a - marginal) ** 2), 4))
+    return {
+        "crossing_conditional_loss": conditional_values,
+        "crossing_marginal_loss": marginal_values,
+        "crossing_constant_gap": [round(c - m, 4) for c, m in
+                                  zip(conditional_values, marginal_values)],
+        "crossing_target_variance": float(targets.var(unbiased=False)),  # The same constant, 9.
+    }
+
+
+def equivalence_check(model, data, num_points):
+    """Theorem 2.1 on the four-Gaussian problem: the gap in value and the agreement in gradient."""
+    x_0, x_1 = sample_coupling(data, num_points)  # [N, d] pairs from the independent coupling.
+    t = torch.rand(num_points, device=x_1.device)  # [N] one uniform time per pair.
+    loss_cfm, loss_fm, gap, cross = objectives(model, x_0, x_1, t)
+    parameters = list(model.parameters())
+    grad_cfm = torch.autograd.grad(loss_cfm, parameters, retain_graph=True)  # Gradient of L_CFM.
+    grad_fm = torch.autograd.grad(loss_fm, parameters)  # Gradient of L_FM, on the same batch.
+    flat_cfm = torch.cat([g.flatten() for g in grad_cfm])  # [P] one vector per objective.
+    flat_fm = torch.cat([g.flatten() for g in grad_fm])  # [P].
+    return {
+        "loss_cfm": float(loss_cfm.detach()),
+        "loss_fm": float(loss_fm.detach()),
+        "loss_gap": float(gap.detach()),  # The constant that separates the two objectives.
+        "loss_cross_term": float(cross.detach()),  # Zero in expectation, so small on this batch.
+        "loss_fm_plus_gap": float((loss_fm + gap).detach()),  # L_CFM once the cross term is added.
+        "gradient_cosine": float(torch.nn.functional.cosine_similarity(flat_cfm, flat_fm, dim=0)),
+        # What separates the two gradients is the gradient of that cross term, which falls like one
+        # over the square root of the batch size.
+        "gradient_relative_difference": float((flat_cfm - flat_fm).norm() / flat_fm.norm()),
+    }
+
+
+# %% 07. Training
+#
+# Section 2.19. One iteration draws a pair and a time, builds the state in closed form, subtracts
+# the endpoints, and takes a gradient step. No ODE is solved while a training example is built.
+def make_loss(model, data, batch_size):
+    """The closure runner.train calls, returning L_CFM on one fresh batch."""
+    def loss_fn(step):
+        x_0, x_1 = sample_coupling(data, batch_size)  # [B, d] a fresh pair per step.
+        t = torch.rand(batch_size, device=x_1.device)  # [B] t ~ Uniform(0, 1), one per pair.
+        x_t, u_t = conditional_path(x_0, x_1, t)  # [B, d] the state and the regression target.
+        return (model(x_t, t) - u_t).square().sum(1).mean()  # ||v_theta(X_t, t) - U_t||^2.
+    return loss_fn
+
+
+# %% 08. Euler sampling
+#
+# Section 2.19. Generation starts from a fresh prior sample and integrates the learned field, so
+# one sample costs as many forward passes as there are steps.
+@torch.no_grad()
+def sample(model, num_samples, steps):
+    """Integrate dX_t/dt = v_theta(X_t, t) from t = 0 to t = 1 on a uniform grid. Returns [N, d]."""
+    device = next(model.parameters()).device  # Use the model's own device.
+    x_t = torch.randn(num_samples, DIM, device=device)  # [N, d] the draw X_0 ~ p_0.
+    dt = 1.0 / steps  # Delta t, the width of one interval.
+    for step in range(steps):  # One network evaluation per step.
+        t = torch.full((num_samples,), step * dt, device=device)  # [N] the shared grid time.
+        x_t = x_t + dt * model(x_t, t)  # Euler: state += time increment times velocity.
+    return x_t
+
+
+# %% 09. The report
+def main():
+    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    parser.add_argument("--steps", type=int, default=TRAIN_STEPS, help="optimizer steps")
+    parser.add_argument("--batch-size", type=int, default=BATCH_SIZE, help="pairs per step")
+    parser.add_argument("--lr", type=float, default=LEARNING_RATE, help="learning rate")
+    parser.add_argument("--width", type=int, default=WIDTH, help="hidden width of the network")
+    parser.add_argument("--samples", type=int, default=NUM_SAMPLES, help="points to generate")
+    parser.add_argument("--sample-steps", type=int, default=SAMPLE_STEPS, help="Euler steps")
+    parser.add_argument("--seed", type=int, default=0, help="random seed")
+    parser.add_argument("--out", type=str, default="", help="directory for the run's output")
+    parser.add_argument("--quiet", action="store_true", help="print only the final report")
+    args = parser.parse_args()
+
+    set_seed(args.seed)  # Seed first, so the same seed gives the same report.
+    data = gaussian_mixture_2d(DATA_POINTS, seed=args.seed)  # [N, d] the target p_1 = p_data.
+    model = VelocityMLP(dim=DIM, width=args.width)  # v_theta(x, t), the learned velocity field.
+
+    report = {"parameters": count_parameters(model)}
+    report.update(worked_pair())  # Section 2.12.
+    report.update(crossing_average())  # Section 2.13.
+    report.update(crossing_gap())  # Section 2.15.
+    # The gradient comparison runs before training, where the two gradients are far from zero.
+    report.update(equivalence_check(model, data, EQUIVALENCE_POINTS))  # Section 2.15.
+
+    losses = train(model, make_loss(model, data, args.batch_size),
+                   steps=args.steps, lr=args.lr, quiet=args.quiet)
+    generated = sample(model, args.samples, args.sample_steps)  # [N, d] fresh points.
+    coverage = mode_report(generated)  # Nearest mode, how many modes, how evenly they are visited.
+    report.update({
+        "final_loss": sum(losses[-100:]) / len(losses[-100:]),  # Averaged over the last steps.
+        "modes_found": coverage["modes_found"],
+        "mode_entropy": coverage["mode_entropy"],  # log 4 = 1.3863 when all four are equally used.
+        "nearest_mode_distance": coverage["nearest_distance"],
+        "mode_fractions": coverage["mode_fractions"],
+        "mean_log_density": mean_log_density(generated),  # Exact log density of the target.
+        "data_mean_log_density": mean_log_density(data),  # The same measurement on real points.
+    })
+
+    print_report("flow matching on four Gaussians", report)
+    save_run(args.out, model=model, config=vars(args), report=report, losses=losses)
+
+
+# %% 10. Run the script
+if __name__ == "__main__":
+    main()

lecture_2/flow_matching_mnist.py

@@ -0,0 +1,230 @@
+"""Flow matching on MNIST, from noise to a handwritten digit in one script.
+
+We draw a Gaussian sample at time zero and a real image at time one, join them by a straight line,
+and train a U-Net to predict the velocity of that line at a random time in between. Generation then
+starts from fresh noise and takes Euler steps along the learned field. Chapter 2, Section 2.20 of
+the notes walks through this file section by section.
+
+    python lecture_2/flow_matching_mnist.py                  # 20 epochs on MNIST, downloads on first use
+    python lecture_2/flow_matching_mnist.py --steps 20       # 20 optimizer steps, for a first look
+    python lecture_2/flow_matching_mnist.py --epochs 1       # a single pass, about two minutes on a CPU
+    python lecture_2/flow_matching_mnist.py --dataset toy    # no download, for a first look or a test
+
+The network class and the checkpoint keys match the published `flow_unet_mnist.pt`, so the saved
+five-epoch run still loads.
+"""
+
+# %% 01. Imports
+import argparse
+from pathlib import Path
+
+import torch  # Tensors, gradients, and optimization.
+import torch.nn as nn  # Layers and model classes.
+import torch.nn.functional as F  # Resize feature maps in the decoder.
+
+
+# %% 02. Configuration
+BATCH_SIZE = 128  # Images per training batch.
+EPOCHS = 20  # Passes through the training set.
+LEARNING_RATE = 2e-4  # Optimizer step size.
+BASE_CHANNELS = 32  # U-Net widths: 32, 64, 128.
+SAMPLE_STEPS = 100  # Euler steps from t=0 to t=1.
+NUM_SAMPLES = 16  # Generate a 4-by-4 image grid.
+SEED = 7  # Seed for initialization and training.
+DATA_DIR = Path("data")  # MNIST download location.
+OUTPUT_DIR = Path("outputs/flow_matching_mnist")  # Checkpoint and generated image folder.
+
+DEVICE = torch.device(
+    "cuda" if torch.cuda.is_available()  # NVIDIA GPU.
+    else "mps" if torch.backends.mps.is_available()  # Apple GPU.
+    else "cpu"  # CPU fallback.
+)
+
+
+# %% 03. Load real images
+def make_loader(dataset="mnist", batch_size=BATCH_SIZE):
+    """Batches of images normalized to [-1, 1], the range the velocity field is trained on."""
+    from torch.utils.data import DataLoader, TensorDataset
+
+    if dataset == "toy":  # Shapes like MNIST, drawn offline, so a test needs no download.
+        from cis6270.data import toy_images
+
+        images = toy_images(num_images=512, size=28, seed=SEED)  # [512, 1, 28, 28].
+        labels = torch.zeros(len(images), dtype=torch.long)  # Unused; generation is unconditional.
+        return DataLoader(TensorDataset(images, labels), batch_size=batch_size, shuffle=True)
+
+    from torchvision import datasets
+    from torchvision.transforms import v2
+
+    transform = v2.Compose([
+        v2.ToImage(),  # Channel-first image: [1, 28, 28].
+        v2.ToDtype(torch.float32, scale=True),  # Integer pixels 0-255 -> floats 0-1.
+        v2.Normalize(mean=(0.5,), std=(0.5,)),  # Real data pixels: [0, 1] -> [-1, 1].
+    ])
+    mnist = datasets.MNIST(root=DATA_DIR, train=True, download=True, transform=transform)
+    return DataLoader(
+        mnist, batch_size=batch_size, shuffle=True,  # Batch shape: [B, 1, 28, 28].
+        num_workers=0, pin_memory=(DEVICE.type == "cuda"),  # Load in the main process.
+    )
+
+
+# %% 04. Sample the conditional path
+def sample_conditional_path(x1):  # x1 contains normalized real images.
+    x0 = torch.randn_like(x1)  # Independent Gaussian prior; same shape as x1.
+    t = torch.rand(x1.shape[0], device=x1.device)  # One random time per image; [B].
+    t_image = t[:, None, None, None]  # [B, 1, 1, 1]; share time across pixels.
+    xt = (1.0 - t_image) * x0 + t_image * x1  # Noise at t=0; real image at t=1.
+    ut = x1 - x0  # Derivative of the straight interpolation.
+    return xt, t, ut  # Network input, time, and target velocity.
+
+
+# %% 05. Define a simple convolution block
+def conv_block(in_channels, out_channels):
+    return nn.Sequential(
+        nn.Conv2d(in_channels, out_channels, 3, padding=1), nn.SiLU(),  # Keep H and W.
+        nn.Conv2d(out_channels, out_channels, 3, padding=1), nn.SiLU(),  # Refine features.
+    )
+
+
+# %% 06. Define the small U-Net
+class FlowUNet(nn.Module):
+    def __init__(self, base_channels=BASE_CHANNELS):
+        super().__init__()
+        c = base_channels  # Base width; c=32 by default.
+        self.encoder1 = conv_block(2, c)  # Image + time -> c channels.
+        self.encoder2 = conv_block(c, 2 * c)  # Encoder: c -> 2c.
+        self.middle = conv_block(2 * c, 4 * c)  # Bottleneck: 2c -> 4c.
+        self.decoder2 = conv_block(6 * c, 2 * c)  # Upsampled 4c + skip 2c -> 2c.
+        self.decoder1 = conv_block(3 * c, c)  # Upsampled 2c + skip c -> c.
+        self.output = nn.Conv2d(c, 1, 1)  # One signed velocity per pixel.
+        self.pool = nn.MaxPool2d(2)  # Halve the image height and width.
+
+    # %% 07. Compute the velocity field
+    def forward(self, xt, t):
+        t_image = t[:, None, None, None].expand_as(xt)  # Constant time channel.
+        x = torch.cat([xt, t_image], dim=1)  # [B, 2, 28, 28].
+        skip1 = self.encoder1(x)  # Keep full-resolution features.
+        skip2 = self.encoder2(self.pool(skip1))  # 28x28 -> 14x14.
+        x = self.middle(self.pool(skip2))  # 14x14 -> 7x7.
+        x = F.interpolate(x, size=skip2.shape[-2:], mode="nearest")  # 7x7 -> 14x14.
+        x = self.decoder2(torch.cat([x, skip2], dim=1))  # Restore encoder detail.
+        x = F.interpolate(x, size=skip1.shape[-2:], mode="nearest")  # 14x14 -> 28x28.
+        x = self.decoder1(torch.cat([x, skip1], dim=1))  # Restore fine detail.
+        v_pred = self.output(x)  # Predict v_theta(xt, t).
+        return v_pred  # [B, 1, 28, 28], matching xt.
+
+
+# %% 08. Train for one epoch
+def train_one_epoch(model, loader, optimizer, max_steps=0):
+    model.train()  # Select training mode.
+    total_loss, num_images = 0.0, 0  # Accumulate an image-weighted epoch loss.
+    for step, (x1, _) in enumerate(loader):  # Ignore digit labels: unconditional generation.
+        if max_steps and step >= max_steps:
+            break  # The step budget the caller left for this epoch is spent.
+        x1 = x1.to(DEVICE)  # Move real images to the model device.
+        xt, t, ut = sample_conditional_path(x1)  # Fresh noise and times for this batch.
+        v_pred = model(xt, t)  # The model receives only xt and t.
+        squared_error = (v_pred - ut).square()  # Velocity error at every pixel.
+        loss = squared_error.flatten(1).sum(1).mean()  # Sum pixels, average images.
+        optimizer.zero_grad(set_to_none=True)  # Clear old gradients.
+        loss.backward()  # Compute gradients with respect to weights.
+        optimizer.step()  # Update the U-Net weights.
+        total_loss += loss.item() * x1.shape[0]  # Weight by the actual batch size.
+        num_images += x1.shape[0]  # Include the smaller final batch.
+    return total_loss / num_images  # Average training loss per image.
+
+
+# %% 09. Generate images with Euler integration
+# Snapshots land every quarter of the default 100-step run, and the ODE states stay unclipped.
+@torch.no_grad()  # Sampling needs no gradient tracking.
+def sample_images(model, num_samples=NUM_SAMPLES, steps=SAMPLE_STEPS):
+    if num_samples < 1 or steps < 1:
+        raise ValueError("num_samples and steps must be positive")
+    model.eval()  # Select evaluation mode.
+    device = next(model.parameters()).device  # Use the model device.
+    xt = torch.randn(num_samples, 1, 28, 28, device=device)  # Initial x0.
+    dt = 1.0 / steps  # Positive time increment.
+    snapshots = [xt.cpu().clone()]  # Keep the initial noise for plotting.
+    for step in range(steps):  # Integrate forward in time.
+        t = torch.full((num_samples,), step * dt, device=device)  # Shared time grid.
+        v_pred = model(xt, t)  # Reevaluate velocity at the current state.
+        xt = xt + dt * v_pred  # Euler: state += time increment * velocity.
+        if (step + 1) % max(steps // 4, 1) == 0 or step + 1 == steps:
+            snapshots.append(xt.cpu().clone())  # Save intermediate and final states.
+    return xt, torch.stack(snapshots, dim=1)  # Raw images; [B, snapshots, 1, 28, 28].
+
+
+# %% 10. Set up the complete run
+def main():
+    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    parser.add_argument("--epochs", type=int, default=EPOCHS)
+    parser.add_argument(
+        "--steps", type=int, default=0,
+        help="optimizer steps in the whole run; 0 runs every requested epoch in full",
+    )
+    parser.add_argument("--batch-size", type=int, default=BATCH_SIZE)
+    parser.add_argument("--dataset", choices=["mnist", "toy"], default="mnist")
+    parser.add_argument("--sample-steps", type=int, default=SAMPLE_STEPS)
+    parser.add_argument("--samples", type=int, default=NUM_SAMPLES)
+    parser.add_argument("--seed", type=int, default=SEED)
+    parser.add_argument("--out", type=str, default=str(OUTPUT_DIR))
+    parser.add_argument("--quiet", action="store_true", help="print only the final report")
+    args = parser.parse_args()
+
+    torch.manual_seed(args.seed)  # Seed training randomness.
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)  # Create the output folder.
+    loader = make_loader(args.dataset, args.batch_size)  # Download, normalize, and batch.
+    model = FlowUNet().to(DEVICE)  # Create the velocity network.
+    optimizer = torch.optim.AdamW(
+        model.parameters(), lr=LEARNING_RATE, weight_decay=0.0,  # Optimize the flow loss.
+    )
+    num_parameters = sum(p.numel() for p in model.parameters())  # Count weights and biases.
+    if not args.quiet:
+        print(f"device={DEVICE} dataset={args.dataset} parameters={num_parameters:,}")
+
+    # %% 11. Train and save the learned parameters
+    losses = []  # Record the training loss after each epoch.
+    budget = args.steps  # Optimizer steps left in the run, counted across all epochs.
+    for epoch in range(1, args.epochs + 1):  # Complete the requested passes through the data.
+        loss = train_one_epoch(model, loader, optimizer, budget)  # Run the training loop.
+        losses.append(loss)  # Keep the loss history.
+        if not args.quiet:
+            print(f"epoch={epoch:02d} squared_norm_loss={loss:.3f}")
+        if args.steps:  # An epoch runs the smaller of the budget and the number of batches.
+            budget -= min(budget, len(loader))  # Optimizer steps this epoch used.
+            if budget <= 0:
+                break  # The run's budget is spent, so the remaining epochs do not run.
+    torch.save({
+        "model": model.state_dict(), "losses": losses,  # Weights and training curve.
+        "base_channels": BASE_CHANNELS,  # Width needed to reconstruct the model.
+        "seed": args.seed,  # Record the training seed.
+        "epochs_completed": len(losses),  # How far this run got.
+    }, out / "flow_unet_mnist.pt")
+
+    # %% 12. Sample and save images
+    torch.manual_seed(args.seed + 1)  # Fix the starting noise for generation.
+    samples, trajectory = sample_images(model, args.samples, args.sample_steps)
+    try:
+        from torchvision.utils import save_image  # Only needed to write the PNGs.
+
+        display_samples = ((samples.cpu() + 1.0) / 2.0).clamp(0.0, 1.0)  # Clip for display only.
+        save_image(display_samples, out / "samples.png", nrow=4)  # Final images.
+        display_path = ((trajectory[:4] + 1.0) / 2.0).clamp(0.0, 1.0)  # First four sample paths.
+        save_image(
+            display_path.flatten(0, 1), out / "trajectory.png",  # One row per sample.
+            nrow=trajectory.shape[1],  # Time increases across columns.
+        )
+        if not args.quiet:
+            print(f"wrote {out}/samples.png and {out}/trajectory.png")
+    except ImportError:  # Without torchvision the run still trains and samples.
+        if not args.quiet:
+            print("torchvision is not installed, so no PNG was written")
+
+    print(f"final loss {losses[-1]:.3f}  sample range "
+          f"[{float(samples.min()):.2f}, {float(samples.max()):.2f}]")
+
+
+# %% 13. Run the script
+if __name__ == "__main__":  # Execute when launched directly.
+    main()  # Load data, train, sample, and save.

lecture_2/interpolants.py

@@ -0,0 +1,210 @@
+"""Three interpolants between the prior and the data.
+
+Replacing the weights 1 - t and t by any pair of schedules alpha_t and beta_t with alpha_0 = 1,
+beta_0 = 0, alpha_1 = 0 and beta_1 = 1 leaves the conditional flow-matching argument intact, and
+the regression target becomes the derivative alpha_dot X_0 + beta_dot X_1. The straight line, the
+t-squared schedule and the trigonometric schedule differ in where the path goes, how fast it gets
+there, and what the intermediate densities look like, so the same network trained on each produces
+different samples.
+
+Chapter 2, Section 2.18 of the notes. Trained on ``cis6270.data.gaussian_mixture_2d``. Prints the
+chapter's table of positions and velocities at t = 0.25, the endpoint check for each schedule, the
+identity covariance of the trigonometric schedule, and the samples each trained field produces at a
+fine and at a coarse step budget.
+
+    python lecture_2/interpolants.py                 # train three fields, about fifteen seconds
+    python lecture_2/interpolants.py --steps 20      # a first look, with no time to learn
+"""
+
+# %% 01. Imports and configuration
+import argparse
+import math
+
+import torch
+
+from cis6270.data import gaussian_mixture_2d
+from cis6270.metrics import mean_log_density, mode_report
+from cis6270.nets import VelocityMLP
+from cis6270.runner import print_report, save_run, set_seed, train
+
+DIM = 2  # d, the width of the state.
+HALF_PI = math.pi / 2.0  # The constant in the trigonometric schedule and in its derivative.
+TRAIN_STEPS = 1000  # Optimizer steps for each of the three fields.
+BATCH_SIZE = 256  # Endpoint pairs per step.
+LEARNING_RATE = 3e-3  # Optimizer step size.
+WIDTH = 128  # Hidden width of the velocity network.
+SAMPLE_STEPS = 64  # Euler steps in the fine sampling budget.
+COARSE_STEPS = 8  # Euler steps in the coarse budget, where a varying velocity costs accuracy.
+NUM_SAMPLES = 1024  # Points generated per schedule and per budget.
+DATA_POINTS = 4096  # Size of the training set drawn from the four-Gaussian target.
+COVARIANCE_POINTS = 8192  # Endpoint pairs used for the empirical covariance of X_t.
+GRID = torch.tensor([0.0, 0.25, 0.50, 0.75, 1.0])  # [5] the chapter's grid of times.
+
+
+# %% 02. The three schedules
+#
+# Section 2.18. Each returns alpha_t, beta_t and their time derivatives, evaluated on a batch of
+# times. The pair of schedules fixes the path and the pair of derivatives fixes the target.
+def linear(t):
+    """alpha_t = 1 - t and beta_t = t, the straight segment traversed at constant speed."""
+    return 1.0 - t, t, -torch.ones_like(t), torch.ones_like(t)  # Four tensors of shape [B].
+
+
+def quadratic(t):
+    """alpha_t = 1 - t^2 and beta_t = t^2, the same segment entered slowly and left quickly."""
+    return 1.0 - t ** 2, t ** 2, -2.0 * t, 2.0 * t  # U_t = 2t (X_1 - X_0), zero speed at t = 0.
+
+
+def trigonometric(t):
+    """alpha_t = cos(pi t / 2) and beta_t = sin(pi t / 2), a quarter circle in the coefficients."""
+    angle = HALF_PI * t  # [B] the angle the coefficient pair has turned through.
+    return (torch.cos(angle), torch.sin(angle),
+            -HALF_PI * torch.sin(angle), HALF_PI * torch.cos(angle))
+
+
+SCHEDULES = {"linear": linear, "quadratic": quadratic, "trigonometric": trigonometric}
+
+
+def interpolate(schedule, x_0, x_1, t):
+    """X_t = alpha_t X_0 + beta_t X_1 and U_t = alpha_dot X_0 + beta_dot X_1. Two [B, d] tensors."""
+    alpha, beta, alpha_dot, beta_dot = schedule(t)  # [B] each, the two weights and their rates.
+    x_t = alpha[:, None] * x_0 + beta[:, None] * x_1  # [B, d] the state on the conditional path.
+    u_t = alpha_dot[:, None] * x_0 + beta_dot[:, None] * x_1  # [B, d] the regression target.
+    return x_t, u_t
+
+
+# %% 03. The endpoints and the velocity at one time
+#
+# Section 2.18. The two ends are correct when alpha_0 = 1, beta_0 = 0, alpha_1 = 0 and beta_1 = 1,
+# and Table 2.1 reads the three paths at X_0 = (2, 0), X_1 = (0, 2) and t = 0.25.
+def endpoint_and_table(name):
+    """The endpoint error of one schedule and its position and velocity at the chapter's input."""
+    x_0 = torch.tensor([[2.0, 0.0]])  # [1, d] the prior endpoint of the table.
+    x_1 = torch.tensor([[0.0, 2.0]])  # [1, d] the data endpoint of the table.
+    schedule = SCHEDULES[name]
+    ends, _ = interpolate(schedule, x_0.expand(2, DIM), x_1.expand(2, DIM),
+                          torch.tensor([0.0, 1.0]))  # [2, d] the states at t = 0 and t = 1.
+    start_error = (ends[0] - x_0[0]).abs().max()  # X_0 is returned at t = 0.
+    finish_error = (ends[1] - x_1[0]).abs().max()  # X_1 is returned at t = 1.
+    x_t, u_t = interpolate(schedule, x_0, x_1, torch.tensor([0.25]))  # One quarter of the interval.
+    return {
+        f"{name}_endpoint_error": float(torch.maximum(start_error, finish_error)),
+        f"{name}_position_at_quarter": [round(float(v), 4) + 0.0 for v in x_t[0]],
+        f"{name}_velocity_at_quarter": [round(float(v), 4) + 0.0 for v in u_t[0]],
+    }
+
+
+# %% 04. Identity covariance
+#
+# Section 2.18. For independent, centered endpoints with identity covariance the interpolation has
+# covariance (alpha_t^2 + beta_t^2) I. That factor is cos^2 + sin^2 = 1 at every time for the
+# trigonometric schedule, and below one between the ends for the other two.
+def covariance_check(name, num_points, seed):
+    """The coefficient squares on the grid, and the largest deviation of Cov(X_t) from the identity."""
+    set_seed(seed)  # Same draw for every schedule, so the comparison is between schedules alone.
+    x_0 = torch.randn(num_points, DIM)  # [N, d] a centered, white prior endpoint.
+    x_1 = torch.randn(num_points, DIM)  # [N, d] a centered, white data endpoint, independent of it.
+    identity = torch.eye(DIM)  # [d, d] the covariance the chapter's hypothesis gives each endpoint.
+    squares, deviations = [], []
+    for t in GRID:  # One time at a time, since the covariance is measured over the batch.
+        times = t.expand(num_points)  # [N] the same time for every pair.
+        alpha, beta, _, _ = SCHEDULES[name](times)
+        squares.append(round(float(alpha[0] ** 2 + beta[0] ** 2), 4))  # The exact variance factor.
+        x_t, _ = interpolate(SCHEDULES[name], x_0, x_1, times)  # [N, d] the interpolated states.
+        covariance = (x_t.T @ x_t) / num_points  # [d, d] the empirical second moment, centered at 0.
+        deviations.append(float((covariance - identity).abs().max()))
+    return {
+        f"{name}_coefficient_squares": squares,
+        f"{name}_covariance_deviation": max(deviations),
+    }
+
+
+# %% 05. Training one field per interpolant
+#
+# Section 2.18 and Section 2.19. The loop is the one of the straight-line case with the target
+# replaced, so any difference in the samples comes from the schedule and not from the setup.
+def train_one(name, data, args):
+    """Train v_theta against the conditional velocity of one schedule and return the loss history."""
+    set_seed(args.seed)  # The same initialization and the same batches for every schedule.
+    model = VelocityMLP(dim=DIM, width=args.width)  # v_theta(x, t).
+    schedule = SCHEDULES[name]
+
+    def loss_fn(step):
+        index = torch.randint(0, len(data), (args.batch_size,))  # [B] which data point to pair.
+        x_1 = data[index]  # [B, d] the data endpoint.
+        x_0 = torch.randn_like(x_1)  # [B, d] the prior endpoint, independent of it.
+        t = torch.rand(args.batch_size)  # [B] t ~ Uniform(0, 1).
+        x_t, u_t = interpolate(schedule, x_0, x_1, t)  # [B, d] the state and the target.
+        return (model(x_t, t) - u_t).square().sum(1).mean()  # ||v_theta(X_t, t) - U_t||^2.
+
+    losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=args.quiet)
+    return model, losses
+
+
+# %% 06. Euler sampling
+#
+# Section 2.19. Every schedule is sampled by the same integrator, so the step budget is the only
+# thing a reader has to hold fixed when comparing them.
+@torch.no_grad()
+def sample(model, num_samples, steps):
+    """Integrate dX_t/dt = v_theta(X_t, t) from t = 0 to t = 1 on a uniform grid. Returns [N, d]."""
+    x_t = torch.randn(num_samples, DIM)  # [N, d] the draw X_0 ~ p_0.
+    dt = 1.0 / steps  # Delta t, the width of one interval.
+    for step in range(steps):  # One network evaluation per step.
+        x_t = x_t + dt * model(x_t, torch.full((num_samples,), step * dt))  # One Euler step.
+    return x_t
+
+
+def measure(model, args):
+    """Mode coverage and mean exact log density at the fine and the coarse step budget."""
+    set_seed(args.seed + 1)  # The same starting noise for every schedule.
+    fine = sample(model, args.samples, args.sample_steps)  # [N, d] points from the fine grid.
+    coarse = sample(model, args.samples, args.coarse_steps)  # [N, d] points from the coarse grid.
+    report = mode_report(fine)
+    return {
+        "modes_found": report["modes_found"],
+        "mode_entropy": report["mode_entropy"],
+        "nearest_distance": report["nearest_distance"],
+        "mean_log_density": mean_log_density(fine),
+        "coarse_log_density": mean_log_density(coarse),
+    }
+
+
+# %% 07. The report
+def main():
+    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    parser.add_argument("--steps", type=int, default=TRAIN_STEPS, help="optimizer steps per field")
+    parser.add_argument("--batch-size", type=int, default=BATCH_SIZE, help="pairs per step")
+    parser.add_argument("--lr", type=float, default=LEARNING_RATE, help="learning rate")
+    parser.add_argument("--width", type=int, default=WIDTH, help="hidden width of the network")
+    parser.add_argument("--samples", type=int, default=NUM_SAMPLES, help="points to generate")
+    parser.add_argument("--sample-steps", type=int, default=SAMPLE_STEPS, help="fine Euler steps")
+    parser.add_argument("--coarse-steps", type=int, default=COARSE_STEPS, help="coarse Euler steps")
+    parser.add_argument("--seed", type=int, default=0, help="random seed")
+    parser.add_argument("--out", type=str, default="", help="directory for the run's output")
+    parser.add_argument("--quiet", action="store_true", help="print only the final report")
+    args = parser.parse_args()
+
+    data = gaussian_mixture_2d(DATA_POINTS, seed=args.seed)  # [N, d] the target p_1 = p_data.
+    report = {}
+    for name in SCHEDULES:  # Linear, quadratic, trigonometric, in the order the chapter gives them.
+        report.update(endpoint_and_table(name))  # Section 2.18, before any training.
+        report.update(covariance_check(name, COVARIANCE_POINTS, args.seed))
+        if not args.quiet:
+            print(f"training the {name} schedule")
+        model, losses = train_one(name, data, args)
+        report[f"{name}_final_loss"] = sum(losses[-100:]) / len(losses[-100:])
+        # The three losses are not comparable with each other, since each regresses on its own
+        # target, and the three measurements below are.
+        for key, value in measure(model, args).items():
+            report[f"{name}_{key}"] = value
+        save_run(f"{args.out}/{name}" if args.out else "", model=model, losses=losses)
+
+    report["data_mean_log_density"] = mean_log_density(data)  # The same measurement on real points.
+    print_report("three interpolants on four Gaussians", report)
+    save_run(args.out, config=vars(args), report=report)
+
+
+# %% 08. Run the script
+if __name__ == "__main__":
+    main()

lecture_2/velocity_fields.py

@@ -0,0 +1,263 @@
+"""Velocity fields, flow maps and the continuity equation.
+
+A velocity field assigns one vector to every state and every time, and following those vectors from
+a starting point traces a trajectory. The trajectories from every starting point together form the
+flow map, and the continuity equation relates the motion of the density to the same field. Nothing
+here is learned, so every number printed has a closed form beside it.
+
+Chapter 2, Sections 2.2 to 2.8 of the notes. Runs on the chapter's analytic fields in one and two
+dimensions. Prints the trajectories, the Euler error at one, two and four steps and the gap closing
+like 1/N, the pushforward of a Gaussian prior, the flux balance on [-1, 1], both sides of the
+continuity equation at x = 2, and the log density along one trajectory.
+
+    python lecture_2/velocity_fields.py                 # the full report
+    python lecture_2/velocity_fields.py --steps 20      # a shorter Euler refinement ladder
+"""
+
+# %% 01. Imports and configuration
+import argparse
+import math
+
+import torch
+
+from cis6270.runner import print_report, save_run, set_seed
+
+torch.set_default_dtype(torch.float64)  # Six printed digits of the Euler ladder need it.
+
+LOG2 = math.log(2.0)  # The rate a of the expansion field v_t(x) = a x, so that e^{at} = 2^t.
+GRID = torch.tensor([0.0, 0.25, 0.50, 0.75, 1.0])  # [5] the chapter's grid of times.
+START_POINTS = torch.tensor([-1.0, 0.0, 1.0])  # [3] the initial states followed through each field.
+EULER_STEPS = 64  # Largest step count in the refinement ladder, halved down to one step.
+NUM_SAMPLES = 4096  # Draws from the prior p_0 used for the pushforward moments.
+
+
+# %% 02. Three velocity fields and the trajectories they produce
+#
+# Section 2.2 and Section 2.3. Each field is a rule at a single instant, written as a function of
+# the state x and the time t, and each has an exact solution to compare a solver against.
+def constant_field(x, t):
+    """v_t(x) = 2, the field that translates every state at the same speed."""
+    return torch.full_like(x, 2.0)  # Same shape as x; the value depends on neither x nor t.
+
+
+def linear_field(x, t, a=LOG2):
+    """v_t(x) = a x, which expands away from the origin for a > 0 and contracts toward it for a < 0."""
+    return a * x  # Speed grows with distance from the origin and vanishes at x = 0.
+
+
+def rotation_field(x, t):
+    """v(x_1, x_2) = (-x_2, x_1), everywhere perpendicular to the position vector."""
+    return torch.stack([-x[..., 1], x[..., 0]], dim=-1)  # [..., 2] same shape as x.
+
+
+def exact_linear(x_0, t, a=LOG2):
+    """x_t = x_0 e^{at}, the solution of dx_t/dt = a x_t obtained by separating variables."""
+    return x_0 * torch.exp(a * t)  # Broadcasts a grid of times against a set of starting points.
+
+
+def trajectories():
+    """The chapter's worked trajectories for translation, expansion, contraction and rotation."""
+    speed = constant_field(torch.tensor(1.0), torch.tensor(0.0))  # The value 2, at any state.
+    translated = 1.0 + speed * GRID  # [5] the integral form x_t = x_0 + 2t from x_0 = 1.
+    expanded = exact_linear(START_POINTS, torch.tensor(1.0))  # [3] states at t = 1, each doubled.
+    contracted = exact_linear(START_POINTS, torch.tensor(1.0), a=-LOG2)  # [3] states each halved.
+    turns = torch.tensor([math.pi / 2, math.pi])  # [2] the quarter turn and the half turn.
+    circle = torch.stack([torch.cos(turns), torch.sin(turns)], dim=-1)  # [2, 2] x_t = (cos t, sin t).
+    # The squared norm is conserved because the field is orthogonal to the position: 2 x . v_t(x) = 0.
+    off_axis = torch.tensor([0.3, 0.7])  # [2] a state away from the axes, to test the conservation.
+    norm_rate = 2.0 * (off_axis * rotation_field(off_axis, torch.tensor(0.0))).sum()  # d||x||^2/dt.
+    return {
+        "constant_field_trajectory": [round(float(v), 4) for v in translated],
+        "expansion_x_half_from_two": float(exact_linear(torch.tensor(2.0), torch.tensor(0.5))),
+        "expansion_endpoints": [round(float(v), 4) + 0.0 for v in expanded],
+        "expansion_outer_spacing": float(expanded[2] - expanded[0]),
+        "contraction_endpoints": [round(float(v), 4) + 0.0 for v in contracted],
+        "contraction_outer_spacing": float(contracted[2] - contracted[0]),
+        "rotation_quarter_turn": [round(float(v), 4) + 0.0 for v in circle[0]],
+        "rotation_half_turn": [round(float(v), 4) + 0.0 for v in circle[1]],
+        "rotation_radius": float(circle.norm(dim=-1).mean()),
+        "rotation_squared_norm_rate": float(norm_rate),
+    }
+
+
+# %% 03. The flow map
+#
+# Section 2.4. Solving the ODE from every starting point at once defines phi_t, which is the identity
+# at t = 0 and, for the expansion field, multiplies the whole line by 2^t.
+def flow_map(x_0, t, a=LOG2):
+    """phi_t(x_0) = x_0 e^{at}, the accumulated transformation after following v_t for a time t."""
+    return exact_linear(x_0, t, a)  # [N] the state each starting point has reached by time t.
+
+
+def flow_map_report():
+    """The identity at t = 0, the Jacobian determinant at t = 1, and noncrossing trajectories."""
+    probe = torch.linspace(-3.0, 3.0, 101)  # [101] states spread over the line.
+    identity_error = (flow_map(probe, torch.tensor(0.0)) - probe).abs().max()  # phi_0 = id.
+    x_0 = torch.tensor(1.3, requires_grad=True)  # A state at which to differentiate phi_1.
+    jacobian = torch.autograd.grad(flow_map(x_0, torch.tensor(1.0)), x_0)[0]  # d phi_1 / d x_0 = 2.
+    # Proposition 2.1: distinct trajectories of a Lipschitz field never share a state at a shared time.
+    gaps = (flow_map(START_POINTS[2], GRID) - flow_map(START_POINTS[0], GRID)).abs()  # [5].
+    return {
+        "flow_map_identity_error": float(identity_error),
+        "flow_map_jacobian_det": float(jacobian),
+        "trajectory_gap_minimum": float(gaps.min()),
+    }
+
+
+# %% 04. Euler's method
+#
+# Section 2.5. Within one step the velocity stays fixed at its value at the start of the step, so
+# on dx_t/dt = x_t, where the velocity grows with the state, every step undershoots.
+def euler(field, x_0, steps):
+    """hat x_{k+1} = hat x_k + Delta t v_{t_k}(hat x_k) on a uniform grid of ``steps`` intervals."""
+    x_hat = x_0.clone()  # The running approximation, started from the exact initial state.
+    dt = 1.0 / steps  # Delta t, the width of one interval of the grid.
+    for k in range(steps):  # One velocity evaluation per step.
+        x_hat = x_hat + dt * field(x_hat, torch.tensor(k * dt))  # The Euler update at grid time t_k.
+    return x_hat
+
+
+def euler_ladder(largest):
+    """Euler on dx_t/dt = x_t from x_0 = 1, at every power of two up to ``largest`` steps."""
+    exact = math.e  # The exact solution x_1 = e^1, the chapter's 2.7183.
+    counts, landings, errors = [], [], []
+    steps = 1
+    largest = max(largest, 4)  # The chapter's three rungs, at one, two and four steps, always run.
+    while steps <= largest:  # 1, 2, 4, 8, ... up to the requested largest step count.
+        landing = float(euler(lambda x, t: x, torch.tensor(1.0), steps))  # The field is v_t(x) = x.
+        counts.append(steps)
+        landings.append(round(landing, 6))
+        errors.append(round(exact - landing, 6))
+        steps *= 2
+    return {
+        "euler_step_counts": counts,
+        "euler_landings": landings,
+        "euler_errors": errors,
+        # A first-order method has error proportional to 1/N, so the error times N settles at e/2.
+        "euler_error_times_steps": [round(e * n, 4) for e, n in zip(errors, counts)],
+        "euler_error_halving_ratios": [round(errors[i] / errors[i + 1], 4)
+                                       for i in range(len(errors) - 1)],
+        "euler_relative_error_four_steps": float((exact - landings[2]) / exact),
+    }
+
+
+# %% 05. A random starting point
+#
+# Section 2.6. The initial state becomes a random variable, each draw follows the same field, and
+# the law at time t is the pushforward p_t = (phi_t)_# p_0.
+def pushforward_moments(num_samples, seed):
+    """Moments of phi_1(X_0) for X_0 ~ N(0, 1) under the expansion field, against N(0, 4)."""
+    set_seed(seed)  # Same seed, same draw, same report.
+    x_0 = torch.randn(num_samples)  # [N] draws from the prior p_0 = N(0, 1).
+    x_1 = flow_map(x_0, torch.tensor(1.0))  # [N] states at t = 1, every one of them doubled.
+    return {
+        "pushforward_mean": float(x_1.mean()),  # Exact value 0, since the flow fixes the origin.
+        "pushforward_variance": float(x_1.var(unbiased=False)),  # Scaling by 2 gives variance 4.
+        "pushforward_exact_variance": 4.0,  # Scaling a centered law by c multiplies variance by c^2.
+    }
+
+
+# %% 06. Probability flux and the continuity equation
+#
+# Section 2.7. The flux j_t(x) = p_t(x) v_t(x) measures probability transported per unit time, and
+# the balance over an interval becomes the continuity equation once the interval shrinks to a point.
+def spreading_density(x, t):
+    """p_t(x) for X_t ~ N(0, 4^t), the law the expansion field produces from a standard normal."""
+    std = torch.pow(torch.tensor(2.0), t)  # 2^t, the standard deviation at time t.
+    return torch.exp(-x ** 2 / (2.0 * std ** 2)) / (std * math.sqrt(2.0 * math.pi))
+
+
+def flux_balance():
+    """The two endpoint fluxes on [-1, 1] at t = 0, and the rate of change of the enclosed mass."""
+    t_0 = torch.tensor(0.0)
+    endpoints = torch.tensor([-1.0, 1.0])  # [2] the interval's endpoints a and b.
+    density = spreading_density(endpoints, t_0)  # [2] p_0(-1) and p_0(1), both 0.2420.
+    flux = density * linear_field(endpoints, t_0)  # [2] j_0(x) = p_0(x) v_0(x).
+    # The same rate read off the distribution: M_t = erf(2^{-t} / sqrt 2), differentiated in t.
+    t = torch.tensor(0.0, requires_grad=True)
+    mass = torch.erf(torch.pow(torch.tensor(2.0), -t) / math.sqrt(2.0))  # P(|X_t| <= 1).
+    mass_rate = torch.autograd.grad(mass, t)[0]  # d M_t / dt at t = 0.
+    return {
+        "density_at_one": float(density[1]),
+        "flux_at_minus_one": float(flux[0]),
+        "flux_at_one": float(flux[1]),
+        "flux_balance": float(flux[0] - flux[1]),  # j_0(a) - j_0(b), the net rate into [-1, 1].
+        "enclosed_mass": float(mass.detach()),
+        "enclosed_mass_rate": float(mass_rate),
+    }
+
+
+def continuity_check():
+    """Both sides of the continuity equation at x = 2 and t = 0, in the expanded product-rule form."""
+    x = torch.tensor(2.0, requires_grad=True)  # The position held fixed while the density changes.
+    t = torch.tensor(0.0, requires_grad=True)  # The time at which both sides are evaluated.
+    p = spreading_density(x, t)  # p_0(2), the chapter's 0.05399.
+    dp_dt = torch.autograd.grad(p, t, create_graph=True)[0]  # Local time derivative of the density.
+    dp_dx = torch.autograd.grad(p, x, create_graph=True)[0]  # Spatial gradient of the density.
+    divergence = torch.autograd.grad(linear_field(x, t), x)[0]  # div v = log 2 in one dimension.
+    transport = linear_field(x, t) * dp_dx  # v_t(x) . grad p_t(x), the transport term.
+    volume = p * divergence  # p_t(x) div v_t(x), the local volume-change term.
+    return {
+        "density_at_two": float(p.detach()),
+        "density_gradient_at_two": float(dp_dx.detach()),
+        "velocity_divergence": float(divergence),
+        "continuity_time_derivative": float(dp_dt.detach()),
+        "continuity_minus_flux_divergence": float(-(transport + volume).detach()),
+    }
+
+
+# %% 07. The density along a moving trajectory
+#
+# Section 2.8. An observer riding with the sample sees no transport term, so the log density changes
+# at the rate -div v_t(x_t), and over a unit of time that agrees with the Jacobian of the flow map.
+def log_density_along_trajectory(x_start=0.7):
+    """The instantaneous change-of-variables equation on the expansion field, both ways."""
+    x_0 = torch.tensor(x_start)  # One starting point, carried from t = 0 to t = 1.
+    x_1 = flow_map(x_0, torch.tensor(1.0))  # The state the trajectory reaches at t = 1.
+    log_p0 = torch.log(spreading_density(x_0, torch.tensor(0.0)))  # log p_0(x_0).
+    log_p1 = torch.log(spreading_density(x_1, torch.tensor(1.0)))  # log p_1(x_1).
+    origin = torch.tensor(0.0)  # The fixed point of the flow, where phi_1(0) = 0.
+    prior_at_origin = float(spreading_density(origin, torch.tensor(0.0)))  # p_0(0) = 0.3989.
+    final_at_origin = float(spreading_density(origin, torch.tensor(1.0)))  # p_1(0) = 0.1995.
+    return {
+        "log_density_rate": -LOG2,  # d log p_t(x_t) / dt = -div v_t(x_t), a constant here.
+        "log_density_change": float(log_p1 - log_p0),  # Measured from the two densities.
+        "log_density_accumulated": -LOG2 * 1.0,  # The integral of -div v_s over s from 0 to 1.
+        "prior_density_at_origin": prior_at_origin,
+        "final_density_at_origin": final_at_origin,
+        "density_ratio_at_origin": prior_at_origin / final_at_origin,  # |det J_{phi_1}| = 2.
+    }
+
+
+# %% 08. The report
+def main():
+    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    parser.add_argument("--steps", type=int, default=EULER_STEPS, help="largest Euler step count")
+    parser.add_argument("--samples", type=int, default=NUM_SAMPLES, help="draws from the prior")
+    parser.add_argument("--seed", type=int, default=0, help="random seed")
+    parser.add_argument("--out", type=str, default="", help="directory for the run's report")
+    parser.add_argument("--quiet", action="store_true", help="print only the final report")
+    args = parser.parse_args()
+
+    report = {}
+    report.update(trajectories())  # Sections 2.2 and 2.3.
+    report.update(flow_map_report())  # Section 2.4.
+    report.update(euler_ladder(args.steps))  # Section 2.5.
+    report.update(pushforward_moments(args.samples, args.seed))  # Section 2.6.
+    report.update(flux_balance())  # Section 2.7.
+    report.update(continuity_check())  # Section 2.7.
+    report.update(log_density_along_trajectory())  # Section 2.8.
+
+    if not args.quiet:  # The Euler ladder, rung by rung, as the chapter develops it.
+        for n, landing, error in zip(report["euler_step_counts"], report["euler_landings"],
+                                     report["euler_errors"]):
+            print(f"  euler steps {n:3d}  landing {landing:.6f}  error {error:.6f}"
+                  f"  error x steps {error * n:.4f}")
+
+    print_report("velocity fields, flow maps and the continuity equation", report)
+    save_run(args.out, report=report)  # Writes report.json when --out names a directory.
+
+
+# %% 09. Run the script
+if __name__ == "__main__":
+    main()

lecture_3/README.md

@@ -0,0 +1,139 @@
+# Lecture 3: Score Matching and Diffusion Models
+
+Chapter 3 of the notes builds a generator out of corruption. A clean sample is pushed toward a
+standard Gaussian by a stochastic differential equation, Ito's lemma and the Fokker-Planck
+equation turn that equation into one for the density, and the variance-preserving choice of drift
+and noise scale makes the density Gaussian in closed form at every level. Reversing the motion
+needs the gradient of the log density, the score, which is learned by regression on the noise that
+was drawn. Discretizing the reverse equation on a grid gives DDPM; rewriting the Laplacian as a
+divergence gives the probability-flow ODE; and one extra gradient inside either sampler gives
+classifier, classifier-free and reward guidance. The last file runs the whole chain in the latent
+space of a protein language model.
+
+The clock runs from the data in every file here: `t = 0` and `k = 0` hold a clean sample and
+corruption grows with `t` and with `k`, up to a standard Gaussian at the terminal level. This is
+the reverse of Chapter 2's convention and it is kept throughout.
+
+| File | Command | What it does |
+| --- | --- | --- |
+| `sde.py` | `python lecture_3/sde.py` | The forward process before any schedule: Brownian increments, the Euler-Maruyama step, the spreading cloud, a restoring drift, Ito's lemma on `phi(x) = x^2`, and the Fokker-Planck equation checked on `N(0, 1+t)`. |
+| `forward_diffusion.py` | `python lecture_3/forward_diffusion.py` | The variance-preserving schedule and the Gaussian corruption it drives: the four noise levels located by inverting `a_t`, the thousand-step grid against the continuous curve, one pixel at step 207, the worked training batch at `abar = 0.8`, and the corruption law measured against `N(a_k x_0, sigma_k^2 I)`. |
+| `score_matching.py` | `python lecture_3/score_matching.py` | Fits explicit, implicit and denoising score matching to the same one-dimensional Gaussian mixture, checks the integration by parts, locates the minimizer of each objective, and estimates a 64-dimensional trace from random directions. |
+| `ddpm.py` | `python lecture_3/ddpm.py` | Trains the chapter's U-Net on 28-by-28 images with `L_simple` and samples with the K-step reverse chain. `--dataset mnist` downloads the digits. |
+| `samplers.py` | `python lecture_3/samplers.py` | Ancestral DDPM, the probability-flow ODE and annealed Langevin on a four-Gaussian target whose corrupted score is exact, so nothing is trained and the three are compared on equal evaluation budgets. |
+| `guidance.py` | `python lecture_3/guidance.py` | Classifier, classifier-free and reward guidance on the same target with one label per mode, with weights on the simplex and a schedule on the strength. |
+| `latent_diffusion.py` | `python lecture_3/latent_diffusion.py` | AMP-Diffusion: masked corruption of residue embeddings, a transformer denoiser, reverse latent sampling, decoding through a language-model head, and a polar-composition constraint at decoding. `--dataset esm2` uses the real encoder. |
+
+Every file takes the shared flags of `cis6270.runner.common_parser`. In `sde.py`,
+`forward_diffusion.py`, `samplers.py` and `guidance.py` nothing is trained, so `--steps` sets the
+number of simulation steps, measured levels or reverse steps; in the other three it sets the
+number of optimizer steps.
+
+## What a default run prints
+
+`sde.py` is deterministic arithmetic plus a Monte Carlo check. The Euler-Maruyama step from
+`x = 2` lands at `2.11` with a drift displacement of `-0.04` against a noise displacement of
+`0.15`, and the typical noise is `2.5` times the drift at this step size. The driftless cloud ends
+with variance `0.2521` against the exact `g^2 T = 0.25`. The Ornstein-Uhlenbeck law at `t = log 2`
+has mean `1` and variance `0.75`, which `16384` simulated paths reproduce as `0.9860` and `0.7458`.
+Ito's lemma returns `E[W_t^2] = 0.2504` at `t = 0.25` against the exact `0.25`, and the same sum
+without the curvature term returns `0.0004`. The Fokker-Planck residual on `N(0, 1+t)` is zero to
+float precision, and the peak of the Gaussian falls from `0.3989` at the rate `-0.1995`, so the
+one-step estimate `0.3790` sits `0.0014` below the exact `0.3804`.
+
+`forward_diffusion.py` locates the four levels of the chapter's table at `t = 0.0589`, `0.2068`,
+`0.4544` and `0.7994`, with `sigma_t` equal to `0.199`, `0.600`, `0.9367` and `0.9992` and `beta(t)`
+equal to `1.2721`, `4.2157`, `9.1420` and `16.0073`; the two curves cross at `t = 0.25896`, where
+both are `0.7071`. On the thousand-step grid `abar_207 = 0.6402`, so `a_207 = 0.8001` against the
+continuous `0.7997` and `sigma_207 = 0.5999` against `0.6004`, and `a_1000 = 0.0064` against
+`0.0066`. One pixel at `x_0 = 0.5` and `eps = -1` lands at `-0.1998`, and by step `500` the signal
+is down to `0.2803`. The worked batch at `abar = 0.8` reproduces `(2.0125, -0.9839, 0.8944)`, a
+loss of `0.02`, score targets `(-1.118, 0.447, -2.236)` and clean estimates `(2.05, -1.05, 0.60)`.
+Corrupting one image `1024` times at each of `16` levels leaves the per-pixel sample mean within
+`0.0257` of `a_k x_0` and the pooled sample variance within `0.0026` of `sigma_k^2`, and both
+gaps fall as one over the square root of the number of draws.
+
+`score_matching.py` fits three networks from the same initialization on the same batches. The
+score errors against the exact field are `0.0337`, `0.1976` and `0.2073` for the explicit, implicit
+and denoising objectives: the explicit fit regresses on the exact target and comes out an order of
+magnitude closer, while the other two land near each other, each paying for the variance its own
+estimator adds. The ordering of those two moves with the seed. Along the family `a s_exact` the
+three minimizers come out `1.0000`, `0.9908` and `1.0053`, which is the claim that the three
+objectives agree on a minimizer, measured and not merely argued. The integration-by-parts identity
+closes to about one percent of its own size on a batch of 4096, the exact 64-dimensional trace
+takes 64 backward passes and returns `-64.00` against the Hutchinson estimate `-64.67` from 32
+passes, and the posterior at `x_t = 7.4` has weights `0.2156` and `0.7844`, clean mean `9.2156`,
+marginal score `-0.0765` by both routes and an irreducible residual of `0.835`.
+
+`ddpm.py` on the offline images drops `L_simple` from `0.169` over the first hundred steps to
+`0.033` over the last hundred, and the generated images land in `[-1.60, 1.50]` with a standard
+deviation of `0.59`, against training data in `[-1, 1]`. The reverse posterior variance at
+`k = 500` is `0.01003` and the ELBO weight it carries is `0.005508`, against `0.625` on the
+two-step hand schedule. That two-step transition gives `8.6461` from both forms of the mean, a
+clean estimate of `10` with interpolation weights `0.4969` and `0.4969`, variance `1/9`, and
+`x_1 = 8.746`.
+
+`samplers.py` gives all three samplers 200 score evaluations. Each finds all four modes with mode
+entropy `1.3854`, `1.3848` and `1.3852` against `log 4 = 1.386`, and the mean distance to the
+nearest mode is `0.4162` for ancestral DDPM, `0.4446` for the probability-flow ODE and `0.4662`
+for annealed Langevin, against `0.4435` for the target itself. The mean log densities are
+`-2.01`, `-2.15` and `-2.24` against the target's own `-2.15`. The worked backward step prints a
+score of `-1.6`, a deterministic velocity of `0.4` reaching `1.16`, and a reverse drift of `2.0`
+reaching `1.0` plus a Gaussian of standard deviation `0.4472`.
+
+`guidance.py` sweeps the weight on the mode at `(2, 2)`. The target fraction runs `0.256`, `0.334`,
+`0.506`, `0.899`, `1.000` at `w = 0, 0.1, 0.25, 0.5, 1`, and the scheduled version of `w = 1`
+reaches `0.991`, slightly weaker because `zeta(t) = 4t(1-t)` switches the condition off at both
+ends of the interval. Reward steering with `lambda = (1, 0)` moves the two right-hand modes from
+`0.491` of the mass unguided to `0.719` at `eta = 0.5` and `0.994` at `eta = 2`, and
+`lambda = (0.7, 0.3)` concentrates `0.740` on the top-right mode. The classifier route and the
+classifier-free route agree on the guided noise to float precision, which is the content of the
+rearrangement of Bayes' rule, and the guided backward step gives the velocity `-0.8` and the state
+`1.28` from the velocity form and the same `-0.8` from the guided noise `0.2`.
+
+`latent_diffusion.py` trains the transformer denoiser on 128 toy embeddings of 16 residues and 32
+coordinates. The masked loss falls from `1.038` to `0.817`. The two numbers that bracket it are
+`1.0`, which is what a denoiser that always returns zero scores on a standardized latent, and
+`0.725`, the posterior variance of `Z_0` given `Z_k` under a Gaussian latent, averaged over a
+uniform step index. Raising `--steps` to 2000 reaches `0.746`, close to that floor, for about
+three times the running time. Decoding the generated latents gives a polar fraction of `0.3047`, and the constraint
+raises it to exactly `0.50` by rewriting the cheapest positions at each sequence until eight are
+polar. The chapter's latent coordinate `z_0 = 1.2` corrupts to `0.7204` at `k = 500`, a prediction
+of `1.05` implies `eps = 0.4438`, and the reverse mean is `0.7193` with standard deviation
+`0.1002`.
+
+## What is simplified
+
+- `ddpm.py` uses a U-Net of base width 8, where the notes uses 64, and trains for a few hundred
+  steps at a learning rate of `2e-3` in place of 20 MNIST epochs at `2e-4`. Both are flags.
+  The class, the schedule, the forward corruption and the reverse step are the chapter's code with
+  its names kept.
+- `forward_diffusion.py` and `ddpm.py` are the two halves of the chapter's implementation section,
+  split so that neither file runs long. No lecture file imports another, so each rebuilds the
+  schedule tensors; the two copies are identical, and `tests/test_lecture_3.py` asserts that they
+  stay so.
+- `samplers.py` and `guidance.py` evaluate the exact score of the corrupted four-Gaussian mixture
+  in place of a trained network, which is what lets the three samplers and the guidance weights be
+  compared with no training error in the way. A counter on the score stands in for the count of
+  network evaluations. Both read the schedule off the continuous curve, `abar_k = exp(-B(t_k))`,
+  and not from a linear sequence of betas, so that every `beta_k` stays below one on the coarse
+  grid these files run; at `K = 1000` the two agree to about `1e-3`.
+- The conditional score in `guidance.py` is exact, so at `w = 1` the condition pins the mode and
+  the mode fractions saturate. The graded part of the sweep therefore sits below one, where a
+  trained conditional model would place it.
+- `score_matching.py` fits a time-independent field at the single level `a_t = 0.8`,
+  `sigma_t = 0.6`, the chapter's running level, in place of one field across all levels.
+  The Hutchinson estimator is written as a vector-Jacobian product; a score field has a symmetric
+  Jacobian, so this equals the Jacobian-vector product of the notes at the same cost.
+- `latent_diffusion.py` runs offline on `cis6270.data.toy_embeddings`, which have 32 coordinates
+  in place of the 1280 of ESM-2. Without the encoder there is no language-model head, so a fixed
+  random projection stands in for it; the decoding path and the composition constraint run
+  unchanged, and the residues they read out carry no biology. `--dataset esm2` uses the real
+  encoder, the real head and the real tokenizer, at `MAX_LEN = 32` in place of the chapter's 64.
+  That path downloads a 2.5 GB checkpoint and widens the denoiser to 1280, which is far too slow
+  for the default `--steps` on a laptop CPU; try 20 first. Generation there fills all `MAX_LEN`
+  positions, including the ones that held padding during training, exactly as the chapter's
+  `sample_latents` does.
+- The reward model of Section 3.46 is a small network on the latent. In `guidance.py` the state is
+  a point in the plane, so the two properties are read off its coordinates directly and keep the
+  same signature, which leaves nothing to train.

lecture_3/ddpm.py

@@ -0,0 +1,315 @@
+"""DDPM, the reverse-time equation written on a finite grid.
+
+A U-Net reads a corrupted image and the step it sits at, and regresses the Gaussian draw that
+corrupted it. The reverse transition rebuilds the previous state from that one prediction, with a
+mean that comes from completing the square and the posterior variance as its spread, and the
+sampler starts at a standard Gaussian and applies that transition K times. The three pieces match
+the derivation one to one, with the chapter's names kept. The schedule and the forward corruption
+are the subject of `forward_diffusion.py`, rebuilt here so that this file stands alone.
+
+Chapter 3 runs its diffusion clock from the data, so step k = 0 holds a clean image and corruption
+grows with k, up to pure noise at k = K.
+
+Chapter 3, Sections 3.19, 3.20 and 3.24 through 3.28 of the notes. Trained on 28-by-28 images, the
+offline toy set by default and MNIST with --dataset mnist. Prints L_simple over the run, the
+reverse posterior variance and the ELBO weight it carries, the chapter's two-step reverse
+transition, and the range of the generated images.
+
+    python lecture_3/ddpm.py                    # 300 steps on the offline images
+    python lecture_3/ddpm.py --steps 20         # a few seconds, for a first look
+    python lecture_3/ddpm.py --dataset mnist    # the real digits, downloaded on first use
+"""
+
+# %% 01. Imports and configuration
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from cis6270.nets import count_parameters
+from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
+
+NUM_STEPS = 1000  # Diffusion steps K.
+BETA_START = 1e-4  # Noise variance added at step 1.
+BETA_END = 0.02  # Noise variance added at step K.
+
+
+# %% 02. The schedule tensors and the forward corruption
+#
+# alpha_k = 1 - beta_k is the one-step retention, abar_k its running product, and beta-tilde_k the
+# variance of the reverse posterior q(x_{k-1} | x_k, x_0). A zero is prefixed to the betas so that
+# Python index k holds beta_k and the cumulative product starts at abar_0 = 1, the clean state.
+# One Gaussian draw then reaches any level, x_k = a_k x_0 + sigma_k eps with a_k = sqrt(abar_k) and
+# sigma_k = sqrt(1 - abar_k).
+
+betas = torch.cat([
+    torch.zeros(1),  # Placeholder so that index k holds beta_k.
+    torch.linspace(BETA_START, BETA_END, NUM_STEPS),  # beta_1 ... beta_K.
+])
+alphas = 1.0 - betas  # One-step signal retention alpha_k.
+alphas_cumprod = torch.cumprod(alphas, dim=0)  # Cumulative retention.
+alphas_cumprod_prev = torch.cat([
+    torch.ones(1),  # abar_0 = 1.
+    alphas_cumprod[:-1],  # Index k now holds abar_{k-1}.
+])
+posterior_variance = (
+    betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
+)  # Reverse variance beta-tilde_k.
+
+
+def extract(schedule, k, shape):
+    # One schedule value per item, reshaped to broadcast over C, H, W.
+    out = schedule.to(k.device).gather(0, k)
+    return out.reshape(k.shape[0], *((1,) * (len(shape) - 1)))
+
+
+def q_sample(x_0, k, noise=None):
+    if noise is None:
+        noise = torch.randn_like(x_0)  # Mean 0, covariance I, shape of x_0.
+    abar = extract(alphas_cumprod, k, x_0.shape)
+    a_k = abar.sqrt()  # Signal coefficient a_k.
+    sigma_k = (1.0 - abar).sqrt()  # Noise standard deviation sigma_k.
+    return a_k * x_0 + sigma_k * noise, noise  # x_k and the exact noise.
+
+
+# %% 03. Sinusoidal step embedding
+#
+# A raw index is a poor input to a linear layer, so the step is spread over many frequencies first.
+
+
+def timestep_embedding(k, dim):
+    half = dim // 2  # Half the coordinates sine, half cosine.
+    freqs = torch.exp(
+        -math.log(10000.0)
+        * torch.arange(half, device=k.device).to(torch.get_default_dtype()) / half
+    )
+    args = k.to(torch.get_default_dtype())[:, None] * freqs[None, :]  # [B, half].
+    return torch.cat([torch.sin(args), torch.cos(args)], dim=1)  # [B, dim].
+
+
+def num_groups(channels):
+    for g in (8, 4, 2, 1):  # Largest valid GroupNorm group count.
+        if channels % g == 0:
+            return g
+
+
+# %% 04. The time-conditioned residual block
+#
+# The step embedding enters as a per-channel offset between the two convolutions, so every layer
+# operates with the current noise level available to it.
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, time_embed_dim):
+        super().__init__()
+        self.norm1 = nn.GroupNorm(num_groups(in_channels), in_channels)
+        self.conv1 = nn.Conv2d(in_channels, out_channels,
+                               kernel_size=3, padding=1)
+        self.time_proj = nn.Linear(time_embed_dim, out_channels)
+        self.norm2 = nn.GroupNorm(num_groups(out_channels), out_channels)
+        self.conv2 = nn.Conv2d(out_channels, out_channels,
+                               kernel_size=3, padding=1)
+        self.skip = (
+            nn.Identity() if in_channels == out_channels
+            else nn.Conv2d(in_channels, out_channels, kernel_size=1)
+        )
+        self.activation = nn.SiLU()
+
+    def forward(self, x, t_emb):
+        h = self.conv1(self.activation(self.norm1(x)))  # GN, SiLU, Conv_1.
+        h = h + self.time_proj(t_emb)[:, :, None, None]  # Offset per channel.
+        h = self.conv2(self.activation(self.norm2(h)))
+        return h + self.skip(x)  # Residual connection.
+
+
+# %% 05. The U-Net
+#
+# Three scales, 28 by 28 down to 7 by 7 and back, with a skip at each resolution. The output has
+# the shape of the input because it predicts one noise coordinate per pixel.
+
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=1, base_channels=64, time_embed_dim=128):
+        super().__init__()
+        self.time_embed_dim = time_embed_dim
+        self.time_mlp = nn.Sequential(  # Learned representation of step k.
+            nn.Linear(time_embed_dim, time_embed_dim),
+            nn.SiLU(),
+            nn.Linear(time_embed_dim, time_embed_dim),
+        )
+        c1 = base_channels  # 64 channels at 28 x 28.
+        c2 = base_channels * 2  # 128 channels at 14 x 14.
+        c3 = base_channels * 4  # 256 channels at 7 x 7.
+
+        self.in_conv = nn.Conv2d(in_channels, c1, kernel_size=3, padding=1)
+        self.down1 = ResidualBlock(c1, c1, time_embed_dim)
+        self.pool1 = nn.Conv2d(c1, c1, kernel_size=3, stride=2, padding=1)
+        self.down2 = ResidualBlock(c1, c2, time_embed_dim)
+        self.pool2 = nn.Conv2d(c2, c2, kernel_size=3, stride=2, padding=1)
+        self.mid = ResidualBlock(c2, c3, time_embed_dim)
+        self.up1 = nn.ConvTranspose2d(c3, c2, kernel_size=4,
+                                      stride=2, padding=1)
+        self.dec1 = ResidualBlock(c2 + c2, c2, time_embed_dim)
+        self.up2 = nn.ConvTranspose2d(c2, c1, kernel_size=4,
+                                      stride=2, padding=1)
+        self.dec2 = ResidualBlock(c1 + c1, c1, time_embed_dim)
+        self.out_conv = nn.Conv2d(c1, in_channels, kernel_size=1)
+
+    def forward(self, x, k):
+        t_emb = self.time_mlp(
+            timestep_embedding(k, self.time_embed_dim)
+        )  # Learned representation of diffusion step k.
+        h = self.in_conv(x)
+        h1 = self.down1(h, t_emb)  # 28 x 28 encoder features.
+        h = self.pool1(h1)  # 28 -> 14.
+        h2 = self.down2(h, t_emb)  # 14 x 14 encoder features.
+        h = self.pool2(h2)  # 14 -> 7.
+        h = self.mid(h, t_emb)  # Compressed representation.
+        h = self.up1(h)  # 7 -> 14.
+        h = self.dec1(torch.cat([h, h2], dim=1), t_emb)  # Join 14 x 14 skip.
+        h = self.up2(h)  # 14 -> 28.
+        h = self.dec2(torch.cat([h, h1], dim=1), t_emb)  # Join 28 x 28 skip.
+        return self.out_conv(h)  # [B, 1, 28, 28] predicted noise.
+
+
+# %% 06. One reverse step and the sampler
+
+
+@torch.no_grad()
+def p_sample(model, x_k, k):
+    eps_theta = model(x_k, k)  # Predicted noise at step k.
+    beta_k = extract(betas, k, x_k.shape)
+    alpha_k = extract(alphas, k, x_k.shape)
+    abar_k = extract(alphas_cumprod, k, x_k.shape)
+    mean = (
+        x_k - beta_k / (1.0 - abar_k).sqrt() * eps_theta
+    ) / alpha_k.sqrt()  # Reverse mean mu_theta.
+    var = extract(posterior_variance, k, x_k.shape)  # beta-tilde_k.
+    z = torch.randn_like(x_k)  # Fresh Gaussian draw.
+    keep = (k > 1).to(torch.get_default_dtype()).reshape(-1, *((1,) * (x_k.dim() - 1)))
+    return mean + keep * var.sqrt() * z  # No noise at the final step.
+
+
+@torch.no_grad()
+def sample(model, num_images, device):
+    x = torch.randn(num_images, 1, 28, 28, device=device)  # X_K ~ N(0, I).
+    for step in range(NUM_STEPS, 0, -1):  # k = K down to k = 1.
+        k = torch.full((num_images,), step, device=device, dtype=torch.long)
+        x = p_sample(model, x, k)
+    return x  # Generated images in the training range [-1, 1].
+
+
+# %% 07. The dropped weight and one complete reverse transition
+#
+# One interior term of the variational bound is beta_k^2 / (2 beta-tilde_k alpha_k (1 - abar_k))
+# multiplied by the squared noise error, and L_simple sets every such factor to one. On the
+# two-step schedule beta_1 = beta_2 = 0.2 the whole transition is small enough to check by hand,
+# and the mean there has two equal forms, one built from the noise prediction and one from the
+# clean estimate it implies.
+
+
+def elbo_weight(beta_k, posterior, alpha_k, abar_k) -> float:
+    """The schedule constant in front of the squared noise error at one step."""
+    return beta_k ** 2 / (2.0 * posterior * alpha_k * (1.0 - abar_k))
+
+
+def two_step_transition(beta: float, x_k: float, eps_theta: float, z: float) -> dict:
+    """The reverse step at k = 2 of the two-step schedule, through both forms of the mean."""
+    alpha = 1.0 - beta  # alpha_2.
+    abar_1, abar_2 = alpha, alpha ** 2  # 0.8 and 0.64 at beta = 0.2.
+    sigma = math.sqrt(1.0 - abar_2)  # 0.6, the corruption standard deviation at k = 2.
+    mean = (x_k - beta / sigma * eps_theta) / math.sqrt(alpha)  # The noise form of mu_theta.
+    x_0_hat = (x_k - sigma * eps_theta) / math.sqrt(abar_2)  # The clean estimate it implies.
+    weights = (math.sqrt(abar_1) * beta / (1.0 - abar_2),  # Weight on the clean estimate.
+               math.sqrt(alpha) * (1.0 - abar_1) / (1.0 - abar_2))  # Weight on x_k.
+    variance = beta * (1.0 - abar_1) / (1.0 - abar_2)  # beta-tilde_2 = 1/9.
+    return {
+        "mean_from_noise": mean, "mean_from_clean": weights[0] * x_0_hat + weights[1] * x_k,
+        "x_0_hat": x_0_hat, "weights": weights, "variance": variance,
+        "x_{k-1}": mean + math.sqrt(variance) * z,
+        "elbo_weight": elbo_weight(beta, variance, alpha, abar_2),
+    }
+
+
+# %% 08. The report
+def main() -> None:
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.add_argument("--dataset", choices=["toy", "mnist"], default="toy")
+    parser.add_argument("--base-channels", type=int, default=8,
+                        help="U-Net base width; the notes uses 64")
+    parser.set_defaults(steps=300, batch_size=64, lr=2e-3, samples=4)
+    args = parser.parse_args()
+    set_seed(args.seed)
+    device = pick_device(args.device)
+
+    if args.dataset == "toy":
+        from cis6270.data import toy_images
+
+        images = toy_images(512, 28, args.seed).to(device)  # [512, 1, 28, 28] offline images.
+        draw = lambda n: images[torch.randint(0, len(images), (n,), device=device)]  # noqa: E731
+    else:
+        from cis6270.data import mnist_loader
+
+        loader = mnist_loader(batch_size=args.batch_size)
+        stream = iter(loader)
+
+        def draw(n):  # One MNIST batch per call, restarting the loader at the end of an epoch.
+            nonlocal stream
+            batch = next(stream, None)
+            if batch is None:
+                stream = iter(loader)
+                batch = next(stream)
+            return batch[0].to(device)
+
+    model = UNet(in_channels=1, base_channels=args.base_channels, time_embed_dim=128).to(device)
+    if not args.quiet:
+        print(f"device={device} dataset={args.dataset} parameters={count_parameters(model):,}")
+
+    def loss_fn(step: int) -> torch.Tensor:
+        x_0 = draw(args.batch_size)  # [B, 1, 28, 28] clean images in [-1, 1].
+        k = torch.randint(1, NUM_STEPS + 1, (x_0.shape[0],), device=device)  # One level per item.
+        x_k, noise = q_sample(x_0, k)  # Corrupt directly to step k.
+        return F.mse_loss(model(x_k, k), noise)  # L_simple on this batch.
+
+    # The notes trains at 2e-4 for 20 epochs; 2e-3 for a few hundred steps is the
+    # classroom version of the same run and reaches a usable sampler in under a minute.
+    losses = train(model, loss_fn, args.steps, args.lr, quiet=args.quiet)
+    images_out = sample(model, args.samples, device)  # [S, 1, 28, 28] after K reverse steps.
+
+    window = min(100, args.steps)  # A short run averages over everything it has.
+    reverse = two_step_transition(0.2, 7.4, -1.0, 0.3)
+    weight_500 = elbo_weight(betas[500], posterior_variance[500],
+                             alphas[500], alphas_cumprod[500])  # The dropped factor at k = 500.
+    report = {
+        "parameters": count_parameters(model),
+        f"loss, first {window} steps": sum(losses[:window]) / window,
+        f"loss, last {window} steps": sum(losses[-window:]) / window,
+        "beta-tilde at k=500": f"{float(posterior_variance[500]):.5f}",  # 0.01003.
+        # The weight L_simple drops, on the thousand-step grid and on the two-step hand schedule.
+        "ELBO weight, k=500 and two-step": [round(float(weight_500), 6),
+                                            round(reverse["elbo_weight"], 4)],  # 0.005508, 0.625.
+        "two-step reverse mean, two forms": [round(reverse["mean_from_noise"], 4),
+                                             round(reverse["mean_from_clean"], 4)],  # 8.646.
+        "two-step x_0 estimate and weights":
+            [round(reverse["x_0_hat"], 4)] + [round(w, 4) for w in reverse["weights"]],
+        "two-step reverse variance and x_1": [round(reverse["variance"], 4),
+                                              round(reverse["x_{k-1}"], 4)],  # 1/9 and 8.746.
+        "generated range": [round(float(images_out.min()), 3), round(float(images_out.max()), 3)],
+        "generated standard deviation": float(images_out.std()),
+    }
+    print_report(f"DDPM on {args.dataset} images", report)
+    saved = save_run(args.out, model=model, config=vars(args), report=report, losses=losses)
+    if saved is not None:
+        try:  # Without torchvision the run still trains, samples and saves its weights.
+            from torchvision.utils import save_image
+
+            display = ((images_out.cpu() + 1.0) / 2.0).clamp(0.0, 1.0)  # [-1, 1] -> [0, 1].
+            save_image(display, saved / "samples.png", nrow=4)  # A grid of the generated images.
+        except ImportError:
+            pass
+
+
+# %% 09. Run the script
+if __name__ == "__main__":
+    main()

lecture_3/forward_diffusion.py

@@ -0,0 +1,204 @@
+"""The schedule and the Gaussian corruption it drives.
+
+The variance-preserving path is fixed once a noising rate is chosen. Its accumulated rate gives a
+signal coefficient and a noise standard deviation at every level, a thousand-step grid tracks the
+continuous curve to three decimal places, and one Gaussian draw is enough to reach any level from
+a clean image. Nothing here learns anything, so every number is arithmetic or a sample average
+against a law written in closed form.
+
+Chapter 3 runs its diffusion clock from the data, so step k = 0 holds a clean image and corruption
+grows with k, up to pure noise at k = K.
+
+Chapter 3, Sections 3.11 and 3.12 of the notes. Runs on the offline 28-by-28 images. Prints the
+four noise levels of the schedule, the discrete grid against the continuous curve, the chapter's
+one pixel at step 207, the worked training batch at abar = 0.8, and the corruption law measured on
+a batch.
+
+    python lecture_3/forward_diffusion.py              # 16 levels, 1024 draws at each
+    python lecture_3/forward_diffusion.py --steps 20   # a few seconds, for a first look
+"""
+
+# %% 01. Imports and configuration
+import math
+
+import torch
+
+from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed
+
+NUM_STEPS = 1000  # Diffusion steps K.
+BETA_START = 1e-4  # Noise variance added at step 1.
+BETA_END = 0.02  # Noise variance added at step K.
+BETA_MIN = 0.1  # beta(0) of the continuous schedule these endpoints discretize.
+BETA_MAX = 20.0  # beta(1) of the same schedule.
+
+# `--steps` is the number of levels at which the corruption law is measured and `--samples` the
+# number of draws taken at each of them; nothing in this file is trained.
+
+
+# %% 02. The schedule tensors
+#
+# alpha_k = 1 - beta_k is the one-step retention and abar_k its running product. A zero is prefixed
+# to the betas so that Python index k holds beta_k and the cumulative product starts at abar_0 = 1,
+# the clean state. The two corruption coefficients are the square roots of abar_k and 1 - abar_k.
+
+betas = torch.cat([
+    torch.zeros(1),  # Placeholder so that index k holds beta_k.
+    torch.linspace(BETA_START, BETA_END, NUM_STEPS),  # beta_1 ... beta_K.
+])
+alphas = 1.0 - betas  # One-step signal retention alpha_k.
+alphas_cumprod = torch.cumprod(alphas, dim=0)  # Cumulative retention abar_k.
+a = alphas_cumprod.sqrt()  # Signal coefficient a_k.
+sigma = (1.0 - alphas_cumprod).sqrt()  # Noise standard deviation sigma_k.
+
+
+# %% 03. Forward corruption
+#
+# x_k = a_k x_0 + sigma_k eps, the Gaussian transition q(x_k | x_0) = N(a_k x_0, sigma_k^2 I) read
+# off at step k. One draw reaches any level, so training never simulates the path.
+
+
+def extract(schedule, k, shape):
+    # One schedule value per item, reshaped to broadcast over C, H, W.
+    out = schedule.to(k.device).gather(0, k)
+    return out.reshape(k.shape[0], *((1,) * (len(shape) - 1)))
+
+
+def q_sample(x_0, k, noise=None):
+    if noise is None:
+        noise = torch.randn_like(x_0)  # Mean 0, covariance I, shape of x_0.
+    abar = extract(alphas_cumprod, k, x_0.shape)
+    a_k = abar.sqrt()  # Signal coefficient a_k.
+    sigma_k = (1.0 - abar).sqrt()  # Noise standard deviation sigma_k.
+    return a_k * x_0 + sigma_k * noise, noise  # x_k and the exact noise.
+
+
+# %% 04. The continuous schedule behind the grid
+#
+# beta(t) = 0.1 + 19.9 t integrates to B(t) = 0.1 t + 9.95 t^2 and a_t = exp(-B(t)/2). Dividing the
+# per-step variances 1e-4 and 0.02 by the step size 1/1000 recovers the two endpoints, so the grid
+# and the curve describe one path. Inverting a_t = exp(-B(t)/2) locates a chosen noise level.
+
+
+def noising_rate(t: float) -> float:
+    """beta(t) = beta_min + (beta_max - beta_min) t, the linear noising rate."""
+    return BETA_MIN + (BETA_MAX - BETA_MIN) * t
+
+
+def continuous_signal(t: float) -> float:
+    """a_t = exp(-B(t)/2) with B(t) = 0.1 t + 9.95 t^2, the curve the discrete betas track."""
+    return math.exp(-(BETA_MIN * t + 0.5 * (BETA_MAX - BETA_MIN) * t ** 2) / 2.0)
+
+
+def level_time(a_t: float) -> dict:
+    """Solve B(t) = -2 log a_t for the time at which the signal coefficient is a_t.
+
+    The quadratic 9.95 t^2 + 0.1 t - B = 0 has one positive root, which is the time the notes
+    tabulates beside the noise standard deviation and the noising rate there.
+    """
+    accumulated = -2.0 * math.log(a_t)  # B(t) at the requested level.
+    quadratic = 0.5 * (BETA_MAX - BETA_MIN)  # Coefficient of t^2 in B(t).
+    t = (-BETA_MIN + math.sqrt(BETA_MIN ** 2 + 4.0 * quadratic * accumulated)) / (2.0 * quadratic)
+    return {"t": t, "sigma_t": math.sqrt(1.0 - a_t ** 2), "beta_t": noising_rate(t)}
+
+
+# %% 05. One training batch at the running level
+#
+# Three clean scalars at a level with abar = 0.8 give the batch the notes carries through the
+# denoising loss by hand: the corrupted values, the squared error against a supposed prediction,
+# the score targets the noise regression never forms, and the clean estimates the prediction
+# implies.
+
+
+def worked_batch(abar: float, x_0, epsilon, predicted) -> dict:
+    """One training batch at a level with the given cumulative retention, as the notes works it."""
+    a_k, sigma_k = math.sqrt(abar), math.sqrt(1.0 - abar)  # 0.8944 and 0.4472 at abar = 0.8.
+    x_0, epsilon, predicted = (torch.tensor(v) for v in (x_0, epsilon, predicted))
+    x_k = a_k * x_0 + sigma_k * epsilon  # [3] the corrupted batch.
+    return {
+        "a_k": a_k, "sigma_k": sigma_k,
+        "x_k": [round(float(v), 4) for v in x_k],
+        "loss": float((predicted - epsilon).square().mean()),  # L_simple on this batch.
+        "score_targets": [round(float(-e / sigma_k), 3) for e in epsilon],
+        "x_0_hat": [round(float((x - sigma_k * p) / a_k), 2) for x, p in zip(x_k, predicted)],
+    }
+
+
+# %% 06. The corruption law, measured
+#
+# q(x_k | x_0) is Gaussian with mean a_k x_0 and variance sigma_k^2 in every coordinate. Corrupting
+# one clean image many times at a fixed level and averaging over the draws therefore recovers
+# a_k x_0 pixel by pixel, and the spread around it recovers sigma_k^2. The largest gap over a
+# spread of levels falls as one over the square root of the number of draws, and a wrong
+# coefficient in either of the two lines of `q_sample` would leave a gap that does not.
+
+
+def measure_law(x_0, levels, repeats: int, generator=None) -> dict:
+    """Corrupt one clean image ``repeats`` times at each level and compare with q(x_k | x_0)."""
+    batch = x_0.expand(repeats, *x_0.shape[1:])  # [R, 1, 28, 28] one image, repeated.
+    mean_gap, variance_gap = 0.0, 0.0
+    for step in levels:
+        k = torch.full((repeats,), int(step), dtype=torch.long)  # [R] the same level throughout.
+        noise = torch.randn(batch.shape, generator=generator)  # [R, 1, 28, 28] fresh draws.
+        x_k, _ = q_sample(batch, k, noise)  # [R, 1, 28, 28] corrupted at this level.
+        centre = float(a[step]) * x_0[0]  # [1, 28, 28] the conditional mean a_k x_0.
+        mean_gap = max(mean_gap, float((x_k.mean(dim=0) - centre).abs().mean()))
+        spread = float(x_k.var(dim=0, unbiased=False).mean())  # Pooled over the pixels.
+        variance_gap = max(variance_gap, abs(spread - float(sigma[step]) ** 2))
+    return {"mean_gap": mean_gap, "variance_gap": variance_gap}
+
+
+# %% 07. The report
+def main() -> None:
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.set_defaults(steps=16, samples=1024)
+    args = parser.parse_args()
+    set_seed(args.seed)
+    device = pick_device(args.device)  # Printed for consistency; this file runs on the CPU.
+    generator = torch.Generator().manual_seed(args.seed)  # Local generator, so the report repeats.
+
+    from cis6270.data import toy_images
+
+    image = toy_images(1, 28, args.seed)  # [1, 1, 28, 28] one offline image in [-1, 1].
+    levels = torch.linspace(1, NUM_STEPS, args.steps).round().long().tolist()  # Spread over K.
+    law = measure_law(image, levels, args.samples, generator)
+
+    # The four levels of the notes' table, located by inverting a_t = exp(-B(t)/2).
+    located = [level_time(level) for level in (0.98, 0.8, 0.35, 0.04)]
+    crossing = level_time(1.0 / math.sqrt(2.0))  # Where a_t and sigma_t meet.
+    batch = worked_batch(0.8, (2.0, -1.0, 0.5), (0.5, -0.2, 1.0), (0.4, -0.1, 0.8))
+
+    report = {
+        "levels measured, draws each": [len(levels), args.samples],
+        "largest gap, sample mean": law["mean_gap"],  # Zero up to the sampling error.
+        "largest gap, sample variance": law["variance_gap"],
+        "times at a_t = 0.98, 0.8, 0.35, 0.04": [round(v["t"], 4) for v in located],
+        "sigma_t at the same four levels": [round(v["sigma_t"], 4) for v in located],
+        "beta(t) at the same four levels": [round(v["beta_t"], 4) for v in located],
+        # Where the two curves cross, and the value they share there.
+        "crossing time and a_t = sigma_t": [round(crossing["t"], 5),
+                                            round(crossing["sigma_t"], 4)],  # 0.25896, 0.7071.
+        # The discrete coefficient against the continuous curve, at step 207 and at the last step.
+        "abar_207": round(float(alphas_cumprod[207]), 4),  # 0.6402.
+        "a_207 discrete and continuous": [round(float(a[207]), 4),
+                                          round(continuous_signal(0.207), 4)],  # 0.8001, 0.7997.
+        "sigma_207 discrete and continuous": [round(float(sigma[207]), 4),
+                                              round(math.sqrt(1.0 - continuous_signal(0.207) ** 2),
+                                                    4)],  # 0.5999, 0.6004.
+        "a_1000 discrete and continuous": [round(float(a[1000]), 4),
+                                           round(continuous_signal(1.0), 4)],  # 0.0064, 0.0066.
+        "one pixel at k=207, x_0=0.5, eps=-1": float(a[207] * 0.5 - sigma[207]),  # -0.1998.
+        "a_500, where noise exceeds signal": round(float(a[500]), 4),  # 0.2803.
+        "worked batch at abar=0.8, x_k": batch["x_k"],  # 2.0125, -0.9839, 0.8944.
+        "worked batch loss": batch["loss"],  # 0.02.
+        "worked batch score targets": batch["score_targets"],  # -1.118, 0.447, -2.236.
+        "worked batch x_0 estimates": batch["x_0_hat"],  # 2.05, -1.05, 0.60.
+    }
+    if not args.quiet:
+        print(f"device={device} levels={args.steps} draws={args.samples} seed={args.seed}")
+    print_report("The variance-preserving corruption", report)
+    save_run(args.out, config=vars(args), report=report)
+
+
+# %% 08. Run the script
+if __name__ == "__main__":
+    main()

lecture_3/guidance.py

@@ -0,0 +1,280 @@
+"""Steering a sampler by adding one gradient to the score.
+
+Bayes' rule splits the conditional score into the unconditional score plus the gradient of the log
+likelihood of the condition, so every form of control in this chapter is one extra term in the
+same place. A classifier supplies that term directly, a conditional and an unconditional
+prediction of one network supply it as their difference, and a differentiable reward supplies it
+as its own gradient, with several objectives combined through weights on the simplex. The target
+is the four-Gaussian mixture with one label per mode, where every score is a closed form, so the
+guidance weight shows up directly in the mode fractions.
+
+Chapter 3 runs its diffusion clock from the data, so t = 0 holds a clean sample and corruption
+grows with t. Generation traverses the same interval from t = 1 back to t = 0.
+
+Chapter 3, Sections 3.37 through 3.45 of the notes. Prints the mode fractions at several guidance
+weights, the chapter's worked guided score and guided backward step, the two candidates under two
+weightings, and the schedule value at t = 0.6.
+
+    python lecture_3/guidance.py              # 200 reverse steps per setting
+    python lecture_3/guidance.py --steps 20   # a coarse grid, for a first look
+"""
+
+# %% 01. Imports and configuration
+import math
+
+import torch
+
+from cis6270.data import GAUSSIAN_CENTERS, GAUSSIAN_STD
+from cis6270.metrics import mode_report
+from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed
+
+BETA_MIN = 0.1  # beta(0) of the linear variance-preserving schedule.
+BETA_MAX = 20.0  # beta(1) of the same schedule.
+TARGET_LABEL = 0  # The mode at (2, 2), which every guided run is steered toward.
+REWARD_SCALE = 2.0  # Distance from the origin to a mode, used to standardize the two rewards.
+
+
+# %% 02. The target, its labels and the exact scores
+#
+# Each mode carries its own label, so p_t(x | c) is a single Gaussian with mean a_t mu_c and
+# variance a_t^2 s^2 + sigma_t^2, and p_t(x) is the equal mixture of the four. Both scores are
+# closed forms, and nothing in this file is trained.
+
+
+def schedule_tensors(num_steps: int) -> dict:
+    """The DDPM tensors on a grid of K steps, read off abar_k = exp(-B(t_k))."""
+    grid = (torch.arange(0, num_steps + 1).to(torch.get_default_dtype())
+            / num_steps)  # [K+1] the times t_k.
+    alphas_cumprod = torch.exp(-(BETA_MIN * grid + 0.5 * (BETA_MAX - BETA_MIN) * grid ** 2))
+    previous = torch.cat([torch.ones(1), alphas_cumprod[:-1]])  # abar_{k-1}.
+    alphas = alphas_cumprod / previous  # One-step retention alpha_k.
+    betas = 1.0 - alphas  # Per-step variance beta_k.
+    return {"betas": betas, "alphas": alphas, "alphas_cumprod": alphas_cumprod,
+            "posterior_variance": betas * (1.0 - previous) / (1.0 - alphas_cumprod)}
+
+
+def component_log_weights(x: torch.Tensor, a_t: float):
+    """Log of each mode's unnormalized contribution at x, and the offsets to the four means."""
+    variance = a_t ** 2 * GAUSSIAN_STD ** 2 + (1.0 - a_t ** 2)  # Width of p_t.
+    offset = x[:, None, :] - a_t * GAUSSIAN_CENTERS.to(x)[None, :, :]  # [B, 4, 2].
+    return -(offset ** 2).sum(dim=2) / (2.0 * variance), offset, variance
+
+
+def unconditional_score(x: torch.Tensor, a_t: float) -> torch.Tensor:
+    """The score of the equal mixture at level a_t. [B, 2] -> [B, 2]."""
+    log_weight, offset, variance = component_log_weights(x, a_t)
+    responsibility = log_weight.softmax(dim=1)[:, :, None]  # [B, 4, 1] posterior over the modes.
+    return (responsibility * (-offset / variance)).sum(dim=1)  # [B, 2].
+
+
+def conditional_score(x: torch.Tensor, a_t: float, label: int) -> torch.Tensor:
+    """The score of p_t(. | c), a single Gaussian centered on a_t mu_c. [B, 2] -> [B, 2]."""
+    _, offset, variance = component_log_weights(x, a_t)
+    return -offset[:, label, :] / variance  # [B, 2].
+
+
+# %% 03. Conditioning the score and the guidance scale
+#
+# s_t(x | c) = s_t(x) + grad log p_t(c | x), and attaching a scale gamma to the second term samples
+# the tilted density p_t(x) p_t(c | x)^gamma. At gamma = 0 generation is unconditional, at
+# gamma = 1 it follows Bayes' rule, and at gamma > 1 the condition is weighted more heavily.
+
+
+def classifier_gradient(x: torch.Tensor, a_t: float, label: int) -> torch.Tensor:
+    """grad log p_t(c | x), taken as the difference of the two scores. [B, 2] -> [B, 2]."""
+    return conditional_score(x, a_t, label) - unconditional_score(x, a_t)
+
+
+def guided_score(score: torch.Tensor, gradient: torch.Tensor, gamma: float) -> torch.Tensor:
+    """s_theta + gamma grad log p_phi(c | x), the guided score of the notes."""
+    return score + gamma * gradient
+
+
+def guided_noise(eps: torch.Tensor, gradient: torch.Tensor, gamma: float, sigma_k: float):
+    """eps_theta - gamma sigma_t grad log p_phi(c | x), the same guidance in the noise."""
+    return eps - gamma * sigma_k * gradient
+
+
+def classifier_free_field(unconditional, conditional, w: float):
+    """F_uncond + w (F_cond - F_uncond), which needs no separate classifier."""
+    return unconditional + w * (conditional - unconditional)
+
+
+# %% 04. Reward guidance and several objectives at once
+#
+# A reward is a differentiable scalar that is larger for more desirable states, and exponential
+# reweighting turns it into a density whose score is s_t(x) + eta grad R_t(x). With M objectives,
+# the weights are clamped to be nonnegative and divided by their sum, which is the simplex
+# condition of the notes, and the gradient is linear in them.
+
+
+def reward_model(z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+    """Two standardized properties of the state, its reach to the right and its reach upward.
+
+    The notes train a small network on the latent and the time here. Two coordinates of a planar
+    state need no network, so the properties are read off directly and keep the same signature.
+    """
+    return z / REWARD_SCALE  # [B, 2] one standardized property per coordinate.
+
+
+def normalize_weights(lambdas, device):
+    """Clamp the weights to be nonnegative and divide by their sum, the simplex condition."""
+    w = torch.as_tensor(lambdas, dtype=torch.get_default_dtype(), device=device)
+    w = w.clamp_min(0.0)
+    return w / w.sum().clamp_min(1e-8)
+
+
+def weighted_reward(reward_model, z, t, lambdas):
+    """R_lambda(z) = sum_m lambda_m r_m(z), one scalar per batch item."""
+    r_hat = reward_model(z, t)                       # [B, M]
+    return (r_hat * normalize_weights(lambdas, z.device)).sum(dim=1)
+
+
+def reward_gradient(reward_model, z, t, lambdas):
+    """grad_z R_lambda(z), taken inside a no-grad sampler. [B, d] -> [B, d]."""
+    with torch.enable_grad():
+        state = z.detach().requires_grad_(True)
+        R = weighted_reward(reward_model, state, t, lambdas).sum()
+        grad, = torch.autograd.grad(R, state)
+    return grad.detach()
+
+
+def scalarize(values, lambdas) -> float:
+    """The weighted sum of normalized objectives, R_lambda(x), on one candidate."""
+    weights = normalize_weights(lambdas, torch.device("cpu"))
+    return float((torch.as_tensor(values, dtype=torch.get_default_dtype()) * weights).sum())
+
+
+# %% 05. Scheduled guidance
+#
+# A property predictor is not equally reliable at every noise level, so the strength is allowed to
+# depend on time, gamma(t) = gamma_max zeta(t). The schedule zeta(t) = 4 t (1 - t) vanishes at both
+# ends of the interval and peaks at one in the middle.
+
+
+def schedule_value(t: float) -> float:
+    """zeta(t) = 4 t (1 - t), between zero and one."""
+    return 4.0 * t * (1.0 - t)
+
+
+# %% 06. The guided sampler
+#
+# The reverse step is the one of Section 3.26 with the guided noise prediction in place of the
+# raw one. Classifier-free guidance is applied first, then the reward term, which is the order the
+# notes implements.
+
+
+@torch.no_grad()
+def sample_diffusion(schedule, num_points, label, w=0.0, eta=0.0, lambdas=(1.0, 0.0),
+                     scheduled=False, generator=None) -> torch.Tensor:
+    """Draw points under classifier-free guidance toward ``label`` and reward steering."""
+    num_steps = len(schedule["betas"]) - 1
+    z = torch.randn(num_points, 2, generator=generator)  # [N, 2] the terminal Gaussian.
+    for step in range(num_steps, 0, -1):  # k = K down to k = 1.
+        t = step / num_steps  # Normalized time, the input the network would receive.
+        abar = float(schedule["alphas_cumprod"][step])
+        a_k, sigma_k = math.sqrt(abar), math.sqrt(1.0 - abar)
+        eps = -sigma_k * unconditional_score(z, a_k)  # [N, 2] unconditional noise.
+        if w > 0.0:
+            conditional = -sigma_k * conditional_score(z, a_k, label)  # [N, 2].
+            strength = w * (schedule_value(t) if scheduled else 1.0)  # gamma(t) = gamma_max zeta.
+            eps = classifier_free_field(eps, conditional, strength)
+        if eta > 0.0:
+            grad_R = reward_gradient(reward_model, z, torch.full((num_points,), t), lambdas)
+            eps = eps - eta * sigma_k * grad_R  # Reward guidance in the noise parameterization.
+        beta_k, alpha_k = float(schedule["betas"][step]), float(schedule["alphas"][step])
+        mean = (z - beta_k / sigma_k * eps) / math.sqrt(alpha_k)  # Reverse mean mu_theta.
+        if step > 1:
+            variance = float(schedule["posterior_variance"][step])
+            z = mean + math.sqrt(variance) * torch.randn(z.shape, generator=generator)
+        else:
+            z = mean
+    return z  # [N, 2].
+
+
+# %% 07. The report
+def main() -> None:
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.set_defaults(steps=200, samples=2048)
+    args = parser.parse_args()
+    set_seed(args.seed)
+    device = pick_device(args.device)  # Two coordinates per point; a CPU is enough.
+    schedule = schedule_tensors(args.steps)
+
+    def run(**settings) -> list:
+        generator = torch.Generator().manual_seed(args.seed)  # The same starting noise each time.
+        points = sample_diffusion(schedule, args.samples, TARGET_LABEL, generator=generator,
+                                  **settings)
+        return mode_report(points)["mode_fractions"]
+
+    # The conditional score here is exact, so at w = 1 the condition already pins the mode and
+    # the fractions saturate. The graded part of the sweep therefore sits below one.
+    report = {}
+    for w in (0.0, 0.1, 0.25, 0.5, 1.0):  # The target mode is the first entry of each list.
+        report[f"mode fractions at w={w}"] = run(w=w)
+    report["mode fractions, w=1 scheduled"] = run(w=1.0, scheduled=True)
+    for eta in (0.5, 2.0):  # Reward steering toward the two right-hand modes.
+        report[f"mode fractions at eta={eta}, lambda=(1,0)"] = run(eta=eta, lambdas=(1.0, 0.0))
+    report["mode fractions at eta=2, lambda=(0.7,0.3)"] = run(eta=2.0, lambdas=(0.7, 0.3))
+
+    # The two routes to one guided noise prediction, on a batch at a middle level.
+    probe = torch.randn(256, 2, generator=torch.Generator().manual_seed(args.seed + 1))
+    a_t, sigma_t = 0.8, 0.6
+    by_classifier = guided_noise(-sigma_t * unconditional_score(probe, a_t),
+                                 classifier_gradient(probe, a_t, TARGET_LABEL), 2.0, sigma_t)
+    by_difference = classifier_free_field(-sigma_t * unconditional_score(probe, a_t),
+                                          -sigma_t * conditional_score(probe, a_t, TARGET_LABEL),
+                                          2.0)
+    gap = float((by_classifier - by_difference).abs().max())
+    report["classifier and classifier-free gap"] = gap
+
+    # The chapter's worked numbers, none of which needs a sampler.
+    score, gradient = torch.tensor([[-1.0, -0.1]]), torch.tensor([[1.3, 1.1]])
+    for gamma in (1.0, 2.0):
+        direction = guided_score(score, gradient, gamma)[0]
+        report[f"guided direction at gamma={gamma}"] = [round(float(v), 4) for v in direction]
+        report[f"its length and angle at gamma={gamma}"] = [
+            round(float(direction.norm()), 4),
+            round(math.degrees(math.atan2(float(direction[1]), float(direction[0]))), 1),
+        ]
+    cfg_eps = float(classifier_free_field(torch.tensor(0.50), torch.tensor(0.10), 2.0))
+    report["classifier-free eps at w=2"] = cfg_eps  # -0.30, past the conditional 0.10.
+    report["its implied score"] = -cfg_eps / 0.5  # 0.6, at sigma_t = 0.5.
+
+    # One guided backward step from the state of the deterministic sampler, through both
+    # parameterizations. A disagreement between them is a missing sigma_t or a missing one half.
+    x_t, beta_t, sigma_t, eps_theta = 1.2, 2.0, 0.5, 0.8  # The state of Section 3.35.
+    gradient, gamma, dt = 0.6, 2.0, 0.1  # The classifier gradient, the scale and the step.
+    unguided = -0.5 * beta_t * (x_t - eps_theta / sigma_t)  # 0.4, the unguided velocity.
+    guided = unguided - 0.5 * gamma * beta_t * gradient  # -0.8, half of gamma g^2 on the gradient.
+    guided_eps = float(guided_noise(torch.tensor(eps_theta), torch.tensor(gradient),
+                                    gamma, sigma_t))  # 0.2, the same guidance in the noise.
+    report["guided velocity and next state"] = [round(guided, 4),
+                                                round(x_t - dt * guided, 4)]  # -0.8 and 1.28.
+    report["guided eps and its velocity"] = [
+        round(guided_eps, 4),  # 0.2, whose implied score is -0.4.
+        round(-0.5 * beta_t * (x_t - guided_eps / sigma_t), 4),  # -0.8, the same velocity.
+    ]
+
+    # Two candidates under two weightings, and two objectives that disagree.
+    candidates = ((1.2, 0.2), (0.5, 1.1))  # Normalized binding and solubility of two peptides.
+    for lambdas in ((0.7, 0.3), (0.3, 0.7)):
+        report[f"candidate rewards, lambda={lambdas}"] = [round(scalarize(v, lambdas), 4)
+                                                          for v in candidates]
+    competing = 0.6 * torch.tensor([1.0, 0.8]) + 0.4 * torch.tensor([0.5, -1.2])
+    report["competing gradients combined"] = [round(float(v), 4) for v in competing]  # (0.8, 0).
+    rewards = torch.tensor([1.2, 0.8, 0.3])
+    report["selection weights at eta=2"] = [round(float(p), 4)
+                                            for p in (2.0 * rewards).softmax(0)]
+    report["schedule value at t=0.6"] = schedule_value(0.6)  # 0.96.
+
+    if not args.quiet:
+        print(f"device={device} points={args.samples} steps={args.steps} "
+              f"target mode={GAUSSIAN_CENTERS[TARGET_LABEL].tolist()}")
+    print_report("Guidance on the four-Gaussian target", report)
+    save_run(args.out, config=vars(args), report=report)
+
+
+# %% 08. Run the script
+if __name__ == "__main__":
+    main()

lecture_3/latent_diffusion.py

@@ -0,0 +1,309 @@
+"""AMP-Diffusion, the same chain run in the latent space of a protein language model.
+
+Gaussian corruption needs a state that can be added to, and an amino acid sequence has no such
+structure. A frozen encoder supplies one. The peptide becomes a tensor of residue embeddings, the
+diffusion runs there, and the language-model head maps the generated latent back to residues.
+Against the image model the state is the latent, a transformer replaces the U-Net, the network
+predicts the clean latent in place of the noise, and a mask marks which positions hold real
+residues.
+
+Chapter 3 runs its diffusion clock from the data, so step k = 0 holds a clean latent and
+corruption grows with k, up to pure noise at k = K.
+
+Chapter 3, Sections 3.21, 3.29 and 3.46 of the notes. Trained on toy embeddings by default, and
+on ESM-2 embeddings of the course peptides with --dataset esm2. Prints the corrupted coordinate at
+k = 500, one reverse step there, the masked loss, and the polar fraction before and after the
+constraint.
+
+    python lecture_3/latent_diffusion.py                 # 800 steps on offline embeddings
+    python lecture_3/latent_diffusion.py --steps 20      # a few seconds, for a first look
+    python lecture_3/latent_diffusion.py --dataset esm2  # real ESM-2 latents and its head
+"""
+
+# %% 01. Imports and configuration
+import math
+
+import torch
+import torch.nn as nn
+
+from cis6270.nets import count_parameters
+from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
+
+MODEL_NAME = "facebook/esm2_t33_650M_UR50D"  # Pretrained ESM-2 encoder.
+EMBED_DIM = 1280  # ESM-2 residue embedding width D.
+MAX_LEN = 32  # Token positions L per padded peptide.
+NUM_LAYERS = 6  # Transformer layers in the denoiser.
+NUM_STEPS = 1000  # Diffusion steps K.
+AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY"  # The 20 standard amino acids.
+POLAR_RESIDUES = set("STNQYCHKRDE")  # Polar or charged side chains.
+POLAR_COLUMNS = [AMINO_ACIDS.index(a) for a in "STNQYCHKRDE"]
+
+
+# %% 02. The schedule tensors
+#
+# The image tensors of Section 3.28, rebuilt here so the file stands alone. Index k holds the
+# quantity at step k, with a zero prefixed so that abar_0 = 1 at the clean latent.
+
+betas = torch.cat([torch.zeros(1), torch.linspace(1e-4, 0.02, NUM_STEPS)])  # beta_k.
+alphas = 1.0 - betas  # One-step retention alpha_k.
+alphas_cumprod = torch.cumprod(alphas, dim=0)  # Cumulative retention abar_k.
+alphas_cumprod_prev = torch.cat([torch.ones(1), alphas_cumprod[:-1]])  # abar_{k-1}.
+posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)  # beta-tilde_k.
+
+
+# %% 03. Masked corruption
+#
+# Z_k = (sqrt(abar_k) Z_0 + sqrt(1 - abar_k) eps) * M. The reshape to [B, 1, 1] broadcasts one
+# coefficient per sequence across all L positions and all D coordinates, and the mask holds padded
+# positions at exactly zero so the denoiser spends no capacity on them.
+
+
+def q_sample_latent(z0, k, mask, noise=None):
+    if noise is None:
+        noise = torch.randn_like(z0)  # [B, L, D] standard Gaussian.
+    abar = alphas_cumprod.to(z0.device)[k].reshape(-1, 1, 1)
+    z_k = abar.sqrt() * z0 + (1.0 - abar).sqrt() * noise
+    return z_k * mask.unsqueeze(-1), noise  # Padding stays exactly zero.
+
+
+def timestep_embedding(k, dim):
+    half = dim // 2  # Half the coordinates sine, half cosine.
+    freqs = torch.exp(-math.log(10000.0)
+                      * torch.arange(half, device=k.device).to(torch.get_default_dtype())
+                      / half)
+    args = k.to(torch.get_default_dtype())[:, None] * freqs[None, :]  # [B, half].
+    return torch.cat([torch.sin(args), torch.cos(args)], dim=1)  # [B, dim].
+
+
+# %% 04. The latent denoiser and its loss
+#
+# The step embedding, a learned position embedding and the latent are added before the transformer
+# acts on them, so the step reaches every residue as one shared offset and the attention layers mix
+# the positions. The loss counts valid residues only.
+
+
+class LatentDenoiser(nn.Module):
+    def __init__(self, embed_dim=EMBED_DIM, num_layers=NUM_LAYERS,
+                 num_heads=8, max_len=MAX_LEN):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.time_mlp = nn.Sequential(  # Same sinusoidal step embedding.
+            nn.Linear(embed_dim, embed_dim),
+            nn.SiLU(),
+            nn.Linear(embed_dim, embed_dim),
+        )
+        self.pos_embed = nn.Embedding(max_len, embed_dim)  # Residue index.
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=embed_dim, nhead=num_heads,
+            dim_feedforward=4 * embed_dim, batch_first=True,
+        )
+        self.transformer = nn.TransformerEncoder(
+            encoder_layer, num_layers=num_layers,
+            enable_nested_tensor=False,  # Keep a padded batch on the ordinary attention path.
+        )
+        self.out = nn.Sequential(nn.LayerNorm(embed_dim),
+                                 nn.Linear(embed_dim, embed_dim))
+
+    def forward(self, z_k, k, attention_mask):
+        t_emb = self.time_mlp(timestep_embedding(k, self.embed_dim))
+        positions = torch.arange(z_k.shape[1], device=z_k.device)
+        h = (z_k
+             + t_emb.unsqueeze(1)  # Shared across positions.
+             + self.pos_embed(positions).unsqueeze(0))  # Shared across batch.
+        h = self.transformer(h, src_key_padding_mask=~attention_mask.bool())
+        return self.out(h)  # Predicted clean latent Z0-hat.
+
+
+def latent_loss(z0_hat, z0, mask):
+    error = (z0_hat - z0) ** 2  # [B, L, D] squared error.
+    error = error * mask.unsqueeze(-1)  # Zero out padded positions.
+    return error.sum() / (mask.sum() * z0.shape[-1])  # Valid tokens only.
+
+
+# %% 05. Reverse latent sampling
+#
+# The network predicts the clean latent, so the noise it implies is recovered by rearranging the
+# corruption equation, and the reverse mean is then the image sampler's expression unchanged.
+
+
+@torch.no_grad()
+def p_sample_latent(model, z_k, k, mask):
+    z0_hat = model(z_k, k, mask)  # Predict the clean latent.
+    abar = alphas_cumprod.to(z_k.device)[k].reshape(-1, 1, 1)
+    eps_hat = (z_k - abar.sqrt() * z0_hat) / (1.0 - abar).sqrt()
+    beta_k = betas.to(z_k.device)[k].reshape(-1, 1, 1)
+    alpha_k = alphas.to(z_k.device)[k].reshape(-1, 1, 1)
+    mean = (z_k - beta_k / (1.0 - abar).sqrt() * eps_hat) / alpha_k.sqrt()
+    var = posterior_variance.to(z_k.device)[k].reshape(-1, 1, 1)
+    xi = torch.randn_like(z_k)
+    keep = (k > 1).to(torch.get_default_dtype()).reshape(-1, 1, 1)
+    return (mean + keep * var.sqrt() * xi) * mask.unsqueeze(-1)
+
+
+@torch.no_grad()
+def sample_latents(model, num_sequences, length, width, device):
+    mask = torch.ones(num_sequences, length, device=device)
+    z = torch.randn(num_sequences, length, width, device=device)
+    for step in range(NUM_STEPS, 0, -1):  # Z_K down to Z_0.
+        k = torch.full((num_sequences,), step, device=device,
+                       dtype=torch.long)
+        z = p_sample_latent(model, z, k, mask)
+    return z
+
+
+# %% 06. Decoding back to amino acids
+#
+# The generated latent is unstandardized first, then read through the head that was trained to
+# score residues from exactly these embeddings. The head also scores special tokens and ambiguity
+# codes, so the columns are restricted to the twenty canonical amino acids before the argmax.
+
+
+@torch.no_grad()
+def decode(z, head, columns, z_mean, z_std):
+    latent = z * z_std + z_mean  # Undo normalization.
+    logits = head(latent)  # Full vocabulary logits.
+    logits = logits.index_select(-1, columns)  # 20 canonical residues.
+    scores, choices = logits.max(dim=-1)  # Argmax at each position.
+    return logits, choices
+
+
+def composition_proxies(seq):
+    """Net charge per residue and polar/charged fraction."""
+    L = max(len(seq), 1)
+    charge = (seq.count("K") + seq.count("R")
+              - seq.count("D") - seq.count("E")) / L
+    polar = sum(1 for a in seq if a in POLAR_RESIDUES) / L
+    return charge, polar
+
+
+# %% 07. The composition constraint at decoding
+#
+# Guidance makes a high polar fraction more likely without guaranteeing it. The requirement is
+# imposed afterwards by substituting residues at the positions where substitution costs the
+# least logit. One projection at the end satisfies the requirement exactly, which a penalty
+# added to the reward along the way does not.
+
+
+def enforce_polar(logits, choices, minimum):
+    cols = torch.tensor(POLAR_COLUMNS, device=logits.device)
+    polar_scores, pick = logits.index_select(-1, cols).max(dim=-1)
+    polar_choices = cols[pick]                      # best polar residue
+    already = torch.isin(choices, cols)             # positions already polar
+    sequences = []
+    for b in range(choices.shape[0]):
+        row = choices[b].clone()
+        missing = minimum - int(already[b].sum())
+        if missing > 0:
+            chosen = logits[b].gather(1, row[:, None]).squeeze(1)
+            cost = chosen - polar_scores[b]         # logit lost by swapping
+            cost = cost.masked_fill(already[b], float("inf"))
+            _, swap = cost.topk(missing, largest=False)
+            row[swap] = polar_choices[b][swap]
+        sequences.append("".join(AMINO_ACIDS[i] for i in row.tolist()))
+    return sequences
+
+
+# %% 08. The latents, standardized
+#
+# ESM-2 coordinates have means and scales of their own, and a corruption that drove such a cloud
+# toward N(0, I) would have to destroy a systematic offset as well as the structure. Each
+# coordinate is therefore centered and scaled over the valid residues before any diffusion runs,
+# and the two statistics are kept for decoding.
+
+
+def load_latents(dataset: str, device, seed: int):
+    """Return the standardized latents, the mask, the decoding head and its residue columns."""
+    if dataset == "toy":
+        from cis6270.data import standardize, toy_embeddings
+
+        z0, mask = toy_embeddings(num_sequences=128, length=16, width=32, seed=seed)
+        z0, z_mean, z_std = standardize(z0, mask)
+        # Without ESM-2 there is no language-model head, so a fixed projection stands in. It makes
+        # the decoding path and the constraint runnable offline; its residues carry no biology.
+        generator = torch.Generator().manual_seed(seed)
+        weight = torch.randn(len(AMINO_ACIDS), z0.shape[-1], generator=generator) / 4.0
+        head = lambda latent: latent @ weight.t().to(latent)  # noqa: E731
+        columns = torch.arange(len(AMINO_ACIDS))
+    else:
+        from cis6270.data import esm2_embeddings, standardize
+
+        z0, mask, tokenizer, head = esm2_embeddings(
+            model_name=MODEL_NAME, max_length=MAX_LEN, device=device,
+        )  # The head is ESM-2's masked language-model head, frozen with the encoder.
+        z0, z_mean, z_std = standardize(z0, mask)
+        columns = torch.tensor([tokenizer.convert_tokens_to_ids(a) for a in AMINO_ACIDS])
+    return z0.to(device), mask.to(device), z_mean.to(device), z_std.to(device), head, columns
+
+
+# %% 09. The report
+def main() -> None:
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.add_argument("--dataset", choices=["toy", "esm2"], default="toy")
+    parser.add_argument("--minimum-polar", type=int, default=8, help="polar residues required")
+    parser.set_defaults(steps=800, batch_size=16, lr=1e-3, samples=8)
+    args = parser.parse_args()
+    set_seed(args.seed)
+    device = pick_device(args.device)
+
+    z0_all, mask_all, z_mean, z_std, head, columns = load_latents(args.dataset, device, args.seed)
+    num_sequences, length, width = z0_all.shape  # [N, L, D].
+    denoiser = LatentDenoiser(embed_dim=width, num_layers=NUM_LAYERS, num_heads=8,
+                              max_len=length).to(device)
+    if not args.quiet:
+        print(f"device={device} dataset={args.dataset} latents={tuple(z0_all.shape)} "
+              f"parameters={count_parameters(denoiser):,}")
+
+    def loss_fn(step: int) -> torch.Tensor:
+        rows = torch.randint(0, num_sequences, (args.batch_size,), device=device)  # [B].
+        z0, mask = z0_all[rows], mask_all[rows]  # [B, L, D] and [B, L].
+        k = torch.randint(1, NUM_STEPS + 1, (z0.shape[0],), device=device)  # One step per item.
+        z_k, _ = q_sample_latent(z0, k, mask)  # Corrupt the clean latent.
+        return latent_loss(denoiser(z_k, k, mask), z0, mask)  # Valid residues only.
+
+    losses = train(denoiser, loss_fn, args.steps, args.lr, quiet=args.quiet)
+    generated = sample_latents(denoiser, args.samples, length, width, device)  # [S, L, D].
+    logits, choices = decode(generated, head, columns.to(device), z_mean, z_std)
+    # The argmax decoding before the constraint, and the same sequences after it.
+    plain = ["".join(AMINO_ACIDS[int(i)] for i in row) for row in choices]
+    constrained = enforce_polar(logits, choices, args.minimum_polar)
+
+    # The chapter's coordinate, corrupted at k = 500 and carried back one reverse step.
+    abar = float(alphas_cumprod[500])
+    a_500, sigma_500 = math.sqrt(abar), math.sqrt(1.0 - abar)  # 0.2803 and 0.9599.
+    z_500 = a_500 * 1.2 + sigma_500 * 0.4  # 0.7204.
+    eps_hat = (z_500 - a_500 * 1.05) / sigma_500  # 0.4438 from the prediction z0_hat = 1.05.
+    mean_500 = (z_500 - float(betas[500]) / sigma_500 * eps_hat) / math.sqrt(float(alphas[500]))
+    # The nine substitution costs of the notes, with two positions still to fill.
+    costs = torch.tensor([0.9, 1.4, 2.6, 3.1, 3.3, 4.0, 4.2, 5.5, 6.1])
+
+    window = min(50, args.steps)  # A short run averages over everything it has.
+    report = {
+        "parameters": count_parameters(denoiser),
+        f"loss, first {window} steps": sum(losses[:window]) / window,
+        f"loss, last {window} steps": sum(losses[-window:]) / window,
+        "valid residues per batch": int(mask_all[:args.batch_size].sum()),
+        "a_500 and sigma_500": [round(a_500, 4), round(sigma_500, 4)],  # 0.2803 and 0.9599.
+        "z_500 from z_0=1.2, eps=0.4": z_500,  # 0.7204.
+        "eps implied by z0_hat=1.05": eps_hat,  # 0.4438, against the draw 0.4.
+        "reverse mean at k=500": mean_500,  # 0.7193.
+        "reverse standard deviation": math.sqrt(float(posterior_variance[500])),  # 0.1002.
+        "clean estimate at a=0.8 from eps=-0.8": (7.4 + 0.6 * 0.8) / 0.8,  # 9.85.
+        "noise recovered from it": (7.4 - 0.8 * 9.85) / 0.6,  # -0.8.
+        "polar fraction, plain decoding": sum(
+            composition_proxies(s)[1] for s in plain) / len(plain),
+        "polar fraction, constrained": sum(
+            composition_proxies(s)[1] for s in constrained) / len(constrained),
+        "net charge per residue, constrained": sum(
+            composition_proxies(s)[0] for s in constrained) / len(constrained),
+        "two cheapest substitutions": [round(float(v), 1)
+                                      for v in costs.topk(2, largest=False)[0]],
+        "logit surrendered": float(costs.topk(2, largest=False)[0].sum()),  # 2.3.
+    }
+    print_report(f"AMP-Diffusion on {args.dataset} latents", report, constrained)
+    save_run(args.out, model=denoiser, config=vars(args), report=report, losses=losses,
+             samples=constrained)
+
+
+# %% 10. Run the script
+if __name__ == "__main__":
+    main()

lecture_3/samplers.py

@@ -0,0 +1,247 @@
+"""Three samplers reading one score.
+
+Once a score is available, generation can take three forms. The reverse-time equation discretized
+on a grid gives ancestral DDPM sampling, which injects fresh noise at every step. Rewriting the
+Laplacian of the Fokker-Planck equation as a divergence gives the probability-flow ODE, which
+reproduces the same marginals with no noise at all. Holding one density fixed and running its
+score as a drift gives Langevin dynamics, which needs a ladder of noise levels to cross between
+modes. The target here is the four-Gaussian mixture, whose corrupted score is a closed form, so
+all three run on the same exact field and nothing is trained.
+
+Chapter 3 runs its diffusion clock from the data, so t = 0 holds a clean sample and corruption
+grows with t. Generation traverses the same interval from t = 1 back to t = 0.
+
+Chapter 3, Sections 3.22, 3.26, 3.32, 3.34 and 3.35 of the notes. Prints the mode coverage of each
+sampler, the number of score evaluations each spends, and the chapter's one backward step from
+x = 1.2 under both dynamics.
+
+    python lecture_3/samplers.py              # 200 steps for each of the three
+    python lecture_3/samplers.py --steps 20   # a coarse grid, for a first look
+"""
+
+# %% 01. Imports and configuration
+import math
+
+import torch
+
+from cis6270.data import GAUSSIAN_CENTERS, GAUSSIAN_STD
+from cis6270.metrics import mean_log_density, mode_report
+from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed
+
+BETA_MIN = 0.1  # beta(0) of the linear variance-preserving schedule.
+BETA_MAX = 20.0  # beta(1) of the same schedule.
+LADDER_START = 0.7  # Noisiest level of the Langevin ladder, in diffusion time.
+LADDER_END = 0.01  # Least noisy level of the ladder.
+
+
+# %% 02. The target and its exact score at any noise level
+#
+# Corrupting a Gaussian mixture by X_t = a_t X_0 + sigma_t eps gives a Gaussian mixture with means
+# a_t mu_m and the common variance a_t^2 s^2 + sigma_t^2, so the score of p_t is a closed form at
+# every level. Every call is counted, because the cost of a sampler is the number of evaluations
+# a trained network would have served.
+
+
+def beta(t: float) -> float:
+    """The noising rate beta(t) = 0.1 + 19.9 t of the linear schedule."""
+    return BETA_MIN + (BETA_MAX - BETA_MIN) * t
+
+
+def signal(t: float) -> float:
+    """a_t = exp(-B(t)/2) with B(t) = 0.1 t + 9.95 t^2, the accumulated noising."""
+    return math.exp(-(BETA_MIN * t + 0.5 * (BETA_MAX - BETA_MIN) * t ** 2) / 2.0)
+
+
+class ExactScore:
+    """The score of p_t at a given signal coefficient, with a counter on the calls."""
+
+    def __init__(self):
+        self.calls = 0  # One call stands for one network evaluation.
+
+    def __call__(self, x: torch.Tensor, a_t: float) -> torch.Tensor:
+        """x: [B, 2], a_t the signal coefficient -> [B, 2]."""
+        self.calls += 1
+        variance = a_t ** 2 * GAUSSIAN_STD ** 2 + (1.0 - a_t ** 2)  # Width of p_t.
+        offset = x[:, None, :] - a_t * GAUSSIAN_CENTERS.to(x)[None, :, :]  # [B, 4, 2].
+        log_weight = -(offset ** 2).sum(dim=2) / (2.0 * variance)  # [B, 4] equal mixture weights.
+        responsibility = log_weight.softmax(dim=1)[:, :, None]  # [B, 4, 1].
+        return (responsibility * (-offset / variance)).sum(dim=1)  # [B, 2].
+
+
+def make_schedule(num_steps: int) -> dict:
+    """The DDPM tensors on a grid of K steps, read off the continuous schedule.
+
+    Setting abar_k = a_{t_k}^2 at t_k = k/K and recovering alpha_k = abar_k / abar_{k-1} matches
+    the continuous marginals at every grid point. At K = 1000 the betas agree with the linear
+    sequence from 1e-4 to 0.02 to the accuracy of log(1 - beta) = -beta, and at a coarse K every
+    beta_k stays below one, which beta(t_k)/K does not.
+    """
+    grid = (torch.arange(0, num_steps + 1).to(torch.get_default_dtype())
+            / num_steps)  # [K+1] the times t_k, with t_0 = 0.
+    alphas_cumprod = torch.exp(
+        -(BETA_MIN * grid + 0.5 * (BETA_MAX - BETA_MIN) * grid ** 2)
+    )  # [K+1] abar_k, with abar_0 = 1 at the clean state.
+    alphas_cumprod_prev = torch.cat([torch.ones(1), alphas_cumprod[:-1]])  # abar_{k-1}.
+    alphas = alphas_cumprod / alphas_cumprod_prev  # One-step retention alpha_k.
+    betas = 1.0 - alphas  # Per-step variance beta_k, zero at the placeholder index 0.
+    return {
+        "betas": betas, "alphas": alphas, "alphas_cumprod": alphas_cumprod,
+        "posterior_variance": betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod),
+    }
+
+
+# %% 03. Ancestral DDPM sampling
+#
+# x_{k-1} = mu_theta(x_k, k) + sqrt(beta-tilde_k) z with the reverse mean built from the noise
+# prediction, eps = -sigma_k s_k. Fresh noise enters at every step except the last.
+
+
+@torch.no_grad()
+def ancestral_sample(score, schedule, num_points: int, generator=None) -> torch.Tensor:
+    """Run the reverse chain from X_K ~ N(0, I) down to X_0. Returns [N, 2]."""
+    num_steps = len(schedule["betas"]) - 1
+    x = torch.randn(num_points, 2, generator=generator)  # [N, 2] the terminal Gaussian.
+    for step in range(num_steps, 0, -1):  # k = K down to k = 1.
+        abar_k = float(schedule["alphas_cumprod"][step])
+        sigma_k = math.sqrt(1.0 - abar_k)  # Corruption standard deviation at this step.
+        eps_theta = -sigma_k * score(x, math.sqrt(abar_k))  # [N, 2] the noise the score implies.
+        mean = (x - float(schedule["betas"][step]) / sigma_k * eps_theta) / math.sqrt(
+            float(schedule["alphas"][step])
+        )  # Reverse mean mu_theta.
+        if step > 1:  # No noise at the final step.
+            variance = float(schedule["posterior_variance"][step])
+            x = mean + math.sqrt(variance) * torch.randn(x.shape, generator=generator)
+        else:
+            x = mean
+    return x  # [N, 2].
+
+
+# %% 04. The probability-flow ODE
+#
+# v_t(x) = -(1/2) beta(t) [x + s_t(x)] has the same marginals with no Brownian increment, and
+# generation integrates it backward, x_{t - dt} = x_t - dt v_t(x_t). Once X_T is drawn the whole
+# trajectory is fixed.
+
+
+@torch.no_grad()
+def probability_flow_sample(score, num_points: int, steps: int, generator=None) -> torch.Tensor:
+    """Euler integration of the probability-flow ODE from t = 1 to t = 0. Returns [N, 2]."""
+    x = torch.randn(num_points, 2, generator=generator)  # [N, 2] the terminal Gaussian.
+    dt = 1.0 / steps  # Positive step, taken in the direction of decreasing t.
+    for step in range(steps, 0, -1):
+        t = step * dt
+        a_t = signal(t)
+        velocity = -0.5 * beta(t) * (x + score(x, a_t))  # [N, 2] forward-time velocity.
+        x = x - dt * velocity  # One backward Euler step.
+    return x  # [N, 2].
+
+
+# %% 05. Langevin dynamics and the noise ladder
+#
+# dX = s(X) d_tau + sqrt(2) dW holds p stationary at one fixed density, and the Euler-Maruyama
+# update is x' = x + eta s(x) + sqrt(2 eta) z. A chain run at the data density alone crosses
+# between modes slowly, so the levels are visited in decreasing order of corruption, with a few
+# iterations at each.
+
+
+@torch.no_grad()
+def annealed_langevin_sample(score, num_points, levels, per_level, scale, generator=None):
+    """Run ``per_level`` Langevin iterations at each level of the ladder. Returns [N, 2]."""
+    x = torch.randn(num_points, 2, generator=generator)  # [N, 2] a Gaussian start.
+    for t in levels:
+        a_t = signal(t)
+        width = a_t ** 2 * GAUSSIAN_STD ** 2 + (1.0 - a_t ** 2)  # Width of p_t at this level.
+        eta = scale * width  # Step size tied to the level, large when the density is broad.
+        for _ in range(per_level):
+            noise = torch.randn(x.shape, generator=generator)  # [N, 2] fresh at every iteration.
+            x = x + eta * score(x, a_t) + math.sqrt(2.0 * eta) * noise
+    return x  # [N, 2] at the least corrupted level of the ladder.
+
+
+def langevin_step(x: float, score_value: float, eta: float, z: float) -> float:
+    """One Langevin update in one coordinate, the chapter's x + eta s(x) + sqrt(2 eta) z."""
+    return x + eta * score_value + math.sqrt(2.0 * eta) * z
+
+
+# %% 06. One backward step under both dynamics
+#
+# The score appears in the reverse stochastic equation with the full g(t)^2 and in the
+# deterministic velocity with half of it, which is the only structural difference between the
+# two samplers.
+
+
+def one_step_comparison(x_t, beta_t, sigma_t, eps_theta, dt) -> dict:
+    """The chapter's state carried one step by the ODE and by the reverse-time equation."""
+    score_value = -eps_theta / sigma_t  # The score the noise prediction implies.
+    velocity = -0.5 * beta_t * (x_t + score_value)  # Probability-flow velocity.
+    reverse_drift = -0.5 * beta_t * x_t - beta_t * score_value  # f - g^2 s.
+    return {
+        "score": score_value, "velocity": velocity, "ode_next": x_t - dt * velocity,
+        "sde_drift": reverse_drift, "sde_mean": x_t - dt * reverse_drift,
+        "sde_noise_std": math.sqrt(beta_t * dt),
+    }
+
+
+# %% 07. The report
+def main() -> None:
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.add_argument("--levels", type=int, default=20, help="rungs of the Langevin ladder")
+    parser.add_argument("--langevin-scale", type=float, default=0.15, help="step size over width")
+    parser.set_defaults(steps=200, samples=2048)
+    args = parser.parse_args()
+    set_seed(args.seed)
+    device = pick_device(args.device)  # Two coordinates per point; a CPU is enough.
+
+    results, counts = {}, {}
+    generator = torch.Generator().manual_seed(args.seed)
+    score = ExactScore()
+    results["ancestral DDPM"] = ancestral_sample(
+        score, make_schedule(args.steps), args.samples, generator,
+    )
+    counts["ancestral DDPM"] = score.calls
+
+    generator = torch.Generator().manual_seed(args.seed)  # The same starting noise for each.
+    score = ExactScore()
+    results["probability flow"] = probability_flow_sample(
+        score, args.samples, args.steps, generator,
+    )
+    counts["probability flow"] = score.calls
+
+    generator = torch.Generator().manual_seed(args.seed)
+    score = ExactScore()
+    per_level = max(args.steps // args.levels, 1)  # Spend the same budget as the other two.
+    ladder = [LADDER_START + (LADDER_END - LADDER_START) * i / (args.levels - 1)
+              for i in range(args.levels)]  # Decreasing corruption.
+    results["annealed Langevin"] = annealed_langevin_sample(
+        score, args.samples, ladder, per_level, args.langevin_scale, generator,
+    )
+    counts["annealed Langevin"] = score.calls
+
+    report = {}
+    for name, points in results.items():
+        measured = mode_report(points)
+        report[f"{name}: score evaluations"] = counts[name]
+        report[f"{name}: modes found"] = measured["modes_found"]
+        report[f"{name}: mode entropy"] = measured["mode_entropy"]  # log 4 = 1.386 when even.
+        report[f"{name}: mode fractions"] = measured["mode_fractions"]
+        report[f"{name}: distance to mode"] = measured["nearest_distance"]  # 0.44 on the target.
+        report[f"{name}: mean log density"] = mean_log_density(points)
+
+    step = one_step_comparison(1.2, 2.0, 0.5, 0.8, 0.1)
+    report["score at the worked state"] = step["score"]  # -1.6.
+    report["probability-flow velocity"] = step["velocity"]  # 0.4.
+    report["ODE next state"] = step["ode_next"]  # 1.16.
+    report["reverse-time drift"] = step["sde_drift"]  # 2.0.
+    report["reverse-time mean"] = step["sde_mean"]  # 1.0.
+    report["reverse-time noise std"] = step["sde_noise_std"]  # 0.4472.
+    report["Langevin step from x=2"] = langevin_step(2.0, -2.0, 0.01, 0.5)  # 2.0507.
+
+    if not args.quiet:
+        print(f"device={device} points={args.samples} steps={args.steps} seed={args.seed}")
+    print_report("Three samplers on one exact score", report)
+    save_run(args.out, config=vars(args), report=report)
+
+
+# %% 08. Run the script
+if __name__ == "__main__":
+    main()

lecture_3/score_matching.py

@@ -0,0 +1,302 @@
+"""Three objectives that fit the same score field.
+
+The loss we would like to minimize, the squared distance from a model to the gradient of the log
+data density, holds a target that a sample does not give us. One integration by parts removes that
+target and leaves the squared norm of the model plus its divergence. One Gaussian corruption
+removes it a second way and leaves a plain regression on the noise we drew. Both forms have the
+minimizer of the first, and all three are fitted here on a one-dimensional Gaussian mixture whose
+exact score is a closed form to measure them against.
+
+Chapter 3 runs its diffusion clock from the data, so t = 0 holds a clean sample and corruption
+grows with t. The level fitted here is the running one of the chapter, a_t = 0.8 and sigma_t = 0.6.
+
+Chapter 3, Sections 3.14 through 3.17 of the notes. Prints the three fitted score errors, the
+integration-by-parts identity, the minimizer of each objective, the Hutchinson trace estimate, and
+the posterior average at x_t = 7.4.
+
+    python lecture_3/score_matching.py              # fit all three objectives
+    python lecture_3/score_matching.py --steps 20   # a few seconds, for a first look
+"""
+
+# %% 01. Imports and configuration
+import math
+
+import torch
+
+from cis6270.nets import VelocityMLP, count_parameters
+from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
+
+A_T = 0.8  # Signal coefficient of the level fitted here.
+SIGMA_T = 0.6  # Noise standard deviation at the same level, with a_t^2 + sigma_t^2 = 1.
+MIXTURE_MEANS = torch.tensor([-2.0, 2.0])  # The two clean modes.
+MIXTURE_WEIGHTS = torch.tensor([0.3, 0.7])  # Their mixture weights.
+MIXTURE_STD = 0.5  # Standard deviation of each clean component.
+TRACE_DIM = 64  # The chapter's 64-dimensional descriptor, used for the trace estimate.
+NAMES = ("explicit", "implicit", "denoising")  # The three objectives, in the order derived.
+
+
+# %% 02. The corrupted mixture and its exact score
+#
+# A Gaussian mixture corrupted by X_t = a_t X_0 + sigma_t eps is a Gaussian mixture again, with
+# means a_t mu_m and variance a_t^2 s^2 + sigma_t^2, so its score is a closed form.
+
+CORRUPTED_MEANS = A_T * MIXTURE_MEANS  # [M] means of p_t.
+CORRUPTED_VARIANCE = A_T ** 2 * MIXTURE_STD ** 2 + SIGMA_T ** 2  # Shared variance of p_t.
+
+
+def sample_clean(num_points: int, generator=None) -> torch.Tensor:
+    """Draw clean samples from the mixture. Returns [N, 1]."""
+    which = torch.multinomial(MIXTURE_WEIGHTS, num_points, True, generator=generator)  # [N].
+    noise = torch.randn(num_points, generator=generator) * MIXTURE_STD  # [N].
+    return (MIXTURE_MEANS[which] + noise)[:, None]  # [N, 1] clean points.
+
+
+def exact_score(x: torch.Tensor) -> torch.Tensor:
+    """The score of the corrupted mixture p_t at each row of x. [N, 1] -> [N, 1]."""
+    offset = x - CORRUPTED_MEANS[None, :]  # [N, M] displacement from each corrupted mean.
+    log_weight = MIXTURE_WEIGHTS.log()[None, :] - offset ** 2 / (2.0 * CORRUPTED_VARIANCE)
+    responsibility = log_weight.softmax(dim=1)  # [N, M] posterior over the components.
+    return (responsibility * (-offset / CORRUPTED_VARIANCE)).sum(dim=1, keepdim=True)  # [N, 1].
+
+
+def model_score(model, x: torch.Tensor) -> torch.Tensor:
+    """The network's score at x. One level is fitted here, so the time input is held at a_t."""
+    return model(x, torch.full((x.shape[0],), A_T))  # [N, 1].
+
+
+def divergence(field, x: torch.Tensor) -> torch.Tensor:
+    """The divergence of a field, which in one coordinate is the ordinary derivative. [N, 1]."""
+    x = x.requires_grad_(True)  # Differentiate with respect to the input, not the weights.
+    d_s, = torch.autograd.grad(field(x).sum(), x, create_graph=True)  # Trace of the Jacobian.
+    return d_s
+
+
+# %% 03. Explicit score matching
+#
+# J(theta) = (1/2) E || s_theta(x) - grad log p(x) ||^2, averaged over samples of p. The target is
+# unavailable in general and available here, so this objective is the reference for the other two.
+
+
+def explicit_loss(model, x: torch.Tensor) -> torch.Tensor:
+    """Half the squared distance between the predicted score and the exact one."""
+    return 0.5 * (model_score(model, x) - exact_score(x)).square().sum(dim=1).mean()
+
+
+# %% 04. The integration by parts that removes the data score
+#
+# E_p[s_theta . grad log p] = -E_p[div s_theta], from p grad log p = grad p followed by one
+# integration by parts per coordinate, the boundary term vanishing for a density with Gaussian
+# tails.
+
+
+def integration_by_parts_gap(field, x: torch.Tensor) -> dict:
+    """The two sides of the identity on one sample, and the gap between them."""
+    left = float((field(x) * exact_score(x)).sum(dim=1).mean().detach())
+    right = float(-divergence(field, x.clone()).sum(dim=1).mean().detach())
+    return {"left": left, "right": right, "gap": abs(left - right)}
+
+
+# %% 05. Implicit score matching with the trace term
+#
+# Substituting the identity leaves J(theta) = E[ (1/2) ||s_theta||^2 + div s_theta ] + C, with C a
+# constant in theta. On the model s_a(x) = -a x with centered data of variance v, the expectation
+# reduces to (1/2) a^2 v - a, whose minimizer is a = 1/v.
+
+
+def implicit_loss(model, x: torch.Tensor) -> torch.Tensor:
+    """The squared norm of the model plus its divergence, with the data score gone."""
+    field = lambda z: model_score(model, z)  # noqa: E731
+    s = field(x)  # [N, 1].
+    return (0.5 * s.square().sum(dim=1) + divergence(field, x.clone()).sum(dim=1)).mean()
+
+
+def one_parameter_loss(a: float, variance: float) -> float:
+    """The implicit loss of s_a(x) = -a x on centered data of the given variance."""
+    return 0.5 * a ** 2 * variance - a  # Minimized at a = 1 / variance.
+
+
+# %% 06. Denoising score matching
+#
+# For a pair drawn as X_t = a_t X_0 + sigma_t eps the conditional score is known exactly,
+# grad log q_t(x | x_0) = -(x - a_t x_0)/sigma_t^2 = -eps/sigma_t. Its posterior average at a fixed
+# state is the marginal score, which is therefore the minimizer of the regression.
+
+
+def denoising_loss(model, x_0: torch.Tensor, epsilon: torch.Tensor) -> torch.Tensor:
+    """Half the squared distance between the predicted score and -eps / sigma_t."""
+    x_t = A_T * x_0 + SIGMA_T * epsilon  # [N, 1] the corrupted state.
+    return 0.5 * (model_score(model, x_t) + epsilon / SIGMA_T).square().sum(dim=1).mean()
+
+
+# %% 07. One minimizer for the three objectives
+#
+# Along the family s_a = a s_exact each objective is a quadratic in a, and its minimizer is that
+# objective's own numerator over the shared denominator E[||s||^2]. The explicit numerator is
+# E[s . grad log p], the implicit one is -E[div s], which the identity of Section 04 makes the
+# same number, and the denoising one is E[s . (-eps/sigma_t)], which is that number again because
+# the posterior average of the conditional score is the marginal score.
+
+
+def minimizers(x_0, epsilon) -> dict:
+    """The minimizing a of each objective along the family a s_exact, on one sample."""
+    x_t = A_T * x_0 + SIGMA_T * epsilon  # [N, 1] samples of p_t and the draws that made them.
+    s = exact_score(x_t)  # [N, 1] the exact field of the family.
+    squared = float(s.square().sum(dim=1).mean())  # E ||s||^2, the shared denominator.
+    data = float((s * exact_score(x_t)).sum(dim=1).mean())  # E[s . grad log p], the explicit one.
+    div_s = float(divergence(exact_score, x_t.clone()).detach().sum(dim=1).mean())  # E div s.
+    target = float((s * (-epsilon / SIGMA_T)).sum(dim=1).mean())  # E[s . (-eps/sigma_t)].
+    return {"explicit": data / squared, "implicit": -div_s / squared,
+            "denoising": target / squared}
+
+
+# %% 08. Random directions for the trace
+#
+# Extracting the d diagonal entries of the Jacobian exactly costs d backward passes, 64 for a
+# learned descriptor and 3072 for a colour image on a 32-by-32 grid. Drawing v with E[v v^T] = I
+# and averaging v^T (grad s) v returns the same number at one pass per direction.
+
+
+def high_dimensional_score(x: torch.Tensor, means: torch.Tensor, variance: float) -> torch.Tensor:
+    """The score of a two-component Gaussian mixture in TRACE_DIM coordinates. [N, d] -> [N, d]."""
+    offset = x[:, None, :] - means[None, :, :]  # [N, 2, d].
+    log_weight = -(offset ** 2).sum(dim=2) / (2.0 * variance)  # [N, 2] equal mixture weights.
+    responsibility = log_weight.softmax(dim=1)[:, :, None]  # [N, 2, 1].
+    return (responsibility * (-offset / variance)).sum(dim=1)  # [N, d].
+
+
+def exact_trace(field, x: torch.Tensor) -> tuple[torch.Tensor, int]:
+    """The divergence coordinate by coordinate, and the number of backward passes it took."""
+    x = x.clone().requires_grad_(True)
+    s = field(x)  # [N, d].
+    total = torch.zeros(x.shape[0])
+    for i in range(x.shape[1]):  # One backward pass per coordinate.
+        column, = torch.autograd.grad(s[:, i].sum(), x, retain_graph=True)
+        total = total + column[:, i]
+    return total.detach(), x.shape[1]
+
+
+def hutchinson_trace(field, x: torch.Tensor, probes: int, generator=None):
+    """Average v^T (grad s) v over ``probes`` random directions, one pass each."""
+    # A score field has a symmetric Jacobian, the Hessian of a log density, so the
+    # vector-Jacobian product taken here equals the Jacobian-vector product of the notes.
+    estimates = []
+    for _ in range(probes):
+        v = torch.randn(x.shape, generator=generator)  # [N, d] direction with E[v v^T] = I.
+        point = x.clone().requires_grad_(True)
+        projected, = torch.autograd.grad((field(point) * v).sum(), point)  # [N, d] one pass.
+        estimates.append((v * projected).sum(dim=1).detach())  # [N] unbiased per direction.
+    stacked = torch.stack(estimates)  # [probes, N].
+    return stacked.mean(dim=0), stacked
+
+
+# %% 09. The posterior average at x_t = 7.4
+#
+# The clean value 10 has prior weight 0.3 and the value 9 has 0.7, and at a_t = 0.8 and
+# sigma_t = 0.6 both have conditional means near the observation 7.4.
+
+
+def posterior_report(x_t: float, candidates, priors, epsilon_theta: float) -> dict:
+    """Posterior weights, the clean estimate and the three predictions at one observation."""
+    candidates, priors = torch.tensor(candidates), torch.tensor(priors)  # [C] and [C].
+    conditional_mean = A_T * candidates  # [C] corruption mean of each candidate.
+    weight = priors * torch.exp(-(x_t - conditional_mean) ** 2 / (2.0 * SIGMA_T ** 2))  # [C].
+    posterior = weight / weight.sum()  # [C] weights after observing x_t.
+    conditional_score = -(x_t - conditional_mean) / SIGMA_T ** 2  # [C] one per candidate.
+    marginal = float((posterior * conditional_score).sum())  # The averaged score.
+    clean_mean = float((posterior * candidates).sum())  # E[X_0 | X_t = x_t].
+    score_hat = -epsilon_theta / SIGMA_T  # The score the chapter's prediction implies.
+    return {
+        "weights": [round(float(w), 4) for w in posterior],
+        "clean_mean": clean_mean,
+        "conditional_scores": [round(float(s), 4) for s in conditional_score],
+        "marginal": [marginal, (A_T * clean_mean - x_t) / SIGMA_T ** 2],  # Two routes, one number.
+        # The posterior variance of the target, which no function of x_t alone can remove.
+        "residual": float((posterior * (marginal - conditional_score) ** 2).sum()),
+        "score_hat": score_hat,
+        # The clean estimate from the noise, from the score, and from an optimal network.
+        "x0": [(x_t - SIGMA_T * epsilon_theta) / A_T, (x_t + SIGMA_T ** 2 * score_hat) / A_T,
+               (x_t + SIGMA_T ** 2 * marginal) / A_T],
+    }
+
+
+# %% 10. The report
+def main() -> None:
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.add_argument("--probes", type=int, default=32, help="random directions for the trace")
+    parser.set_defaults(steps=250, batch_size=256, width=64, lr=1e-2)
+    args = parser.parse_args()
+    set_seed(args.seed)
+    device = pick_device(args.device)  # The fields here are one-dimensional; a CPU is enough.
+    generator = torch.Generator().manual_seed(args.seed + 1)  # Draws for the measurements below.
+
+    objectives = {}
+    for name in NAMES:
+        set_seed(args.seed)  # Identical initialization for the three fits.
+        model = VelocityMLP(dim=1, width=args.width, depth=2)
+        # The training batches restart with the model, so the only difference between the three
+        # runs is the objective itself.
+        batches = torch.Generator().manual_seed(args.seed + 3)
+
+        def loss_fn(step: int, name=name, model=model, generator=batches) -> torch.Tensor:
+            x_0 = sample_clean(args.batch_size, generator)  # [B, 1] clean batch.
+            epsilon = torch.randn(x_0.shape, generator=generator)  # [B, 1] corruption draw.
+            if name == "denoising":
+                return denoising_loss(model, x_0, epsilon)
+            x_t = A_T * x_0 + SIGMA_T * epsilon  # [B, 1] samples of p_t for the other two.
+            return explicit_loss(model, x_t) if name == "explicit" else implicit_loss(model, x_t)
+
+        if not args.quiet:
+            print(f"{name} score matching")
+        objectives[name] = (model, train(model, loss_fn, args.steps, args.lr, quiet=args.quiet))
+
+    # Every fitted field is measured against the same exact score on one held-out batch.
+    clean = sample_clean(4096, generator)  # [4096, 1].
+    noise = torch.randn(clean.shape, generator=generator)  # [4096, 1].
+    evaluation = A_T * clean + SIGMA_T * noise  # [4096, 1] samples of p_t.
+    target = exact_score(evaluation)
+    errors = {n: float((model_score(m, evaluation).detach() - target).square().mean().sqrt())
+              for n, (m, _) in objectives.items()}
+
+    identity = integration_by_parts_gap(lambda z: model_score(objectives["explicit"][0], z),
+                                        evaluation)
+    wide = torch.Generator().manual_seed(args.seed + 2)  # The family scan has its own draws.
+    scan_clean = sample_clean(32768, wide)  # [32768, 1].
+    scan = minimizers(scan_clean, torch.randn(scan_clean.shape, generator=wide))
+
+    means = torch.randn(2, TRACE_DIM, generator=generator) * 1.5  # [2, d] two distant modes.
+    points = torch.randn(8, TRACE_DIM, generator=generator)  # [8, d] evaluation points.
+    field = lambda z: high_dimensional_score(z, means, 1.0)  # noqa: E731
+    exact, passes = exact_trace(field, points)
+    estimate, draws = hutchinson_trace(field, points, args.probes, generator)
+
+    posterior = posterior_report(7.4, (10.0, 9.0), (0.3, 0.7), epsilon_theta=-0.8)
+    report = {
+        "parameters per model": count_parameters(objectives["explicit"][0]),
+        "score error, the three fits": [round(errors[n], 4) for n in NAMES],
+        "E[s . grad log p] and -E[div s]": [round(identity[k], 4) for k in ("left", "right")],
+        "integration by parts gap": identity["gap"],
+        "minimizing a, the three forms": [round(scan[n], 4) for n in NAMES],  # All near 1.
+        # The one-parameter fit at a = 0.5 and a = 1 on unit variance, then at a = 1/v with v = 4.
+        "one-parameter losses": [round(one_parameter_loss(a, v), 4)
+                                 for a, v in ((0.5, 1.0), (1.0, 1.0), (0.25, 4.0))],
+        "exact trace, backward passes": passes,  # 64, one per coordinate.
+        "exact trace, mean": float(exact.mean()),
+        "Hutchinson trace, mean": float(estimate.mean()),
+        "Hutchinson standard error": float(draws.std(dim=0).mean() / math.sqrt(args.probes)),
+        "posterior weights at x_t=7.4": posterior["weights"],  # 0.2156 and 0.7844.
+        "posterior mean of x_0": posterior["clean_mean"],  # 9.2156.
+        "conditional scores": posterior["conditional_scores"],  # 1.6667 and -0.5556.
+        "marginal score, two routes": [round(v, 4) for v in posterior["marginal"]],  # -0.0765.
+        "irreducible residual": posterior["residual"],  # 0.835.
+        "score from eps_theta=-0.8": posterior["score_hat"],  # 1.3333.
+        "x_0 estimates": [round(v, 4) for v in posterior["x0"]],  # 9.85, 9.85 and 9.22.
+    }
+    if not args.quiet:
+        print(f"\ndevice={device} level a_t={A_T} sigma_t={SIGMA_T} seed={args.seed}")
+    print_report("Three objectives, one score field", report)
+    save_run(args.out, config=vars(args), report=report, losses=objectives["denoising"][1])
+
+
+# %% 11. Run the script
+if __name__ == "__main__":
+    main()

lecture_3/sde.py

@@ -0,0 +1,249 @@
+"""The forward corruption as a stochastic differential equation.
+
+A deterministic flow is fixed by its velocity at every point. A random path needs one more
+ingredient, the perturbation applied over each small interval. A Brownian increment of an interval
+of length h has size sqrt(h), one Euler-Maruyama step adds that increment to the Euler step of
+Chapter 2, a population started at one point spreads at a rate the noise scale sets, with a
+restoring drift that spread stays bounded, and under Ito's lemma and the Fokker-Planck equation
+a single drift and noise scale become an equation for the density. No network appears here.
+
+Chapter 3 runs its diffusion clock from the data, so t = 0 holds a clean sample and corruption
+grows with t. Every time in this file is read that way.
+
+Chapter 3, Sections 3.2 through 3.10 of the notes. Prints the chapter's worked step from x = 2, the
+Ornstein-Uhlenbeck mean and variance at t = log 2, the Ito check on E[W_t^2], and the
+Fokker-Planck estimate of the Gaussian peak at t = 0.1.
+
+    python lecture_3/sde.py                 # 400 Euler-Maruyama steps, 16384 paths
+    python lecture_3/sde.py --steps 20      # a coarse grid, for a first look
+"""
+
+# %% 01. Imports and configuration
+import math
+
+import torch
+
+from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed
+
+# `--steps` is the number of Euler-Maruyama steps on the grid and `--samples` the number of paths;
+# nothing in this file is trained, so those two flags carry the whole cost of a run.
+TERMINAL_TIME = 1.0  # T, the end of the simulated interval.
+
+
+# %% 02. Drift and Brownian increments
+#
+# dX_t = f(X_t, t) dt + g(t) dW_t. The drift f fixes the average displacement over a short interval
+# and the Brownian increment dW_t spreads the possible next states around that average. Setting
+# g = 0 recovers the ordinary differential equation of Chapter 2, with f as the velocity field.
+
+
+def brownian_increment(h: float, shape: tuple[int, ...], generator=None) -> torch.Tensor:
+    """One Brownian increment over an interval of length h, drawn as sqrt(h) * epsilon."""
+    epsilon = torch.randn(shape, generator=generator)  # [*shape] standard Gaussian draw.
+    return math.sqrt(h) * epsilon  # Covariance h I, so each coordinate has variance h.
+
+
+# %% 03. The size of a Brownian increment
+#
+# Multiplying a standard Gaussian by sqrt(h) multiplies its covariance by h, so a coordinate of the
+# increment has standard deviation sqrt(h). Halving the interval divides the typical displacement
+# by sqrt(2) = 1.414, and the ratio sqrt(h)/h = 1/sqrt(h) grows without bound as the grid refines,
+# which is why dW_t appears in the equation and never a velocity dW_t/dt.
+
+
+def increment_scale(h: float, f_value: float = 1.0, g_value: float = 1.0) -> dict:
+    """The standard deviation of one coordinate of the increment, and two ratios built from it."""
+    std = math.sqrt(h)  # Typical displacement of the noise over an interval of length h.
+    return {
+        "std": std,
+        "std_over_h": std / h,  # Grows without bound as the grid refines.
+        "noise_over_drift": g_value * std / (abs(f_value) * h),  # g / (|f| sqrt(h)).
+    }
+
+
+# %% 04. The Euler-Maruyama step
+#
+# x_{k+1} = x_k + f(x_k, t_k) h + g(t_k) sqrt(h) epsilon_k. The drift term is the Euler step of
+# Chapter 2 unchanged and the noise term is the only addition, with the exact law of a Brownian
+# increment over an interval of length h.
+
+
+def euler_maruyama_step(x_k, f_value, g_value, h: float, epsilon):
+    """One step of the scheme, with the Gaussian draw epsilon_k supplied by the caller."""
+    drift_part = f_value * h  # Displacement the drift alone would produce.
+    noise_part = g_value * math.sqrt(h) * epsilon  # Displacement from this draw.
+    return x_k + drift_part + noise_part, drift_part, noise_part
+
+
+# %% 05. Fresh randomness and the spreading cloud
+#
+# Each step draws a new epsilon_k, so two runs from the same state give two trajectories. Started
+# from a common point x, the next state is Gaussian with mean x + f h and covariance g^2 h I, so
+# the drift moves the center and the noise sets the width. Without a drift the variances of
+# independent increments add, so after N steps the variance is N g^2 h = g^2 T and is unbounded.
+
+
+def driftless_cloud(num_paths: int, steps: int, g_value: float, terminal: float, generator=None):
+    """Simulate f = 0 on a grid of ``steps`` intervals and return the terminal states. [N]."""
+    h = terminal / steps  # Step size, a time.
+    x = torch.zeros(num_paths)  # [N] every path starts at the origin.
+    for _ in range(steps):  # Independent increments accumulate additively in variance.
+        x = x + g_value * brownian_increment(h, (num_paths,), generator)
+    return x  # [N] terminal states, with variance g^2 T in the limit of many paths.
+
+
+# %% 06. A restoring drift
+#
+# dX_t = -kappa X_t dt + g dW_t. Multiplying by the integrating factor e^{kappa t} clears the
+# drift, and under the Ito isometry the variance of the remaining stochastic integral is an
+# ordinary integral, so X_t given x_0 is Gaussian with mean e^{-kappa t} x_0 and variance
+# g^2 (1 - e^{-2 kappa t})/(2 kappa). The stationary variance g^2/(2 kappa) is the level at
+# which the inward pull of the drift offsets the outward spread of the noise.
+
+
+def ornstein_uhlenbeck_law(x_0: float, kappa: float, g_value: float, t: float) -> dict:
+    """The exact conditional mean and variance of the Ornstein-Uhlenbeck process at time t."""
+    decay = math.exp(-kappa * t)  # Fraction of the initial signal that remains.
+    stationary = g_value ** 2 / (2.0 * kappa)  # Variance the process settles on.
+    variance = stationary * (1.0 - decay ** 2)  # Rises from zero toward the stationary value.
+    return {"mean": decay * x_0, "variance": variance, "std": math.sqrt(variance),
+            "stationary_variance": stationary}
+
+
+def simulate_ornstein_uhlenbeck(num_paths, steps, x_0, kappa, g_value, t, generator=None):
+    """Euler-Maruyama on the same process, for comparison with the closed form. Returns [N]."""
+    h = t / steps
+    x = torch.full((num_paths,), float(x_0))  # [N] all paths start at x_0.
+    for _ in range(steps):
+        x = x - kappa * x * h + g_value * brownian_increment(h, (num_paths,), generator)
+    return x  # [N] states at time t.
+
+
+# %% 07. Ito's lemma
+#
+# d phi(X_t) = (phi_t + f phi_x + (1/2) g^2 phi_xx) dt + g phi_x dW_t. The squared increment
+# (Delta X)^2 tends to g^2 Delta t and does not vanish, so the second-order Taylor term is as
+# large as the first-order drift term and belongs in the dt part. The chapter's test function is
+# phi(x) = x^2 on X_t = W_t, where phi_x = 2x and phi_xx = 2, giving d(W_t^2) = dt + 2 W_t dW_t.
+
+
+def ito_drift(f_value, g_value, phi_x, phi_xx):
+    """The dt bracket of Ito's lemma for a probe with no explicit time dependence."""
+    return f_value * phi_x + 0.5 * g_value ** 2 * phi_xx  # Drift term plus curvature correction.
+
+
+def check_ito_on_square(num_paths: int, steps: int, terminal: float, generator=None) -> dict:
+    """Integrate d(W_t^2) = dt + 2 W_t dW_t on a grid and compare with the exact E[W_t^2] = t."""
+    h = terminal / steps
+    w = torch.zeros(num_paths)  # [N] Brownian position.
+    accumulated = torch.zeros(num_paths)  # [N] right-hand side of Ito's lemma, integrated.
+    without_curvature = torch.zeros(num_paths)  # [N] the same sum with the dt term dropped.
+    for _ in range(steps):
+        dw = brownian_increment(h, (num_paths,), generator)  # [N] one increment.
+        accumulated = accumulated + ito_drift(0.0, 1.0, 2.0 * w, 2.0) * h + 2.0 * w * dw
+        without_curvature = without_curvature + 2.0 * w * dw  # Ordinary chain rule.
+        w = w + dw  # Advance the path itself.
+    return {
+        "realized_mean_square": float((w ** 2).mean()),  # Direct average of W_T^2.
+        "ito_mean": float(accumulated.mean()),  # Same quantity through Ito's lemma.
+        "no_curvature_mean": float(without_curvature.mean()),  # Near zero, the wrong answer.
+        "exact_mean_square": terminal,  # E[W_T^2] = T.
+    }
+
+
+# %% 08. The Fokker-Planck equation
+#
+# d_t p_t = -div(f p_t) + (1/2) g^2 laplacian(p_t). With f = 0 and g = 1 started from N(0, 1), the
+# solution stays Gaussian with variance 1 + t, which gives both sides in closed form. Autograd
+# supplies the two derivatives, so the identity is checked and not assumed.
+
+
+def spreading_gaussian_log_density(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+    """log N(x; 0, 1 + t), the exact solution of d_t p = (1/2) d_x^2 p from N(0, 1)."""
+    variance = 1.0 + t  # Variance grows at the rate the diffusion injects it.
+    return -0.5 * torch.log(2.0 * math.pi * variance) - x ** 2 / (2.0 * variance)
+
+
+def fokker_planck_residual(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+    """d_t p - (1/2) d_x^2 p on the closed-form solution, which is zero at every (x, t)."""
+    x = x.clone().requires_grad_(True)  # [M] evaluation points.
+    t = t.clone().requires_grad_(True)  # [M] evaluation times.
+    p = spreading_gaussian_log_density(x, t).exp()  # [M] the density itself.
+    ones = torch.ones_like(p)
+    dp_dt, dp_dx = torch.autograd.grad(p, (t, x), ones, create_graph=True)  # First derivatives.
+    d2p_dx2, = torch.autograd.grad(dp_dx, x, torch.ones_like(dp_dx))  # Second derivative in x.
+    return (dp_dt - 0.5 * d2p_dx2).detach()  # The equation with f = 0 and g = 1.
+
+
+def peak_extrapolation(step: float) -> dict:
+    """The chapter's first-order estimate of the peak density after a short interval."""
+    p_0 = 1.0 / math.sqrt(2.0 * math.pi)  # p_0(0), the height of the standard Gaussian.
+    curvature = -p_0  # p_0''(x) = (x^2 - 1) p_0(x), so p_0''(0) = -p_0(0).
+    rate = 0.5 * 1.0 * curvature  # d_t p_t(0) at t = 0, with g = 1.
+    linear = p_0 + step * rate  # One-step extrapolation from the rate.
+    exact = 1.0 / math.sqrt(2.0 * math.pi * (1.0 + step))  # Peak of N(0, 1 + step).
+    return {"p_0": p_0, "rate": rate, "linear": linear, "exact": exact, "gap": exact - linear}
+
+
+# %% 09. The report
+def main() -> None:
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.set_defaults(steps=400, samples=16384)
+    args = parser.parse_args()
+    set_seed(args.seed)
+    device = pick_device(args.device)  # Printed for consistency; this file runs on the CPU.
+    generator = torch.Generator().manual_seed(args.seed)  # Local generator, so the report repeats.
+
+    # The chapter's step: x_k = 2, f = -1, g = 0.5, h = 0.04 and the draw epsilon_k = 1.5.
+    h = 0.04
+    landed, drift_part, noise_part = euler_maruyama_step(2.0, -1.0, 0.5, h, 1.5)
+    scale = increment_scale(h, f_value=-1.0, g_value=0.5)
+    three_draws = [euler_maruyama_step(2.0, -1.0, 0.5, h, e)[0] for e in (1.5, -0.5, 0.2)]
+
+    # The cloud without a drift: variance g^2 T after the whole interval.
+    terminal_states = driftless_cloud(args.samples, args.steps, 0.5, TERMINAL_TIME, generator)
+
+    # The Ornstein-Uhlenbeck example: x_0 = 2, kappa = 1, g = sqrt(2), evaluated at t = log 2.
+    ou_time = math.log(2.0)
+    ou = ornstein_uhlenbeck_law(2.0, 1.0, math.sqrt(2.0), ou_time)
+    ou_paths = simulate_ornstein_uhlenbeck(
+        args.samples, args.steps, 2.0, 1.0, math.sqrt(2.0), ou_time, generator,
+    )
+
+    ito = check_ito_on_square(args.samples, args.steps, 0.25, generator)  # t = 0.25.
+    grid = torch.linspace(-3.0, 3.0, 61)  # [61] evaluation points for the density equation.
+    residual = fokker_planck_residual(grid, torch.full_like(grid, 0.5))  # At t = 0.5.
+    peak = peak_extrapolation(0.1)
+
+    report = {
+        "euler_maruyama x_{k+1}": landed,  # 2.11 in the notes.
+        "drift displacement f h": drift_part,  # -0.04.
+        "noise displacement g sqrt(h) eps": noise_part,  # 0.15.
+        "typical noise over drift": scale["noise_over_drift"],  # 2.5 at this step size.
+        "increment std at h=0.04": scale["std"],  # 0.2.
+        "three endpoints, mean": sum(three_draws) / 3.0,  # 2.00 against the drift's 1.96.
+        "driftless terminal variance": float(terminal_states.var(unbiased=False)),  # g^2 T = 0.25.
+        "OU mean at t=log 2": ou["mean"],  # 1, half the initial signal.
+        "OU variance at t=log 2": ou["variance"],  # 0.75 of the stationary 1.
+        "OU simulated mean": float(ou_paths.mean()),
+        "OU simulated variance": float(ou_paths.var(unbiased=False)),
+        "E[W_t^2] at t=0.25, exact": ito["exact_mean_square"],  # 0.25.
+        "E[W_t^2] simulated directly": ito["realized_mean_square"],
+        "E[W_t^2] through Ito's lemma": ito["ito_mean"],
+        "E[W_t^2] without the curvature": ito["no_curvature_mean"],  # Near zero.
+        "Fokker-Planck residual, largest": float(residual.abs().max()),  # Zero up to float error.
+        "peak p_0(0)": peak["p_0"],  # 0.3989.
+        "peak rate at t=0": peak["rate"],  # -0.1995.
+        "peak linear estimate at t=0.1": peak["linear"],  # 0.3790.
+        "peak exact at t=0.1": peak["exact"],  # 0.3804.
+        "second-order gap": peak["gap"],  # 0.0014.
+    }
+    if not args.quiet:
+        print(f"device={device} steps={args.steps} paths={args.samples} seed={args.seed}")
+    print_report("Forward corruption without a network", report)
+    save_run(args.out, config=vars(args), report=report)
+
+
+# %% 10. Run the script
+if __name__ == "__main__":
+    main()

lecture_4/README.md

@@ -0,0 +1,108 @@
+# Lecture 4: Discrete Diffusion
+
+Chapter 4 keeps every state inside a finite vocabulary and replaces the continuous displacement of
+Chapter 3 by a decision about which tokens to erase, reveal or replace. The chapter builds masked
+diffusion by hand from a survival schedule, then rebuilds it as one member of a family: a rate
+matrix says what a token can do in the next instant, the master equation follows the marginals, the
+matrix exponential gives the corruption kernel at any time, and time reversal turns the forward
+rates into the rates that generate. Seven published models come out of that family by changing the
+rate matrix, the time variable, what the network predicts, or the state space itself, and a
+further group keeps the rates and changes the order in which positions are filled, the weighting
+of the loss, or the distribution to sample from. Every file here runs on the shared synthetic DNA
+of `cis6270.data`, with `t = 0` clean and `t = 1` all masks, survival probability `alpha_t`, and
+rates `R_t(x, y)` in the row convention `p_dot = p R`.
+
+## The files
+
+| File | Command | What it trains and generates |
+| --- | --- | --- |
+| `ctmc.py` | `python lecture_4/ctmc.py` | Nothing is trained. Works every number of the CTMC sections: the two rules, one row for DNA, waiting times, the master equation, the matrix exponential, the `R_pi` family, eigenvalues, detailed balance and time reversal. |
+| `mdlm.py` | `python lecture_4/mdlm.py` | Trains the denoiser on the weighted cross-entropy at masked positions, generates DNA from all masks. The chapter's thirty-line model. |
+| `uniform.py` | `python lecture_4/uniform.py` | Trains against the rate loss under uniform replacement, generates by tau-leaping from uniformly random tokens. |
+| `d3pm.py` | `python lecture_4/d3pm.py` | Trains against the discrete-time variational bound plus the auxiliary cross-entropy, generates by `T` ancestral steps. |
+| `sedd.py` | `python lecture_4/sedd.py` | Trains a network to output the concrete score through denoising score entropy, generates from the predicted ratios. |
+| `block.py` | `python lecture_4/block.py` | Trains masked diffusion inside blocks with a clean prefix, generates block by block. `--block` moves the dial between MDLM and autoregression. |
+| `simplex.py` | `python lecture_4/simplex.py` | Trains a denoiser that reads a Dirichlet belief at every position, generates by thinning the belief, drawing an innovation and mixing the two. |
+| `guidance.py` | `python lecture_4/guidance.py` | Trains a label-dropout denoiser and a classifier on noisy sequences, then samples unguided, with exact classifier guidance, with its gradient approximation, and classifier free. |
+| `planner.py` | `python lecture_4/planner.py` | Trains with the planner weights inside the bound, samples the same weights in a uniform order and in the planner's order with a remasking corrector. |
+| `search.py` | `python lecture_4/search.py` | Trains masked diffusion, then searches the tree of partial sequences under the h-transform and keeps a Pareto archive over two objectives. |
+
+Two constructions of the chapter have no file here, and the chapter develops both in full. Per-token
+schedules, MD4 and GenMD4, give every clean token its own rate `beta_a(t)` and change nothing but
+the weight, which becomes `w_a(t) = -alpha_dot_t(a)/(1 - alpha_t(a))` inside the sum over tokens.
+A2D2
+lets the length change by adding an insertion jump beside the unmasking jump, which enlarges the
+state space and keeps the master equation, the reversal and the rate loss exactly as they are.
+
+Every file takes the shared flags of `cis6270.runner.common_parser`, so `--steps`, `--samples`,
+`--sample-steps`, `--seed` and `--quiet` mean the same thing everywhere. `ctmc.py` trains nothing,
+so its `--steps` is the number of Euler steps in the matrix-exponential comparison and its
+`--samples` the number of simulated paths. `d3pm.py` takes
+`T` from `--timesteps`, since the same number fixes the corruption and the ancestral sampler.
+`search.py` reads `--sample-steps` as the length of one rollout, which is short because every leaf
+of the tree is scored by its own rollouts.
+
+## What a default run produces
+
+All numbers below come from `--seed 0` on a CPU. The worked numbers are exact and repeat on any
+machine; the trained numbers move a little with the hardware's floating point, and the sample
+statistics move with the seed. The training data have GC fraction 0.577 and motif fraction 0.777,
+and a sampler drawing its bases uniformly would land near 0.05 on the motif, since ACGT fits in a
+length-16 sequence at thirteen places and each one has probability `1/256`.
+
+| File | Worked numbers it prints | Trained numbers |
+| --- | --- | --- |
+| `ctmc.py` | exit rate 4 out of A, `P(A, A) = 0.6` at `h = 0.1`, Euler step `(0.16, 0.22, 0.28, 0.34)`, `exp(tau R_pi)` at `0.625` and `0.125`, Euler products `0.480`, `0.570`, `0.616`, reverse rates `2` and `0.5` out of C | none |
+| `mdlm.py` | `r = 0.5`, reverse law `(0.3, 0.2, 0, 0, 0.5)`, `w(0.5) = 2`, loss term `2.0996`, rate error `1.0217` | bound 1.25 nats per token, GC 0.56, motif 0.69 |
+| `uniform.py` | kernel `0.625` and `0.125`, clean factor 5, posterior `(0.375, 0.625)`, conditional rates 5 and 0.2, posterior average **2**, model rates `2.76` and `1.16`, rate error `0.754` | rate loss 19.1 nats per sequence, GC 0.56, motif 0.50 |
+| `d3pm.py` | `alpha_bar_2 = 0.56`, `Q_bar_2` at `0.67` and `0.11`, the step-1 posterior `(0.5795, 0.3523, 0.0341, 0.0341)`, Gaussian row `0.4088` and `0.0912`, `exp(0.3 R)` first row `(0.6747, 0.1964, 0.1122, 0.0167)` | bound 18.2 nats per sequence, GC 0.57, motif 0.72 |
+| `sedd.py` | target ratio 2, score entropy `0.0107`, `0.3863`, `0.1891` | score entropy 18.8 nats per sequence, mean predicted ratio 1.26 toward a conditional target of 5, GC 0.55, motif 0.33 |
+| `block.py` | worked bound `2.7567` nats on ACGTAC, `integral of w(t)(1 - alpha_t) = 1` | bound 1.21 nats per token at `L' = 4` and 1.29 at `L' = 1`, GC 0.57, motif 0.67 |
+| `simplex.py` | `c_t = 8` and `c_s = 20`, parameters `(5, 1, 1, 1)` and `(17, 1, 1, 1)`, standard deviation `0.161` on A, increment `(12, 0, 0, 0)`, `W ~ Beta(8, 12)` with mean `0.4`, simulated standard deviation `0.0779` against the exact `0.0779` | cross-entropy 0.42 nats per token, GC 0.57, motif 0.34 |
+| `guidance.py` | `h = (2, 1, 0.5, 0.5)`, guided rates `(1, 1.5, 0.5)`, at `gamma = 2` rates `(0.5, 4.5, 0.5)` with exit rate `5.5` and share `0.8182` on G | GC 0.58 unguided, 0.68 with the exact classifier, 0.63 with the gradient approximation, 0.74 classifier free |
+| `planner.py` | reveals positions 3 and 6, remasks position 5, both corrector fluxes 1, PAPL weights `(0.4, 2.0, 0.6, 1.0)`, inner sums `2.9` and `2.12` | motif 0.61 in a uniform order against 0.86 in the planner's order, GC 0.58 and 0.51 |
+| `search.py` | `h(z) = 0.35`, `h(z') = 0.75`, factor `2.143`, a rate of `1.2` becomes `2.571`, the archive keeps P1 and P2 | 144 rollouts scored, an archive of 3 sequences, best GC 0.88 and best motif match 1.0 |
+
+A default run of `ctmc.py` takes about two seconds, and every other file takes between thirty and
+forty seconds of processor time on one core. `--steps 20` runs each file in a few seconds.
+
+`mdlm.py --steps 3000` lowers the bound from 1.25 to 1.22 nats per token in about three times the
+time, and leaves the motif fraction inside the spread it already has across seeds, which at seeds
+0, 1 and 2 is 0.69, 0.45 and 0.72. The one-position bound of `block.py --block 1` is the exact
+autoregressive log-likelihood with no slack, and at this budget it still measures above the
+`L' = 4` bound, because the estimator draws one block of sixteen and multiplies by sixteen, which
+is sixteen times noisier at the same number of draws, and because a denoiser restricted to a left
+prefix has a harder prediction problem than one that reads both sides.
+`sedd.py` sits below the other samplers on the motif because its sampler stops at the time floor,
+as the list below says. `simplex.py --sample-steps 256` raises its motif fraction from 0.34 to
+0.38, so the reverse grid is not what holds that number down.
+
+## What is simplified
+
+- The time weight `w(t) = 1/t` is unbounded at `t = 0`, so every file draws its times above a small
+  floor. The chapter's integrand is correct as written; the floor keeps one draw from dominating a
+  batch.
+- The samplers that use tau-leaping rescale a step whose total jump probability would exceed one.
+  `uniform.py` then finishes with the exact posterior step down to `t = 0` that the chapter derives
+  for the `pi` family. `sedd.py` predicts ratios, with no posterior over clean tokens to step
+  through, so it stops at the time floor, where `alpha_t = 0.92` leaves about eight percent of the
+  positions redrawn from `pi`.
+- `uniform.py` and `sedd.py` use the constant schedule `beta(t) = 4`, which is the schedule of the
+  chapter's worked example and reaches `alpha_t = 1/2` at `t = (log 2)/4`. Its survival probability
+  at `t = 1` is `0.018` and not zero, and that residue is the `KL(q_1 || p_1)` term of the bound.
+- `d3pm.py` samples one step index per sequence and multiplies by `T`, which estimates the sum over
+  steps without evaluating all of it.
+- `block.py` samples one block per sequence and multiplies by the number of blocks, for the same
+  reason. Training a block of length `L'` hides the suffix behind the mask in place of a
+  block-causal attention mask.
+- `simplex.py` holds the concentration constant at `gamma = 1`, realizes the sum over clean tokens
+  by one categorical draw as the chapter does, and reads the final token as the corner the belief
+  has reached. Its loss is the chapter's cross-entropy, which bounds no likelihood, so its nats
+  per token do not compare with the bounds above.
+- `guidance.py` uses a mean-pooling classifier on one-hot inputs, because the gradient
+  approximation needs a derivative with respect to the sequence and the label of this data is a
+  thresholded mean over positions.
+- `planner.py` reveals the number of positions the schedule calls for at each step, which is the
+  deterministic version of the chapter's independent reveal draws, and it remasks a fixed count.
+- `search.py` scores `--sample-steps` rollout steps per completion and expands a fixed number of
+  children, so the tree stays small enough to build in a few seconds.

lecture_4/block.py

@@ -0,0 +1,180 @@
+"""Block diffusion on DNA, with the block size as the dial between diffusion and autoregression.
+
+The chain rule over blocks is exact, and each factor is a distribution over one block of tokens
+conditioned on a clean prefix, which is the kind of distribution masked diffusion models. At one
+block the bound is the MDLM loss, at blocks of a single position it is the autoregressive negative
+log-likelihood with no slack, and every size in between interpolates.
+
+Chapter 4, "Blocks between diffusion and autoregression" and "The block size as a dial between two
+models". Time runs from t = 0 clean to t = 1 all masks, the survival probability is
+alpha_t = 1 - t, and the rates R_t(x, y) act on probability rows, p_dot = p R. Trained on the
+synthetic DNA of cis6270.data. Prints the worked bound of 2.7567 nats, the one-position identity,
+and the sample statistics.
+
+    python lecture_4/block.py                 # train and generate with blocks of four
+    python lecture_4/block.py --block 1       # the autoregressive end of the dial
+    python lecture_4/block.py --steps 20      # a few seconds, for a first look
+"""
+
+# %% 01. Imports and configuration
+import math  # The worked bound is a sum of logarithms.
+
+import torch  # Tensors, gradients and the categorical draws.
+import torch.nn.functional as F  # Cross-entropy at every position.
+
+from cis6270.data import DNA_SIZE, decode_dna, dna_splits  # ACGT, K = 4 and the shared sequences.
+from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
+from cis6270.nets import SequenceDenoiser, count_parameters  # The denoiser d_theta(. | z, t).
+from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
+
+MASK = DNA_SIZE  # A, C, G, T are 0 to 3, and token 4 is the mask m.
+VOCAB = DNA_SIZE + 1  # The masked alphabet V_m has K + 1 = 5 symbols.
+TIME_FLOOR = 1e-3  # The reveal rate w(t) = 1/t is unbounded at t = 0.
+
+
+# %% 02. One block, with a clean prefix and a masked suffix
+#
+# The rate matrix acts only on the current block. Everything before it is clean and conditions the
+# prediction, and everything after it is hidden, so the denoiser never reads the future.
+def block_input(x0: torch.Tensor, block: torch.Tensor, size: int, t: torch.Tensor) -> torch.Tensor:
+    """The network input for one block: clean prefix, corrupted block, masked suffix. [B, L]."""
+    position = torch.arange(x0.shape[1], device=x0.device)[None, :]  # [1, L] the column index.
+    start = (block * size)[:, None]  # [B, 1] where the current block begins.
+    inside = (position >= start) & (position < start + size)  # [B, L] the current block.
+    keep = torch.rand(x0.shape, device=x0.device) < (1.0 - t)  # [B, L] the survival draw.
+    corrupted = torch.where(keep, x0, MASK)  # [B, L] masking applied inside the block.
+    return torch.where(position < start, x0, torch.where(inside, corrupted, MASK))  # [B, L].
+
+
+def block_loss(model: torch.nn.Module, x0: torch.Tensor, size: int) -> torch.Tensor:
+    """One Monte Carlo sample of the block bound, with one block and one time per sequence."""
+    blocks = math.ceil(x0.shape[1] / size)  # B, the number of blocks the sequence splits into.
+    block = torch.randint(0, blocks, (x0.shape[0],), device=x0.device)  # [B] which factor to sample.
+    t = torch.rand(x0.shape[0], device=x0.device).clamp(min=TIME_FLOOR)  # [B] one noise level each.
+    z = block_input(x0, block, size, t[:, None])  # [B, L] the denoiser's input.
+    position = torch.arange(x0.shape[1], device=x0.device)[None, :]  # [1, L].
+    start = (block * size)[:, None]  # [B, 1].
+    inside = (position >= start) & (position < start + size)  # [B, L] the current block.
+    logits = model(z, t)  # [B, L, 5] the clean-token prediction at every position.
+    logits[..., MASK] = -1e9  # The zero-mask rule: no posterior mass belongs on m.
+    ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none")  # [B, L].
+    scored = ce * ((z == MASK) & inside)  # [B, L] only the masked positions of the current block.
+    # Drawing one block uniformly estimates a sum over blocks, so the term is scaled by that count.
+    return (blocks * (1.0 / t)[:, None] * scored).sum(-1).mean()  # Nats per sequence.
+
+
+# %% 03. The block size as a dial
+#
+# At a block of one position the time integral collapses. The position is masked with probability
+# 1 - alpha_t and weighted by w(t) = -alpha_dot/(1 - alpha_t), so the integrand is -alpha_dot times
+# the cross-entropy, and the remaining integral is alpha_0 - alpha_1 = 1.
+def one_position_integral(grid: int = 2048) -> float:
+    """The integral of w(t)(1 - alpha_t) over [0, 1], which the collapse makes exactly one."""
+    t = (torch.arange(grid).double() + 0.5) / grid  # [grid] midpoints of a uniform partition.
+    return float(((1.0 / t) * t).mean())  # w(t)(1 - alpha_t) = 1 at every t on the linear schedule.
+
+
+def one_position_estimate(probability: float, draws: int, generator) -> float:
+    """The same integral by Monte Carlo, which is how the training loss evaluates it."""
+    t = torch.rand(draws, generator=generator).clamp(min=TIME_FLOOR)  # [N] uniform times.
+    masked = torch.rand(draws, generator=generator) < t  # [N] the single position is masked or not.
+    return float(((1.0 / t) * masked * -math.log(probability)).mean())  # Should land on -log d.
+
+
+# %% 04. Generating block by block
+@torch.no_grad()  # Generation needs no gradient tracking.
+def sample(model: torch.nn.Module, batch: int, length: int, size: int,
+           steps: int = 32) -> torch.Tensor:
+    """Fill one block at a time by masked diffusion, with the finished blocks as a clean prefix."""
+    device = next(model.parameters()).device
+    z = torch.full((batch, length), MASK, device=device)  # [B, L] every position starts masked.
+    ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device)  # The reverse grid per block.
+    position = torch.arange(length, device=device)[None, :]  # [1, L] the column index.
+    for start in range(0, length, size):  # Blocks are completed in order, left to right.
+        inside = (position >= start) & (position < start + size)  # [1, L] the current block.
+        for t, s in zip(ts[:-1], ts[1:]):
+            logits = model(z, t.expand(batch))  # [B, L, 5] the prefix conditions every prediction.
+            logits[..., MASK] = -1e9  # The zero-mask rule, at sampling time.
+            x0 = torch.distributions.Categorical(logits=logits).sample()  # [B, L] candidate tokens.
+            unmask = torch.rand(z.shape, device=device) < (t - s) / t  # [B, L] the reveal draw.
+            z = torch.where((z == MASK) & unmask & inside, x0, z)  # Reveal inside this block only.
+        logits = model(z, ts[-1].expand(batch))  # The last step of the block runs down to s = 0.
+        logits[..., MASK] = -1e9
+        final = torch.distributions.Categorical(logits=logits).sample()  # [B, L].
+        z = torch.where((z == MASK) & inside, final, z)  # The block is clean before the next starts.
+    return z  # [B, L] clean DNA.
+
+
+# %% 05. The report
+def main() -> None:
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.add_argument("--block", type=int, default=4, help="L', the block length")
+    # Generation runs the whole reverse grid once per block, so --sample-steps stays at 16 while
+    # training rises to 1000 steps and the report measures 64 sequences, inside about a minute.
+    parser.set_defaults(steps=1000, sample_steps=16, samples=64)
+    args = parser.parse_args()
+    set_seed(args.seed)
+    device = pick_device(args.device)
+    talk = (lambda *a: None) if args.quiet else print  # --quiet keeps only the report block.
+
+    # Section 02, one draw of the bound on ACGTAC with blocks of two at t = 0.5.
+    weight = 1.0 / 0.5  # w(0.5) = 2 on the linear schedule.
+    contributions = [  # Block 1 masked twice, block 2 once, block 3 once.
+        weight * (-math.log(0.5) - math.log(0.8)),
+        weight * -math.log(0.7),
+        weight * -math.log(0.9),
+    ]
+    talk(f"worked block bound  {[round(v, 4) for v in contributions]}  total "
+         f"{sum(contributions):.4f} nats on ACGTAC")
+
+    # Section 03, the one-position identity that makes blocks of length one autoregressive.
+    generator = torch.Generator().manual_seed(args.seed)  # A fixed stream for the estimate.
+    exact = one_position_integral()  # The integral of w(t)(1 - alpha_t), which is one.
+    estimate = one_position_estimate(0.7, 20000, generator)  # The same number by Monte Carlo.
+    talk(f"one-position dial   integral {exact:.4f}  Monte Carlo bound {estimate:.4f}  against "
+         f"-log 0.7 = {-math.log(0.7):.4f}")
+
+    # Train the denoiser on the shared synthetic DNA at the chosen block size.
+    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
+    train_tokens = splits["train_tokens"].to(device)  # [N, L] clean sequences over ACGT.
+    valid_tokens = splits["valid_tokens"].to(device)  # [M, L] held out for the bound below.
+    model = SequenceDenoiser(VOCAB, args.length, width=args.width).to(device)
+    blocks = math.ceil(args.length / args.block)  # How many factors the chain rule has.
+    talk(f"device={device} parameters={count_parameters(model):,} "
+         f"blocks={blocks} of length {args.block}")
+
+    def step_loss(step: int) -> torch.Tensor:
+        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
+        return block_loss(model, train_tokens[rows], args.block)  # One clean batch into one scalar.
+
+    losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
+    with torch.no_grad():
+        bound = sum(float(block_loss(model, valid_tokens, args.block)) for _ in range(8)) / 8.0
+
+    torch.manual_seed(args.seed + 1)  # Fix the reveal draws, so the samples repeat.
+    samples = sample(model, args.samples, args.length, args.block, args.sample_steps)
+    report = {
+        "worked bound on ACGTAC": sum(contributions),
+        "first block of the worked bound": contributions[0],
+        "integral of w(t)(1 - alpha_t)": exact,
+        "one-position Monte Carlo bound": estimate,
+        "blocks": blocks,
+        "block length": args.block,
+        "training bound, last 50": sum(losses[-50:]) / len(losses[-50:]),
+        "validation bound, nats/seq": bound,
+        "bound, nats per token": bound / args.length,
+        "GC fraction": gc_fraction(samples),
+        "GC fraction in the data": gc_fraction(train_tokens),
+        "motif fraction": motif_fraction(samples),
+        "motif fraction in the data": motif_fraction(train_tokens),
+        "distinct samples": unique_fraction(samples),
+        "letter entropy, nats": token_entropy(samples, DNA_SIZE),
+    }
+    print_report(f"block: diffusion in blocks of {args.block}", report, decode_dna(samples))
+    save_run(args.out, model=model, config=vars(args), report=report,
+             losses=losses, samples=decode_dna(samples))
+
+
+# %% 06. Run the script
+if __name__ == "__main__":  # Execute when launched directly.
+    main()  # Work the block bound, train at one block size, and generate block by block.

lecture_4/ctmc.py

@@ -0,0 +1,278 @@
+"""Continuous-time Markov chains on one DNA base, the mathematics with no learning in it.
+
+A rate matrix says what a base can do in the next instant, and the master equation, the matrix
+exponential, the stationary distribution and the time reversal all follow from it. Nothing here is
+trained: every number this file prints is a number the chapter works by hand.
+
+Chapter 4, "The two rules a rate matrix obeys" through "Running the jumps backward". Time runs from
+t = 0 clean to t = 1 all masks, the survival probability is alpha_t, and the rates R_t(x, y) act on
+probability rows, p_dot = p R. Prints every worked quantity of those sections.
+
+    python lecture_4/ctmc.py                   # the worked numbers, 10 Euler steps
+    python lecture_4/ctmc.py --steps 20        # a finer Euler grid in Section 09
+"""
+
+# %% 01. Imports and configuration
+import math  # Logarithms and exponentials of the worked numbers.
+
+import torch  # Matrices, eigenvalues and the matrix exponential.
+
+from cis6270.data import DNA, DNA_SIZE  # The alphabet ACGT and K = 4.
+from cis6270.runner import common_parser, print_report, save_run, set_seed  # Flags and the report.
+
+MASK_STATES = DNA_SIZE + 1  # A, C, G, T and m, the five states of the masked alphabet.
+
+
+# %% 02. The two rules a rate matrix obeys
+def rate_matrix(off_diagonal: torch.Tensor) -> torch.Tensor:
+    """Fill the diagonal of a nonnegative array so that R_t(x, x) = -sum_{y != x} R_t(x, y)."""
+    R = off_diagonal.clone().double()  # [K, K] the chosen jump rates.
+    R.fill_diagonal_(0.0)  # The diagonal adds no information of its own.
+    return R - torch.diag(R.sum(dim=1))  # [K, K] every row now sums to zero.
+
+
+def obeys_the_two_rules(R: torch.Tensor, tol: float = 1e-9) -> bool:
+    """Nonnegative off the diagonal, and zero row sums, which is the definition."""
+    off = R - torch.diag(torch.diag(R))  # [K, K] the diagonal removed.
+    return bool((off >= -tol).all()) and bool(R.sum(dim=1).abs().max() < tol)
+
+
+def exit_rates(R: torch.Tensor) -> torch.Tensor:
+    """lambda_t(x) = -R_t(x, x), the total rate of leaving each state."""
+    return -torch.diag(R)  # [K] one exit rate per state.
+
+
+# %% 03. One row of a rate matrix for DNA
+#
+# A jumps to C at rate 2 and to G and to T at rate 1 each. The chapter gives that one row; the
+# other three repeat the pattern with the next base in ACGT order taking the rate 2.
+DNA_RATES = rate_matrix(torch.tensor([
+    [0.0, 2.0, 1.0, 1.0],  # A: rate 2 toward C, rate 1 toward G and T.
+    [1.0, 0.0, 2.0, 1.0],  # C: rate 2 toward G.
+    [1.0, 1.0, 0.0, 2.0],  # G: rate 2 toward T.
+    [2.0, 1.0, 1.0, 0.0],  # T: rate 2 toward A.
+]))
+
+
+def short_interval_row(R: torch.Tensor, state: int, h: float) -> torch.Tensor:
+    """One row of P_{t,t+h} = I + h R_t + o(h), the first-order transition probabilities."""
+    identity_row = torch.zeros(R.shape[0], dtype=R.dtype)  # [K] the row of I at this state.
+    identity_row[state] = 1.0  # Staying put is the only certainty at h = 0.
+    return identity_row + h * R[state]  # [K] probabilities, valid while the diagonal stays positive.
+
+
+def largest_valid_step(R: torch.Tensor) -> float:
+    """The bound h * sum_{y != x} R_t(x, y) <= 1, read at the state with the largest exit rate."""
+    return float(1.0 / exit_rates(R).max())  # One bad state is enough to give a negative entry.
+
+
+# %% 04. Waiting times and destinations
+#
+# A state x holds for an exponential time of rate lambda(x), because S(s + h) = S(s)(1 - lambda h)
+# gives dS/ds = -lambda S. The destination is then the row of rates normalized to sum to one.
+def sample_path(R: torch.Tensor, start: int, horizon: float, generator) -> list[tuple[float, int]]:
+    """Gillespie simulation: an exponential wait, then a destination drawn from the normalized row."""
+    time_now, state, path = 0.0, start, [(0.0, start)]  # The path begins at the starting base.
+    while True:
+        rate = float(exit_rates(R)[state])  # lambda(x) for the base we are sitting on.
+        u = float(torch.rand((), generator=generator))  # One uniform variate for the wait.
+        time_now += -math.log(u) / rate  # Inverting S(s) = exp(-lambda s) at u.
+        if time_now > horizon:
+            return path  # The last jump falls outside the interval we are watching.
+        row = R[state].clamp(min=0.0)  # [K] the nonnegative rates out of this base.
+        state = int(torch.multinomial(row / row.sum(), 1, generator=generator))  # The destination.
+        path.append((time_now, state))  # Record when the jump happened and where it went.
+
+
+# %% 05. The master equation
+def master_equation(p: torch.Tensor, R: torch.Tensor) -> torch.Tensor:
+    """p_dot = p R, the Kolmogorov forward equation written on probability rows."""
+    return p @ R  # [K] the y-th entry collects the y-th column of R, every rate pointing into y.
+
+
+def arrivals_and_departures(p: torch.Tensor, R: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """The same derivative split into the flux j_t(x -> y) arriving at y and the flux leaving y."""
+    flux = p[:, None] * R  # [K, K] j_t(x -> y) = p_t(x) R_t(x, y), off the diagonal.
+    arriving = flux.sum(dim=0) - torch.diag(flux)  # [K] sum over x != y of p_t(x) R_t(x, y).
+    leaving = -torch.diag(flux)  # [K] p_t(y) times the exit rate of y.
+    return arriving, leaving  # Their difference is the master equation.
+
+
+# %% 06. A worked evolution of four DNA probabilities
+#
+# Every base jumps to every other at rate 1, so each diagonal entry is -3 and each coordinate obeys
+# p_dot(y) = 1 - 4 p(y), whose solution is 1/4 + (p_0(y) - 1/4) exp(-4t).
+UNIT_RATES = rate_matrix(torch.ones(DNA_SIZE, DNA_SIZE))  # [4, 4] all pairwise rates equal to 1.
+START = torch.tensor([0.1, 0.2, 0.3, 0.4], dtype=torch.float64)  # p_t over (A, C, G, T).
+
+
+def exact_evolution(p0: torch.Tensor, t: float) -> torch.Tensor:
+    """The closed-form solution of p_dot(y) = 1 - 4 p(y) for the all-ones rate matrix."""
+    return 0.25 + (p0 - 0.25) * math.exp(-4.0 * t)  # [K] each coordinate decays toward 1/4.
+
+
+# %% 07. The matrix exponential and its composition over two intervals
+def transition_matrix(R: torch.Tensor, elapsed: float) -> torch.Tensor:
+    """P_{0,t} = exp(t R), the solution of the master equation for a fixed rate matrix."""
+    return torch.matrix_exp(elapsed * R)  # [K, K] rows sum to one and no entry is negative.
+
+
+def euler_product(R: torch.Tensor, elapsed: float, n: int) -> torch.Tensor:
+    """(I + (t/n) R)^n, the product of n short-interval steps, which approaches exp(t R)."""
+    step = torch.eye(R.shape[0], dtype=R.dtype) + (elapsed / n) * R  # [K, K] one first-order step.
+    return torch.linalg.matrix_power(step, n)  # [K, K] the steps compose by matrix multiplication.
+
+
+# %% 08. The R_pi family and its closed-form kernel
+#
+# At rate one the current token is discarded and a fresh one drawn from pi, so every row of the rate
+# matrix is pi with the current state subtracted. Because M = 1 pi^T is idempotent, the exponential
+# collapses to a mixture of staying put and drawing afresh.
+def r_pi(pi: torch.Tensor) -> torch.Tensor:
+    """R_pi = 1 pi^T - I, the one-free-distribution family."""
+    return torch.ones(len(pi), 1, dtype=pi.dtype) @ pi[None, :] - torch.eye(len(pi), dtype=pi.dtype)
+
+
+def pi_kernel(pi: torch.Tensor, alpha: float) -> torch.Tensor:
+    """Q(alpha) = alpha I + (1 - alpha) 1 pi^T, the kernel at survival probability alpha."""
+    M = torch.ones(len(pi), 1, dtype=pi.dtype) @ pi[None, :]  # [K, K] every row equals pi.
+    return alpha * torch.eye(len(pi), dtype=pi.dtype) + (1.0 - alpha) * M  # [K, K] a convex mixture.
+
+
+# %% 09. Eigenvalues, stationarity and detailed balance
+def detailed_balance_gap(R: torch.Tensor, pi: torch.Tensor) -> float:
+    """The largest mismatch in pi(x) R(x, y) = pi(y) R(y, x), zero when the pair is reversible."""
+    flux = pi[:, None] * R  # [K, K] the stationary flux along each edge.
+    return float((flux - flux.T).abs().max())  # Zero exactly when every edge balances.
+
+
+CYCLE = rate_matrix(torch.tensor([  # A -> C -> G -> A at rate 1, with no jumps the other way.
+    [0.0, 1.0, 0.0],
+    [0.0, 0.0, 1.0],
+    [1.0, 0.0, 0.0],
+]))
+
+
+# %% 10. Time reversal
+def reverse_rates(R: torch.Tensor, p: torch.Tensor) -> torch.Tensor:
+    """R_hat_t(y, x) = R_t(x, y) p_t(x) / p_t(y), the rates of the process watched backward."""
+    reversed_rates = (R * p[:, None] / p[None, :]).T  # [K, K] transpose moves (x, y) to (y, x).
+    reversed_rates.fill_diagonal_(0.0)  # The diagonal is fixed by the row-sum rule, not by the ratio.
+    return reversed_rates - torch.diag(reversed_rates.sum(dim=1))  # [K, K] a rate matrix again.
+
+
+# %% 11. The report
+def main() -> None:
+    parser = common_parser(__doc__.split("\n")[0])
+    # Nothing trains here, so --steps is the Euler grid of Section 09 and --samples the number of
+    # Gillespie paths of Section 04. The flags that describe a network go unread.
+    parser.set_defaults(steps=10, samples=512)
+    args = parser.parse_args()
+    set_seed(args.seed)  # The same seed gives the same sampled path and the same histogram.
+    generator = torch.Generator().manual_seed(args.seed)  # Local stream for the simulations.
+    talk = (lambda *a: None) if args.quiet else print  # --quiet keeps only the report block.
+
+    # Sections 02 and 03, the row for A and its first-order transition probabilities.
+    row_probabilities = short_interval_row(DNA_RATES, 0, 0.1)  # [4] at h = 0.1.
+    exact_stay = math.exp(-4.0 * 0.1)  # The probability of never leaving A over that interval.
+    talk(f"rate row for A      {[round(float(v), 2) for v in DNA_RATES[0]]}  rules hold "
+         f"{obeys_the_two_rules(DNA_RATES)}  exit rate {float(exit_rates(DNA_RATES)[0]):.0f}")
+    talk(f"P(A, .) at h = 0.1  {[round(float(v), 2) for v in row_probabilities]}  "
+         f"exact stay {exact_stay:.4f}  valid while h <= {largest_valid_step(DNA_RATES):.2f}, "
+         f"since at h = 0.3 the staying entry is "
+         f"{float(short_interval_row(DNA_RATES, 0, 0.3)[0]):.1f}")
+
+    # Section 04, the waiting time and the destination, first from the two uniforms of the chapter.
+    wait = -math.log(0.42) / 4.0  # u = 0.42 inverted through S(s) = exp(-4s).
+    destination = DNA[int(torch.cumsum(DNA_RATES[0].clamp(min=0) / 4.0, 0).lt(0.31).sum())]
+    path = sample_path(DNA_RATES, 0, 1.0, generator)  # One sampled trajectory over the unit interval.
+    ends = torch.tensor([sample_path(DNA_RATES, 0, 1.0, generator)[-1][1]
+                         for _ in range(args.samples)])  # [N] the base each path finishes on.
+    histogram = torch.bincount(ends, minlength=DNA_SIZE).double() / args.samples  # [4] proportions.
+    exact_end = transition_matrix(DNA_RATES, 1.0)[0]  # [4] the exact law of the base at t = 1.
+    talk(f"wait from u = 0.42  {wait:.4f}  destination from v = 0.31  {destination}  "
+         f"mean {1 / 4:.2f}  median {math.log(2) / 4:.4f}")
+    talk(f"sampled path        {' -> '.join(f'{DNA[s]}@{time:.3f}' for time, s in path)}")
+    talk(f"{args.samples} paths at t = 1   {[round(float(v), 3) for v in histogram]}  against the "
+         f"exact {[round(float(v), 3) for v in exact_end]}")
+
+    # Sections 05 and 06, the master equation and the worked evolution.
+    derivative = master_equation(START, UNIT_RATES)  # [4] p_dot = (0.6, 0.2, -0.2, -0.6).
+    arriving, leaving = arrivals_and_departures(START, UNIT_RATES)  # The two halves of that number.
+    euler_step = START + 0.1 * derivative  # One step of size h = 0.1, inside the bound h <= 1/3.
+    talk(f"p_dot at A          {float(arriving[0]):.1f} in minus {float(leaving[0]):.1f} out "
+         f"= {float(derivative[0]):.1f}   p_dot {[round(float(v), 1) for v in derivative]}")
+    talk(f"Euler step h = 0.1  {[round(float(v), 2) for v in euler_step]}  exact "
+         f"{[round(float(v), 4) for v in exact_evolution(START, 0.1)]}  "
+         f"total change {float(derivative.sum()):.0e}")
+    talk(f"exact at t = 1      {[round(float(v), 4) for v in exact_evolution(START, 1.0)]}")
+
+    # Section 07, the matrix exponential and the Chapman-Kolmogorov composition.
+    whole = transition_matrix(UNIT_RATES, 0.7)  # P_{s,u} over an interval of length 0.7.
+    first, second = transition_matrix(UNIT_RATES, 0.3), transition_matrix(UNIT_RATES, 0.4)
+    halves = first @ second  # P_{s,t} P_{t,u}, the two intervals composed.
+    talk(f"exp(tR) rows        sum to {float(whole.sum(1).min()):.6f}, smallest entry "
+         f"{float(whole.min()):.6f}, composition gap {float((whole - halves).abs().max()):.2e}")
+
+    # Section 08, the R_pi family, its closed form and the multiplication of survival probabilities.
+    uniform = torch.full((DNA_SIZE,), 1.0 / DNA_SIZE, dtype=torch.float64)  # pi uniform over ACGT.
+    tau = math.log(2.0)  # The clock the linear schedule reaches at t = 0.5, where alpha = 1/2.
+    closed_form = pi_kernel(uniform, math.exp(-tau))  # alpha I + (1 - alpha) 1 pi^T.
+    series = transition_matrix(r_pi(uniform), tau)  # The same matrix from the series.
+    composed = pi_kernel(uniform, 0.8) @ pi_kernel(uniform, 0.5 / 0.8)  # Q(0.8) Q(0.625).
+    talk(f"closed form vs exp  gap {float((closed_form - series).abs().max()):.2e}   "
+         f"Q(0.8) Q(0.625) at (A, A) {float(composed[0, 0]):.5f} = Q(0.5) entry "
+         f"{float(pi_kernel(uniform, 0.5)[0, 0]):.5f}")
+
+    # Section 09, the matrix exponential worked for DNA, and the masking kernel beside it.
+    euler_entries = {n: float(euler_product(r_pi(uniform), tau, n)[0, 0]) for n in (1, 2, args.steps)}
+    mask_pi = torch.zeros(MASK_STATES, dtype=torch.float64)  # pi = delta_m on the five states.
+    mask_pi[DNA_SIZE] = 1.0  # All of the target probability sits on the mask.
+    mask_kernel = pi_kernel(mask_pi, 0.5)  # [5, 5] the masking kernel at alpha = 1/2.
+    talk(f"exp(tau R_pi)       diagonal {float(closed_form[0, 0]):.3f}, off-diagonal "
+         f"{float(closed_form[0, 1]):.3f}, Euler " +
+         ", ".join(f"n={n}: {v:.3f}" for n, v in euler_entries.items()) +
+         f";  masking A row {[round(float(v), 2) for v in mask_kernel[0]]}, m row "
+         f"{[round(float(v), 2) for v in mask_kernel[DNA_SIZE]]}")
+
+    # Section 10, the eigenvalues and the stationary distribution, then detailed balance.
+    eigenvalues = sorted(float(v.real) for v in torch.linalg.eigvals(r_pi(uniform)))  # 0 and -1.
+    cycle_pi = torch.full((3,), 1.0 / 3.0, dtype=torch.float64)  # Uniform over the three bases.
+    talk(f"eigenvalues of R_pi {[round(v, 4) for v in eigenvalues]}, pi R_pi "
+         f"{float((uniform @ r_pi(uniform)).abs().max()):.2e}, uniform flux "
+         f"{float((uniform[:, None] * r_pi(uniform))[0, 1]):.4f} each way with gap "
+         f"{detailed_balance_gap(r_pi(uniform), uniform):.2e}")
+    talk(f"cycle of three      stationary gap {float((cycle_pi @ CYCLE).abs().max()):.2e}, "
+         f"edge gap {detailed_balance_gap(CYCLE, cycle_pi):.4f}")
+
+    # Section 11, the reverse rates and the marginals they retrace.
+    marginal = torch.tensor([0.5, 0.25, 0.125, 0.125], dtype=torch.float64)  # p_t at alpha_t = 1/2.
+    unit_forward = rate_matrix(torch.ones(DNA_SIZE, DNA_SIZE))  # Forward rate 1 between any two bases.
+    backward = reverse_rates(unit_forward, marginal)  # [4, 4] R_hat_t(y, x).
+    talk(f"reverse row for C   {[round(float(v), 3) for v in backward[1]]}  exit rate "
+         f"{float(exit_rates(backward)[1]):.1f}  mean holding time "
+         f"{1 / float(exit_rates(backward)[1]):.4f}")
+
+    report = {
+        "exit rate out of A": float(exit_rates(DNA_RATES)[0]),
+        "P(A, A) at h = 0.1": float(row_probabilities[0]),
+        "largest valid step": largest_valid_step(DNA_RATES),
+        "wait from u = 0.42": wait,
+        "p_dot at A": float(derivative[0]),
+        "Euler step at A": float(euler_step[0]),
+        "composition gap": float((whole - halves).abs().max()),
+        "exp(tau R_pi) at (A, A)": float(closed_form[0, 0]),
+        "exp(tau R_pi) off diagonal": float(closed_form[0, 1]),
+        f"Euler product, n = {args.steps}": euler_entries[args.steps],
+        "detailed balance gap, cycle": detailed_balance_gap(CYCLE, cycle_pi),
+        "reverse rate C to A": float(backward[1, 0]),
+        "reverse rate C to G": float(backward[1, 2]),
+        "p R_hat plus p R": float((marginal @ backward + marginal @ unit_forward).abs().max()),
+    }
+    print_report("ctmc: rates, marginals and their reversal", report)
+    save_run(args.out, config=vars(args), report=report)
+
+
+# %% 12. Run the script
+if __name__ == "__main__":  # Execute when launched directly.
+    main()  # Work every number of the sections above.

lecture_4/d3pm.py

@@ -0,0 +1,234 @@
+"""D3PM on DNA, discrete time with one transition matrix per step.
+
+Corruption multiplies a one-hot token by a transition matrix once per step, and the product of
+those matrices stays inside the pi family, so the cumulative kernel is closed form. Reversing one
+finite step is Bayes' rule on three factors, and the continuous-time rate matrices of the rest of
+the chapter are the limit of these products.
+
+Chapter 4, "Discrete time, as in D3PM", "Reversing one finite step" and "Rate matrices chosen for
+the data". Time runs from t = 0 clean to t = T noise, the cumulative survival probability is
+alpha_bar_t, and rates act on probability rows, p_dot = p R. Trained on the synthetic DNA of
+cis6270.data. Prints the worked cumulative kernel, the worked reverse step, and the sample
+statistics.
+
+    python lecture_4/d3pm.py                  # train and generate
+    python lecture_4/d3pm.py --steps 20       # a few seconds, for a first look
+"""
+
+# %% 01. Imports and configuration
+import math  # The continuous-time limit is an exponential.
+
+import torch  # Tensors, gradients and the categorical draws.
+
+from cis6270.data import DNA_SIZE, decode_dna, dna_splits  # ACGT, K = 4 and the shared sequences.
+from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
+from cis6270.nets import SequenceDenoiser, count_parameters  # The denoiser over clean tokens.
+from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
+
+AUXILIARY = 0.01  # lambda_aux, the coefficient of the cross-entropy D3PM adds to the bound.
+UNIFORM = torch.full((DNA_SIZE,), 1.0 / DNA_SIZE)  # pi uniform over the four bases.
+
+
+# %% 02. The transition matrices and their product
+#
+# Q_t = (1 - beta_t) I + beta_t 1 pi^T is the one-step version of the pi family, and two members
+# multiply to a third whose survival probability is the product of the two.
+def one_step(beta: float) -> torch.Tensor:
+    """Q_t, with beta the probability of discarding the token and redrawing from pi. Returns [K, K]."""
+    return (1.0 - beta) * torch.eye(DNA_SIZE) + beta * UNIFORM.expand(DNA_SIZE, DNA_SIZE)
+
+
+def cumulative(alpha_bar: float) -> torch.Tensor:
+    """Q_bar_t = alpha_bar_t I + (1 - alpha_bar_t) 1 pi^T, the product of the first t matrices."""
+    return alpha_bar * torch.eye(DNA_SIZE) + (1.0 - alpha_bar) * UNIFORM.expand(DNA_SIZE, DNA_SIZE)
+
+
+def schedule(timesteps: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """A linear cumulative survival alpha_bar_t = 1 - t/T, and the per-step beta_t it implies."""
+    alpha_bar = 1.0 - torch.arange(timesteps + 1).to(torch.get_default_dtype()) / timesteps  # [T + 1] with alpha_bar_0 = 1.
+    beta = 1.0 - alpha_bar[1:] / alpha_bar[:-1].clamp(min=1e-12)  # [T] the redraw probability.
+    return alpha_bar, beta  # The two descriptions of the same corruption.
+
+
+# %% 03. Reversing one finite step
+#
+# The numerator holds the two halves of the path through the candidate x_{t-1}, one from the clean
+# token forward to the candidate and one from the candidate forward to the observation.
+def exact_posterior(observed: torch.Tensor, clean_distribution: torch.Tensor,
+                    step_matrix: torch.Tensor, previous: torch.Tensor) -> torch.Tensor:
+    """q(x_{t-1} | x_t, x_0) when clean_distribution is a point mass, averaged over it otherwise.
+
+    observed: [B, L] token ids. clean_distribution: [B, L, K] over candidate clean tokens.
+    step_matrix and previous: [B, K, K], one Q_t and one Q_bar_{t-1} per sequence.
+    Returns [B, L, K], normalized over the candidate token at step t - 1.
+    """
+    length = observed.shape[1]
+    columns = step_matrix[:, None].expand(-1, length, -1, -1)  # [B, L, K, K] the step matrix.
+    to_observation = columns.gather(  # [B, L, K] the column of Q_t at the observed token.
+        3, observed[:, :, None, None].expand(-1, -1, DNA_SIZE, 1)).squeeze(-1)
+    from_clean = torch.bmm(clean_distribution, previous)  # [B, L, K] the row of Q_bar_{t-1}.
+    joint = to_observation * from_clean  # [B, L, K] the two halves of the path, multiplied.
+    return joint / joint.sum(-1, keepdim=True).clamp(min=1e-20)  # [B, L, K] a distribution.
+
+
+# %% 04. Rate matrices chosen for the data
+#
+# Nothing in the five facts requires that every pair of tokens be treated alike. A discretized
+# Gaussian keeps a pixel near its own intensity, and a nearest-neighbour graph does the same for
+# tokens.
+def gaussian_row(levels: int = 5, center: int = 2, sigma: float = 1.0) -> torch.Tensor:
+    """The off-diagonal weights out of one intensity level, normalized to a total rate of one."""
+    offsets = torch.arange(levels).to(torch.get_default_dtype()) - center  # [levels] the distance to each level.
+    weights = torch.exp(-offsets ** 2 / (2.0 * sigma ** 2))  # [levels] nearby levels weigh more.
+    weights[center] = 0.0  # The diagonal is fixed by the row-sum rule.
+    return weights / weights.sum()  # [levels] a row of rates summing to one.
+
+
+def graph_rates(adjacency: torch.Tensor) -> torch.Tensor:
+    """R(i, j) = A_ij off the diagonal and R(i, i) = -sum_l A_il, from a token-similarity graph."""
+    R = adjacency.clone()  # [K, K] the graph weights.
+    R.fill_diagonal_(0.0)  # Only the off-diagonal entries are rates.
+    return R - torch.diag(R.sum(dim=1))  # [K, K] every row sums to zero.
+
+
+# %% 05. The discrete-time bound
+#
+# The loss is the variational bound of Section "A variational bound for the whole path", one KL
+# divergence per step and per position, plus the auxiliary cross-entropy on the clean tokens.
+def d3pm_loss(model: torch.nn.Module, x0: torch.Tensor, matrices: tuple) -> torch.Tensor:
+    """One sample of L_vb + lambda_aux CE, with one step index and one corruption per sequence."""
+    steps, cumulatives, alpha_bar = matrices  # [T, K, K], [T + 1, K, K] and [T + 1].
+    timesteps = len(steps)
+    index = torch.randint(1, timesteps + 1, (x0.shape[0],), device=x0.device)  # [B] the step t.
+    probabilities = cumulatives[index][  # [B, L, K] the row of Q_bar_t at each clean token.
+        torch.arange(len(x0), device=x0.device)[:, None], x0]
+    xt = torch.distributions.Categorical(probs=probabilities).sample()  # [B, L] the noisy tokens.
+
+    time = index.to(torch.get_default_dtype()) / timesteps  # [B] the step index as a time in [0, 1] for the network.
+    denoiser = model(xt, time).softmax(-1)  # [B, L, K] the posterior over clean tokens.
+    one_hot = torch.nn.functional.one_hot(x0, DNA_SIZE).to(torch.get_default_dtype())  # [B, L, K] the true clean token.
+    step_matrix, previous = steps[index - 1], cumulatives[index - 1]  # [B, K, K] each.
+    target = exact_posterior(xt, one_hot, step_matrix, previous)  # [B, L, K] the exact reverse step.
+    predicted = exact_posterior(xt, denoiser, step_matrix, previous)  # [B, L, K] the model's own.
+    divergence = (target * (target.clamp(min=1e-20)  # KL(q || p_theta) at every position.
+                            / predicted.clamp(min=1e-20)).log()).sum((1, 2))  # [B].
+    cross_entropy = -(one_hot * denoiser.clamp(min=1e-20).log()).sum((1, 2))  # [B] the auxiliary term.
+    # Drawing one step index uniformly estimates a sum of T terms, so the divergence is scaled by T.
+    return (timesteps * divergence + AUXILIARY * cross_entropy).mean()  # Nats per sequence.
+
+
+# %% 06. Ancestral sampling
+@torch.no_grad()  # Generation needs no gradient tracking.
+def sample(model: torch.nn.Module, batch: int, length: int, matrices: tuple) -> torch.Tensor:
+    """Walk the chain from x_T drawn from pi down to x_0. Returns [B, L] over ACGT."""
+    steps, cumulatives, _ = matrices
+    device = next(model.parameters()).device
+    timesteps = len(steps)
+    x = torch.distributions.Categorical(  # [B, L] the prior, which is pi at every position.
+        probs=UNIFORM.to(device).expand(batch, length, DNA_SIZE)).sample()
+    for t in range(timesteps, 0, -1):  # One finite reverse step at a time.
+        time = torch.full((batch,), t / timesteps, device=device)  # [B] the step as a time.
+        denoiser = model(x, time).softmax(-1)  # [B, L, K] the prediction of the clean token.
+        posterior = exact_posterior(  # [B, L, K] the reverse step, averaged over the denoiser.
+            x, denoiser, steps[t - 1].expand(batch, -1, -1), cumulatives[t - 1].expand(batch, -1, -1))
+        x = torch.distributions.Categorical(probs=posterior).sample()  # [B, L] the earlier tokens.
+    return x  # [B, L] clean DNA.
+
+
+# %% 07. The report
+def main() -> None:
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.add_argument("--timesteps", type=int, default=32, help="T, the number of discrete steps")
+    # T fixes both the corruption and the ancestral sampler, so this file reads it from --timesteps
+    # and leaves --sample-steps alone. A default run trains 1000 steps on 64 sequences, inside
+    # about a minute.
+    parser.set_defaults(steps=1000, samples=64)
+    args = parser.parse_args()
+    set_seed(args.seed)
+    device = pick_device(args.device)
+    talk = (lambda *a: None) if args.quiet else print  # --quiet keeps only the report block.
+
+    # Section 02, the worked cumulative kernel with beta_1 = 0.2 and beta_2 = 0.3.
+    first, second = one_step(0.2), one_step(0.3)  # [K, K] the two transition matrices.
+    product = first @ second  # [K, K] the cumulative kernel after two steps.
+    alpha_bar_two = 0.8 * 0.7  # The product of the per-step survival probabilities.
+    talk(f"Q_bar_2             alpha_bar {alpha_bar_two:.2f}  diagonal {float(product[0, 0]):.2f}  "
+         f"off-diagonal {float(product[0, 1]):.2f}  closed-form gap "
+         f"{float((product - cumulative(alpha_bar_two)).abs().max()):.2e}")
+
+    # Section 02 again, the continuous-time limit of the same product at a constant rate.
+    limits = {n: (1.0 - 1.0 / n) ** n for n in (10, 100, 1000)}  # The products climb toward exp(-1).
+    talk("constant rate 1     " + "  ".join(f"T={n}: {v:.4f}" for n, v in limits.items()) +
+         f"  exp(-1) = {math.exp(-1):.4f}")
+
+    # Section 03, reversing one finite step from an observed C under a clean A.
+    clean_a = torch.tensor([[[1.0, 0.0, 0.0, 0.0]]])  # [1, 1, K] the clean base A.
+    observed_c = torch.tensor([[1]])  # [1, 1] the base observed at step 2.
+    history = exact_posterior(observed_c, clean_a, second[None], first[None])[0, 0]  # [K] at step 1.
+    talk(f"observed C, clean A column of Q_2 {[round(float(v), 3) for v in second.T[1]]}  row of "
+         f"Q_bar_1 {[round(float(v), 2) for v in first[0]]}")
+    talk(f"q(x_1 | x_2, x_0)   {[round(float(v), 4) for v in history]}  "
+         f"evidence {float(product[0, 1]):.2f}")
+
+    # Section 04, two rate matrices built for the data they corrupt.
+    neighbours = gaussian_row()  # [5] the row out of a mid-grey pixel.
+    adjacency = torch.tensor([  # A to C, C to G and G to T at weight 1, A to G at weight 0.5.
+        [0.0, 1.0, 0.5, 0.0], [1.0, 0.0, 1.0, 0.0],
+        [0.5, 1.0, 0.0, 1.0], [0.0, 0.0, 1.0, 0.0]])
+    R = graph_rates(adjacency)  # [K, K] the rate matrix of the token graph.
+    graph_kernel = torch.matrix_exp(0.3 * R)  # [K, K] its transition matrix at clock 0.3.
+    flux = UNIFORM[:, None] * R  # [K, K] the stationary flux along each edge.
+    talk(f"gaussian row        neighbour {float(neighbours[1]):.4f}  two away "
+         f"{float(neighbours[0]):.4f}  ratio {float(neighbours[1] / neighbours[0]):.1f}")
+    talk(f"exp(0.3 R) first row {[round(float(v), 4) for v in graph_kernel[0]]}  detailed balance "
+         f"gap {float((flux - flux.T).abs().max()):.2e}")
+
+    # Train the denoiser against the discrete-time bound.
+    alpha_bar, beta = schedule(args.timesteps)
+    matrices = (torch.stack([one_step(float(b)) for b in beta]).to(device),
+                torch.stack([cumulative(float(a)) for a in alpha_bar]).to(device),
+                alpha_bar.to(device))
+    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
+    train_tokens = splits["train_tokens"].to(device)  # [N, L] clean sequences over ACGT.
+    valid_tokens = splits["valid_tokens"].to(device)  # [M, L] held out for the bound below.
+    model = SequenceDenoiser(DNA_SIZE, args.length, width=args.width).to(device)
+    talk(f"device={device} parameters={count_parameters(model):,} T={args.timesteps}")
+
+    def step_loss(step: int) -> torch.Tensor:
+        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
+        return d3pm_loss(model, train_tokens[rows], matrices)  # One clean batch into one scalar.
+
+    losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
+    with torch.no_grad():
+        held_out = sum(float(d3pm_loss(model, valid_tokens, matrices)) for _ in range(4)) / 4.0
+
+    torch.manual_seed(args.seed + 1)  # Fix the ancestral draws, so the samples repeat.
+    samples = sample(model, args.samples, args.length, matrices)  # [B, L] over ACGT.
+    report = {
+        "alpha_bar after two steps": alpha_bar_two,
+        "Q_bar_2 on the diagonal": float(product[0, 0]),
+        "Q_bar_2 off the diagonal": float(product[0, 1]),
+        "product at T = 1000": limits[1000],
+        "q(x_1 = A | C, A)": float(history[0]),
+        "q(x_1 = C | C, A)": float(history[1]),
+        "gaussian neighbour weight": float(neighbours[1]),
+        "gaussian far weight": float(neighbours[0]),
+        "exp(0.3 R) at (A, C)": float(graph_kernel[0, 1]),
+        "graph detailed balance gap": float((flux - flux.T).abs().max()),
+        "training bound, last 50": sum(losses[-50:]) / len(losses[-50:]),
+        "validation bound, nats/seq": held_out,
+        "GC fraction": gc_fraction(samples),
+        "GC fraction in the data": gc_fraction(train_tokens),
+        "motif fraction": motif_fraction(samples),
+        "motif fraction in the data": motif_fraction(train_tokens),
+        "distinct samples": unique_fraction(samples),
+        "letter entropy, nats": token_entropy(samples, DNA_SIZE),
+    }
+    print_report("d3pm: discrete-time diffusion on DNA", report, decode_dna(samples))
+    save_run(args.out, model=model, config=vars(args), report=report,
+             losses=losses, samples=decode_dna(samples))
+
+
+# %% 08. Run the script
+if __name__ == "__main__":  # Execute when launched directly.
+    main()  # Work the finite-step numbers, train against the bound, and generate.

lecture_4/guidance.py

@@ -0,0 +1,308 @@
+"""Guiding masked diffusion on DNA with Doob's h-transform.
+
+Conditioning a jump process on an event observed at the end of its path multiplies every rate by
+the ratio of the value function at the destination to the value at the current state. The reverse
+process is itself such a transform, with h_t = p_t/pi, and a second transform by a classifier aims
+generation at a property. Three estimates of that classifier appear here: the exact one, its
+first-order expansion in the one-hot input, and the conditional denoiser of classifier-free
+guidance.
+
+Chapter 4, "Conditioning a jump process: Doob's h-transform", "The reverse process as an
+h-transform" and "Guiding generation toward a property". Time runs from t = 0 clean to t = 1 all
+masks, the survival probability is alpha_t = 1 - t, and the rates R_t(x, y) act on probability
+rows, p_dot = p R. Trained on the synthetic DNA of cis6270.data, whose label is a GC fraction of at
+least 0.6. Prints the worked guided jump and the GC fraction under each kind of guidance.
+
+    python lecture_4/guidance.py              # train, then sample four ways
+    python lecture_4/guidance.py --gamma 3    # a sharper transform
+    python lecture_4/guidance.py --steps 20   # a few seconds, for a first look
+"""
+
+# %% 01. Imports and configuration
+import torch  # Tensors, gradients and the categorical draws.
+import torch.nn as nn  # The classifier's layers.
+import torch.nn.functional as F  # Cross-entropy and the one-hot encoding.
+
+from cis6270.data import DNA_SIZE, decode_dna, dna_splits  # ACGT, K = 4 and the shared sequences.
+from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
+from cis6270.nets import TimeEmbedding, SequenceDenoiser, count_parameters
+from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
+
+MASK = DNA_SIZE  # A, C, G, T are 0 to 3, and token 4 is the mask m.
+VOCAB = DNA_SIZE + 1  # The masked alphabet V_m has K + 1 = 5 symbols.
+NULL = 2  # The null label the conditional denoiser drops to, one past the two real labels.
+DROP = 0.1  # How often training replaces the label by the null one.
+TIME_FLOOR = 1e-3  # The reveal rate w(t) = 1/t is unbounded at t = 0.
+
+
+# %% 02. Doob's h-transform
+#
+# R^h_t(x, y) = R_t(x, y) h_t(y)/h_t(x). A jump that raises the chance of the event is faster by
+# the factor by which it raises it, the graph of allowed jumps is unchanged, and a constant h
+# changes nothing, which is the statement that conditioning on a certain event does nothing.
+def h_transform(R: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
+    """Reweight every off-diagonal rate by h(y)/h(x), then refill the diagonal. Returns [K, K]."""
+    transformed = R * (h[None, :] / h[:, None])  # [K, K] one ratio per edge.
+    transformed.fill_diagonal_(0.0)  # The diagonal is fixed by the row-sum rule.
+    return transformed - torch.diag(transformed.sum(dim=1))  # [K, K] a rate matrix again.
+
+
+# %% 03. The reverse process as an h-transform
+#
+# For a forward process in detailed balance with pi, the reverse rate equals the forward rate times
+# h_t(x)/h_t(y) with h_t = p_t/pi, so generation is the noise process aimed at the states that are
+# more likely under the data than under pure noise.
+def value_from_marginal(p: torch.Tensor, pi: torch.Tensor) -> torch.Tensor:
+    """h_t = p_t / pi, defined wherever pi is positive, as uniform noise is and masking is not."""
+    return p / pi  # [K] the excess of the data marginal over the noise distribution.
+
+
+# %% 04. A classifier on noisy sequences
+#
+# The value function for a property is h_t(x) = p(c | X_t = x), estimated by a classifier trained at
+# every noise level. The input is a one-hot sequence, so a gradient with respect to it exists and
+# Section 06 can take one.
+class NoisyClassifier(nn.Module):
+    """A small network from a one-hot noisy sequence and a time to the logit of p(c | z, t)."""
+
+    def __init__(self, length: int, width: int = 128):
+        super().__init__()
+        self.embed = nn.Linear(VOCAB, width, bias=False)  # One-hot in, so the input has a gradient.
+        self.position = nn.Embedding(length, width)  # Where each token sits.
+        self.time = TimeEmbedding(width)  # How far along the corruption we are.
+        self.body = nn.Sequential(nn.Linear(width, width), nn.SiLU(), nn.Linear(width, width),
+                                  nn.SiLU())
+        self.head = nn.Linear(width, 1)  # One logit per sequence.
+
+    def forward(self, one_hot: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        """one_hot: [B, L, V], t: [B] -> [B] the logit of the property."""
+        positions = torch.arange(one_hot.shape[1], device=one_hot.device)  # [L].
+        h = (self.embed(one_hot) + self.position(positions)[None]
+             + self.time(t)[:, None, :])  # [B, L, width].
+        return self.head(self.body(h).mean(dim=1))[:, 0]  # [B] pooled over the positions.
+
+
+def as_one_hot(z: torch.Tensor) -> torch.Tensor:
+    """The sequence over the masked alphabet as floats, so a gradient in the input exists."""
+    return F.one_hot(z, VOCAB).to(torch.get_default_dtype())  # [B, L, V] one row per position.
+
+
+def classifier_loss(classifier: nn.Module, x0: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+    """Binary cross-entropy on sequences corrupted to a random noise level."""
+    t = torch.rand(x0.shape[0], device=x0.device).clamp(min=TIME_FLOOR)  # [B] one level each.
+    keep = torch.rand(x0.shape, device=x0.device) < (1.0 - t)[:, None]  # [B, L] the survival draw.
+    z = torch.where(keep, x0, MASK)  # [B, L] a partially masked sequence.
+    logit = classifier(as_one_hot(z), t)  # [B] the predicted log odds of the property.
+    return F.binary_cross_entropy_with_logits(logit, labels.to(torch.get_default_dtype()))
+
+
+# %% 05. Exact classifier guidance
+#
+# The guided rate has p(c | z^{i <- a}) in it at every position and every token, so the exact
+# version costs one classifier evaluation per candidate. A length-16 DNA sequence has 64 of them.
+def exact_ratio(classifier: nn.Module, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+    """p(c | z^{i <- a}) / p(c | z) at every position and clean token. Returns [B, L, K]."""
+    batch, length = z.shape
+    one_hot = as_one_hot(z)  # [B, L, V] the current state.
+    here = torch.sigmoid(classifier(one_hot, t))[:, None, None]  # [B, 1, 1] the value at z.
+    candidates = one_hot[:, None, None].repeat(1, length, DNA_SIZE, 1, 1)  # [B, L, K, L, V].
+    index = torch.arange(length, device=z.device)  # [L] the position being changed.
+    candidates[:, index, :, index] = torch.eye(VOCAB, device=z.device)[:DNA_SIZE]  # Set position i.
+    flat = candidates.reshape(batch * length * DNA_SIZE, length, VOCAB)  # One batch of candidates.
+    time = t[:, None, None].expand(batch, length, DNA_SIZE).reshape(-1)  # [B L K] the same time.
+    value = torch.sigmoid(classifier(flat, time)).view(batch, length, DNA_SIZE)  # [B, L, K].
+    return value / here.clamp(min=1e-8)  # [B, L, K] the h-transform factor on each jump.
+
+
+# %% 06. The gradient approximation
+#
+# One backward pass through the classifier gives every ratio at once, because a first-order
+# expansion of log p(c | .) in the one-hot input differs between two candidates only in the two
+# coordinates that change.
+def gradient_ratio(classifier: nn.Module, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+    """exp(g_{i,a} - g_{i,z^i}) with g the gradient of log p(c | z) in the one-hot. [B, L, K]."""
+    one_hot = as_one_hot(z).requires_grad_(True)  # [B, L, V] the point of the expansion.
+    with torch.enable_grad():  # The sampler runs under no_grad, and this one pass needs gradients.
+        log_value = F.logsigmoid(classifier(one_hot, t)).sum()  # One backward covers the batch.
+        gradient, = torch.autograd.grad(log_value, one_hot)  # [B, L, V] the expansion's slopes.
+    current = gradient.gather(2, z[:, :, None])  # [B, L, 1] the coordinate that is being left.
+    return (gradient[..., :DNA_SIZE] - current).exp()  # [B, L, K] the approximated ratio.
+
+
+# %% 07. The guided sampler
+#
+# The reveal rate at a masked position is w(t) d_theta(a), so the h-transform multiplies each of
+# those rates by its own ratio. Both the probability of revealing anything and the letter that
+# appears change, and a ratio of one at every token returns the unguided sampler exactly.
+def denoiser_probabilities(model: nn.Module, z: torch.Tensor, time: torch.Tensor,
+                           labels: torch.Tensor | None) -> torch.Tensor:
+    """d_theta(. | z, t) over the four bases, with the zero-mask rule applied. Returns [B, L, K]."""
+    logits = model(z, time, labels)  # [B, L, 5] the clean-token prediction.
+    logits[..., MASK] = -1e9  # The zero-mask rule: no posterior mass belongs on m.
+    return logits.softmax(-1)[..., :DNA_SIZE]  # [B, L, K], which sums to one over ACGT.
+
+
+def guided_weights(model: nn.Module, z: torch.Tensor, time: torch.Tensor,
+                   labels: torch.Tensor | None, ratio_fn, gamma: float) -> torch.Tensor:
+    """The reveal rates over w(t): the denoiser times its h-transform ratio. Returns [B, L, K]."""
+    weights = denoiser_probabilities(model, z, time, labels)  # [B, L, K] the unguided rates.
+    if ratio_fn is None:
+        return weights  # [B, L, K] summing to one, which is the unguided exit rate w(t).
+    return weights * ratio_fn(z, time).clamp(min=1e-8) ** gamma  # [B, L, K] the h-transform.
+
+
+@torch.no_grad()  # Only the gradient ratio reopens the graph, and it does so locally.
+def sample(model: nn.Module, labels: torch.Tensor | None, batch: int, length: int,
+           steps: int = 32, ratio_fn=None, gamma: float = 1.0) -> torch.Tensor:
+    """Generate from all masks, with the reveal rates reweighted by ratio_fn. Returns [B, L]."""
+    device = next(model.parameters()).device
+    z = torch.full((batch, length), MASK, device=device)  # [B, L] the all-mask start.
+    ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device)  # The reverse grid.
+    for t, s in zip(ts[:-1], ts[1:]):
+        time = t.expand(batch)  # [B] the shared time of this step.
+        weights = guided_weights(model, z, time, labels, ratio_fn, gamma)  # [B, L, K].
+        total = weights.sum(-1, keepdim=True)  # [B, L, 1] the transformed exit rate over w(t).
+        reveal = ((t - s) / t * total).clamp(max=1.0)  # [B, L, 1] the probability of acting at all.
+        drawn = torch.distributions.Categorical(probs=weights / total).sample()  # [B, L].
+        open_now = torch.rand(z.shape, device=device)[:, :, None] < reveal  # [B, L, 1].
+        z = torch.where((z == MASK) & open_now[..., 0], drawn, z)  # Carry-over unmasking.
+    # The last step runs down to s = 0, where the schedule reveals every remaining mask. The same
+    # transform picks the letter, so the guidance reaches the positions that waited to the end.
+    last = ts[-1].expand(batch)  # [B] the final time of the grid.
+    weights = guided_weights(model, z, last, labels, ratio_fn, gamma)  # [B, L, K].
+    drawn = torch.distributions.Categorical(  # [B, L] one letter per position.
+        probs=weights / weights.sum(-1, keepdim=True)).sample()
+    return torch.where(z == MASK, drawn, z)  # [B, L] clean DNA.
+
+
+# %% 08. The report
+def main() -> None:
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.add_argument("--gamma", type=float, default=2.0, help="exponent on the h ratio")
+    # Two networks train and four samplers run, so each network gets 1000 steps and each sampler
+    # measures 64 sequences, which is the longest default run of the lecture.
+    parser.set_defaults(steps=1000, sample_steps=32, samples=64)
+    args = parser.parse_args()
+    set_seed(args.seed)
+    device = pick_device(args.device)
+    talk = (lambda *a: None) if args.quiet else print  # --quiet keeps only the report block.
+
+    # Sections 02 and 03, the reverse rates read off h_t = p_t / pi at alpha_t = 1/2.
+    marginal = torch.tensor([0.5, 0.25, 0.125, 0.125])  # p_t over (A, C, G, T).
+    pi = torch.full((DNA_SIZE,), 0.25)  # Uniform noise.
+    unit = torch.ones(DNA_SIZE, DNA_SIZE) - torch.diag(torch.full((DNA_SIZE,), DNA_SIZE * 1.0))
+    values = value_from_marginal(marginal, pi)  # [4] h_t = (2, 1, 0.5, 0.5).
+    reverse = h_transform(unit, values)  # [4, 4] the reverse rates, at unit forward rates.
+    conditioned = value_from_marginal(torch.tensor([0.625, 0.125, 0.125, 0.125]), pi)  # Clean A.
+    talk(f"h from the marginal {[round(float(v), 2) for v in values]}  reverse row for C "
+         f"{[round(float(v), 2) for v in reverse[1]]}")
+    talk(f"h given a clean A   {[round(float(v), 2) for v in conditioned]}  rate C to A "
+         f"{float(conditioned[0] / conditioned[1]):.0f}  given a clean C "
+         f"{float(conditioned[1] / conditioned[0]):.1f}")
+
+    # Section 05 on paper, the worked guided jump out of C with a classifier on the four bases.
+    unguided = torch.tensor([2.0, 0.5, 0.5])  # The reverse rates toward A, G and T.
+    classifier_value = torch.tensor([0.1, 0.6, 0.2]) / 0.2  # p(c | destination) over p(c | C).
+    guided = unguided * classifier_value  # [3] the guided rates.
+    sharper = unguided * classifier_value ** 2  # [3] the same transform at gamma = 2.
+    shares = [round(float(v), 4) for v in guided / guided.sum()]  # Where the jump most likely goes.
+    talk(f"guided jump         {[round(float(v), 2) for v in guided]}  exit "
+         f"{float(guided.sum()):.0f}  destinations {shares}")
+    talk(f"at gamma = 2        {[round(float(v), 2) for v in sharper]}  exit "
+         f"{float(sharper.sum()):.1f}  destinations "
+         f"{[round(float(v), 4) for v in sharper / sharper.sum()]}  holding "
+         f"{1 / float(sharper.sum()):.4f}")
+
+    # Train the conditional denoiser and the classifier on the shared synthetic DNA.
+    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
+    tokens = splits["train_tokens"].to(device)  # [N, L] clean sequences over ACGT.
+    labels = splits["train_labels"].to(device)  # [N] one when the GC fraction is at least 0.6.
+    model = SequenceDenoiser(VOCAB, args.length, width=args.width, num_classes=2).to(device)
+    classifier = NoisyClassifier(args.length, args.width).to(device)
+    share = float(labels.to(torch.get_default_dtype()).mean())  # How much of the data is GC-rich.
+    talk(f"device={device} denoiser={count_parameters(model):,} "
+         f"classifier={count_parameters(classifier):,} labelled {share:.2f}")
+
+    def denoiser_loss(step: int) -> torch.Tensor:
+        """The MDLM loss, with the label replaced by the null one a tenth of the time."""
+        rows = torch.randint(0, len(tokens), (args.batch_size,), device=device)
+        x0, y = tokens[rows], labels[rows].clone()
+        y[torch.rand(len(y), device=device) < DROP] = NULL  # Dropping the label trains both heads.
+        t = torch.rand(len(x0), device=device).clamp(min=TIME_FLOOR)  # [B] one noise level each.
+        zt = torch.where(torch.rand(x0.shape, device=device) < (1.0 - t)[:, None], x0, MASK)
+        logits = model(zt, t, y)  # [B, L, 5].
+        logits[..., MASK] = -1e9  # The zero-mask rule.
+        ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none")  # [B, L].
+        return ((1.0 / t)[:, None] * ce * (zt == MASK)).sum(-1).mean()
+
+    def labelled_batch(step: int) -> torch.Tensor:
+        """One batch of clean sequences and their labels, for the classifier."""
+        rows = torch.randint(0, len(tokens), (args.batch_size,), device=device)
+        return classifier_loss(classifier, tokens[rows], labels[rows])
+
+    losses = train(model, denoiser_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
+    classifier_losses = train(classifier, labelled_batch, steps=args.steps, lr=args.lr,
+                              quiet=args.quiet)
+
+    # Section 06, the exact ratios against their first-order approximation on one noisy batch.
+    probe_t = torch.full((32,), 0.5, device=device)  # [32] halfway through the corruption.
+    probe = torch.where(torch.rand(tokens[:32].shape, device=device) < 0.5, tokens[:32], MASK)
+    with torch.no_grad():
+        exact = exact_ratio(classifier, probe, probe_t)  # [32, L, K] every candidate evaluated.
+    approximate = gradient_ratio(classifier, probe, probe_t)  # [32, L, K] one backward pass.
+    gap = exact.clamp(min=1e-6).log() - approximate.clamp(min=1e-6).log()  # [32, L, K] log ratios.
+    agreement = float(gap.abs().mean())  # How far the expansion sits from the exact ratio.
+
+    # Sample four ways, from the same weights and the same starting state.
+    wanted = torch.ones(args.samples, dtype=torch.long, device=device)  # Ask for the GC-rich label.
+    null = torch.full_like(wanted, NULL)  # The label the unconditional head reads.
+
+    def free_ratio(z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        """d_theta(a | z, c) / d_theta(a | z), which is the classifier ratio implied by Bayes."""
+        # Both denoisers go through the zero-mask rule first. The loss overwrites the mask logit,
+        # so it never trains, and a softmax that kept it would leave a stray factor in the ratio.
+        conditional = denoiser_probabilities(model, z, t, wanted)  # [B, L, K] with the label.
+        unconditional = denoiser_probabilities(model, z, t, null)  # [B, L, K] with the null label.
+        return conditional / unconditional.clamp(min=1e-8)  # [B, L, K] the h-transform factor.
+
+    runs = {}
+    for name, ratio in (
+        ("unguided", None),
+        ("exact classifier", lambda z, t: exact_ratio(classifier, z, t)),
+        ("gradient approximation", lambda z, t: gradient_ratio(classifier, z, t)),
+        ("classifier free", free_ratio),
+    ):
+        torch.manual_seed(args.seed + 1)  # The same reveal draws for every run.
+        runs[name] = sample(model, null, args.samples, args.length, args.sample_steps,
+                            ratio, args.gamma)
+
+    report = {
+        "h at A from the marginal": float(values[0]),
+        "reverse rate C to A": float(reverse[1, 0]),
+        "guided rate C to A": float(guided[0]),
+        "guided rate C to G": float(guided[1]),
+        "guided exit rate": float(guided.sum()),
+        "guided exit rate at gamma = 2": float(sharper.sum()),
+        "share on G at gamma = 2": float((sharper / sharper.sum())[1]),
+        "denoiser loss, last 50": sum(losses[-50:]) / len(losses[-50:]),
+        "classifier loss, last 50": sum(classifier_losses[-50:]) / len(classifier_losses[-50:]),
+        "mean log-ratio gap, exact vs gradient": agreement,
+        "GC fraction in the data": gc_fraction(tokens),
+        "motif fraction in the data": motif_fraction(tokens),
+        "gamma": args.gamma,
+    }
+    for name, drawn in runs.items():
+        report[f"GC fraction, {name}"] = gc_fraction(drawn)
+        # The label is a GC threshold, so the motif is the structure guidance leaves alone.
+        report[f"motif fraction, {name}"] = motif_fraction(drawn)
+        report[f"distinct, {name}"] = unique_fraction(drawn)
+    report["letter entropy, exact classifier"] = token_entropy(runs["exact classifier"], DNA_SIZE)
+    print_report("guidance: h-transforms of masked diffusion", report,
+                 decode_dna(runs["exact classifier"]))
+    save_run(args.out, model=model, config=vars(args), report=report, losses=losses,
+             samples=decode_dna(runs["exact classifier"]))
+
+
+# %% 09. Run the script
+if __name__ == "__main__":  # Execute when launched directly.
+    main()  # Work the guided jump, train the two networks, and sample four ways.

lecture_4/mdlm.py

@@ -0,0 +1,190 @@
+"""Masked diffusion on DNA, the whole model in thirty lines.
+
+Corruption replaces each base by the mask with probability 1 - alpha_t, a network reads the masked
+sequence and predicts the clean token at every position, and generation reveals masked positions
+with the probability the schedule allows. The corruption, the loss and the sampler below are the
+three listings of Section "The whole model for DNA in thirty lines", expanded with comments.
+
+Chapter 4, "Masked diffusion on one token" through "The whole model for DNA in thirty lines". Time
+runs from t = 0 clean to t = 1 all masks, the survival probability is alpha_t = 1 - t, and the rates
+R_t(x, y) act on probability rows, p_dot = p R. Trained on the synthetic DNA of cis6270.data. Prints
+the chapter's worked reverse step and loss term, then the training loss and the sample statistics.
+
+    python lecture_4/mdlm.py                  # train and generate, about a minute on a CPU
+    python lecture_4/mdlm.py --steps 20       # a few seconds, for a first look
+"""
+
+# %% 01. Imports and configuration
+import math  # The worked numbers of the chapter are logarithms.
+
+import torch  # Tensors, gradients and the categorical draws.
+import torch.nn.functional as F  # Cross-entropy at every position.
+from torch.distributions import Categorical  # One draw per position from the denoiser.
+
+from cis6270.data import DNA_SIZE, decode_dna, dna_splits  # ACGT, K = 4 and the shared sequences.
+from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
+from cis6270.nets import SequenceDenoiser, count_parameters  # The denoiser d_theta(. | z, t).
+from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
+
+MASK = DNA_SIZE  # A, C, G, T are 0 to 3, and token 4 is the mask m.
+VOCAB = DNA_SIZE + 1  # The masked alphabet V_m has K + 1 = 5 symbols.
+TIME_FLOOR = 1e-3  # The weight 1/t is unbounded at t = 0, so draws start just above it.
+
+
+# %% 02. The corruption
+#
+# Fact 2 of the construction: the forward marginal q_t(z^i | x_0^i) = alpha_t delta(z^i, x_0^i) +
+# (1 - alpha_t) delta(z^i, m) is closed form, so any noise level is reachable in one draw.
+def corrupt(x0: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+    """Mask each position independently with probability 1 - alpha_t. x0: [B, L], t: [B, 1]."""
+    alpha = 1.0 - t  # The linear schedule, alpha_0 = 1 and alpha_1 = 0.
+    keep = torch.rand(x0.shape, device=x0.device) < alpha  # [B, L] one Bernoulli survival draw.
+    return torch.where(keep, x0, MASK)  # [B, L] the clean token where it survived, m where it did not.
+
+
+# %% 03. The weighted cross-entropy at masked positions
+#
+# The per-step KL divergence collapses to -r log d_theta(a), and summing the steps into an integral
+# replaces r by the reveal rate w(t) = -alpha_dot / (1 - alpha) = 1/t on the linear schedule.
+def loss(model: torch.nn.Module, x0: torch.Tensor) -> torch.Tensor:
+    """One Monte Carlo sample of L_mask, with one time and one corruption per sequence."""
+    t = torch.rand(x0.shape[0], device=x0.device).clamp(min=TIME_FLOOR)  # [B] one noise level each.
+    zt = corrupt(x0, t[:, None])  # [B, L] the time spreads across the positions of its own sequence.
+    logits = model(zt, t)  # [B, L, 5] the denoiser reads the whole noisy sequence.
+    logits[..., MASK] = -1e9  # The zero-mask rule: no posterior mass belongs on m.
+    ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none")  # [B, L] -log d_theta(x_0^i).
+    w = 1.0 / t  # [B] the reveal rate, which is -alpha_dot / (1 - alpha) on this schedule.
+    return (w[:, None] * ce * (zt == MASK)).sum(-1).mean()  # Visible positions contribute zero.
+
+
+# %% 04. The sampler
+#
+# A masked position is revealed with probability r_{s,t} = (alpha_s - alpha_t)/(1 - alpha_t), which
+# is (t - s)/t here, and the letter that appears is one draw from the denoiser.
+@torch.no_grad()  # Generation needs no gradient tracking.
+def sample(model: torch.nn.Module, batch: int, length: int, steps: int = 32) -> torch.Tensor:
+    """Generate from all masks down to clean DNA. Returns [B, L] token ids over ACGT."""
+    device = next(model.parameters()).device  # Sample on the device the weights sit on.
+    z = torch.full((batch, length), MASK, device=device)  # [B, L] the all-mask start, pi = delta_m.
+    ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device)  # The reverse grid, t down to 0.
+    for t, s in zip(ts[:-1], ts[1:]):  # One interval [s, t] at a time.
+        logits = model(z, t.expand(batch))  # [B, L, 5] the clean-token prediction at this level.
+        logits[..., MASK] = -1e9  # The zero-mask rule again, at sampling time.
+        x0 = Categorical(logits=logits).sample()  # [B, L] one candidate clean token per position.
+        unmask = torch.rand(z.shape, device=device) < (t - s) / t  # [B, L] the reveal draw.
+        z = torch.where((z == MASK) & unmask, x0, z)  # Carry-over unmasking: a revealed base stays.
+    logits = model(z, ts[-1].expand(batch))  # The last step runs down to s = 0, where r = 1.
+    logits[..., MASK] = -1e9  # The same rule once more.
+    return torch.where(z == MASK, Categorical(logits=logits).sample(), z)  # [B, L] clean DNA.
+
+
+# %% 05. A complete reverse step, worked
+#
+# With alpha_s = 0.7 and alpha_t = 0.4, which the linear schedule reaches at s = 0.3 and t = 0.6,
+# the reveal probability is half and the denoiser splits that half between A and C.
+def reveal_probability(alpha_s: float, alpha_t: float) -> float:
+    """r_{s,t} = (alpha_s - alpha_t)/(1 - alpha_t), the lost survival mass over the mask mass."""
+    return (alpha_s - alpha_t) / (1.0 - alpha_t)  # A function of the schedule and of nothing else.
+
+
+def reverse_step_distribution(alpha_s: float, alpha_t: float, denoiser: list[float]) -> list[float]:
+    """The reverse law over (A, C, G, T, m): r times the denoiser, with 1 - r left on the mask."""
+    r = reveal_probability(alpha_s, alpha_t)  # The probability of acting at all.
+    return [r * d for d in denoiser] + [1.0 - r]  # [5] the five outcomes of one step.
+
+
+# %% 06. One term of the loss, and the rate error it comes from
+#
+# At t = 0.5 the weight is w(0.5) = 2, and the clean sequence ACGT corrupted to AmGm contributes the
+# negative log probabilities of C and T. The chapter reaches the same weighted cross-entropy a
+# second time from the rate error, on the reveal figure's denoiser: summing l(A, B) over every
+# destination at one masked position leaves w(t) times the negative log probability of the clean
+# token there, once the staying terms restore the normalization of the denoiser.
+def reveal_rate(t: float) -> float:
+    """w(t) = -alpha_dot / (1 - alpha_t), which is 1/t for the linear schedule."""
+    return 1.0 / t  # The rate at which a masked position reveals something, whatever the denoiser.
+
+
+def rate_error(target: float, model_rate: float) -> float:
+    """l(A, B) = A log(A/B) - A + B, the Bregman divergence that the path KL divergence produces."""
+    if target == 0.0:
+        return model_rate  # l(0, B) = B, which is the staying term on its own.
+    return target * math.log(target / model_rate) - target + model_rate  # Zero exactly at B = A.
+
+
+# %% 07. The report
+def main() -> None:
+    parser = common_parser(__doc__.split("\n")[0])
+    # The motif needs a longer budget than the GC fraction, so the default trains for 1000 steps
+    # and measures on 64 sequences. The whole run stays inside about a minute on one core.
+    parser.set_defaults(steps=1000, sample_steps=64, samples=64)
+    args = parser.parse_args()
+    set_seed(args.seed)  # The same seed gives the same training batches and the same samples.
+    device = pick_device(args.device)
+    talk = (lambda *a: None) if args.quiet else print  # --quiet keeps only the report block.
+
+    # Section 05, the worked reverse step at alpha_s = 0.7 and alpha_t = 0.4.
+    r = reveal_probability(0.7, 0.4)  # 0.3 of recovered mass over 0.6 of mask mass.
+    step_law = reverse_step_distribution(0.7, 0.4, [0.6, 0.4, 0.0, 0.0])  # [5] over (A, C, G, T, m).
+    talk(f"reverse step        r = {r:.1f}  p(A, C, G, T, m) = {[round(v, 2) for v in step_law]}  "
+         f"sum {sum(step_law):.1f}")
+
+    # Section 06, one term of the loss at t = 0.5 on ACGT corrupted to AmGm.
+    weight = reveal_rate(0.5)  # w(0.5) = 2.
+    masked_ce = -math.log(0.7) - math.log(0.5)  # The clean tokens C and T at positions 2 and 4.
+    worked_term = weight * masked_ce  # One sample of the objective.
+    position_two = [0.2, 0.6, 0.1, 0.1]  # The reveal figure's denoiser at the second position.
+    reveal_rates = [weight * d for d in position_two]  # [4] rates that sum to w(t) = 2.
+    collapse = rate_error(weight, reveal_rates[1]) + sum(  # The clean token is C here.
+        rate_error(0.0, rate) for i, rate in enumerate(reveal_rates) if i != 1)
+    talk(f"loss term at t=0.5  w = {weight:.0f}  masked CE = {masked_ce:.5f}  "
+         f"w CE = {worked_term:.5f}")
+    talk(f"reveal rates        {[round(v, 1) for v in reveal_rates]}  sum {sum(reveal_rates):.0f}  "
+         f"rate error {collapse:.4f} = w(-log 0.6) = {weight * -math.log(0.6):.4f}")
+
+    # Train the denoiser on the shared synthetic DNA.
+    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
+    train_tokens = splits["train_tokens"].to(device)  # [N, L] clean sequences over ACGT.
+    valid_tokens = splits["valid_tokens"].to(device)  # [M, L] held out for the bound below.
+    model = SequenceDenoiser(VOCAB, args.length, width=args.width).to(device)
+    talk(f"device={device} parameters={count_parameters(model):,} "
+         f"train={len(train_tokens)} length={args.length}")
+
+    def step_loss(step: int) -> torch.Tensor:
+        """One optimizer step draws its own batch of clean sequences."""
+        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
+        return loss(model, train_tokens[rows])  # [B, L] clean sequences into one scalar.
+
+    losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
+
+    # The validation bound, averaged over several corruptions of the held-out sequences.
+    with torch.no_grad():
+        bound = sum(float(loss(model, valid_tokens)) for _ in range(4)) / 4.0  # Nats per sequence.
+
+    # Generate, then measure the structure the data was built with.
+    torch.manual_seed(args.seed + 1)  # Fix the reveal draws, so the samples repeat.
+    samples = sample(model, args.samples, args.length, args.sample_steps)  # [B, L] over ACGT.
+    report = {
+        "reveal probability r": r,
+        "reverse step on the mask": step_law[-1],
+        "weight w(0.5)": weight,
+        "worked loss term": worked_term,
+        "rate error at position 2": collapse,
+        "training loss, last 50": sum(losses[-50:]) / len(losses[-50:]),
+        "validation bound, nats/seq": bound,
+        "bound, nats per token": bound / args.length,
+        "GC fraction": gc_fraction(samples),
+        "GC fraction in the data": gc_fraction(train_tokens),
+        "motif fraction": motif_fraction(samples),
+        "motif fraction in the data": motif_fraction(train_tokens),
+        "distinct samples": unique_fraction(samples),
+        "letter entropy, nats": token_entropy(samples, DNA_SIZE),
+    }
+    print_report("mdlm: masked diffusion on DNA", report, decode_dna(samples))
+    save_run(args.out, model=model, config=vars(args), report=report,
+             losses=losses, samples=decode_dna(samples))
+
+
+# %% 08. Run the script
+if __name__ == "__main__":  # Execute when launched directly.
+    main()  # Work the numbers, train the denoiser, and generate.

lecture_4/planner.py

@@ -0,0 +1,238 @@
+"""Choosing the order of the reveals in masked diffusion on DNA.
+
+Every masked position leaves at the same total rate w(t), so which positions a step reveals is free
+and the marginal reverse rates stay where they were. A planner spends that freedom on the positions
+the denoiser is most sure about, a corrector returns the least supported visible tokens to the
+mask, and the loss can be reweighted toward the positions the sampler reaches first.
+
+Chapter 4, "Choosing the order of the reveals", "One sampling step with a planner" and "Bringing
+the planner into training". Time runs from t = 0 clean to t = 1 all masks, the survival probability
+is alpha_t = 1 - t, and the rates R_t(x, y) act on probability rows, p_dot = p R. Trained on the
+synthetic DNA of cis6270.data. Prints the worked selection, the worked PAPL weights, and the
+sample statistics under both samplers.
+
+    python lecture_4/planner.py                     # train with planner weights, sample both ways
+    python lecture_4/planner.py --remask 0          # confidence order with no corrector
+    python lecture_4/planner.py --steps 20          # a few seconds, for a first look
+"""
+
+# %% 01. Imports and configuration
+import torch  # Tensors, gradients and the categorical draws.
+import torch.nn.functional as F  # Cross-entropy at every position.
+
+from cis6270.data import DNA_SIZE, decode_dna, dna_splits  # ACGT, K = 4 and the shared sequences.
+from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
+from cis6270.nets import SequenceDenoiser, count_parameters  # The denoiser d_theta(. | z, t).
+from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
+
+MASK = DNA_SIZE  # A, C, G, T are 0 to 3, and token 4 is the mask m.
+VOCAB = DNA_SIZE + 1  # The masked alphabet V_m has K + 1 = 5 symbols.
+TIME_FLOOR = 1e-3  # The reveal rate w(t) = 1/t is unbounded at t = 0.
+
+
+# %% 02. The planner score and the selection it drives
+#
+# The score condenses the distribution at a position into one number. The probability of the most
+# likely clean token is high where the denoiser has committed and low where it is spread out, and
+# the negative entropy of the same distribution is the other natural choice.
+def planner_score(probabilities: torch.Tensor, kind: str = "confidence") -> torch.Tensor:
+    """One number per position from d_theta(. | z, t). probabilities: [B, L, K] -> [B, L]."""
+    if kind == "entropy":
+        return (probabilities * probabilities.clamp(min=1e-20).log()).sum(-1)  # Negative entropy.
+    return probabilities.max(-1).values  # The probability of the most likely clean token.
+
+
+def topk_where(score: torch.Tensor, eligible: torch.Tensor, count: int) -> torch.Tensor:
+    """A boolean mask selecting the ``count`` highest-scoring eligible positions in each row."""
+    picked = torch.zeros_like(eligible)  # [B, L] nothing selected yet.
+    if count <= 0:
+        return picked  # Taking none is a valid choice, and it recovers the plain sampler.
+    ranked = score.masked_fill(~eligible, float("-inf"))  # Ineligible positions never win.
+    index = ranked.topk(min(count, score.shape[1]), dim=1).indices  # [B, count] the best positions.
+    return picked.scatter_(1, index, True) & eligible  # Rows with too few eligible slots take fewer.
+
+
+# %% 03. One sampling step with a planner
+#
+# The first four lines are the sampler of the implementation section: one forward pass, the
+# zero-mask rule, a softmax, and one categorical draw per position. Nothing in them depends on the
+# planner, so a trained checkpoint can be sampled this way with no retraining.
+def planner_step(model: torch.nn.Module, z: torch.Tensor, t: torch.Tensor,
+                 n_unmask: int, n_remask: int, kind: str = "confidence") -> torch.Tensor:
+    """One reverse step that reveals the best ``n_unmask`` masks and remasks the worst tokens."""
+    logits = model(z, t)  # [B, L, 5] the clean-token prediction at this noise level.
+    logits[..., MASK] = -1e9  # The zero-mask rule: no posterior mass belongs on m.
+    p = logits.softmax(-1)  # [B, L, 5] the denoiser as probabilities.
+    x0 = torch.distributions.Categorical(probs=p).sample()  # [B, L] one candidate per position.
+
+    score = planner_score(p[..., :DNA_SIZE], kind)  # [B, L] how sure the denoiser is here.
+    masked = z == MASK  # [B, L] the positions a reveal may act on.
+    pick = topk_where(score, masked, n_unmask)  # [B, L] the best masked positions.
+    z = torch.where(pick, x0, z)  # Unmask the best.
+
+    back = topk_where(-score, ~masked, n_remask)  # [B, L] the worst visible positions.
+    return torch.where(back, torch.full_like(z, MASK), z)  # Remask the worst, which is the corrector.
+
+
+# %% 04. The corrector and its flux balance
+#
+# C_t = R_t + R_hat_t is in detailed balance with p_t, so running it at a fixed time holds the
+# marginal in place. Under masking its forward part returns a visible token to the mask at rate
+# beta(t) and its reverse part reveals a masked position at rate w(t).
+def corrector_fluxes(t: float) -> tuple[float, float]:
+    """alpha_t beta(t) and (1 - alpha_t) w(t), the two fluxes across the visible-masked boundary."""
+    alpha = 1.0 - t  # The linear schedule.
+    return alpha * (1.0 / alpha), (1.0 - alpha) * (1.0 / t)  # Both equal -alpha_dot, which is one.
+
+
+# %% 05. Bringing the planner into training
+#
+# The MDLM loss weights every masked position equally, which is right when the order of the reveals
+# is uniform. Once a planner chooses the order, the weight lambda_i = M rho_i puts the gradient on
+# the positions the sampler will reach first, and it averages to one over the masked positions.
+def planner_weights(score: torch.Tensor, masked: torch.Tensor, strength: float) -> torch.Tensor:
+    """lambda_i = 1 + strength (M rho_i - 1), with rho the planner's softmax over masked positions."""
+    # A large negative number in place of -inf, so a row with nothing masked still has a softmax.
+    ranked = score.masked_fill(~masked, -1e4)  # [B, L] only masked positions compete.
+    rho = ranked.softmax(dim=1)  # [B, L] the planner's distribution over which position to fill.
+    count = masked.sum(dim=1, keepdim=True).clamp(min=1)  # [B, 1] M, the number of masked positions.
+    return 1.0 + strength * (count * rho - 1.0)  # [B, L] a weight of one at strength zero.
+
+
+def papl_loss(model: torch.nn.Module, x0: torch.Tensor, strength: float) -> torch.Tensor:
+    """The MDLM loss with the planner weight inside the sum over masked positions."""
+    t = torch.rand(x0.shape[0], device=x0.device).clamp(min=TIME_FLOOR)  # [B] one noise level each.
+    keep = torch.rand(x0.shape, device=x0.device) < (1.0 - t)[:, None]  # [B, L] the survival draw.
+    zt = torch.where(keep, x0, MASK)  # [B, L] the corrupted sequence.
+    logits = model(zt, t)  # [B, L, 5] the clean-token prediction.
+    logits[..., MASK] = -1e9  # The zero-mask rule.
+    ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none")  # [B, L].
+    masked = zt == MASK  # [B, L] where the loss lives.
+    score = planner_score(logits.softmax(-1)[..., :DNA_SIZE]).detach()  # No gradient through these.
+    weights = planner_weights(score, masked, strength).clamp(min=0.0)  # [B, L] lambda_i.
+    return ((1.0 / t)[:, None] * weights * ce * masked).sum(-1).mean()  # Nats per sequence.
+
+
+# %% 06. Sampling with and without the planner
+#
+# The number revealed per step follows the schedule: after a step down to s, the masked count is
+# the one the survival probability calls for, which is s L on the linear schedule.
+@torch.no_grad()  # Generation needs no gradient tracking.
+def sample(model: torch.nn.Module, batch: int, length: int, steps: int = 32,
+           n_remask: int = 0, kind: str = "confidence", planner: bool = True) -> torch.Tensor:
+    """Generate from all masks. With planner false the reveals are drawn uniformly at random."""
+    device = next(model.parameters()).device
+    z = torch.full((batch, length), MASK, device=device)  # [B, L] the all-mask start.
+    ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device)  # The reverse grid.
+    for t, s in zip(ts[:-1], ts[1:]):
+        target = int(round(float(s) * length))  # The masked count the schedule calls for at s.
+        current = int(z.eq(MASK).sum(dim=1).max())  # The largest masked count in the batch.
+        n_unmask = max(current - target, 0) + n_remask  # Enough to reach the target after remasking.
+        if planner:
+            z = planner_step(model, z, t.expand(batch), n_unmask, n_remask, kind)
+        else:  # The uniform order of the plain sampler, with the same number of reveals.
+            z = _random_step(model, z, t.expand(batch), n_unmask)
+    logits = model(z, ts[-1].expand(batch))  # The last step runs down to s = 0, where r = 1.
+    logits[..., MASK] = -1e9
+    return torch.where(z == MASK, torch.distributions.Categorical(logits=logits).sample(), z)
+
+
+def _random_step(model: torch.nn.Module, z: torch.Tensor, t: torch.Tensor,
+                 n_unmask: int) -> torch.Tensor:
+    """The same step with the masked positions ranked by a uniform draw in place of the score."""
+    logits = model(z, t)  # [B, L, 5].
+    logits[..., MASK] = -1e9  # The zero-mask rule.
+    x0 = torch.distributions.Categorical(logits=logits).sample()  # [B, L].
+    noise = torch.rand(z.shape, device=z.device)  # [B, L] an order with no preference.
+    return torch.where(topk_where(noise, z == MASK, n_unmask), x0, z)
+
+
+# %% 07. The report
+def main() -> None:
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.add_argument("--remask", type=int, default=1, help="positions the corrector sends back")
+    parser.add_argument("--planner-strength", type=float, default=1.0,
+                        help="0 gives the MDLM loss, 1 gives the PAPL weights")
+    parser.add_argument("--score", choices=["confidence", "entropy"], default="confidence")
+    # Two samplers run on the same weights, so the default trains 1000 steps and measures 64
+    # sequences in each order, which keeps the run inside about a minute.
+    parser.set_defaults(steps=1000, sample_steps=32, samples=64)
+    args = parser.parse_args()
+    set_seed(args.seed)
+    device = pick_device(args.device)
+    talk = (lambda *a: None) if args.quiet else print  # --quiet keeps only the report block.
+
+    # Section 02, the worked selection on a six-position sequence.
+    scores = torch.tensor([[0.42, 0.88, 0.91, 0.55, 0.31, 0.78]])  # [1, 6] one score per position.
+    masked = torch.tensor([[True, False, True, True, False, True]])  # Positions 1, 3, 4 and 6.
+    revealed = topk_where(scores, masked, 2).nonzero()[:, 1] + 1  # The two best masked positions.
+    returned = topk_where(-scores, ~masked, 1).nonzero()[:, 1] + 1  # The worst visible position.
+    talk(f"planner selection   reveals {sorted(int(i) for i in revealed)}  remasks "
+         f"{sorted(int(i) for i in returned)}")
+
+    # Section 04, the two fluxes across the boundary between visible and masked.
+    talk("corrector flux      " + "  ".join(
+        f"t={t}: beta flux {a:.1f}, w flux {b:.1f}" for t in (0.5, 0.25)
+        for a, b in [corrector_fluxes(t)]))
+
+    # Section 05, the worked PAPL weights on a state with four masked positions.
+    rho = torch.tensor([[0.10, 0.50, 0.15, 0.25]])  # [1, 4] the planner's distribution.
+    cross_entropies = torch.tensor([[1.2, 0.3, 0.9, 0.5]])  # [1, 4] nats at each masked position.
+    lam = 4.0 * rho  # [1, 4] lambda_i = M rho_i, which averages to one.
+    talk(f"PAPL weights        {[round(float(v), 1) for v in lam[0]]}  sum {float(lam.sum()):.1f}  "
+         f"MDLM inner sum {float(cross_entropies.sum()):.2f}  PAPL inner sum "
+         f"{float((lam * cross_entropies).sum()):.2f}")
+
+    # Train the denoiser with the planner weights inside the loss.
+    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
+    train_tokens = splits["train_tokens"].to(device)  # [N, L] clean sequences over ACGT.
+    valid_tokens = splits["valid_tokens"].to(device)  # [M, L] held out for the loss below.
+    model = SequenceDenoiser(VOCAB, args.length, width=args.width).to(device)
+    talk(f"device={device} parameters={count_parameters(model):,} strength={args.planner_strength}")
+
+    def step_loss(step: int) -> torch.Tensor:
+        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
+        return papl_loss(model, train_tokens[rows], args.planner_strength)
+
+    losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
+    with torch.no_grad():
+        weighted = sum(float(papl_loss(model, valid_tokens, args.planner_strength))
+                       for _ in range(4)) / 4.0
+        plain = sum(float(papl_loss(model, valid_tokens, 0.0)) for _ in range(4)) / 4.0
+
+    # Sample the same weights twice, once in a uniform order and once in the planner's order.
+    runs = {}
+    for name, planner, remask in (("uniform order", False, 0),
+                                  ("planner order", True, args.remask)):
+        torch.manual_seed(args.seed + 1)  # The same draws for both samplers.
+        runs[name] = sample(model, args.samples, args.length, args.sample_steps,
+                            remask, args.score, planner)
+
+    report = {
+        "revealed positions": sorted(int(i) for i in revealed),
+        "remasked position": sorted(int(i) for i in returned),
+        "flux at t = 0.5": corrector_fluxes(0.5)[0],
+        "flux at t = 0.25": corrector_fluxes(0.25)[1],
+        "PAPL weight at the planner's pick": float(lam[0, 1]),
+        "MDLM inner sum": float(cross_entropies.sum()),
+        "PAPL inner sum": float((lam * cross_entropies).sum()),
+        "training loss, last 50": sum(losses[-50:]) / len(losses[-50:]),
+        "validation PAPL loss": weighted,
+        "validation MDLM loss": plain,
+        "corrector remasks per step": args.remask,
+    }
+    for name, drawn in runs.items():
+        report[f"GC fraction, {name}"] = gc_fraction(drawn)
+        report[f"motif fraction, {name}"] = motif_fraction(drawn)
+        report[f"distinct, {name}"] = unique_fraction(drawn)
+    report["GC fraction in the data"] = gc_fraction(train_tokens)
+    report["motif fraction in the data"] = motif_fraction(train_tokens)
+    report["letter entropy, planner order"] = token_entropy(runs["planner order"], DNA_SIZE)
+    print_report("planner: the order of the reveals", report, decode_dna(runs["planner order"]))
+    save_run(args.out, model=model, config=vars(args), report=report,
+             losses=losses, samples=decode_dna(runs["planner order"]))
+
+
+# %% 08. Run the script
+if __name__ == "__main__":  # Execute when launched directly.
+    main()  # Work the selection, train with the planner weights, and sample both ways.

lecture_4/search.py

@@ -0,0 +1,250 @@
+"""Guiding masked diffusion on DNA by search over the tree of partial sequences.
+
+For a property computed by an external scorer there is no classifier at intermediate noise levels,
+so the value function is estimated by finishing the sequence and scoring the result. The partial
+sequences form a tree whose edges are single reveals, selection walks it by an upper-confidence
+rule, a rollout scores a new leaf, and the reward propagates back up the branch. Two objectives
+that no single number orders are kept in a Pareto archive.
+
+Chapter 4, "Guiding generation by search" and "Search over the tree of partial sequences". The
+fine-tuning of "Keeping the results of the search" is left to the chapter. Time runs from t = 0
+clean to t = 1 all masks, the survival probability is alpha_t = 1 - t, and the rates R_t(x, y)
+act on probability rows, p_dot = p R.
+Trained on the synthetic DNA of cis6270.data. Prints the worked value estimate, the worked archive,
+and the archive the search finds.
+
+    python lecture_4/search.py                  # train, then search
+    python lecture_4/search.py --searches 64    # a larger tree
+    python lecture_4/search.py --sample-steps 8 # longer rollouts behind every value
+    python lecture_4/search.py --steps 20       # a few seconds, for a first look
+"""
+
+# %% 01. Imports and configuration
+import math  # The upper-confidence rule is a square root of a logarithm.
+
+import torch  # Tensors, gradients and the categorical draws.
+import torch.nn.functional as F  # Cross-entropy at every position.
+
+from cis6270.data import DNA_SIZE, MOTIF, decode_dna, dna_splits, encode_dna
+from cis6270.metrics import gc_fraction, motif_fraction, unique_fraction
+from cis6270.nets import SequenceDenoiser, count_parameters  # The denoiser d_theta(. | z, t).
+from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
+
+MASK = DNA_SIZE  # A, C, G, T are 0 to 3, and token 4 is the mask m.
+VOCAB = DNA_SIZE + 1  # The masked alphabet V_m has K + 1 = 5 symbols.
+TIME_FLOOR = 1e-3  # The reveal rate w(t) = 1/t is unbounded at t = 0.
+EXPLORATION = 1.0  # c_uct, the weight on the visit-count term of the selection rule.
+
+
+# %% 02. The two toy objectives
+#
+# Binding, solubility and permeability do not reduce to one number without a weighting that fixes
+# the trade-off in advance. Two scores stand in for them here, a GC fraction and the best partial
+# match to the planted motif, both in [0, 1] and neither implying the other.
+def objectives(tokens: torch.Tensor) -> torch.Tensor:
+    """Returns [N, 2]: the GC fraction of each sequence and its best match to the motif."""
+    gc = ((tokens == 1) | (tokens == 2)).to(torch.get_default_dtype()).mean(dim=1)  # [N] C and G.
+    motif = encode_dna([MOTIF])[0].to(tokens.device)  # [4] the planted motif ACGT.
+    windows = tokens.unfold(1, len(motif), 1)  # [N, L - 3, 4] every window of the motif's length.
+    hits = (windows == motif).to(torch.get_default_dtype()).mean(dim=2)  # [N, L - 3] per window.
+    match = hits.max(dim=1).values  # [N] the best window of each sequence.
+    return torch.stack([gc, match], dim=1)  # [N, 2] one row per sequence.
+
+
+# %% 03. Rollouts as a value estimate
+#
+# h_t(z) is the probability that a sequence continuing from z ends with the property, so running
+# the sampler to completion N times and averaging the reward estimates it without bias, with a
+# variance that falls like 1/N.
+@torch.no_grad()  # Rollouts need no gradient tracking.
+def complete(model: torch.nn.Module, z: torch.Tensor, t: float, steps: int = 3) -> torch.Tensor:
+    """Run the masked sampler from z at time t down to clean DNA. z: [N, L] -> [N, L]."""
+    device = z.device
+    ts = torch.linspace(t, TIME_FLOOR, steps + 1, device=device)  # The remaining reverse grid.
+    for now, nxt in zip(ts[:-1], ts[1:]):
+        logits = model(z, now.expand(len(z)))  # [N, L, 5] the clean-token prediction.
+        logits[..., MASK] = -1e9  # The zero-mask rule: no posterior mass belongs on m.
+        x0 = torch.distributions.Categorical(logits=logits).sample()  # [N, L].
+        unmask = torch.rand(z.shape, device=device) < (now - nxt) / now  # [N, L] the reveal draw.
+        z = torch.where((z == MASK) & unmask, x0, z)  # Carry-over unmasking.
+    logits = model(z, ts[-1].expand(len(z)))  # The last step runs down to s = 0, where r = 1.
+    logits[..., MASK] = -1e9
+    return torch.where(z == MASK, torch.distributions.Categorical(logits=logits).sample(), z)
+
+
+def value_estimate(model: torch.nn.Module, state: torch.Tensor, t: float, rollouts: int,
+                   steps: int = 3) -> tuple[torch.Tensor, torch.Tensor]:
+    """Average the reward over N completions of one partial sequence. Returns [2] and [N, L]."""
+    batch = state[None, :].repeat(rollouts, 1)  # [N, L] the same state, N independent endings.
+    finished = complete(model, batch, 1.0 if t >= 1.0 else t, steps)  # [N, L] clean sequences.
+    return objectives(finished).mean(dim=0), finished  # [2] the value, and what produced it.
+
+
+# %% 04. The Pareto archive
+#
+# A candidate worse on every axis than one already kept is dominated and discarded, and the rest
+# stay, since each of them beats every other on at least one axis.
+class ParetoArchive:
+    """Every nondominated sequence the search has produced, with its objective vector."""
+
+    def __init__(self) -> None:
+        self.entries: list[tuple[tuple[int, ...], tuple[float, ...]]] = []
+
+    def add(self, sequence: tuple[int, ...], scores: tuple[float, ...]) -> None:
+        """Keep the candidate when nothing in the archive beats it on every axis."""
+        for _, kept in self.entries:
+            if all(k >= s for k, s in zip(kept, scores)) and kept != scores:
+                return  # Dominated, so it adds nothing.
+        self.entries = [(q, kept) for q, kept in self.entries  # Drop what the newcomer dominates.
+                        if not (all(s >= k for s, k in zip(scores, kept)) and kept != scores)]
+        if sequence not in {q for q, _ in self.entries}:
+            self.entries.append((sequence, scores))
+
+
+# %% 05. The tree of partial sequences
+#
+# Selection walks from the root to a leaf by the upper-confidence rule, expansion adds children
+# drawn from the denoiser perturbed by Gumbel noise, a rollout scores one of them, and the reward
+# propagates back up the branch so that every node on it accumulates visits.
+class Node:
+    """A partially unmasked sequence, with the visits and rewards the search has accumulated."""
+
+    def __init__(self, state: torch.Tensor, depth: int, parent=None):
+        self.state, self.depth, self.parent = state, depth, parent
+        self.children: list[Node] = []
+        self.visits, self.total = 0, 0.0  # The visit count and the summed scalar reward.
+
+    @property
+    def mean(self) -> float:
+        """The mean reward at this node, which the selection rule weighs against its visit count."""
+        return self.total / self.visits if self.visits else 0.0
+
+    def best_child(self):
+        """argmax of r_bar_j + c sqrt(log n / n_j), the usual upper-confidence rule."""
+        return max(self.children, key=lambda child: child.mean + EXPLORATION * math.sqrt(
+            math.log(max(self.visits, 1)) / child.visits) if child.visits else float("inf"))
+
+
+@torch.no_grad()  # The expansion needs no gradient tracking.
+def expand(model: torch.nn.Module, node: Node, t: float, branch: int) -> None:
+    """Add children by revealing one masked position, drawn from the denoiser plus Gumbel noise."""
+    masked = (node.state == MASK).nonzero().flatten()  # [m] the positions a reveal may act on.
+    if len(masked) == 0:
+        return  # A complete sequence has no children.
+    logits = model(node.state[None], torch.tensor([t], device=node.state.device))[0]  # [L, 5].
+    logits[..., MASK] = -1e9  # The zero-mask rule.
+    gumbel = -torch.log(-torch.log(torch.rand_like(logits).clamp(min=1e-20)))  # [L, 5].
+    perturbed = (logits + gumbel).masked_fill(  # Taking maxima of this samples the denoiser.
+        (node.state != MASK)[:, None], float("-inf"))
+    flat = perturbed[:, :DNA_SIZE].flatten()  # [L K] one score per position and clean token.
+    for index in flat.topk(min(branch, int(len(masked)) * DNA_SIZE)).indices:  # Distinct reveals.
+        child_state = node.state.clone()  # [L] the parent with one more position filled.
+        child_state[int(index) // DNA_SIZE] = int(index) % DNA_SIZE
+        node.children.append(Node(child_state, node.depth + 1, node))
+
+
+def backup(node: Node, reward: float) -> None:
+    """Add one rollout's reward to every node on the branch from here to the root."""
+    while node is not None:
+        node.visits, node.total = node.visits + 1, node.total + reward
+        node = node.parent
+
+
+# %% 06. The report
+def main() -> None:
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.add_argument("--searches", type=int, default=24, help="tree expansions")
+    parser.add_argument("--branch", type=int, default=3, help="children added at each expansion")
+    parser.add_argument("--rollouts", type=int, default=6, help="completions scored per leaf")
+    # Every leaf is scored by its own rollouts, so --sample-steps is the length of one rollout and
+    # stays short. Training rises to 1000 steps, which keeps the run inside about a minute.
+    parser.set_defaults(steps=1000, sample_steps=3)
+    args = parser.parse_args()
+    set_seed(args.seed)
+    device = pick_device(args.device)
+    talk = (lambda *a: None) if args.quiet else print  # --quiet keeps only the report block.
+
+    # Section 03, the worked value estimate and the reweighting it produces.
+    at_parent = float(torch.tensor([0.1, 0.8, 0.3, 0.2]).mean())  # Four rollouts from z.
+    at_child = float(torch.tensor([0.7, 0.9, 0.6, 0.8]).mean())  # Four rollouts from its child.
+    factor = at_child / at_parent  # The h-transform factor on that edge.
+    sibling = 0.2 / at_parent  # A child whose rollouts averaged 0.2 is slowed by this factor.
+    talk(f"value estimates     h(z) = {at_parent:.2f}  h(z') = {at_child:.2f}  factor "
+         f"{factor:.3f}  a model rate of 1.2 becomes {1.2 * factor:.3f}, and a sibling at 0.2 is "
+         f"slowed by {sibling:.3f}")
+
+    # Section 04, the worked archive on three candidates scored on two axes.
+    worked = ParetoArchive()
+    for name, scores in (("P1", (0.8, 0.3)), ("P2", (0.5, 0.7)), ("P3", (0.4, 0.2))):
+        worked.add(tuple(ord(c) for c in name), scores)  # The name stands in for a sequence.
+    kept = [chr(q[0]) + chr(q[1]) for q, _ in worked.entries]  # Which candidates survived.
+    talk(f"worked archive      keeps {kept}, so P3 is dominated by P1 on both axes")
+
+    # Train the masked diffusion model the search will guide.
+    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
+    train_tokens = splits["train_tokens"].to(device)  # [N, L] clean sequences over ACGT.
+    model = SequenceDenoiser(VOCAB, args.length, width=args.width).to(device)
+    talk(f"device={device} parameters={count_parameters(model):,} searches={args.searches}")
+
+    def step_loss(step: int) -> torch.Tensor:
+        """The masked diffusion loss, unchanged: the search reweights rates, not weights."""
+        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
+        x0 = train_tokens[rows]
+        t = torch.rand(len(x0), device=device).clamp(min=TIME_FLOOR)  # [B] one noise level each.
+        zt = torch.where(torch.rand(x0.shape, device=device) < (1.0 - t)[:, None], x0, MASK)
+        logits = model(zt, t)  # [B, L, 5].
+        logits[..., MASK] = -1e9  # The zero-mask rule.
+        ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none")  # [B, L].
+        return ((1.0 / t)[:, None] * ce * (zt == MASK)).sum(-1).mean()
+
+    losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
+
+    # Run the search, keeping every completed rollout that is not dominated.
+    torch.manual_seed(args.seed + 1)  # Fix the Gumbel draws and the rollouts.
+    root = Node(torch.full((args.length,), MASK, device=device), 0)
+    archive, scored = ParetoArchive(), []  # The archive and every rollout's objective vector.
+    for _ in range(args.searches):
+        node = root  # Selection walks from the root down the tree already built.
+        while node.children:
+            node = node.best_child()
+        time_here = max(1.0 - node.depth / args.length, TIME_FLOOR)  # The level at this depth.
+        expand(model, node, time_here, args.branch)  # Expansion adds children by one reveal each.
+        leaf = node.children[0] if node.children else node  # Score the first new child.
+        value, finished = value_estimate(model, leaf.state, time_here, args.rollouts,
+                                         args.sample_steps)
+        backup(leaf, float(value.mean()))  # One scalar reward guides the selection rule.
+        for sequence, score in zip(finished.tolist(), objectives(finished).tolist()):
+            archive.add(tuple(sequence), tuple(round(v, 4) for v in score))  # Keep the best front.
+            scored.append(score)
+
+    found = torch.tensor([s for _, s in archive.entries])  # [A, 2] the archive's objectives.
+    everything = torch.tensor(scored)  # [R, 2] every rollout the search scored.
+    best = torch.tensor([list(q) for q, _ in archive.entries], device=device)  # [A, L] sequences.
+    report = {
+        "value at the parent": at_parent,
+        "value at the child": at_child,
+        "reweighting factor": factor,
+        "guided rate from 1.2": 1.2 * factor,
+        "factor at a sibling scoring 0.2": sibling,
+        "worked archive size": len(worked.entries),
+        "training loss, last 50": sum(losses[-50:]) / len(losses[-50:]),
+        "rollouts scored": len(scored),
+        "archive size": len(archive.entries),
+        "best GC in the archive": float(found[:, 0].max()),
+        "best motif match in the archive": float(found[:, 1].max()),
+        "mean GC over all rollouts": float(everything[:, 0].mean()),
+        "mean motif match over all rollouts": float(everything[:, 1].mean()),
+        "rollout length, steps": args.sample_steps,
+        "GC fraction in the data": gc_fraction(train_tokens),
+        "motif fraction in the data": motif_fraction(train_tokens),
+        "motif fraction in the archive": motif_fraction(best),
+        "distinct in the archive": unique_fraction(best),
+    }
+    print_report("search: a Pareto archive over partial sequences", report, decode_dna(best))
+    save_run(args.out, model=model, config=vars(args), report=report,
+             losses=losses, samples=decode_dna(best))
+
+
+# %% 07. Run the script
+if __name__ == "__main__":  # Execute when launched directly.
+    main()  # Work the value estimate, train the model, and search under the h-transform.

lecture_4/sedd.py

@@ -0,0 +1,176 @@
+"""SEDD on DNA, a network that predicts the ratios the reverse rates are built from.
+
+The reverse rate at a position is a forward rate times the ratio p_t(z with position i set to a)
+divided by p_t(z), and the two models before this one reached that ratio through a denoiser over
+clean tokens and a sum against the forward kernel. Here the network outputs the ratio itself, and
+the rate error becomes score entropy, whose minimum over the prediction sits at the true ratio.
+
+Chapter 4, "The ratio we have to learn" and "Predicting the ratios directly", on the uniform member
+of the pi family. Time runs from t = 0 clean to t = 1 noise, the survival probability is
+alpha_t = exp(-beta t), and the rates R_t(x, y) act on probability rows, p_dot = p R. Trained on the
+synthetic DNA of cis6270.data. Prints the worked score entropy at three predictions and the sample
+statistics.
+
+    python lecture_4/sedd.py                  # train and generate
+    python lecture_4/sedd.py --steps 20       # a few seconds, for a first look
+"""
+
+# %% 01. Imports and configuration
+import math  # The worked score entropies are logarithms.
+
+import torch  # Tensors, gradients and the categorical draws.
+
+from cis6270.data import DNA_SIZE, decode_dna, dna_splits  # ACGT, K = 4 and the shared sequences.
+from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
+from cis6270.nets import SequenceDenoiser, count_parameters  # Reused as L heads of K outputs.
+from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
+
+BETA = 4.0  # The constant noise rate, so alpha_t = exp(-4t) and the forward rate is BETA/K.
+TIME_FLOOR = 2e-2  # The true ratios grow like 1/(1 - alpha_t), so sampling stops just above zero.
+MAX_RATIO = 1e4  # The predicted ratios are exponentials, kept finite for the sampler.
+
+
+# %% 02. The concrete score
+#
+# s_t(z)_{i,a} = p_t(z^{i <- a}) / p_t(z) holds L(K - 1) numbers per sequence, one per position and
+# per alternative token, which is exactly the shape of a network with L heads of K logits.
+def concrete_score(model: torch.nn.Module, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+    """s_theta(z)_{i,a}, positive and with no normalization constraint. Returns [B, L, K]."""
+    return model(z, t).clamp(max=math.log(MAX_RATIO)).exp()  # [B, L, K] exponentiate the logits.
+
+
+def kernel_ratio(x0: torch.Tensor, z: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
+    """The conditional target r = q_t(a | x_0) / q_t(z^i | x_0), one per destination. [B, L, K]."""
+    uniform = (1.0 - alpha) / DNA_SIZE  # [B, 1] the redrawn share at this noise level.
+    numerator = uniform[..., None] + alpha[..., None] * torch.nn.functional.one_hot(
+        x0, DNA_SIZE)  # [B, L, K] q_t(a | x_0) at every destination a.
+    denominator = numerator.gather(2, z[:, :, None])  # [B, L, 1] the same kernel at the current token.
+    return numerator / denominator  # [B, L, K] the ratio the network has to reproduce.
+
+
+# %% 03. Denoising score entropy
+#
+# Writing the target rate as lambda r and the model rate as lambda s turns the rate error into
+# lambda [s - r log s + r(log r - 1)]. The first two terms hold every dependence on the prediction,
+# a linear term against a logarithmic one, so the minimum over s at fixed r sits at s = r.
+def score_entropy(target: torch.Tensor, predicted: torch.Tensor,
+                  forward_rate: float = 1.0) -> torch.Tensor:
+    """l(lambda r, lambda s), elementwise, which vanishes at s = r."""
+    constant = target * (target.clamp(min=1e-20).log() - 1.0)  # r(log r - 1), free of the prediction.
+    return forward_rate * (predicted - target * predicted.clamp(min=1e-20).log() + constant)
+
+
+def sedd_loss(model: torch.nn.Module, x0: torch.Tensor) -> torch.Tensor:
+    """One Monte Carlo sample of the score entropy, summed over destinations and positions."""
+    t = torch.rand(x0.shape[0], 1, device=x0.device).clamp(min=TIME_FLOOR)  # [B, 1] one time each.
+    alpha = torch.exp(-BETA * t)  # [B, 1] the survival probability at that time.
+    keep = torch.rand(x0.shape, device=x0.device) < alpha  # [B, L] survival of the original base.
+    fresh = torch.randint(0, DNA_SIZE, x0.shape, device=x0.device)  # [B, L] a draw from pi.
+    z = torch.where(keep, x0, fresh)  # [B, L] uniform replacement, with no marker of the damage.
+    target = kernel_ratio(x0, z, alpha)  # [B, L, K] the ratio attached to this clean sequence.
+    predicted = concrete_score(model, z, t[:, 0])  # [B, L, K] the network's own ratios.
+    current = torch.nn.functional.one_hot(z, DNA_SIZE).bool()  # [B, L, K] the token already in place.
+    error = score_entropy(target, predicted, BETA / DNA_SIZE)  # [B, L, K] one term per jump.
+    return error.masked_fill(current, 0.0).sum((1, 2)).mean()  # Nats per sequence.
+
+
+# %% 04. Sampling from the predicted ratios
+#
+# The reverse rate is the forward rate of the jump that would undo it, times the predicted ratio,
+# and tau-leaping turns each rate into a jump probability over a step of length h.
+@torch.no_grad()  # Generation needs no gradient tracking.
+def sample(model: torch.nn.Module, batch: int, length: int, steps: int = 64) -> torch.Tensor:
+    """Generate from uniformly random tokens down to clean DNA. Returns [B, L] over ACGT."""
+    device = next(model.parameters()).device
+    z = torch.randint(0, DNA_SIZE, (batch, length), device=device)  # [B, L] a draw from pi.
+    ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device)  # The reverse grid.
+    for t, s in zip(ts[:-1], ts[1:]):
+        rates = (BETA / DNA_SIZE) * concrete_score(model, z, t.expand(batch))  # [B, L, K].
+        jump = (t - s) * rates  # [B, L, K] the tau-leaping probability of each jump.
+        jump.scatter_(2, z[:, :, None], 0.0)  # A jump to the current base is not a jump.
+        jump = jump / jump.sum(-1, keepdim=True).clamp(min=1.0)  # Rescale an oversized step.
+        stay = (1.0 - jump.sum(-1, keepdim=True)).clamp(min=0.0)  # [B, L, 1] the chance of staying.
+        choice = torch.distributions.Categorical(
+            probs=torch.cat([jump, stay], dim=-1)).sample()  # [B, L] destination, or K for staying.
+        z = torch.where(choice == DNA_SIZE, z, choice.clamp(max=DNA_SIZE - 1))  # [B, L].
+    return z  # [B, L] clean DNA, since uniform replacement never leaves the alphabet.
+
+
+# %% 05. The report
+def main() -> None:
+    parser = common_parser(__doc__.split("\n")[0])
+    # The ratios need a longer budget than the worked numbers do, so the default trains for 1000
+    # steps and measures on 64 sequences, which keeps the run inside about a minute.
+    parser.set_defaults(steps=1000, sample_steps=64, samples=64)
+    args = parser.parse_args()
+    set_seed(args.seed)
+    device = pick_device(args.device)
+    talk = (lambda *a: None) if args.quiet else print  # --quiet keeps only the report block.
+
+    # Section 02, the target ratio of the running example, reached through the denoiser route.
+    posterior = torch.tensor([0.375, 0.625, 0.0, 0.0])  # p(x_0 | X_t = C) at alpha_t = 1/2.
+    per_clean = torch.tensor([5.0, 0.2, 1.0, 1.0])  # q_t(A | b) / q_t(C | b) for b over ACGT.
+    target_ratio = float(posterior @ per_clean)  # The ratio in the A direction, which comes out 2.
+
+    # Section 03, the score entropy at three predictions of that ratio.
+    errors = {s: float(score_entropy(torch.tensor(target_ratio), torch.tensor(s)))
+              for s in (1.8, 1.0, 3.0)}  # A ten percent error against two larger ones.
+    talk(f"target ratio        {target_ratio:.1f} from 0.375(5) + 0.625(0.2)")
+    at_target = float(score_entropy(torch.tensor(target_ratio), torch.tensor(target_ratio)))
+    talk("score entropy       " + "  ".join(f"s={s}: {v:.4f}" for s, v in errors.items()) +
+         f"  at s = r it is {at_target:.1e}")
+
+    # Train the network to output the ratios.
+    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
+    train_tokens = splits["train_tokens"].to(device)  # [N, L] clean sequences over ACGT.
+    valid_tokens = splits["valid_tokens"].to(device)  # [M, L] held out for the loss below.
+    model = SequenceDenoiser(DNA_SIZE, args.length, width=args.width).to(device)
+    talk(f"device={device} parameters={count_parameters(model):,} ratios per sequence "
+         f"{args.length * (DNA_SIZE - 1)}")
+
+    def step_loss(step: int) -> torch.Tensor:
+        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
+        return sedd_loss(model, train_tokens[rows])  # One clean batch into one scalar.
+
+    losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
+    with torch.no_grad():
+        held_out = sum(float(sedd_loss(model, valid_tokens)) for _ in range(4)) / 4.0
+
+    # The ratios the trained network predicts at a half-corrupted batch, against their targets.
+    with torch.no_grad():
+        check_t = torch.full((len(valid_tokens), 1), math.log(2.0) / BETA, device=device)
+        check_alpha = torch.exp(-BETA * check_t)  # alpha_t = 1/2 at this time.
+        keep = torch.rand(valid_tokens.shape, device=device) < check_alpha
+        noisy = torch.where(keep, valid_tokens,
+                            torch.randint(0, DNA_SIZE, valid_tokens.shape, device=device))
+        predicted = concrete_score(model, noisy, check_t[:, 0])  # [M, L, K] the predicted ratios.
+        exact = kernel_ratio(valid_tokens, noisy, check_alpha)  # [M, L, K] the conditional targets.
+        toward_clean = exact > 1.0  # The jumps that point back at the clean base.
+        mean_predicted = float(predicted[toward_clean].mean())  # Should sit below the target 5.
+
+    torch.manual_seed(args.seed + 1)  # Fix the jump draws, so the samples repeat.
+    samples = sample(model, args.samples, args.length, args.sample_steps)  # [B, L] over ACGT.
+    report = {
+        "target ratio at C toward A": target_ratio,
+        "score entropy at s = 1.8": errors[1.8],
+        "score entropy at s = 1.0": errors[1.0],
+        "score entropy at s = 3.0": errors[3.0],
+        "training score entropy, last 50": sum(losses[-50:]) / len(losses[-50:]),
+        "validation score entropy": held_out,
+        "mean predicted ratio toward clean": mean_predicted,
+        "conditional target toward clean": float(exact[toward_clean].mean()),
+        "GC fraction": gc_fraction(samples),
+        "GC fraction in the data": gc_fraction(train_tokens),
+        "motif fraction": motif_fraction(samples),
+        "motif fraction in the data": motif_fraction(train_tokens),
+        "distinct samples": unique_fraction(samples),
+        "letter entropy, nats": token_entropy(samples, DNA_SIZE),
+    }
+    print_report("sedd: score entropy on DNA", report, decode_dna(samples))
+    save_run(args.out, model=model, config=vars(args), report=report,
+             losses=losses, samples=decode_dna(samples))
+
+
+# %% 06. Run the script
+if __name__ == "__main__":  # Execute when launched directly.
+    main()  # Work the score entropy, train the ratios, and generate.

lecture_4/simplex.py

@@ -0,0 +1,246 @@
+"""Simplex diffusion on DNA, a belief at every position in place of a token.
+
+Every model before this one draws a token at each reverse step and discards the rest of the
+denoiser's distribution, so a position the model was unsure about looks the same as one it was
+certain about. Here the state at a position is a probability vector over the four bases, the
+corruption is a Dirichlet around the uniform-replacement kernel, and one reverse step thins the
+current belief, draws an innovation and mixes the two.
+
+Chapter 4, "From tokens to beliefs" and "Thinning, an innovation, and a mixture". Time runs from
+t = 0 clean to t = 1 noise, the survival probability is alpha_t = 1 - t, the concentration is
+c_t = gamma K / (1 - alpha_t), and the rates R_t(x, y) act on probability rows, p_dot = p R.
+Trained on the synthetic DNA of cis6270.data. Prints the worked Dirichlet parameters, the worked
+thinning weight and its Monte Carlo check, then the sample statistics.
+
+    python lecture_4/simplex.py                    # train and generate
+    python lecture_4/simplex.py --mixture-draws 10000  # a coarser check of the mixture
+    python lecture_4/simplex.py --steps 20         # a few seconds, for a first look
+"""
+
+# %% 01. Imports and configuration
+import torch  # Tensors, gradients and the Dirichlet, Beta and categorical draws.
+import torch.nn as nn  # The denoiser whose input layer is the one thing that changes.
+import torch.nn.functional as F  # Cross-entropy at every position.
+
+from cis6270.data import DNA_SIZE, decode_dna, dna_splits  # ACGT, K = 4 and the shared sequences.
+from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
+from cis6270.nets import SequenceDenoiser, count_parameters  # The transformer the belief feeds.
+from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
+
+GAMMA = 1.0  # The constant of the concentration schedule c_t = gamma K / (1 - alpha_t).
+TIME_FLOOR = 1e-3  # The concentration grows like 1/(1 - alpha_t), so the grid stops above zero.
+
+
+# %% 02. The forward marginal as a Dirichlet
+#
+# Fact 2 calls for a forward marginal in closed form, and the natural requirement is that the
+# belief fluctuate around the marginal the token models already use,
+# E[P_t | x_0] = alpha_t onehot(x_0) + (1 - alpha_t) pi, with the spread set by c_t.
+def concentration(alpha: float | torch.Tensor) -> float | torch.Tensor:
+    """c_t = gamma K / (1 - alpha_t), the schedule that makes the increment nonnegative."""
+    return GAMMA * DNA_SIZE / (1.0 - alpha)  # 8 at alpha_t = 0.5 and 20 at alpha_s = 0.8.
+
+
+def belief_mean(alpha: float, x0: int) -> torch.Tensor:
+    """E[P_t | x_0] over (A, C, G, T), the uniform-replacement kernel of the token models. [K]."""
+    mean = torch.full((DNA_SIZE,), (1.0 - alpha) / DNA_SIZE)  # [K] the share from pi.
+    mean[x0] = mean[x0] + alpha  # The clean token keeps the survival mass on top of its share.
+    return mean  # [K] a distribution, since alpha + (1 - alpha) = 1.
+
+
+def dirichlet_parameters(alpha: float, x0: int) -> torch.Tensor:
+    """c_t E[P_t | x_0], the Dirichlet parameters at one position. Returns [K]."""
+    return concentration(alpha) * belief_mean(alpha, x0)  # (5, 1, 1, 1) at alpha_t = 0.5.
+
+
+def parameter_field(x0: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
+    """The same parameters for a whole batch. x0: [B, L], alpha: [B, 1] -> [B, L, K].
+
+    Every coordinate other than the clean token has parameter gamma, constant in time, and the
+    clean coordinate has gamma (1 + K alpha_t/(1 - alpha_t)), the factor of the uniform rates.
+    """
+    off = torch.full(x0.shape + (DNA_SIZE,), GAMMA, device=x0.device)  # [B, L, K] the pi share.
+    clean = GAMMA * DNA_SIZE * alpha / (1.0 - alpha)  # [B, 1] the survival mass times c_t.
+    return off.scatter_add(2, x0[:, :, None], clean[..., None].expand_as(x0[:, :, None]))
+
+
+def corrupt(x0: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
+    """Draw P_t ~ Dir(c_t E[P_t | x_0]) at each position. x0: [B, L], alpha: [B, 1] -> [B, L, K]."""
+    return torch.distributions.Dirichlet(parameter_field(x0, alpha)).sample()  # [B, L, K] beliefs.
+
+
+# %% 03. Thinning, an innovation, and a mixture
+#
+# A Dirichlet with parameter a + b equals W V + (1 - W) U for V ~ Dir(a), U ~ Dir(b) and
+# W ~ Beta(sum a, sum b). With this concentration schedule the increment b sits entirely on the
+# clean coordinate, so the innovation U is the one-hot vector of the clean token.
+def increment(alpha_t: float, alpha_s: float, x0: int) -> torch.Tensor:
+    """b = c_s E[P_s | x_0] - c_t E[P_t | x_0], nonnegative for alpha_s > alpha_t. Returns [K]."""
+    return dirichlet_parameters(alpha_s, x0) - dirichlet_parameters(alpha_t, x0)  # (12, 0, 0, 0).
+
+
+def thinning_law(alpha_t: float, alpha_s: float) -> tuple[float, float]:
+    """The two Beta parameters, sum a = c_t and sum b = c_s - c_t, which are 8 and 12 here."""
+    return float(concentration(alpha_t)), float(concentration(alpha_s) - concentration(alpha_t))
+
+
+def reverse_step(belief: torch.Tensor, clean: torch.Tensor,
+                 alpha_t: float, alpha_s: float) -> torch.Tensor:
+    """P_s = W P_t + (1 - W) onehot(x_0), one Beta draw per position. belief: [B, L, K]."""
+    first, second = thinning_law(alpha_t, alpha_s)  # The two parameters of the thinning weight.
+    w = torch.distributions.Beta(  # [B, L] one weight per position, with mean c_t / c_s.
+        torch.tensor(first, device=belief.device),
+        torch.tensor(second, device=belief.device)).sample(belief.shape[:2])
+    innovation = F.one_hot(clean, DNA_SIZE).to(belief.dtype)  # [B, L, K] the corner it moves to.
+    return w[..., None] * belief + (1.0 - w[..., None]) * innovation  # [B, L, K] the mixture.
+
+
+# %% 04. The denoiser that reads a belief
+#
+# The only architectural change is at the input. A token is looked up in the embedding table, and
+# a belief is the same K rows of that table combined in its own proportions, so a one-hot input
+# recovers the denoiser of every other file in this lecture exactly.
+class BeliefDenoiser(nn.Module):
+    """d_theta(. | P_t, t), the shared sequence denoiser reading K numbers at every position."""
+
+    def __init__(self, length: int, width: int = 128):
+        super().__init__()
+        self.inner = SequenceDenoiser(DNA_SIZE, length, width=width)  # Embeddings and transformer.
+
+    def forward(self, beliefs: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        """beliefs: [B, L, K], t: [B] -> logits over the clean tokens [B, L, K]."""
+        inner = self.inner
+        positions = torch.arange(beliefs.shape[1], device=beliefs.device)  # [L].
+        h = (beliefs @ inner.token_embedding.weight  # [B, L, D] the belief mixes the K embeddings.
+             + inner.position_embedding(positions)[None, :, :]  # [1, L, D] where it sits.
+             + inner.time_embedding(t)[:, None, :])  # [B, 1, D] how far along the path we are.
+        return inner.head(inner.transformer(h))  # [B, L, K] every position sees every other.
+
+
+# %% 05. The loss
+#
+# Fact 5 is the cross-entropy whose optimum is the posterior over clean tokens. The reverse step
+# carries the denoiser's answer through the mixture of Section 03, whose weight already comes from
+# the concentration schedule, so the time derivative of alpha_t multiplies nothing here.
+def loss(model: nn.Module, x0: torch.Tensor) -> torch.Tensor:
+    """One Monte Carlo sample of E[-log d_theta(x_0 | P_t, t)], summed over positions."""
+    t = torch.rand(x0.shape[0], 1, device=x0.device).clamp(min=TIME_FLOOR)  # [B, 1] one time each.
+    alpha = 1.0 - t  # [B, 1] the linear schedule, alpha_0 = 1 and alpha_1 = 0.
+    beliefs = corrupt(x0, alpha)  # [B, L, K] one Dirichlet draw per position.
+    logits = model(beliefs, t[:, 0])  # [B, L, K] the denoiser reads the whole field of beliefs.
+    ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none")  # [B, L].
+    return ce.sum(-1).mean()  # Nats per sequence, with every position counted.
+
+
+# %% 06. The sampler
+#
+# Generation starts from Dir(gamma, ..., gamma), which at alpha_1 = 0 is the law the corruption
+# ends at. Each step draws one clean token from the denoiser, thins the current belief and mixes
+# the two, so the belief moves through the interior and reaches a corner at the last step.
+@torch.no_grad()  # Generation needs no gradient tracking.
+def sample(model: nn.Module, batch: int, length: int, steps: int = 64) -> torch.Tensor:
+    """Generate from beliefs scattered around pi down to clean DNA. Returns [B, L] over ACGT."""
+    device = next(model.parameters()).device
+    start = torch.full((batch, length, DNA_SIZE), GAMMA, device=device)  # [B, L, K] at alpha_1 = 0.
+    belief = torch.distributions.Dirichlet(start).sample()  # [B, L, K] scattered around pi.
+    ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device)  # The reverse grid, t down to 0.
+    for t, s in zip(ts[:-1], ts[1:]):  # One interval [s, t] at a time.
+        logits = model(belief, t.expand(batch))  # [B, L, K] the clean-token prediction.
+        clean = torch.distributions.Categorical(logits=logits).sample()  # [B, L] the innovation.
+        belief = reverse_step(belief, clean, float(1.0 - t), float(1.0 - s))  # [B, L, K].
+    return belief.argmax(-1)  # [B, L] the corner each position has reached.
+
+
+# %% 07. The report
+def main() -> None:
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.add_argument("--mixture-draws", type=int, default=400_000,
+                        help="draws in the check of the thinning representation")
+    # The belief carries the denoiser's uncertainty between steps, which needs enough reverse steps
+    # to show. A default run trains 1000 steps and measures 64 sequences, inside about a minute.
+    parser.set_defaults(steps=1000, sample_steps=64, samples=64)
+    args = parser.parse_args()
+    set_seed(args.seed)  # The same seed gives the same training batches and the same samples.
+    device = pick_device(args.device)
+    talk = (lambda *a: None) if args.quiet else print  # --quiet keeps only the report block.
+
+    # Section 02, the worked Dirichlet at the two noise levels of the chapter, with a clean A.
+    at_t = dirichlet_parameters(0.5, 0)  # [K] (5, 1, 1, 1) at the later of the two times.
+    at_s = dirichlet_parameters(0.8, 0)  # [K] (17, 1, 1, 1) at the earlier one.
+    total_t, total_s = at_t.sum(), at_s.sum()  # c_t and c_s, the totals that set every variance.
+    spread_t = (at_t * (total_t - at_t) / (total_t ** 2 * (total_t + 1))).sqrt()  # [K] exact.
+    spread_s = (at_s * (total_s - at_s) / (total_s ** 2 * (total_s + 1))).sqrt()  # [K] exact.
+    talk(f"alpha_t = 0.5       c_t {concentration(0.5):.0f}  mean "
+         f"{[round(float(v), 3) for v in belief_mean(0.5, 0)]}  parameters "
+         f"{[round(float(v), 0) for v in at_t]}  std on A {float(spread_t[0]):.3f}")
+    talk(f"alpha_s = 0.8       c_s {concentration(0.8):.0f}  mean "
+         f"{[round(float(v), 2) for v in belief_mean(0.8, 0)]}  parameters "
+         f"{[round(float(v), 0) for v in at_s]}")
+
+    # Section 03, the increment, the thinning weight, and the mean the mixture lands on.
+    step = increment(0.5, 0.8, 0)  # [K] all of the increase sits on the clean coordinate.
+    first, second = thinning_law(0.5, 0.8)  # Beta(8, 12), whose mean is 0.4.
+    mean_weight = first / (first + second)  # c_t / c_s, the average share the old belief keeps.
+    mixed = mean_weight * belief_mean(0.5, 0) + (1.0 - mean_weight) * torch.eye(DNA_SIZE)[0]  # [K].
+    talk(f"increment           {[round(float(v), 0) for v in step]}  W ~ Beta({first:.0f}, "
+         f"{second:.0f})  mean {mean_weight:.1f}  mixture mean "
+         f"{[round(float(v), 2) for v in mixed]}")
+
+    # The same mixture by simulation, against the exact standard deviations of Dir(17, 1, 1, 1).
+    thinned = torch.distributions.Dirichlet(at_t).sample((args.mixture_draws,))  # [N, K] draw V.
+    weight = torch.distributions.Beta(  # [N] the Beta draw of the representation.
+        torch.tensor(first), torch.tensor(second)).sample((args.mixture_draws,))
+    drawn = weight[:, None] * thinned + (1.0 - weight[:, None]) * torch.eye(DNA_SIZE)[0]  # [N, K].
+    talk(f"mixture check       {args.mixture_draws} draws  std "
+         f"{[round(float(v), 3) for v in drawn.std(0)]}  against the exact "
+         f"{[round(float(v), 3) for v in spread_s]}")
+
+    # Train the denoiser on the shared synthetic DNA. Reseeding here keeps the training batches
+    # independent of how many draws the check above consumed.
+    set_seed(args.seed)
+    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
+    train_tokens = splits["train_tokens"].to(device)  # [N, L] clean sequences over ACGT.
+    valid_tokens = splits["valid_tokens"].to(device)  # [M, L] held out for the loss below.
+    model = BeliefDenoiser(args.length, width=args.width).to(device)
+    talk(f"device={device} parameters={count_parameters(model):,} "
+         f"train={len(train_tokens)} length={args.length}")
+
+    def step_loss(step_index: int) -> torch.Tensor:
+        """One optimizer step draws its own batch of clean sequences."""
+        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
+        return loss(model, train_tokens[rows])  # [B, L] clean sequences into one scalar.
+
+    losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
+    with torch.no_grad():
+        held_out = sum(float(loss(model, valid_tokens)) for _ in range(4)) / 4.0
+
+    # Generate, then measure the structure the data was built with.
+    torch.manual_seed(args.seed + 1)  # Fix the Beta and categorical draws, so the samples repeat.
+    samples = sample(model, args.samples, args.length, args.sample_steps)  # [B, L] over ACGT.
+    report = {
+        "c_t at alpha_t = 0.5": float(concentration(0.5)),
+        "Dirichlet parameter on A at 0.5": float(at_t[0]),
+        "standard deviation on A at 0.5": float(spread_t[0]),
+        "Dirichlet parameter on A at 0.8": float(at_s[0]),
+        "increment on A": float(step[0]),
+        "mean thinning weight": mean_weight,
+        "simulated standard deviation on A": float(drawn.std(0)[0]),
+        "exact standard deviation on A": float(spread_s[0]),
+        "training loss, last 50": sum(losses[-50:]) / len(losses[-50:]),
+        "validation loss, nats/seq": held_out,
+        "loss, nats per token": held_out / args.length,
+        "GC fraction": gc_fraction(samples),
+        "GC fraction in the data": gc_fraction(train_tokens),
+        "motif fraction": motif_fraction(samples),
+        "motif fraction in the data": motif_fraction(train_tokens),
+        "distinct samples": unique_fraction(samples),
+        "letter entropy, nats": token_entropy(samples, DNA_SIZE),
+    }
+    print_report("simplex: beliefs in place of tokens", report, decode_dna(samples))
+    save_run(args.out, model=model, config=vars(args), report=report,
+             losses=losses, samples=decode_dna(samples))
+
+
+# %% 08. Run the script
+if __name__ == "__main__":  # Execute when launched directly.
+    main()  # Work the Dirichlet numbers, train the belief denoiser, and generate.

lecture_4/uniform.py

@@ -0,0 +1,240 @@
+"""Uniform replacement on DNA, the member of the pi family with no mask.
+
+Corruption redraws a base from the uniform distribution over ACGT, so a noisy base may be the
+original or a replacement, and no marker of the damage appears in the sequence. Every pair of
+bases is then joined by a reverse rate, the loss keeps the staying term that masking cancelled, and
+a sampler can revise a base it has already written.
+
+Chapter 4, "A rate matrix with one free distribution", "The posterior average of conditional reverse
+rates" and "Uniform replacement and the rates it produces". Time runs from t = 0 clean to t = 1
+noise, the survival probability is alpha_t = exp(-beta t), and the rates R_t(x, y) act on
+probability rows, p_dot = p R. Trained on the synthetic DNA of cis6270.data. Prints the worked
+posterior average of 2, the worked rate error, and the sample statistics.
+
+    python lecture_4/uniform.py               # train and generate
+    python lecture_4/uniform.py --steps 20    # a few seconds, for a first look
+"""
+
+# %% 01. Imports and configuration
+import math  # The worked rate errors are logarithms.
+
+import torch  # Tensors, gradients and the categorical draws.
+
+from cis6270.data import DNA_SIZE, decode_dna, dna_splits  # ACGT, K = 4 and the shared sequences.
+from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
+from cis6270.nets import SequenceDenoiser, count_parameters  # The denoiser d_theta(. | z, t).
+from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
+
+BETA = 4.0  # The constant noise rate of the chapter's worked example, so alpha_t = exp(-4t).
+TIME_FLOOR = 2e-2  # The reverse rates grow like 1/(1 - alpha_t), so sampling stops just above zero.
+
+
+# %% 02. The uniform kernel and the rates that produce it
+#
+# Under R_pi = 1 pi^T - I with pi uniform, every pair of distinct bases is joined at rate beta/K,
+# and the matrix exponential is the kernel q_t(x | x_0) = alpha_t delta(x, x_0) + (1 - alpha_t)/K.
+def survival(t: torch.Tensor | float) -> torch.Tensor | float:
+    """alpha_t = exp(-tau(t)) for the constant schedule beta(t) = BETA, so tau(t) = BETA t."""
+    return torch.exp(-BETA * t) if torch.is_tensor(t) else math.exp(-BETA * t)
+
+
+def kernel(alpha: float | torch.Tensor) -> torch.Tensor:
+    """Q(alpha) over (A, C, G, T), whose (x_0, x) entry is q_t(x | x_0). Returns [K, K]."""
+    uniform = (1.0 - alpha) / DNA_SIZE  # The share of the redrawn mass that lands on each base.
+    return alpha * torch.eye(DNA_SIZE, dtype=torch.float64) + uniform  # [K, K] rows sum to one.
+
+
+def corrupt(x0: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+    """Redraw each position from the uniform distribution with probability 1 - alpha_t."""
+    keep = torch.rand(x0.shape, device=x0.device) < survival(t)  # [B, L] survival of the original.
+    fresh = torch.randint(0, DNA_SIZE, x0.shape, device=x0.device)  # [B, L] a draw from pi.
+    return torch.where(keep, x0, fresh)  # [B, L] nothing marks which positions were redrawn.
+
+
+# %% 03. The conditional reverse rates and their three cases
+#
+# Dividing the conditional reverse rate by the noise rate beta/K leaves one factor, which is
+# 1 + K alpha_t/(1 - alpha_t) toward the clean token, its reciprocal away from it, and one between
+# two tokens that are both wrong.
+def clean_factor(alpha: float) -> float:
+    """1 + K alpha_t / (1 - alpha_t), the factor 5 of the worked example at alpha_t = 1/2."""
+    return 1.0 + DNA_SIZE * alpha / (1.0 - alpha)  # The ratio q_t(x_0 | x_0) / q_t(x | x_0).
+
+
+def conditional_rates(alpha: float, current: int, clean: int) -> torch.Tensor:
+    """R_hat_t(c, a | x_0) for every destination a, in units where the forward rate is one."""
+    Q = kernel(alpha)  # [K, K] the forward kernel at this survival probability.
+    return Q[clean] / Q[clean, current]  # [K] q_t(a | x_0) divided by q_t(c | x_0).
+
+
+# %% 04. A worked posterior average on one base
+#
+# With a clean base A at probability 0.75 and C at 0.25, an observed C at alpha_t = 1/2 has
+# posterior 0.375 and 0.625, and the two conditional rates 5 and 0.2 average to the marginal rate.
+def posterior_over_clean(alpha: float, observed: int, prior: torch.Tensor) -> torch.Tensor:
+    """p(x_0 = b | X_t = observed), Bayes' rule with the forward kernel as the likelihood."""
+    weights = prior * kernel(alpha)[:, observed]  # [K] prior times q_t(observed | b).
+    return weights / weights.sum()  # [K] normalized over the candidate clean bases.
+
+
+def marginal_reverse_rate(alpha: float, observed: int, destination: int,
+                          prior: torch.Tensor) -> float:
+    """The posterior average of the conditional reverse rates, which is the marginal reverse rate."""
+    posterior = posterior_over_clean(alpha, observed, prior)  # [K] the weights of the average.
+    targets = kernel(alpha)[:, destination] / kernel(alpha)[:, observed]  # [K] one target per clean b.
+    return float(posterior @ targets)  # The two together give the rate the sampler uses.
+
+
+# %% 05. The rate error and the staying term
+#
+# l(A, B) = A log(A/B) - A + B is the Bregman divergence the path KL divergence produces. The term
+# -A + B comes from the probability of staying put, and under uniform noise it does not cancel.
+def rate_error(target: torch.Tensor, model_rate: torch.Tensor) -> torch.Tensor:
+    """Elementwise l(A, B), nonnegative and zero exactly at B = A."""
+    logarithmic = torch.where(target > 0, target * (target.clamp(min=1e-12) / model_rate).log(),
+                              torch.zeros_like(target))  # A log(A/B), with l(0, B) = B.
+    return logarithmic - target + model_rate  # The staying term is the -A + B part.
+
+
+def ratio_table(current: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
+    """q_t(a | b) / q_t(c | b) for every clean base b and destination a. Returns [B, L, K, K]."""
+    uniform = (1.0 - alpha) / DNA_SIZE  # [B, 1] the redrawn share at this noise level.
+    Q = (alpha[..., None, None] * torch.eye(DNA_SIZE, device=current.device)
+         + uniform[..., None, None])  # [B, 1, K, K] the kernel q_t(a | b) at each sequence's time.
+    denominator = Q.expand(-1, current.shape[1], -1, -1).gather(
+        3, current[:, :, None, None].expand(-1, -1, DNA_SIZE, 1))  # [B, L, K, 1] q_t(c | b).
+    return Q / denominator  # [B, L, K, K] the only quantity the rates need from the kernel.
+
+
+def model_rates(denoiser: torch.Tensor, ratios: torch.Tensor) -> torch.Tensor:
+    """R_hat^theta_t(c, a) = (beta/K) sum_b d_theta(b) q_t(a | b) / q_t(c | b). Returns [B, L, K]."""
+    averaged = torch.einsum("nlb,nlba->nla", denoiser, ratios)  # [B, L, K] averaged over clean bases.
+    return (BETA / DNA_SIZE) * averaged  # [B, L, K] the forward rate times the posterior average.
+
+
+def conditional_targets(x0: torch.Tensor, ratios: torch.Tensor) -> torch.Tensor:
+    """R_hat_t(c, a | x_0), the same formula with the posterior collapsed onto the clean token."""
+    rows = x0[:, :, None, None].expand(-1, -1, 1, DNA_SIZE)  # [B, L, 1, K] select the row b = x_0.
+    return (BETA / DNA_SIZE) * ratios.gather(2, rows).squeeze(2)  # [B, L, K] one target per jump.
+
+
+def rate_loss(model: torch.nn.Module, x0: torch.Tensor) -> torch.Tensor:
+    """One Monte Carlo sample of L_cond, summed over destinations and positions."""
+    t = torch.rand(x0.shape[0], 1, device=x0.device).clamp(min=TIME_FLOOR)  # [B, 1] one time each.
+    alpha = survival(t)  # [B, 1] the survival probability at that time.
+    zt = corrupt(x0, t)  # [B, L] the noisy sequence, with no marker of which positions moved.
+    ratios = ratio_table(zt, alpha)  # [B, L, K, K] the kernel ratios at this noise level.
+    denoiser = model(zt, t[:, 0]).softmax(-1)  # [B, L, K] the posterior over clean tokens.
+    predicted = model_rates(denoiser, ratios).clamp(min=1e-8)  # [B, L, K] nonnegative by construction.
+    target = conditional_targets(x0, ratios)  # [B, L, K] the rate attached to this clean sequence.
+    staying = torch.nn.functional.one_hot(zt, DNA_SIZE).bool()  # [B, L, K] the destination c itself.
+    return rate_error(target, predicted).masked_fill(staying, 0.0).sum((1, 2)).mean()
+
+
+# %% 06. Sampling by tau-leaping
+#
+# Over a step of length h a position moves to a with probability h times its reverse rate, and the
+# remaining probability keeps it where it is. The step is valid while h times the total rate is one.
+@torch.no_grad()  # Generation needs no gradient tracking.
+def sample(model: torch.nn.Module, batch: int, length: int, steps: int = 64) -> torch.Tensor:
+    """Generate from uniformly random tokens down to clean DNA. Returns [B, L] over ACGT."""
+    device = next(model.parameters()).device
+    z = torch.randint(0, DNA_SIZE, (batch, length), device=device)  # [B, L] a draw from pi.
+    ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device)  # The reverse grid.
+    for t, s in zip(ts[:-1], ts[1:]):
+        time_row = t.expand(batch, 1)  # [B, 1] the shared time of this step.
+        denoiser = model(z, time_row[:, 0]).softmax(-1)  # [B, L, K] the clean-token prediction.
+        rates = model_rates(denoiser, ratio_table(z, survival(time_row)))  # [B, L, K] per destination.
+        jump = (t - s) * rates  # [B, L, K] the tau-leaping probabilities of this step.
+        jump.scatter_(2, z[:, :, None], 0.0)  # A jump to the current base is not a jump.
+        total = jump.sum(-1, keepdim=True).clamp(min=1.0)  # Rescale when the step would exceed one.
+        jump = jump / total  # [B, L, K] now a sub-probability vector at every position.
+        stay = (1.0 - jump.sum(-1, keepdim=True)).clamp(min=0.0)  # [B, L, 1] the chance of staying.
+        choice = torch.distributions.Categorical(
+            probs=torch.cat([jump, stay], dim=-1)).sample()  # [B, L] destination, or K for staying.
+        z = torch.where(choice == DNA_SIZE, z, choice.clamp(max=DNA_SIZE - 1))  # [B, L].
+    final = model(z, ts[-1].expand(batch)).softmax(-1)  # The exact posterior step down to t = 0.
+    return torch.distributions.Categorical(probs=final).sample()  # [B, L] clean DNA.
+
+
+# %% 07. The report
+def main() -> None:
+    parser = common_parser(__doc__.split("\n")[0])
+    # The rate loss needs a longer budget than the worked numbers do, so the default trains for
+    # 1000 steps and measures on 64 sequences, which keeps the run inside about a minute.
+    parser.set_defaults(steps=1000, sample_steps=64, samples=64)
+    args = parser.parse_args()
+    set_seed(args.seed)
+    device = pick_device(args.device)
+    talk = (lambda *a: None) if args.quiet else print  # --quiet keeps only the report block.
+
+    # Sections 03 and 04, the three cases and the worked posterior average on one base.
+    prior = torch.tensor([0.75, 0.25, 0.0, 0.0], dtype=torch.float64)  # Clean A at 0.75, C at 0.25.
+    posterior = posterior_over_clean(0.5, 1, prior)  # [4] after observing a C at alpha_t = 1/2.
+    targets = conditional_rates(0.5, current=1, clean=0)  # [4] the rates attached to a clean A.
+    average = marginal_reverse_rate(0.5, 1, 0, prior)  # The posterior average, which comes out 2.
+    marginal = prior @ kernel(0.5)  # [4] p_t over the four bases, which is (0.5, 0.25, 0.125, 0.125).
+    talk(f"kernel at alpha=0.5 clean {float(kernel(0.5)[0, 0]):.3f}  other "
+         f"{float(kernel(0.5)[0, 1]):.3f}  clean factor {clean_factor(0.5):.0f}")
+    talk(f"posterior after C   {[round(float(v), 3) for v in posterior]}  conditional rates "
+         f"{float(targets[0]):.0f} under a clean A and "
+         f"{float(conditional_rates(0.5, 1, 1)[0]):.1f} under a clean C")
+    talk(f"posterior average   {average:.1f}  marginals {[round(float(v), 3) for v in marginal]}  "
+         f"ratio {float(marginal[0] / marginal[1]):.1f}  to G and T "
+         f"{marginal_reverse_rate(0.5, 1, 2, prior):.1f}")
+
+    # Section 05, the rate error on one base, where the staying term keeps the sum from collapsing.
+    worked_denoiser = torch.tensor([[[0.5, 0.3, 0.1, 0.1]]], dtype=torch.float64)  # [1, 1, 4].
+    worked_current = torch.tensor([[1]])  # The current base is C.
+    worked_alpha = torch.tensor([[0.5]], dtype=torch.float64)  # alpha_t = 1/2.
+    worked_ratios = ratio_table(worked_current, worked_alpha)  # [1, 1, K, K] at alpha_t = 1/2.
+    predicted = model_rates(worked_denoiser, worked_ratios)[0, 0] / (BETA / DNA_SIZE)
+    target = conditional_rates(0.5, current=1, clean=0)  # The clean base is A.
+    error = float(rate_error(target[[0, 2, 3]], predicted[[0, 2, 3]]).sum())  # Skip the current base.
+    talk(f"model rates from d  {[round(float(v), 2) for v in predicted]}  targets "
+         f"{[round(float(v), 1) for v in target]}  rate error {error:.3f}")
+
+    # Train the denoiser against the conditional rate targets.
+    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
+    train_tokens = splits["train_tokens"].to(device)  # [N, L] clean sequences over ACGT.
+    valid_tokens = splits["valid_tokens"].to(device)  # [M, L] held out for the loss below.
+    model = SequenceDenoiser(DNA_SIZE, args.length, width=args.width).to(device)
+    talk(f"device={device} parameters={count_parameters(model):,} train={len(train_tokens)}")
+
+    def step_loss(step: int) -> torch.Tensor:
+        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
+        return rate_loss(model, train_tokens[rows])  # One clean batch into one scalar.
+
+    losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
+    with torch.no_grad():
+        held_out = sum(float(rate_loss(model, valid_tokens)) for _ in range(4)) / 4.0
+
+    torch.manual_seed(args.seed + 1)  # Fix the jump draws, so the samples repeat.
+    samples = sample(model, args.samples, args.length, args.sample_steps)  # [B, L] over ACGT.
+    report = {
+        "q_t(x_0 | x_0) at alpha=0.5": float(kernel(0.5)[0, 0]),
+        "clean factor at alpha=0.5": clean_factor(0.5),
+        "posterior on A after a C": float(posterior[0]),
+        "conditional rate under A": float(targets[0]),
+        "conditional rate under C": float(conditional_rates(0.5, 1, 1)[0]),
+        "posterior average rate": average,
+        "rate from the marginals": float(marginal[0] / marginal[1]),
+        "worked model rate to A": float(predicted[0]),
+        "worked rate error": error,
+        "training rate loss, last 50": sum(losses[-50:]) / len(losses[-50:]),
+        "validation rate loss": held_out,
+        "GC fraction": gc_fraction(samples),
+        "GC fraction in the data": gc_fraction(train_tokens),
+        "motif fraction": motif_fraction(samples),
+        "motif fraction in the data": motif_fraction(train_tokens),
+        "distinct samples": unique_fraction(samples),
+        "letter entropy, nats": token_entropy(samples, DNA_SIZE),
+    }
+    print_report("uniform: replacement diffusion on DNA", report, decode_dna(samples))
+    save_run(args.out, model=model, config=vars(args), report=report,
+             losses=losses, samples=decode_dna(samples))
+
+
+# %% 08. Run the script
+if __name__ == "__main__":  # Execute when launched directly.
+    main()  # Work the posterior average, train against the rate loss, and generate.

lecture_5/README.md

@@ -0,0 +1,97 @@
+# Lecture 5: Discrete Flow Matching (Chapter 5)
+
+Chapter 4 fixed a corruption process and solved for its reverse. This lecture reverses that order.
+We pick the probability path first, write down a velocity that generates it, and read the sampler
+off the velocity. On tokens the velocity is a rate matrix read from the destination's side, the
+master equation is the condition under which it generates a path, and the mixture path
+`(1 - kappa_t) delta_{x_0} + kappa_t delta_{x_1}` at each position has a conditional velocity that a
+posterior average turns into the schedule factor times a denoiser. Cross-entropy on the data token
+recovers that denoiser. Everything after that is a choice made inside the same construction: a
+schedule, a corrector, a kinetic-optimal path, edits that change the length, a relaxation of the
+token into a point of the simplex or of a sphere, a guidance tilt, or a different coupling.
+
+Two conventions hold in every file and are stated at the top of each. Time runs from the source at
+`t = 0` to the data at `t = 1`, so `kappa_t` here equals the survival probability `alpha_{1-t}` of
+Chapter 4. The probability velocity `u_t(y, x)` is the same number Chapter 4 writes as the rate
+`R_t(x, y)`, with the two arguments exchanged.
+
+## The files
+
+Everything runs on the synthetic DNA of `cis6270.data`: 512 sequences of length 16 over A, C, G, T,
+with a GC-rich background and the motif ACGT planted in three quarters of them at one of three
+positions. Those sequences have GC fraction 0.577, motif fraction 0.777 and token entropy 1.374 of
+`log 4 = 1.386`, which are the numbers a trained sampler is measured against.
+
+| File | Command | What it does |
+| --- | --- | --- |
+| `dfm.py` | `python lecture_5/dfm.py` | The coupling, the mixture path, the conditional velocity, the posterior average, the cross-entropy loss and the two-part Euler rule. Trains a denoiser and generates DNA from either source, `--source uniform` or `--source mask`. |
+| `velocity.py` | `python lecture_5/velocity.py` | The mathematics with nothing learned: a jump rate written as a velocity, the two rules, the master equation, the worked DNA velocity, survival and hazard, and the continuous and discrete velocities side by side. |
+| `corrector.py` | `python lecture_5/corrector.py` | Velocities with zero net flux, the three-state example, the corrector for the mask source, the jump counts at three corrector strengths, and the parallel-step error. |
+| `kinetic_optimal.py` | `python lecture_5/kinetic_optimal.py` | The flux and its energy, the velocity of least energy, the kinetic optimality of the mixture path, the transition distance on DNA, the worked metric path, and the energy of the linear and geodesic schedules. |
+| `edit_flows.py` | `python lecture_5/edit_flows.py` | Insertions, deletions and substitutions; alignments and the column-wise path; the posterior over two alignments; the rate loss; one sampler step; and the length statistics of a run on DNA pairs of different lengths. |
+| `dirichlet.py` | `python lecture_5/dirichlet.py` | The Dirichlet family as a conditional path, the speed from the flux across a cut, the two-letter closed form, the worked speeds on a DNA base, and the posterior average. Trains a denoiser on the simplex and integrates to `t = 8`. |
+| `fisher.py` | `python lecture_5/fisher.py` | The Fisher information as a metric, the square-root map to the sphere, geodesics, the tangent projection and the regression loss. Trains a velocity field and samples with exponential-map steps. |
+| `gumbel.py` | `python lecture_5/gumbel.py` | The Gumbel-max trick, the cooling temperature, the annealed path, the replicator equation, the inference-time velocity, and training against the marginal velocity. |
+| `guidance.py` | `python lecture_5/guidance.py` | Guidance as a tilt of the rates, the h-transform, the reachable set, the guidance strength, the classifier-free form, the Taylor expansion of the log ratio, and the guided sampler. Steers toward the GC label with `--gamma`. |
+| `multi_objective.py` | `python lecture_5/multi_objective.py` | The Pareto front, MOG-DFM's rank, direction and cone with the controller that widens it, and pCoMole's worst-objective utility with a rollout of the base flow on DNA. |
+| `rectify.py` | `python lecture_5/rectify.py` | Total correlation on two bases, ReDi's recoupling rounds run on the exact sampler, and AReUReDi's Tchebycheff tilt with its Metropolis acceptance step. |
+
+Every file takes the shared flags of `cis6270.runner`. `--steps 20 --quiet` runs each one in a few
+seconds, and `--steps` controls the training steps in the five files that train, the grid in the
+five that integrate, and the number of rollouts in `multi_objective.py`.
+
+## The numbers a default run produces
+
+Seed 0, one core. The five files that train take under half a minute each, and the other six take
+one to seven seconds. The worked numbers are exact and repeat on any machine; the trained numbers
+move a little with the hardware's floating point.
+
+| File | Headline numbers |
+| --- | --- |
+| `dfm.py` | loss 0.876 nats, GC fraction 0.578, motif fraction 0.672, token entropy 1.374 of `log 4 = 1.386` |
+| `velocity.py` | `p_dot = (-0.1, -0.2, 0.7, -0.4)`, Euler step `(0.09, 0.18, 0.37, 0.36)`, hazard 2 at `t = 0.5` and 10 at `t = 0.9`, Euler error 0.0002 against the closed form |
+| `corrector.py` | masked fraction 0.500 at `t = 0.5` for every `eta`, jumps per position 1.000, 2.991 and 10.934 at `eta = 0, 2, 10`, pair survival 0.133 measured against the predicted 0.135 at `eta = 2` |
+| `kinetic_optimal.py` | rate out of C 5.780, `p_dot = (-0.3997, -0.4774, 1.3546, -0.4774)` from the velocity and from the path, energy 27.63 and 41.45 for the linear schedule, `pi^2 = 9.8696` for the geodesic one |
+| `edit_flows.py` | averaged edit rates 0.667, 1.333 and 0.667 for a total of 2.667, example loss 3.446, untouched 0.8452 against the first-order 0.8350, DNA length 13.945 at `t = 0.5` against the predicted 13.941 |
+| `dirichlet.py` | speeds 0.232, 0.199, 0.159, 0.104 at `t = 1`, continuity residual `2.5e-08`, loss 1.027 nats, GC fraction 0.532, motif fraction 0.547 |
+| `fisher.py` | Fisher-Rao distance 1.485 against the Euclidean 0.141 and 0.013 on the two moves, arc 0.3713 from the midpoint, loss 0.278, GC fraction 0.576, motif fraction 0.203 |
+| `gumbel.py` | path `(0.0528, 0.0214, 0.8781, 0.0477)` at `t = 1`, replicator velocity `(-0.136, -0.217, 0.500, -0.147)`, loss 0.672 nats, GC fraction 0.572, motif fraction 0.516 |
+| `guidance.py` | guided destination `(0.029, 0.203, 0.478, 0.290)`, exit rate 2.50 at the exact `h`, GC fraction 0.531 unguided, 0.657 at `gamma = 2` and 0.698 at `gamma = 5` |
+| `multi_objective.py` | combined score `(-1.511, 1.541, 2.297, -2.327)`, cone angles 47.7, 37.9 and 107.7 degrees, selection `(0.371, 0.629, 0)` after the chapter's rollouts, and on DNA a base selection of `(0.417, 0.583)` that the rollout value turns into `(0.962, 0.038)` |
+| `rectify.py` | total correlation 1.386, 1.101, 0.843, 0.696 and 0.616 over rounds 0 to 4 against the limit 0.520; one-step success 25.0%, 40.3%, 52.5%, 58.0% and 60.4% against the limit 62.5% |
+
+Four thousand training steps, about two minutes a file, buy a sampler that reaches the planted
+motif more often: `dfm.py` goes from 0.672 to 0.734 against the data's 0.777, and `fisher.py` from
+0.203 to 0.453. Nothing else moves much.
+
+`tests/test_lecture_5.py` asserts the chapter's worked numbers against these files, including the
+posterior 0.625 and 0.375 at a masked position, the entropy 0.6616 that bounds the loss there, the
+velocity `(0.2, 1.2, -1.6, 0.2)` out of G, the metric-path rates 0.898 and 4.882 out of C, the
+Dirichlet speeds, the Fisher-Rao arc, the Gumbel-softmax path at four times, the guided destination
+distribution, the Pareto front, and the ReDi rounds.
+
+## What is simplified
+
+- The corrector simulation uses 40,000 positions where the chapter uses 200,000. Pass
+  `--positions 200000` to reproduce the chapter's figures exactly.
+- The parallel-step section measures the survival probability of a jointly revealed pair, which is
+  the quantity the chapter computes in closed form. At a fine grid the end-to-end error of the
+  corrector is not lower than at `eta = 0`, because the admissible step shrinks with `eta`.
+- `edit_flows.py` runs the column-wise path on one sampled alignment per pair, which the chapter's
+  marginalization theorem permits, and trains nothing. The edit intensities of the sampler section
+  are the chapter's own numbers, not a network's.
+- `dirichlet.py` evaluates the derivative of the regularized incomplete beta function by a central
+  difference on `scipy.special.betainc`, in double precision whatever the ambient dtype, since the
+  two values agree to six decimals. The continuity check differentiates the flux with a five-point
+  stencil.
+- `rectify.py` runs ReDi with the exact marginal velocity of the current coupling on a
+  two-position state space, so no network is trained. The AReUReDi section works one Metropolis
+  iteration, not a full refinement chain.
+- `guidance.py` trains the denoiser and the property model under one optimizer, and the property is
+  the GC label of `cis6270.data`, not a biological measurement.
+- `multi_objective.py` trains nothing. Its two DNA objectives are the GC fraction and the AT
+  fraction of a finished sequence, its feasible set is the sequences carrying the planted motif,
+  and the base flow a rollout continues is the mask-source sampler with the data's letter frequency
+  as its denoiser.
+- The relaxed paths replace the token lookup of `cis6270.nets.SequenceDenoiser` with a linear map on
+  the relaxed token, and leave the rest of the network untouched.

lecture_5/corrector.py

@@ -0,0 +1,251 @@
+"""Correctors: extra motion that leaves the probability path where it is.
+
+The master equation is linear in the velocity, so a second velocity with no net flux into any state
+can be added to the first without moving the marginals. Detailed balance is the easy way to arrange
+that, and the mixture path already contains such a term, the difference between the forward velocity
+and the same path run backward. Adding it with weight eta makes a position unmask more often and
+remask in between, which gives the sampler more chances to replace a token the rest of the sequence
+has ruled out.
+
+Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
+u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
+
+Chapter 5, Sections 5.13 and 5.14. The denoiser of the long simulation is the letter frequency of the
+synthetic DNA of ``cis6270.data``. Prints the zero-flux check, the per-step unmasking table, the jump
+counts at three corrector strengths, the parallel-step error at 64 steps, and the chance a jointly
+revealed pair survives to t = 1 untouched.
+
+    python lecture_5/corrector.py                     # 600 Euler steps on 40,000 positions
+    python lecture_5/corrector.py --eta 10            # one corrector strength in the step table
+    python lecture_5/corrector.py --steps 20 --quiet  # a coarse grid, visibly biased jump counts
+"""
+
+# %% 01. Imports and configuration
+import math
+
+import torch
+
+from cis6270.data import DNA, DNA_SIZE, dna_splits
+from cis6270.runner import common_parser, print_report, save_run, set_seed
+
+EXACT = torch.double  # The worked arithmetic keeps full precision whatever the default dtype.
+MASK = DNA_SIZE  # Index 4 is the mask symbol m, outside the four-letter alphabet.
+
+
+# %% 02. Velocities that move samples without moving the path
+#
+# A velocity v_t whose net flux into every state is zero can be added to a velocity that generates
+# p_t, and the sum generates the same path. Detailed balance, equal flux in both directions along
+# every edge, is sufficient. Reversing every term turns the net flux into p_t(y) times a column of
+# rates, and those rates sum to zero.
+
+
+def net_flux(v, p):
+    """sum_x v_t(y, x) p_t(x) at every state y. v: [K, K], p: [K] -> [K]."""
+    return v @ p  # Zero at every entry when v moves no probability.
+
+
+def detailed_balance_velocity(p, chosen):
+    """Build v_t from chosen forward rates, with each reverse rate fixed by detailed balance.
+
+    ``chosen`` maps an edge (x, y) to the rate v_t(y, x). The reverse rate is then
+    v_t(x, y) = v_t(y, x) p_t(x) / p_t(y), which equalizes the flux along that edge.
+    """
+    v = torch.zeros(len(p), len(p), dtype=p.dtype)  # [K, K], destination indexing the row.
+    for (x, y), rate in chosen.items():
+        v[y, x] = rate  # The rate we pick, out of x toward y.
+        v[x, y] = rate * p[x] / p[y]  # The reverse rate detailed balance forces.
+    return v - torch.diag(v.sum(dim=0))  # Each column now sums to zero.
+
+
+# %% 03. A zero-flux velocity on three states
+#
+# On the states A, C, G with p_t = (0.2, 0.3, 0.5) we can pick the rate 3 from A to C and the rate 1
+# from A to G, and detailed balance returns 2 and 0.4. Adding this to a velocity that sends every
+# state to G at rate 1 raises the total rate of leaving A from 1 to 5 while p_dot stays put.
+
+
+def three_state_example():
+    """Section 5.13.2, returning the two velocities, the path and the two rates of change."""
+    p = torch.tensor([0.2, 0.3, 0.5], dtype=EXACT)  # [3] over A, C, G.
+    v = detailed_balance_velocity(p, {(0, 1): 3.0, (0, 2): 1.0})  # A to C at 3, A to G at 1.
+    u = torch.zeros(3, 3, dtype=EXACT)  # The velocity that does move the path.
+    for x in range(2):
+        u[2, x], u[x, x] = 1.0, -1.0  # A and C jump to G at rate 1; G never leaves.
+    return p, u, v, net_flux(u, p), net_flux(u + v, p)  # The last two agree.
+
+
+# %% 04. The corrector for the mask source
+#
+# With the linear schedule, the weights alpha_t = 1 + eta t and beta_t = eta t, and the mask source,
+# the backward velocity is the point mass on the mask and the combination collapses to two brackets.
+# A masked position unmasks at total rate (1 + eta t)/(1 - t), split over tokens by the denoiser, and
+# a visible position remasks at rate eta.
+
+
+def corrector_velocity(p1, x_i, t, eta):
+    """bar_u^i_t(., x) over the five symbols A, C, G, T, m. p1: [4] -> [5]."""
+    current = torch.zeros(DNA_SIZE + 1, dtype=EXACT)
+    current[x_i] = 1.0  # delta on the token the position holds, which may be the mask.
+    denoiser = torch.cat([p1.to(EXACT), torch.zeros(1, dtype=EXACT)])  # The data token is never m.
+    to_mask = torch.zeros(DNA_SIZE + 1, dtype=EXACT)
+    to_mask[MASK] = 1.0  # delta_m, the destination of a remasking.
+    forward = (1.0 + eta * t) / (1.0 - t) * (denoiser - current)  # The first bracket.
+    return forward + eta * (to_mask - current)  # The second bracket, the circulation back to m.
+
+
+def admissible_step(t, eta):
+    """The largest Euler step with a redraw probability at most one, h <= (1 - t)/(1 + eta t)."""
+    return (1.0 - t) / (1.0 + eta * t)  # 0.1, 0.0357 and 0.01 at t = 0.9 for eta = 0, 2, 10.
+
+
+# %% 05. Counting the jumps
+#
+# Writing M(t) for the probability that the position is masked, the two flows cancel so that
+# M_dot(t) = -1 for every eta, and M(t) = 1 - t whatever the corrector strength. Integrating the
+# unmasking rate against M and the remasking rate against 1 - M gives 1 + eta/2 unmaskings and eta/2
+# remaskings, for 1 + eta jumps per position.
+
+
+def masked_fraction_derivative(t, eta, masked=None):
+    """M_dot(t) = -M (1 + eta t)/(1 - t) + (1 - M) eta, which equals -1 at M = 1 - t."""
+    masked = 1.0 - t if masked is None else masked  # The solution we expect.
+    return -masked * (1.0 + eta * t) / (1.0 - t) + (1.0 - masked) * eta
+
+
+def simulate_jumps(p1, eta, positions, steps, generator=None):
+    """Run ``positions`` independent positions from the mask and count their jumps.
+
+    Returns the fraction still masked at t = 0.5 and the mean number of jumps per position.
+    """
+    state = torch.full((positions,), MASK, dtype=torch.long)  # Every position masked at t = 0.
+    jumps = torch.zeros(positions, dtype=torch.int32)  # One counter per position.
+    h, halfway = 1.0 / steps, None  # The increment, and the masked fraction recorded at t = 0.5.
+    for k in range(steps):
+        t = k * h
+        masked = state == MASK  # [P] which positions are still at the mask symbol.
+        if halfway is None and t >= 0.5:
+            halfway = float(masked.to(torch.get_default_dtype()).mean())  # M(0.5), which the chapter puts at 1 - t.
+        up = min(1.0, h * (1.0 + eta * t) / max(1.0 - t, 1e-12))  # Unmasking probability this step.
+        down = h * eta  # Remasking probability this step.
+        coin = torch.rand(positions, generator=generator)
+        unmask = masked & (coin < up)  # Masked positions that reveal a token.
+        remask = (~masked) & (coin < down)  # Visible positions that return to the mask.
+        revealed = int(unmask.sum())  # Draw a token only where one is needed.
+        if revealed:
+            tokens = torch.multinomial(p1, revealed, replacement=True, generator=generator)
+            state = state.masked_scatter(unmask, tokens)  # Split over tokens by the denoiser.
+        state = state.masked_fill(remask, MASK)  # The circulation back to the mask symbol.
+        jumps += (unmask | remask).to(torch.int32)  # Every reveal and remasking counts as a jump.
+    return halfway, float(jumps.double().mean())  # Near 1 - t = 0.5, and near 1 + eta.
+
+
+# %% 06. The parallel-step error
+#
+# Let the data be uniform over AA and TT and the source be all masks. With both positions masked the
+# denoiser at each is (0.5, 0.5), and a step that reveals both at once draws them independently, so
+# half of those draws land outside the data. With eta > 0 a pair revealed jointly is remasked at rate
+# 2 eta and redrawn with the other position visible, where the denoiser is a point mass.
+
+
+def joint_reveal_rate(steps):
+    """The probability that the two reveal times land in the same cell of a uniform grid, h."""
+    return 1.0 / steps  # Sum over cells of h squared, which is 1/64 at 64 steps.
+
+
+def pair_survival(eta, tau=0.5):
+    """exp(-2 eta (1 - tau)), the chance a pair revealed at tau is untouched at t = 1."""
+    return math.exp(-2.0 * eta * (1.0 - tau))  # 0.135 at eta = 2 and 4.54e-5 at eta = 10.
+
+
+def simulate_survival(eta, samples, steps, tau=0.5, generator=None):
+    """The fraction of pairs revealed at ``tau`` that neither position remasks before t = 1.
+
+    Each of the two positions returns to the mask at rate eta, so the first remasking of either
+    happens at rate 2 eta, and a pair that is remasked is redrawn with the other position visible.
+    """
+    remaining = max(1, int(steps * (1.0 - tau)))  # Grid cells left between tau and t = 1.
+    h = (1.0 - tau) / remaining  # The increment, matching the grid the sampler runs on.
+    untouched = torch.ones(samples, dtype=torch.bool)  # Pairs no remasking has reached.
+    for _ in range(remaining):
+        coin = torch.rand(samples, 2, generator=generator)  # One clock per position.
+        untouched &= ~(coin < h * eta).any(dim=1)  # Either position remasking ends the survival.
+    return float(untouched.to(torch.get_default_dtype()).mean())  # Approaches exp(-2 eta (1 - tau)).
+
+
+def simulate_pair(eta, samples, steps, generator=None):
+    """The AA-or-TT two-position problem under the corrector. Returns the fraction outside the data."""
+    state = torch.full((samples, 2), MASK, dtype=torch.long)  # Both positions masked at t = 0.
+    h = 1.0 / steps
+    for k in range(steps):
+        t, last = k * h, k == steps - 1  # The last step reveals whatever is left.
+        masked = state == MASK  # [S, 2] read before either update, so the two are simultaneous.
+        other = state.flip(1)  # What the opposite position holds, which the denoiser reads.
+        guess = torch.randint(0, 2, state.shape, generator=generator)  # The (0.5, 0.5) denoiser.
+        drawn = torch.where(other != MASK, other, guess)  # A point mass once the partner is visible.
+        up = 1.0 if last else min(1.0, h * (1.0 + eta * t) / max(1.0 - t, 1e-12))
+        down = 0.0 if last else h * eta  # No remasking on the step that ends the run.
+        reveal = masked & (torch.rand(state.shape, generator=generator) < up)
+        remask = (~masked) & (torch.rand(state.shape, generator=generator) < down)
+        state = torch.where(reveal, drawn, state)  # Unmask, in parallel at both positions.
+        state = torch.where(remask, torch.full_like(state, MASK), state)  # The circulation back to m.
+    return float((state[:, 0] != state[:, 1]).double().mean())  # Outside {AA, TT}.
+
+
+# %% 07. The report
+def main():
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.add_argument("--eta", type=float, default=2.0, help="corrector strength in the step table")
+    parser.add_argument("--positions", type=int, default=40_000, help="positions in the simulation")
+    parser.add_argument("--pairs", type=int, default=20_000, help="sequences in the pair simulation")
+    args = parser.parse_args()
+    set_seed(args.seed)
+
+    p, u, v, with_u, with_both = three_state_example()
+    report = {
+        "zero_flux": [round(float(f), 6) for f in net_flux(v, p)],  # Zero at A, C and G.
+        "p_dot_u": [round(float(f), 4) for f in with_u],  # (-0.2, -0.3, 0.5).
+        "p_dot_u_plus_v": [round(float(f), 4) for f in with_both],  # The same three numbers.
+        "rate_out_of_A_u": round(float(-u[0, 0]), 4),  # 1 under the velocity alone.
+        "rate_out_of_A_both": round(float(-(u + v)[0, 0]), 4),  # 5 once the corrector is added.
+    }
+
+    # Section 5.14.3: one masked position at t = 0.4 with the chapter's denoiser and h = 0.01.
+    p1 = torch.tensor([0.1, 0.6, 0.2, 0.1])  # [4] over A, C, G, T.
+    bar_u = corrector_velocity(p1, MASK, t=0.4, eta=args.eta)  # [5] rates out of the mask.
+    seen = torch.zeros(DNA_SIZE)
+    seen[1] = 1.0  # At a visible position the true denoiser is the point mass on the visible token.
+    report["unmask_rate"] = round(float(-bar_u[MASK]), 4)  # (1 + eta t)/(1 - t); 3.0 at eta = 2.
+    report["remask_rate"] = round(float(corrector_velocity(seen, 1, 0.4, args.eta)[MASK]), 4)
+    report["step_probabilities"] = [round(float(0.01 * r), 5) for r in bar_u[:DNA_SIZE]]
+    report["stay_masked"] = round(float(1.0 + 0.01 * bar_u[MASK]), 5)  # 0.97 at eta = 2.
+    report["remask_probability"] = round(0.01 * args.eta, 4)  # eta h, which is 0.02 at eta = 2.
+    report["masked_derivative"] = round(masked_fraction_derivative(0.4, args.eta), 6)  # -1.
+    report["admissible_step_at_0.9"] = round(admissible_step(0.9, args.eta), 4)  # 0.0357 at eta = 2.
+
+    # The long simulation, with the letter frequency of the synthetic DNA as the denoiser.
+    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
+    counts = torch.bincount(splits["train_tokens"].flatten(), minlength=DNA_SIZE).double()
+    frequency = counts / counts.sum()  # [4] the GC-rich background the lecture shares.
+    for eta in (0.0, 2.0, 10.0):
+        half, mean = simulate_jumps(frequency, eta, args.positions, args.steps)
+        report[f"masked_at_0.5_eta_{eta:g}"] = round(half, 4)  # 1 - t = 0.5 at every eta.
+        report[f"jumps_eta_{eta:g}"] = round(mean, 3)  # 1 + eta.
+
+    # Section 5.14.4: the two positions that have to agree.
+    report["joint_reveal_at_64"] = round(joint_reveal_rate(64), 4)  # 0.0156.
+    report["error_at_64_predicted"] = round(0.5 * joint_reveal_rate(64), 4)  # Half of those are wrong.
+    report["error_at_64_measured"] = f"{simulate_pair(0.0, args.pairs, 64):.5f}"  # Near 0.0078.
+    report["error_at_grid_measured"] = f"{simulate_pair(0.0, args.pairs, args.steps):.5f}"
+    for eta in (0.0, 2.0, 10.0):
+        predicted = pair_survival(eta)  # exp(-2 eta (1 - tau)) at the midpoint tau = 0.5.
+        report[f"survival_eta_{eta:g}"] = f"{predicted:.3e}"
+        if predicted * args.pairs >= 20:  # Below that, the simulation has nothing to resolve.
+            measured = simulate_survival(eta, args.pairs, args.steps)  # The same number, simulated.
+            report[f"survival_eta_{eta:g}_measured"] = f"{measured:.3e}"
+    print_report("correctors", report)
+    save_run(args.out, config=vars(args), report=report)
+
+
+if __name__ == "__main__":
+    main()

lecture_5/dfm.py

@@ -0,0 +1,237 @@
+"""Discrete flow matching, the whole construction in one file.
+
+We pick the probability path first, write down a velocity that generates it, and read the sampler off
+the velocity. Each position follows the mixture path between a source token and a data token, its
+conditional velocity is a schedule factor times the difference of two point masses, and the posterior
+average of those conditional velocities is the same factor times a denoiser. Cross-entropy on the
+data token recovers that denoiser, and one Euler step redraws a position from it with probability
+lambda_t. Chapter 4 fixed the corruption and solved for its reverse; here the corruption is a
+consequence of the path we chose.
+
+Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
+u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
+
+Chapter 5, Sections 5.5 to 5.11. Trained on the synthetic DNA of ``cis6270.data``. Prints the
+chapter's worked numbers, the training loss and the measurements of a generated batch.
+
+    python lecture_5/dfm.py                     # train and generate from the uniform source
+    python lecture_5/dfm.py --source mask       # the all-mask source instead
+    python lecture_5/dfm.py --steps 20 --quiet  # a few seconds, for a first look
+"""
+
+# %% 01. Imports and configuration
+
+import torch
+import torch.nn.functional as F
+from torch.distributions import Categorical
+
+from cis6270.data import DNA_SIZE, decode_dna, dna_splits, encode_dna
+from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
+from cis6270.nets import SequenceDenoiser, count_parameters
+from cis6270.runner import common_parser, print_report, save_run, set_seed, train
+
+
+# %% 02. The class of Section 5.11, with both sources
+#
+# The constructor holds the vocabulary size V and the schedule kappa_t with its derivative
+# kappa_dot_t, passed as functions of a time tensor. The defaults are the linear schedule
+# kappa_t = t and kappa_dot_t = 1. A mask source returns the index of the mask symbol at every
+# position, with the network's vocabulary raised by one to make room for it.
+
+
+class DiscreteFlowMatching:
+    """The source, the mixture path, the cross-entropy loss and one sampling step."""
+
+    def __init__(self, V, kappa=lambda t: t, kappa_dot=lambda t: torch.ones_like(t), mask=None):
+        self.V, self.k, self.kd = V, kappa, kappa_dot  # Alphabet size and the schedule.
+        self.mask = mask  # Mask token id under the mask source, None under the uniform source.
+
+    def source(self, shape, device="cpu"):
+        if self.mask is None:
+            return torch.randint(0, self.V, shape, device=device)  # Uniform source, [B, L].
+        return torch.full(shape, self.mask, dtype=torch.long, device=device)  # All-mask source.
+
+    def corrupt(self, x1, t):
+        """One draw from p_t(. | x_0, x_1) by the coin flip of Section 5.5.5. x1: [B, L], t: [B]."""
+        x0 = self.source(x1.shape, x1.device)  # The other end of the pair; [B, L].
+        heads = torch.rand(x1.shape, device=x1.device) < self.k(t)[:, None]  # Heads with prob kappa_t.
+        return torch.where(heads, x1, x0)  # The data token on heads and the source token on tails.
+
+
+# %% 03. The conditional velocity
+#
+# Differentiating the conditional path at one position gives kappa_dot_t times the difference of the
+# two indicators, and dividing that flux by the probability 1 - kappa_t still sitting on the source
+# token turns it into a rate out of the token the position currently holds.
+
+
+def conditional_velocity(z, x1, t, V=DNA_SIZE):
+    """u_t(a, z | x_1) = kappa_dot/(1 - kappa) [delta_{x_1}(a) - delta_z(a)]. Returns [V]."""
+    factor = 1.0 / (1.0 - t)  # Schedule factor on the linear schedule, kappa_dot_t = 1.
+    return factor * (F.one_hot(torch.tensor(x1), V) - F.one_hot(torch.tensor(z), V)).double()
+
+
+# %% 04. The posterior average, collapsed onto the data token
+#
+# The conditional velocity depends on the pair only through the data token at this position, so every
+# pair sharing that token contributes the same term and the sum over pairs collapses onto the four
+# tokens. What is left is the denoiser p_{1|t}(. | x), the only quantity a network has to learn.
+
+
+def marginal_velocity(p1, x_i, t, V=DNA_SIZE):
+    """u^i_t(., x) = kappa_dot/(1 - kappa) [p_{1|t}(. | x) - delta_{x^i}]. p1: [V] -> [V]."""
+    factor = 1.0 / (1.0 - t)  # Same schedule factor, now multiplying a distribution.
+    return factor * (p1 - F.one_hot(torch.tensor(x_i), V).double())  # Signed, sums to zero.
+
+
+def mask_posterior(data, weights, visible, kappa_t):
+    """The posterior over the data token at the masked position, computed by enumeration.
+
+    ``data`` is a list of candidate sequences and ``weights`` their probabilities under q. Position 0
+    is visible and holds ``visible``; position 1 is masked. The time factor is common to every
+    candidate and cancels, which is Proposition 6 of Discrete Flow Matching.
+    """
+    masked, seen = 1, 1  # One masked position and one visible position in this example.
+    factor = (1.0 - kappa_t) ** masked * kappa_t ** seen  # Common to every surviving candidate.
+    unnormalized = torch.zeros(DNA_SIZE, dtype=torch.double)  # One entry per base at position 1.
+    for sequence, weight in zip(data, weights):
+        if sequence[0] == visible:  # A candidate disagreeing with the visible token is excluded.
+            unnormalized[encode_dna([sequence])[0, 1]] += factor * weight
+    return unnormalized / unnormalized.sum()  # [V] the denoiser at the masked position.
+
+
+# %% 05. The cross-entropy loss and its minimizer
+#
+# The loss scores the true data token at every position, with no time weight and no rate in it, so the
+# sampler is chosen after training. Averaging over the data token given X_t splits the expected loss
+# into the entropy of the true denoiser and a KL divergence from it to the network, and only the
+# divergence depends on theta.
+
+
+def dfm_loss(logits, x1):
+    """Cross-entropy of Section 5.9, with the class dimension second. logits: [B, L, V]."""
+    return F.cross_entropy(logits.transpose(1, 2), x1)  # Mean over the batch and over positions.
+
+
+def entropy_and_gap(posterior, prediction):
+    """The minimum of the loss at one position and the excess a wrong prediction pays, in nats."""
+    entropy = -(posterior * posterior.log()).sum()  # H(p_{1|t}), the floor no network beats.
+    cross = -(posterior * prediction.log()).sum()  # The cross-entropy the prediction scores.
+    return float(entropy), float(cross - entropy)  # The gap is the KL divergence to the prediction.
+
+
+# %% 06. The Euler step as a two-part rule
+#
+# Substituting the marginal velocity into the per-position Euler step collects the two terms into a
+# mixture of the current point mass and the denoiser. With probability lambda_t the position redraws
+# from the denoiser, and with probability 1 - lambda_t it keeps the token it holds.
+
+
+def redraw_probability(t, h):
+    """lambda_t = h kappa_dot_t / (1 - kappa_t), the redraw probability of Section 5.8."""
+    return h / (1.0 - t)  # Linear schedule; above one the sampler redraws with certainty.
+
+
+@torch.no_grad()
+def dfm_step(proc, model, xt, t, h):
+    """One two-part step at every position at once. xt: [B, L], t: [B], h a scalar."""
+    p1 = model(xt, t)[..., :proc.V].softmax(-1)  # [B, L, V] denoiser, the only network call.
+    rate = proc.kd(t) / (1 - proc.k(t)).clamp(min=1e-4)  # kappa_dot/(1 - kappa), kept finite at t=1.
+    redraw = torch.rand(xt.shape, device=xt.device) < (h * rate)[:, None]  # [B, L] Bernoulli.
+    x1 = Categorical(probs=p1).sample()  # [B, L] one candidate token per position.
+    return torch.where(redraw, x1, xt)  # Keep the candidate where the coin came up true.
+
+
+# %% 07. Training and generation
+#
+# A training step draws one time per sequence, corrupts the batch at that time, and scores the true
+# tokens. Generation walks a uniform grid of ``sample_steps`` times from t = 0 to t = 1, and the grid
+# is free to change because training never chose a step size.
+
+
+@torch.no_grad()
+def generate(proc, model, batch, length, steps, device="cpu"):
+    """Run the sampler from the source to t = 1 and return the tokens. [B, L]."""
+    x = proc.source((batch, length), device)  # Draw from the source distribution p.
+    grid = torch.linspace(0.0, 1.0, steps + 1, device=device)  # The sampling times.
+    for t, s in zip(grid[:-1], grid[1:]):
+        x = dfm_step(proc, model, x, t.expand(batch), s - t)  # One forward pass per step.
+    return x
+
+
+# %% 08. The chapter's worked numbers
+def worked_numbers():
+    """The arithmetic Sections 5.5.5, 5.7.4, 5.8, 5.9.2 and 5.9.3 carry out by hand."""
+    report = {}
+
+    # Section 5.5.5: x_1 = ACGT, a uniform source draw x_0 = GTTA, kappa_t = 0.5, coins HTHT.
+    x1, x0 = encode_dna(["ACGT"])[0], encode_dna(["GTTA"])[0]  # [4] each.
+    heads = torch.tensor([True, False, True, False])  # The four coins of the example.
+    xt = torch.where(heads, x1, x0)  # Data token on heads, source token on tails.
+    report["mixture_draw"] = decode_dna(xt[None])[0]  # ATGA.
+
+    # Section 5.7.4: one position holding G, denoiser (0.1, 0.6, 0.2, 0.1), t = 0.5.
+    p1 = torch.tensor([0.1, 0.6, 0.2, 0.1], dtype=torch.double)  # [V] over A, C, G, T.
+    u = marginal_velocity(p1, x_i=2, t=0.5)  # Factor 1/0.5 = 2 times the bracket.
+    report["velocity_at_G"] = [round(float(v), 4) for v in u]  # (0.2, 1.2, -1.6, 0.2).
+    report["velocity_sums_to"] = f"{float(u.sum()):.1e}"  # Zero, which is the second rule.
+    step = F.one_hot(torch.tensor(2), DNA_SIZE).double() + 0.1 * u  # One Euler step of size h = 0.1.
+    report["euler_step_at_G"] = [round(float(v), 4) for v in step]  # (0.02, 0.12, 0.84, 0.02).
+    report["redraw_probability"] = round(redraw_probability(0.5, 0.1), 4)  # lambda_t = 0.2.
+
+    # Section 5.9.2: the noisy input ATGA scored against the target ACGT.
+    scored = torch.tensor([0.7, 0.2, 0.6, 0.1], dtype=torch.double)  # One probability per position.
+    report["example_loss_nats"] = round(float(-scored.log().sum()), 3)  # 4.780 nats.
+
+    # Section 5.9.3: length two, the mask source, q(AC) = 0.5, q(AG) = 0.3, q(GC) = 0.2.
+    candidates, weights = ["AC", "AG", "GC"], [0.5, 0.3, 0.2]
+    half = mask_posterior(candidates, weights, visible="A", kappa_t=0.5)  # Denoiser at kappa = 0.5.
+    late = mask_posterior(candidates, weights, visible="A", kappa_t=0.8)  # The same at kappa = 0.8.
+    report["mask_posterior_C_G"] = [round(float(half[1]), 4), round(float(half[2]), 4)]
+    report["mask_posterior_unchanged"] = bool(torch.allclose(half, late))  # The time cancels.
+    live = half[half > 0]  # The two bases the posterior keeps.
+    entropy, gap = entropy_and_gap(live, torch.tensor([0.5, 0.5], dtype=torch.double))
+    report["posterior_entropy_nats"] = round(entropy, 4)  # 0.6616, the floor on the loss.
+    report["kl_gap_to_uniform"] = round(gap, 4)  # 0.0316 above that floor.
+    return report
+
+
+# %% 09. The report
+def main():
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.add_argument("--source", choices=["uniform", "mask"], default="uniform")
+    parser.set_defaults(steps=1000, width=64, samples=64)
+    args = parser.parse_args()
+    set_seed(args.seed)
+
+    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
+    data = splits["train_tokens"]  # [N, L] the shared synthetic DNA.
+    mask = DNA_SIZE if args.source == "mask" else None  # Index 4 is the mask symbol m.
+    vocab = DNA_SIZE + 1 if mask is not None else DNA_SIZE  # The network's vocabulary.
+    proc = DiscreteFlowMatching(DNA_SIZE, mask=mask)  # Linear schedule by default.
+    model = SequenceDenoiser(vocab, args.length, width=args.width, depth=2)
+
+    def loss_fn(step):
+        rows = torch.randint(0, len(data), (args.batch_size,))  # One minibatch of data sequences.
+        x1 = data[rows]  # [B, L] the data end of the pair.
+        t = torch.rand(args.batch_size)  # One time per sequence, uniform on [0, 1].
+        xt = proc.corrupt(x1, t)  # [B, L] a draw from the conditional path.
+        return dfm_loss(model(xt, t)[..., :DNA_SIZE], x1)  # Score the true token at every position.
+
+    if not args.quiet:
+        print(f"source={args.source} vocab={vocab} parameters={count_parameters(model):,}")
+    losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=args.quiet)
+    tokens = generate(proc, model, args.samples, args.length, args.sample_steps)
+
+    report = worked_numbers()
+    report["final_loss_nats"] = sum(losses[-20:]) / len(losses[-20:])  # Averaged over the last steps.
+    report["gc_fraction"] = gc_fraction(tokens)  # The training data sit at 0.577.
+    report["motif_fraction"] = motif_fraction(tokens)  # Sequences carrying the planted ACGT.
+    report["unique_fraction"] = unique_fraction(tokens)  # Catches a collapsed sampler.
+    report["token_entropy_nats"] = token_entropy(tokens, DNA_SIZE)  # At most log 4 = 1.386.
+    print_report(f"discrete flow matching, {args.source} source", report, decode_dna(tokens))
+    save_run(args.out, model, vars(args), report, losses, decode_dna(tokens))
+
+
+if __name__ == "__main__":
+    main()

lecture_5/dirichlet.py

@@ -0,0 +1,262 @@
+"""Dirichlet Flow Matching: a token relaxed into a point of the simplex.
+
+A straight line toward a corner shrinks its support as t grows, and past t = 1/k at most k - 1
+corners remain possible, which makes the denoiser target discontinuous in time and in space.
+Replacing the line by the Dirichlet family Dir(1 + t e_i) keeps every corner possible at every finite
+time. The velocity of that path points at the data corner and changes only its speed, and matching
+the probability flux across the cut x_i = c solves for that speed.
+
+Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
+u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y). On this path the label x_1
+still means the data sample while its subscript stops meaning time one, since time runs to infinity
+and inference stops at t = 8.
+
+Chapter 5, Sections 5.19.3 and 5.20.1 to 5.20.5. Trained on the synthetic DNA of ``cis6270.data``.
+Prints the worked speeds, the continuity check, the posterior average, and the measurements of a
+generated batch.
+
+    python lecture_5/dirichlet.py                     # train, then 100 Euler steps to t = 8
+    python lecture_5/dirichlet.py --steps 20 --quiet  # a few seconds, for a first look
+"""
+
+# %% 01. Imports and configuration
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from scipy.special import betainc, betaln
+
+from cis6270.data import DNA_SIZE, decode_dna, dna_splits
+from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
+from cis6270.nets import SequenceDenoiser, count_parameters
+from cis6270.runner import common_parser, print_report, save_run, set_seed, train
+
+T_MAX = 8.0  # The inference horizon, where the data coordinate has mean (t+1)/(t+K) = 0.75 for DNA.
+EDGE = 1e-6  # The speed divides by (1 - x_i) and by x_i, so a coordinate is held off both ends.
+
+
+# %% 02. The Dirichlet family as a conditional path
+#
+# With the concentration 1 + t in the data coordinate and 1 everywhere else, every exponent except
+# the i-th is zero and the density collapses to a single factor. At t = 0 the path is the uniform
+# distribution on the simplex, and for every finite t it is positive wherever x_i is, so no corner is
+# ever ruled out.
+
+
+def dirichlet_density(x, i, t, K=DNA_SIZE):
+    """p_t(x | x_1 = e_i) = Gamma(t + K)/Gamma(t + 1) x_i^t."""
+    prefactor = math.exp(math.lgamma(t + K) - math.lgamma(t + 1.0))  # Gamma(4) = 6 at t = 0.
+    return prefactor * float(x[i]) ** t  # Constant over the simplex at t = 0.
+
+
+def sample_dirichlet(targets, t, K=DNA_SIZE):
+    """Draw x ~ Dir(1 + t e_i) at every position. targets: [B, L], t: [B] -> [B, L, K]."""
+    alpha = torch.ones(*targets.shape, K) + t[:, None, None] * F.one_hot(targets, K)  # [B, L, K].
+    gamma = torch._standard_gamma(alpha)  # One Gamma draw per letter, with unit rate.
+    return gamma / gamma.sum(-1, keepdim=True)  # Normalizing the Gamma draws gives a Dirichlet.
+
+
+def straight_line_corners(x, t):
+    """How many corners the straight-line path still allows at x, which is the count of x_i >= t."""
+    return int((x >= t).sum())  # Proposition 5.5 bounds this by k - 1 once t > 1/k.
+
+
+# %% 03. The velocity by matching flux across a cut
+#
+# Cutting the simplex at x_i = c, the mass below the cut is the Beta(t+1, K-1) distribution function
+# at c, and its rate of change is the derivative of that function in its first parameter. Setting the
+# loss rate equal to minus the outward flux and solving for the speed leaves one scalar per
+# coordinate, with no closed form at general K.
+
+
+def d_betainc_da(x, a, b, eps=1e-6):
+    """tilde I_x(a, b), the derivative of the regularized incomplete beta function in a."""
+    return (betainc(a + eps, b, x) - betainc(a - eps, b, x)) / (2.0 * eps)  # A central difference.
+
+
+def speed(x_i, t, K=DNA_SIZE):
+    """C(x_i, t) = -tilde I_{x_i}(t+1, K-1) B(t+1, K-1) / ((1 - x_i)^{K-1} x_i^t).
+
+    The data coordinate may be a scalar or a tensor of any shape, since every factor broadcasts,
+    and a whole batch of relaxed tokens costs the two calls to the incomplete beta function that
+    one costs. The central difference subtracts two numbers that agree to about six decimals, and
+    at single precision nothing of the difference survives, so this one quantity is computed in
+    double whatever the ambient dtype and the caller casts the result back.
+    """
+    a, b = t + 1.0, float(K - 1)  # The two Beta parameters of the data coordinate.
+    beta = math.exp(betaln(a, b))  # The complete beta function, 1/12 at a = 2 and b = 3.
+    held = torch.as_tensor(x_i, dtype=torch.double).clamp(EDGE, 1.0 - EDGE)  # Off both endpoints.
+    derivative = torch.as_tensor(d_betainc_da(held.numpy(), a, b))  # Negative, so the speed is not.
+    return -derivative * beta / ((1.0 - held) ** (K - 1) * held ** t)  # Positive.
+
+
+def conditional_velocity(x, i, t):
+    """u_t(x | x_1 = e_i) = C(x_i, t)(e_i - x), tangent to the simplex. x: [K] -> [K]."""
+    direction = F.one_hot(torch.tensor(i), len(x)).to(x) - x  # Coordinates summing to zero.
+    return float(speed(float(x[i]), t, len(x))) * direction  # The speed scales the direction.
+
+
+# %% 04. Two letters, in closed form
+#
+# With K = 2 the second Beta parameter is one, the incomplete beta function is x^a, and the speed
+# comes out in elementary functions. Differentiating the density in time and the flux in space then
+# cancels term by term, which is the continuity equation at K = 2.
+
+
+def speed_two_letters(x, t):
+    """C(x_1, t) = -x_1 log x_1 / ((t + 1)(1 - x_1)), the K = 2 case of the general formula."""
+    return -x * math.log(x) / ((t + 1.0) * (1.0 - x))  # Positive on the whole open interval.
+
+
+def continuity_terms(x, i, t, h=1e-2):
+    """partial_t p and div(p u) in the chart (x_1, ..., x_{K-1}), whose sum the path drives to zero."""
+    x = x.to(torch.double)  # A stencil on top of a numerical derivative needs the extra digits.
+    K = len(x)
+
+    def density(point):
+        return dirichlet_density(point, i, t, K)  # The Dirichlet density at a simplex point.
+
+    def chart(point, coordinate):
+        return density(point) * float(conditional_velocity(point, i, t)[coordinate])  # p u_j.
+
+    time = density(x) * (_digamma(t + K) - _digamma(t + 1.0) + math.log(float(x[i])))  # partial_t p.
+    divergence = 0.0
+    for j in range(K - 1):  # The chart drops the last coordinate, which the sum constraint fixes.
+        step = torch.zeros(K, dtype=x.dtype)
+        step[j], step[K - 1] = 1.0, -1.0  # A direction inside the simplex, keeping the sum at one.
+        divergence += (  # A five-point stencil, since the speed itself is a numerical derivative.
+            -chart(x + 2 * h * step, j) + 8 * chart(x + h * step, j)
+            - 8 * chart(x - h * step, j) + chart(x - 2 * h * step, j)
+        ) / (12.0 * h)
+    return time, divergence  # Equal and opposite to the precision of the difference.
+
+
+def _digamma(z):
+    """The digamma function, used only to differentiate the Gamma prefactor in time."""
+    return float(torch.digamma(torch.tensor(z, dtype=torch.double)))
+
+
+# %% 05. A worked speed on a DNA base
+#
+# A coordinate already holding a large share has a larger speed, so the flow concentrates gradually
+# and the coordinate in the lead moves fastest toward its corner.
+
+
+def speed_vector(x, t):
+    """C(x_k, t) at every letter of one relaxed base. x: [K] -> [K]."""
+    return speed(x, t, x.shape[-1]).to(x)  # One vectorized evaluation, cast back to x's dtype.
+
+
+# %% 06. The posterior average and the loss
+#
+# The data corner is unknown at generation time, so the four conditional velocities combine with the
+# denoiser as weights. Each one is tangent to the simplex and a weighted average of tangent vectors
+# is tangent, so the learned field needs no constraint. Cross-entropy on the data token trains the
+# denoiser, and the conditional velocity appears only inside the sampler.
+
+
+def marginal_velocity(x, p_hat, t):
+    """v_t(x) = sum_k u_t(x | e_k) p_{1|t}(e_k | x). x: [K], p_hat: [K] -> [K]."""
+    columns = torch.stack([conditional_velocity(x, k, t) for k in range(len(x))], dim=1)  # [K, K].
+    return columns @ p_hat  # The posterior average of the four conditional velocities.
+
+
+def batched_velocity(x, p_hat, t):
+    """The same average at every position of a batch. x, p_hat: [B, L, K] -> [B, L, K]."""
+    speeds = speed(x, t, x.shape[-1]).to(x)  # [B, L, K] one speed per coordinate.
+    weighted = speeds * p_hat  # [B, L, K] the coefficient of each conditional velocity.
+    return weighted - weighted.sum(-1, keepdim=True) * x  # sum_k w_k (e_k - x), collected.
+
+
+def relaxed_denoiser(length, width, K=DNA_SIZE, depth=2):
+    """The shared sequence denoiser with its lookup table replaced by a map on the relaxed token."""
+    model = SequenceDenoiser(K, length, width=width, depth=depth)
+    model.token_embedding = nn.Linear(K, width, bias=False)  # A point of the simplex, not an index.
+    return model
+
+
+# %% 07. Training on DNA, and sampling by Euler steps
+#
+# Training draws t from an exponential with rate one, places a point on the conditional path, and
+# scores the true letter. Generation starts from the uniform Dirichlet and integrates the marginal
+# velocity to t = 8, projecting back onto the simplex after each step.
+
+
+@torch.no_grad()
+def generate(model, batch, length, steps, K=DNA_SIZE):
+    """Integrate the marginal velocity from the uniform simplex to t = T_MAX. Returns [B, L]."""
+    gamma = torch._standard_gamma(torch.ones(batch, length, K))  # Dir(1) is the uniform source.
+    x = gamma / gamma.sum(-1, keepdim=True)  # [B, L, K].
+    h = T_MAX / steps  # The time increment, 0.08 over 100 steps.
+    for k in range(steps):
+        t = k * h
+        p_hat = model(x, torch.full((batch,), t)).softmax(-1)  # [B, L, K] the denoiser.
+        x = x + h * batched_velocity(x, p_hat, t)  # One Euler step along the marginal velocity.
+        x = x.clamp(min=1e-8)  # A large step can push a coordinate below zero.
+        x = x / x.sum(-1, keepdim=True)  # Back onto the simplex.
+    return x.argmax(-1)  # Read the token off the largest coordinate.
+
+
+# %% 08. The report
+def main():
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.set_defaults(steps=1000, width=64, samples=64, sample_steps=100)
+    args = parser.parse_args()
+    set_seed(args.seed)
+
+    report = {
+        "density_at_t_0": round(dirichlet_density(torch.full((4,), 0.25), 0, 0.0), 4),  # 6.
+        "corners_left_at_0.3": straight_line_corners(torch.tensor([0.34, 0.33, 0.33, 0.0]), 0.3),
+        "corners_left_at_0.6": straight_line_corners(torch.tensor([0.70, 0.30, 0.0, 0.0]), 0.6),
+        "speed_K2_closed": f"{speed_two_letters(0.3, 1.0):.7f}",  # 0.2579942.
+        "speed_K2_general": f"{float(speed(0.3, 1.0, K=2)):.7f}",  # The same to six decimals.
+    }
+    point = torch.tensor([0.5, 0.3, 0.2], dtype=torch.double)  # A base of three letters, at t = 0.8.
+    time_term, space_term = continuity_terms(point, 0, 0.8)
+    report["continuity_time_term"] = round(time_term, 4)  # 0.6355.
+    report["continuity_space_term"] = round(space_term, 4)  # -0.6355.
+    report["continuity_residual"] = f"{time_term + space_term:.2e}"  # Their sum, near zero.
+
+    # Section 5.20.4: K = 4, t = 1 and the relaxed base (0.40, 0.30, 0.20, 0.10).
+    x = torch.tensor([0.40, 0.30, 0.20, 0.10], dtype=torch.double)
+    report["speeds"] = [round(float(c), 3) for c in speed_vector(x, 1.0)]  # 0.232, 0.199, 0.159, 0.104.
+    toward_a = conditional_velocity(x, 0, 1.0)  # [4] the velocity toward the A corner.
+    report["velocity_toward_A"] = [round(float(v), 3) for v in toward_a]
+    report["velocity_sums_to"] = f"{float(toward_a.sum()):.1e}"  # Zero, so the step stays on it.
+
+    # Section 5.20.5: the same base with the denoiser (0.60, 0.25, 0.10, 0.05).
+    p_hat = torch.tensor([0.60, 0.25, 0.10, 0.05], dtype=torch.double)
+    v = marginal_velocity(x, p_hat, 1.0)
+    report["marginal_velocity"] = [round(float(c), 4) for c in v]  # (0.0553, -0.0133, -0.0262, -0.0158).
+    report["euler_step_h_0.1"] = [round(float(c), 4) for c in x + 0.1 * v]  # (0.4055, ...).
+    report["step_sums_to"] = round(float((x + 0.1 * v).sum()), 9)  # One, still on the simplex.
+
+    # Training on the shared DNA, then generation.
+    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
+    data = splits["train_tokens"]  # [N, L].
+    model = relaxed_denoiser(args.length, args.width)
+
+    def loss_fn(step):
+        rows = torch.randint(0, len(data), (args.batch_size,))  # One minibatch of data sequences.
+        x1 = data[rows]  # [B, L] the letters the denoiser has to recover.
+        t = torch.distributions.Exponential(1.0).sample((args.batch_size,))  # t ~ Exp(1).
+        xt = sample_dirichlet(x1, t)  # [B, L, K] a point on the conditional path.
+        return F.cross_entropy(model(xt, t).transpose(1, 2), x1)  # Score the true letter.
+
+    if not args.quiet:
+        print(f"parameters={count_parameters(model):,} horizon={T_MAX}")
+    losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=args.quiet)
+    tokens = generate(model, args.samples, args.length, args.sample_steps)
+
+    report["final_loss_nats"] = sum(losses[-20:]) / len(losses[-20:])
+    report["gc_fraction"] = gc_fraction(tokens)  # The training data sit at 0.577.
+    report["motif_fraction"] = motif_fraction(tokens)  # Sequences carrying the planted ACGT.
+    report["unique_fraction"] = unique_fraction(tokens)  # Catches a collapsed sampler.
+    report["token_entropy_nats"] = token_entropy(tokens, DNA_SIZE)  # At most log 4 = 1.386.
+    print_report("Dirichlet flow matching", report, decode_dna(tokens))
+    save_run(args.out, model, vars(args), report, losses, decode_dna(tokens))
+
+
+if __name__ == "__main__":
+    main()

lecture_5/edit_flows.py

@@ -0,0 +1,295 @@
+"""Edit Flows: paths whose sequences change length.
+
+Every path so far keeps a single length, because a substitution changes a token and never adds or
+removes one. Enlarging the set of moves to insertions, deletions and substitutions removes that
+constraint. An alignment pads the two ends with a blank until they have the same number of columns,
+the mixture construction then applies column by column, and the alignment is an auxiliary variable we
+marginalize exactly as we marginalized the data sequence. The loss scores rates in place of a
+distribution over tokens, and its minimizer is the rate at which an edit makes progress.
+
+Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
+u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
+
+Chapter 5, Section 5.17. Runs on the synthetic DNA of ``cis6270.data``, trimmed to varying lengths.
+Prints the alignment of the chapter, the posterior over two alignments, the averaged edit rates, the
+loss of one example, one sampler step, and the length statistics of a run on DNA.
+
+    python lecture_5/edit_flows.py                     # 600 sampler steps on 256 DNA pairs
+    python lecture_5/edit_flows.py --steps 20 --quiet  # a coarse grid, a visibly coarser length
+"""
+
+# %% 01. Imports and configuration
+import math
+
+import torch
+
+from cis6270.data import DNA, DNA_SIZE, dna_splits, encode_dna
+from cis6270.runner import common_parser, print_report, save_run, set_seed
+
+BLANK = DNA_SIZE  # Index 4 is the blank symbol epsilon that padding writes into a column.
+
+
+# %% 02. Three edits and the rate split
+#
+# At each position of the current sequence there are three edits, and each rate splits into an
+# intensity and a distribution over tokens. A deletion writes no token, so the network emits three
+# intensities and two token distributions per position, which is 3 + 2K numbers in place of the K of
+# the denoiser.
+
+
+def edit_rates(intensities, tokens=DNA_SIZE):
+    """Count the numbers a network emits per position, and the total rate of leaving the sequence."""
+    total = sum(sum(kind) for kind in intensities.values())  # -u_t(x, x), the total exit rate.
+    return 3 + 2 * tokens, total  # Three intensities and two token distributions per position.
+
+
+# %% 03. Alignments and the column-wise path
+#
+# An alignment puts the positions of the two sequences in correspondence by padding both with the
+# blank. A blank in the target column is a deletion, a blank in the source column an insertion, and
+# two different tokens a substitution. On the aligned pair each column follows the mixture path
+# independently, and removing the blanks from z_t gives the sequence x_t the model reads.
+
+
+def align(x0, x1):
+    """A minimal alignment of two token lists, as the padded pair (z_0, z_1) of equal length."""
+    n, m = len(x0), len(x1)
+    cost = [[0] * (m + 1) for _ in range(n + 1)]  # Edit distance with unit costs.
+    for i in range(n + 1):
+        cost[i][0] = i  # Deleting the first i tokens of the source.
+    for j in range(m + 1):
+        cost[0][j] = j  # Inserting the first j tokens of the target.
+    for i in range(1, n + 1):
+        for j in range(1, m + 1):
+            cost[i][j] = min(cost[i - 1][j] + 1, cost[i][j - 1] + 1,
+                             cost[i - 1][j - 1] + int(x0[i - 1] != x1[j - 1]))
+    z0, z1, i, j = [], [], n, m
+    while i or j:  # One backtrace through the table, taking the diagonal where it is minimal.
+        if i and j and cost[i][j] == cost[i - 1][j - 1] + int(x0[i - 1] != x1[j - 1]):
+            z0.append(x0[i - 1]), z1.append(x1[j - 1]), (i := i - 1), (j := j - 1)
+        elif i and cost[i][j] == cost[i - 1][j] + 1:
+            z0.append(x0[i - 1]), z1.append(BLANK), (i := i - 1)  # A deletion column.
+        else:
+            z0.append(BLANK), z1.append(x1[j - 1]), (j := j - 1)  # An insertion column.
+    return z0[::-1], z1[::-1]  # [C] and [C] over the same number of columns.
+
+
+def column_kinds(z0, z1):
+    """Name the edit each column owes: ``del``, ``ins``, ``sub``, or ``none`` when the two agree."""
+    kinds = []
+    for a, b in zip(z0, z1):
+        kinds.append("none" if a == b else "del" if b == BLANK else "ins" if a == BLANK else "sub")
+    return kinds  # [C] one label per column.
+
+
+def draw_state(z0, z1, finished):
+    """z_t from a choice of which columns already hold their target, with the blanks removed."""
+    held = [b if done else a for a, b, done in zip(z0, z1, finished)]  # [C] the aligned state.
+    return [token for token in held if token != BLANK]  # [|x_t|] the sequence the model reads.
+
+
+def expected_length(x0, x1, kappa_t):
+    """E[|x_t|] = (1 - kappa_t)|x_0| + kappa_t|x_1|, the interpolation one position up."""
+    return (1.0 - kappa_t) * len(x0) + kappa_t * len(x1)  # Linear from |x_0| to |x_1|.
+
+
+# %% 04. Averaging over the alignments we cannot see
+#
+# Many alignments produce the same pair, and x_t does not record which one the path is on. Bayes'
+# rule over the alignments weights each by the probability that its column draw produces the observed
+# x_t, and the rates average against those weights. No single alignment has to be chosen.
+
+
+def alignment_posterior(alignments, observed, kappa_t):
+    """p_t(A | x_t) over a list of aligned pairs, by enumerating the column draws of each."""
+    weights = []
+    for z0, z1 in alignments:
+        owed = [i for i, kind in enumerate(column_kinds(z0, z1)) if kind != "none"]  # Unfinished.
+        total = 0.0
+        for mask in range(1 << len(owed)):  # Every subset of the owed columns already finished.
+            finished = [True] * len(z0)
+            probability = 1.0
+            for bit, column in enumerate(owed):
+                done = bool(mask >> bit & 1)
+                finished[column] = done
+                probability *= kappa_t if done else 1.0 - kappa_t  # The column-wise mixture path.
+            if draw_state(z0, z1, finished) == observed:
+                total += probability  # This draw is consistent with what we see.
+        weights.append(total)
+    norm = sum(weights)
+    return [w / norm for w in weights]  # One posterior weight per alignment.
+
+
+def averaged_rates(alignments, observed, kappa_t, kappa_dot=1.0):
+    """The marginal rate of each edit that finishes a column, averaged over the alignments.
+
+    ``observed`` is the draw in which every column that owes an edit still holds its source symbol,
+    which is the state the worked example of Section 5.17.2 reads the rates at.
+    """
+    rate = kappa_dot / (1.0 - kappa_t)  # The weight on every unfinished column.
+    posterior, out = alignment_posterior(alignments, observed, kappa_t), {}
+    for (z0, z1), weight in zip(alignments, posterior):
+        held = [b if kind == "none" else a  # A column owing nothing is finished at every time.
+                for a, b, kind in zip(z0, z1, column_kinds(z0, z1))]
+        for i, (a, b, kind) in enumerate(zip(z0, z1, column_kinds(z0, z1))):
+            seen = sum(1 for token in held[:i] if token != BLANK)  # Position inside x_t.
+            if kind == "sub":
+                name = f"sub({seen + 1},{DNA[b]})"  # Replace the token at that position.
+            elif kind == "ins":
+                name = f"ins({seen},{DNA[b]})"  # Place a token to the right of that position.
+            elif kind == "del":
+                name = f"del({seen + 1})"  # Remove the token at that position.
+            else:
+                continue  # Nothing is owed here.
+            out[name] = out.get(name, 0.0) + weight * rate
+    return posterior, out
+
+
+# %% 05. The training loss
+#
+# The first term is the sum of the model's intensities and pushes all of them down. The second is a
+# weighted cross-entropy on the edits that move x_t toward x_1 and pushes those up. Collecting the
+# terms belonging to one edit gives lambda - w log lambda, whose stationary point is lambda = w.
+
+
+def edit_loss(total_intensity, owed, kappa_t, kappa_dot=1.0):
+    """L = sum of all intensities - sum over unfinished columns of the weight times log lambda."""
+    weight = kappa_dot / (1.0 - kappa_t)  # kappa_dot/(1 - kappa), which is 2 at t = 0.5.
+    return total_intensity - sum(weight * math.log(lam) for lam in owed)  # The model's own rates.
+
+
+def loss_minimizer(weight, grid):
+    """The intensity minimizing lambda - w log lambda on a grid, which the derivative puts at w."""
+    values = [lam - weight * math.log(lam) for lam in grid]  # The contribution of one edit.
+    return grid[min(range(len(grid)), key=values.__getitem__)]  # Near w, with w/lambda^2 > 0.
+
+
+# %% 06. One step of the sampler
+#
+# Each position inserts with probability h lambda_ins, and deletes or substitutes with probability
+# h(lambda_del + lambda_sub), the two sharing that draw in proportion to their intensities. Every
+# chosen edit is applied at once, so the step has the same second-order error as the parallel
+# step on a fixed length.
+
+
+def step_probabilities(intensities, h):
+    """The per-position insertion and delete-or-substitute probabilities of one Euler step."""
+    insert = [h * lam for lam in intensities["ins"]]  # One coin per position.
+    change = [h * (d + s) for d, s in zip(intensities["del"], intensities["sub"])]  # A second coin.
+    return insert, change
+
+
+def untouched_probability(intensities, h):
+    """The chance the step applies no edit anywhere, and the first-order estimate beside it."""
+    insert, change = step_probabilities(intensities, h)
+    exact = math.prod(1.0 - p for p in insert + change)  # Every coin coming up tails.
+    _, total = edit_rates(intensities)
+    return exact, 1.0 - h * total  # The gap is the second-order term from two edits at once.
+
+
+def expected_length_change(intensities, h):
+    """h times the insertion intensities minus the deletion intensities, in tokens."""
+    return h * (sum(intensities["ins"]) - sum(intensities["del"]))  # Nothing fixes the length.
+
+
+# %% 07. Edits on the synthetic DNA
+#
+# Each pair of sequences is aligned once, every column that owes an edit finishes at rate
+# kappa_dot/(1 - kappa), and the sequence the model reads is the aligned state with its blanks
+# removed. The realized length wanders while its mean follows the interpolation.
+
+
+def simulate_columns(pairs, steps, generator=None):
+    """Run the column-wise path on a list of aligned pairs. Returns the lengths and the end states."""
+    width = max(len(z0) for z0, _ in pairs)  # Pad every alignment to the widest one.
+    src = torch.full((len(pairs), width), BLANK)
+    tgt = torch.full((len(pairs), width), BLANK)
+    valid = torch.zeros(len(pairs), width, dtype=torch.bool)
+    for row, (z0, z1) in enumerate(pairs):
+        src[row, :len(z0)] = torch.tensor(z0)
+        tgt[row, :len(z1)] = torch.tensor(z1)
+        valid[row, :len(z0)] = True  # Columns beyond the alignment take no part.
+    done = (src == tgt) & valid  # A column whose two symbols agree owes no edit.
+    h, halfway = 1.0 / steps, None
+    for k in range(steps):
+        t = k * h
+        if halfway is None and t >= 0.5:
+            held = torch.where(done, tgt, src)  # [P, C] the aligned state at this time.
+            halfway = float(((held != BLANK) & valid).sum(dim=1).double().mean())
+        lam = 1.0 if k == steps - 1 else min(1.0, h / max(1.0 - t, 1e-12))  # lambda_t per column.
+        coin = torch.rand(src.shape, generator=generator) < lam
+        done = done | (coin & valid)  # An unfinished column takes its target symbol.
+    finished = float((done | ~valid).all(dim=1).double().mean())  # Fraction of pairs that arrived.
+    return halfway, finished
+
+
+# %% 08. The report
+def main():
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.add_argument("--pairs", type=int, default=256, help="DNA pairs aligned and run")
+    args = parser.parse_args()
+    set_seed(args.seed)
+
+    # Section 5.17.1: the chapter's alignment of ACGT with AGAC over five columns.
+    x0, x1 = encode_dna(["ACGT"])[0].tolist(), encode_dna(["AGAC"])[0].tolist()
+    z0 = x0 + [BLANK]  # (A, C, G, T, eps).
+    z1 = [x1[0], BLANK] + x1[1:]  # (A, eps, G, A, C).
+    kinds = column_kinds(z0, z1)
+    draw = [True, False, True, True, True]  # Column 2 behind, columns 4 and 5 already at the target.
+    report = {
+        "column_kinds": kinds,  # none, del, none, sub, ins.
+        "draw_at_kappa_0.5": "".join(DNA[token] for token in draw_state(z0, z1, draw)),  # ACGAC.
+        "draw_probability": round(0.5 ** sum(kind != "none" for kind in kinds), 4),  # 0.125.
+        "expected_length": round(expected_length(x0, x1, 0.5), 4),  # 4, both ends having length 4.
+        "numbers_per_position": edit_rates({"ins": [0.0], "del": [0.0], "sub": [0.0]})[0],  # 11.
+    }
+
+    # Section 5.17.2: the pair AC and AAC, which admits two alignments on three columns.
+    short, long = encode_dna(["AC"])[0].tolist(), encode_dna(["AAC"])[0].tolist()
+    first = ([short[0], short[1], BLANK], long)  # Substitute C with A, then insert C.
+    second = ([short[0], BLANK, short[1]], long)  # Insert A, and leave the last column alone.
+    observed = short  # The noisy sequence x_t = AC.
+    for kappa_t in (0.25, 0.5, 0.75):
+        weights = alignment_posterior([first, second], observed, kappa_t)
+        report[f"posterior_A1_kappa_{kappa_t}"] = round(weights[0], 4)  # (1-kappa)/(2-kappa).
+    posterior, rates = averaged_rates([first, second], observed, 0.5)
+    for name in sorted(rates):
+        report[f"rate_{name}"] = round(rates[name], 3)  # 0.667, 1.333 and 0.667.
+    report["total_target_rate"] = round(sum(rates.values()), 3)  # 2.667.
+
+    # Section 5.17.3: one unfinished column, a model intensity of 0.8, and 3.0 over all edits.
+    report["example_loss"] = round(edit_loss(3.0, [0.8], kappa_t=0.5), 4)  # 3.446.
+    grid = [0.01 * i for i in range(1, 801)]  # Intensities between 0.01 and 8.
+    report["loss_minimizer"] = round(loss_minimizer(2.0, grid), 2)  # 2, the weight the edit receives.
+
+    # Section 5.17.4: one whole step on ACG with the chapter's nine intensities.
+    intensities = {"ins": [0.2, 0.3, 0.4], "del": [0.1, 0.9, 0.2], "sub": [0.4, 0.3, 0.5]}
+    exact, first_order = untouched_probability(intensities, 0.05)
+    insert, change = step_probabilities(intensities, 0.05)
+    report["insert_probabilities"] = [round(p, 3) for p in insert]  # 0.010, 0.015, 0.020.
+    report["change_probabilities"] = [round(p, 3) for p in change]  # 0.025, 0.060, 0.035.
+    report["deletion_given_change"] = round(0.9 / (0.9 + 0.3), 4)  # 0.75, three times in four.
+    report["untouched"] = round(exact, 4)  # 0.8452.
+    report["untouched_first_order"] = round(first_order, 4)  # 0.835.
+    report["second_order_gap"] = round(exact - first_order, 4)  # 0.0102.
+    report["length_change"] = round(expected_length_change(intensities, 0.05), 4)  # -0.015.
+
+    # The same path on the shared DNA, trimmed so that the two ends have different lengths.
+    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
+    tokens = splits["train_tokens"]
+    cuts = torch.randint(args.length - 4, args.length + 1, (2 * args.pairs,))  # Varying lengths.
+    rows = [tokens[i % len(tokens), :cuts[i]].tolist() for i in range(2 * args.pairs)]
+    pairs = [align(rows[2 * i], rows[2 * i + 1]) for i in range(args.pairs)]
+    predicted = sum(expected_length(rows[2 * i], rows[2 * i + 1], 0.5) for i in range(args.pairs))
+    measured, arrived = simulate_columns(pairs, args.steps)
+    report["dna_columns_mean"] = round(sum(len(z0) for z0, _ in pairs) / len(pairs), 3)
+    report["dna_length_at_0.5"] = round(measured, 3)  # Measured over the run.
+    report["dna_length_predicted"] = round(predicted / args.pairs, 3)  # The interpolation.
+    report["dna_pairs_arrived"] = round(arrived, 4)  # Every column finishes by t = 1.
+    report["dna_example"] = " -> ".join("".join(DNA[t] for t in row) for row in rows[:2])
+    print_report("edit flows", report)
+    save_run(args.out, config=vars(args), report=report)
+
+
+if __name__ == "__main__":
+    main()

lecture_5/fisher.py

@@ -0,0 +1,257 @@
+"""Fisher-Flow: the simplex measured by the metric statistics already uses for distributions.
+
+The straight-line distance on the simplex measures absolute differences in probability, so dividing a
+rare letter's probability by ten counts for almost nothing. The Fisher information measures how
+distinguishable two nearby distributions are, and by Chentsov's theorem it is the only metric under
+which relabeling the letters or merging two of them is an isometry. Taking the square root of every
+coordinate turns it into the ordinary metric on a sphere, where the conditional path is a
+great-circle arc and the velocity is the remaining arc divided by the remaining time.
+
+Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
+u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
+
+Chapter 5, Sections 5.21.1 to 5.21.5, with the listing of Section 5.21.7. Trained on the synthetic
+DNA of ``cis6270.data``. Prints the two distances, the worked arc on a DNA base, the tangent
+projection, and the measurements of a generated batch.
+
+    python lecture_5/fisher.py                     # train, then 64 exponential-map steps
+    python lecture_5/fisher.py --steps 20 --quiet  # a few seconds, for a first look
+"""
+
+# %% 01. Imports and configuration
+import math
+
+import torch
+import torch.nn as nn
+
+from cis6270.data import DNA_SIZE, decode_dna, dna_splits
+from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
+from cis6270.nets import SequenceDenoiser, count_parameters
+from cis6270.runner import common_parser, print_report, save_run, set_seed, train
+
+SMOOTHING = 0.03  # Label smoothing that holds a data corner off the boundary, where the metric blows up.
+
+
+# %% 02. Small probabilities under the straight-line distance
+#
+# Moving a base from (0.5, 0.5) to (0.4, 0.6) shifts a tenth of the mass between two evenly matched
+# letters. Moving from (0.01, 0.99) to (0.001, 0.999) divides the first probability by ten, which
+# matters for a model of sequences, and the Euclidean distance calls the second move eleven times
+# smaller.
+
+
+def euclidean_distance(p, q):
+    """The straight-line distance between two probability vectors."""
+    return float((p - q).norm())  # 0.1414 and 0.0127 for the two moves above.
+
+
+def fisher_rao_distance(p, q):
+    """d_FR(p, q) = 2 arccos(<sqrt(p), sqrt(q)>), at most pi for any two distributions."""
+    inner = float((p.sqrt() * q.sqrt()).sum().clamp(max=1.0))  # The angle between square roots.
+    return 2.0 * math.acos(inner)  # 0.201 and 0.137 for the same two moves.
+
+
+# %% 03. The Fisher information as a local metric
+#
+# With the probabilities as the parameters, the derivative of the log likelihood of outcome i with
+# respect to p^j is delta_ij/p^j, and the expectation over outcomes collapses the two deltas onto the
+# diagonal. Dividing by p^i makes a change in a rare coordinate count for more.
+
+
+def fisher_matrix(p):
+    """g_jk(p) = delta_jk / p^j, the Fisher information of a categorical distribution. [K, K]."""
+    return torch.diag(1.0 / p)  # Diagonal, with the reciprocal of each probability.
+
+
+def fisher_inner(u, v, p):
+    """<u, v>_p = sum_i u^i v^i / p^i, the Fisher-Rao inner product of two tangent vectors at p."""
+    return float((u * v / p).sum())  # Tangent means the coordinates sum to zero.
+
+
+# %% 04. The square-root map to the sphere
+#
+# The simplex constraint becomes a norm constraint, so the square roots lie on the unit sphere, in
+# the octant with every coordinate nonnegative. Differentiating the square root shows the squared
+# Euclidean length there is one quarter of the Fisher-Rao length, so the map scaled by two is an
+# isometry.
+
+
+def to_sphere(p):
+    """A probability vector becomes a unit vector. The clamp holds a zero off the boundary."""
+    return p.clamp(min=1e-8).sqrt()  # simplex -> sphere.
+
+
+def to_simplex(x):
+    """A unit vector with nonnegative coordinates becomes a probability vector."""
+    return x ** 2  # sphere -> simplex.
+
+
+def smooth(one_hot, amount=SMOOTHING):
+    """Move a corner inside the simplex, which the metric needs since it diverges at the boundary."""
+    K = one_hot.shape[-1]
+    return one_hot * (1.0 - K * amount) + amount  # (0.91, 0.03, 0.03, 0.03) at amount 0.03.
+
+
+# %% 05. Geodesics on the sphere
+#
+# The shortest path between two points of a sphere is the arc of the great circle through them, and
+# the conditional path is that arc traversed at constant speed. The logarithm map has length equal to
+# the remaining arc, so dividing by the remaining time leaves the speed constant along the path.
+
+
+def log_map(x, y):
+    """The tangent vector at x reaching y in unit time, of length equal to the arc from x to y."""
+    c = (x * y).sum(-1, keepdim=True).clamp(max=1 - 1e-6)  # The cosine of the angle between them.
+    v = y - c * x  # orthogonal to x.
+    return v * torch.arccos(c) / v.norm(dim=-1, keepdim=True)  # Rescaled to the arc length.
+
+
+def exp_map(x, w):
+    """The point reached by moving along the sphere from x in the direction w for a distance ||w||."""
+    n = w.norm(dim=-1, keepdim=True).clamp(min=1e-8)  # The length of the tangent vector.
+    return torch.cos(n) * x + torch.sin(n) * (w / n)  # Unit length by cos^2 + sin^2 = 1.
+
+
+def geodesic(x0, x1, t):
+    """x_t = exp_{x_0}(t log_{x_0}(x_1)), the point a fraction t of the way along the arc."""
+    return exp_map(x0, t * log_map(x0, x1))  # the conditional path.
+
+
+def conditional_velocity(xt, x1, t):
+    """u_t(x_t | x_1) = log_{x_t}(x_1)/(1 - t), the remaining arc over the remaining time."""
+    return log_map(xt, x1) / (1 - t)  # Constant in length along the whole path.
+
+
+# %% 06. The tangent projection and the regression loss
+#
+# A linear output layer emits an arbitrary vector of R^K, and a velocity on the sphere has to be
+# orthogonal to the current point. Subtracting the component along x enforces that in one line, and
+# the minimizer of the squared loss is the average of the conditional velocities passing through the
+# point, which is the averaging every path in this chapter performs.
+
+
+def project(xt, raw):
+    """v_theta = raw - <x_t, raw> x_t, which is orthogonal to x_t since ||x_t|| = 1."""
+    return raw - (xt * raw).sum(-1, keepdim=True) * xt  # Any architecture can be used.
+
+
+def loss(model, x0, x1, t):
+    """The squared distance between the projected network output and the geodesic target."""
+    xt = geodesic(x0, x1, t[:, None, None])  # [B, L, K] a point on the arc at the sampled time.
+    u = log_map(xt, x1) / (1 - t[:, None, None])  # [B, L, K] the remaining arc over the time left.
+    raw = model(xt, t)  # [B, L, K] unconstrained.
+    v = raw - (xt * raw).sum(-1, keepdim=True) * xt  # The tangent part of that output.
+    return ((v - u) ** 2).sum(-1).mean()  # Summed over letters, averaged over positions and batch.
+
+
+def velocity_network(length, width, K=DNA_SIZE, depth=2):
+    """The shared sequence denoiser with its lookup table replaced by a map on the relaxed token.
+
+    The head emits a vector of R^K here, which ``project`` turns into a velocity on the sphere.
+    """
+    model = SequenceDenoiser(K, length, width=width, depth=depth)
+    model.token_embedding = nn.Linear(K, width, bias=False)  # A point on the sphere, not an index.
+    return model
+
+
+# %% 07. Training on DNA, and sampling along the sphere
+#
+# Source points are uniform draws from the simplex, data points are smoothed corners, and the pair is
+# joined by the arc between their square roots. Generation starts from a fresh source point and takes
+# exponential-map steps, so every state stays on the sphere and no renormalization is needed.
+
+
+def source_points(batch, length, K=DNA_SIZE):
+    """Uniform draws from the simplex, mapped to the sphere. Returns [B, L, K]."""
+    gamma = torch._standard_gamma(torch.ones(batch, length, K))  # Dir(1) over the simplex.
+    return to_sphere(gamma / gamma.sum(-1, keepdim=True))  # [B, L, K] on the unit sphere.
+
+
+@torch.no_grad()
+def generate(model, batch, length, steps, K=DNA_SIZE):
+    """Integrate the learned field from the source to t = 1 with exponential-map steps. [B, L]."""
+    x = source_points(batch, length, K)  # [B, L, K].
+    h = 1.0 / steps
+    for k in range(steps):
+        t = torch.full((batch,), k * h)
+        v = project(x, model(x, t))  # [B, L, K] the tangent part of the network output.
+        x = exp_map(x, h * v)  # Move along the sphere, which keeps the norm at one.
+    return to_simplex(x).argmax(-1)  # Read the token off the largest coordinate.
+
+
+# %% 08. The report
+def main():
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.set_defaults(steps=1000, width=64, samples=64)
+    args = parser.parse_args()
+    set_seed(args.seed)
+
+    near = torch.tensor([0.5, 0.5]), torch.tensor([0.4, 0.6])  # Two evenly matched letters.
+    rare = torch.tensor([0.01, 0.99]), torch.tensor([0.001, 0.999])  # One rare letter, divided by ten.
+    report = {
+        "euclidean_near": round(euclidean_distance(*near), 4),  # 0.1414.
+        "euclidean_rare": round(euclidean_distance(*rare), 4),  # 0.0127, eleven times smaller.
+        "fisher_rao_near": round(fisher_rao_distance(*near), 3),  # 0.201.
+        "fisher_rao_rare": round(fisher_rao_distance(*rare), 3),  # 0.137, comparable.
+        "two_letters_are_orthogonal": round(
+            fisher_rao_distance(torch.tensor([1.0, 0.0]), torch.tensor([0.0, 1.0])), 4),  # pi.
+    }
+
+    # Section 5.21.3: a uniform base and a smoothed A, taken to the sphere.
+    p0 = torch.full((DNA_SIZE,), 1.0 / DNA_SIZE)  # The uniform base.
+    p1 = smooth(torch.eye(DNA_SIZE)[0])  # (0.91, 0.03, 0.03, 0.03).
+    x0, x1 = to_sphere(p0), to_sphere(p1)
+    report["fisher_diagonal"] = [round(float(v), 2) for v in torch.diagonal(fisher_matrix(p0))]
+    report["sphere_x1"] = [round(float(v), 3) for v in x1]  # (0.954, 0.173, 0.173, 0.173).
+    report["inner_product"] = round(float((x0 * x1).sum()), 3)  # 0.737.
+    report["angle"] = round(math.acos(float((x0 * x1).sum())), 3)  # 0.743.
+    report["fisher_rao_distance"] = round(fisher_rao_distance(p0, p1), 3)  # 1.485, twice the angle.
+
+    # Section 5.21.4: the midpoint of the arc, and one exponential-map step from it.
+    mid = geodesic(x0, x1, torch.tensor(0.5))
+    report["midpoint_sphere"] = [round(float(v), 3) for v in mid]  # (0.780, 0.361, 0.361, 0.361).
+    report["midpoint_simplex"] = [round(float(v), 3) for v in to_simplex(mid)]  # (0.609, 0.130, ...).
+    remaining = log_map(mid, x1)
+    report["log_map_from_mid"] = [round(float(v), 4) for v in remaining]  # (0.2323, -0.1672, ...).
+    report["remaining_arc"] = round(float(remaining.norm()), 4)  # 0.3713, half the total.
+    report["constant_speed"] = round(float(conditional_velocity(mid, x1, 0.5).norm()), 3)  # 0.743.
+    stepped = exp_map(mid, 0.1 * conditional_velocity(mid, x1, 0.5))  # Delta = 0.1 from t = 0.5.
+    report["step_sphere"] = [round(float(v), 3) for v in stepped]  # (0.824, 0.327, 0.327, 0.327).
+    report["step_simplex"] = [round(float(v), 3) for v in to_simplex(stepped)]  # (0.680, 0.107, ...).
+    report["step_norm"] = round(float(stepped.norm()), 10)  # One, because the step followed the arc.
+
+    # Section 5.21.5: the tangent projection of one raw network output.
+    raw = torch.tensor([0.3, -0.2, 0.5, 0.1])
+    report["raw_inner_product"] = round(float((mid * raw).sum()), 4)  # 0.3785.
+    tangent = project(mid, raw)
+    report["projected"] = [round(float(v), 3) for v in tangent]  # (0.005, -0.337, 0.363, -0.037).
+    report["projection_inner"] = f"{float((mid * tangent).sum()):.1e}"  # Zero to machine precision.
+
+    # Training on the shared DNA, then generation.
+    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
+    data = splits["train_tokens"]  # [N, L].
+    model = velocity_network(args.length, args.width)
+
+    def loss_fn(step):
+        rows = torch.randint(0, len(data), (args.batch_size,))  # One minibatch of data sequences.
+        target = to_sphere(smooth(torch.eye(DNA_SIZE)[data[rows]]))  # [B, L, K] smoothed corners.
+        start = source_points(args.batch_size, args.length)  # [B, L, K] the independent coupling.
+        t = torch.rand(args.batch_size) * 0.98  # Held off t = 1, where the target divides by zero.
+        return loss(model, start, target, t)
+
+    if not args.quiet:
+        print(f"parameters={count_parameters(model):,} smoothing={SMOOTHING}")
+    losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=args.quiet)
+    tokens = generate(model, args.samples, args.length, args.sample_steps)
+
+    report["final_loss"] = sum(losses[-20:]) / len(losses[-20:])
+    report["gc_fraction"] = gc_fraction(tokens)  # The training data sit at 0.577.
+    report["motif_fraction"] = motif_fraction(tokens)  # Sequences carrying the planted ACGT.
+    report["unique_fraction"] = unique_fraction(tokens)  # Catches a collapsed sampler.
+    report["token_entropy_nats"] = token_entropy(tokens, DNA_SIZE)  # At most log 4 = 1.386.
+    print_report("Fisher-Flow", report, decode_dna(tokens))
+    save_run(args.out, model, vars(args), report, losses, decode_dna(tokens))
+
+
+if __name__ == "__main__":
+    main()

lecture_5/guidance.py

@@ -0,0 +1,234 @@
+"""Guidance: steering a trained flow toward a property by tilting its rates.
+
+The distribution we want is the data distribution reweighted by the likelihood of a property, whose
+normalizer runs over every sequence. A rate is a ratio, so that constant cancels, and the conditioned
+process multiplies each rate by the ratio of a value function at the destination to its value at the
+current state. A first-order expansion of the logarithm of that ratio replaces one property
+evaluation per destination with one backward pass, and an exponent on the ratio gives one strength
+parameter, which the classifier-free form reaches from two trained velocities and no classifier.
+
+Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
+u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
+
+Chapter 5, Sections 5.24 to 5.29, with the listing of Section 5.33. Several objectives at once are in
+``lecture_5/multi_objective.py``. Trained on the synthetic DNA of ``cis6270.data``, with the GC label
+as the property. Prints the worked guided base, the guidance strength, the Taylor tilt, and the GC
+fraction of a batch generated with and without guidance.
+
+    python lecture_5/guidance.py                     # train, then sample at gamma = 0 and gamma = 2
+    python lecture_5/guidance.py --gamma 5           # a harder tilt toward the GC-rich label
+    python lecture_5/guidance.py --steps 20 --quiet  # a few seconds, for a first look
+"""
+
+# %% 01. Imports and configuration
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from cis6270.data import DNA_SIZE, decode_dna, dna_splits
+from cis6270.metrics import gc_fraction, motif_fraction, token_entropy
+from cis6270.nets import SequenceDenoiser, TimeEmbedding, count_parameters
+from cis6270.runner import common_parser, print_report, save_run, set_seed, train
+
+V = DNA_SIZE + 1  # Four bases plus the mask symbol, which the sampler starts from.
+MASK = DNA_SIZE  # Index 4 is the mask symbol m.
+
+
+def k(t):
+    """The schedule kappa_t, linear here."""
+    return t
+
+
+def kd(t):
+    """Its derivative kappa_dot_t, which the linear schedule holds at one."""
+    return torch.ones_like(t)
+
+
+# %% 02. Guidance as a tilt of the rates
+#
+# The conditioned velocity is the unguided one multiplied by the ratio of the property probability at
+# the destination to its value at the current state. A move toward a destination the property model
+# scores higher happens sooner, and a destination with rate zero stays at rate zero, so the tilt
+# reweights the support of the model without enlarging it.
+
+
+def tilt(rates, destination_scores, current_score, gamma=1.0):
+    """u^(gamma)_t(y | x, c) = [p(c | y, t)/p(c | x, t)]^gamma u_t(y, x)."""
+    return rates * (destination_scores / current_score) ** gamma  # gamma = 1 is the exact process.
+
+
+def h_value(denoiser, destination_scores):
+    """h_t(x) = E[p(c | X_1) | X_t = x], the posterior average of the property likelihood."""
+    return float((denoiser * destination_scores).sum())  # The exact value at a masked position.
+
+
+def classifier_free(conditional, unconditional, gamma):
+    """u_c^gamma u_u^(1-gamma), the geometric mixture that reaches the same one-parameter family."""
+    return conditional ** gamma * unconditional ** (1.0 - gamma)  # Linear in the logarithms.
+
+
+def reachable(positions, letters):
+    """D(S - 1) + 1, the number of states one jump can reach, since exactly one position changes."""
+    return positions * (letters - 1) + 1  # 2001 for a 100-residue peptide over 21 symbols.
+
+
+# %% 03. A first-order expansion of the log ratio
+#
+# Embedding a sequence as one-hot vectors makes a single-position change a sparse difference, and
+# expanding the log property probability around the current sequence collapses the inner product onto
+# the difference of two gradient coordinates. One forward and one backward pass then serve every
+# destination.
+
+
+def taylor_log_ratio(gradient, current):
+    """[grad log p_phi]_{i,b} - [grad log p_phi]_{i,a} at one position. gradient: [K] -> [K]."""
+    return gradient - gradient[current]  # Zero at the letter in place, whose ratio is one.
+
+
+# %% 04. The guided sampler
+#
+# One guided step differs from the unguided one at a single line, the factor exp(gamma adv). The
+# one-hot encoding has gradients switched on, so the property network can be differentiated against
+# letters that are not currently in place, and summing over the batch before differentiating shares
+# one backward pass across it. The denoiser is read over the four bases alone, since the data token
+# is never the mask, which holds the rate back into the mask at zero and the exit rate at
+# kappa_dot/(1 - kappa) when gamma is zero.
+
+
+def guided_step(model, prop, xt, t, h, gamma=1.0):
+    """One Euler step of the tilted velocity at every position at once. xt: [B, D] -> [B, D]."""
+    oh = F.one_hot(xt, V).to(torch.get_default_dtype()).requires_grad_()
+    logp = prop(oh, t).sum()  # property on noisy x
+    g = torch.autograd.grad(logp, oh)[0]  # one backward pass
+    adv = g - (oh * g).sum(-1, keepdim=True)  # Taylor ratio
+    with torch.no_grad():
+        bases = model(xt, t)[..., :DNA_SIZE].softmax(-1)  # denoiser
+        p1 = F.pad(bases, (0, V - DNA_SIZE))  # the data token is never the mask
+        rate = kd(t) / (1 - k(t)).clamp(min=1e-4)
+        u = rate[:, None, None] * p1 * torch.exp(gamma * adv)  # every rate tilted
+        u.scatter_(-1, xt[..., None], 0.0)  # no self-transition
+        jump = torch.rand(xt.shape) < h * u.sum(-1)
+        y = torch.distributions.Categorical(probs=u).sample()
+    return torch.where(jump, y, xt)
+
+
+class PropertyModel(nn.Module):
+    """log p_phi(c | x, t) from a one-hot sequence and a time, trained on noisy intermediate states."""
+
+    def __init__(self, length, width=64):
+        super().__init__()
+        self.time_embedding = TimeEmbedding(width, scale=1.0)  # The property depends on the time.
+        self.body = nn.Sequential(
+            nn.Linear(length * V + width, width), nn.SiLU(), nn.Linear(width, 1),
+        )
+
+    def forward(self, oh, t):
+        """oh: [B, D, V] one-hot and differentiable, t: [B] -> [B] log probability of the property."""
+        h = torch.cat([oh.flatten(1), self.time_embedding(t)], dim=1)  # [B, D V + width].
+        return F.logsigmoid(self.body(h)).squeeze(-1)  # [B] log p_phi(c | x, t).
+
+
+# %% 05. Training the denoiser and the property model
+#
+# The denoiser is the one built earlier in this chapter, trained from the mask source. The property
+# model reads the same noisy sequences and predicts the GC label of the clean one, so it is a
+# function of an intermediate state and not of a clean sequence.
+
+
+def corrupt(x1, t):
+    """A draw from the conditional path with the all-mask source. x1: [B, D], t: [B] -> [B, D]."""
+    heads = torch.rand(x1.shape) < k(t)[:, None]  # Heads with probability kappa_t at each position.
+    return torch.where(heads, x1, torch.full_like(x1, MASK))  # Data token on heads, mask on tails.
+
+
+@torch.no_grad()
+def generate(model, prop, batch, length, steps, gamma):
+    """Run the guided sampler from all masks to t = 1. Returns [B, D]."""
+    x = torch.full((batch, length), MASK, dtype=torch.long)  # The all-mask source.
+    grid = torch.linspace(0.0, 1.0, steps + 1)
+    for t, s in zip(grid[:-1], grid[1:]):
+        with torch.enable_grad():  # The property gradient needs a graph inside the no-grad sampler.
+            x = guided_step(model, prop, x, t.expand(batch), float(s - t), gamma)
+    p1 = model(x, torch.ones(batch))[..., :DNA_SIZE].softmax(-1)  # The data token is never the mask.
+    drawn = torch.distributions.Categorical(probs=p1).sample()  # lambda_t reaching one on the last step.
+    return torch.where(x == MASK, drawn, x)  # Any position still masked is drawn from the denoiser.
+
+
+# %% 06. The report
+def main():
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.add_argument("--gamma", type=float, default=2.0, help="guidance strength")
+    parser.set_defaults(steps=1000, width=64, samples=64)
+    args = parser.parse_args()
+    set_seed(args.seed)
+
+    # Section 5.27: a masked base at t = 0.6 with the chapter's denoiser and property scores.
+    p1 = torch.tensor([0.1, 0.2, 0.3, 0.4], dtype=torch.double)  # The denoiser over A, C, G, T.
+    rates = p1 / (1.0 - 0.6)  # The unguided unmasking rates, 2.5 times the denoiser.
+    scores = torch.tensor([0.10, 0.35, 0.55, 0.25], dtype=torch.double)  # p(c | a, t).
+    guided = tilt(rates, scores, 0.20)  # The property model scores the masked state at 0.20.
+    report = {
+        "unguided_rates": [round(float(r), 4) for r in rates],  # (0.25, 0.50, 0.75, 1.00).
+        "unguided_exit_rate": round(float(rates.sum()), 4),  # 2.50.
+        "likelihood_ratios": [round(float(r), 4) for r in scores / 0.20],  # (0.50, 1.75, 2.75, 1.25).
+        "guided_rates": [round(float(r), 4) for r in guided],  # (0.125, 0.875, 2.063, 1.250).
+        "guided_exit_rate": round(float(guided.sum()), 4),  # 4.3125.
+        "guided_destination": [round(float(r), 3) for r in guided / guided.sum()],  # (0.029, ...).
+        "bayes_products": [round(float(r), 3) for r in p1 * scores],  # (0.010, 0.070, 0.165, 0.100).
+        "exact_h": round(h_value(p1, scores), 4),  # 0.345.
+        "exit_rate_at_exact_h": round(float(tilt(rates, scores, h_value(p1, scores)).sum()), 4),  # 2.5.
+        "reachable_states": reachable(100, 21),  # 2001 destinations for a 100-residue peptide.
+    }
+
+    # Section 5.28: the same base at guidance strength two, and the classifier-free form.
+    harder = tilt(rates, scores, 0.20, gamma=2.0)
+    report["gamma_2_rates"] = [round(float(r), 4) for r in harder]  # (0.0625, 1.531, 5.672, 1.563).
+    report["gamma_2_exit_rate"] = round(float(harder.sum()), 4)  # 8.828.
+    report["gamma_2_destination"] = [round(float(r), 3) for r in harder / harder.sum()]
+    mixed = classifier_free(tilt(rates, scores, 0.20), rates, gamma=2.0)
+    report["classifier_free_agrees"] = bool(torch.allclose(mixed, harder))  # The same family.
+
+    # Section 5.29: the Taylor tilt against the exact ratio.
+    gradient = torch.tensor([-2.10, -1.15, -0.75, -1.30, -1.61], dtype=torch.double)  # With the mask.
+    approximate = taylor_log_ratio(gradient, MASK)[:DNA_SIZE]  # The mask is the letter in place.
+    report["taylor_log_ratios"] = [round(float(r), 3) for r in approximate]  # (-0.49, 0.46, ...).
+    report["exact_log_ratios"] = [round(float(r), 3) for r in (scores / 0.20).log()]  # (-0.693, ...).
+    taylor_rates = rates * approximate.exp()
+    report["taylor_destination"] = [round(float(r), 3) for r in taylor_rates / taylor_rates.sum()]
+
+    # Training the denoiser and the property model on the shared DNA, then sampling twice.
+    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
+    data, labels = splits["train_tokens"], splits["train_labels"]  # [N, L] and [N] the GC label.
+    model = SequenceDenoiser(V, args.length, width=args.width, depth=2)
+    prop = PropertyModel(args.length, width=args.width)
+
+    def loss_fn(step):
+        rows = torch.randint(0, len(data), (args.batch_size,))  # One minibatch of data sequences.
+        x1 = data[rows]  # [B, L] the sequences the denoiser has to recover.
+        y = labels[rows].to(torch.get_default_dtype())  # [B] the property to steer toward.
+        t = torch.rand(args.batch_size)  # One time per sequence.
+        xt = corrupt(x1, t)  # [B, L] a draw from the conditional path.
+        denoise = F.cross_entropy(model(xt, t)[..., :DNA_SIZE].transpose(1, 2), x1)
+        logp = prop(F.one_hot(xt, V).to(torch.get_default_dtype()), t)  # [B] the property on x_t.
+        property_loss = -(y * logp + (1 - y) * torch.log1p(-logp.exp().clamp(max=1 - 1e-6))).mean()
+        return denoise + property_loss  # The two networks share one optimizer here.
+
+    if not args.quiet:
+        print(f"denoiser={count_parameters(model):,} property={count_parameters(prop):,}")
+    losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=args.quiet,
+                   extra_parameters=list(prop.parameters()))
+    plain = generate(model, prop, args.samples, args.length, args.sample_steps, gamma=0.0)
+    steered = generate(model, prop, args.samples, args.length, args.sample_steps, gamma=args.gamma)
+
+    report["final_loss"] = sum(losses[-20:]) / len(losses[-20:])
+    report["gc_fraction_unguided"] = gc_fraction(plain)  # The training data sit at 0.577.
+    report[f"gc_fraction_gamma_{args.gamma:g}"] = gc_fraction(steered)  # Guidance raises it.
+    report["motif_fraction_unguided"] = motif_fraction(plain)
+    report["token_entropy_unguided"] = token_entropy(plain, DNA_SIZE)
+    print_report("guidance", report, decode_dna(steered))
+    save_run(args.out, model, vars(args), report, losses, decode_dna(steered))
+
+
+if __name__ == "__main__":
+    main()

lecture_5/gumbel.py

@@ -0,0 +1,226 @@
+"""Gumbel-Softmax Flow Matching: the path comes from the way a token is sampled.
+
+Sampling a categorical is a deterministic function of noise. We add a Gumbel draw to each log
+probability and take the largest. Replacing the argmax by a softmax turns the winner into a point of
+the simplex, and cooling the temperature slides that point from the center toward a corner. With the
+noise held fixed the trajectory is a smooth curve, its time derivative is the replicator equation of
+evolutionary dynamics, and setting the noise to zero at inference collapses the velocity to a
+multiple of the direction toward the data corner.
+
+Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
+u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
+
+Chapter 5, Sections 5.22.1 to 5.22.6. Trained on the synthetic DNA of ``cis6270.data``. Prints the
+Gumbel-max frequencies, the annealed path at four times, the replicator velocity, the marginal
+velocity late in a trajectory, and the measurements of a generated batch.
+
+    python lecture_5/gumbel.py                     # train, then 64 Euler steps on the simplex
+    python lecture_5/gumbel.py --steps 20 --quiet  # a few seconds, for a first look
+"""
+
+# %% 01. Imports and configuration
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from cis6270.data import DNA_SIZE, decode_dna, dna_splits
+from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
+from cis6270.nets import SequenceDenoiser, count_parameters
+from cis6270.runner import common_parser, print_report, save_run, set_seed, train
+
+TAU_MAX = 10.0  # The temperature at t = 0, hot enough that every letter holds nearly the same mass.
+LAMBDA = 3.0  # The cooling rate, so tau(1) = 10 exp(-3) = 0.498.
+BETA = 2.0  # The noise scale; a larger beta makes the paths from one data token spread less.
+
+
+# %% 02. The Gumbel-max trick
+#
+# A Gumbel variable is the negative logarithm of an exponential clock, so the trick is Chapter 4's
+# race between clocks written with logarithms, where the winner is an argmax and the rates are
+# logits. All of the randomness sits in the noise, and the logits are ours to move.
+
+
+def gumbel(shape, generator=None):
+    """g ~ Gumbel(0, 1), as minus the log of minus the log of a uniform draw."""
+    u = torch.rand(shape, generator=generator).clamp(min=1e-12, max=1 - 1e-12)
+    return -torch.log(-torch.log(u))  # The distribution function is exp(-exp(-x)).
+
+
+def gumbel_max(probabilities, draws, generator=None):
+    """Frequencies of argmax_k [log pi_k + g_k] over ``draws`` samples. Returns [K]."""
+    logits = probabilities.log()[None, :] + gumbel((draws, len(probabilities)), generator)
+    counts = torch.bincount(logits.argmax(-1), minlength=len(probabilities)).double()
+    return counts / counts.sum()  # Approaches pi itself.
+
+
+# %% 03. The softmax relaxation and a cooling temperature
+#
+# A softmax of the noisy logits divided by a temperature is a point of the simplex. As the
+# temperature falls the point approaches the one-hot of the argmax, and as it grows the point
+# approaches the center. The temperature never reaches zero, so the path stays in the interior and
+# its density stays defined.
+
+
+def temperature(t, tau_max=TAU_MAX, rate=LAMBDA):
+    """tau(t) = tau_max exp(-lambda t), which is 10 at t = 0 and 0.498 at t = 1."""
+    return tau_max * math.exp(-rate * t)  # 2.231 at t = 0.5 and 0.672 at t = 0.9.
+
+
+# %% 04. The annealed path toward a data token
+#
+# Putting the data token into the logits and letting the cooling do the rest gives the conditional
+# path. The noise is drawn once and held for the whole trajectory, so the only quantity that changes
+# with time is the temperature and the point moves continuously toward the data corner.
+
+
+def annealed_path(logits, t):
+    """psi_t = softmax((e_k + g/beta)/tau(t)), with the logits already assembled. Returns [..., K]."""
+    return torch.softmax(logits / temperature(t), dim=-1)  # A point of the simplex at every time.
+
+
+def path_logits(targets, noise, K=DNA_SIZE, beta=BETA):
+    """e_k + g/beta at every position. targets: [B, L], noise: [B, L, K] -> [B, L, K]."""
+    return F.one_hot(targets, K) + noise / beta  # The data indicator plus the scaled Gumbel draw.
+
+
+# %% 05. Differentiating the path: the replicator equation
+#
+# The only time dependence sits in one over tau, which grows at relative rate lambda. Composing with
+# the softmax Jacobian gives the current mass times the gap between a letter's logit and the
+# mass-weighted average of all of them, which is the replicator equation with fitness
+# lambda ell_i/tau(t).
+
+
+def replicator_velocity(x, logits, t, rate=LAMBDA):
+    """u_{t,i} = (lambda/tau(t)) x_i (ell_i - sum_j x_j ell_j). x, logits: [..., K] -> [..., K]."""
+    average = (x * logits).sum(-1, keepdim=True)  # The mass-weighted average of the logits.
+    return (rate / temperature(t)) * x * (logits - average)  # Coordinates summing to zero.
+
+
+# %% 06. The inference-time velocity
+#
+# The sampler has only the current point, so the noise is set to zero. The logits become the data
+# indicator, the bracket becomes the gap to the data coordinate, and the velocity collapses to a
+# multiple of the direction toward the corner, with the factor x_k(1 - x_k) in the k-th component.
+
+
+def inference_velocity(x, k, t, rate=LAMBDA):
+    """u_t(x | x_1 = e_k) = (lambda/tau(t)) x_k (e_k - x). x: [..., K] -> [..., K]."""
+    corner = F.one_hot(torch.tensor(k), x.shape[-1]).to(x)  # The data corner.
+    return (rate / temperature(t)) * x[..., k, None] * (corner - x)  # Zero once x_k reaches zero.
+
+
+def marginal_velocity(x, p_hat, t, rate=LAMBDA):
+    """The posterior average over corners. x, p_hat: [..., K] -> [..., K]."""
+    weight = (rate / temperature(t)) * x * p_hat  # [..., K] the coefficient of each corner.
+    return weight - weight.sum(-1, keepdim=True) * x  # sum_k w_k (e_k - x), collected.
+
+
+# %% 07. Training and the marginal velocity
+#
+# Cross-entropy on the data token trains the denoiser, as it did for the mixture paths and for the
+# Dirichlet path, and the conditional velocity appears nowhere in the gradient. At sampling time the
+# marginal velocity is the posterior average, the fourth appearance of the same rule in this chapter.
+
+
+def relaxed_denoiser(length, width, K=DNA_SIZE, depth=2):
+    """The shared sequence denoiser with its lookup table replaced by a map on the relaxed token."""
+    model = SequenceDenoiser(K, length, width=width, depth=depth)
+    model.token_embedding = nn.Linear(K, width, bias=False)  # A point of the simplex, not an index.
+    return model
+
+
+@torch.no_grad()
+def generate(model, batch, length, steps, K=DNA_SIZE):
+    """Integrate the marginal velocity from the center of the simplex to t = 1. Returns [B, L]."""
+    x = annealed_path(gumbel((batch, length, K)) / BETA, 0.0)  # The source, with no data indicator.
+    h = 1.0 / steps
+    for k in range(steps):
+        t = k * h
+        p_hat = model(x, torch.full((batch,), t)).softmax(-1)  # [B, L, K] the denoiser.
+        x = x + h * marginal_velocity(x, p_hat, t)  # One Euler step along the marginal velocity.
+        x = x.clamp(min=1e-8)  # A large step can drive a coordinate below zero.
+        x = x / x.sum(-1, keepdim=True)  # Project back onto the simplex.
+    return x.argmax(-1)  # Read the token off the largest coordinate.
+
+
+# %% 08. The report
+def main():
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.add_argument("--draws", type=int, default=400_000, help="draws in the Gumbel-max check")
+    parser.set_defaults(steps=1000, width=64, samples=64)
+    args = parser.parse_args()
+    set_seed(args.seed)
+
+    pi = torch.tensor([0.5, 0.2, 0.2, 0.1])  # The categorical of the chapter's quick check.
+    report = {
+        "gumbel_max_frequencies": [round(float(v), 4) for v in gumbel_max(pi, args.draws)],
+        "target_probabilities": [round(float(v), 4) for v in pi],  # What the frequencies approach.
+        "tau_at_0": round(temperature(0.0), 4),  # 10.
+        "tau_at_1": round(temperature(1.0), 4),  # 0.498.
+    }
+
+    # Section 5.22.3: one base annealed toward G with the chapter's fixed noise.
+    scaled = torch.tensor([0.0, -0.45, 0.40, -0.05])  # g/beta over A, C, G, T.
+    logits = F.one_hot(torch.tensor(2), DNA_SIZE) + scaled  # The data indicator at G plus the noise.
+    report["path_logits"] = [round(float(v), 2) for v in logits]  # (0, -0.45, 1.40, -0.05).
+    for t in (0.0, 0.5, 0.9, 1.0):
+        point = annealed_path(logits, t)
+        report[f"path_at_{t}"] = [round(float(v), 4) for v in point]
+
+    # Section 5.22.4: the replicator velocity at the midpoint.
+    x = annealed_path(logits, 0.5)  # (0.2142, 0.1751, 0.4012, 0.2095).
+    report["tau_at_0.5"] = round(temperature(0.5), 4)  # 2.231.
+    report["rate_over_tau"] = round(LAMBDA / temperature(0.5), 3)  # 1.345.
+    report["average_logit"] = round(float((x * logits).sum()), 4)  # 0.4724.
+    noisy = replicator_velocity(x, logits, 0.5)
+    report["replicator_velocity"] = [round(float(v), 3) for v in noisy]  # (-0.136, -0.217, 0.5, -0.147).
+    report["replicator_sums_to"] = f"{float(noisy.sum()):.1e}"  # Zero, so the step stays on the simplex.
+
+    # Section 5.22.5: the same point with the noise set to zero.
+    clean = inference_velocity(x, 2, 0.5)
+    report["inference_velocity"] = [round(float(v), 3) for v in clean]  # (-0.116, -0.094, 0.323, ...).
+
+    # Section 5.22.6: one Euler step late in the trajectory.
+    late = torch.tensor([0.138, 0.084, 0.647, 0.131])  # The base at t = 0.8.
+    p_hat = torch.tensor([0.05, 0.10, 0.70, 0.15])  # The denoiser output there.
+    report["tau_at_0.8"] = round(temperature(0.8), 4)  # 0.9072.
+    coefficients = (LAMBDA / temperature(0.8)) * late * p_hat
+    report["corner_coefficients"] = [round(float(v), 4) for v in coefficients]  # (0.0228, ..., 1.4976).
+    v = marginal_velocity(late, p_hat, 0.8)
+    report["marginal_velocity"] = [round(float(c), 3) for c in v]  # (-0.200, -0.108, 0.454, -0.146).
+    report["euler_step_h_0.01"] = [round(float(c), 4) for c in late + 0.01 * v]  # (0.1360, ...).
+
+    # Training on the shared DNA, then generation.
+    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
+    data = splits["train_tokens"]  # [N, L].
+    model = relaxed_denoiser(args.length, args.width)
+
+    def loss_fn(step):
+        rows = torch.randint(0, len(data), (args.batch_size,))  # One minibatch of data sequences.
+        x1 = data[rows]  # [B, L] the letters the denoiser has to recover.
+        t = torch.rand(args.batch_size)  # One time per sequence, uniform on [0, 1].
+        noise = gumbel((args.batch_size, args.length, DNA_SIZE))  # Held fixed along one trajectory.
+        cooled = path_logits(x1, noise) / torch.tensor(
+            [temperature(float(s)) for s in t])[:, None, None]  # Each sequence at its own temperature.
+        xt = cooled.softmax(-1)  # [B, L, K] a point on the conditional path.
+        return F.cross_entropy(model(xt, t).transpose(1, 2), x1)  # -log <x_theta, x_1>.
+
+    if not args.quiet:
+        print(f"parameters={count_parameters(model):,} tau_max={TAU_MAX} lambda={LAMBDA}")
+    losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=args.quiet)
+    tokens = generate(model, args.samples, args.length, args.sample_steps)
+
+    report["final_loss_nats"] = sum(losses[-20:]) / len(losses[-20:])
+    report["gc_fraction"] = gc_fraction(tokens)  # The training data sit at 0.577.
+    report["motif_fraction"] = motif_fraction(tokens)  # Sequences carrying the planted ACGT.
+    report["unique_fraction"] = unique_fraction(tokens)  # Catches a collapsed sampler.
+    report["token_entropy_nats"] = token_entropy(tokens, DNA_SIZE)  # At most log 4 = 1.386.
+    print_report("Gumbel-Softmax flow matching", report, decode_dna(tokens))
+    save_run(args.out, model, vars(args), report, losses, decode_dna(tokens))
+
+
+if __name__ == "__main__":
+    main()

lecture_5/kinetic_optimal.py

@@ -0,0 +1,231 @@
+"""Kinetic-optimal paths: choosing the velocity, and then the path, by minimizing an energy.
+
+Several velocities generate the same path and several paths connect the same endpoints, so a number
+is needed to choose among them. Written in the flux, both constraints are linear and the energy is
+convex, and the minimizer is the probability of the destination times the positive part of a
+difference of logarithmic growth rates. The mixture velocity of this chapter turns out to be that
+minimizer for the path it generates. Picking a distance on the tokens and cooling a softmax of
+negative distances then produces a path that moves probability only toward tokens closer to the data
+token.
+
+Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
+u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
+
+Chapter 5, Sections 5.16.1 to 5.16.5 and 5.16.7, with the geodesic schedule of 5.16.6. The data marginal
+of the mask-path check is the letter frequency of the synthetic DNA of ``cis6270.data``. Prints the
+kinetic-optimal rates, the worked metric path for DNA, and the energy of two schedules.
+
+    python lecture_5/kinetic_optimal.py                     # 600 quadrature points per integral
+    python lecture_5/kinetic_optimal.py --steps 20 --quiet  # a coarse rule, visibly less accurate
+"""
+
+# %% 01. Imports and configuration
+import math
+
+import torch
+
+from cis6270.data import DNA, DNA_SIZE, dna_splits
+from cis6270.runner import common_parser, print_report, save_run, set_seed
+
+
+torch.set_default_dtype(torch.double)  # Everything here is arithmetic, so keep full precision.
+MASK = DNA_SIZE  # Index 4 is the mask symbol m, which the mask path visits.
+TRANSITION = torch.tensor([[0, 2, 1, 2], [2, 0, 2, 1], [1, 2, 0, 2], [2, 1, 2, 0]], dtype=torch.double)
+# d = 1 for a transition, A for G or C for T, and d = 2 for a transversion across the two classes.
+
+
+# %% 02. The flux and its energy
+#
+# The flux along an edge is the velocity times the probability of the state it leaves, and the master
+# equation is linear in it. The cost is a weight over the probability of the current state, times the
+# squared flux, which in the velocity reads as an expected squared rate under p_t.
+
+
+def flux(u, p):
+    """j_t(y, x) = u_t(y, x) p_t(x). u: [K, K], p: [K] -> [K, K]."""
+    return u * p[None, :]  # Each column scaled by the probability of its current state.
+
+
+def energy_density(u, p, weight=None):
+    """sum_{y != x} w_t(y, x) u_t(y, x)^2 p_t(x), the cost at one time written in the velocity."""
+    off = u - torch.diag(torch.diagonal(u))  # The diagonal holds no flux of its own.
+    weight = torch.ones_like(u) if weight is None else weight  # w = 1 gives the expected squared rate.
+    return float((weight * off ** 2 * p[None, :]).sum())
+
+
+# %% 03. The velocity of least energy
+#
+# Minimizing the energy edge by edge under the master equation, with the weight 1/p_t(y) that makes
+# the cost symmetric in the two ends of an edge, the multiplier comes out as the logarithmic growth
+# rate and the flux becomes p_t(y) p_t(x) times the positive part of its difference. Dividing by
+# p_t(x) returns the velocity.
+
+
+def kinetic_optimal(p, p_dot):
+    """u*_t(y, x) = p_t(y) [d/dt log p_t(y) - d/dt log p_t(x)]_+. p: [K], p_dot: [K] -> [K, K]."""
+    growth = torch.where(p > 0, p_dot / p.clamp(min=1e-300), torch.zeros_like(p))  # d/dt log p_t.
+    difference = growth[:, None] - growth[None, :]  # [K, K] destination minus current state.
+    u = p[:, None] * difference.clamp(min=0.0)  # Probability moves only toward the faster grower.
+    u = u * (p[None, :] > 0)  # A state with no probability sends none.
+    u.fill_diagonal_(0.0)
+    return u - torch.diag(u.sum(dim=0))  # The two rules fix the diagonal.
+
+
+# %% 04. Kinetic optimality of the mixture path
+#
+# On the conditional mixture path the two growth rates are kappa_dot/kappa at the data token and
+# minus kappa_dot/(1 - kappa) at the source token, and their difference collapses to the conditional
+# velocity built earlier in this chapter. The marginal mask path behaves the same way.
+
+
+def mask_path(q, kappa_t):
+    """p_t and p_dot_t over (A, C, G, T, m) for the mask source with data marginal q. Returns two [5]."""
+    p = torch.cat([kappa_t * q, torch.tensor([1.0 - kappa_t])])  # Tokens share kappa_t, m holds 1-kappa.
+    p_dot = torch.cat([q.clone(), torch.tensor([-1.0])])  # kappa_dot = 1 on the linear schedule.
+    return p, p_dot
+
+
+def mixture_rate(kappa_t, kappa_dot=1.0):
+    """kappa_dot_t/(1 - kappa_t), the conditional velocity out of the source token."""
+    return kappa_dot / (1.0 - kappa_t)  # 2 at t = 0.5 on the linear schedule.
+
+
+# %% 05. A distance on the tokens
+#
+# The formula for u* accepts any path, so we can design one. Cooling a softmax of negative distances
+# concentrates the conditional path on the data token, and the normalizing constant cancels between
+# the two logarithms, leaving the probability of the destination times beta_dot times the distance
+# the move gains.
+
+
+def metric_path(x1, beta_t, distance=TRANSITION):
+    """p_t(. | x_1) = softmax(-beta_t d(., x_1)). Returns [K]."""
+    return torch.softmax(-beta_t * distance[:, x1], dim=0)  # Uniform at beta = 0, a point mass later.
+
+
+def metric_velocity(x1, beta_t, beta_dot, distance=TRANSITION):
+    """u*_t(y, x | x_1) = p_t(y | x_1) beta_dot_t [d(x, x_1) - d(y, x_1)]_+. Returns [K, K]."""
+    p = metric_path(x1, beta_t, distance)  # [K] the path at this time.
+    d = distance[:, x1]  # [K] distance from each token to the data token.
+    gain = (d[None, :] - d[:, None]).clamp(min=0.0)  # [K, K] what the move from x to y gains.
+    u = p[:, None] * beta_dot * gain  # Faster when the step in distance is larger.
+    u.fill_diagonal_(0.0)
+    return u - torch.diag(u.sum(dim=0))  # Columns sum to zero.
+
+
+# %% 06. A worked metric path for DNA
+#
+# A transition swaps a purine for a purine or a pyrimidine for a pyrimidine, and a transversion
+# crosses between the two classes, so d = 1 and d = 2 record which mistakes are small. With
+# beta_t = t/(1 - t) and the data token G, the path at t = 0.5 favors the purine A over the two
+# pyrimidines, and the rates out of C run only toward A and G.
+
+
+def beta_schedule(t):
+    """beta_t = t/(1 - t) with beta_dot_t = 1/(1 - t)^2, which satisfies beta_0 = 0."""
+    return t / (1.0 - t), 1.0 / (1.0 - t) ** 2  # 1 and 4 at t = 0.5.
+
+
+def path_derivative(x1, t, h=1e-6, distance=TRANSITION):
+    """d p_t/dt by a central difference, for checking the velocity against the path it generates."""
+    forward, backward = metric_path(x1, beta_schedule(t + h)[0], distance), metric_path(
+        x1, beta_schedule(t - h)[0], distance)
+    return (forward - backward) / (2.0 * h)  # [K], compared against u* p_t.
+
+
+# %% 07. The energy of two schedules
+#
+# On a two-state conditional mixture path the energy is the time integral of kappa_dot squared over
+# kappa times one minus kappa. The linear schedule has a non-integrable singularity at each endpoint,
+# so its energy grows with the truncation and never settles. The geodesic schedule has the constant
+# integrand pi squared, which is the squared arc length of a quarter great circle.
+
+
+def schedule_energy(kappa, kappa_dot, steps, cut=1e-6):
+    """The integral of kappa_dot^2/(kappa (1 - kappa)) over [cut, 1 - cut], by a substituted midpoint.
+
+    The grid is uniform in s with t = 1/(1 + exp(-s)), whose Jacobian t(1 - t) cancels the endpoint
+    singularity exactly for the linear schedule and tames it for any other.
+    """
+    span = math.log((1.0 - cut) / cut)  # The s range that the truncation [cut, 1 - cut] maps to.
+    h, total = 2.0 * span / steps, 0.0
+    for k in range(steps):
+        s = -span + (k + 0.5) * h  # Midpoint of the k-th cell.
+        t = 1.0 / (1.0 + math.exp(-s))  # Back to the time variable.
+        k_t = min(max(kappa(t), 1e-12), 1.0 - 1e-12)  # Hold the schedule off both endpoints.
+        total += h * t * (1.0 - t) * kappa_dot(t) ** 2 / (k_t * (1.0 - k_t))
+    return total
+
+
+LINEAR = (lambda t: t, lambda t: 1.0)  # kappa_t = t.
+GEODESIC = (lambda t: math.sin(math.pi * t / 2) ** 2,  # The arc of Section 5.16.6 read as a schedule.
+            lambda t: (math.pi / 2) * math.sin(math.pi * t))
+
+
+def geodesic_schedule(t, source_probability=0.0):
+    """kappa_t(x_1) = 1 - sin^2((1 - t) Omega)/sin^2 Omega with Omega = arccos(sqrt(p(x_1)))."""
+    omega = math.acos(math.sqrt(source_probability))  # pi/2 under the mask source.
+    return 1.0 - math.sin((1.0 - t) * omega) ** 2 / math.sin(omega) ** 2
+
+
+# %% 08. The report
+def main():
+    parser = common_parser(__doc__.split("\n")[0])
+    args = parser.parse_args()
+    set_seed(args.seed)
+
+    # Section 5.16.3: the mask path at t = 0.5 with the chapter's data marginal.
+    q = torch.tensor([0.1, 0.6, 0.2, 0.1])  # [4] over A, C, G, T.
+    p, p_dot = mask_path(q, kappa_t=0.5)
+    u_star = kinetic_optimal(p, p_dot)
+    report = {
+        "mask_path_ACGTm": [round(float(v), 4) for v in p],  # (0.05, 0.30, 0.10, 0.05, 0.50).
+        "growth_at_mask": round(float(p_dot[MASK] / p[MASK]), 4),  # -2.
+        "growth_at_token": round(float(p_dot[1] / p[1]), 4),  # +2.
+        "rate_mask_to_C": round(float(u_star[1, MASK]), 4),  # 0.30 times 4, which is 1.2.
+        "mixture_rate_to_C": round(mixture_rate(0.5) * float(q[1]), 4),  # 2 times 0.6, the same.
+        "rate_between_tokens": round(float(u_star[1, 2]), 6),  # Zero, the no-revision mask path.
+        "generates_the_path": bool(torch.allclose(u_star @ p, p_dot, atol=1e-9)),
+    }
+
+    # The same check with the letter frequency of the synthetic DNA as the data marginal.
+    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
+    counts = torch.bincount(splits["train_tokens"].flatten(), minlength=DNA_SIZE).double()
+    frequency = counts / counts.sum()  # [4] the GC-rich background the lecture shares.
+    p_dna, p_dot_dna = mask_path(frequency, kappa_t=0.5)
+    u_dna = kinetic_optimal(p_dna, p_dot_dna)
+    report["dna_marginal"] = [round(float(v), 4) for v in frequency]
+    report["dna_rate_mask_to_G"] = round(float(u_dna[2, MASK]), 4)  # 2 times the frequency of G.
+    report["dna_energy_density"] = round(energy_density(u_dna, p_dna), 4)  # Expected squared rate.
+    report["dna_flux_mask_to_G"] = round(float(flux(u_dna, p_dna)[2, MASK]), 4)  # Rate times p_t(m).
+
+    # Section 5.16.5: the metric path toward G at t = 0.5.
+    g, c = DNA.index("G"), DNA.index("C")
+    beta_t, beta_dot = beta_schedule(0.5)  # 1 and 4.
+    p_metric = metric_path(g, beta_t)
+    u_metric = metric_velocity(g, beta_t, beta_dot)
+    report["metric_path_at_0.5"] = [round(float(v), 4) for v in p_metric]
+    report["rate_C_to_A"] = round(float(u_metric[0, c]), 3)  # 0.898.
+    report["rate_C_to_G"] = round(float(u_metric[g, c]), 3)  # 4.882.
+    report["rate_A_to_G"] = round(float(u_metric[g, 0]), 3)  # 2.441.
+    report["rate_out_of_C"] = round(float(-u_metric[c, c]), 3)  # 5.780.
+    step = torch.eye(DNA_SIZE)[c] + 0.05 * u_metric[:, c]  # One Euler step of size h = 0.05 from C.
+    report["euler_step_from_C"] = [round(float(v), 4) for v in step]
+    report["p_dot_from_velocity"] = [round(float(v), 4) for v in u_metric @ p_metric]
+    report["p_dot_from_path"] = [round(float(v), 4) for v in path_derivative(g, 0.5)]
+
+    # Section 5.16.7: the energy of the two schedules, and the geodesic schedule itself.
+    report["energy_linear_1e-6"] = round(schedule_energy(*LINEAR, args.steps, 1e-6), 2)  # 27.63.
+    report["energy_linear_1e-9"] = round(schedule_energy(*LINEAR, args.steps, 1e-9), 2)  # 41.45.
+    report["energy_geodesic"] = round(schedule_energy(*GEODESIC, args.steps, 1e-6), 4)  # 9.8696.
+    report["pi_squared"] = round(math.pi ** 2, 4)  # The squared arc length of a quarter circle.
+    report["geodesic_kappa_0.25"] = round(geodesic_schedule(0.25), 4)  # 0.1464.
+    report["geodesic_kappa_0.75"] = round(geodesic_schedule(0.75), 4)  # 0.8536.
+    report["geodesic_rate_0.9"] = round(math.pi * math.tan(math.pi * 0.9 / 2), 4)  # 19.84.
+    report["uniform_source_kappa_0.5"] = round(geodesic_schedule(0.5, 0.25), 4)  # 0.6667.
+    print_report("kinetic-optimal paths", report)
+    save_run(args.out, config=vars(args), report=report)
+
+
+if __name__ == "__main__":
+    main()

lecture_5/multi_objective.py

@@ -0,0 +1,241 @@
+"""Several objectives at once: MOG-DFM's cone and pCoMole's rollout.
+
+A peptide is scored on binding, solubility and half-life together, and no dataset labels one
+sequence for all of them, so the single value function of Discrete Guidance is out of reach.
+MOG-DFM drops the ratio form and builds a tilt out of a rank inside each objective and a direction
+against a weight vector, with a cone that removes the candidates pointing the wrong way. pCoMole
+keeps the ratio and replaces the property likelihood by an exponentiated worst-objective utility,
+whose value at a candidate is the average terminal reward over rollouts of the base flow.
+
+Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
+u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
+
+Chapter 5, Sections 5.30 to 5.32. The single-objective tilt is in ``lecture_5/guidance.py``. The
+rollout runs on the synthetic DNA of ``cis6270.data``, scored by GC fraction against AT fraction.
+Prints the Pareto front, the chapter's MOG-DFM step with its cone, the chapter's pCoMole rollout,
+and a rollout on DNA whose value reverses the base rates.
+
+    python lecture_5/multi_objective.py                     # 600 rollouts per candidate on DNA
+    python lecture_5/multi_objective.py --steps 4000        # a tighter estimate of the second value
+    python lecture_5/multi_objective.py --steps 20 --quiet  # twenty rollouts, for a first look
+"""
+
+# %% 01. Imports and configuration
+import math
+
+import torch
+
+from cis6270.data import DNA, DNA_SIZE, MOTIF, decode_dna, dna_splits
+from cis6270.runner import common_parser, print_report, save_run, set_seed
+
+MASK = DNA_SIZE  # Index 4 is the mask symbol m, which a candidate still holds off its own position.
+OMEGA = torch.tensor([0.5, 0.5], dtype=torch.double)  # The trade-off the DNA rollout asks for.
+REFERENCE = torch.tensor([0.4, 0.3], dtype=torch.double)  # The levels r_k the objectives sit above.
+BETA = 2.0  # The inverse temperature of the terminal reward.
+
+
+# %% 02. Several objectives at once
+#
+# With several objectives there is no single best sequence. A candidate is on the Pareto front when
+# no other candidate matches it everywhere and beats it somewhere, and a weight vector picks one
+# point of that front by collapsing the scores to a single weighted number.
+
+
+def dominated(scores, row):
+    """True when some other candidate is at least as good everywhere and better somewhere."""
+    others = torch.cat([scores[:row], scores[row + 1:]])  # Every candidate except this one.
+    weak = (others >= scores[row]).all(dim=1)  # At least as good on every objective.
+    strict = (others > scores[row]).any(dim=1)  # Strictly better on at least one.
+    return bool((weak & strict).any())
+
+
+def pareto_front(scores):
+    """The indices of the candidates no other candidate dominates. scores: [M, N] -> list."""
+    return [row for row in range(len(scores)) if not dominated(scores, row)]
+
+
+# %% 03. MOG-DFM: rank, direction and a cone
+#
+# Without a joint likelihood there is no ratio to compute, so every candidate letter is scored twice.
+# Ranking inside each objective removes the units, the inner product with the weight vector measures
+# direction, and the two z-scored terms exponentiate into a tilt on the pretrained rate.
+
+
+def rank_scores(changes):
+    """I_n(y^i, x), the rank of each candidate's change within each objective, over the alphabet size."""
+    order = changes.argsort(dim=0).argsort(dim=0) + 1  # Rank one for the smallest change.
+    return order.to(changes.dtype) / changes.shape[0]  # [K, N] in {1/K, ..., 1}.
+
+
+def z_score(values):
+    """Norm[.], the z-score across the candidate letters at one position."""
+    spread = values.std(unbiased=False).clamp(min=1e-12)  # The population deviation, held off zero.
+    return (values - values.mean()) / spread  # Zero everywhere when no candidate stands out.
+
+
+def combined_score(changes, importance, weights, lam=1.0):
+    """Delta S = Norm[mean_n i_n I_n] + lambda Norm[Delta s . omega]. changes: [K, N] -> [K]."""
+    ranks = (rank_scores(changes) * importance).mean(dim=1)  # Importance-weighted average of ranks.
+    direction = changes @ weights  # [K] the directional score of each candidate.
+    return z_score(ranks) + lam * z_score(direction), ranks, direction
+
+
+# %% 04. Accepting a move only inside a cone
+#
+# A z-scored score puts about half the candidates above zero, so the tilt alone promotes the least
+# bad letter at a position where every letter moves away from the requested trade-off. The angle
+# between a candidate's score change and omega removes those, and a controller on the logarithm of
+# the angle holds the surviving fraction near one minus the target rejection rate.
+
+
+def cone_angles(changes, weights):
+    """alpha^i = arccos(Delta s . omega / (||Delta s|| ||omega||)), in degrees."""
+    cosine = (changes @ weights) / (changes.norm(dim=1) * weights.norm())  # [K].
+    return torch.rad2deg(torch.arccos(cosine.clamp(-1.0, 1.0))), cosine
+
+
+def update_cone(phi, rejected, target=0.2, gain=0.5, low=30.0, high=90.0):
+    """Phi_{t+h} = clip(Phi_t exp(eta (r - tau)), Phi_min, Phi_max), proportional control on log Phi."""
+    return min(max(phi * math.exp(gain * (rejected - target)), low), high)  # Wider when it rejects more.
+
+
+# %% 05. pCoMole: a reward in place of a likelihood
+#
+# Scoring a sequence by its weakest objective stops the sampler from trading one property away, and a
+# small multiple of the sum breaks ties. The exponentiated utility times a feasibility indicator is
+# the terminal reward, and the value of a candidate is the average reward over rollouts of the base
+# flow started from it.
+
+
+def utility(scores, weights, reference, rho=0.1):
+    """U(x) = min_k omega_k (f_k(x) - r_k) + rho sum_k omega_k (f_k(x) - r_k)."""
+    gaps = weights * (scores - reference)  # [N] the weighted gap on each objective.
+    return float(gaps.min() + rho * gaps.sum())  # The minimum, with the tie-breaking correction.
+
+
+def terminal_reward(scores, weights, reference, beta, feasible, rho=0.1):
+    """G(x) = exp(beta U(x)) 1[x in F], which is zero for an infeasible sequence."""
+    return math.exp(beta * utility(scores, weights, reference, rho)) * float(feasible)
+
+
+def rollout_value(rollouts, feasible, weights, reference, beta, rho=0.1):
+    """hat h_t(y_c), the average terminal reward over R continuations of the base flow."""
+    rewards = [terminal_reward(s, weights, reference, beta, f, rho) for s, f in zip(rollouts, feasible)]
+    return sum(rewards) / len(rewards), rewards  # A candidate with no feasible ending scores zero.
+
+
+# %% 06. A rollout on the synthetic DNA
+#
+# The base flow of a rollout is the mask-source sampler of this chapter with the letter frequency of
+# the data as its denoiser, so a continuation costs no network call. Two objectives that pull
+# against each other are the GC fraction and the AT fraction of the finished sequence, and the
+# feasible set is the sequences that carry the planted motif.
+
+
+def continuations(partial, frequency, rollouts, generator=None):
+    """Fill every masked position of a candidate from the base flow. partial: [L] -> [R, L]."""
+    grid = partial[None, :].repeat(rollouts, 1)  # [R, L] one copy of the candidate per rollout.
+    masked = grid == MASK  # [R, L] the positions the flow still owes a letter.
+    draws = torch.multinomial(frequency, int(masked.sum()), replacement=True, generator=generator)
+    return grid.masked_scatter(masked, draws)  # [R, L] the terminal sequences of the rollouts.
+
+
+def dna_objectives(sequences):
+    """The GC fraction and the AT fraction of each sequence, which sum to one. [R, L] -> [R, 2]."""
+    gc = ((sequences == 1) | (sequences == 2)).to(torch.get_default_dtype()).mean(dim=1)  # C and G.
+    return torch.stack([gc, 1.0 - gc], dim=1)  # [R, 2] raising one lowers the other.
+
+
+def carries_motif(sequences):
+    """One per sequence that holds the planted motif, which is the hard constraint here. [R] ints."""
+    return [int(MOTIF in sequence) for sequence in decode_dna(sequences)]
+
+
+def dna_rollout_value(partial, frequency, rollouts, weights, reference, beta, generator=None):
+    """hat h_t at one candidate, from ``rollouts`` continuations of the base flow. Returns a float."""
+    finished = continuations(partial, frequency, rollouts, generator)  # [R, L].
+    scores = dna_objectives(finished)  # [R, 2] the two objectives at the end of each rollout.
+    value, _ = rollout_value(list(scores), carries_motif(finished), weights, reference, beta)
+    return value  # Zero when no continuation of this candidate is feasible.
+
+
+# %% 07. The report
+def main():
+    parser = common_parser(__doc__.split("\n")[0])
+    args = parser.parse_args()
+    set_seed(args.seed)
+
+    # Section 5.30: four candidates on two objectives, and the two weight vectors.
+    candidates = torch.tensor([[0.8, 0.3], [0.6, 0.7], [0.5, 0.2], [0.6, 0.3]], dtype=torch.double)
+    front = pareto_front(candidates)
+    report = {
+        "pareto_front": front,  # Candidates 1 and 2; 3 and 4 are dominated.
+        "weighted_at_0.6_0.4": [round(float(s), 3) for s in candidates @ torch.tensor(
+            [0.6, 0.4], dtype=torch.double)],  # 0.60 for the first and 0.64 for the second.
+        "weighted_at_0.8_0.2": [round(float(s), 3) for s in candidates @ torch.tensor(
+            [0.8, 0.2], dtype=torch.double)],  # 0.70 and 0.62, which reverses the order.
+    }
+
+    # Section 5.31.2: one MOG-DFM step at a position holding A, with three candidate letters.
+    changes = torch.tensor([[0.0, 0.0], [0.20, -0.05], [0.05, 0.15], [-0.10, 0.08]], dtype=torch.double)
+    importance = torch.tensor([1.2, 0.8], dtype=torch.double)
+    weights = torch.tensor([0.6, 0.4], dtype=torch.double)
+    delta_s, ranks, direction = combined_score(changes, importance, weights)
+    report["rank_average"] = [round(float(r), 3) for r in ranks]  # (0.500, 0.700, 0.850, 0.450).
+    report["directional_scores"] = [round(float(d), 3) for d in direction]  # (0, 0.100, 0.090, -0.028).
+    report["combined_score"] = [round(float(s), 3) for s in delta_s]  # (-1.511, 1.541, 2.297, -2.327).
+    posterior = torch.tensor([0.3, 0.3, 0.3, 0.1], dtype=torch.double)  # The denoiser at t = 0.7.
+    base = posterior[1:] / (1.0 - 0.7)  # The pretrained rates out of A toward C, G and T.
+    report["base_rates"] = [round(float(r), 4) for r in base]  # (1.000, 1.000, 0.333).
+    tilted = base * delta_s[1:].exp()
+    report["guided_rates"] = [round(float(r), 3) for r in tilted]  # (4.667, 9.948, 0.033).
+    angles, cosines = cone_angles(changes[1:], weights)
+    report["cone_cosines"] = [round(float(c), 4) for c in cosines]  # (0.6727, 0.7894, -0.3032).
+    report["cone_angles_deg"] = [round(float(a), 1) for a in angles]  # (47.7, 37.9, 107.7).
+    inside = angles <= 60.0  # The cone test at Phi = 60 degrees.
+    survivors = tilted[inside]
+    report["after_cone"] = [round(float(r), 3) for r in survivors / survivors.sum()]  # (0.319, 0.681).
+    rejected = float((~inside).to(torch.get_default_dtype()).mean())  # One candidate of three.
+    report["next_cone_angle"] = round(update_cone(60.0, rejected), 1)  # 64.1.
+
+    # Section 5.32.2: three candidate edits, each rolled out four times.
+    omega, reference = torch.tensor([0.5, 0.5], dtype=torch.double), torch.zeros(2, dtype=torch.double)
+    edits = [
+        ([(0.8, 0.4), (0.6, 0.6), (0.9, 0.1), (0.5, 0.5)], [1, 1, 0, 1], 0.6),
+        ([(1.0, 0.0), (0.9, 0.1), (0.95, 0.05), (1.0, 0.1)], [1, 1, 1, 1], 1.2),
+        ([(0.4, 0.4)] * 4, [0, 0, 0, 0], 2.0),
+    ]
+    values = [rollout_value([torch.tensor(s, dtype=torch.double) for s in rollouts],
+                            feasible, omega, reference, 2.0)[0] for rollouts, feasible, _ in edits]
+    base_rates = torch.tensor([rate for _, _, rate in edits], dtype=torch.double)
+    products = torch.tensor(values) * base_rates  # The base rate times the continuation value.
+    report["rollout_values"] = [round(v, 3) for v in values]  # (1.390, 1.181, 0).
+    report["rollout_products"] = [round(float(p), 3) for p in products]  # (0.834, 1.417, 0).
+    report["rollout_selection"] = [round(float(p), 3) for p in products / products.sum()]
+    report["rollout_base_selection"] = [round(float(p), 3) for p in base_rates / base_rates.sum()]
+
+    # The same rollout on the shared DNA, from a partly revealed sequence with ACG already in place.
+    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
+    counts = torch.bincount(splits["train_tokens"].flatten(), minlength=DNA_SIZE).double()
+    frequency = counts / counts.sum()  # [4] the denoiser of the base flow, GC-rich like the data.
+    # The base rates below are kappa_dot/(1 - kappa) times that denoiser, read at t = 0.3.
+    revealed = torch.full((args.length,), MASK, dtype=torch.long)
+    revealed[2:5] = torch.tensor([0, 1, 2])  # Positions 3, 4 and 5 hold A, C and G.
+    finish = {letter: revealed.clone() for letter in "TC"}
+    for letter, candidate in finish.items():
+        candidate[5] = DNA.index(letter)  # Position 6 takes T, completing ACGT, or C, which does not.
+    seen = {letter: dna_rollout_value(candidate, frequency, args.steps, OMEGA, REFERENCE, BETA)
+            for letter, candidate in finish.items()}
+    rates = torch.tensor([float(frequency[DNA.index(letter)]) for letter in finish]) / (1.0 - 0.3)
+    guided = rates * torch.tensor([seen[letter] for letter in finish])  # The base rate times hat h.
+    report["dna_rollout_value_T"] = round(seen["T"], 4)  # Every continuation carries the motif.
+    report["dna_rollout_value_C"] = round(seen["C"], 4)  # Only the ones that plant it elsewhere.
+    report["dna_base_rates_T_C"] = [round(float(r), 4) for r in rates]  # The GC-rich data prefers C.
+    report["dna_base_selection"] = [round(float(r), 4) for r in rates / rates.sum()]
+    report["dna_guided_selection"] = [round(float(r), 4) for r in guided / guided.sum()]
+    print_report("several objectives at once", report)
+    save_run(args.out, config=vars(args), report=report)
+
+
+if __name__ == "__main__":
+    main()

lecture_5/rectify.py

@@ -0,0 +1,293 @@
+"""Rectification and AReUReDi: changing the coupling so that fewer steps suffice.
+
+A one-step sampler draws every position from its own marginal, so whatever structure the data has
+between positions is lost in that draw. The error it makes is the conditional total correlation of
+the coupling, which is a property of how the pairs were drawn and not of the data. Training on the
+current pairs, sampling from the model, and keeping the pairs it produces lowers that total
+correlation, and the discrete counterpart of the transport-cost bound says it never raises it.
+AReUReDi then composes a rectified flow with a Metropolis refinement whose proposal starts from the
+model's own per-position distribution.
+
+Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
+u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
+
+Chapter 5, Sections 5.36 to 5.41. The second recoupling runs on adjacent base pairs of the synthetic
+DNA of ``cis6270.data``. Prints the total correlation of each round, the Hamming distance and the
+one-step success rate, and the AReUReDi proposal with its acceptance ratio.
+
+    python lecture_5/rectify.py                     # four rounds, 600 quadrature points each
+    python lecture_5/rectify.py --rounds 6          # closer to the limiting coupling
+    python lecture_5/rectify.py --steps 20 --quiet  # a coarse integration, visibly less accurate
+"""
+
+# %% 01. Imports and configuration
+import math
+
+import torch
+
+from cis6270.data import DNA, DNA_SIZE, dna_splits
+from cis6270.runner import common_parser, print_report, save_run, set_seed
+
+
+torch.set_default_dtype(torch.double)  # The recoupling is an exact integration, so keep precision.
+LENGTH = 2  # Two positions, so the state space has 16 sequences and can be enumerated.
+STATES = DNA_SIZE ** LENGTH
+LETTER = torch.stack([  # [L, 16] the letter each sequence index holds at each position.
+    torch.arange(STATES) // DNA_SIZE, torch.arange(STATES) % DNA_SIZE,
+])
+
+
+# %% 02. The pairing of source and data
+#
+# Since the first Euler step the joint distribution of the next state has been replaced by the
+# product of the per-position marginals. With small steps that costs almost nothing. One large step
+# draws every position at once from a product, and the divergence between the joint and that product
+# is the conditional total correlation of the coupling.
+
+
+def position_marginals(row):
+    """The per-position marginals of one distribution over sequences. row: [16] -> [L, K]."""
+    out = torch.zeros(LENGTH, DNA_SIZE)
+    for i in range(LENGTH):
+        out[i].index_add_(0, LETTER[i], row)  # Collect the mass of every sequence by its letter.
+    return out
+
+
+def total_correlation(coupling, source):
+    """TC_pi(X_1 | X_0), the expected divergence from the joint to the product of its marginals."""
+    total = 0.0
+    for s in range(STATES):
+        marginal = position_marginals(coupling[s])  # [L, K] given this source sequence.
+        product = marginal[0][LETTER[0]] * marginal[1][LETTER[1]]  # [16] the factorized draw.
+        live = coupling[s] > 0
+        total += float(source[s] * (coupling[s][live] * (coupling[s][live] / product[live]).log()).sum())
+    return total  # Zero exactly when the positions are independent given the source.
+
+
+def one_step_success(coupling, source, support):
+    """The chance a single factorized draw lands inside the data, averaged over source sequences."""
+    total = 0.0
+    for s in range(STATES):
+        marginal = position_marginals(coupling[s])
+        total += float(source[s] * sum(marginal[0][a] * marginal[1][b] for a, b in support))
+    return total  # 1/4 under the independent coupling on the complementary pairs.
+
+
+def hamming(coupling, source):
+    """The expected number of positions at which a source and its paired data sequence differ."""
+    distance = (LETTER[:, :, None] != LETTER[:, None, :]).sum(dim=0).to(coupling)  # [16, 16].
+    return float((source[:, None] * coupling * distance).sum())  # 1.5 under independent pairing.
+
+
+# %% 03. Total correlation on two bases
+#
+# Let the data be the four complementary pairs, each with probability one quarter, and the source be
+# uniform over all sixteen ordered pairs. Each position of the data alone is uniform over the four
+# bases, so the product of the two marginals spreads one sixteenth everywhere while the joint sits on
+# four states with one quarter apiece.
+
+
+def complementary_pairs():
+    """The data distribution over AT, TA, CG and GC, and the indices of those four sequences."""
+    support = [(DNA.index(a), DNA.index(b)) for a, b in ("AT", "TA", "CG", "GC")]
+    q = torch.zeros(STATES)
+    for a, b in support:
+        q[a * DNA_SIZE + b] = 0.25  # Each complementary pair with probability one quarter.
+    return q, support
+
+
+def adjacent_pairs(tokens):
+    """The empirical distribution of adjacent base pairs in the training DNA. Returns [16]."""
+    left, right = tokens[:, :-1].reshape(-1), tokens[:, 1:].reshape(-1)  # Every adjacent pair.
+    counts = torch.bincount(left * DNA_SIZE + right, minlength=STATES).to(torch.get_default_dtype())
+    return counts / counts.sum()  # [16] over the sixteen ordered pairs.
+
+
+# %% 04. ReDi: one recoupling round
+#
+# The model here is the entire many-step sampler, and the next coupling pairs a source sample with
+# the sequence the sampler produces from it. The velocity is the true marginal velocity of the
+# current coupling, and the master equation is integrated without discretization error, so the only
+# approximation left is the factorization across the two positions.
+
+
+def generator(joint, t):
+    """The factorized velocity at time t, in the variable u with t = 1 - exp(-u). Returns [16, 16]."""
+    eye = torch.eye(DNA_SIZE)
+    mix = (1 - t) * eye[:, :, None] + t * eye[:, None, :]  # [K, K, K] the mixture path at one position.
+    per = [mix[:, LETTER[i][:, None], LETTER[i][None, :]] for i in range(LENGTH)]  # [K, 16, 16] each.
+    weight = joint[None, None, :, :] * per[0][:, None, :, :] * per[1][None, :, :, :]
+    weight = weight.reshape(STATES, STATES, STATES)  # [x, x_0, x_1] the unnormalized posterior.
+    norm = weight.sum(dim=(1, 2)).clamp(min=1e-300)
+    posterior = weight.sum(dim=1) / norm[:, None]  # [x, x_1] the posterior over the data sequence.
+    rates, columns = torch.zeros(STATES, STATES), torch.arange(STATES)
+    for i in range(LENGTH):
+        denoiser = torch.zeros(STATES, DNA_SIZE)  # [x, a] the denoiser at position i.
+        denoiser.index_add_(1, LETTER[i], posterior)
+        for a in range(DNA_SIZE):
+            moved = LETTER.clone()
+            moved[i] = a  # [L, 16] every sequence with position i replaced by the letter a.
+            destination = moved[0] * DNA_SIZE + moved[1]  # [16] the index of that sequence.
+            keep = LETTER[i] != a  # A sequence already holding a at position i sends nothing.
+            rates[destination[keep], columns[keep]] += denoiser[keep, a]
+    return rates - torch.diag(rates.sum(dim=0))  # In u-time the schedule factor is one.
+
+
+def recouple(coupling, source, steps, horizon=14.0):
+    """Run the exact sampler from every source sequence and return the coupling it produces.
+
+    The substitution t = 1 - exp(-u) holds the schedule factor at one, so the integration stays
+    accurate as t approaches the horizon, where 1 - t is about 1e-6. Each cell takes a Heun step,
+    which is second order and reaches the chapter's figures at a few hundred cells.
+    """
+    joint = source[:, None] * coupling  # pi(x_0, x_1) over pairs.
+    probability = torch.eye(STATES)  # Row s is the state distribution started at the source s.
+    du = horizon / steps
+    for k in range(steps):
+        left = generator(joint, 1.0 - math.exp(-k * du))  # The velocity at the start of the cell.
+        right = generator(joint, 1.0 - math.exp(-(k + 1) * du))  # And at its end.
+        first = probability @ left.T  # The Euler slope.
+        second = (probability + du * first) @ right.T  # The slope the Euler step lands on.
+        probability = probability + 0.5 * du * (first + second)  # Their average.
+        probability = probability.clamp(min=0.0)
+        probability = probability / probability.sum(dim=1, keepdim=True)
+    return probability  # [16, 16] the next coupling, one row per source sequence.
+
+
+# %% 05. AReUReDi: the Tchebycheff score and the tilt
+#
+# Scoring a sequence by its weakest weighted objective stops the sampler from trading one property
+# away. The target is the base model times the exponentiated score, and the refinement anneals the
+# strength linearly over its iterations, so early ones explore and late ones commit.
+
+
+def tchebycheff(scores, weights):
+    """S_omega(x) = min_n omega_n s_n(x), the smallest weighted objective. scores: [M, N] -> [M]."""
+    return (scores * weights).min(dim=-1).values  # Raising a weight lowers the pressure on it.
+
+
+def anneal(iteration, total, low=1.0, high=20.0):
+    """eta_t = eta_min + (eta_max - eta_min) t/(T - 1), which is 1 + t at the chapter's settings."""
+    return low + (high - low) * iteration / (total - 1)  # 12 at the twelfth iteration of twenty.
+
+
+# %% 06. The proposal and its acceptance step
+#
+# The proposal starts from the model's own per-position distribution and reweights it by a balancing
+# function of how much the tilt would change. The identity g(u) = u g(1/u) collapses the acceptance
+# ratio to a ratio of two one-position normalizers, which costs one sum over the alphabet.
+
+
+def barker(u):
+    """g(u) = u/(1 + u), which satisfies g(u) = u g(1/u)."""
+    return u / (1.0 + u)  # The square root is the other common choice.
+
+
+def proposal(model_distribution, tilts, current):
+    """q_i(. | x) and its normalizer Z(x). tilts: [K] the value of W at each letter."""
+    ratios = tilts / tilts[current]  # r_i(y; x), the tilt at the proposed letter over the current one.
+    weighted = model_distribution * barker(ratios)  # The model's prior times the balancing function.
+    return weighted / weighted.sum(), float(weighted.sum()), ratios  # q, Z(x) and the ratios.
+
+
+def acceptance(z_current, z_proposed):
+    """alpha_i(x, x') = min(1, Z(x)/Z(x')), what the balancing identity leaves of the ratio."""
+    return min(1.0, z_current / z_proposed)  # One in both directions when the normalizers agree.
+
+
+# %% 07. Concentration on the maximizers
+#
+# Keeping only the maximizers in the denominator and bounding the numerator by the score gap gives an
+# exponential bound on the probability of anything outside the maximizing set, so the samples
+# concentrate there as the strength grows.
+
+
+def tilted_distribution(base, scores, eta):
+    """pi_{eta, omega}(x) proportional to p_1(x) exp(eta S_omega(x))."""
+    weighted = base * (eta * scores).exp()
+    return weighted / weighted.sum()  # [M] over the sequences under consideration.
+
+
+def concentration_bound(base, scores, eta):
+    """exp(-eta Delta) p_1(x) over the base mass of the maximizers, for a sequence outside the set."""
+    best = scores.max()
+    gap = float(best - scores[scores < best].max())  # Delta_omega, the gap to the next best score.
+    return math.exp(-eta * gap) / float(base[scores == best].sum()), gap
+
+
+# %% 08. The report
+def main():
+    parser = common_parser(__doc__.split("\n")[0])
+    parser.add_argument("--rounds", type=int, default=4, help="rectification rounds")
+    args = parser.parse_args()
+    set_seed(args.seed)
+
+    source = torch.full((STATES,), 1.0 / STATES)  # Uniform over all sixteen ordered pairs.
+    q, support = complementary_pairs()
+    coupling = q[None, :].repeat(STATES, 1)  # The independent coupling pairs every source with q.
+    report = {
+        "tc_round_0": round(total_correlation(coupling, source), 3),  # log 4 = 1.386.
+        "log_4": round(math.log(4), 3),  # What the enumeration of the four pairs gives.
+        "hamming_round_0": round(hamming(coupling, source), 3),  # 1.5.
+        "success_round_0": round(one_step_success(coupling, source, support), 4),  # 0.25.
+    }
+    for round_index in range(1, args.rounds + 1):
+        coupling = recouple(coupling, source, args.steps)  # One application of the operator.
+        report[f"tc_round_{round_index}"] = round(total_correlation(coupling, source), 3)
+        report[f"hamming_round_{round_index}"] = round(hamming(coupling, source), 3)
+        report[f"success_round_{round_index}"] = round(one_step_success(coupling, source, support), 4)
+        if round_index == 1:  # Where the first round's gain comes from, read at two sources.
+            columns = [a * DNA_SIZE + b for a, b in support]  # The four complementary pairs.
+            for name in ("AA", "AT"):  # One source outside the data and one inside it.
+                row = DNA.index(name[0]) * DNA_SIZE + DNA.index(name[1])
+                report[f"pairing_from_{name}"] = [round(float(coupling[row, c]), 3) for c in columns]
+    report["tc_limit"] = round(0.75 * math.log(2), 3)  # Twelve sources split in half, 0.520 nats.
+    report["success_limit"] = round(10.0 / 16.0, 4)  # 0.625, four exact sources and twelve halves.
+    report["hamming_limit"] = round(12.0 / 16.0, 4)  # 0.75, the least any coupling can reach.
+
+    # The same measurement on adjacent base pairs of the synthetic DNA, under independent pairing.
+    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
+    dna = adjacent_pairs(splits["train_tokens"])  # [16] the empirical distribution of a base pair.
+    dna_coupling = dna[None, :].repeat(STATES, 1)  # Every source paired with the same distribution.
+    live = dna[dna > 0]
+    report["dna_pair_entropy"] = round(float(-(live * live.log()).sum()), 4)  # At most log 16 = 2.77.
+    report["dna_total_correlation"] = round(total_correlation(dna_coupling, source), 4)
+    report["dna_fraction_of_entropy"] = round(
+        total_correlation(dna_coupling, source) / float(-(live * live.log()).sum()), 4)
+
+    # Section 5.41.2: one refinement iteration at the twelfth of twenty, with two objectives.
+    scores = torch.tensor([[0.60, 0.80], [0.70, 0.76], [0.50, 0.90], [0.64, 0.62]])
+    weights = torch.tensor([0.5, 0.5])
+    eta = anneal(11, 20)  # The twelfth iteration, counting from zero, where 1 + t equals 12.
+    tchebycheff_scores = tchebycheff(scores, weights)
+    tilts = (eta * tchebycheff_scores).exp()  # W = exp(eta S) at that strength.
+    report["annealed_eta"] = round(eta, 4)  # 1 + t over twenty iterations.
+    report["tchebycheff"] = [round(float(s), 3) for s in tchebycheff_scores]  # (0.300, ...).
+    report["tilts"] = [round(float(w), 2) for w in tilts]  # (36.60, 66.69, 20.09, 41.26).
+    model_distribution = torch.tensor([0.40, 0.25, 0.20, 0.15])  # The rectified model at this position.
+    q_current, z_current, ratios = proposal(model_distribution, tilts, 0)  # The position holds A.
+    report["ratios_from_A"] = [round(float(r), 3) for r in ratios]  # (1.000, 1.822, 0.549, 1.127).
+    report["barker_from_A"] = [round(float(barker(r)), 3) for r in ratios]  # (0.500, 0.646, ...).
+    report["normalizer_at_A"] = round(z_current, 4)  # 0.5118.
+    report["proposal_at_A"] = [round(float(p), 3) for p in q_current]  # (0.391, 0.315, 0.138, 0.155).
+    _, z_proposed, _ = proposal(model_distribution, tilts, 1)  # The proposed sequence holds C.
+    report["normalizer_at_C"] = round(z_proposed, 4)  # 0.3704.
+    report["accept_A_to_C"] = round(acceptance(z_current, z_proposed), 3)  # One.
+    report["accept_C_to_A"] = round(acceptance(z_proposed, z_current), 3)  # 0.724.
+
+    # Section 5.41.3: the tilted distribution over four sequences, and the bound on a non-maximizer.
+    base = torch.tensor([0.50, 0.25, 0.15, 0.10])
+    objective = torch.tensor([0.20, 0.30, 0.45, 0.45])
+    for strength in (20.0, 60.0):
+        tilted = tilted_distribution(base, objective, strength)
+        bound, gap = concentration_bound(base, objective, strength)
+        report[f"mass_on_maximizers_eta_{strength:g}"] = f"{float(tilted[2:].sum()):.6f}"
+        report[f"second_sequence_eta_{strength:g}"] = f"{float(tilted[1]):.6f}"
+        report[f"bound_eta_{strength:g}"] = f"{bound * float(base[1]):.6f}"
+    report["score_gap"] = round(gap, 4)  # 0.15.
+    print_report("rectification", report)
+    save_run(args.out, config=vars(args), report=report)
+
+
+if __name__ == "__main__":
+    main()