diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..0baf5b6b288d5a661aec811f1d9e45774e06dde9 --- /dev/null +++ b/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. diff --git a/MIGRATION.md b/MIGRATION.md new file mode 100644 index 0000000000000000000000000000000000000000..6898af1aebad61dab8a83ddfe1f121a2693ee046 --- /dev/null +++ b/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. diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..8a6cd1dc6fac56815d8e6a3bb11ad9e509804807 --- /dev/null +++ b/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. diff --git a/cis6270/__init__.py b/cis6270/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ce5bc20588d6a7b004de914a2749ad074be434de --- /dev/null +++ b/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" diff --git a/cis6270/data.py b/cis6270/data.py new file mode 100644 index 0000000000000000000000000000000000000000..a80bd048a9e843d199fab9b843b869af85156670 --- /dev/null +++ b/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 diff --git a/cis6270/metrics.py b/cis6270/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..a1b3e8056c1ac9836d16e9e3c78813f9cbba038e --- /dev/null +++ b/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. diff --git a/cis6270/nets.py b/cis6270/nets.py new file mode 100644 index 0000000000000000000000000000000000000000..9c811c79db621b64541c06e545f23c7a3fe2cc7b --- /dev/null +++ b/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) diff --git a/cis6270/runner.py b/cis6270/runner.py new file mode 100644 index 0000000000000000000000000000000000000000..b03c1650ce59b619ea8c8be1a88e781dbe478538 --- /dev/null +++ b/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}") diff --git a/lecture_1/README.md b/lecture_1/README.md new file mode 100644 index 0000000000000000000000000000000000000000..3c583b584c01fb9188204249ad33d05a5d2086ef --- /dev/null +++ b/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. diff --git a/lecture_1/calculus.py b/lecture_1/calculus.py new file mode 100644 index 0000000000000000000000000000000000000000..d6c8035d88b7fbe8fa041f6e545fd8037353737d --- /dev/null +++ b/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. diff --git a/lecture_1/fields.py b/lecture_1/fields.py new file mode 100644 index 0000000000000000000000000000000000000000..3ca8e806495afd7e9f4d37c9f6f24a63ebd2b7cc --- /dev/null +++ b/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. diff --git a/lecture_1/information.py b/lecture_1/information.py new file mode 100644 index 0000000000000000000000000000000000000000..cd1215fd5cd5bf2436eb1ca8446fe34c9aba0840 --- /dev/null +++ b/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. diff --git a/lecture_1/linear_algebra.py b/lecture_1/linear_algebra.py new file mode 100644 index 0000000000000000000000000000000000000000..b0e770f0a5cb29381c870aa3e35f52d2126f5b8e --- /dev/null +++ b/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. diff --git a/lecture_1/optimization.py b/lecture_1/optimization.py new file mode 100644 index 0000000000000000000000000000000000000000..d3b531d5bfdfdb34022aa1312df78aa88774537f --- /dev/null +++ b/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. diff --git a/lecture_1/probability.py b/lecture_1/probability.py new file mode 100644 index 0000000000000000000000000000000000000000..4a0ecc968588f65e3d5d794e01a59f2465f14eb7 --- /dev/null +++ b/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. diff --git a/lecture_2/README.md b/lecture_2/README.md new file mode 100644 index 0000000000000000000000000000000000000000..9d96c7aac12a8f52a3cde336e68ce30a4a2b012c --- /dev/null +++ b/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. diff --git a/lecture_2/flow_matching_2d.py b/lecture_2/flow_matching_2d.py new file mode 100644 index 0000000000000000000000000000000000000000..738a326c34d8787d00d3151e0028e3bbb463759d --- /dev/null +++ b/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() diff --git a/lecture_2/flow_matching_mnist.py b/lecture_2/flow_matching_mnist.py new file mode 100644 index 0000000000000000000000000000000000000000..8a6c8754f807a4f5759c9d869e4848038af52833 --- /dev/null +++ b/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. diff --git a/lecture_2/interpolants.py b/lecture_2/interpolants.py new file mode 100644 index 0000000000000000000000000000000000000000..be2bca47c03b713036e0ab2aac4586dbaeac1c6e --- /dev/null +++ b/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() diff --git a/lecture_2/velocity_fields.py b/lecture_2/velocity_fields.py new file mode 100644 index 0000000000000000000000000000000000000000..67f2c3491f4b601777a69c0b28f9eab51a1e0f9a --- /dev/null +++ b/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() diff --git a/lecture_3/README.md b/lecture_3/README.md new file mode 100644 index 0000000000000000000000000000000000000000..cc3673b8e86eff4ec060fe9926a22616dae99f83 --- /dev/null +++ b/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. diff --git a/lecture_3/ddpm.py b/lecture_3/ddpm.py new file mode 100644 index 0000000000000000000000000000000000000000..f3482fe9fb1eeaf9dcc8fd2d888252a5d6e520ff --- /dev/null +++ b/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() diff --git a/lecture_3/forward_diffusion.py b/lecture_3/forward_diffusion.py new file mode 100644 index 0000000000000000000000000000000000000000..ec68f21a0a5443f3c02fb86f5fdf0d9282efd003 --- /dev/null +++ b/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() diff --git a/lecture_3/guidance.py b/lecture_3/guidance.py new file mode 100644 index 0000000000000000000000000000000000000000..55ce908264602c389af030b5dfd7c04e8837fdfd --- /dev/null +++ b/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() diff --git a/lecture_3/latent_diffusion.py b/lecture_3/latent_diffusion.py new file mode 100644 index 0000000000000000000000000000000000000000..9773653a1242e7c35a6b30efa2ef677d06038ee5 --- /dev/null +++ b/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() diff --git a/lecture_3/samplers.py b/lecture_3/samplers.py new file mode 100644 index 0000000000000000000000000000000000000000..dd3722c39edfdf7044da392f8163f99ae20af55e --- /dev/null +++ b/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() diff --git a/lecture_3/score_matching.py b/lecture_3/score_matching.py new file mode 100644 index 0000000000000000000000000000000000000000..3d02b5dfcaf2860a9904f0461492698a3cdd6a78 --- /dev/null +++ b/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() diff --git a/lecture_3/sde.py b/lecture_3/sde.py new file mode 100644 index 0000000000000000000000000000000000000000..5c0644f825d0d62a27697939e8228498a9bf52d8 --- /dev/null +++ b/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() diff --git a/lecture_4/README.md b/lecture_4/README.md new file mode 100644 index 0000000000000000000000000000000000000000..88e0f33d528db845de562e06396c1a64b0411706 --- /dev/null +++ b/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. diff --git a/lecture_4/block.py b/lecture_4/block.py new file mode 100644 index 0000000000000000000000000000000000000000..8634dabaded189dd58c5264c044facab439a4bcc --- /dev/null +++ b/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. diff --git a/lecture_4/ctmc.py b/lecture_4/ctmc.py new file mode 100644 index 0000000000000000000000000000000000000000..09c67ef4d1784efd7b658d7faca7728e08704060 --- /dev/null +++ b/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. diff --git a/lecture_4/d3pm.py b/lecture_4/d3pm.py new file mode 100644 index 0000000000000000000000000000000000000000..a45fc870694a46858c8d015ab18a48a61cc58135 --- /dev/null +++ b/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. diff --git a/lecture_4/guidance.py b/lecture_4/guidance.py new file mode 100644 index 0000000000000000000000000000000000000000..01e4f9f2250cdf13a832a29cf4d6b9616af09b1f --- /dev/null +++ b/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. diff --git a/lecture_4/mdlm.py b/lecture_4/mdlm.py new file mode 100644 index 0000000000000000000000000000000000000000..ff6e20da580b0d8d9fb8a22a52217d54539af0e1 --- /dev/null +++ b/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. diff --git a/lecture_4/planner.py b/lecture_4/planner.py new file mode 100644 index 0000000000000000000000000000000000000000..0b1c727303eeeebb187711b866f90aeacf9e396c --- /dev/null +++ b/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. diff --git a/lecture_4/search.py b/lecture_4/search.py new file mode 100644 index 0000000000000000000000000000000000000000..923ec962e19fbedb307c467c4d5a0da48a9eefd3 --- /dev/null +++ b/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. diff --git a/lecture_4/sedd.py b/lecture_4/sedd.py new file mode 100644 index 0000000000000000000000000000000000000000..58b0a32ad15f4bda75c891affb4af816c586bcdc --- /dev/null +++ b/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. diff --git a/lecture_4/simplex.py b/lecture_4/simplex.py new file mode 100644 index 0000000000000000000000000000000000000000..dece25155f466ea0e588eeb9c5ede698e1ae45ef --- /dev/null +++ b/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. diff --git a/lecture_4/uniform.py b/lecture_4/uniform.py new file mode 100644 index 0000000000000000000000000000000000000000..896cc86d88bbb34c2cd1bf3ed807fe1c23f44649 --- /dev/null +++ b/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. diff --git a/lecture_5/README.md b/lecture_5/README.md new file mode 100644 index 0000000000000000000000000000000000000000..55d17d20a97c7eb0233e79fd75bfdd4f32874b3c --- /dev/null +++ b/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. diff --git a/lecture_5/corrector.py b/lecture_5/corrector.py new file mode 100644 index 0000000000000000000000000000000000000000..3b8e086f462ab14f62360ba2d8a37bb79ec3db96 --- /dev/null +++ b/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() diff --git a/lecture_5/dfm.py b/lecture_5/dfm.py new file mode 100644 index 0000000000000000000000000000000000000000..ea0f5db17501d4fe0740fbd302c94b124443e9ad --- /dev/null +++ b/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() diff --git a/lecture_5/dirichlet.py b/lecture_5/dirichlet.py new file mode 100644 index 0000000000000000000000000000000000000000..35ab43e9a8300a630bfb91b065e585af4aa44d4d --- /dev/null +++ b/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() diff --git a/lecture_5/edit_flows.py b/lecture_5/edit_flows.py new file mode 100644 index 0000000000000000000000000000000000000000..b189e5391aed11e256bb1149ecefce9b0d95535e --- /dev/null +++ b/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() diff --git a/lecture_5/fisher.py b/lecture_5/fisher.py new file mode 100644 index 0000000000000000000000000000000000000000..14e5e9a826b330bd473704d34bc4d456af5da8e3 --- /dev/null +++ b/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(), 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): + """_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, 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() diff --git a/lecture_5/guidance.py b/lecture_5/guidance.py new file mode 100644 index 0000000000000000000000000000000000000000..624e53bd90bb64aadc30a3b0b35bfbc9222a7f48 --- /dev/null +++ b/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() diff --git a/lecture_5/gumbel.py b/lecture_5/gumbel.py new file mode 100644 index 0000000000000000000000000000000000000000..7dc0134f1fb888cdf66521816d5c034bcc20015a --- /dev/null +++ b/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 . + + 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() diff --git a/lecture_5/kinetic_optimal.py b/lecture_5/kinetic_optimal.py new file mode 100644 index 0000000000000000000000000000000000000000..f57977701151fe6304db128fe0814550ef5b072f --- /dev/null +++ b/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() diff --git a/lecture_5/multi_objective.py b/lecture_5/multi_objective.py new file mode 100644 index 0000000000000000000000000000000000000000..22676a529a1229bc6e79385daef0aa363f061aa0 --- /dev/null +++ b/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() diff --git a/lecture_5/rectify.py b/lecture_5/rectify.py new file mode 100644 index 0000000000000000000000000000000000000000..3e5df98ac8ba64aa0fdc2eb1848bee3b07bf4d05 --- /dev/null +++ b/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() diff --git a/lecture_5/velocity.py b/lecture_5/velocity.py new file mode 100644 index 0000000000000000000000000000000000000000..2027415bf18f4bd966db171af3c15d09ace85d8b --- /dev/null +++ b/lecture_5/velocity.py @@ -0,0 +1,212 @@ +"""A jump rate written as a velocity, and the master equation it satisfies. + +Nothing here is learned. A velocity on a finite set is a rate matrix read from the destination's +side, two rules fix which tables of numbers qualify, and the master equation says when such a table +generates a probability path. The schedule factor of the mixture path turns out to be a hazard rate +whose survival curve is the path itself, and the discrete velocity has the same shape as the +continuous one of Chapter 2, with a posterior distribution in place of a posterior mean. + +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.2, 5.3, 5.6.3 and 5.8.1. The starting distribution of the long integration is +the letter frequency of the synthetic DNA of ``cis6270.data``. Prints the chapter's worked velocity, +its Euler step, the hazard at two times, and the two velocities side by side. + + python lecture_5/velocity.py # 600 Euler steps in the long integration + python lecture_5/velocity.py --steps 20 --quiet # a coarser grid, visibly less accurate +""" + +# %% 01. Imports and configuration +import math + +import torch +import torch.nn.functional as F + +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) # Every number here is arithmetic, so keep full precision. + + + +# %% 02. A jump rate written as a velocity +# +# Chapter 4 wrote R_t(x, y) for the rate of jumping from x to y. Discrete flow matching writes the +# same number with the destination first, u_t(y, x) = R_t(x, y). Fixing x and letting y vary gives +# one vector attached to the state x, with one component per destination. + + +def velocity_from_rates(rates): + """Transpose a Chapter 4 rate matrix R_t(x, y) into a Chapter 5 velocity u_t(y, x).""" + return rates.transpose(0, 1) # [K, K] with the destination now indexing the row. + + +def jump_to_one_state(target, size=DNA_SIZE, rate=1.0): + """Every state other than ``target`` leaves toward it at ``rate``, and ``target`` never leaves.""" + u = torch.zeros(size, size) # [K, K], row = destination, column = current state. + for x in range(size): + if x == target: + continue + u[target, x] = rate # The one positive off-diagonal entry in column x. + u[x, x] = -rate # The diagonal is minus the total rate of leaving x. + return u + + +# %% 03. The two rules a velocity obeys +# +# Off the diagonal the entries are nonnegative, since probability cannot flow out of a state at a +# negative rate, and every column sums to zero, since probability leaving one state has to arrive at +# another. The two rules fix the diagonal once the off-diagonal entries are chosen. + + +def obeys_the_two_rules(u, tolerance=1e-9): + """True when every off-diagonal entry is nonnegative and every column sums to zero.""" + off = u - torch.diag(torch.diagonal(u)) # The matrix with the diagonal removed. + return bool((off >= -tolerance).all() and (u.sum(dim=0).abs() < tolerance).all()) + + +def set_diagonal(u): + """u_t(x, x) = -sum_{y != x} u_t(y, x), the total rate of leaving x.""" + u = u.clone() + u.fill_diagonal_(0.0) # Keep the off-diagonal entries and discard whatever the diagonal held. + return u - torch.diag(u.sum(dim=0)) # Each column now sums to zero. + + +# %% 04. The master equation in velocity form +# +# The rate of change of p_t(y) is the velocity into y times the probability of each current state, +# summed over states. Written as a matrix equation it is p_dot = u p, with p a column vector and no +# transpose, which is the column convention of this chapter. + + +def master_equation(u, p): + """p_dot_t(y) = sum_x u_t(y, x) p_t(x). u: [K, K], p: [K] -> [K].""" + return u @ p # The matrix multiplies the probability vector from the left. + + +def arrivals_and_departures(u, p, y): + """The flux arriving at y from every other state, and the flux leaving y, as two numbers.""" + others = [x for x in range(len(p)) if x != y] # Every state except y itself. + arriving = sum(float(u[y, x] * p[x]) for x in others) # u_t(y, x) p_t(x) along each edge in. + leaving = sum(float(u[x, y] * p[y]) for x in others) # u_t(x, y) p_t(y) along each edge out. + return arriving, leaving # Their difference is p_dot_t(y). + + +def integrate(u, p, steps): + """Carry p forward by ``steps`` Euler steps over one unit of time under a constant velocity.""" + h = 1.0 / steps # The time increment of one step. + for _ in range(steps): + p = p + h * master_equation(u, p) # p_{t+h} = p_t + h p_dot_t. + return p # [K] after one unit of time. + + +# %% 05. One velocity on DNA, with numbers +# +# Let the token be a single DNA base, let every base other than G jump to G at rate 1, and let G +# never leave. Probability moves into G and the total stays at one, which the two rules guarantee. + + +def worked_master_equation(): + """Section 5.3.2 with p_t = (0.1, 0.2, 0.3, 0.4) and the all-to-G velocity.""" + u = jump_to_one_state(DNA.index("G")) # [4, 4] the matrix printed in the chapter. + p = torch.tensor([0.1, 0.2, 0.3, 0.4]) # [4] over A, C, G, T. + p_dot = master_equation(u, p) # (-0.1, -0.2, 0.7, -0.4). + arriving, leaving = arrivals_and_departures(u, p, DNA.index("G")) # 0.7 in and 0 out. + return u, p, p_dot, arriving, leaving + + +# %% 06. Survival and hazard +# +# The schedule factor kappa_dot_t/(1 - kappa_t) is a hazard rate, and the probability that a position +# still holds its source token is its survival curve. Separating variables gives S(t) = 1 - kappa_t, +# the same weight the mixture path puts on the source token. + + +def hazard(t): + """kappa_dot_t/(1 - kappa_t) on the linear schedule, which is 1/(1 - t).""" + return 1.0 / (1.0 - t) # 1 at t = 0, 2 at t = 0.5, 10 at t = 0.9. + + +def survival(steps, upper=0.9): + """Integrate dS/dt = -hazard(t) S(t) from S(0) = 1 and return S at ``upper``.""" + h, s = upper / steps, 1.0 # One increment and the initial survival probability. + for k in range(steps): + s = s - h * hazard(k * h) * s # Euler on the survival equation. + return s # Exactly 1 - kappa_t = 1 - upper, since the two factors telescope at every grid. + + +# %% 07. The continuous and discrete velocities side by side +# +# In continuous space the velocity is the schedule factor times the posterior mean minus the current +# point. On tokens it is the same factor times the posterior distribution minus the point mass on the +# current token. One step of size h displaces the point by h u_t(x); on tokens it redraws from the +# denoiser with probability lambda_t, and the two agree in expectation on the one-hot embedding. + + +def continuous_velocity(x, x1_hat, t): + """u_t(x) = kappa_dot/(1 - kappa) [x_hat_{1|t}(x) - x] at a point of R^d. x: [d] -> [d].""" + return hazard(t) * (x1_hat - x) # The displacement that remains, divided by the time remaining. + + +def discrete_velocity(x_i, p1, t, size=DNA_SIZE): + """u^i_t(., x) = kappa_dot/(1 - kappa) [p_{1|t}(. | x) - delta_{x^i}]. Returns [K].""" + return hazard(t) * (p1 - F.one_hot(torch.tensor(x_i), size).double()) # Signed, sums to zero. + + +def side_by_side(x_i, p1, t, h): + """The continuous Euler step at the one-hot point and the discrete redraw distribution.""" + x = F.one_hot(torch.tensor(x_i), len(p1)).double() # The current token as a corner of R^K. + moved = x + h * continuous_velocity(x, p1, t) # Displace by h times the velocity. + lam = h * hazard(t) # lambda_t, the redraw probability of the two-part rule. + drawn = (1.0 - lam) * x + lam * p1 # Keep the current token, or redraw from the denoiser. + return moved, drawn, lam # The two coincide, which is the last row of Table 5.1. + + +# %% 08. The report +def main(): + parser = common_parser(__doc__.split("\n")[0]) + args = parser.parse_args() + set_seed(args.seed) + + u, p, p_dot, arriving, leaving = worked_master_equation() + report = { + "two_rules_hold": obeys_the_two_rules(u), # Nonnegative off the diagonal, columns sum to zero. + "diagonal_recovered": bool(torch.allclose(u, set_diagonal(u))), # The rules fix the diagonal. + "p_dot": [round(float(v), 4) for v in p_dot], # (-0.1, -0.2, 0.7, -0.4). + "p_dot_sums_to": round(float(p_dot.sum()), 6), # Zero, so probability is conserved. + "flux_into_G": round(arriving, 4), # 0.7, the three incoming rates against their probabilities. + "flux_out_of_G": round(leaving, 4), # 0, since G has no outgoing rate. + "euler_step_h_0.1": [round(float(v), 4) for v in p + 0.1 * p_dot], # (0.09, 0.18, 0.37, 0.36). + "hazard_at_0.5": round(hazard(0.5), 4), # 2. + "hazard_at_0.9": round(hazard(0.9), 4), # 10. + "survival_at_0.9": f"{survival(args.steps):.6f}", # 1 - 0.9 = 0.1 at any grid. + "survival_exact": 0.1, # S(t) = 1 - kappa_t. + } + + # Section 5.8.1: one position holding G with the denoiser of the chapter. + p1 = torch.tensor([0.1, 0.6, 0.2, 0.1]) # [4] over A, C, G, T. + moved, drawn, lam = side_by_side(DNA.index("G"), p1, t=0.5, h=0.1) + report["discrete_velocity_at_G"] = [round(float(v), 4) for v in discrete_velocity(2, p1, 0.5)] + report["continuous_step"] = [round(float(v), 4) for v in moved] # (0.02, 0.12, 0.84, 0.02). + report["discrete_redraw"] = [round(float(v), 4) for v in drawn] # The identical vector. + report["redraw_probability"] = round(lam, 4) # lambda_t = 0.2. + + # The same velocity run for a whole unit of time, starting from the data's letter frequency. + splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed) + counts = torch.bincount(splits["train_tokens"].flatten(), minlength=DNA_SIZE).double() + start = counts / counts.sum() # [4] the empirical frequency of A, C, G, T in the training DNA. + final = integrate(u, start, args.steps) # Euler on the master equation, args.steps increments. + exact = start.clone() + exact[:] = start * math.exp(-1.0) # Every base other than G decays as exp(-t). + exact[DNA.index("G")] = 1.0 - (start.sum() - start[DNA.index("G")]) * math.exp(-1.0) + report["dna_start"] = [round(float(v), 4) for v in start] # The GC-rich background of the data. + report["dna_after_one_unit"] = [round(float(v), 4) for v in final] # Euler at args.steps. + report["dna_exact"] = [round(float(v), 4) for v in exact] # The closed-form solution. + report["euler_error"] = round(float((final - exact).abs().max()), 6) # Shrinks with args.steps. + print_report("a jump rate written as a velocity", report) + save_run(args.out, config=vars(args), report=report) + + +if __name__ == "__main__": + main() diff --git a/lecture_6/README.md b/lecture_6/README.md new file mode 100644 index 0000000000000000000000000000000000000000..30d8acf99ca43e10da1c2748e33ce3b24a921c85 --- /dev/null +++ b/lecture_6/README.md @@ -0,0 +1,150 @@ +# Lecture 6: Flow Maps + +Chapter 6 of the notes, `book/chapters/ch6_flow_maps.tex`. + +Chapters 2 to 5 learned a local change and applied it again and again, so the number of steps set +the cost of a sample. This lecture learns the map between two times directly. The map +`X_{s,t}(x)` is fixed by three axioms, identity, cocycle and tangent, and every file writes it the +same way, as the state plus the elapsed time times an average velocity the network produces, which +makes the identity axiom hold for any weights at all. Differentiating the cocycle in the end time +gives the Lagrangian identity, in the start time the Eulerian identity, and leaving it finite +gives the composition identity. Each one becomes a squared residual, each is paired with a flow +matching anchor on the diagonal, and the rest of the lecture carries the construction into five +settings: a point in space, a distribution over tokens, a Brownian path, a dimension that grows, +and a conditioning state. + +Every file is standalone, runs offline on a laptop CPU, and prints the numbers the chapter works +by hand. + +## The files + +| File | Command | What it builds | +| --- | --- | --- | +| `flow_map.py` | `python lecture_6/flow_map.py` | No learning. The chapter's velocity, the map by RK4, the three axioms, the inverse, the Jacobian and its differential equation, Liouville's formula, density transport and the generator | +| `residuals.py` | `python lecture_6/residuals.py` | The Lagrangian, Eulerian and composition residuals and the diagonal anchor, from the chapter's listing, with four models trained on the four Gaussian blobs | +| `meanflow.py` | `python lecture_6/meanflow.py` | MeanFlow: the Eulerian identity in the average-velocity parameterization, with the total derivative from one forward-mode pass | +| `shortcut.py` | `python lecture_6/shortcut.py` | Shortcut models: the composition identity indexed by step size, on a dyadic ladder from 1/128 to 1 | +| `consistency.py` | `python lecture_6/consistency.py` | Consistency models on the edge `t = 1` with the boundary coefficients, and consistency trajectory models on the whole triangle | +| `fmlm.py` | `python lecture_6/fmlm.py` | Flow Map Language Models: the two-time denoiser, the convex form, the cocycle in denoiser coordinates, the two gradients a softmax head sees, two KL terms and the decoding clock, on DNA | +| `categorical_map.py` | `python lecture_6/categorical_map.py` | Categorical Flow Maps: the Lagrangian identity with a softmax head, the teacher gap and drift, and the bound that replaces the residual, on phrases | +| `discrete_map.py` | `python lecture_6/discrete_map.py` | Discrete Flow Maps: why the derivative targets leave the simplex, and the repair in logit space, on DNA | +| `stochastic_map.py` | `python lecture_6/stochastic_map.py` | Strong stochastic flow maps: the solution map of an SDE, the weak map's floor, the Legendre coefficients and the Chen relations | +| `expanding.py` | `python lecture_6/expanding.py` | Expanding flow maps: a state whose length grows, local clocks, the composition identity under growth, and one draw of the insertion noise | +| `posterior_map.py` | `python lecture_6/posterior_map.py` | Diamond Maps and Meta Flow Maps: the inner clock, posterior samples, the value function from K samples, its gradient and the importance weight | + +Every file takes the shared flags of `cis6270.runner`. `--steps 20 --quiet` finishes in a few +seconds. In `flow_map.py`, `stochastic_map.py` and `posterior_map.py` nothing is trained, so +`--steps` sets the Runge-Kutta substep count, the Monte Carlo draw count and the inner substep +count instead, which each file's docstring says. + +## What a default run prints + +The chapter's worked numbers come out exactly, and the test suite asserts them. + +| Quantity | Chapter | Printed | +| --- | --- | --- | +| `X_{0,1}(0.7)`, the cocycle through 0.5, the inverse, `det J_{0,1}` | 4.1, 4.1, 0.7, 3 | 4.100000, 4.100000, 0.700000, 3.000000 | +| `log det J_{0,1}` by Liouville, and the density ratio | 1.0986, 3.000 | 1.0986, 3.0000 | +| Stay-still residuals at `x = 1`, `s = 0`, `t = 1` | 1, 1, 0 | 1.0000, 1.0000, 0.0000 | +| MeanFlow on `dx/dt = x`: `1.718x = x + 0.718x` | 1.718, 1, 0.718 | 1.7183, 1.0000, 0.7183 | +| Shortcut doubling: two steps of 0.25 against one of 0.5 | 1.297443 | 1.297443 | +| Consistency distillation residual at `h = 0.01` | 8.189e-05 | 8.189e-05 | +| FMLM composed target, its four KL terms and their sum | (0.06, 0.12, 0.76, 0.06), 0.030747 | the same, 0.030747 | +| The two gradients in the logits at that target and student | (0.04, 0.03, -0.11, 0.04) and (0.0099, 0.01335, -0.03315, 0.0099) | the same | +| The same two once the student saturates to `q_1 = 1e-06` | -0.0600 and -6.10e-08 | -0.0600 and -4.6171e-08, see below | +| The decoding clock at 0.60, 0.75, 0.80, 0.85 | 0.006, 0.141, 0.456, 0.931 | 0.0057, 0.1411, 0.4561, 0.9313 | +| Four steps on the clock | 0.773490, 0.804481, 0.828095 | the same | +| Gamma on the last three times, and over the long first leg | 0.494620, 0.830144 | the same | +| Categorical gap and drift, and the bound around a zero residual | antiparallel, 0.199 | 1e-34 residual, 0.1989 | +| The same two terms averaged over three thousand states | 0.393, 0.095 | 0.3955, 0.0963 | +| Discrete derivative target, three entries negative | (-0.3333, -0.6667, 1.3333, -0.3333) | the same | +| The logit repair of (-0.284, 0.828, 0.456) | (0.0194, 0.6259, 0.3547) | the same | +| The weak map's floor `2K_1` | 6.389 | 6.3886 by Monte Carlo | +| Chen fusion of four half-interval draws | 0.50, -0.20 | 0.50, -0.20 | +| Local clocks at `t = 0.8`, and omega at (0.3, 0.6, 0.9) | 0.800, 0.692, 0.333 and 0.125 | the same | +| Jensen gap, and `E[V_K]` at K = 1, 2, 4, 8 | 0.434 and 1.000, 1.217, 1.340, 1.393 | the same | +| Meta Flow Maps fixed point against the per-sample ratio | 0.625 and 0.540 | the same | + +One of those disagrees with the notes. Section 6.27.4 quotes a saturated squared-error gradient of +`-6.10e-08`. Its own expression, the gradient of half the squared error through the softmax +Jacobian, + + d/dz_i (1/2) sum_j (q_j - p_j)^2 = q_i (q_i - p_i) - q_i sum_j q_j (q_j - p_j) , + +evaluated at `p = (0.06, 0.12, 0.76, 0.06)` and the student moved to `q_1 = 1e-06` with the other +three renormalized, gives `-4.6171e-08`. The first term on its own is `-6.00e-08`, so the quoted +number is that term with the second one dropped. `fmlm.py` prints the full expression, takes both +gradients by automatic differentiation, and reproduces the unsaturated pair of Section 6.27.4 +exactly. The point the chapter draws from the number is unaffected, since `-0.0600` against +`-4.6171e-08` is still six orders of magnitude. + +The trained numbers from one default run on a laptop CPU, with seed 0. + +| Run | Headline | +| --- | --- | +| `residuals.py`, 600 steps each | distance to its own ODE 0.900 for the diagonal alone, 0.481, 0.376 and 0.541 with the Lagrangian, Eulerian and composition residuals; fraction on a mode 0.44 against 0.74, 0.66 and 0.70 | +| `meanflow.py`, 2000 steps | nearest mode 0.627 at one step, 0.522 at two, 0.505 at four, all four modes found | +| `shortcut.py`, 2000 steps | nearest mode 0.605 at one step and 0.508 at eight; the doubling residual rises from below 1e-04 at `d = 1/128` to 0.034 at `d = 1/2` | +| `consistency.py`, 1500 steps each | nearest mode 0.513 for the consistency model at one step, 0.515, 0.420 and 0.470 for the trajectory model at one, two and four | +| `fmlm.py`, 1500 steps | GC 0.586 at one step against 0.581 in the training data, motif 0.63, token entropy 1.371 of a possible 1.386, every sequence distinct | +| `categorical_map.py`, 1000 steps each | every phrase in the language at both step counts; trained on the residual, 0.56 of one-step phrases distinct at entropy 2.75; trained on the bound, 0.33 at 2.54, and 0.12 at 2.05 by four steps | +| `discrete_map.py`, 1500 steps | GC 0.609 at one step, motif 0.81, entropy 1.362, every sequence distinct; the unrepaired additive target is negative on 0.05% of its entries | +| `stochastic_map.py` | the strong map's terminal cloud sits 8.8e-04 from the finest crossing at two hops, 1.1e-04 at four, 1.4e-05 at eight and 2e-06 at sixteen, the weak map's 3.06 at every hop count | +| `expanding.py`, 1500 steps | GC 0.617 at one step, 0.621 on coordinates born at the source and 0.613 on coordinates born later; motif 0.75 | +| `posterior_map.py` | on the blobs, the exact value -0.1127, the 16-sample estimate -0.1088, the single endpoint -0.1628 | + +The DNA files draw 1024 training sequences of the default length 16, which have GC fraction 0.581 +and motif fraction 0.749. They read 64 samples, as the phrase file does, so a reported fraction +carries about 0.06 of sampling noise and its second digit moves between seeds. + +Longer runs, where they change the picture: + +- `residuals.py --steps 2400` moves the four distances to 1.033, 0.171, 0.247 and 0.511 and the + four mode fractions to 0.27, 0.82, 0.71 and 0.79, against Table 6.2's 0.896, 0.292, 0.242 and + 0.325 at 0.12, 0.84, 0.82 and 0.84. The gap between the anchor alone and the anchor plus any + residual widens with the budget, and the three residuals stay inside each other's spread. +- `meanflow.py --steps 8000` brings the one-step nearest mode from 0.627 to 0.437, below its own + two-step and four-step figures, which is the one-step map overtaking the grid. +- `categorical_map.py --steps 3000` widens the gap between the two students, to 0.60 of one-step + phrases distinct for the residual against 0.27 for the bound, and 0.54 against 0.09 at four + steps. The bounded student loses variety as its budget grows, which is what Section 6.28.4 + predicts, since the bound is minimized by a prediction that does not move with the end time. +- The DNA runs are at their plateau by their default budgets. At `--steps 6000` the one-step motif + fraction of `fmlm.py` reads 0.61 and of `discrete_map.py` 0.69, inside the sampling noise of the + default run. + +Timings at the default flags, on one core. Three or four seconds for `flow_map.py`, +`stochastic_map.py` and `posterior_map.py`, which train nothing, and twelve to twenty-seven +seconds for the eight that do, including `residuals.py` with its four models and +`categorical_map.py` and `consistency.py` with their two. + +## What is simplified + +- The four blob models of `residuals.py` train for 600 steps on a four-layer network, so the + distances are larger than Table 6.2 of the chapter, which trains longer. The ordering the + chapter draws, the anchor alone against the anchor plus any residual, comes out at this budget. +- `shortcut.py` evaluates the zero-step term and the doubling term on the whole batch and weights + them three to one, in place of the chapter's split of the batch in that proportion. The two have + the same expectation and the weighted form costs one network call more per step. +- `consistency.py` distils from the exact marginal velocity of the training set, which is a + softmax over the data points, in place of a separately trained flow matching teacher. The + teacher's discretization error is then the only error the student inherits, which is the point + of Section 6.20. +- `posterior_map.py` builds the inner map from the exact posterior over the training points, in + place of training a conditional family, so the inner velocity is a closed form and the value, + its gradient and the Jensen gap are all available to compare against. +- The decoding clock of `fmlm.py` is computed by quadrature on the one-position integral, which is + exact for the Gaussian interpolant and says nothing about correlations between positions. +- `categorical_map.py` and `discrete_map.py` keep the end time below one, because the drift term + of each loss divides by the remaining time. `fmlm.py` has no such term, so it draws a quarter of + each batch at `t = 1`, the slice Section 6.26 asks for by name. +- `stochastic_map.py` carries three Legendre coefficients and fuses them with the Chen relations + written out for `n = 0, 1, 2`. A fourth coefficient needs the constants `c_{3,m}`, which the + chapter does not list. +- Every network is `cis6270.nets.TwoTimeMLP`. The discrete files flatten the one-hot lift into its + input and read its output as per-position logits, so one architecture covers both halves of the + lecture and any difference in the samples comes from the loss. +- Each file sets `torch.set_num_threads(1)`. These networks are small enough that splitting one + matrix product across cores costs more than it saves. `tests/conftest.py` sets the same thing + for the test session, so the line in each file is the one that acts when a reader runs the + script directly. diff --git a/lecture_6/categorical_map.py b/lecture_6/categorical_map.py new file mode 100644 index 0000000000000000000000000000000000000000..dc58e1a8dbe772467d18540f36cd32facf2a1999 --- /dev/null +++ b/lecture_6/categorical_map.py @@ -0,0 +1,244 @@ +"""Categorical Flow Maps, the Lagrangian identity on the simplex and the bound that replaces it. + +A softmax head makes the two-time prediction a distribution at every pair of times, and the +Lagrangian identity rewritten in those coordinates reads pi_{s,t} = pi_{t,t}(X_{s,t}) minus a +coefficient times the end-time derivative of the prediction. The derivative has no sign +constraint, so the residual splits into a teacher gap, which is a difference of two distributions, +and a drift, which is the scaled derivative. A parallelogram bound and Pinsker's inequality turn +the squared residual into four KL divergences plus two squared drifts, and at the exact map the +two pieces are antiparallel, so the bound is strictly positive where the residual is zero. + +Chapter 6, Section 6.28 of the notes. Trained on the four-token phrases of +``cis6270.data.phrases``. Prints the chapter's worked cancellation on three tokens, where the gap +is the exact negative of the drift and the bound reads 0.199 around a residual of zero, and then +two models trained on the residual and on its bound. + + python lecture_6/categorical_map.py # two models, 1000 steps each + python lecture_6/categorical_map.py --steps 20 # a few seconds, for a first look +""" + +# %% 01. Imports and configuration +import torch + +from cis6270.data import phrases +from cis6270.metrics import token_entropy, unique_fraction +from cis6270.nets import TwoTimeMLP, count_parameters +from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train + + +torch.set_num_threads(1) # These networks are small, and splitting one matrix product costs more +# than it saves. ``tests/conftest.py`` sets the same thing for the test session, so the line +# here is the one that acts when a reader runs this script directly. + +BOOK_PRIOR = (0.5, 0.3, 0.2) # The three data probabilities of the chapter's worked example. +BOOK_SOURCE = (-0.5739, -0.8344, -0.0584) # The source draw I_0 the chapter starts from. +BOOK_STATES = 3000 # States the chapter averages the two terms of the bound over. + + +# %% 02. The exact model on three tokens +# +# With a Gaussian source and the linear interpolant, the posterior over the data token is a +# softmax whose only state dependence is the coordinate of that token. The velocity follows from +# D_t(x) = x + (1-t) b(x,t), and the flow integrates to twelve digits with Runge-Kutta. +def exact_denoiser(x, t, prior): + """pi_{t,t}(x), the posterior over the data token given the state. [B, K].""" + return torch.softmax(prior.log() + t * x / (1.0 - t) ** 2, dim=-1) + + +def exact_velocity(x, t, prior): + """b(x,t) = (pi_{t,t}(x) - x)/(1 - t) for the linear interpolant. [B, K].""" + return (exact_denoiser(x, t, prior) - x) / (1.0 - t) + + +def exact_flow(x, s, t, prior, n: int = 2000): + """X_{s,t}(x) by fourth-order Runge-Kutta on the exact velocity.""" + h = (t - s) / n + for k in range(n): + u = s + k * h + k1 = exact_velocity(x, u, prior) + k2 = exact_velocity(x + 0.5 * h * k1, u + 0.5 * h, prior) + k3 = exact_velocity(x + 0.5 * h * k2, u + 0.5 * h, prior) + k4 = exact_velocity(x + h * k3, u + h, prior) + x = x + (h / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4) + return x + + +def exact_two_time(x, s, t, prior): + """pi_{s,t}(x) read off the convex parameterization, [(1-s) X_{s,t}(x) - (1-t) x]/(t-s).""" + return ((1.0 - s) * exact_flow(x, s, t, prior) - (1.0 - t) * x) / (t - s) + + +# %% 03. The coefficient, the gap and the drift +# +# gamma_{s,t} = (t-s)(1-t)/(1-s) vanishes on the diagonal and at t = 1 and never exceeds 1/4 on +# the square. The residual scaled by the remaining time splits as (1-t) r = a + b, with a the gap +# between the student's prediction and the prediction at the transported state. +def gamma_coefficient(s, t): + """gamma_{s,t} = (t-s)(1-t)/(1-s).""" + return (t - s) * (1.0 - t) / (1.0 - s).clamp(min=1e-6) + + +def end_time_derivative(x, s, t, prior, n: int = 2000): + """d pi_{s,t}/dt from the integral representation, (1-s)/((t-s)(1-t)) times the bracket.""" + transported = exact_flow(x, s, t, prior, n) + prediction = ((1.0 - s) * transported - (1.0 - t) * x) / (t - s) + at_output = exact_denoiser(transported, t, prior) + return (1.0 - s) / ((t - s) * (1.0 - t)) * (at_output - prediction), prediction, at_output + + +# %% 04. The worked cancellation of Section 6.28 +# +# The derivative above comes from the closed form of Section 6.26, which already carries the +# cancellation inside it, so the central difference of pi_{s,t} in t below is the entry that checks +# the two pieces against each other without assuming the identity. +def worked_cancellation(s: float = 0.2, t: float = 0.6): + """The three-token example, with every number integrated or differentiated here.""" + prior = torch.tensor(BOOK_PRIOR, dtype=torch.float64) + source = torch.tensor(BOOK_SOURCE, dtype=torch.float64) + x = (1.0 - s) * source + s * torch.eye(3, dtype=torch.float64)[0] # Data token is the first. + derivative, prediction, at_output = end_time_derivative(x, s, t, prior) + gamma = (t - s) * (1.0 - t) / (1.0 - s) + gap = prediction - at_output # a_{s,t}, a difference of two simplex points. + drift = gamma * derivative # b_{s,t}, the scaled derivative of the prediction. + divergence = (at_output * (at_output / prediction).log()).sum() # KL(teacher || student). + difference = (exact_two_time(x, s, t + 1e-5, prior) + - exact_two_time(x, s, t - 1e-5, prior)) / 2e-5 # A central check. + return { + "state x": [round(v, 4) for v in x.tolist()], + "X_s,t(x)": [round(v, 4) for v in exact_flow(x, s, t, prior).tolist()], + "pi_s,t(x)": [round(v, 4) for v in prediction.tolist()], + "pi_t,t(X_s,t(x))": [round(v, 4) for v in at_output.tolist()], + "gamma_s,t": gamma, + "d pi / dt": [round(v, 4) for v in derivative.tolist()], + "central difference": [round(v, 4) for v in difference.tolist()], + "gap a": [round(v, 3) for v in gap.tolist()], + "drift b": [round(v, 3) for v in drift.tolist()], + "residual |a + b|^2": float(((gap + drift) ** 2).sum()), + "4 KL": 4.0 * float(divergence), + "2 |b|^2": 2.0 * float((drift ** 2).sum()), + "the bound": 4.0 * float(divergence) + 2.0 * float((drift ** 2).sum()), + "teacher entropy": -float((at_output * at_output.log()).sum()), + "cross-entropy": -float((at_output * prediction.log()).sum()), + } + + +def averaged_bound(count: int = BOOK_STATES, s: float = 0.2, t: float = 0.6, seed: int = 0, + n: int = 400): + """The same two terms averaged over states drawn from the interpolant at the same times.""" + prior = torch.tensor(BOOK_PRIOR, dtype=torch.float64) + generator = torch.Generator().manual_seed(seed) + which = torch.multinomial(prior.expand(count, 3), 1, generator=generator).squeeze(1) + x1 = torch.eye(3, dtype=torch.float64)[which] # [N, 3] the lifted data token. + x0 = torch.randn(count, 3, generator=generator, dtype=torch.float64) # [N, 3] the source. + x = (1.0 - s) * x0 + s * x1 # [N, 3] the interpolant at the start time. + derivative, prediction, at_output = end_time_derivative(x, s, t, prior, n) + drift = ((t - s) * (1.0 - t) / (1.0 - s)) * derivative + divergence = (at_output * (at_output / prediction).log()).sum(-1).mean() + return {"averaged 4 KL": 4.0 * float(divergence), + "averaged 2 |b|^2": 2.0 * float((drift ** 2).sum(-1).mean()), + "averaged residual": float(((prediction - at_output + drift) ** 2).sum(-1).mean())} + + +# %% 05. The student, and the two losses it can be trained on +def predict(model, x, s, t, length, vocab): + """pi_{s,t}(x) at every position, a softmax head over the vocabulary.""" + return torch.softmax(model(x.flatten(1), s, t).view(-1, length, vocab), dim=-1) + + +def pieces(model, xs, s, t, length, vocab): + """The gap a and the drift b of the split, with the teacher detached.""" + f = lambda tt: predict(model, xs, s, tt, length, vocab) + prediction, derivative = torch.func.jvp(f, (t,), (torch.ones_like(t),)) # One forward mode pass. + keep, move = (1.0 - t) / (1.0 - s), (t - s) / (1.0 - s) + transported = keep[:, None, None] * xs + move[:, None, None] * prediction # X_{s,t}(x). + teacher = predict(model, transported.detach(), t, t, length, vocab).detach() # pi_{t,t}. + drift = gamma_coefficient(s, t)[:, None, None] * derivative + return prediction - teacher, drift, teacher, prediction + + +def make_loss(model, data, length, vocab, batch_size, device, bounded: bool): + eye = torch.eye(vocab, device=device) + + def loss_fn(step: int) -> torch.Tensor: + index = torch.randint(0, len(data), (batch_size,), device=device) + x1, ids = eye[data[index]], data[index] # [B, L, K] and [B, L]. + x0 = torch.randn_like(x1) # [B, L, K] the Gaussian source. + t_diag = torch.rand(batch_size, device=device) + xt = (1 - t_diag[:, None, None]) * x0 + t_diag[:, None, None] * x1 + anchor = -(predict(model, xt, t_diag, t_diag, length, vocab).clamp_min(1e-9).log() + .gather(-1, ids[..., None])).mean() # Cross-entropy on the diagonal. + + s = 0.9 * torch.rand(batch_size, device=device) # [B] start time, kept off t = 1. + t = s + (0.95 - s) * torch.rand(batch_size, device=device) # [B] end time below one. + xs = (1 - s[:, None, None]) * x0 + s[:, None, None] * x1 + gap, drift, teacher, prediction = pieces(model, xs, s, t, length, vocab) + if bounded: # Four KL divergences plus two squared drifts, the bound of Section 6.28. + divergence = (teacher * (teacher.clamp_min(1e-9) + / prediction.clamp_min(1e-9)).log()).sum(-1).mean() + return anchor + 4.0 * divergence + 2.0 * (drift ** 2).sum(-1).mean() + return anchor + ((gap + drift) ** 2).sum(-1).mean() # The residual itself. + + return loss_fn + + +# %% 06. Sampling, and what a valid phrase is +@torch.no_grad() +def sample(model, count, length, vocab, steps, device): + """Walk a uniform grid with the convex form and read the phrase off the last prediction.""" + x = torch.randn(count, length, vocab, device=device) + grid = torch.linspace(0.0, 1.0, steps + 1) + for a, b in zip(grid[:-1], grid[1:]): + s, t = a.expand(count).to(device), b.expand(count).to(device) + prediction = predict(model, x, s, t, length, vocab) + keep, move = (1.0 - t) / (1.0 - s), (t - s) / (1.0 - s) + x = keep[:, None, None] * x + move[:, None, None] * prediction + return x.argmax(-1) # [N, L]. + + +def valid_fraction(drawn, group: int = 4): + """A phrase is in the language when position i holds a token from the i-th quarter.""" + slots = torch.arange(drawn.shape[1], device=drawn.device) # [L]. + return float(((drawn // group) == slots[None, :]).all(dim=1).to(torch.get_default_dtype()).mean()) + + +# %% 07. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=1000, batch_size=128, samples=256, sample_steps=4) + args = parser.parse_args() + set_seed(args.seed) + device = pick_device(args.device) + tokens, vocabulary = phrases(num_phrases=1024, seed=args.seed) + data, length, vocab = tokens.to(device), tokens.shape[1], len(vocabulary) + + report = worked_cancellation() + report.update(averaged_bound(seed=args.seed)) # The chapter's three thousand states. + if not args.quiet: + print(f"gap {report['gap a']} and drift {report['drift b']} cancel to " + f"{report['residual |a + b|^2']:.2e}, bound {report['the bound']:.3f}") + + samples: list[str] = [] + for name, bounded in (("the residual", False), ("the bound", True)): + set_seed(args.seed) # Both students start from the same weights and see the same batches. + model = TwoTimeMLP(dim=length * vocab, width=args.width).to(device) + if not args.quiet: + print(f"\ntraining on {name} parameters={count_parameters(model):,}") + losses = train(model, make_loss(model, data, length, vocab, args.batch_size, device, + bounded), steps=args.steps, lr=args.lr, quiet=args.quiet) + for steps in (1, args.sample_steps): + torch.manual_seed(args.seed + 1) + drawn = sample(model, args.samples, length, vocab, steps, device) + report[f"{name}, {steps} step(s): valid"] = valid_fraction(drawn) + report[f"{name}, {steps} step(s): distinct"] = unique_fraction(drawn) + report[f"{name}, {steps} step(s): entropy"] = token_entropy(drawn, vocab) + report[f"{name}: final loss"] = sum(losses[-20:]) / len(losses[-20:]) + drawn = sample(model, 4, length, vocab, args.sample_steps, device) + samples.append(f"{name:<14s} " + " | ".join( + " ".join(vocabulary[int(i)] for i in row) for row in drawn[:2])) + + print_report("Categorical Flow Maps, Chapter 6", report, samples) + save_run(args.out, config=vars(args), report=report, samples=samples) + + +if __name__ == "__main__": + main() diff --git a/lecture_6/consistency.py b/lecture_6/consistency.py new file mode 100644 index 0000000000000000000000000000000000000000..657d3d9b5ceed69b5c43a43544df40f0ded6b1b5 --- /dev/null +++ b/lecture_6/consistency.py @@ -0,0 +1,231 @@ +"""Consistency models and consistency trajectory models, one time fixed and then both free. + +A consistency model fixes the end time at the data end and learns the top edge of the triangle of +start and end times, with the identity axiom imposed at the corner by the boundary coefficients +c_skip and c_out. A consistency trajectory model puts the second time back and learns the whole +triangle, writing the map as a weighted average of the input and a head that predicts a +destination. Both are distilled here from a teacher velocity available in closed form. + +Chapter 6, Section 6.20 of the notes. Trained on the four Gaussian blobs of +``cis6270.data.gaussian_mixture_2d``. Prints the boundary parameterization on dx/dt = x, where the +model returns e^{0.5}, e^{0.1} and 1 at three start times, the teacher's error floor of +8.189e-05, and the destination head of a consistency trajectory model at two pairs of times. + + python lecture_6/consistency.py # two models, 1500 steps each + python lecture_6/consistency.py --steps 20 # a few seconds, for a first look +""" + +# %% 01. Imports and configuration +import math + +import torch + +from cis6270.data import gaussian_mixture_2d +from cis6270.metrics import mean_log_density, mode_report +from cis6270.nets import TwoTimeMLP, count_parameters +from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train + + +torch.set_num_threads(1) # These networks are small, and splitting one matrix product costs more +# than it saves. ``tests/conftest.py`` sets the same thing for the test session, so the line +# here is the one that acts when a reader runs this script directly. + +TEACHER_STEP = 0.02 # h, the length of the teacher's solver step inside the distillation loss. + + +# %% 02. The teacher velocity, in closed form on the empirical data +# +# The posterior over the data point given the interpolant is a softmax over the training set, so +# the exact marginal velocity of the four blobs needs no training. A trained flow matching model +# would serve in its place, with its own error on top of the solver's. +def teacher_denoiser(x, t, data): + """D_t(x) = E[X_1 | X_t = x], the posterior mean over the training points. [B, 2].""" + spread = (1.0 - t).clamp(min=1e-4)[:, None] # [B, 1] the width of the conditional. + squared = ((x[:, None, :] - t[:, None, None] * data[None]) ** 2).sum(-1) # [B, N]. + weights = torch.softmax(-squared / (2.0 * spread ** 2), dim=1) # [B, N] posterior over data. + return weights @ data # [B, 2] the posterior mean. + + +def teacher_velocity(x, t, data): + """b(x,t) = (D_t(x) - x)/(1 - t) for the linear interpolant. [B, 2].""" + return (teacher_denoiser(x, t, data) - x) / (1.0 - t).clamp(min=1e-4)[:, None] + + +def solver_step(x, t, h, data): + """One Euler step of the teacher, which the chapter writes as Phi_h.""" + return x + h * teacher_velocity(x, t, data) # [B, 2] the state at time t + h. + + +# %% 03. The boundary parameterization, checked on the scalar flow +# +# f_theta(x, s) = c_skip(s) x + c_out(s) F_theta(x, s) with c_skip = 1 and c_out = 1 - s. At the +# data end the output coefficient is zero and the model returns its input, so no training signal +# is spent on a condition we already know. +def scalar_consistency(x_value: float = 1.0): + """The velocity head, the output coefficient and the model at three start times.""" + out = {} + for s in (0.5, 0.9, 1.0): + remaining = 1.0 - s # c_out(s), which vanishes at the data end. + head = x_value * math.expm1(remaining) / remaining if remaining > 0 else x_value + out[f"c_out({s})"] = remaining + out[f"v_{s},1(1)"] = head + out[f"f({s})"] = x_value + remaining * head # Equals e^{1-s} on this flow. + return out + + +def teacher_floor(x_value: float = 1.0, s: float = 0.5, h: float = 0.01): + """The distillation residual at the exact map, which is the teacher's own Euler error.""" + exact = lambda x, t: x * math.exp(1.0 - t) # f(x, t) = x e^{1-t} on dx/dt = x. + moved = x_value + h * x_value # One teacher Euler step, which undershoots e^{h}. + residual = exact(x_value, s) - exact(moved, s + h) + local = math.exp(h) - 1.0 - h # The local error of that step. + return {"f(1, 0.5)": exact(x_value, s), "f(1.01, 0.51)": exact(moved, s + h), + "distillation residual": residual, "local Euler error": local, + "carried to the data end": local * math.exp(1.0 - s - h)} + + +# %% 04. Consistency trajectory models, with both times free +# +# X_theta(x, s, t) = a_{s,t} x + (1 - a_{s,t}) g_theta(x, s, t) with a_{s,t} = (1-t)/(1-s), which +# equals one on the diagonal, so the map returns its input there for any head whatever. +def ctm_weight(s, t): + """a_{s,t} = (1-t)/(1-s), the share of the input kept by the step.""" + return (1.0 - t) / (1.0 - s).clamp(min=1e-4) + + +def ctm_map(model, x, s, t): + """The weighted average of the input and the destination head.""" + a = ctm_weight(s, t)[:, None] # [B, 1]. + return a * x + (1.0 - a) * model(x, s, t) # [B, 2]. + + +def scalar_ctm(x_value: float = 1.0): + """The head solved from the exact map at two pairs of times, and the two-route sampler.""" + out = {} + for s, t in ((0.0, 0.5), (0.0, 1.0)): + a = (1.0 - t) / (1.0 - s) + reached = x_value * math.exp(t - s) # The exact map of dx/dt = x. + out[f"a_{s},{t}"] = a + out[f"g_{s},{t}"] = (reached - a * x_value) / (1.0 - a) # Solve the weighted average. + short = x_value * math.exp(0.5) # X_{0,0.5}(1), the direct short step. + long_then_back = (x_value * math.exp(1.0)) * math.exp(-0.5) # X_{1,0.5}(X_{0,1}(1)). + out["short step"], out["long jump and return"] = short, long_then_back + out["two routes apart"] = abs(short - long_then_back) + return out + + +# %% 05. The two losses +# +# Both compare the model at one time against the model at a state the teacher has advanced by h, +# with the target detached, which is the frozen copy of the weights the chapter writes as theta +# minus. The consistency model holds the end time at one; the trajectory model draws it. +def make_consistency_loss(model, data, batch_size, device): + ones = torch.ones(batch_size, device=device) + + def loss_fn(step: int) -> torch.Tensor: + index = torch.randint(0, len(data), (batch_size,), device=device) + x1, x0 = data[index], torch.randn(batch_size, 2, device=device) + s = (1.0 - TEACHER_STEP) * torch.rand(batch_size, device=device) # [B] with s + h <= 1. + xs = (1 - s[:, None]) * x0 + s[:, None] * x1 # [B, 2] the interpolant at the start time. + with torch.no_grad(): # The second leg is the teacher's, and no gradient reaches it. + advanced = solver_step(xs, s, TEACHER_STEP, data) + target = model.flow_map(advanced, s + TEACHER_STEP, ones) # f_{theta-}(Phi_h(x_s), s+h). + return ((model.flow_map(xs, s, ones) - target) ** 2).mean() + + return loss_fn + + +def make_trajectory_loss(model, data, batch_size, device): + + def loss_fn(step: int) -> torch.Tensor: + index = torch.randint(0, len(data), (batch_size,), device=device) + x1, x0 = data[index], torch.randn(batch_size, 2, device=device) + s = (1.0 - TEACHER_STEP) * torch.rand(batch_size, device=device) # [B] start time. + t = s + TEACHER_STEP + (1.0 - s - TEACHER_STEP) * torch.rand(batch_size, device=device) + xs = (1 - s[:, None]) * x0 + s[:, None] * x1 # [B, 2]. + with torch.no_grad(): + advanced = solver_step(xs, s, TEACHER_STEP, data) + target = ctm_map(model, advanced, s + TEACHER_STEP, t) # The same end time, one h later. + return ((ctm_map(model, xs, s, t) - target) ** 2).mean() + + return loss_fn + + +# %% 06. Sampling +@torch.no_grad() +def consistency_sample(model, x0): + """One call of the edge map, from the source straight to the data end.""" + zeros = torch.zeros(len(x0), device=x0.device) + return model.flow_map(x0, zeros, torch.ones_like(zeros)) # [N, 2]. + + +@torch.no_grad() +def trajectory_sample(model, x0, steps: int): + """Cross the triangle in equal pieces, one evaluation of the weighted average per piece.""" + x = x0 + grid = torch.linspace(0.0, 1.0, steps + 1, device=x0.device) # [steps+1] the time grid. + for s, t in zip(grid[:-1], grid[1:]): + x = ctm_map(model, x, s.expand(len(x)), t.expand(len(x))) + return x # [N, 2]. + + +# %% 07. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=1500, batch_size=128, samples=512) + args = parser.parse_args() + set_seed(args.seed) + device = pick_device(args.device) + data = gaussian_mixture_2d(num_points=1024, seed=args.seed).to(device) # [1024, 2]. + + report = scalar_consistency() + report.update(teacher_floor()) + report.update(scalar_ctm()) + if not args.quiet: + print(f"boundary parameterization: f(0.5) = {report['f(0.5)']:.6f}, " + f"f(0.9) = {report['f(0.9)']:.6f}, f(1.0) = {report['f(1.0)']:.6f}") + print(f"teacher floor {report['distillation residual']:.6e} against the carried local " + f"error {report['carried to the data end']:.6e}") + + samples: list[str] = [] + set_seed(args.seed) + consistency = TwoTimeMLP(dim=2, width=args.width).to(device) + if not args.quiet: + print(f"\ntraining the consistency model parameters={count_parameters(consistency):,}") + edge_losses = train(consistency, + make_consistency_loss(consistency, data, args.batch_size, device), + steps=args.steps, lr=args.lr, quiet=args.quiet) + report["consistency final loss"] = sum(edge_losses[-20:]) / len(edge_losses[-20:]) + torch.manual_seed(args.seed + 1) + x0 = torch.randn(args.samples, 2, device=device) # [N, 2] the shared source points. + points = consistency_sample(consistency, x0) + stats = mode_report(points) + report["consistency one-step nearest"] = stats["nearest_distance"] + report["consistency one-step log density"] = mean_log_density(points.cpu()) + samples.append(f"consistency, 1 step: nearest {stats['nearest_distance']:.3f} " + f"modes {stats['modes_found']} entropy {stats['mode_entropy']:.3f}") + + set_seed(args.seed) + trajectory = TwoTimeMLP(dim=2, width=args.width).to(device) + if not args.quiet: + print(f"\ntraining the trajectory model parameters={count_parameters(trajectory):,}") + losses = train(trajectory, make_trajectory_loss(trajectory, data, args.batch_size, device), + steps=args.steps, lr=args.lr, quiet=args.quiet) + for steps in (1, 2, 4): + points = trajectory_sample(trajectory, x0, steps) + stats = mode_report(points) + report[f"trajectory {steps}-step nearest"] = stats["nearest_distance"] + samples.append(f"trajectory, {steps} step(s): nearest {stats['nearest_distance']:.3f} " + f"modes {stats['modes_found']} entropy {stats['mode_entropy']:.3f}") + report["trajectory one step against four"] = float( + (trajectory_sample(trajectory, x0, 1) - trajectory_sample(trajectory, x0, 4)).norm(dim=1) + .mean()) + report["trajectory final loss"] = sum(losses[-20:]) / len(losses[-20:]) + + print_report("Consistency and consistency trajectory models, Chapter 6", report, samples) + save_run(args.out, model=trajectory, config=vars(args), report=report, losses=losses, + samples=samples) + + +if __name__ == "__main__": + main() diff --git a/lecture_6/discrete_map.py b/lecture_6/discrete_map.py new file mode 100644 index 0000000000000000000000000000000000000000..3401bbff59035e978348d76bc1927d48b6c19915 --- /dev/null +++ b/lecture_6/discrete_map.py @@ -0,0 +1,218 @@ +"""Discrete Flow Maps, the derivative identities repaired in logit space. + +A curve that stays on the simplex has coordinates summing to one, so its velocity has coordinates +summing to zero, and a target built from that velocity has a negative entry as soon as one +coordinate of the prediction is smaller than the derivative pulling it down. Written in logits the +same derivative becomes a positive multiplicative factor, dividing by it keeps every coordinate +positive, and a softmax normalizes what is left, so the Kullback-Leibler divergence against the +repaired target is defined at every state. + +Chapter 6, Sections 6.26 and 6.29 of the notes. Trained on the synthetic DNA of +``cis6270.data.synthetic_dna``. Prints the derivative target of Section 6.26, three of whose four +entries are negative, the three-token repair, where the additive target (-0.284, 0.828, 0.456) +becomes the distribution (0.0194, 0.6259, 0.3547), and samples from the trained map. + + python lecture_6/discrete_map.py # 1500 steps, about half a minute + python lecture_6/discrete_map.py --steps 20 # a few seconds, for a first look +""" + +# %% 01. Imports and configuration +import torch + +from cis6270.data import DNA_SIZE, decode_dna, synthetic_dna +from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction +from cis6270.nets import TwoTimeMLP, count_parameters +from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train + + +torch.set_num_threads(1) # These networks are small, and splitting one matrix product costs more +# than it saves. ``tests/conftest.py`` sets the same thing for the test session, so the line +# here is the one that acts when a reader runs this script directly. + +FLOOR = 1e-3 # The factor is clipped away from zero before the logarithm, as the chapter notes. + + +# %% 02. Why the derivative target leaves the simplex +# +# d delta_{s,t}/dt = (1-s)/((t-s)(1-t)) [D_t(X_{s,t}(x)) - delta_{s,t}(x)]. The bracket is a +# difference of two simplex points, so its coordinates sum to zero, and the prefactor only scales +# it up. Nothing about the alphabet or the prior changes the sign pattern. +def derivative_target(prediction, at_output, s: float, t: float): + """The end-time derivative of the two-time denoiser, a direction and not a distribution.""" + prefactor = (1.0 - s) / ((t - s) * (1.0 - t)) # 6.6667 at s = 0.2, t = 0.8. + return prefactor * (at_output - prediction), prefactor + + +def worked_derivative(): + """The chapter's DNA example at s = 0.2, t = 0.8, whose target has three negative entries.""" + prediction = torch.tensor([0.10, 0.20, 0.60, 0.10]) # delta_{0.2,0.8} over (A, C, G, T). + at_output = torch.tensor([0.05, 0.10, 0.80, 0.05]) # D_{0.8} at the transported state. + target, prefactor = derivative_target(prediction, at_output, 0.2, 0.8) + return {"derivative prefactor": prefactor, + "bracket": [round(v, 4) for v in (at_output - prediction).tolist()], + "derivative target": [round(v, 4) for v in target.tolist()], + "target sum": float(target.sum()), + "negative entries": int((target < 0).sum())} + + +# %% 03. Differentiating a softmax +# +# The Jacobian of the softmax is diag(psi) - psi psi^T, so contracting it with the logit +# derivative gives psi times the centered derivative. The centring makes the coordinates of +# d psi/dt sum to zero and makes the identity insensitive to the shift invariance of the softmax. +def centered(prediction, logit_derivative): + """c_{s,t} = d z/dt - 1, the quantity the whole identity is built on.""" + mean = (prediction * logit_derivative).sum(-1, keepdim=True) # . + return logit_derivative - mean # [..., K] with = 0. + + +# %% 04. The identity as a multiplicative relation, and the repaired teacher +# +# psi_{t,t}(X_{s,t}(x)) = psi_{s,t}(x) * (1 + gamma c_{s,t}(x)) coordinatewise. Solving for the +# student divides by a strictly positive factor, and in logits the division is a subtraction of +# its logarithm, so the target is a softmax of an expression the network already produces. +def gamma_lagrangian(s, t): + """gamma_{s,t} = (t-s)(1-t)/(1-s), which never exceeds 1/4 on the square.""" + return (t - s) * (1.0 - t) / (1.0 - s).clamp(min=1e-6) + + +def gamma_eulerian(s, t): + """gamma^E_{s,t} = (1-s)(t-s)/(1-t), larger than the Lagrangian one and growing toward t = 1.""" + return (1.0 - s) * (t - s) / (1.0 - t).clamp(min=1e-6) + + +def repaired_teacher(teacher_logits, factor): + """Softmax(z_{t,t}(X_{s,t}(x)) - log(1 + gamma c)), a distribution by construction.""" + return torch.softmax(teacher_logits - factor.clamp(min=FLOOR).log(), dim=-1) + + +# %% 05. The repair worked on three tokens +def worked_repair(s: float = 0.2, t: float = 0.6): + """The chapter's three-token repair, with the additive target beside the multiplicative one.""" + prediction = torch.tensor([0.2, 0.5, 0.3]) # psi_{s,t}(x). + logit_derivative = torch.tensor([9.60, -2.28, -2.60]) # d z_{s,t}/dt. + gamma = (t - s) * (1.0 - t) / (1.0 - s) + c = centered(prediction, logit_derivative) # Already centered on this example. + derivative = prediction * c # d psi/dt, whose coordinates sum to zero. + teacher = torch.tensor([0.10, 0.60, 0.30]) # psi_{t,t}(X_{s,t}(x)) during training. + additive = teacher - gamma * derivative # Read straight off the identity. + factor = 1.0 + gamma * c # The multiplicative factor, positive in every coordinate. + quotient = teacher / factor # The teacher divided coordinatewise. + return { + "weighted average of dz/dt": float((prediction * logit_derivative).sum()), + "centered derivative": [round(v, 4) for v in c.tolist()], + "d psi/dt": [round(v, 4) for v in derivative.tolist()], + "d psi/dt sum": float(derivative.sum()), + "additive target": [round(v, 4) for v in additive.tolist()], + "additive target sum": float(additive.sum()), + "additive negative entries": int((additive < 0).sum()), + "multiplicative factor": [round(v, 4) for v in factor.tolist()], + "quotient": [round(v, 4) for v in quotient.tolist()], + "quotient mass": float(quotient.sum()), + "repaired target": [round(v, 4) for v in (quotient / quotient.sum()).tolist()], + "log of the factor": [round(v, 4) for v in factor.log().tolist()], + "gamma_s,t": gamma, + "gamma^E_s,t": (1.0 - s) * (t - s) / (1.0 - t), + } + + +# %% 06. The student, and the loss built on the repaired target +def logits_of(model, x, s, t, length, vocab): + """z_{s,t}(x), one vector of length K per position.""" + return model(x.flatten(1), s, t).view(-1, length, vocab) # [B, L, K]. + + +def make_loss(model, data, length, vocab, batch_size, device, counter): + eye = torch.eye(vocab, device=device) + + def loss_fn(step: int) -> torch.Tensor: + index = torch.randint(0, len(data), (batch_size,), device=device) + ids = data[index] # [B, L] the data tokens. + x1, x0 = eye[ids], torch.randn(batch_size, length, vocab, device=device) + t_diag = torch.rand(batch_size, device=device) + xt = (1 - t_diag[:, None, None]) * x0 + t_diag[:, None, None] * x1 + diagonal_logits = logits_of(model, xt, t_diag, t_diag, length, vocab) + anchor = torch.nn.functional.cross_entropy( # KL to the one-hot token, the diagonal anchor. + diagonal_logits.reshape(-1, vocab), ids.reshape(-1)) + + s = 0.9 * torch.rand(batch_size, device=device) # [B] start time. + t = s + (0.95 - s) * torch.rand(batch_size, device=device) # [B] end time below one. + xs = (1 - s[:, None, None]) * x0 + s[:, None, None] * x1 # [B, L, K]. + f = lambda tt: logits_of(model, xs, s, tt, length, vocab) + z, dz = torch.func.jvp(f, (t,), (torch.ones_like(t),)) # One forward-mode pass in t. + prediction = torch.softmax(z, dim=-1) # psi_{s,t}(x). + gamma = gamma_lagrangian(s, t)[:, None, None] + factor = 1.0 + gamma * centered(prediction, dz) # The multiplicative factor. + keep, move = (1.0 - t) / (1.0 - s), (t - s) / (1.0 - s) + transported = keep[:, None, None] * xs + move[:, None, None] * prediction # X_{s,t}(x). + with torch.no_grad(): # The teacher is the same network read on the diagonal, detached. + teacher_logits = logits_of(model, transported.detach(), t, t, length, vocab) + target = repaired_teacher(teacher_logits, factor.detach()) + additive = torch.softmax(teacher_logits, -1) - gamma * prediction * centered( + prediction, dz) # The unrepaired target, measured and not trained on. + counter["negative"] += float((additive < 0).to(torch.get_default_dtype()).mean()) + counter["batches"] += 1.0 + student = torch.log_softmax(z, dim=-1) # The student in log space. + return anchor + (target * ((target.clamp_min(1e-9)).log() - student)).sum(-1).mean() + + return loss_fn + + +# %% 07. Sampling +@torch.no_grad() +def sample(model, count, length, vocab, steps, device): + """Walk a uniform grid with the convex form and read the alphabet off the last prediction.""" + x = torch.randn(count, length, vocab, device=device) + grid = torch.linspace(0.0, 1.0, steps + 1) + for a, b in zip(grid[:-1], grid[1:]): + s, t = a.expand(count).to(device), b.expand(count).to(device) + prediction = torch.softmax(logits_of(model, x, s, t, length, vocab), dim=-1) + keep, move = (1.0 - t) / (1.0 - s), (t - s) / (1.0 - s) + x = keep[:, None, None] * x + move[:, None, None] * prediction + return x.argmax(-1) # [N, L]. + + +# %% 08. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=1500, batch_size=128, samples=64, length=16, sample_steps=4) + args = parser.parse_args() + set_seed(args.seed) + device = pick_device(args.device) + tokens, _ = synthetic_dna(num_sequences=1024, length=args.length, seed=args.seed) + data = tokens.to(device) # [1024, L]. + + report = worked_derivative() + report.update(worked_repair()) + if not args.quiet: + print(f"derivative target {report['derivative target']} has " + f"{report['negative entries']} negative entries and sums to " + f"{report['target sum']:.4f}") + print(f"additive {report['additive target']} repairs to {report['repaired target']}") + + counter = {"negative": 0.0, "batches": 0.0} + model = TwoTimeMLP(dim=args.length * DNA_SIZE, width=args.width).to(device) + if not args.quiet: + print(f"\ntraining with the repaired target parameters={count_parameters(model):,}") + losses = train(model, make_loss(model, data, args.length, DNA_SIZE, args.batch_size, device, + counter), steps=args.steps, lr=args.lr, quiet=args.quiet) + report["additive target negative share"] = counter["negative"] / max(counter["batches"], 1.0) + + samples: list[str] = [] + for steps in (1, args.sample_steps): + torch.manual_seed(args.seed + 1) # The same source noise at every step count. + drawn = sample(model, args.samples, args.length, DNA_SIZE, steps, device) + report[f"{steps}-step GC fraction"] = gc_fraction(drawn) + report[f"{steps}-step motif fraction"] = motif_fraction(drawn) + report[f"{steps}-step token entropy"] = token_entropy(drawn, DNA_SIZE) + report[f"{steps}-step unique fraction"] = unique_fraction(drawn) + samples.append(f"{steps} step(s): {decode_dna(drawn[:1])[0]} GC {gc_fraction(drawn):.3f}") + report["training GC fraction"] = gc_fraction(data) + report["final loss"] = sum(losses[-20:]) / len(losses[-20:]) + + print_report("Discrete Flow Maps, Chapter 6", report, samples) + save_run(args.out, model=model, config=vars(args), report=report, losses=losses, samples=samples) + + +if __name__ == "__main__": + main() diff --git a/lecture_6/expanding.py b/lecture_6/expanding.py new file mode 100644 index 0000000000000000000000000000000000000000..e3bb16e3737dcd74bc0076c70cb71d9ec2be528c --- /dev/null +++ b/lecture_6/expanding.py @@ -0,0 +1,217 @@ +"""Expanding flow maps, where the state gains coordinates as time runs. + +The Jacobian of a map between states of different length is rectangular, so it has no determinant +and no inverse, and Liouville's formula, the exact likelihood and the inverse map all stop +applying. Splitting the move into an expansion that fills new coordinates with source noise and a +transport that acts at the longer length restores a square Jacobian. Each coordinate then carries +its own clock, rescaled to the interval it occupies, and the composition identity is still +available, because a finite difference between two maps needs no derivative in time. + +Chapter 6, Section 6.32 of the notes. Trained on the synthetic DNA of +``cis6270.data.synthetic_dna``, grown from half its length to all of it. Prints the local clocks +of the chapter's figure, which read 0.800, 0.692 and 0.333 at a global time of 0.8, the +composition weight 0.125 at (0.3, 0.6, 0.9), and samples from the trained map. + + python lecture_6/expanding.py # 1500 steps, about half a minute + python lecture_6/expanding.py --steps 20 # a few seconds, for a first look +""" + +# %% 01. Imports and configuration +import math + +import torch + +from cis6270.data import DNA_SIZE, decode_dna, synthetic_dna +from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction +from cis6270.nets import TwoTimeMLP, count_parameters +from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train + + +torch.set_num_threads(1) # These networks are small, and splitting one matrix product costs more +# than it saves. ``tests/conftest.py`` sets the same thing for the test session, so the line +# here is the one that acts when a reader runs this script directly. + +BOOK_BIRTHS = (0.0, 0.0, 0.35, 0.7) # The insertion times of the chapter's four-tile figure. + + +# %% 02. Local clocks +# +# A coordinate inserted at t_ins has 1 - t_ins of global time left, so its own clock runs at +# 1/(1 - t_ins) and reaches one exactly at the data end. Before its insertion the clock reads +# zero and the coordinate holds its source value. +def local_time(t, birth): + """t_i = (t - t_ins)/(1 - t_ins), clamped to [0, 1] so an unborn coordinate reads zero.""" + return ((t - birth) / (1.0 - birth).clamp(min=1e-6)).clamp(0.0, 1.0) + + +def worked_clocks(t: float = 0.8): + """The clocks, the speeds and the interpolant shares of the chapter's figure at one time.""" + births = torch.tensor(BOOK_BIRTHS) + clocks = local_time(torch.tensor(t), births) # [4]. + speeds = 1.0 / (1.0 - births) # [4] dt_i/dt. + return { + "insertion times": [round(v, 2) for v in births.tolist()], + "local clocks at t=0.8": [round(v, 3) for v in clocks.tolist()], + "clock speeds": [round(v, 3) for v in speeds.tolist()], + "source share": [round(v, 3) for v in (1.0 - clocks).tolist()], + "last clock from its speed": float(speeds[-1] * (1.0 - t)), + "live coordinates at 0.8": int((births <= t).sum()), + "live coordinates at 0.6": int((births <= 0.6).sum()), + } + + +# %% 03. The composition weight under growth +# +# The convexity weight is a ratio of remaining times and makes no reference to the dimension, so +# the derivation at fixed length goes through once every leg has been lifted to the final length. +# A coordinate born at zero carries the chapter's omega; a later birth carries its own. +def omega(s, u, t): + """omega_{s,u,t} = (u-s)(1-t)/((t-s)(1-u)), the share of the earlier leg.""" + return (u - s) * (1.0 - t) / (((t - s) * (1.0 - u)).clamp(min=1e-9)) + + +def worked_omega(s: float = 0.3, u: float = 0.6, t: float = 0.9): + """The weight at the chapter's triple, for coordinates born at three different times.""" + out = {"omega at (0.3, 0.6, 0.9)": float(omega(*(torch.tensor([v]) for v in (s, u, t))))} + for birth in (0.0, 0.35, 0.7): + b = torch.tensor([birth]) + local = [local_time(torch.tensor([v]), b) for v in (s, u, t)] + out[f"omega for a birth at {birth}"] = float(omega(*local)) + return out + + +# %% 04. A state whose length grows, and where the fresh noise enters +# +# The source values of every coordinate that will ever appear are drawn once, at the final +# length, and a leg that stops before an insertion leaves that coordinate at its source value. +# One draw then serves all three maps of the composition identity. +def lifted_state(source, data_onehot, t, births): + """The state at global time t, each coordinate interpolated on its own clock. [B, L, K].""" + clock = local_time(t[:, None], births[None, :])[:, :, None] # [B, L, 1]. + return (1.0 - clock) * source + clock * data_onehot # [B, L, K]. + + +def expansion_log_density(source, births, s: float, t: float): + """The source log density of the coordinates inserted inside (s, t], the expansion term.""" + inserted = (births > s) & (births <= t) # [L] which coordinates appear on this interval. + per_coordinate = (-0.5 * source ** 2 - 0.5 * math.log(2.0 * math.pi)).sum(-1) # [B, L]. + return float(per_coordinate[:, inserted].sum(-1).mean()), int(inserted.sum()) + + +# %% 05. The model, a two-time denoiser that also reads each coordinate's clock +def predict(model, x, s, t, births, length, vocab): + """psi_{s,t}(x) at every position, with each coordinate's own clock as an extra channel. + + The network has the same width in and out, so every position carries K + 1 channels, the K + coordinates of the state and the local clock. The clock channel of the output is dropped. + """ + clock = local_time(s[:, None], births[None, :])[:, :, None] # [B, L, 1] where each stands. + features = torch.cat([x, clock], dim=-1).flatten(1) # [B, L*(K+1)]. + logits = model(features, s, t).view(-1, length, vocab + 1)[..., :vocab] # [B, L, K]. + return torch.softmax(logits, dim=-1) # [B, L, K] on the simplex at every position. + + +def step_map(x, prediction, s, t, births): + """The convex form, with the weights read on each coordinate's own clock.""" + s_local = local_time(s[:, None], births[None, :])[:, :, None] # [B, L, 1]. + t_local = local_time(t[:, None], births[None, :])[:, :, None] # [B, L, 1]. + keep = (1.0 - t_local) / (1.0 - s_local).clamp(min=1e-6) + return keep * x + (1.0 - keep) * prediction # [B, L, K]. + + +def make_loss(model, data, births, length, vocab, batch_size, device): + eye = torch.eye(vocab, device=device) + + def loss_fn(step: int) -> torch.Tensor: + index = torch.randint(0, len(data), (batch_size,), device=device) + ids = data[index] # [B, L]. + x1 = eye[ids] # [B, L, K] the data, lifted. + source = torch.randn(batch_size, length, vocab, device=device) # Drawn once, at full length. + t_diag = torch.rand(batch_size, device=device) + xt = lifted_state(source, x1, t_diag, births) + diagonal = -(predict(model, xt, t_diag, t_diag, births, length, vocab).clamp_min(1e-9) + .log().gather(-1, ids[..., None])).mean() # Cross-entropy on the diagonal. + + s = torch.rand(batch_size, device=device) + t = s + (1.0 - s) * torch.rand(batch_size, device=device) + u = s + (t - s) * torch.rand(batch_size, device=device) + xs = lifted_state(source, x1, s, births) # The same source draw serves all three legs. + with torch.no_grad(): + first = predict(model, xs, s, u, births, length, vocab) + second = predict(model, step_map(xs, first, s, u, births), u, t, births, length, vocab) + s_l, u_l, t_l = (local_time(v[:, None], births[None, :]) for v in (s, u, t)) + weight = omega(s_l, u_l, t_l)[:, :, None] # [B, L, 1] one weight per coordinate. + target = weight * first + (1.0 - weight) * second # Still on the simplex. + student = predict(model, xs, s, t, births, length, vocab).clamp_min(1e-9) + return diagonal + (target * (target.clamp_min(1e-9) / student).log()).sum(-1).mean() + + return loss_fn + + +# %% 06. Sampling +@torch.no_grad() +def sample(model, count, births, length, vocab, steps, device): + """Walk a uniform grid, with every coordinate stepping on its own clock.""" + x = torch.randn(count, length, vocab, device=device) # The source values, drawn once. + grid = torch.linspace(0.0, 1.0, steps + 1) + for a, b in zip(grid[:-1], grid[1:]): + s, t = a.expand(count).to(device), b.expand(count).to(device) + x = step_map(x, predict(model, x, s, t, births, length, vocab), s, t, births) + return x.argmax(-1) # [N, L]. + + +# %% 07. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=1500, batch_size=128, samples=64, length=16, sample_steps=4) + args = parser.parse_args() + set_seed(args.seed) + device = pick_device(args.device) + tokens, _ = synthetic_dna(num_sequences=1024, length=args.length, seed=args.seed) + data = tokens.to(device) # [1024, L] the full-length sequences. + + # Half the coordinates exist at the source; the rest are inserted at evenly spaced times. + half = args.length // 2 + births = torch.cat([torch.zeros(half, device=device), + torch.linspace(0.1, 0.8, args.length - half, device=device)]) # [L]. + + report = worked_clocks() + report.update(worked_omega()) + report["coordinates at the source"] = int((births == 0).sum()) + report["coordinates at the data"] = args.length + source = torch.randn(256, args.length, DNA_SIZE) + added, count = expansion_log_density(source, births.cpu(), 0.0, 0.5) + report["coordinates inserted in (0, 0.5]"] = count + report["source log density of those"] = added + if not args.quiet: + print(f"local clocks at t = 0.8: {report['local clocks at t=0.8']} with speeds " + f"{report['clock speeds']}") + print(f"omega at (0.3, 0.6, 0.9) = {report['omega at (0.3, 0.6, 0.9)']:.3f}") + + model = TwoTimeMLP(dim=args.length * (DNA_SIZE + 1), width=args.width).to(device) + if not args.quiet: + print(f"\ntraining the expanding map parameters={count_parameters(model):,}") + losses = train(model, make_loss(model, data, births, args.length, DNA_SIZE, args.batch_size, + device), steps=args.steps, lr=args.lr, quiet=args.quiet) + + samples: list[str] = [] + early, late = births == 0, births > 0 # Coordinates born at the source and born later. + for steps in (1, args.sample_steps): + torch.manual_seed(args.seed + 1) + drawn = sample(model, args.samples, births, args.length, DNA_SIZE, steps, device) + report[f"{steps}-step GC fraction"] = gc_fraction(drawn) + report[f"{steps}-step GC, born at the source"] = gc_fraction(drawn[:, early]) + report[f"{steps}-step GC, born later"] = gc_fraction(drawn[:, late]) + report[f"{steps}-step motif fraction"] = motif_fraction(drawn) + report[f"{steps}-step token entropy"] = token_entropy(drawn, DNA_SIZE) + report[f"{steps}-step unique fraction"] = unique_fraction(drawn) + samples.append(f"{steps} step(s): {decode_dna(drawn[:1])[0]} GC {gc_fraction(drawn):.3f}") + report["training GC fraction"] = gc_fraction(data) + report["final loss"] = sum(losses[-20:]) / len(losses[-20:]) + + print_report("Expanding flow maps, Chapter 6", report, samples) + save_run(args.out, model=model, config=vars(args), report=report, losses=losses, samples=samples) + + +if __name__ == "__main__": + main() diff --git a/lecture_6/flow_map.py b/lecture_6/flow_map.py new file mode 100644 index 0000000000000000000000000000000000000000..5b1088828bce658ef631338875ba70f03bdf5fd4 --- /dev/null +++ b/lecture_6/flow_map.py @@ -0,0 +1,206 @@ +"""The flow map of an ordinary differential equation, with no learning anywhere. + +A flow map X_{s,t}(x) returns the state at time t of the trajectory that passes through x at time +s, so one evaluation replaces the N Euler steps of Chapter 2. Here we integrate a velocity we +already know, the Gaussian flow of the chapter, and check against it every property the rest of +the lecture trains for: the three axioms, the inverse, the Jacobian and its differential equation, +Liouville's formula, the transport of a density, and the generator. + +Chapter 6, Sections 6.3 to 6.12 of the notes. Runs on the scalar Gaussian flow and on a batch from +``cis6270.data.gaussian_mixture_2d``. Prints the four numbers of the chapter's listing, 4.100000 +for the map, the cocycle agreeing with it, 0.700000 for the inverse and a Jacobian determinant of +3.000000, together with log 3, the density ratio and the log likelihood. + + python lecture_6/flow_map.py # 200 Runge-Kutta substeps per crossing + python lecture_6/flow_map.py --steps 20 # 20 substeps, a second, and the same report +""" + +# %% 01. Imports and configuration +import math + +import torch +from torch.func import jacrev # Differentiates the whole integration in its starting point. + +from cis6270.data import GAUSSIAN_STD, gaussian_mixture_2d +from cis6270.runner import common_parser, print_report, save_run, set_seed + +torch.set_num_threads(1) # The integration is a scalar loop over a small batch, so one thread is +# the fastest here. ``tests/conftest.py`` sets the same thing for the test session, so the line here +# is the one that acts when a reader runs this script directly. +torch.set_default_dtype(torch.float64) # The axioms are checked below the float32 noise floor. + + +SUBSTEPS = 200 # n, the Runge-Kutta substeps per crossing. No optimizer runs in this file. +START = 0.7 # The state x at which the chapter evaluates every check. + + +# %% 02. The velocity field of the Gaussian flow +# +# Source N(0,1) and data N(2,9), paired by x_1 = 2 + 3 x_0 and joined by a straight line. The +# velocity is a translation at speed 2 plus an expansion whose strength 2/(1+2t) decays as the +# cloud widens. +def b(x, t): # velocity field + return 2.0 + 2.0 * (x - 2.0 * t) / (1.0 + 2.0 * t) # Section 6.9, b(x,t). + + +def exact_map(x, s, t): + """X_{s,t}(x) = 2t + (1+2t)(x - 2s)/(1 + 2s), the closed form the integrator has to reproduce.""" + return 2.0 * t + (1.0 + 2.0 * t) * (x - 2.0 * s) / (1.0 + 2.0 * s) # [B] or scalar. + + +# %% 03. The flow map by fourth-order Runge-Kutta +# +# Four velocity evaluations per substep, at its start, twice at its midpoint and at its end, +# combined with the weights 1, 2, 2, 1 over six. The step h is negative when t is below s, so the +# same function integrates backward with no change. +def flow(x, s, t, n=SUBSTEPS): # X_{s,t}(x) by RK4 + h = (t - s) / n + for k in range(n): + u = s + k * h + k1 = b(x, u) + k2 = b(x + 0.5 * h * k1, u + 0.5 * h) + k3 = b(x + 0.5 * h * k2, u + 0.5 * h) + k4 = b(x + h * k3, u + h) + x = x + (h / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4) + return x + + +# %% 04. The three axioms, checked numerically +# +# Identity, cocycle and tangent characterize the map (Proposition 6.1). The first holds by +# definition, the second by uniqueness of the trajectory, the third by the defining equation. +def check_axioms(x, n): + """Return the residual of each axiom at the state x, each of which the true map drives to zero.""" + identity = flow(x, 0.5, 0.5, n) - x # X_{t,t} = id. + cocycle = flow(flow(x, 0.0, 0.5, n), 0.5, 1.0, n) - flow(x, 0.0, 1.0, n) # X_{t,u} o X_{s,t}. + h = 1e-6 # One-sided difference quotient at the diagonal. + tangent = (flow(x, 0.5, 0.5 + h, n) - x) / h - b(x, 0.5) # d/dt X at t = s equals b(x,s). + return float(identity.abs()), float(cocycle.abs()), float(tangent.abs()) + + +# %% 05. The Jacobian and the equation it satisfies +# +# dJ/dt = grad b(X_{s,t}(x), t) J, starting from J_{s,s} = I. In one dimension grad b is the single +# derivative 2/(1+2t), and the solution from zero is J_{0,t} = 1 + 2t. +def jacobian_by_ode(x, s, t, n): + """Integrate the state and its Jacobian together, with one Euler step each on the same grid.""" + jacobian = torch.ones_like(x) # J_{s,s} = I, the identity axiom differentiated in x. + h = (t - s) / n + for k in range(n): + u = s + k * h + grad_b = 2.0 / (1.0 + 2.0 * u) # d b / d x at time u, the divergence in one dimension. + jacobian = jacobian + h * grad_b * jacobian # Linear in J, with the coefficient frozen. + x = x + h * b(x, u) # The trajectory the coefficient is read along. + return jacobian + + +def liouville_log_det(s, t, n): + """Integrate d/dt log det J = div b, which on this flow is the number log 3 over [0, 1].""" + h = (t - s) / n + grid = s + h * (torch.arange(n) + 0.5) # Midpoints of the substeps. + return float((2.0 / (1.0 + 2.0 * grid)).sum() * h) # Riemann sum of the divergence. + + +# %% 06. Transporting a density along the flow +# +# p_t(X_{s,t}(x)) |det J_{s,t}(x)| = p_s(x). The flow stretches every length by 3 over the unit +# interval, so it divides every density by 3. +def normal_density(x, mean, std): + """The scalar Gaussian density, used at both ends of the transport statement.""" + return math.exp(-((x - mean) ** 2) / (2.0 * std ** 2)) / (std * math.sqrt(2.0 * math.pi)) + + +def mixture_log_density_2d(points, scale, shift): + """Log density of the four-Gaussian mixture after the affine map y = shift + scale * x. + + The pushforward of a Gaussian mixture under an affine map is a Gaussian mixture with the + centers and the widths carried through, which gives the transported density in closed form. + """ + from cis6270.data import GAUSSIAN_CENTERS + + centers = shift + scale * GAUSSIAN_CENTERS.to(points) # [4, 2] transported centers. + std = scale * GAUSSIAN_STD # Each component is widened by the same factor. + squared = ((points[:, None, :] - centers) ** 2).sum(dim=-1) # [N, 4]. + component = -0.5 * squared / std ** 2 - math.log(2.0 * math.pi * std ** 2) # [N, 4]. + return torch.logsumexp(component + math.log(0.25), dim=-1) # [N] mixture log density. + + +# %% 07. The generator of the flow +# +# (L_t f)(x) = b(x,t) . grad f(x), and the rate of change of f along a trajectory is the generator +# applied to f and read at the moving point. +def generator(f, grad_f, x, t): + """(L_t f)(x) for a scalar test function, which in one dimension is b times the derivative.""" + return b(x, t) * grad_f(x) # Section 6.11, the generator. + + +# %% 08. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=SUBSTEPS, samples=512) + args = parser.parse_args() # --steps is the substep count n here, since nothing is trained. + set_seed(args.seed) + n = max(args.steps, 4) # Runge-Kutta needs a handful of substeps to be worth its name. + + # %% 09. The chapter's four lines + x = torch.tensor(START) # The single starting point every check below uses. + one_crossing = flow(x, 0.0, 1.0, n) # 2 + 3x = 4.1. + two_crossings = flow(flow(x, 0.0, 0.5, n), 0.5, 1.0, n) # The cocycle. + there_and_back = flow(flow(x, 0.0, 1.0, n), 1.0, 0.0, n) # The inverse. + det_jacobian = jacrev(lambda z: flow(z, 0.0, 1.0, n))(x) # det J = 3. + if not args.quiet: + print(f"X_0,1({START}) = {float(one_crossing):.6f} closed form " + f"{float(exact_map(x, 0.0, 1.0)):.6f}") + print(f"cocycle through 0.5 = {float(two_crossings):.6f}") + print(f"inverse X_1,0 o X_0,1 = {float(there_and_back):.6f}") + print(f"det J_0,1 = {float(det_jacobian):.6f}") + + identity, cocycle, tangent = check_axioms(x, n) + + # %% 10. The Jacobian two ways, and Liouville's formula + jacobian_ode = jacobian_by_ode(x, 0.0, 1.0, 4000) # One Euler solve of dJ/dt = grad b . J. + log_det = liouville_log_det(0.0, 1.0, 4000) # Integral of the divergence over [0, 1]. + + # %% 11. Transport, in one dimension and in two + image = float(one_crossing) # X_{0,1}(0.7) = 4.1. + p_source = normal_density(START, 0.0, 1.0) # p_0(0.7). + p_data = normal_density(image, 2.0, 3.0) # p_1(4.1). + log_likelihood = math.log(p_source) - log_det # Section 6.10, with the divergence integral. + + points = gaussian_mixture_2d(num_points=args.samples, seed=args.seed) # [N, 2] source cloud. + moved = exact_map(points, 0.0, 1.0) # The same affine flow applied to both coordinates. + log_p0 = mixture_log_density_2d(points, 1.0, 0.0) # [N] density before transport. + log_p1 = mixture_log_density_2d(moved, 3.0, 2.0) # [N] density of the transported mixture. + transport_gap = float((log_p1 + 2.0 * math.log(3.0) - log_p0).abs().max()) # det J = 3^2 = 9. + + # %% 12. The generator along one trajectory + f, grad_f = (lambda z: z ** 2), (lambda z: 2.0 * z) # A test function and its derivative. + h = 1e-5 + along = (f(flow(x, 0.0, 0.5 + h, n)) - f(flow(x, 0.0, 0.5 - h, n))) / (2.0 * h) # d/dt f(X). + generator_gap = abs(float(along) - float(generator(f, grad_f, flow(x, 0.0, 0.5, n), 0.5))) + + report = { + "X_0,1(0.7)": float(one_crossing), + "cocycle through u=0.5": float(two_crossings), + "inverse back to the start": float(there_and_back), + "det J_0,1 by jacrev": float(det_jacobian), + "det J_0,1 by its own ODE": float(jacobian_ode), + "log det J_0,1 by Liouville": log_det, + "identity axiom residual": identity, + "cocycle axiom residual": cocycle, + "tangent axiom residual": tangent, + "p_0(0.7)": p_source, + "p_1(4.1)": p_data, + "density ratio p_0/p_1": p_source / p_data, + "log p_1(4.1) by transport": log_likelihood, + "log p_1(4.1) directly": math.log(p_data), + "2-D transport gap": transport_gap, + "generator gap": generator_gap, + "substeps n": n, + } + print_report("Flow map of the Gaussian flow, Chapter 6", report) + save_run(args.out, config=vars(args), report=report) + + +if __name__ == "__main__": + main() diff --git a/lecture_6/fmlm.py b/lecture_6/fmlm.py new file mode 100644 index 0000000000000000000000000000000000000000..8b495b2836f7cb7138c17e73ba562c1300b6308d --- /dev/null +++ b/lecture_6/fmlm.py @@ -0,0 +1,291 @@ +"""Flow Map Language Models, the composition identity in denoiser coordinates. + +Each token is lifted to a one-hot vector, the source is Gaussian over that space, and the state +moves through R^{L x K}, so no atom has to be split. Stepping by the full remaining time 1 - s in +place of the elapsed time t - s turns the average velocity into a two-time denoiser that lands on +the simplex, the map becomes a convex combination of staying put and jumping to the prediction, +and the cocycle becomes a convex combination of two shorter predictions, which a KL divergence +can measure. + +Chapter 6, Sections 6.25 to 6.27 of the notes. Trained on the synthetic DNA of +``cis6270.data.synthetic_dna``. Prints the convex weights, the composition weight gamma = 0.2 at +(0.2, 0.5, 0.8), the worked KL of 0.030747 nats, the two gradients that decide the loss, the +decoding clock of a 30,522-token vocabulary and the four steps it produces, and samples from the +trained map. + + python lecture_6/fmlm.py # 1500 steps, about half a minute + python lecture_6/fmlm.py --steps 20 # a few seconds, for a first look +""" + +# %% 01. Imports and configuration +import math + +import torch + +from cis6270.data import DNA_SIZE, decode_dna, synthetic_dna +from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction +from cis6270.nets import TwoTimeMLP, count_parameters +from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train + + +torch.set_num_threads(1) # These networks are small, and splitting one matrix product costs more +# than it saves. ``tests/conftest.py`` sets the same thing for the test session, so the line +# here is the one that acts when a reader runs this script directly. + +BOOK_VOCAB = 30522 # The vocabulary the chapter puts on the clock, a standard tokenizer. +ENDPOINT_FRACTION = 0.25 # Share of each batch drawn with t = 1, the slice the sampler reads. + + +# %% 02. The convex form and the composition weight +# +# X_{s,t}(x) = (1-t)/(1-s) x + (t-s)/(1-s) delta_{s,t}(x). Both weights are non-negative for +# s <= t <= 1 and they sum to one, so each step is a convex combination of staying put and jumping +# to the predicted endpoint. At t = s the pair is (1, 0) and the identity axiom holds. +def convex_weights(s, t): + """The pair ((1-t)/(1-s), (t-s)/(1-s)) that the step mixes the state and the prediction with.""" + remaining = (1.0 - s).clamp(min=1e-6) + return (1.0 - t) / remaining, (t - s) / remaining + + +def composition_weight(s, u, t): + """gamma = (1-t)(u-s)/((1-u)(t-s)), the share of the long prediction made over the first leg.""" + return (1.0 - t) * (u - s) / (((1.0 - u) * (t - s)).clamp(min=1e-6)) + + +# %% 03. The cocycle in denoiser coordinates, checked on the scalar flow +# +# delta_{s,t}(x) = gamma delta_{s,u}(x) + (1-gamma) delta_{u,t}(X_{s,u}(x)). On dx/dt = x the +# two-time denoiser is x[1 + (1-s)(e^{t-s}-1)/(t-s)], so every term is available in closed form. +def scalar_denoiser(x, s, t): + """delta_{s,t}(x) on the flow dx/dt = x.""" + return x * (1.0 + (1.0 - s) * math.expm1(t - s) / (t - s)) + + +def scalar_cocycle(x: float = 1.0, s: float = 0.2, u: float = 0.5, t: float = 0.8): + """Every term of the denoiser cocycle at the chapter's triple of times.""" + gamma = (1.0 - t) * (u - s) / ((1.0 - u) * (t - s)) + first = scalar_denoiser(x, s, u) # The prediction over the first leg. + reached = x * math.exp(u - s) # X_{s,u}(x), the state the first leg arrives at. + second = scalar_denoiser(reached, u, t) # The prediction over the second leg. + long = scalar_denoiser(x, s, t) # The prediction over the whole interval. + keep, move = (1.0 - t) / (1.0 - s), (t - s) / (1.0 - s) + return {"gamma": gamma, "delta_s,u": first, "delta_u,t": second, "delta_s,t": long, + "two legs combined": gamma * first + (1.0 - gamma) * second, + "convex weights": (round(keep, 4), round(move, 4)), + "convex form": keep * x + move * long, "X_s,t(x)": x * math.exp(t - s)} + + +# %% 04. One term of the loss, on the DNA alphabet +# +# The composed target is a convex combination of two simplex points, so it is a distribution and +# the divergence against it needs no projection. Its gradient pushes mass onto the token with more +# weight in the two-leg decomposition than in the long prediction. +def worked_kl(): + """The chapter's four-token example, with the target composed from two student predictions.""" + first = torch.tensor([0.10, 0.20, 0.60, 0.10]) # delta_{0.2,0.5} over (A, C, G, T). + second = torch.tensor([0.05, 0.10, 0.80, 0.05]) # delta_{0.5,0.8} at the state the leg reaches. + gamma = (1.0 - 0.8) * (0.5 - 0.2) / ((1.0 - 0.5) * (0.8 - 0.2)) + target = gamma * first + (1.0 - gamma) * second # [4] the composed target. + student = torch.tensor([0.10, 0.15, 0.65, 0.10]) # The student's long prediction. + terms = target * (target / student).log() # [4] the four terms of the divergence. + return {"worked gamma": gamma, "composed target": [round(v, 4) for v in target.tolist()], + "target mass": float(target.sum()), "worked KL": float(terms.sum()), + "the four terms": [round(v, 6) for v in terms.tolist()]} + + +# %% 05. The two gradients a softmax head sees +# +# With a softmax head the KL gradient in the logits is q - p, so the size of the correction in +# coordinate i is the error in coordinate i. The squared error passes through the softmax Jacobian +# q_j(delta_ij - q_i) and carries a factor q_i in every term, so a token the model has learned to +# rule out gets almost no correction. Both gradients are taken by automatic differentiation here, +# which checks the two expressions of Section 6.27 as well as evaluating them. +def logit_gradients(target, student): + """The gradient of KL(p || q) and of half the squared error in the logits of q. Two [K].""" + for_kl = student.log().clone().requires_grad_(True) # Logits whose softmax is the student. + divergence = (target * (target / torch.softmax(for_kl, -1)).log()).sum() + kl_gradient, = torch.autograd.grad(divergence, for_kl) # [K] equal to q - p. + for_squared = student.log().clone().requires_grad_(True) # A second graph on the same logits. + squared = 0.5 * ((torch.softmax(for_squared, -1) - target) ** 2).sum() + squared_gradient, = torch.autograd.grad(squared, for_squared) # [K] with a q_i throughout. + return kl_gradient, squared_gradient + + +def gradient_comparison(floor: float = 1e-6): + """The two gradients at the chapter's student, and again once that student has saturated.""" + target = torch.tensor([0.06, 0.12, 0.76, 0.06]) # p, the composed target of Section 6.27. + student = torch.tensor([0.10, 0.15, 0.65, 0.10]) # q, the student's long prediction. + kl, squared = logit_gradients(target, student) + ratio = (squared / kl).abs() # Between a quarter and a half, coordinate by coordinate. + rest = student[1:] * (1.0 - floor) / student[1:].sum() # The other three keep their shares. + saturated = torch.cat([torch.full_like(student[:1], floor), rest]) # [K], summing to one. + kl_low, squared_low = logit_gradients(target, saturated) + return { + "KL gradient": [round(v, 5) for v in kl.tolist()], + "squared-error gradient": [round(v, 5) for v in squared.tolist()], + "squared over KL": [round(v, 4) for v in ratio.tolist()], + "saturated student": [round(v, 6) for v in saturated.tolist()], + "saturated KL gradient": float(kl_low[0]), + # The surviving squared-error gradient is of order 1e-08, which four printed decimals + # cannot show, so the report carries it in units of 1e-08. + "saturated squared-error gradient, in 1e-08": 1e8 * float(squared_low[0]), + "orders of magnitude between them": math.log10( + abs(float(kl_low[0])) / abs(float(squared_low[0]))), + } + + +# %% 06. The decoding clock +# +# A position decodes correctly when its data coordinate is the largest of the |V| coordinates, +# a shifted standard normal against the maximum of the rest. The reparameterized time tau is the +# fraction of the initial decoding error already removed, and it is the clock the sampler uses. +def decoding_error(t: float, vocab_size: int, nodes: int = 4001) -> float: + """P_e(t), the probability that one position decodes to the wrong token at time t.""" + if t >= 1.0: + return 0.0 + shift = t / (1.0 - t) # a = t/(1-t), the gap the data coordinate opens. + z = torch.linspace(-12.0, 14.0, nodes, dtype=torch.float64) # Quadrature grid over the draw. + density = torch.exp(-0.5 * z * z) / math.sqrt(2.0 * math.pi) # phi(z). + cumulative = 0.5 * torch.erfc(-(z + shift) / math.sqrt(2.0)) # Phi(z + a). + log_all = (vocab_size - 1) * cumulative.clamp(min=1e-300).log() # All the others fall short. + return 1.0 - float(torch.trapz(density * log_all.exp(), z)) + + +def clock(vocab_size: int): + """tau(t) = (P_e(0) - P_e(t))/P_e(0), as a callable, with the source error rate beside it.""" + start = decoding_error(0.0, vocab_size) + return (lambda t: (start - decoding_error(t, vocab_size)) / start), start + + +def even_in_tau(vocab_size: int, pieces: int = 4): + """Solve tau(t) = k/pieces by bisection and return the grid with its two endpoints.""" + tau, _ = clock(vocab_size) + grid = [0.0] + for k in range(1, pieces): + low, high = 0.0, 1.0 - 1e-6 + for _ in range(60): # Bisection; tau is increasing in t. + middle = 0.5 * (low + high) + low, high = (middle, high) if tau(middle) < k / pieces else (low, middle) + grid.append(0.5 * (low + high)) + return grid + [1.0] + + +# %% 07. The model, a two-time denoiser with a softmax head +def predict(model, x, s, t, length, vocab): + """delta_{s,t}(x) at every position, a point of the simplex produced by a softmax head.""" + logits = model(x.flatten(1), s, t).view(-1, length, vocab) # [B, L, K] the network output. + return torch.softmax(logits, dim=-1) # [B, L, K] on the simplex at every position. + + +def step_map(x, prediction, s, t): + """The convex form, which is the flow map written in denoiser coordinates.""" + keep, move = convex_weights(s, t) + return keep[:, None, None] * x + move[:, None, None] * prediction # [B, L, K]. + + +def make_loss(model, data, length, vocab, batch_size, device): + eye = torch.eye(vocab, device=device) # The one-hot lift of the alphabet. + + def loss_fn(step: int) -> torch.Tensor: + index = torch.randint(0, len(data), (batch_size,), device=device) + x1 = eye[data[index]] # [B, L, K] the data, lifted. + x0 = torch.randn_like(x1) # [B, L, K] the Gaussian source. + t_diag = torch.rand(batch_size, device=device) # [B] a time for the diagonal term. + xt = (1 - t_diag[:, None, None]) * x0 + t_diag[:, None, None] * x1 # [B, L, K]. + diagonal = -(predict(model, xt, t_diag, t_diag, length, vocab).clamp_min(1e-9).log() + .gather(-1, data[index][..., None])).mean() # Cross-entropy to the token. + + s = torch.rand(batch_size, device=device) # [B] start time. + t = s + (1.0 - s) * torch.rand(batch_size, device=device) # [B] end time. + # The slice t = 1 is the one-step case the sampler reads, and a continuous draw never + # lands on it, so a fixed share of the batch is set there on purpose (Section 6.26). + endpoint = torch.rand(batch_size, device=device) < ENDPOINT_FRACTION # [B] which rows. + t = torch.where(endpoint, torch.ones_like(t), t) + u = s + (t - s) * torch.rand(batch_size, device=device) # [B] the split point. + xs = (1 - s[:, None, None]) * x0 + s[:, None, None] * x1 # [B, L, K] the interpolant. + with torch.no_grad(): # Both legs and the intermediate state are frozen targets. + first = predict(model, xs, s, u, length, vocab) + second = predict(model, step_map(xs, first, s, u), u, t, length, vocab) + gamma = composition_weight(s, u, t)[:, None, None] + target = gamma * first + (1.0 - gamma) * second # [B, L, K], still on the simplex. + student = predict(model, xs, s, t, length, vocab).clamp_min(1e-9) + off_diagonal = (target * (target.clamp_min(1e-9) / student).log()).sum(-1).mean() + return diagonal + off_diagonal + + return loss_fn + + +# %% 08. Sampling on a grid of times +@torch.no_grad() +def sample(model, count, length, vocab, grid, device): + """Walk the grid with the convex form, and read the alphabet off the last prediction.""" + x = torch.randn(count, length, vocab, device=device) # [N, L, K] the Gaussian source. + for a, b in zip(grid[:-1], grid[1:]): + s = torch.full((count,), a, device=device) + t = torch.full((count,), b, device=device) + x = step_map(x, predict(model, x, s, t, length, vocab), s, t) + return x.argmax(-1) # [N, L] token ids, read by an argmax at every position. + + +# %% 09. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=1500, batch_size=128, samples=64, length=16) + args = parser.parse_args() + set_seed(args.seed) + device = pick_device(args.device) + tokens, _ = synthetic_dna(num_sequences=1024, length=args.length, seed=args.seed) + data = tokens.to(device) # [1024, L] the training sequences. + + report = scalar_cocycle() + report.update(worked_kl()) + report.update(gradient_comparison()) + if not args.quiet: + print(f"at the saturated student the KL gradient is " + f"{report['saturated KL gradient']:.4f} and the squared-error gradient " + f"{1e-8 * report['saturated squared-error gradient, in 1e-08']:.3e}") + book_tau, book_start = clock(BOOK_VOCAB) + for t in (0.60, 0.75, 0.80, 0.85): + report[f"tau({t:.2f}) at |V|=30522"] = book_tau(t) + book_grid = even_in_tau(BOOK_VOCAB) + report["four steps on the clock"] = [round(t, 6) for t in book_grid] + rounded = [round(t, 3) for t in book_grid] # The chapter quotes the grid to three decimals. + report["convex weights on that grid"] = [ + (round((1 - b) / (1 - a), 4), round((b - a) / (1 - a), 4)) + for a, b in zip(rounded[:-1], rounded[1:])] + report["gamma on the last three"] = float(composition_weight( + *(torch.tensor([v]) for v in rounded[1:4]))) + report["gamma over the long first leg"] = float(composition_weight( + *(torch.tensor([v]) for v in (0.0, rounded[1], rounded[2])))) + report["tau on a uniform grid"] = [round(book_tau(t), 5) for t in (0.25, 0.5, 0.75)] + if not args.quiet: + print(f"clock at |V|=30522: P_e(0) = {book_start:.6f}, four steps at " + f"{[round(t, 3) for t in book_grid]}") + + model = TwoTimeMLP(dim=args.length * DNA_SIZE, width=args.width).to(device) + if not args.quiet: + print(f"\ntraining the two-time denoiser parameters={count_parameters(model):,}") + losses = train(model, make_loss(model, data, args.length, DNA_SIZE, args.batch_size, device), + steps=args.steps, lr=args.lr, quiet=args.quiet) + + dna_grid = even_in_tau(DNA_SIZE) # The same clock, read at the four-letter alphabet. + report["DNA clock grid"] = [round(t, 4) for t in dna_grid] + samples: list[str] = [] + for name, grid in (("one step", [0.0, 1.0]), ("four, uniform", [0.0, 0.25, 0.5, 0.75, 1.0]), + ("four, on the clock", dna_grid)): + torch.manual_seed(args.seed + 1) # The same source noise for every schedule. + drawn = sample(model, args.samples, args.length, DNA_SIZE, grid, device) + report[f"{name}: GC fraction"] = gc_fraction(drawn) + report[f"{name}: motif fraction"] = motif_fraction(drawn) + report[f"{name}: token entropy"] = token_entropy(drawn, DNA_SIZE) + report[f"{name}: unique fraction"] = unique_fraction(drawn) + samples.append(f"{name:<20s} {decode_dna(drawn[:1])[0]} GC {gc_fraction(drawn):.3f}") + report["training GC fraction"] = gc_fraction(data) + report["final loss"] = sum(losses[-20:]) / len(losses[-20:]) + + print_report("Flow Map Language Models, Chapter 6", report, samples) + save_run(args.out, model=model, config=vars(args), report=report, losses=losses, samples=samples) + + +if __name__ == "__main__": + main() diff --git a/lecture_6/meanflow.py b/lecture_6/meanflow.py new file mode 100644 index 0000000000000000000000000000000000000000..56a0468d35ad0a59d534618a57a36c5854961217 --- /dev/null +++ b/lecture_6/meanflow.py @@ -0,0 +1,175 @@ +"""MeanFlow, the Eulerian identity written in the average-velocity parameterization. + +Substituting X_{s,t}(x) = x + (t-s) v_{s,t}(x) into the Eulerian identity leaves one equation, +v_{s,t}(x) = b(x,s) + (t-s) d/ds v_{s,t}(x), whose correction term is a total derivative along the +trajectory. One forward-mode pass along the tangent (b, 1, 0) in the arguments (x, s, t) computes +that derivative without forming a Jacobian, and the target uses the realized velocity x_1 - x_0 of +the interpolant in place of b, so a step costs one forward pass and one derivative. + +Chapter 6, Sections 6.13, 6.15 and 6.18 of the notes. Trained on the four Gaussian blobs of +``cis6270.data.gaussian_mixture_2d``. Prints the identity worked by hand on dx/dt = x, where the +average velocity is 1.718x, the instantaneous velocity is x and the correction is 0.718x, and then +one-step, two-step and four-step samples from the trained map. + + python lecture_6/meanflow.py # 2000 steps, about twenty seconds + python lecture_6/meanflow.py --steps 20 # a few seconds, for a first look +""" + +# %% 01. Imports and configuration +import math + +import torch + +from cis6270.data import gaussian_mixture_2d +from cis6270.metrics import mean_log_density, mode_report +from cis6270.nets import TwoTimeMLP, count_parameters +from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train + + +torch.set_num_threads(1) # These networks are small, and splitting one matrix product costs more +# than it saves. ``tests/conftest.py`` sets the same thing for the test session, so the line +# here is the one that acts when a reader runs this script directly. + +DIAGONAL_FRACTION = 0.25 # Share of each batch drawn with t = s, where the identity is flow matching. + + +# %% 02. The average velocity of the scalar flow, in closed form +# +# The exact map of dx/dt = x multiplies the state by e^{t-s}, so the displacement over the elapsed +# time is x (e^{t-s} - 1)/(t-s). On the diagonal the limit is x, which is b(x,t). +def scalar_average_velocity(x, s, t): + """v_{s,t}(x) = x (e^{t-s} - 1)/(t - s) on the flow dx/dt = x.""" + elapsed = t - s + return x * torch.expm1(elapsed) / elapsed # expm1 keeps the limit accurate near the diagonal. + + +# %% 03. The identity, checked on the scalar flow by hand +# +# v_{s,t}(x) = b(x,s) + (t-s) d/ds v_{s,t}(x), with d/ds = partial_s + b(x,s) . grad_x. At s = 0, +# t = 1 the partial derivative is -x, the spatial gradient is e - 1, and the transport term is +# therefore (e-1)x, so the correction is x(e-2) = 0.718x. +def scalar_check(x_value: float = 1.0): + """Return every term of the identity at s = 0, t = 1, each one computed and not quoted.""" + x = torch.tensor(x_value) + s, t = torch.tensor(0.0), torch.tensor(1.0) + average = scalar_average_velocity(x, s, t) # v_{0,1}(x) = x(e-1). + instant = x # b(x, 0) = x on this flow. + partial_s = torch.func.jvp( # Derivative in the start time alone, with x and t held fixed. + lambda ss: scalar_average_velocity(x, ss, t), (s,), (torch.ones_like(s),))[1] + gradient_x = torch.func.jvp( # Derivative in the state alone, the other half of d/ds. + lambda xx: scalar_average_velocity(xx, s, t), (x,), (torch.ones_like(x),))[1] + transport = instant * gradient_x # b(x,s) . grad_x v, the second piece of the total derivative. + total = partial_s + transport # d/ds v_{s,t}(x) = x(e - 2). + return { + "v_0,1(1)": float(average), + "b(1, 0)": float(instant), + "partial_s v": float(partial_s), + "grad_x v_0,1": float(gradient_x), + "d/ds v_0,1": float(total), + "identity gap": float((average - (instant + (t - s) * total)).abs()), + "one map evaluation": float(x + (t - s) * average), + "one Euler step": float(x + (t - s) * instant), + "e": math.e, + } + + +# %% 04. The Gaussian flow, where the correction vanishes +# +# A linear interpolant between paired Gaussians produces straight trajectories, and on a straight +# trajectory the average velocity equals the initial velocity, so the total derivative is zero. +def gaussian_check(x_value: float = 0.7): + """v_{0,t}(x) = 2 + 2x for every end time, and b(x,0) is the same number.""" + x = torch.tensor(x_value) + average = torch.stack([ # Three end times, all giving the same average velocity. + (2.0 * t + (1.0 + 2.0 * t) * x - x) / t for t in (0.25, 0.5, 1.0)]) + instant = 2.0 + 2.0 * (x - 0.0) / 1.0 # b(x, 0) on the Gaussian flow. + return {"v_0,t(0.7) spread": float(average.max() - average.min()), + "v_0,1(0.7)": float(average[-1]), "b(0.7, 0)": float(instant)} + + +# %% 05. The MeanFlow target, by one forward-mode pass +# +# The tangent (x_1 - x_0, 1, 0) in the arguments (x, s, t) contracts the full Jacobian of v into +# the total derivative along the trajectory: the realized velocity in the state slot, one in the +# start-time slot, and zero in the end-time slot because the end time is held fixed. +def mean_flow_target(model, xs, s, t, realized): + """b(x,s) + (t-s) d/ds v_{s,t}(x), detached, which is the right-hand side of the identity.""" + f = lambda xx, ss: model(xx, ss, t) + _, total = torch.func.jvp(f, (xs, s), (realized, torch.ones_like(s))) # [B, d] d/ds v. + return (realized + (t - s)[:, None] * total).detach() # [B, d] the target, held fixed. + + +def make_loss(model, data, batch_size, device): + """Draw a coupling, draw times from the triangle, and regress v onto its own identity.""" + + def loss_fn(step: int) -> torch.Tensor: + index = torch.randint(0, len(data), (batch_size,), device=device) + x1 = data[index] # [B, 2] the data draw. + x0 = torch.randn_like(x1) # [B, 2] the source draw. + s = torch.rand(batch_size, device=device) # [B] start time. + t = s + (1.0 - s) * torch.rand(batch_size, device=device) # [B] end time, s <= t <= 1. + diagonal = torch.rand(batch_size, device=device) < DIAGONAL_FRACTION # [B] which rows anchor. + t = torch.where(diagonal, s, t) # On those rows the correction term has a factor of zero. + xs = (1 - s[:, None]) * x0 + s[:, None] * x1 # [B, 2] the interpolant I_s. + realized = x1 - x0 # [B, 2] the derivative of the interpolant, available in closed form. + target = mean_flow_target(model, xs, s, t, realized) + return ((model(xs, s, t) - target) ** 2).mean() + + return loss_fn + + +# %% 06. Sampling, at any step count the reader chooses +@torch.no_grad() +def sample(model, x0, steps: int): + """Cross [0, 1] in equal pieces, one evaluation of the map per piece.""" + x = x0 + grid = torch.linspace(0.0, 1.0, steps + 1, device=x0.device) # [steps+1] the time grid. + for s, t in zip(grid[:-1], grid[1:]): + x = model.flow_map(x, s.expand(len(x)), t.expand(len(x))) # X_{s,t}(x). + return x # [N, 2] generated points. + + +# %% 07. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=2000, batch_size=128, samples=512) + args = parser.parse_args() + set_seed(args.seed) + device = pick_device(args.device) + data = gaussian_mixture_2d(num_points=2048, seed=args.seed).to(device) # [2048, 2]. + + report = scalar_check() + report.update(gaussian_check()) + if not args.quiet: + print(f"on dx/dt = x at s=0, t=1: v = {report['v_0,1(1)']:.6f} = b {report['b(1, 0)']:.6f}" + f" + correction {report['d/ds v_0,1']:.6f}") + print(f"one map evaluation {report['one map evaluation']:.6f} against one Euler step " + f"{report['one Euler step']:.6f}, exact e = {math.e:.6f}") + + model = TwoTimeMLP(dim=2, width=args.width).to(device) + if not args.quiet: + print(f"\ntraining MeanFlow parameters={count_parameters(model):,}") + losses = train(model, make_loss(model, data, args.batch_size, device), + steps=args.steps, lr=args.lr, quiet=args.quiet) + + torch.manual_seed(args.seed + 1) + x0 = torch.randn(args.samples, 2, device=device) # [N, 2] the source points every run shares. + samples: list[str] = [] + for steps in (1, 2, 4): + points = sample(model, x0, steps) + stats = mode_report(points) + report[f"{steps}-step nearest distance"] = stats["nearest_distance"] + report[f"{steps}-step mode entropy"] = stats["mode_entropy"] + report[f"{steps}-step log density"] = mean_log_density(points.cpu()) + samples.append(f"{steps} step(s): nearest {stats['nearest_distance']:.3f} " + f"modes {stats['modes_found']} entropy {stats['mode_entropy']:.3f}") + one, four = sample(model, x0, 1), sample(model, x0, 4) + report["one step against four"] = float((one - four).norm(dim=1).mean()) + report["final loss"] = sum(losses[-20:]) / len(losses[-20:]) + + print_report("MeanFlow, Chapter 6", report, samples) + save_run(args.out, model=model, config=vars(args), report=report, losses=losses, samples=samples) + + +if __name__ == "__main__": + main() diff --git a/lecture_6/posterior_map.py b/lecture_6/posterior_map.py new file mode 100644 index 0000000000000000000000000000000000000000..6f649ba372d2042d5ad436b5a1a5f06f15955c8a --- /dev/null +++ b/lecture_6/posterior_map.py @@ -0,0 +1,221 @@ +"""Diamond Maps and Meta Flow Maps, the posterior over endpoints that guidance needs. + +Steering a flow toward a reward needs the value function, the log expected exponential reward over +the terminal states reachable from the current one, and that average runs over a posterior. A +deterministic map returns one endpoint, so the only estimate it supports is the reward at that +endpoint, which is low by the Jensen gap. Giving the map a fresh noise argument and indexing it by +the state it is conditioned on turns it into a family of inner flow maps, one per conditioning +pair, whose draws are posterior samples. + +Chapter 6, Sections 6.33 and 6.34 of the notes. Runs on a two-point reward problem and on the four +Gaussian blobs of ``cis6270.data.gaussian_mixture_2d``. Prints the Jensen gap of 0.434, the bias +of the K-sample value at K = 1, 2, 4 and 8, the softmax-weighted value gradient of -0.4483, the +importance weight that pulls a renoised proposal from 0.420 back to 0.0131, and the stop-gradient +fixed point 0.625 of Meta Flow Maps. + + python lecture_6/posterior_map.py # 16 inner samples per state + python lecture_6/posterior_map.py --steps 20 # fewer inner substeps, for a first look +""" + +# %% 01. Imports and configuration +import math + +import torch + +from cis6270.data import gaussian_mixture_2d +from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed + +torch.set_num_threads(1) # The work here is small, and one thread keeps the timing stable. +# ``tests/conftest.py`` sets the same thing for the test session, so the line here is the one +# that acts when a reader runs this script directly. +torch.set_default_dtype(torch.float64) # Exact posteriors, compared against sample averages. + + +INNER_STEPS = 48 # Substeps of the inner flow, which carries noise to a posterior sample. + + +# %% 02. The tilt, the value function and the Jensen gap +# +# The tilted target is the data times the exponential reward, and the same reweighting acts on the +# posterior over endpoints at every intermediate state. A deterministic map returns the posterior +# mean, so its estimate of the value is the reward there, and the exponential is convex, so that +# estimate is always the low one. +def two_point(reward_gap: float = 2.0, x: float = 0.7, t: float = 0.8): + """The chapter's two-point example, with z in {0, 2}, r(z) = z and a linear schedule.""" + weights = torch.tensor([1.0, math.exp(reward_gap)]) # e^{r(0)} and e^{r(2)}. + tilted = weights / weights.sum() # The tilted target over the two points. + far = math.exp(-((x - t * 2.0) ** 2) / (2.0 * (1.0 - t) ** 2)) # Unnormalized, at z = 2. + near = math.exp(-(x ** 2) / (2.0 * (1.0 - t) ** 2)) # Unnormalized, at z = 0. + posterior = torch.tensor([near, far]) / (near + far) # The posterior over endpoints at x. + value = math.log(float((posterior * weights).sum())) # V_t(x) at that state. + tilted_posterior = posterior * weights / float((posterior * weights).sum()) + exact = math.log(0.5 + 0.5 * math.exp(reward_gap)) # V over the data itself. + return { + "tilted target": [round(v, 4) for v in tilted.tolist()], + "posterior at x=0.7": [round(v, 4) for v in posterior.tolist()], + "tilted posterior": [round(v, 4) for v in tilted_posterior.tolist()], + "V_t(0.7)": value, + "V over the data": exact, + "single-endpoint estimate": 1.0 * reward_gap / 2.0, # r at the posterior mean z = 1. + "Jensen gap": exact - reward_gap / 2.0, + } + + +def value_bias(reward_gap: float = 2.0): + """E[V_K] over K independent posterior draws, which approaches the exact value like 1/K.""" + exact = math.log(0.5 + 0.5 * math.exp(reward_gap)) + out = {} + for k in (1, 2, 4, 8): + expectation = sum( + math.comb(k, j) * 2.0 ** (-k) + * math.log(((k - j) + j * math.exp(reward_gap)) / k) for j in range(k + 1)) + out[f"E[V_{k}]"] = expectation + out[f"bias at K={k}"] = exact - expectation + return out + + +# %% 03. The value gradient from K samples +# +# Differentiating the logarithm of an average of exponentials produces a softmax over the rewards, +# so a sample whose reward is far above the others dominates the direction and a set with equal +# rewards gives their plain average. +def value_gradient(rewards, gradients): + """The softmax weights over the rewards, and the weighted average of the reward gradients.""" + weights = torch.softmax(rewards, dim=0) # [K]. + return weights, float((weights * gradients).sum()), float(gradients.mean()) + + +# %% 04. The inner clock +# +# For a fixed conditioning pair the map is an ordinary two-time flow map on an inner time axis +# running from the noise draw at 0 to a posterior sample at 1. Its velocity is the flow matching +# velocity of the posterior, which is available in closed form over the training points. +def posterior_weights(x, t, data): + """p_{1|t}(z_i | x), a softmax over the training points under the linear interpolant. [B, N].""" + spread = max(1.0 - t, 1e-6) + squared = ((x[:, None, :] - t * data[None]) ** 2).sum(-1) # [B, N]. + return torch.softmax(-squared / (2.0 * spread ** 2), dim=1) # [B, N]. + + +def inner_map(noise, weights, data, substeps: int): + """Y_{0,1}(eps | x, t), the inner flow from a Gaussian draw to a posterior sample. [B, K, 2].""" + y = noise # [B, K, 2] one inner state per sample. + step = 1.0 / substeps + for i in range(substeps): + a = i * step # The inner time. + squared = ((y[:, :, None, :] - a * data[None, None]) ** 2).sum(-1) # [B, K, N]. + inner = torch.softmax(torch.log(weights[:, None, :].clamp_min(1e-30)) + - squared / (2.0 * max(1.0 - a, 1e-6) ** 2), dim=-1) # [B, K, N]. + denoiser = inner @ data # [B, K, 2] the inner posterior mean. + y = y + step * (denoiser - y) / max(1.0 - a, 1e-6) # Euler on the inner velocity. + return y + + +# %% 05. Posterior samples from a map we already have +# +# We can renoise the current state back to an earlier time with fresh noise and push it through +# the existing map, which discards information and widens the law of the endpoint. The correction +# is an importance weight, written as a shift of the reward by the log likelihood of the current +# state given that endpoint. +def renoise_example(x: float = 0.7, t: float = 0.8, earlier: float = 0.6): + """The chapter's renoised proposal, and the importance weight that recovers the posterior.""" + scale = earlier / t # alpha_{t'}/alpha_t on the linear schedule. + variance = (1.0 - earlier) ** 2 - scale ** 2 * (1.0 - t) ** 2 # The variance added back. + mean = scale * x # The center of the proposal. + boundary = (earlier * 2.0 / 2.0 - mean) / math.sqrt(variance) # Where the map switches mode. + proposal_far = 0.5 * math.erfc(boundary / math.sqrt(2.0)) # Share of draws reaching z = 2. + penalty_far = (x - t * 2.0) ** 2 / (2.0 * (1.0 - t) ** 2) # Log likelihood of x given z = 2. + penalty_near = x ** 2 / (2.0 * (1.0 - t) ** 2) # The same for z = 0. + far = proposal_far * math.exp(-penalty_far) + near = (1.0 - proposal_far) * math.exp(-penalty_near) + return {"renoise scale": scale, "renoise variance": variance, + "renoise standard deviation": math.sqrt(variance), + "proposal weight on the far mode": proposal_far, + "penalty at z=2": penalty_far, "penalty at z=0": penalty_near, + "corrected weight on the far mode": far / (far + near)} + + +# %% 06. Meta Flow Maps, the stop-gradient form +# +# The controlled velocity needs a ratio of two expectations, and regressing onto a per-sample +# ratio is biased at every sample size. Multiplying through by the denominator leaves a condition +# with one stop-gradient whose fixed point is the ratio of the expectations. +def stop_gradient_fixed_point(weights, gradients, diffusion_squared: float = 2.0): + """The fixed point of the loss, the per-sample ratio it is compared against, and the gap.""" + exact = 0.5 * diffusion_squared * float(gradients.mean()) / float(weights.mean()) + naive = float((gradients / weights).mean()) * 0.5 * diffusion_squared + covariance = float(gradients.mean()) - float((gradients / weights).mean() * weights.mean()) + return {"E[w]": float(weights.mean()), "E[G]": float(gradients.mean()), + "stop-gradient fixed point": exact, "per-sample ratio average": naive, + "gap": naive - exact, "covariance of G/w and w": covariance, + "gap from the covariance": -covariance / float(weights.mean())} + + +# %% 07. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=INNER_STEPS, samples=16) + args = parser.parse_args() # --steps is the inner substep count here, since nothing trains. + set_seed(args.seed) + pick_device(args.device) # Every number here is double precision, which not every + # accelerator carries, so the work stays on the CPU whatever --device names. + substeps = max(args.steps, 8) + + report = two_point() + report.update(value_bias()) + weights, weighted, plain = value_gradient(torch.tensor([0.2, 1.5, 0.9]), + torch.tensor([0.4, -1.2, 0.5])) + report["softmax weights"] = [round(v, 4) for v in weights.tolist()] + report["weighted value gradient"] = weighted + report["plain average gradient"] = plain + report.update(renoise_example()) + report.update(stop_gradient_fixed_point(torch.tensor([1.0, 2.0, 5.0]), + torch.tensor([0.4, 1.0, 3.6]))) + if not args.quiet: + print(f"Jensen gap {report['Jensen gap']:.3f} between the exact value " + f"{report['V over the data']:.4f} and the single endpoint " + f"{report['single-endpoint estimate']:.3f}") + print(f"renoised proposal {report['proposal weight on the far mode']:.3f} corrected to " + f"{report['corrected weight on the far mode']:.4f} against the posterior " + f"{report['posterior at x=0.7'][1]:.4f}") + + # %% 08. The same quantities on the four Gaussian blobs + data = gaussian_mixture_2d(num_points=256, seed=args.seed).double() # [N, 2] the endpoints. + outer_t = 0.6 # The outer time the sampler is conditioned at. + torch.manual_seed(args.seed + 1) + index = torch.randint(0, len(data), (32,)) + states = ((1.0 - outer_t) * torch.randn(32, 2) + outer_t * data[index]) # [32, 2] contexts. + states.requires_grad_(True) + exact_weights = posterior_weights(states, outer_t, data) # [32, N] the exact posterior. + reward = data[:, 0] # r(z) = the first coordinate, a reward with a clear direction. + exact_value = torch.log((exact_weights * reward.exp()[None, :]).sum(1)) # [32] V_t(x). + exact_grad = torch.autograd.grad(exact_value.sum(), states, retain_graph=True)[0] # [32, 2]. + + noise = torch.randn(32, args.samples, 2) # [32, K, 2] inner draws, one per posterior sample. + drawn = inner_map(noise, exact_weights, data, substeps) # [32, K, 2] posterior samples. + sampled_reward = drawn[:, :, 0] # [32, K] the reward at each sample. + estimate = torch.logsumexp(sampled_reward, dim=1) - math.log(args.samples) # [32] V hat. + sample_grad = torch.autograd.grad(estimate.sum(), states, retain_graph=True)[0] # [32, 2]. + deterministic = (exact_weights @ reward) # [32] the reward at the posterior mean. + + exact_value, estimate = exact_value.detach(), estimate.detach() # Past the two gradients. + deterministic, drawn = deterministic.detach(), drawn.detach() + report["exact V, averaged"] = float(exact_value.mean()) + report[f"V from {args.samples} samples"] = float(estimate.mean()) + report["single endpoint, averaged"] = float(deterministic.mean()) + report["Jensen gap on the blobs"] = float((exact_value - deterministic).mean()) + report["posterior mean error"] = float( + (drawn.mean(1) - exact_weights.detach() @ data).norm(dim=1).mean()) + report["value gradient error"] = float((sample_grad - exact_grad).norm(dim=1).mean()) + report["exact value gradient norm"] = float(exact_grad.norm(dim=1).mean()) + states = states.detach() + samples = [f"state ({float(x[0]):+.2f}, {float(x[1]):+.2f}) exact V {float(v):.3f} " + f"{args.samples} samples {float(e):.3f} one endpoint {float(d):.3f}" + for x, v, e, d in zip(states[:4], exact_value[:4], estimate[:4], deterministic[:4])] + + print_report("Diamond Maps and Meta Flow Maps, Chapter 6", report, samples) + save_run(args.out, config=vars(args), report=report, samples=samples) + + +if __name__ == "__main__": + main() diff --git a/lecture_6/residuals.py b/lecture_6/residuals.py new file mode 100644 index 0000000000000000000000000000000000000000..2ca5d8e01c0b30adda725a2a94bc20ec3c38d1af --- /dev/null +++ b/lecture_6/residuals.py @@ -0,0 +1,204 @@ +"""The three residuals of a flow map, and the diagonal anchor added to each one. + +The cocycle differentiated in the end time is the Lagrangian identity, differentiated in the start +time it is the Eulerian identity, and left finite it is the composition identity. Each becomes a +loss by squaring the amount the model violates it. None of the three mentions the data, so each is +satisfied by a model that is self-consistent and wrong, and a flow matching loss on the diagonal +supplies the one piece of information about the target. + +Chapter 6, Sections 6.15 to 6.17 and 6.23 of the notes. Trained on the four Gaussian blobs +of ``cis6270.data.gaussian_mixture_2d``. Prints the stay-still residuals of Section 6.17, which are +1, 1 and 0, and then four trained models side by side in one-step sampling. + + python lecture_6/residuals.py # four models, 600 steps each + python lecture_6/residuals.py --steps 20 # a few seconds, for a first look +""" + +# %% 01. Imports and configuration +import torch + +from cis6270.data import GAUSSIAN_STD, gaussian_mixture_2d +from cis6270.metrics import mode_report +from cis6270.nets import TwoTimeMLP, count_parameters +from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train + + +torch.set_num_threads(1) # These networks are small, and splitting one matrix product costs more +# than it saves. ``tests/conftest.py`` sets the same thing for the test session, so the line +# here is the one that acts when a reader runs this script directly. + +ON_MODE = 3.0 * GAUSSIAN_STD # A sample counts as on a mode within three data standard deviations. +ODE_STEPS = 128 # Substeps used to integrate the model's own diagonal velocity accurately. + + +# %% 02. The model, whose identity axiom is free +# +# The network produces the average velocity v_{s,t}(x), and the map is the state plus the elapsed +# time times that velocity. At t = s the second term vanishes whatever the weights are, so no term +# of the loss is spent on the identity axiom. +class Map(torch.nn.Module): + """X_{s,t}(x) = x + (t - s) v_{s,t}(x), with v the two-time network of ``cis6270.nets``.""" + + def __init__(self, dim: int = 2, width: int = 128): + super().__init__() + self.net = TwoTimeMLP(dim=dim, width=width) # Reads x, s and t; returns a vector. + + def v(self, x, s, t): # the network + return self.net(x, s, t) # [B, d] the average velocity over [s, t]. + + def X(self, x, s, t): # identity axiom, free + return self.net.flow_map(x, s, t) # [B, d] = x + (t - s) v_{s,t}(x). + + +# %% 03. The diagonal anchor, which is the tangent axiom +# +# I_t = (1-t) x_0 + t x_1 and its derivative x_1 - x_0 are available in closed form, so the anchor +# costs one network call and no simulation. Its minimizer is E[I_dot | I_t], the true velocity. +def anchor(m, x0, x1, t): # tangent axiom + xt = (1 - t[:, None]) * x0 + t[:, None] * x1 # [B, d] the interpolant at time t. + return ((m.v(xt, t, t) - (x1 - x0)) ** 2).mean() + + +# %% 04. The Lagrangian residual +# +# d/dt X = b at the output. The closure holds the state and the start time fixed, and a tangent of +# one in t returns the map and its end-time derivative from a single forward-mode pass. The target +# is the model's own diagonal velocity, detached, so the diagonal trains only through the anchor. +def lagrangian(m, xs, s, t): # d/dt X = b at the output + f = lambda tt: m.X(xs, s, tt) + X, dXdt = torch.func.jvp(f, (t,), (torch.ones_like(t),)) + return ((dXdt - m.v(X, t, t).detach()) ** 2).mean() + + +# %% 05. The Eulerian residual +# +# d/ds X + b . grad X = 0. Both the state and the start time are open in the closure, and the +# tangent (v, 1) makes the one returned directional derivative the whole left-hand side. The +# identity is already written as zero, so the residual needs no target at all. +def eulerian(m, xs, s, t): # d/ds X + b . grad X = 0 + vs = m.v(xs, s, s).detach() + f = lambda xx, ss: m.X(xx, ss, t) + _, dX = torch.func.jvp(f, (xs, s), (vs, torch.ones_like(s))) + return (dX ** 2).mean() + + +# %% 06. The composition residual +# +# Two short maps against one long one, with no derivative of any kind. The detach covers the whole +# composed target, so the two legs are frozen and only the long map is corrected. +def composition(m, xs, s, u, t): # two short maps, one long one + tgt = m.X(m.X(xs, s, u), u, t).detach() + return ((m.X(xs, s, t) - tgt) ** 2).mean() + + +# %% 07. The stay-still solution each residual has to rule out +# +# The average velocity equals the true b when the two times coincide and zero everywhere else, so +# the map is the identity for every pair of times and the sampler generates the source. The +# identity axiom, the cocycle and the anchor all hold on it, and the two differential residuals +# separate it from the truth by a full unit. +class StayStill(torch.nn.Module): + """v_{s,t}(x) = b(x,t) on the diagonal and 0 off it, on the scalar flow dx/dt = x.""" + + def v(self, x, s, t): + on_diagonal = (s == t)[:, None] # [B, 1] true where the two times meet. + return torch.where(on_diagonal, x, torch.zeros_like(x)) # b(x,t) = x on this flow. + + def X(self, x, s, t): + return x # Off the diagonal (t-s) times zero, on it zero times b, so the map is the identity. + + +def stay_still_residuals(): + """The three residuals of Section 6.17 at x = 1, s = 0, t = 1, which are 1, 1 and 0.""" + m = StayStill() + x = torch.ones(1, 1) # [1, 1] the state x = 1. + s, t, u = torch.zeros(1), torch.ones(1), torch.full((1,), 0.5) + diagonal_gap = float(((m.v(x, t, t) - x) ** 2).mean()) # The anchor, at its minimum. + return { + "stay-still Lagrangian": float(lagrangian(m, x, s, t)), + "stay-still Eulerian": float(eulerian(m, x, s, t)), + "stay-still composition": float(composition(m, x, s, u, t)), + "stay-still diagonal gap": diagonal_gap, + } + + +# %% 08. One training step, the same shape whichever residual we add +def make_loss(model, data, residual, batch_size, device): + """Draw a coupling, draw times from the triangle s <= t, and add the residual to the anchor.""" + + def loss_fn(step: int) -> torch.Tensor: + index = torch.randint(0, len(data), (batch_size,), device=device) + x1 = data[index] # [B, 2] the data draw. + x0 = torch.randn_like(x1) # [B, 2] the independent source draw. + s = torch.rand(batch_size, device=device) # [B] start time. + t = s + (1.0 - s) * torch.rand(batch_size, device=device) # [B] end time, s <= t <= 1. + u = s + (t - s) * torch.rand(batch_size, device=device) # [B] intermediate time. + total = anchor(model, x0, x1, torch.rand(batch_size, device=device)) # The data term. + if residual is None: + return total # The diagonal alone, which is flow matching read in one step. + xs = (1 - s[:, None]) * x0 + s[:, None] * x1 # [B, 2] the interpolant at the start time. + if residual is composition: + return total + residual(model, xs, s, u, t) + return total + residual(model, xs, s, t) + + return loss_fn + + +# %% 09. One-step sampling, and the model's own ordinary differential equation +@torch.no_grad() +def evaluate(model, x0): + """One-step output, the endpoint of the model's own diagonal velocity, and the mode statistics.""" + from cis6270.data import GAUSSIAN_CENTERS + + zeros = torch.zeros(len(x0), device=x0.device) + one_step = model.X(x0, zeros, torch.ones_like(zeros)) # [N, 2] a single evaluation of the map. + x, h = x0.clone(), 1.0 / ODE_STEPS # Integrate dx/dt = v_{t,t}(x) with the same weights. + for k in range(ODE_STEPS): + t = torch.full((len(x0),), k * h, device=x0.device) + x = x + h * model.v(x, t, t) # Euler on the diagonal velocity the anchor trains. + distance = float((one_step - x).norm(dim=1).mean()) # Column two of Table 6.2. + nearest = torch.cdist(one_step, GAUSSIAN_CENTERS.to(one_step)).min(dim=1).values # [N]. + on_mode = float((nearest < ON_MODE).to(torch.get_default_dtype()).mean()) # Column three of Table 6.2. + return one_step, distance, on_mode, mode_report(one_step) + + +# %% 10. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(samples=512) + args = parser.parse_args() + set_seed(args.seed) + device = pick_device(args.device) + data = gaussian_mixture_2d(num_points=2048, seed=args.seed).to(device) # [2048, 2] the target. + + report = stay_still_residuals() # Section 6.17, before any training. + if not args.quiet: + print(f"stay-still at x=1, s=0, t=1: Lagrangian {report['stay-still Lagrangian']:.3f} " + f"Eulerian {report['stay-still Eulerian']:.3f} " + f"composition {report['stay-still composition']:.3f}") + + recipes = [("the diagonal alone", None), ("+ Lagrangian", lagrangian), + ("+ Eulerian", eulerian), ("+ composition", composition)] + samples: list[str] = [] + for name, residual in recipes: + set_seed(args.seed) # Same starting weights, same batches and same noise for all four. + model = Map(dim=2, width=args.width).to(device) + if not args.quiet: + print(f"\ntraining {name} parameters={count_parameters(model):,}") + losses = train(model, make_loss(model, data, residual, args.batch_size, device), + steps=args.steps, lr=args.lr, quiet=args.quiet) + torch.manual_seed(args.seed + 1) # The same source points for every model. + x0 = torch.randn(args.samples, 2, device=device) + _, distance, on_mode, stats = evaluate(model, x0) + report[f"{name}: distance to its own ODE"] = distance + report[f"{name}: fraction on a mode"] = on_mode + report[f"{name}: final loss"] = sum(losses[-20:]) / len(losses[-20:]) + samples.append(f"{name:<20s} distance {distance:.3f} on a mode {on_mode:.2f} " + f"modes found {stats['modes_found']}") + + print_report("Three residuals and the diagonal anchor, Chapter 6", report, samples) + save_run(args.out, config=vars(args), report=report, samples=samples) + + +if __name__ == "__main__": + main() diff --git a/lecture_6/shortcut.py b/lecture_6/shortcut.py new file mode 100644 index 0000000000000000000000000000000000000000..90ce57f56a00271d9c46162dc4163c5cd2cbc610 --- /dev/null +++ b/lecture_6/shortcut.py @@ -0,0 +1,184 @@ +"""Shortcut models, the composition identity indexed by step size. + +The second input of the network is the length of the step and not the time at which it ends, and +the identity is enforced only between a step and the step twice as long. One doubling constraint +suffices because the ladder chains: with a finest step of 1/128 and seven doublings, the one-step +map at level seven rests on seven relations and one regression against the data. + +Chapter 6, Section 6.20 of the notes. Trained on the four Gaussian blobs of +``cis6270.data.gaussian_mixture_2d``. Prints the doubling worked on dx/dt = x, where two steps of +0.25 average to 1.297443 and the direct step of 0.5 returns the same number, then the dyadic ladder +and the trained model's doubling residual at every rung. + + python lecture_6/shortcut.py # 2000 steps, about twenty seconds + python lecture_6/shortcut.py --steps 20 # a few seconds, for a first look +""" + +# %% 01. Imports and configuration +import math + +import torch + +from cis6270.data import gaussian_mixture_2d +from cis6270.metrics import mean_log_density, mode_report +from cis6270.nets import TwoTimeMLP, count_parameters +from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train + + +torch.set_num_threads(1) # These networks are small, and splitting one matrix product costs more +# than it saves. ``tests/conftest.py`` sets the same thing for the test session, so the line +# here is the one that acts when a reader runs this script directly. + +FINEST = 1.0 / 128.0 # d_0, the finest step of the ladder. +LEVELS = 7 # log2(128) doublings, so the step sizes are d_k = 2^k d_0 for k = 0 to 7. +ZERO_FRACTION = 0.75 # Three quarters of each batch is drawn at step size zero. + + +# %% 02. The network in step-size coordinates +# +# The substitution s_theta(x, t, d) = v_{t, t+d}(x) turns the published network into the two-time +# average velocity of this chapter, so the same weights serve both readings. +def shortcut(model, x, t, d): + """s_theta(x, t, d), the average velocity over the step of length d starting at time t.""" + return model(x, t, t + d) # [B, 2] = v_{t, t+d}(x). + + +def take_step(model, x, t, d): + """One step of length d, which is the flow map X_{t, t+d} in these coordinates.""" + return x + d[:, None] * shortcut(model, x, t, d) # [B, 2]. + + +# %% 03. The doubling identity, worked on the scalar flow +# +# s(x, t, 2d) = half s(x, t, d) + half s(x', t+d, d) with x' = x + d s(x, t, d). On dx/dt = x the +# average velocity is x (e^d - 1)/d, and the identity holds at every step length at once. +def scalar_doubling(x_value: float = 1.0, d: float = 0.25): + """Every term of the doubling identity on dx/dt = x, each one computed here.""" + first = x_value * math.expm1(d) / d # s(1, 0, 0.25). + reached = x_value + d * first # x', which is e^{0.25} to six figures. + second = reached * math.expm1(d) / d # s(x', 0.25, 0.25), faster because the state has grown. + doubled = x_value * math.expm1(2.0 * d) / (2.0 * d) # s(1, 0, 0.5), computed directly. + return { + "s(1, 0, d)": first, + "state after the first step": reached, + "s(x', d, d)": second, + "average of the two": 0.5 * (first + second), + "s(1, 0, 2d) direct": doubled, + "doubling gap": abs(0.5 * (first + second) - doubled), + "state after the doubled step": x_value + 2.0 * d * doubled, + "exp(2d)": math.exp(2.0 * d), + } + + +# %% 04. The dyadic ladder +# +# Eight step sizes from 1/128 up to 1, each training pair tying level k to level k+1. A state at +# level zero is already trained against the diagonal by the anchor, so the correctness of the +# one-step map rests on a chain of seven relations and one regression. +def ladder(): + """The step sizes d_k = 2^k d_0, from the finest to the whole interval.""" + return [FINEST * 2.0 ** k for k in range(LEVELS + 1)] # [1/128, ..., 1.0]. + + +# %% 05. The loss +# +# At step size zero the target is the flow matching target of Chapter 2, and above it the batch +# draws a level of the ladder, forms the two short steps and regresses the doubled step onto their +# average, with the target detached. The chapter puts three quarters of the batch on the zero step +# and the rest on the doubling. Both terms are evaluated on the whole batch here and weighted in +# that proportion, which has the same expectation at one network call more per step. +def make_loss(model, data, batch_size, device): + sizes = torch.tensor(ladder()[:LEVELS], device=device) # Levels whose doubling still fits in [0,1]. + + def loss_fn(step: int) -> torch.Tensor: + index = torch.randint(0, len(data), (batch_size,), device=device) + x1 = data[index] # [B, 2] the data draw. + x0 = torch.randn_like(x1) # [B, 2] the source draw. + t = torch.rand(batch_size, device=device) # [B] a time for the zero-step rows. + xt = (1 - t[:, None]) * x0 + t[:, None] * x1 # [B, 2] the interpolant. + zero = torch.zeros(batch_size, device=device) + anchor = ((shortcut(model, xt, t, zero) - (x1 - x0)) ** 2).mean() # Step size zero. + + d = sizes[torch.randint(0, LEVELS, (batch_size,), device=device)] # [B] one rung each. + start = (1.0 - 2.0 * d) * torch.rand(batch_size, device=device) # [B] t with t + 2d <= 1. + xs = (1 - start[:, None]) * x0 + start[:, None] * x1 # [B, 2] the interpolant at t. + with torch.no_grad(): # The two legs are frozen; only the doubled step is corrected. + first = shortcut(model, xs, start, d) + reached = xs + d[:, None] * first + second = shortcut(model, reached, start + d, d) + target = 0.5 * (first + second) + doubled = ((shortcut(model, xs, start, 2.0 * d) - target) ** 2).mean() + return ZERO_FRACTION * anchor + (1.0 - ZERO_FRACTION) * doubled + + return loss_fn + + +# %% 06. Sampling, and the residual at every rung +@torch.no_grad() +def sample(model, x0, steps: int): + """Cross [0, 1] in equal steps of length 1/steps, one network evaluation per step.""" + x, d = x0, torch.full((len(x0),), 1.0 / steps, device=x0.device) + for k in range(steps): + t = torch.full((len(x0),), k / steps, device=x0.device) + x = take_step(model, x, t, d) + return x # [N, 2]. + + +@torch.no_grad() +def rung_residuals(model, x0): + """The squared doubling gap at each level of the ladder, measured on source points.""" + out = {} + for k, d in enumerate(ladder()[:LEVELS]): + step = torch.full((len(x0),), d, device=x0.device) + start = torch.zeros(len(x0), device=x0.device) + first = shortcut(model, x0, start, step) + second = shortcut(model, x0 + step[:, None] * first, start + step, step) + gap = shortcut(model, x0, start, 2.0 * step) - 0.5 * (first + second) + out[f"rung {k}: d = 1/{int(round(1.0 / d))}"] = float((gap ** 2).mean()) + return out + + +# %% 07. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=2000, batch_size=128, samples=512) + args = parser.parse_args() + set_seed(args.seed) + device = pick_device(args.device) + data = gaussian_mixture_2d(num_points=2048, seed=args.seed).to(device) # [2048, 2]. + + report = scalar_doubling() + report["ladder levels"] = LEVELS + report["finest step d_0"] = FINEST + if not args.quiet: + print(f"on dx/dt = x: two steps of 0.25 average to {report['average of the two']:.6f}, " + f"one step of 0.5 returns {report['s(1, 0, 2d) direct']:.6f}") + print(f"ladder {[round(d, 5) for d in ladder()]}") + + model = TwoTimeMLP(dim=2, width=args.width).to(device) + if not args.quiet: + print(f"\ntraining the shortcut model parameters={count_parameters(model):,}") + losses = train(model, make_loss(model, data, args.batch_size, device), + steps=args.steps, lr=args.lr, quiet=args.quiet) + + torch.manual_seed(args.seed + 1) + x0 = torch.randn(args.samples, 2, device=device) # [N, 2] the shared source points. + report.update(rung_residuals(model, x0)) + samples: list[str] = [] + for steps in (1, 2, 4, 8): + points = sample(model, x0, steps) + stats = mode_report(points) + report[f"{steps}-step nearest distance"] = stats["nearest_distance"] + report[f"{steps}-step log density"] = mean_log_density(points.cpu()) + samples.append(f"{steps} step(s): nearest {stats['nearest_distance']:.3f} " + f"modes {stats['modes_found']} entropy {stats['mode_entropy']:.3f}") + report["one step against eight"] = float( + (sample(model, x0, 1) - sample(model, x0, 8)).norm(dim=1).mean()) + report["final loss"] = sum(losses[-20:]) / len(losses[-20:]) + + print_report("Shortcut models, Chapter 6", report, samples) + save_run(args.out, model=model, config=vars(args), report=report, losses=losses, samples=samples) + + +if __name__ == "__main__": + main() diff --git a/lecture_6/stochastic_map.py b/lecture_6/stochastic_map.py new file mode 100644 index 0000000000000000000000000000000000000000..075d2acbae7d5f3616cf4a0fad2b6c8fdd720d86 --- /dev/null +++ b/lecture_6/stochastic_map.py @@ -0,0 +1,218 @@ +"""Strong stochastic flow maps, where the driving path is an argument of the map. + +A map that matches only the transition kernel draws fresh noise on every leg, so one full step and +two half steps follow different trajectories and their composition residual has a floor at twice +the variance of the interval. A strong map takes the path itself, projected onto the first few +shifted Legendre polynomials, and the Chen relations build the coefficients of a long interval +from the coefficients of its halves, so all three legs of the cocycle are driven by one +realization. + +Chapter 6, Section 6.31 of the notes. Runs on the linear equation dX = X dt + sigma dW +and on a cloud of starting states from ``cis6270.data.gaussian_mixture_2d``. Prints the weak +map's floor of 2K_1 = 6.389, the Legendre variances 1, 1/3 and 1/5, the Chen fusion that turns +four half-interval draws into (0.50, -0.20), and the pathwise error against the coefficient count. + + python lecture_6/stochastic_map.py # 400,000 Monte Carlo draws, a few seconds + python lecture_6/stochastic_map.py --steps 20 # the floor of 1,000 draws, for a first look +""" + +# %% 01. Imports and configuration +import math + +import torch + +from cis6270.data import gaussian_mixture_2d +from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed + +torch.set_num_threads(1) # The work here is elementwise, and one thread keeps the timing stable. +# ``tests/conftest.py`` sets the same thing for the test session, so the line here is the one +# that acts when a reader runs this script directly. +torch.set_default_dtype(torch.float64) # The checks below sit well under the float32 noise floor. + + +SIGMA = 1.0 # The noise level of dX = X dt + sigma dW. +ORDER = 3 # Coefficients carried per interval, the number used in practice. +FINEST = 32 # Subintervals the path is drawn on before the Chen relations fuse it upward. + + +# %% 02. The transition kernel, and the weak map built from it +# +# Starting at x at time s, the state at t is e^{t-s} x plus a centered Gaussian whose variance is +# K_{t-s}. A weak map reproduces that law on every interval with one fresh draw. +def kernel_variance(d): + """K_d = sigma^2 (e^{2d} - 1)/2, the variance the equation accumulates over a length d.""" + return 0.5 * SIGMA ** 2 * (math.exp(2.0 * d) - 1.0) + + +def weak_step(x, d, noise): + """The weak map, which takes a state and one fresh standard normal draw.""" + return math.exp(d) * x + math.sqrt(kernel_variance(d)) * noise + + +# %% 03. The cost of matching only the transition kernel +# +# The variances compose correctly, so nothing is wrong at the level of distributions. At the level +# of samples the three draws are independent, so the squared gap between one full step and two +# half steps averages to twice the full-interval variance however well the map is trained. +def weak_floor(x: float = 0.3, u: float = 0.5, t: float = 1.0, draws: int = 400_000, seed: int = 0): + """The worked comparison of Section 6.31, and a Monte Carlo average of the floor.""" + full = weak_step(x, t, 0.8) # One step with the chapter's third draw. + half = weak_step(weak_step(x, u, 0.5), t - u, -1.2) # Two legs with the first two draws. + carried = math.exp(2.0 * (t - u)) * kernel_variance(u) + kernel_variance(t - u) + matching = (math.exp(t - u) * math.sqrt(kernel_variance(u)) * 0.5 + - math.sqrt(kernel_variance(t - u)) * 1.2) / math.sqrt(kernel_variance(t)) + generator = torch.Generator().manual_seed(seed) + e1, e2, e3 = (torch.randn(draws, generator=generator) for _ in range(3)) + gap = (math.sqrt(kernel_variance(t)) * e3 + - math.exp(t - u) * math.sqrt(kernel_variance(u)) * e1 + - math.sqrt(kernel_variance(t - u)) * e2) # [draws] the difference of the two routes. + return { + "K_1": kernel_variance(t), "sqrt K_1": math.sqrt(kernel_variance(t)), + "sqrt K_0.5": math.sqrt(kernel_variance(u)), + "one full step": full, "two half steps": half, "gap": full - half, + "squared gap": (full - half) ** 2, + "carried variance equals K_1": carried, "floor 2 K_1": 2.0 * kernel_variance(t), + "the draw that would close it": matching, + "Monte Carlo floor": float((gap ** 2).mean()), + } + + +# %% 04. The path as an input, through the shifted Legendre polynomials +# +# I^{(n)}_{s,t} integrates the n-th shifted Legendre polynomial of the normalized time against the +# increment. The zeroth coefficient is the plain increment, the first separates a path that did +# its travelling early from one that did it late, and distinct coefficients are independent. +def legendre(n: int, r): + """The n-th shifted Legendre polynomial on [0, 1].""" + return [torch.ones_like(r), 2.0 * r - 1.0, 6.0 * r * r - 6.0 * r + 1.0][n] + + +def draw_coefficients(shape, length: float, generator): + """I^{(n)} ~ N(0, (t-s)/(2n+1)), drawn directly, with no path simulated.""" + scale = torch.tensor([math.sqrt(length / (2 * n + 1)) for n in range(ORDER)]) # [ORDER]. + return torch.randn(*shape, ORDER, generator=generator) * scale # [..., ORDER]. + + +def kernel_coefficients(length: float, nodes: int = 8001): + """a_n(h) in e^{h(1-r)} = sum_n a_n P_n(r), which turns the Ito integral into a finite sum.""" + r = torch.linspace(0.0, 1.0, nodes) + weight = torch.exp(length * (1.0 - r)) # The kernel e^{t-u} written in normalized time. + return torch.stack([(2 * n + 1) * torch.trapz(weight * legendre(n, r), r) + for n in range(ORDER)]) # [ORDER]. + + +def strong_step(x, length: float, coefficients): + """The solution map over one interval, driven by the coefficients of the path on it.""" + kernel = kernel_coefficients(length).to(x) # [ORDER]. + return math.exp(length) * x + SIGMA * (coefficients * kernel).sum(-1) # Same shape as x. + + +# %% 05. The Chen relations +# +# A shifted Legendre polynomial of the long interval expands in the shifted Legendre basis of each +# half, so the coefficients fuse by a fixed linear rule with signs alternating in the parity of +# n + m. The first rule says the increments add, and the minus sign in the second is the oddness +# of P_1 about the midpoint. +def chen(first, second): + """Fuse the coefficients of two equal halves into the coefficients of their union.""" + return torch.stack([ + first[..., 0] + second[..., 0], # The increments add. + -0.5 * (first[..., 0] - second[..., 0]) + 0.5 * (first[..., 1] + second[..., 1]), + 0.25 * (first[..., 2] + second[..., 2]) - 0.75 * (first[..., 1] - second[..., 1]), + ], dim=-1) + + +def fuse_to(pieces, target: int): + """Halve the number of intervals by Chen fusion until ``target`` intervals remain.""" + while pieces.shape[-2] > target: + pieces = chen(pieces[..., 0::2, :], pieces[..., 1::2, :]) # Fuse neighbouring pairs. + return pieces # [..., target, ORDER]. + + +def worked_chen(): + """The chapter's four half-interval draws, fused into the coefficients of the unit interval.""" + halves = torch.tensor([[0.80, 0.20, 0.0], [-0.30, 0.50, 0.0]]) # [2, ORDER], the two legs. + fused = chen(halves[0], halves[1]) + return {"I0 over [0,1]": float(fused[0]), "I1 over [0,1]": float(fused[1]), + "variance of I1 at h=1": 1.0 / 3.0, + "fused variance at h=1": 0.25 * 0.5 + 0.25 * 0.5 + 0.25 / 6.0 + 0.25 / 6.0} + + +# %% 06. The pathwise error against the number of coefficients +# +# A fine simulation gives both the exact Ito integral and the coefficients of the same path, so +# the truncation can be measured on the realization it was taken from. +def truncation_error(paths: int, grid: int, seed: int): + """Mean absolute pathwise error of a one-step map reading 1, 2 and 3 coefficients.""" + generator = torch.Generator().manual_seed(seed) + step = 1.0 / grid + time = torch.arange(grid) * step # [grid] left endpoints of the substeps. + increments = torch.randn(paths, grid, generator=generator) * math.sqrt(step) # [P, grid]. + exact = (torch.exp(1.0 - time)[None, :] * increments).sum(1) # [P] the Ito integral. + measured = torch.stack([(legendre(n, time)[None, :] * increments).sum(1) + for n in range(ORDER)], dim=1) # [P, ORDER]. + kernel = kernel_coefficients(1.0) # [ORDER]. + out = {f"coefficient variance I{n}": float(measured[:, n].var()) for n in range(ORDER)} + for keep in range(1, ORDER + 1): + approximate = (measured[:, :keep] * kernel[:keep]).sum(1) # [P]. + out[f"pathwise error with {keep}"] = float((exact - approximate).abs().mean()) + return out + + +# %% 07. One path, several step counts +# +# A strong map crosses one trajectory of one realization in any number of hops, so its terminal +# state barely moves with the step count. A weak map redraws at every change of the step count, +# so its terminal state moves by the floor of Section 6.31. +def cross(x, coefficients, steps: int, weak: bool, generator=None): + """Cross [0, 1] in ``steps`` equal hops, either along the fused path or with fresh draws.""" + pieces = fuse_to(coefficients, steps) # [..., steps, ORDER] on this grid. + length = 1.0 / steps + for k in range(steps): + if weak: + noise = torch.randn(x.shape, generator=generator) # A fresh draw for every leg. + x = weak_step(x, length, noise) + else: + x = strong_step(x, length, pieces[..., k, :]) + return x + + +# %% 08. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=400_000, samples=512) + args = parser.parse_args() # --steps is the Monte Carlo draw count here, since nothing trains. + set_seed(args.seed) + pick_device(args.device) # Every number here is double precision, which not every + # accelerator carries, so the work stays on the CPU whatever --device names. + + report = weak_floor(draws=max(args.steps, 1000), seed=args.seed) + report.update(worked_chen()) + report.update(truncation_error(paths=1024, grid=4096, seed=args.seed)) + if not args.quiet: + print(f"weak floor 2K_1 = {report['floor 2 K_1']:.3f}, Monte Carlo " + f"{report['Monte Carlo floor']:.3f}") + print(f"pathwise error {report['pathwise error with 1']:.4f} with the increment alone " + f"against {report['pathwise error with 3']:.6f} with three coefficients") + + # %% 09. The same cloud at four step counts + points = gaussian_mixture_2d(num_points=args.samples, seed=args.seed).double() # [N, 2]. + generator = torch.Generator().manual_seed(args.seed + 1) + path = draw_coefficients((args.samples, 2, FINEST), 1.0 / FINEST, generator) # [N, 2, F, 3]. + reference = cross(points, path, FINEST, weak=False) # The finest crossing, as the yardstick. + samples: list[str] = [] + for steps in (2, 4, 8, 16): # Short of FINEST, which would be the yardstick against itself. + strong = cross(points, path, steps, weak=False) # One realization, several hop counts. + weak_generator = torch.Generator().manual_seed(args.seed + 2) + weak = cross(points, path, steps, weak=True, generator=weak_generator) + report[f"{steps} hops: strong spread"] = float((strong - reference).norm(dim=1).mean()) + report[f"{steps} hops: weak spread"] = float((weak - reference).norm(dim=1).mean()) + samples.append(f"{steps:2d} hops strong {float((strong - reference).norm(dim=1).mean()):.6f}" + f" weak {float((weak - reference).norm(dim=1).mean()):.3f}") + + print_report("Strong stochastic flow maps, Chapter 6", report, samples) + save_run(args.out, config=vars(args), report=report, samples=samples) + + +if __name__ == "__main__": + main() diff --git a/lecture_7/README.md b/lecture_7/README.md new file mode 100644 index 0000000000000000000000000000000000000000..995a991c10ef9daf8d993dfcf517a9ef39316acc --- /dev/null +++ b/lecture_7/README.md @@ -0,0 +1,166 @@ +# Lecture 7: Optimal Transport and Schrodinger Bridges + +Chapter 7 of the notes. + +Every model of Chapters 2 through 6 trains on a pair of samples, and in almost all of them the two +were drawn independently. This lecture takes the pairing as the quantity to be chosen. The first +four files price a pairing, relax Monge's map to Kantorovich's coupling, certify the optimum with +prices, smooth the problem with an entropy term that Sinkhorn's algorithm solves by rescaling rows +and columns, and read transport as motion in time, where the minimizer moves every particle in a +straight line. The next five replace the cost matrix by the endpoint law of a reference process, +which turns the same scaling into a Schrodinger bridge: the chain rule puts reference bridges +between every pair of ends, Doob's h-transform makes the bridge a process we can run forward, and +three learning methods differ only in what they simulate and what they regress. Three more carry +the construction to a finite state space, where a rate matrix replaces the drift, where a molecular +graph factorizes into forty-five categorical features, and where discrete time makes the whole path +law a table of sixteen numbers. The last file covers a population that splits into weighted branches +and a system of interacting particles whose own dynamics is the reference. + +Almost every number below is computed in closed form, which is what the chapter does. Two files +train. + +| file | command | what it computes | +| ---- | ------- | ---------------- | +| `optimal_transport.py` | `python lecture_7/optimal_transport.py` | the four points, the two-by-two plan and its prices, the sorted pairing on three points, and the Gaussian map | +| `sinkhorn.py` | `python lecture_7/sinkhorn.py` | the Gibbs kernel, the first sweep by hand, the converged entropic plan, and the plan at four values of epsilon | +| `displacement.py` | `python lecture_7/displacement.py` | the kinetic action of one particle and of a population, the displacement interpolation, and its velocity field | +| `ot_flow_matching.py` | `python lecture_7/ot_flow_matching.py` | **trains** two velocity networks on the four-Gaussian target, one per coupling, and measures the straightness of each | +| `static_bridge.py` | `python lecture_7/static_bridge.py` | the eight path probabilities, the chain rule, proportional fitting, the tilted matrices, and the bridge path law | +| `dynamic_bridge.py` | `python lecture_7/dynamic_bridge.py` | Girsanov on a constant control, the backward and Hamilton-Jacobi-Bellman residuals, and the Brownian bridge | +| `dsb.py` | `python lecture_7/dsb.py` | the two half-steps of proportional fitting on path laws, and the diffusion reversal on a Gaussian | +| `dsbm.py` | `python lecture_7/dsbm.py` | the averaged bridge drift at one point, and one round of Markovian fitting on the chain | +| `sf2m.py` | `python lecture_7/sf2m.py` | the two closed-form regression targets at a bridge sample, and the entropic plan of a batch | +| `discrete_bridge.py` | `python lecture_7/discrete_bridge.py` | the two-state continuous-time bridge, its Euler sampler and its reverse rates, and the cost computed twice | +| `ddsbm.py` | `python lecture_7/ddsbm.py` | the single bond with three types, the rates pinned at a terminal graph, and the counting over forty-five features | +| `csbm.py` | `python lecture_7/csbm.py` | the two projections in discrete time, and D-IMF on sixteen paths against the reference's mixing | +| `branched.py` | `python lecture_7/branched.py` | the two-branch energies, the tilt of a reference, and **trains** a bias force on two interacting particles | + +Every file takes `--seed` and `--quiet`, and `--steps` means the quantity that file iterates: the +number of quantile levels, Sinkhorn sweeps, times on a grid, optimizer steps, fitting rounds or +Euler steps. The slowest default run is about seven seconds; all thirteen run on a laptop CPU with +no download. + +## What a default run produces + +`optimal_transport.py`. The ordered pairing of the four points costs 1 against 5 for the crossed +one and 3 for the independent coupling. The two-by-two optimum is `[[0.3, 0.3], [0, 0.4]]` at cost +3.4, the prices `f = (0, -8)` and `g = (1, 9)` have slack `[[0, 0], [8, 0]]` and dual value 3.4, so +the duality gap is zero and `W_2 = 1.8439`. The six permutations of the three-point example cost +11/3, 17/3, 35/3, 15, 77/3 and 27, and the quantile integral on 3072 levels returns 3.6667. +Transporting N(0,1) to N(3,4) gives `W_2^2 = 10` under the map `T(x) = 3 + 2x`, and the same map on +a stratified sample gives 9.9996. + +`sinkhorn.py`. The kernel is `[[0.606531, 0.011109], [0.606531, 0.606531]]`. The first sweep gives +`u = (0.971440, 0.329744)` and `v = (0.38013, 3.32081)`, and the marginal error falls 0.680, 0.382, +0.155, 0.047, 0.012, reaching 1e-9 after 17 sweeps. The entropic plan is +`[[0.293122, 0.306878], [0.006878, 0.393122]]` at cost 3.45502 against 3.4, and the regularized +objective is 1.208 at it, 1.222 at the exact optimum and 1.792 at independence. At epsilon = 10, +2, 0.5 and 0.1 the cost is 4.0540165, 3.4550204, 3.4000004 and 3.4000000, in 7, 17, 28 and 80 +sweeps. + +`displacement.py`. A steady particle has action 4 and a dawdling one 8, for the same endpoints. The +average squared displacement is 1, 3 and 5 under the ordered, independent and crossed pairings. At +t = 1/2 the displacement interpolation is N(1.5, 2.25) and the mixture of the two ends peaks at +0.13, with total variation 0.257 between them, and `W_2(rho_0, rho_1/2) = 1.5811`, half of 3.1623. +The straight paths have action 9.9987, which is their average squared displacement. The action +density is 9.9994 at every time on the grid, the stratified estimate of the exact 10, and the +continuity residual is 1.1e-16. + +`ot_flow_matching.py`. On the two-point batch the independent coupling has average squared velocity +3 and the optimal one 1. Over training the batch coupling cost is 10.31 against 3.25, the learned +paths have straightness 1.80 against 0.065, and the one-step error is 2.56 against 0.41. Both +models find all four modes. At `--steps 1500`, about eighteen seconds, the optimal coupling reaches +straightness 0.024 and one-step error 0.28 while the independent one does not improve. + +`static_bridge.py`. The eight reference paths sum to one with `R(1,1,2) = 0.105`. The chain rule +splits 0.2180 into 0.0201 and 0.1979. The reference lands 0.73 away from the requested terminal +law, and proportional fitting reaches `f = (1.072002, 0.937019)` and `g = (0.350517, 1.843637)`, +with the terminal error falling by 187.5 per round. The bridge coupling is +`[[0.114605, 0.385395], [0.085395, 0.414605]]` at `KL = 0.2820` against 0.2847 for independence. +The tilted matrices carry (0.5, 0.5) to (0.449213, 0.550787) to (0.2, 0.8), the path ratios take +four values, and `KL(P* || R) = 0.281961`. + +`dynamic_bridge.py`. The relative entropy of a constant control u = 2 is 2 from the control energy, +2 from the two endpoint laws and 1.9965 from the simulated paths. The backward residual is 2.2e-16 +and the Hamilton-Jacobi-Bellman residual 8.9e-16, and `-grad V` equals the bridge drift to 4.4e-16. +At t = 0.25 the bridge has mean 0.5 and variance 0.1875, and at the state 0.6732 its drift is +1.7691 against the straight-line velocity 2. At t = 1/2 the simulated mean is 1.0095 and the +simulated variance 0.2563, against the exact 1 and 0.25. + +`dsb.py`. The reversal of the last step is `[[0.681416, 0.318584], [0.379310, 0.620690]]`. The +backward chain started from (0.2, 0.8) reaches (0.466585, 0.533415), which is 0.0668 from the source +law, and the second half-step returns the terminal law (0.201967, 0.798033), off by 0.0039. That is +the round-1 error of proportional fitting on the coupling, and the error falls by 187.5 per round. +Run backward from its own terminal law, the diffusion reversal lands at variance 0.9836 against the +exact 1. + +`dsbm.py`. Two bridges through x = 1 at t = 0.5 have drifts 4 and -2 and average 1, where the +derivative of the squared error vanishes. One round of Markovian fitting gives +`M_0 = [[0.591425, 0.408575], [0.307692, 0.692308]]` and +`M_1 = [[0.298457, 0.701543], [0.119588, 0.880412]]`, whose endpoint coupling has exact marginals +and brings the gap from 0.0584 to 0.0077. The gap then falls by 7.37 per round and nine rounds reach +1e-9. + +`sf2m.py`. At t = 0.25 with z = 0.4 the sample is 0.6732, the drift 1.7691, the score -0.9238, the +flow target 2.2309 and the scaled score exactly -0.4. At t = 0.75 the coefficient changes sign and +the flow target is 1.7691, so the two lie symmetrically about the straight-line velocity 2. + +`discrete_bridge.py`. Halving every rate costs 0.1534 and doubling costs 0.3863. The two-state +bridge has `g = (0.3349, 1.9643)`, `phi_1/2 = (0.7999, 1.2667)` and the generator +`[[-1.5836, 1.5836], [0.9472, -0.9472]]` at t = 1/2, where the marginal is (0.4668, 0.5332) and the +reverse rates 1.3862 and 1.0821 match the forward flux at 0.7392. The Euler sampler's terminal +error is 0.0223, 0.0106, 0.0051 and 0.0020 as the step is halved twice and then taken to a hundred +steps. The cost is 0.3235 on the endpoints and 0.3235 on the paths, with the integrand 0.1300 at +t = 1/2, and it rises 0.3148, 0.3235, 0.3337, 0.3348 toward the limit 0.3348 as the reference +speeds up. + +`ddsbm.py`. The bond reference has `beta = 0.9163` and off-diagonal rate 0.3054, and the first +column scaling is (0.2083, 0.7143, 2.9167). Two hundred sweeps give `f = (1.1779, 0.9506, 0.5062)` +and `g = (0.1842, 0.6918, 3.0004)`, a plan that moves 0.4948 of the mass from no bond to double and +leaves 0.2612 unchanged, between 0.18 for independence and 0.4 for the edit-minimizing plan. At +r = 0.2, 0.4, 0.6, 0.8 and 0.99 the mass that stays is 0.2250, 0.2612, 0.2926, 0.3248 and 0.3812. +The survival ratio reaches 1/2 at t = 0.2435, where the pinned rates are 1.2217, 0.0764 and 0.3054, +a factor of 16 between a move onto the terminal type and a move off it, and one such edge against a +network rate of 1 contributes 0.0229. With two terminal types the network learns 0.8552 at a loss +of 0.1357, against 0.1467 at 1 and 0.2395 at 0.5. Nine heavy atoms give 45 features and 9.2e+27 +graphs, at a total jump rate of 24.74 per unit time. + +`csbm.py`. With alpha = 0.35 and three transitions, `A^3 = [[0.5135, 0.4865], [0.4865, 0.5135]]` +and the exact bridge coupling is `[[0.10432, 0.39568], [0.09568, 0.40432]]`. The path 1,1,1,1 has +reference bridge weight 0.5348, so its probability is 0.0535 at the start and 0.0558 under the +bridge. The two laws agree at the interior times, (0.4754, 0.5246) and (0.4107, 0.5893), and differ +by 0.0173 over the paths. One round of D-IMF brings that to 2.9e-04 and later rounds divide it by +58.9, reaching 2.4e-11 after five. The last conditional (0.3160, 0.6840) against a network at +(0.25, 0.75) costs 0.0110 nats. At alpha = 0.2, 0.35 and 0.4 the starting gap is 0.13589, 0.01728 +and 0.00512 and the rounds divide it by 5, 59 and 284; on 8, 16 and 64 paths it is 0.05743, 0.01728 +and 0.00156 and the rounds divide it by 7, 59 and 4674. The ordered reference out of the middle of +five categories is (0.0545, 0.2442, 0.4026, 0.2442, 0.0545). + +`branched.py`. The two branches have weighted energies 1.7 and 0.3 for a total of 2.0, and with the +second target at -4 the linear schedule costs 2.9 against 2.6 for one that releases mass late. The +bent interpolant has energy 2.1667 against 2 for the straight line. The two-outcome tilt has +`Z = 0.109`, raises the chance of reaching B from 0.1 to 0.917, and costs 1.836 nats, with +`J(b*) = -log Z = 2.216`. Training the bias force raises the mean terminal reward from -21.9 under +the reference to -6.9. A longer run does not improve that, because the importance weights +concentrate as the control grows and the effective sample size falls with it. + +## What is simplified + +- Every example runs on the chapter's own two states, three bond types, four points or two + particles. Nothing here is a benchmark, and no file touches a real dataset. +- The minibatch coupling of `ot_flow_matching.py` is an exact assignment through + `scipy.optimize.linear_sum_assignment`, which needs the two sides to be the same size. The + chapter's general `n` by `m` problem with unequal masses appears only in `sinkhorn.py`. +- The two training files use small networks and short runs, so their reports show the direction of + an effect and not its converged size. +- `dsb.py` and `dsbm.py` carry out their projections exactly on the full table of paths. The + networks the chapter describes, which replace those projections by a regression when the state + space is large, are not trained here. +- `discrete_bridge.py` and `ddsbm.py` work one categorical feature at a time. `ddsbm.py` counts the + forty-five features of a molecular graph and the rate of the product process, and no graph is + built. +- `csbm.py` runs both projections exactly on the sixteen-entry table. The network that replaces the + conditionals, and the simulation that replaces the exact endpoint coupling, are not trained here. +- `branched.py` uses constant growth rates taken from the worked example, and trains only the bias + force of the entangled construction. The four-stage training of the branched construction and the + state cost `V_t` are not fitted. diff --git a/lecture_7/branched.py b/lecture_7/branched.py new file mode 100644 index 0000000000000000000000000000000000000000..e80cebbfa0160b915cd5cc72cdf58a1f28b07140 --- /dev/null +++ b/lecture_7/branched.py @@ -0,0 +1,344 @@ +"""Two bridges that change the shape of the problem: branching, and interacting particles. + +A population leaving one source does not always end in one place, and Branched Schrodinger Bridge +Matching gives each outcome its own controlled diffusion and its own share of the mass, with a +learned growth rate moving mass between branches. Entangled Schrodinger Bridge Matching keeps one +population and changes the reference, taking the molecular dynamics of the system itself and +tilting it at the final time by a target density, which a bias force on every particle reproduces. + +Chapter 7, Sections 7.26 to 7.27.3 of the notes. Runs the two-branch example in closed form and +trains a bias force on two interacting particles in the plane. Prints the weighted energies of the +two branch schedules, the interpolant and its energy, the reweighting arithmetic of the two-outcome +example, one Euler-Maruyama step, and the reward the trained force reaches. + + python lecture_7/branched.py # 100 optimizer steps on the bias force + python lecture_7/branched.py --steps 20 # a few seconds, for a first look +""" + +# %% 01. Imports and configuration +from __future__ import annotations + +import math + +import torch +import torch.nn as nn + +from cis6270.nets import TimeEmbedding, count_parameters +from cis6270.runner import ( + common_parser, pick_device, print_report, save_run, set_seed, train, +) + +PARTICLES = 2 # n, the number of interacting particles. +DIMENSION = 2 # Each particle lives in the plane, so the projector below has something to remove. +MASS, FRICTION, TEMPERATURE = 1.0, 1.0, 1.0 # m_i, gamma and k_B tau, all in reduced units. +NOISE = math.sqrt(2.0 * FRICTION * TEMPERATURE / MASS) # The noise scale of the reference. +TARGET_WIDTH = 0.4 # The width of the target density around the region B. +PATHS = 64 # Trajectories simulated per optimizer step. +SIM_STEPS = 30 # Euler-Maruyama steps per trajectory. + + +# %% 02. A population that splits +# +# Branch k is a controlled diffusion carrying the fraction w^k_t of the population, and the weights +# change at a learned growth rate with sum_k g^k_t = 0, so a branch grows only by taking mass from +# another. Section 7.26. + + +def branch_weights(growth, t: torch.Tensor) -> torch.Tensor: + """w^1_t = integral of the growth rate from 0 to t, with branch 0 holding the remainder.""" + return growth(t) # The weight of branch 1; branch 0 carries 1 - w^1_t. + + +def weighted_energy(drift: float, weight, points: int = 2001) -> float: + """The time integral of (1/2) ||u||^2 times the weight of the branch, by the trapezoid rule.""" + t = torch.linspace(0.0, 1.0, points) + density = 0.5 * drift ** 2 * weight(t) # The control energy priced by the mass that spends it. + return float(torch.trapezoid(density, t)) + + +def branch_sample(weight_one: float, generator: torch.Generator) -> int: + """Draw a branch with probability w^k_t, which is how a sample of the population is taken.""" + return int(torch.rand(1, generator=generator) < weight_one) # 1 for branch 1, 0 for branch 0. + + +# %% 03. The interpolant +# +# A network correction multiplied by t(1 - t) leaves both endpoints fixed for every value of the +# parameters, so the endpoint constraints hold during training with no penalty term. With no state +# cost the energy of a bent path exceeds the energy of the straight line, and the training drives +# the correction to zero. Section 7.26.2. + + +def interpolant(x0: float, x1: float, correction: float, t: torch.Tensor) -> torch.Tensor: + """x_t = (1 - t) x_0 + t x_1 + t(1 - t) phi, with phi here a constant.""" + return (1.0 - t) * x0 + t * x1 + t * (1.0 - t) * correction + + +def interpolant_velocity(x0: float, x1: float, correction: float, t: torch.Tensor) -> torch.Tensor: + """x_dot = (x_1 - x_0) + (1 - 2t) phi, the derivative of the interpolant in t.""" + return (x1 - x0) + (1.0 - 2.0 * t) * correction + + +def interpolant_energy(x0: float, x1: float, correction: float, points: int = 2001) -> float: + """The integral of one half the squared velocity of the interpolant over the unit interval.""" + t = torch.linspace(0.0, 1.0, points) + return float(torch.trapezoid(0.5 * interpolant_velocity(x0, x1, correction, t) ** 2, t)) + + +# %% 04. The reference of Entangled Schrodinger Bridge Matching +# +# The reference is the dynamics the system actually follows, underdamped Langevin with the physical +# force, the friction and the thermal noise. Section 7.27. + + +def potential(positions: torch.Tensor) -> torch.Tensor: + """U(R), a double well on each particle plus a spring that couples the two. [B, n, d] -> [B].""" + radial = ((positions ** 2).sum(dim=-1) - 1.0) ** 2 # [B, n] a well at every unit circle. + coupling = 0.5 * ((positions[:, 0] - positions[:, 1]) ** 2).sum(dim=-1) # [B] the spring. + return radial.sum(dim=-1) + coupling # [B]. + + +def physical_force(positions: torch.Tensor) -> torch.Tensor: + """-grad_R U(R), differentiated by hand so the simulation needs no autograd. [B, n, d].""" + radial = -4.0 * ((positions ** 2).sum(dim=-1, keepdim=True) - 1.0) * positions # [B, n, d]. + separation = positions[:, 0] - positions[:, 1] # [B, d] the displacement across the spring. + spring = torch.stack([-separation, separation], dim=1) # [B, n, d] equal and opposite. + return radial + spring + + +def force_check(positions: torch.Tensor) -> float: + """The largest gap between the force written by hand and the gradient autograd computes.""" + r = positions.detach().requires_grad_(True) + gradient = torch.autograd.grad(potential(r).sum(), r)[0] + return float((physical_force(positions) + gradient).abs().max()) + + +TARGET = torch.tensor([1.0, 0.0]) # The centre of the region B, the same for both particles. + + +def log_target(positions: torch.Tensor) -> torch.Tensor: + """r(X_1) = log pi_B(R_1), zero inside the region and negative outside. [B, n, d] -> [B].""" + offset = positions - TARGET.to(positions) # [B, n, d] displacement from the target. + return -(offset ** 2).sum(dim=(1, 2)) / (2.0 * TARGET_WIDTH ** 2) # [B]. + + +def target_direction(positions: torch.Tensor) -> torch.Tensor: + """s_hat_i, the unit vector along the gradient of log pi_B for particle i. [B, n, d].""" + gradient = -(positions - TARGET.to(positions)) / TARGET_WIDTH ** 2 # [B, n, d]. + return gradient / gradient.norm(dim=-1, keepdim=True).clamp(min=1e-8) + + +# %% 05. The bias force and its parameterization +# +# A nonnegative scalar along the direction of steepest increase of the target density, plus a second +# output with its component along that direction removed, so the push toward the target is never +# negative and is read off directly as alpha. Section 7.27.3. + + +def compose_bias(alpha: torch.Tensor, h: torch.Tensor, s_hat: torch.Tensor) -> torch.Tensor: + """b = alpha s_hat + (I - s_hat s_hat^T) h, the parameterization of the force on each particle.""" + along = (h * s_hat).sum(dim=-1, keepdim=True) * s_hat # [B, n, d] the removed component. + return alpha.unsqueeze(-1) * s_hat + (h - along) # [B, n, d]. + + +class BiasForce(nn.Module): + """A network from the positions, the velocities and the time to one force per particle.""" + + def __init__(self, width: int = 64): + super().__init__() + self.time_embedding = TimeEmbedding(width, scale=1.0) + self.input = nn.Linear(2 * PARTICLES * DIMENSION, width) # Positions and velocities. + self.body = nn.Sequential(nn.Linear(width, width), nn.SiLU(), + nn.Linear(width, width), nn.SiLU()) + self.head = nn.Linear(width, PARTICLES * (1 + DIMENSION)) # alpha and h per particle. + + def forward(self, positions: torch.Tensor, velocities: torch.Tensor, + t: torch.Tensor) -> torch.Tensor: + """positions, velocities: [B, n, d], t: [B] -> the bias force [B, n, d].""" + state = torch.cat([positions.flatten(1), velocities.flatten(1)], dim=1) # [B, 2nd]. + hidden = self.body(self.input(state) + self.time_embedding(t)) # [B, width]. + output = self.head(hidden).view(-1, PARTICLES, 1 + DIMENSION) # [B, n, 1 + d]. + alpha = nn.functional.softplus(output[..., 0]) # [B, n] nonnegative, so b never points away. + return compose_bias(alpha, output[..., 1:], target_direction(positions)) + + +# %% 06. Simulating the controlled dynamics +# +# The control enters the velocity equation divided by the mass, against the noise scale, so in units +# of the noise it is b / (m sigma) and Girsanov prices it at one half its squared size. +# Sections 7.27.1 and 7.27.2. + + +@torch.no_grad() +def roll_out(force: BiasForce | None, paths: int, steps: int, device: torch.device, + generator: torch.Generator) -> dict: + """Run the Langevin dynamics with the given bias force, storing the states and the increments.""" + dt = 1.0 / steps + positions = torch.full((paths, PARTICLES, DIMENSION), -1.0, device=device) + positions[:, :, 1] = 0.0 # Both particles start at the far well. + velocities = torch.randn(paths, PARTICLES, DIMENSION, generator=generator).to(device) * \ + math.sqrt(TEMPERATURE / MASS) # Maxwell-Boltzmann velocities at the reference temperature. + stored_positions, stored_velocities, stored_u, increments = [], [], [], [] + energy = torch.zeros(paths, device=device) # (1/2) int ||u||^2 dt along each path. + cross = torch.zeros(paths, device=device) # int u . dW along each path. + for step in range(steps): + t = torch.full((paths,), step * dt, device=device) + u = (force(positions, velocities, t) / (MASS * NOISE)) if force is not None else \ + torch.zeros_like(positions) # The control in units of the noise. + noise = torch.randn(paths, PARTICLES, DIMENSION, + generator=generator).to(device) * math.sqrt(dt) + stored_positions.append(positions.clone()) + stored_velocities.append(velocities.clone()) + stored_u.append(u.clone()) + increments.append(noise.clone()) + energy = energy + 0.5 * (u ** 2).sum(dim=(1, 2)) * dt + cross = cross + (u * noise).sum(dim=(1, 2)) + acceleration = physical_force(positions) / MASS - FRICTION * velocities # The reference. + velocities = velocities + (acceleration + NOISE * u) * dt + NOISE * noise + positions = positions + velocities * dt # The position follows the velocity. + return { + "positions": torch.stack(stored_positions), "velocities": torch.stack(stored_velocities), + "u_bar": torch.stack(stored_u), "noise": torch.stack(increments), + "reward": log_target(positions), "energy": energy, "cross": cross, "dt": dt, + } + + +def importance_weights(rollout: dict) -> torch.Tensor: + """softmax over the batch of r(X_1) + log dR / dP^bbar, which cancels the unknown Z.""" + log_weight = rollout["reward"] - rollout["cross"] - rollout["energy"] + return torch.softmax(log_weight, dim=0) # [P] normalized weights, summing to one. + + +def cross_entropy_loss(force: BiasForce, rollout: dict, weights: torch.Tensor) -> torch.Tensor: + """The weighted average of -log dP^{b_theta} / dR along the stored paths, the loss of 7.27.3.""" + steps, paths = rollout["u_bar"].shape[0], rollout["u_bar"].shape[1] + dt = rollout["dt"] + clock = torch.arange(steps, device=weights.device).to(torch.get_default_dtype()) * dt + times = clock[:, None].expand(steps, paths) # [T, P] the time of every stored state. + u = force(rollout["positions"].flatten(0, 1), rollout["velocities"].flatten(0, 1), + times.flatten()).view_as(rollout["u_bar"]) / (MASS * NOISE) # [T, P, n, d]. + drift = (u * (rollout["u_bar"] * dt + rollout["noise"])).sum(dim=(0, 2, 3)) # [P]. + energy = 0.5 * (u ** 2).sum(dim=(0, 2, 3)) * dt # [P] the control energy of each path. + return (weights * (energy - drift)).sum() # Weighted by the importance weights. + + +# %% 07. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=100, lr=3e-3, width=64, samples=256) + args = parser.parse_args() + set_seed(args.seed) + device = pick_device(args.device) + loud = not args.quiet + generator = torch.Generator().manual_seed(args.seed) + + linear = lambda t: 0.3 * t # The growth rate 0.3, so w^1_t rises at a constant rate. + late = lambda t: 0.3 * t ** 2 # The growth rate 0.6 t, which starts slowly. + near = {"linear": weighted_energy(2.0, lambda t: 1.0 - linear(t)) + + weighted_energy(-2.0, linear)} + far = {"linear": weighted_energy(2.0, lambda t: 1.0 - linear(t)) + weighted_energy(-4.0, linear), + "late": weighted_energy(2.0, lambda t: 1.0 - late(t)) + weighted_energy(-4.0, late)} + half = float(branch_weights(linear, torch.tensor(0.5))) + if loud: + print("Two branches from the point 0, ending at +2 with weight 0.7 and at -2 with 0.3") + print(f" weights at t = 1/2: {1 - half:.2f} and {half:.2f}; at t = 1: " + f"{1 - float(branch_weights(linear, torch.tensor(1.0))):.1f} and " + f"{float(branch_weights(linear, torch.tensor(1.0))):.1f}") + print(f" weighted energies {weighted_energy(2.0, lambda t: 1.0 - linear(t)):.1f} and " + f"{weighted_energy(-2.0, linear):.1f}, total {near['linear']:.1f}") + print(f" with the second target at -4: linear schedule {far['linear']:.1f}, " + f"late schedule {far['late']:.1f}, saving {far['linear'] - far['late']:.1f}") + + bent = interpolant_energy(0.0, 2.0, 1.0) + straight = interpolant_energy(0.0, 2.0, 0.0) + grid = torch.tensor([0.0, 0.5, 1.0]) + if loud: + print(f"\nThe interpolant with a correction of 1 passes through " + f"{[round(float(z), 2) for z in interpolant(0.0, 2.0, 1.0, grid)]}") + speed = [round(float(z), 0) for z in interpolant_velocity(0.0, 2.0, 1.0, grid)] + print(f" its velocity is {speed} and its energy {bent:.4f}, against " + f"{straight:.1f} for the straight line") + + reference_mass, inside, outside = 0.1, 1.0, 0.01 # The two-outcome example of Section 7.27.1. + normalizer = reference_mass * inside + (1.0 - reference_mass) * outside + tilted_mass = reference_mass * inside / normalizer + reward_average = tilted_mass * math.log(inside) + (1.0 - tilted_mass) * math.log(outside) + tilt_entropy = reward_average - math.log(normalizer) + if loud: + print(f"\nA reference reaching B one time in ten, tilted by pi_B = 1 inside and 0.01 outside") + print(f" Z = {normalizer:.3f}, P*(B) = {tilted_mass:.3f}, " + f"a factor of {tilted_mass / reference_mass:.2f}") + print(f" E[r] = {reward_average:.3f}, -log Z = {-math.log(normalizer):.3f}, " + f"KL(P* || R) = {tilt_entropy:.3f}, J(b*) = {tilt_entropy - reward_average:.3f}") + + step_drift = (-0.3 + 0.5) * 0.01 # The physical force plus the bias, over one step. + step_friction = -1.0 * 0.8 * 0.01 # Minus the friction times the current velocity. + step_noise = NOISE * math.sqrt(0.01) * 0.4 # The noise scale times a normal draw of 0.4. + unit_control = 0.5 / NOISE # The bias force measured in units of the noise. + plane = compose_bias(torch.tensor([[0.5]]), torch.tensor([[[0.3, -0.4]]]), + torch.tensor([[[1.0, 0.0]]]))[0, 0] # The example of Section 7.27.3. + log_weights = torch.tensor([-1.0, 0.0, 1.0]) + softmax = torch.softmax(log_weights, dim=0) + if loud: + print(f"\nOne Euler-Maruyama step at dt = 0.01, with a noise scale of {NOISE:.4f}") + print(f" dv = {step_drift:.4f} {step_friction:+.4f} {step_noise:+.4f} = " + f"{step_drift + step_friction + step_noise:.4f}, dr = {0.8 * 0.01:.4f}") + print(f" the control in units of the noise is {unit_control:.4f}, contributing " + f"{0.5 * unit_control ** 2 * 0.01:.6f} per step and " + f"{0.5 * unit_control ** 2 * 0.01 * 100:.4f} over a hundred") + print(f" b = {[round(float(z), 1) for z in plane]}, with alpha " + f"{float((plane * torch.tensor([1.0, 0.0])).sum()):.1f} along s_hat " + f"and length {float(plane.norm()):.3f}") + print(f" three log-weights -1, 0, 1 normalize to {[round(float(z), 3) for z in softmax]}") + + force = BiasForce(args.width).to(device) + reference_run = roll_out(None, args.samples, SIM_STEPS, device, generator) + before = roll_out(force, args.samples, SIM_STEPS, device, generator) + + def loss_fn(step: int) -> torch.Tensor: + rollout = roll_out(force, PATHS, SIM_STEPS, device, generator) # Simulate under b_bar. + return cross_entropy_loss(force, rollout, importance_weights(rollout)) + + if loud: + print(f"\ntraining the bias force, {count_parameters(force):,} parameters, " + f"{PATHS} paths of {SIM_STEPS} steps per optimizer step") + losses = train(force, loss_fn, steps=args.steps, lr=args.lr, quiet=not loud, log_every=50) + after = roll_out(force, args.samples, SIM_STEPS, device, generator) + effective = float(1.0 / (importance_weights(after) ** 2).sum() / args.samples) + + report = { + "weight of branch 1 at t = 1/2": half, + "weighted energy, branch 0": weighted_energy(2.0, lambda t: 1.0 - linear(t)), + "weighted energy, branch 1": weighted_energy(-2.0, linear), + "total, both ends at distance 2": near["linear"], + "total with the far target, linear": far["linear"], + "total with the far target, late": far["late"], + "a branch drawn at t = 1/2": branch_sample(half, generator), + "interpolant energy, correction 1": bent, + "interpolant energy, straight": straight, + "Z of the two outcomes": normalizer, + "P*(B)": tilted_mass, + "the factor the tilt applies": tilted_mass / reference_mass, + "E_{P*}[r]": reward_average, + "KL(P* || R)": tilt_entropy, + "-log Z, the value of J(b*)": -math.log(normalizer), + "one velocity step": step_drift + step_friction + step_noise, + "control in units of the noise": unit_control, + "control energy over a hundred steps": 0.5 * unit_control ** 2 * 0.01 * 100, + "the composed force": [round(float(z), 2) for z in plane], + "its length": float(plane.norm()), + "softmax of -1, 0, 1": [round(float(z), 3) for z in softmax], + "mean reward, reference": float(reference_run["reward"].mean()), + "mean reward, before training": float(before["reward"].mean()), + "mean reward, after training": float(after["reward"].mean()), + "control energy after training": float(after["energy"].mean()), + "effective sample size, fraction": effective, + "force written by hand against autograd": force_check(after["positions"][-1]), + "final loss": sum(losses[-10:]) / len(losses[-10:]), + } + print_report("Branched and Entangled bridges, Chapter 7 Sections 7.26 to 7.27.3", report) + save_run(args.out, config=vars(args), report=report) # Written only with --out. + + +if __name__ == "__main__": + main() diff --git a/lecture_7/csbm.py b/lecture_7/csbm.py new file mode 100644 index 0000000000000000000000000000000000000000..7cd9634ed0db307ab72fd0bb624809c45111fc0b --- /dev/null +++ b/lecture_7/csbm.py @@ -0,0 +1,254 @@ +"""Categorical Schrodinger Bridge Matching: D-IMF on a path law small enough to write out. + +With N intermediate times a path is N + 2 states and the whole path law is a finite table, so both +projections have one-line expressions. The reciprocal projection holds the endpoint coupling and +redraws the interior from the reference, the Markov projection keeps the one-step conditionals and +multiplies them into a chain, and the bridge is the only law in both sets. On two states and three +transitions the sixteen paths fit in a list, so the alternation runs exactly. + +Chapter 7, Section 7.25 of the notes. Runs on S = 2, four time points, the uniform reference with +alpha = 0.35, the source law (0.5, 0.5) and the requested law (0.2, 0.8). Prints the uniform and +the ordered reference row, the exact bridge coupling, the conditionals of the first Markov +projection, the gap after each round, and how that gap moves with alpha and with the path length. + + python lecture_7/csbm.py # 5 rounds of D-IMF, the five the chapter reports + python lecture_7/csbm.py --steps 1 # the single round the chapter writes out +""" + +# %% 01. Imports and configuration +from __future__ import annotations + +import torch + +from cis6270.runner import common_parser, print_report, save_run, set_seed + +DTYPE = torch.float64 # The gap reaches 1e-11 in five rounds, well below float32. + +ALPHA = 0.35 # The probability of changing category in one step. +TRANSITIONS = 3 # N + 1, so four time points and sixteen paths. +PI_0 = torch.tensor([0.5, 0.5], dtype=DTYPE) # The source law. +PI_1 = torch.tensor([0.2, 0.8], dtype=DTYPE) # The requested law at the last time. + + +# %% 02. The reference over S categories +# +# The uniform reference holds the category with probability 1 - alpha and spreads the rest evenly +# over the alternatives. For categories that have an order a discretized Gaussian centred on the +# current one is available instead, under which a neighbour is likelier than a distant category. +# Section 7.25. + + +def uniform_reference(alpha: float, size: int) -> torch.Tensor: + """A_n(x, y) = 1 - alpha on the diagonal and alpha / (S - 1) off it. [S, S].""" + off = alpha / (size - 1) # The share each alternative receives. + return torch.full((size, size), off, dtype=DTYPE) + torch.eye(size, dtype=DTYPE) * ( + 1.0 - alpha - off) + + +def ordered_reference(width: float, size: int) -> torch.Tensor: + """Rows proportional to exp(-(y - x)^2 / 2 width^2), under which a nearby category is likelier.""" + offset = torch.arange(size, dtype=DTYPE)[None, :] - torch.arange(size, dtype=DTYPE)[:, None] + weights = torch.exp(-offset ** 2 / (2.0 * width ** 2)) # [S, S] before normalizing. + return weights / weights.sum(dim=1, keepdim=True) + + +def sinkhorn(matrix: torch.Tensor, start: torch.Tensor, end: torch.Tensor, sweeps: int = 400 + ) -> torch.Tensor: + """Scale diag(p_0) A^N to the two marginals, which gives the exact bridge coupling.""" + f = torch.ones_like(start) + g = torch.ones_like(end) + for _ in range(sweeps): + g = end / (matrix.T @ f) # Column sums become exactly p_1. + f = start / (matrix @ g) # Row sums become exactly p_0. + return f[:, None] * matrix * g[None, :] + + +# %% 03. The two projections in discrete time +# +# A path law is a tensor with one axis per time. The reciprocal projection multiplies the endpoint +# coupling by the reference bridge weight of every path, and the Markov projection reads off the +# one-step conditionals and multiplies them back into a chain. Section 7.25.1. + + +def reference_joint(matrix: torch.Tensor, steps: int) -> torch.Tensor: + """The product of the transition probabilities along a path, as a tensor with one axis per time.""" + joint = matrix.clone() # [x_0, x_1]. + for _ in range(steps - 1): + joint = joint.unsqueeze(-1) * matrix # Append one more transition on a new axis. + return joint # [S] * (steps + 1). + + +def bridge_table(matrix: torch.Tensor, steps: int) -> torch.Tensor: + """R(x_interior | x_0, x_N), the reference bridge weight of every path given its two ends.""" + whole = torch.linalg.matrix_power(matrix, steps) # [S, S] the reference over the interval. + size = len(matrix) + shape = [size] + [1] * (steps - 1) + [size] # Broadcast over the interior axes. + return reference_joint(matrix, steps) / whole.reshape(shape) + + +def reciprocal_projection(coupling: torch.Tensor, matrix: torch.Tensor, steps: int) -> torch.Tensor: + """proj_R(P) = P(x_0, x_N) R(x_in | x_0, x_N), the endpoint coupling with reference bridges.""" + size = len(matrix) + shape = [size] + [1] * (steps - 1) + [size] + return coupling.reshape(shape) * bridge_table(matrix, steps) + + +def markov_projection(law: torch.Tensor) -> tuple[list[torch.Tensor], torch.Tensor]: + """proj_M(P) = P(x_0) prod_n P(x_{t_n} | x_{t_{n-1}}), as the conditionals and the coupling.""" + conditionals = [] + for step in range(law.dim() - 1): + axes = [d for d in range(law.dim()) if d not in (step, step + 1)] + pair = law.sum(dim=axes) # [S, S] the joint of two consecutive states. + conditionals.append(pair / pair.sum(dim=1, keepdim=True)) + joint = torch.diag(law.sum(dim=list(range(1, law.dim())))) # [S, S] the law at time 0. + for matrix in conditionals: + joint = joint @ matrix + return conditionals, joint + + +def transition_loss(true: torch.Tensor, network: torch.Tensor) -> float: + """KL(P(x_t | x_s) || m_theta(x_t | x_s)) on one transition out of one state, in nats.""" + return float((true * (true / network).log()).sum()) + + +# %% 04. D-IMF on sixteen paths +# +# Alternating the two projections is iterative Markovian fitting with discrete time. The starting +# law is the independent coupling filled in with reference bridges, so both requested marginals are +# already exact and only the joint law over paths moves. Section 7.25.2. + + +def d_imf(coupling: torch.Tensor, matrix: torch.Tensor, steps: int, rounds: int, + target: torch.Tensor) -> tuple[list[float], list[torch.Tensor]]: + """Run the alternation, returning the gap after each round and the first round's conditionals.""" + gaps: list[float] = [] + first: list[torch.Tensor] = [] + for index in range(rounds): + law = reciprocal_projection(coupling, matrix, steps) # Reference bridges in the interior. + conditionals, coupling = markov_projection(law) # The chain and its endpoint coupling. + if index == 0: + first = conditionals + gaps.append(float((reciprocal_projection(coupling, matrix, steps) - target).abs().sum())) + return gaps, first + + +def starting_gap(matrix: torch.Tensor, steps: int, start: torch.Tensor, + end: torch.Tensor) -> tuple[float, torch.Tensor, torch.Tensor]: + """The gap of the independent coupling from the bridge, with both path laws.""" + whole = torch.linalg.matrix_power(matrix, steps) + plan = sinkhorn(torch.diag(start) @ whole, start, end) + target = reciprocal_projection(plan, matrix, steps) + initial = reciprocal_projection(start[:, None] * end[None, :], matrix, steps) + return float((initial - target).abs().sum()), initial, target + + +# %% 05. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=5) # --steps is the number of D-IMF rounds. + args = parser.parse_args() + set_seed(args.seed) + loud = not args.quiet + rounds = max(args.steps, 1) + + A = uniform_reference(ALPHA, len(PI_0)) + whole = torch.linalg.matrix_power(A, TRANSITIONS) + ordered = ordered_reference(1.0, 5) # Five ordered categories, a width of one category. + if loud: + print(f"S = 2, {TRANSITIONS + 1} time points, alpha = {ALPHA}, " + f"{len(PI_0) ** (TRANSITIONS + 1)} paths") + print(f" A = {[[round(float(z), 2) for z in row] for row in A]}, " + f"A^{TRANSITIONS} = {[[round(float(z), 4) for z in row] for row in whole]}") + print(f" the ordered reference out of the middle of five categories: " + f"{[round(float(z), 4) for z in ordered[2]]}") + + plan = sinkhorn(torch.diag(PI_0) @ whole, PI_0, PI_1) + gap_0, initial, target = starting_gap(A, TRANSITIONS, PI_0, PI_1) + interior = [initial.sum(dim=[d for d in range(TRANSITIONS + 1) if d != k]) for k in (1, 2)] + independent = PI_0[:, None] * PI_1[None, :] # Where the alternation starts. + if loud: + print(f"\npi* = {[[round(float(z), 5) for z in row] for row in plan]}, against the " + f"independent coupling " + f"{[[round(float(z), 1) for z in row] for row in independent]}") + print(f" the path 1,1,1,1 has reference bridge weight " + f"{float(bridge_table(A, TRANSITIONS)[0, 0, 0, 0]):.4f}, so its probability is " + f"{float(initial[0, 0, 0, 0]):.4f} at the start and {float(target[0, 0, 0, 0]):.4f} " + f"under the bridge") + print(f" the path 2,2,2,2 has {float(initial[1, 1, 1, 1]):.4f} and " + f"{float(target[1, 1, 1, 1]):.4f}") + print(f" the two laws agree at the interior times, " + f"{[[round(float(z), 4) for z in m] for m in interior]}, and differ by " + f"{gap_0:.4f} over the paths") + + gaps, conditionals = d_imf(independent, A, TRANSITIONS, rounds, target) + # Eight rounds reach machine precision, past which a ratio of two gaps measures rounding, so + # the report quotes the first one. + ratios = [gaps[k] / gaps[k + 1] for k in range(len(gaps) - 1) if gaps[k + 1] > 1e-14] + per_round = ratios[0] if ratios else gap_0 / gaps[0] # One round alone gives only the first. + network = torch.tensor([0.25, 0.75], dtype=DTYPE) # A network that is wrong by a little. + one_term = transition_loss(conditionals[-1][0], network) + if loud: + print(f"\nthe {len(conditionals)} conditionals of the first Markov projection") + for matrix in conditionals: + print(f" {[[round(float(z), 4) for z in row] for row in matrix]}") + print(f" gap {gap_0:.4f} at the start, then {['%.1e' % x for x in gaps]}") + print(f" ratios {[round(r, 1) for r in ratios]}") + print(f" the true conditional {[round(float(z), 4) for z in conditionals[-1][0]]} against " + f"a network at {[float(z) for z in network]} costs {one_term:.4f} nats") + + # A reference that mixes well and a longer chain both leave the bridge less to do, which the + # starting gap records and the rate of convergence follows. Two rounds separate the two ratios + # from the float64 floor, which the sixty-four-path column reaches after three. + mixing = {} + for alpha in (0.2, ALPHA, 0.4): + matrix = uniform_reference(alpha, len(PI_0)) + start, _, goal = starting_gap(matrix, TRANSITIONS, PI_0, PI_1) + later, _ = d_imf(independent, matrix, TRANSITIONS, 2, goal) + mixing[alpha] = (start, later[0] / later[1]) + lengths = {} + for steps in (2, TRANSITIONS, 5): + start, _, goal = starting_gap(A, steps, PI_0, PI_1) + later, _ = d_imf(independent, A, steps, 2, goal) + lengths[steps] = (start, later[0] / later[1]) + if loud: + print("\nalpha starting gap gap divided per round") + for alpha, (start, ratio) in mixing.items(): + print(f" {alpha:<6g} {start:<14.5f} {ratio:.0f}") + print("\npaths starting gap gap divided per round") + for steps, (start, ratio) in lengths.items(): + print(f" {len(PI_0) ** (steps + 1):<6d} {start:<14.5f} {ratio:.0f}") + + report = { + "A^3": [round(float(z), 4) for z in whole.flatten()], + "the ordered reference row": [round(float(z), 4) for z in ordered[2]], + "pi*_11": float(plan[0, 0]), + "pi*_12": float(plan[0, 1]), + "reference bridge weight of 1,1,1,1": float(bridge_table(A, TRANSITIONS)[0, 0, 0, 0]), + "P^0(1,1,1,1)": float(initial[0, 0, 0, 0]), + "P*(1,1,1,1)": float(target[0, 0, 0, 0]), + "P^0(2,2,2,2)": float(initial[1, 1, 1, 1]), + "P*(2,2,2,2)": float(target[1, 1, 1, 1]), + "marginal at the first interior time": [round(float(z), 4) for z in interior[0]], + "marginal at the second": [round(float(z), 4) for z in interior[1]], + "gap at the start": gap_0, + "gap after one round": gaps[0], + "gap after the last round": gaps[-1], + "rounds run": rounds, + "gap divided per round": per_round, + "first conditional, out of state 1": [round(float(z), 4) for z in conditionals[0][0]], + "last conditional, out of state 1": [round(float(z), 4) for z in conditionals[-1][0]], + "one transition against a wrong network": one_term, + "starting gap at alpha = 0.2": mixing[0.2][0], + "rounds divide the gap by, at alpha = 0.2": mixing[0.2][1], + "rounds divide the gap by, at alpha = 0.4": mixing[0.4][1], + "starting gap on eight paths": lengths[2][0], + "starting gap on sixty-four paths": lengths[5][0], + "rounds divide the gap by, on eight paths": lengths[2][1], + "rounds divide the gap by, on sixty-four paths": lengths[5][1], + } + print_report("Categorical Schrodinger Bridge Matching, Chapter 7 Section 7.25", report) + save_run(args.out, config=vars(args), report=report) # Written only with --out. + + +if __name__ == "__main__": + main() diff --git a/lecture_7/ddsbm.py b/lecture_7/ddsbm.py new file mode 100644 index 0000000000000000000000000000000000000000..5e30038a9e6b8236b81ce861483cea5a59f36de1 --- /dev/null +++ b/lecture_7/ddsbm.py @@ -0,0 +1,243 @@ +"""Discrete Diffusion Schrodinger Bridge Matching: one categorical feature of a molecular graph. + +A molecule is a graph with a categorical feature on every atom and on every pair of atoms, and the +reference resamples each feature independently, so one feature with S categories specifies the +whole construction. Scaling the reference kernel to the two laws gives the bridge coupling for that +feature, and pinning a path at a known terminal graph gives the rate a network regresses onto. +Both stay usable on a state space of 10^27 graphs because the reference is a product over the +features, so every sum over graphs becomes a sum over the categories of one of them. + +Chapter 7, Section 7.24 of the notes. Runs on one bond with three types, no bond, single and +double, at the survival probability r = 0.4. Prints the kernel and its generator, the bridge +coupling and the mass it leaves unchanged, the pinned rates and the posterior average a network +learns, and the counting over the forty-five features of a nine-atom graph. + + python lecture_7/ddsbm.py # 200 Sinkhorn sweeps on the bond coupling + python lecture_7/ddsbm.py --steps 1 # the first sweep alone, computed by hand +""" + +# %% 01. Imports and configuration +from __future__ import annotations + +import math + +import torch + +from cis6270.runner import common_parser, print_report, save_run, set_seed + +DTYPE = torch.float64 # The chapter quotes four figures of the coupling and of the pinned rates. + +PRIOR = torch.full((3,), 1.0 / 3.0, dtype=DTYPE) # m, uniform over the three bond types. +BOND_0 = torch.tensor([0.7, 0.2, 0.1], dtype=DTYPE) # pi_0, mostly no bond. +BOND_1 = torch.tensor([0.1, 0.2, 0.7], dtype=DTYPE) # pi_1, mostly double. +SURVIVAL = 0.4 # r = alpha_bar(1) / alpha_bar(0) over the whole unit interval. +ATOMS = 9 # Heavy atoms in the molecule the chapter counts. +ATOM_TYPES, BOND_TYPES = 5, 4 # Categories of a node feature and of a bond feature. + + +# %% 02. The reference process on one feature +# +# The kernel keeps the category with probability r and redraws it from the prior m otherwise. Over +# a short step r = 1 - h beta(t), so the rate out of x toward a different category y is the +# resampling rate beta(t) times the chance m(y) that the redraw lands on y. Section 7.24.1. + + +def replacement_kernel(r: float, prior: torch.Tensor) -> torch.Tensor: + """K_st(x, y) = r 1_{x=y} + (1 - r) m(y), the uniform-replacement kernel with a general prior.""" + return r * torch.eye(len(prior), dtype=DTYPE) + (1.0 - r) * prior[None, :] # [S, S]. + + +def replacement_generator(beta: float, prior: torch.Tensor) -> torch.Tensor: + """Q_t(x, y) = beta(t) m(y) off the diagonal, with the diagonal -beta(t) (1 - m(x)).""" + rates = beta * prior[None, :].repeat(len(prior), 1) # [S, S]. + rates = rates - torch.diag(rates.diagonal()) # Clear the diagonal before closing the row. + return rates - torch.diag(rates.sum(dim=1)) + + +def resampling_rate(r: float) -> float: + """beta = -log alpha_bar(1) for the schedule alpha_bar(t) = r^t, which is constant in time.""" + return -math.log(r) + + +# %% 03. A bridge for a single bond +# +# Scaling diag(pi_0) K_01 to the two marginals gives the coupling. The mass it leaves on the +# diagonal sits between the mass an independent coupling happens to leave and the largest any +# coupling can leave, and r fixes where between the two it lands. Section 7.24.2. + + +def sinkhorn(matrix: torch.Tensor, start: torch.Tensor, end: torch.Tensor, sweeps: int + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[float]]: + """Scale a positive matrix to the two marginals, returning f, g, the plan and the column errors.""" + f = torch.ones_like(start) + g = torch.ones_like(end) + errors: list[float] = [] + for _ in range(sweeps): + g = end / (matrix.T @ f) # Column sums become exactly pi_1. + f = start / (matrix @ g) # Row sums become exactly pi_0, and the columns move again. + errors.append(float(((f[:, None] * matrix * g[None, :]).sum(dim=0) - end).abs().sum())) + return f, g, f[:, None] * matrix * g[None, :], errors + + +def held_mass(plan: torch.Tensor) -> float: + """The trace of the plan, the fraction of bonds whose type the bridge does not change.""" + return float(plan.diagonal().sum()) + + +# %% 04. Training with the terminal graph fixed +# +# Pinning the endpoint is an h-transform with the reference probability of arriving at x_1 in the +# place of the potential, so a move onto the terminal type has its rate multiplied and a move off +# it divided. Several terminal types can occur at one state, and the minimizer of the rate error +# over a single number is their posterior average. Section 7.24.3. + + +def pinned_rate(generator: torch.Tensor, kernel: torch.Tensor, x: int, y: int, x1: int) -> float: + """Q_t(x, y; x_1) = Q_t(x, y) K_t1(y, x_1) / K_t1(x, x_1), the rate of the bridge ending at x_1.""" + return float(generator[x, y] * kernel[y, x1] / kernel[x, x1]) + + +def half_survival_time(r: float) -> float: + """The time where alpha_bar(1) / alpha_bar(t) = r^{1-t} equals 1/2, on alpha_bar(t) = r^t.""" + return 1.0 - math.log(0.5) / math.log(r) # At r = 0.4 this is log 0.8 / log 0.4 = 0.2435. + + +def rate_error(controlled: float, network: float) -> float: + """Q^pin log(Q^pin / Q^theta) - Q^pin + Q^theta, the Bregman divergence of u log u on one edge.""" + return controlled * math.log(controlled / network) - controlled + network + + +def posterior_loss(rates: list[float], weights: list[float], network: float) -> float: + """The expected rate error at a state whose terminal type is uncertain, over the posterior.""" + return sum(w * rate_error(q, network) for q, w in zip(rates, weights)) + + +# %% 05. The factorization over a graph +# +# The reference applies the same kernel independently at every node and every pair of nodes, so the +# generator changes one feature at a time, the total jump rate out of a graph is the sum of the +# per-feature rates, and in the pinned rate every factor cancels but the one of the feature that +# moves. Section 7.24.4. + + +def feature_count(atoms: int) -> tuple[int, int, int]: + """The node features, the bond features and their total, for a molecule with this many atoms.""" + bonds = atoms * (atoms - 1) // 2 # One feature per unordered pair. + return atoms, bonds, atoms + bonds + + +def state_space_size(atoms: int, atom_types: int, bond_types: int) -> float: + """atom_types^nodes times bond_types^bonds, the number of graphs the sums would run over.""" + nodes, bonds, _ = feature_count(atoms) + return float(atom_types ** nodes) * float(bond_types ** bonds) + + +def total_jump_rate(features: int, beta: float, categories: int) -> float: + """The sum of the per-feature rates out of a graph, which fixes the step size of the sampler.""" + return features * beta * (1.0 - 1.0 / categories) # A redraw lands elsewhere with 1 - m(x). + + +# %% 06. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=200) # --steps is the number of Sinkhorn sweeps on the bond coupling. + args = parser.parse_args() + set_seed(args.seed) + loud = not args.quiet + sweeps = max(args.steps, 1) + + kernel = replacement_kernel(SURVIVAL, PRIOR) + beta = resampling_rate(SURVIVAL) + rates = replacement_generator(beta, PRIOR) + reference = torch.diag(BOND_0) @ kernel + if loud: + print(f"One bond with three types at r = {SURVIVAL}, beta = {beta:.4f}, " + f"off-diagonal rate {float(rates[0, 1]):.4f}") + print(f" K_01 = {[[round(float(z), 1) for z in row] for row in kernel]}") + print(f" diag(pi_0) K_01 has column sums " + f"{[round(float(z), 2) for z in reference.sum(dim=0)]}, so the first g is " + f"{[round(float(z), 4) for z in BOND_1 / reference.sum(dim=0)]}") + + f, g, plan, errors = sinkhorn(reference, BOND_0, BOND_1, sweeps) + independent = float((BOND_0 * BOND_1).sum()) + largest = float(torch.minimum(BOND_0, BOND_1).sum()) + held = {r: held_mass(sinkhorn(torch.diag(BOND_0) @ replacement_kernel(r, PRIOR), + BOND_0, BOND_1, 200)[2]) + for r in (0.2, 0.4, 0.6, 0.8, 0.99)} + if loud: + print(f"\nAfter {sweeps} sweeps, column error {errors[-1]:.2e}") + print(f" f = {[round(float(z), 4) for z in f]}, " + f"g = {[round(float(z), 4) for z in g]}") + print(f" pi* = {[[round(float(z), 3) for z in row] for row in plan]}") + print(f" mass that stays: bridge {held_mass(plan):.4f}, independent {independent:.2f}, " + f"largest possible {largest:.1f}") + print(f" diagonal mass at r = 0.2, 0.4, 0.6, 0.8, 0.99: " + f"{[round(held[r], 4) for r in (0.2, 0.4, 0.6, 0.8, 0.99)]}") + + pinned_time = half_survival_time(SURVIVAL) + pinned_kernel = replacement_kernel(0.5, PRIOR) + toward = pinned_rate(rates, pinned_kernel, 0, 2, 2) # A move onto the terminal type. + away = pinned_rate(rates, pinned_kernel, 2, 0, 2) # A move off it. + between = pinned_rate(rates, pinned_kernel, 0, 1, 2) # Between two non-terminal types. + one_edge = rate_error(toward, 1.0) # The term one edge contributes at a network rate of 1. + weights = [0.6, 0.4] # The posterior over the two terminal types at this state. + posterior = weights[0] * toward + weights[1] * between + losses = {q: posterior_loss([toward, between], weights, q) for q in (posterior, 1.0, 0.5)} + if loud: + print(f"\nAt t = {pinned_time:.4f} the survival ratio to time 1 is 1/2, so K_t1 is " + f"{float(pinned_kernel[0, 0]):.4f} on the diagonal and " + f"{float(pinned_kernel[0, 1]):.4f} off it") + print(f" pinned rates: {toward:.4f} toward the terminal type, {away:.4f} away from it " + f"and {between:.4f} between two others, a factor of {toward / away:.0f}") + print(f" one edge at a network rate of 1 contributes {one_edge:.4f}") + print(f" with two terminal types the network learns {posterior:.4f}, at a loss of " + f"{losses[posterior]:.4f}, against {losses[1.0]:.4f} at 1 and {losses[0.5]:.4f} at 0.5") + + nodes, bonds, features = feature_count(ATOMS) + graphs = state_space_size(ATOMS, ATOM_TYPES, BOND_TYPES) + jump_rate = total_jump_rate(bonds, beta, BOND_TYPES) + if loud: + print(f"\n{ATOMS} heavy atoms: {nodes} node features and {bonds} bond features, " + f"{features} in all") + print(f" {ATOM_TYPES}^{nodes} x {BOND_TYPES}^{bonds} = {graphs:.1e} graphs, " + f"which no sum reaches") + print(f" the total jump rate out of a graph is {bonds} x {beta:.4f} x " + f"(1 - 1/{BOND_TYPES}) = {jump_rate:.2f} per unit time") + + report = { + "beta of alpha_bar(t) = 0.4^t": beta, + "off-diagonal reference rate": float(rates[0, 1]), + "K_01 on the diagonal": float(kernel[0, 0]), + "K_01 off the diagonal": float(kernel[0, 1]), + "rows of the generator sum to": float(rates.sum(dim=1).abs().max()), + "column sums of diag(pi_0) K_01": [round(float(z), 2) for z in reference.sum(dim=0)], + "the first g": [round(float(z), 4) for z in BOND_1 / reference.sum(dim=0)], + "sweeps run": sweeps, + "column error after the last sweep": errors[-1], + "f": [round(float(z), 4) for z in f], + "g": [round(float(z), 4) for z in g], + "bond plan, no bond to double": float(plan[0, 2]), + "bond mass that stays": held_mass(plan), + "bond mass, independent coupling": independent, + "bond mass, largest possible": largest, + "diagonal mass at r = 0.2": held[0.2], + "diagonal mass at r = 0.99": held[0.99], + "time where the survival ratio is 1/2": pinned_time, + "pinned rate toward the terminal": toward, + "pinned rate away from it": away, + "pinned rate between two others": between, + "one edge at a network rate of 1": one_edge, + "posterior average of the rates": posterior, + "loss at the posterior average": losses[posterior], + "loss at 1": losses[1.0], + "loss at 0.5": losses[0.5], + "node, bond and total features": [nodes, bonds, features], + "graphs on nine heavy atoms": f"{graphs:.1e}", + "total jump rate out of a graph": jump_rate, + } + print_report("Discrete Diffusion Schrodinger Bridge Matching, Chapter 7 Section 7.24", report) + save_run(args.out, config=vars(args), report=report) # Written only with --out. + + +if __name__ == "__main__": + main() diff --git a/lecture_7/discrete_bridge.py b/lecture_7/discrete_bridge.py new file mode 100644 index 0000000000000000000000000000000000000000..03a20c4fe27c746e4d835727b93f302283194750 --- /dev/null +++ b/lecture_7/discrete_bridge.py @@ -0,0 +1,267 @@ +"""The Schrodinger bridge over a finite set: rates in place of drifts. + +On a finite state space there is no small displacement, so the control is a reweighting of jump +rates and relative entropy on paths is a time integral of a Bregman divergence between two rate +matrices. Everything else carries over. The chain rule splits the problem into an endpoint term and +a bridge term, the two potentials are the Sinkhorn scalings of the reference endpoint law, and +under Doob's h-transform each reference rate is multiplied by the ratio of a backward potential +at the destination to its value at the current state. + +Chapter 7, Sections 7.21 to 7.23.3 of the notes. Runs the two-state continuous-time example with +the generator Q, the source law (0.5, 0.5) and the requested terminal law (0.2, 0.8). Prints the +bridge rates at three times, the reverse rates, the Euler sampler's error as the step shrinks, +and the cost computed on the endpoints and on the paths. + + python lecture_7/discrete_bridge.py # 10 Euler steps, then 20, 40 and 100 + python lecture_7/discrete_bridge.py --steps 20 # a finer sampler, for a first look +""" + +# %% 01. Imports and configuration +from __future__ import annotations + +import math + +import torch + +from cis6270.runner import common_parser, print_report, save_run, set_seed + +DTYPE = torch.float64 # Matrix exponentials and a time integral, both in double precision. + +Q = torch.tensor([[-1.0, 1.0], [1.5, -1.5]], dtype=DTYPE) # The reference generator. +PI_0 = torch.tensor([0.5, 0.5], dtype=DTYPE) # The source law. +PI_1 = torch.tensor([0.2, 0.8], dtype=DTYPE) # The requested terminal law. +STATIONARY = torch.tensor([0.6, 0.4], dtype=DTYPE) # Proportional to (1.5, 1), the rows of Q. +GRID = 1001 # Points in the Simpson rule for the time integral of the rate error. + + +# %% 02. Relative entropy between two jump processes +# +# Collecting the one-step relative entropies and letting the step go to zero leaves a time integral +# of the Bregman divergence of u log u between the two rates, which is nonnegative and vanishes +# only when they agree. The divergence is not symmetric in the ratio. Section 7.21.1. + + +def rate_error(controlled: torch.Tensor, reference: torch.Tensor) -> torch.Tensor: + """Q^theta log(Q^theta / Q) - Q^theta + Q, the Bregman divergence of u log u, entry by entry.""" + safe = controlled.clamp(min=1e-300) # A zero rate contributes the reference rate alone. + return controlled * (safe / reference).log() - controlled + reference + + +def scaling_cost(c: float) -> float: + """c log c - c + 1, the cost per unit of reference rate of scaling every rate by c.""" + return c * math.log(c) - c + 1.0 + + +# %% 03. The Schrodinger problem for a chain +# +# The endpoint law of a chain with a time-independent generator is diag(pi_0^ref) e^Q, and the two +# potentials are the Sinkhorn scalings of that matrix to the two marginals. Section 7.21.2. + + +def endpoint_law(generator: torch.Tensor, start: torch.Tensor) -> torch.Tensor: + """R_01 = diag(pi_0) e^Q, the reference probability of starting at x and ending at y.""" + return torch.diag(start) @ torch.matrix_exp(generator) # [S, S]. + + +def sinkhorn(matrix: torch.Tensor, start: torch.Tensor, end: torch.Tensor, sweeps: int = 400 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Scale a positive matrix to the two marginals, returning f, g and the plan f R_01 g.""" + f = torch.ones_like(start) + g = torch.ones_like(end) + for _ in range(sweeps): + g = end / (matrix.T @ f) # Column sums become exactly pi_1. + f = start / (matrix @ g) # Row sums become exactly pi_0. + return f, g, f[:, None] * matrix * g[None, :] + + +def relative_entropy(p: torch.Tensor, q: torch.Tensor) -> float: + """Sum over entries of p log(p / q), for two positive matrices or vectors.""" + mask = p > 0 + return float((p[mask] * (p[mask] / q[mask]).log()).sum()) + + +# %% 04. Rates reweighted by an h-transform +# +# The backward potential is g averaged over the reference from the current time to time 1, and each +# off-diagonal rate is multiplied by the ratio of that potential at the destination to its value at +# the current state. The diagonal is then set so that the row sums to zero. Section 7.22. + + +def backward_potential(generator: torch.Tensor, g: torch.Tensor, t: float) -> torch.Tensor: + """phi_t = e^{Q(1-t)} g, the reference average of g at the final state seen from time t.""" + return torch.matrix_exp(generator * (1.0 - t)) @ g # [S]. + + +def tilted_generator(generator: torch.Tensor, phi: torch.Tensor) -> torch.Tensor: + """Q*(x, y) = Q(x, y) phi(y) / phi(x) off the diagonal, with the diagonal closing each row.""" + tilted = generator * phi[None, :] / phi[:, None] # [S, S] every entry reweighted. + tilted.fill_diagonal_(0.0) # Discard the reweighted diagonal, which has no meaning. + return tilted - torch.diag(tilted.sum(dim=1)) # Each row now sums to zero. + + +def bridge_marginal(generator: torch.Tensor, start: torch.Tensor, f: torch.Tensor, + g: torch.Tensor, t: float) -> torch.Tensor: + """p*_t = psi_t phi_t, the forward potential carried under the reference times the backward one.""" + psi = (start * f) @ torch.matrix_exp(generator * t) # [S] the forward factor. + return psi * backward_potential(generator, g, t) # [S] their entrywise product. + + +# %% 05. Sampling the bridge with Euler steps +# +# Over a step the chain moves to a different state with probability dt Q*(x, y) and stays with the +# remaining probability, which is the Euler sampler of Chapter 4 applied to a generator that +# changes with time. Section 7.23.1. + + +def euler_marginal(generator: torch.Tensor, start: torch.Tensor, f: torch.Tensor, + g: torch.Tensor, steps: int) -> tuple[torch.Tensor, list[float]]: + """Propagate p <- p (I + dt Q*_t) and record the probability of leaving state 1 at each step.""" + dt = 1.0 / steps + p = start.clone() + leaving = [] + identity = torch.eye(len(start), dtype=DTYPE) + for step in range(steps): + tilted = tilted_generator(generator, backward_potential(generator, g, step * dt)) + leaving.append(float(tilted[0, 1] * dt)) # The jump probability out of state 1. + p = p @ (identity + dt * tilted) # One Euler step of the master equation. + return p, leaving + + +# %% 06. The same bridge run backwards +# +# Matching the probability flux on an edge in both directions fixes the reverse rate, which is the +# forward rate times the ratio of the marginal at the two ends of the edge. Section 7.23.2. + + +def reverse_rates(tilted: torch.Tensor, marginal: torch.Tensor) -> torch.Tensor: + """Q_hat(y, x) = Q(x, y) p_t(x) / p_t(y), the rates under which the flux on an edge matches.""" + return (tilted * marginal[:, None] / marginal[None, :]).T # [S, S], reversed in both indices. + + +# %% 07. The cost of the bridge, computed twice +# +# On the endpoints the cost is the relative entropy of the plan from the reference endpoint law. On +# the paths it is the time integral of the rate error, weighted by the marginal of the bridge at +# that time. The two agree. Section 7.23.3. + + +def path_cost(generator: torch.Tensor, start: torch.Tensor, f: torch.Tensor, g: torch.Tensor, + points: int = GRID) -> tuple[float, float]: + """Integrate the weighted rate error over [0, 1] by Simpson's rule. + + Returns the integral and the value of the integrand at t = 1/2, the one the chapter prints. + """ + off = ~torch.eye(len(generator), dtype=torch.bool) # [S, S] true away from the diagonal. + times = torch.linspace(0.0, 1.0, points, dtype=DTYPE) + values = [] + for t in times.tolist(): + tilted = tilted_generator(generator, backward_potential(generator, g, t)) + marginal = bridge_marginal(generator, start, f, g, t) + error = torch.zeros_like(tilted) + error[off] = rate_error(tilted[off], generator[off]) # Only the edges carry a cost. + values.append(float((marginal * error.sum(dim=1)).sum())) + weights = torch.ones(points, dtype=DTYPE) + weights[1:-1:2] = 4.0 # Simpson's rule on an even number of intervals. + weights[2:-1:2] = 2.0 + integral = float((weights * torch.tensor(values, dtype=DTYPE)).sum() / (3.0 * (points - 1))) + return integral, values[(points - 1) // 2] + + +# %% 08. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=10) # --steps is the coarsest number of Euler steps in the table. + args = parser.parse_args() + set_seed(args.seed) + loud = not args.quiet + coarse = max(args.steps, 2) + + if loud: + print("Scaling every reference rate by c costs c log c - c + 1 per unit of rate") + print(f" halving {scaling_cost(0.5):.4f}, doubling {scaling_cost(2.0):.4f}") + + R01 = endpoint_law(Q, PI_0) + f, g, plan = sinkhorn(R01, PI_0, PI_1) + phi_half = backward_potential(Q, g, 0.5) + tilted_half = tilted_generator(Q, phi_half) + tilted_zero = tilted_generator(Q, backward_potential(Q, g, 0.0)) + tilted_one = tilted_generator(Q, g) + marginal_half = bridge_marginal(Q, PI_0, f, g, 0.5) + marginal_one = bridge_marginal(Q, PI_0, f, g, 1.0) + reverse = reverse_rates(tilted_half, marginal_half) + if loud: + print(f"\ne^Q = {[[round(float(z), 4) for z in row] for row in torch.matrix_exp(Q)]}") + print(f"R_01 = {[[round(float(z), 4) for z in row] for row in R01]}, " + f"column sums {[round(float(z), 4) for z in R01.sum(0)]}") + print(f" f = {[round(float(z), 4) for z in f]}, g = {[round(float(z), 4) for z in g]}, " + f"g ratio {float(g[1] / g[0]):.3f}") + print(f" pi* = {[[round(float(z), 4) for z in row] for row in plan]}") + print(f" phi_1/2 = {[round(float(z), 4) for z in phi_half]}, " + f"ratio {float(phi_half[1] / phi_half[0]):.2f}") + print(f" Q*_1/2 = {[[round(float(z), 4) for z in row] for row in tilted_half]}") + print(f" rates out of the two states: at t = 0 " + f"{float(tilted_zero[0, 1]):.4f} and {float(tilted_zero[1, 0]):.4f}, " + f"at t = 1 {float(tilted_one[0, 1]):.4f} and {float(tilted_one[1, 0]):.4f}") + print(f" marginal at t = 1/2 {[round(float(z), 4) for z in marginal_half]}, " + f"at t = 1 {[round(float(z), 4) for z in marginal_one]}") + print(f" reverse rates at t = 1/2: {float(reverse[1, 0]):.4f} from 2 back to 1, " + f"{float(reverse[0, 1]):.4f} from 1 back to 2") + print(f" the flux on the edge: {float(marginal_half[0] * tilted_half[0, 1]):.4f} " + f"against {float(marginal_half[1] * reverse[1, 0]):.4f}") + + sampler = {} + for steps in (coarse, 2 * coarse, 4 * coarse, 10 * coarse): + terminal, leaving = euler_marginal(Q, PI_0, f, g, steps) + sampler[steps] = (terminal, float((terminal - PI_1).abs().sum() / 2), leaving) + if loud: + print("\nsteps terminal law error first and last jump probability") + for steps, (terminal, error, leaving) in sampler.items(): + print(f" {steps:<5d} ({float(terminal[0]):.4f}, {float(terminal[1]):.4f})" + f" {error:.4f} {leaving[0]:.4f}, {leaving[-1]:.4f}") + + endpoint_cost = relative_entropy(plan, R01) + integral, at_half = path_cost(Q, PI_0, f, g) + speeds = {} + for lam in (0.5, 1.0, 2.0, 4.0): + scaled_law = endpoint_law(lam * Q, PI_0) + _, _, scaled_plan = sinkhorn(scaled_law, PI_0, PI_1) + speeds[lam] = relative_entropy(scaled_plan, scaled_law) + limit = relative_entropy(PI_1, STATIONARY) + if loud: + print(f"\nKL(pi* || R_01) = {endpoint_cost:.4f}") + print(f" the integrand at t = 1/2 is {at_half:.4f}, and its integral is {integral:.4f}") + print(f" at lambda = 1/2, 1, 2, 4 the cost is " + f"{[round(speeds[k], 4) for k in (0.5, 1.0, 2.0, 4.0)]}") + print(f" the limit KL(pi_1 || (0.6, 0.4)) is {limit:.4f}") + + report = { + "cost of halving every rate": scaling_cost(0.5), + "cost of doubling every rate": scaling_cost(2.0), + "R_01 column sums": [round(float(z), 4) for z in R01.sum(dim=0)], + "f": [round(float(z), 4) for z in f], + "g": [round(float(z), 4) for z in g], + "pi*_11": float(plan[0, 0]), + "phi at t = 1/2": [round(float(z), 4) for z in phi_half], + "bridge rate out of state 1": float(tilted_half[0, 1]), + "bridge rate out of state 2": float(tilted_half[1, 0]), + "rate out of state 1 at t = 1": float(tilted_one[0, 1]), + "marginal at t = 1/2": [round(float(z), 4) for z in marginal_half], + "marginal at t = 1": [round(float(z), 4) for z in marginal_one], + "reverse rate, 2 back to 1": float(reverse[1, 0]), + "Euler error at the coarsest step": sampler[coarse][1], + "Euler error, step halved": sampler[2 * coarse][1], + "Euler error, step quartered": sampler[4 * coarse][1], + "Euler error at ten times the steps": sampler[10 * coarse][1], + "KL(pi* || R_01)": endpoint_cost, + "rate error integrand at t = 1/2": at_half, + "its integral over [0, 1]": integral, + "cost at lambda = 1/2": speeds[0.5], + "cost at lambda = 4": speeds[4.0], + "KL(pi_1 || the stationary law)": limit, + } + print_report("The bridge over a finite set, Chapter 7 Sections 7.21 to 7.23.3", report) + save_run(args.out, config=vars(args), report=report) # Written only with --out. + + +if __name__ == "__main__": + main() diff --git a/lecture_7/displacement.py b/lecture_7/displacement.py new file mode 100644 index 0000000000000000000000000000000000000000..3069526fae92c02d8653bebf467a13f92f7ba57d --- /dev/null +++ b/lecture_7/displacement.py @@ -0,0 +1,237 @@ +"""Transport as motion in time: the kinetic action, straight lines, and displacement interpolation. + +A plan records which source point ends where and nothing about the trajectory between them. The +Benamou-Brenier formulation replaces the plan by a density and a velocity field and measures the +transport by the energy the particles expend, and its minimizer moves every particle in a straight +line at constant speed with the endpoints paired by the optimal plan. + +Chapter 7, Sections 7.7 to 7.7.5 of the notes. Runs on the one-particle example of Section 7.7.1 +and on the Gaussian pair N(0,1) to N(3,4), both in closed form. Prints the action of a steady and a +dawdling particle, the two sides of the averaged bound, the interpolated density against the mixture +of the endpoints, and the velocity field with its action density. + + python lecture_7/displacement.py # 1024 times on the grid + python lecture_7/displacement.py --steps 20 # a coarse time grid, for a first look +""" + +# %% 01. Imports and configuration +from __future__ import annotations + +import math + +import torch + +from cis6270.runner import common_parser, print_report, save_run, set_seed + +DTYPE = torch.float64 # Closed forms throughout, so the printed digits are the chapter's. + +MEAN_0, STD_0 = 0.0, 1.0 # The source N(0, 1). +MEAN_1, STD_1 = 3.0, 2.0 # The target N(3, 4), whose standard deviation is 2. + + +# %% 02. One particle +# +# The squared displacement of a particle is at most the integral of its squared speed, by +# Cauchy-Schwarz against the constant function 1 on an interval of length 1, with equality when the +# velocity is constant in time. Section 7.7.1. + + +def action_of_a_path(positions: torch.Tensor, times: torch.Tensor) -> float: + """The integral of the squared speed along a trajectory sampled at the given times.""" + velocity = (positions[1:] - positions[:-1]) / (times[1:] - times[:-1]) # [n-1] piecewise speed. + return float((velocity ** 2 * (times[1:] - times[:-1])).sum()) # Integral of ||x_dot||^2. + + +def steady_particle(times: torch.Tensor) -> torch.Tensor: + """A particle covering the distance 2 at the constant speed 2, so its action is 4.""" + return 2.0 * times # x(t) = 2t. + + +def dawdling_particle(times: torch.Tensor) -> torch.Tensor: + """A particle that waits until t = 1/2 and then covers the same distance at speed 4.""" + return torch.clamp(4.0 * (times - 0.5), min=0.0) # Zero on [0, 1/2], then 4(t - 1/2). + + +# %% 03. Averaging the bound over the particles +# +# The joint law of the two endpoints of a particle is a coupling of mu and nu, so the action is at +# least the average squared displacement under that coupling, which is at least W_2^2. +# Section 7.7.2. + + +def average_squared_displacement(x: torch.Tensor, y: torch.Tensor, pi: torch.Tensor) -> float: + """The average of ||x_1 - x_0||^2 over endpoint pairs drawn from the coupling pi.""" + return float((pi * (x[:, None] - y[None, :]) ** 2).sum()) # Sum over routes of mass times cost. + + +# %% 04. Straight lines attain the bound +# +# Drawing the endpoint pair from the optimal coupling and moving along the segment between them +# makes the velocity constant, so Cauchy-Schwarz is tight for every particle, and the coupling is +# optimal, so the second inequality is tight as well. Section 7.7.3. + + +def straight_path(x0: torch.Tensor, x1: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + """x(t) = (1 - t) x_0 + t x_1, whose velocity is the constant x_1 - x_0.""" + return (1.0 - t) * x0 + t * x1 # Same shape as x0 and x1 broadcast against t. + + +# %% 05. Displacement interpolation +# +# Reading off the law of the particles at an intermediate time gives the displacement interpolation. +# Under the optimal map X_1 = 3 + 2 X_0 of the Gaussian example, X_t collects into 3t + (1 + t) X_0, +# so rho_t is Gaussian with mean 3t and standard deviation 1 + t. Section 7.7.4. + + +def displacement_law(t: float) -> tuple[float, float]: + """The mean and standard deviation of rho_t for the Gaussian pair, both linear in t.""" + return 3.0 * t, 1.0 + t # Mean slides from 0 to 3, standard deviation from 1 to 2. + + +def normal_density(x: torch.Tensor, mean: float, std: float) -> torch.Tensor: + """The density of N(mean, std^2) evaluated at every entry of x.""" + return torch.exp(-0.5 * ((x - mean) / std) ** 2) / (std * math.sqrt(2.0 * math.pi)) + + +def mixture_density(x: torch.Tensor) -> torch.Tensor: + """Half the source density plus half the target density, the other path between the two ends.""" + return 0.5 * normal_density(x, MEAN_0, STD_0) + 0.5 * normal_density(x, MEAN_1, STD_1) + + +def gaussian_w2(m0: float, s0: float, m1: float, s1: float) -> float: + """W_2 between two Gaussians on the line, the closed form of Section 7.6.""" + return math.sqrt((m0 - m1) ** 2 + (s0 - s1) ** 2) + + +# %% 06. The velocity field behind the displacement path +# +# Each particle obeys X_t = 3t + (1 + t) X_0 with the constant velocity 3 + X_0. Solving the +# trajectory for X_0 and substituting makes the velocity a function of the current position alone. +# Section 7.7.5. + + +def velocity_field(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + """v_t(x) = 3 + (x - 3t) / (1 + t), the mean speed plus a stretch about the mean.""" + return 3.0 + (x - 3.0 * t) / (1.0 + t) # Broadcasts over a grid of states and times. + + +def continuity_residual(x: float, t: float) -> float: + """The left side of d_t rho_t + d_x (rho_t v_t) = 0, computed by automatic differentiation.""" + xs = torch.tensor(x, dtype=DTYPE, requires_grad=True) + ts = torch.tensor(t, dtype=DTYPE, requires_grad=True) + rho = normal_density(xs, 3.0 * ts, 1.0 + ts) # rho_t written with the two linear schedules. + flux = rho * velocity_field(xs, ts) # The probability flux rho_t v_t. + d_rho_dt = torch.autograd.grad(rho, ts, retain_graph=True)[0] # Keep the graph for the flux. + d_flux_dx = torch.autograd.grad(flux, xs)[0] + return float((d_rho_dt + d_flux_dx).detach()) + + +def action_density(t: float, levels: int) -> float: + """The integral of rho_t(x) v_t(x)^2 over x, which is E[(3 + X_0)^2] at every time.""" + mean, std = displacement_law(t) + q = (torch.arange(levels, dtype=DTYPE) + 0.5) / levels # Quantile levels of rho_t. + x = mean + std * torch.special.ndtri(q) # [levels] a stratified sample of rho_t. + speed = velocity_field(x, torch.tensor(t, dtype=DTYPE)) # [levels]. + return float((speed ** 2).mean()) # The average of the squared speed under rho_t. + + +# %% 07. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=1024) # --steps is the number of times on the grid. + args = parser.parse_args() + set_seed(args.seed) + loud = not args.quiet + n = max(args.steps, 2) + times = torch.linspace(0.0, 1.0, n + 1, dtype=DTYPE) # [n+1] the time grid. + + steady = action_of_a_path(steady_particle(times), times) + dawdle = action_of_a_path(dawdling_particle(times), times) + if loud: + print(f"One particle from 0 to 2 on a grid of {n} steps") + print(f" constant speed 2: action {steady:.3f}, squared displacement {2.0 ** 2:.0f}") + print(f" waiting then speed 4: action {dawdle:.3f}, the same two endpoints") + + x = torch.tensor([0.0, 2.0], dtype=DTYPE) # Four points, mass 1/2 each. + y = torch.tensor([1.0, 3.0], dtype=DTYPE) + ordered = torch.tensor([[0.5, 0.0], [0.0, 0.5]], dtype=DTYPE) + crossed = torch.tensor([[0.0, 0.5], [0.5, 0.0]], dtype=DTYPE) + independent = torch.full((2, 2), 0.25, dtype=DTYPE) + w22_points = average_squared_displacement(x, y, ordered) # The optimal coupling on the line. + if loud: + print(f"\nAveraging over particles, four points with mass 1/2") + for name, pi in (("ordered", ordered), ("independent", independent), ("crossed", crossed)): + print(f" {name:<12s} E||x_1 - x_0||^2 = " + f"{average_squared_displacement(x, y, pi):.1f} >= W_2^2 = {w22_points:.1f}") + + # Straight lines under the optimal coupling, with the endpoints of the Gaussian example. + q = (torch.arange(n, dtype=DTYPE) + 0.5) / n + x0 = MEAN_0 + STD_0 * torch.special.ndtri(q) # [n] a stratified sample of N(0, 1). + x1 = MEAN_1 + (STD_1 / STD_0) * (x0 - MEAN_0) # The Brenier map 3 + 2x, the optimal pairing. + paths = straight_path(x0[None, :], x1[None, :], times[:, None]) # [n+1, n] every trajectory. + speeds = (paths[1:] - paths[:-1]) / (times[1:] - times[:-1])[:, None] # [n, n] piecewise speed. + action = float((speeds ** 2 * (times[1:] - times[:-1])[:, None]).sum(dim=0).mean()) + squared_displacement = float(((x1 - x0) ** 2).mean()) + if loud: + print("\nStraight lines under the optimal coupling") + print(f" action averaged over the particles {action:.4f} " + f"average squared displacement {squared_displacement:.4f}") + + mean_half, std_half = displacement_law(0.5) + grid = torch.linspace(-6.0, 9.0, 15001, dtype=DTYPE) # A dense grid for the two densities. + mixture = mixture_density(grid) + peak = float(grid[int(mixture.argmax())]) + interpolated = normal_density(grid, mean_half, std_half) + separation = float(0.5 * (interpolated - mixture).abs().sum() * (grid[1] - grid[0])) + half_distance = gaussian_w2(MEAN_0, STD_0, mean_half, std_half) + whole_distance = gaussian_w2(MEAN_0, STD_0, MEAN_1, STD_1) + if loud: + print(f"\nAt t = 1/2 the displacement interpolation is N({mean_half:.1f}, " + f"{std_half ** 2:.2f}) with standard deviation {std_half:.1f}") + print(f" the mixture of the two ends peaks at x = {peak:.2f} and is not a Gaussian") + print(f" total variation between the two candidate densities {separation:.3f}") + print(f" W_2(rho_0, rho_1/2) = {half_distance:.3f} = " + f"{half_distance / whole_distance:.1f} x W_2(mu, nu) = {whole_distance:.3f}") + + mean_speed = float(velocity_field(torch.tensor(3.0 * 0.5, dtype=DTYPE), + torch.tensor(0.5, dtype=DTYPE))) # A particle at the mean. + above_speed = float(velocity_field(torch.tensor(3.0 * 0.5 + 1.5, dtype=DTYPE), + torch.tensor(0.5, dtype=DTYPE))) # One standard deviation up. + residual = max(abs(continuity_residual(xx, tt)) + for xx in (-1.0, 0.0, 1.5, 4.0) for tt in (0.1, 0.5, 0.9)) + densities = [action_density(float(t), 2048) for t in torch.linspace(0.0, 1.0, 5, dtype=DTYPE)] + landing = x0 + 1.0 * velocity_field(x0, torch.zeros((), dtype=DTYPE)) # One Euler step of size 1. + one_step = float((landing - x1).abs().max()) # The step lands on the Brenier map 3 + 2 x_0. + if loud: + print(f"\nv_t(x) = 3 + (x - 3t) / (1 + t)") + print(f" at the mean, speed {mean_speed:.0f}; one standard deviation above, " + f"speed {above_speed:.0f}") + print(f" action density at five times {[round(d, 4) for d in densities]}") + print(f" largest continuity residual on the test states {residual:.2e}") + + report = { + "action of the steady particle": steady, + "action of the dawdling particle": dawdle, + "squared displacement of both": 4.0, + "E||x1-x0||^2, ordered": average_squared_displacement(x, y, ordered), + "E||x1-x0||^2, independent": average_squared_displacement(x, y, independent), + "E||x1-x0||^2, crossed": average_squared_displacement(x, y, crossed), + "action of the straight paths": action, + "W_2^2 of the Gaussian pair": whole_distance ** 2, + "rho_1/2 mean, standard deviation": f"{mean_half:.1f}, {std_half:.1f}", + "peak of the mixture": peak, + "total variation of the two paths": separation, + "W_2(rho_0, rho_1/2)": half_distance, + "W_2(mu, nu)": whole_distance, + "speed at the mean": mean_speed, + "speed one deviation above": above_speed, + "action density, every time": max(densities), + "continuity equation residual": residual, + "one Euler step of size 1 minus T": one_step, + } + print_report("Transport as motion, Chapter 7 Sections 7.7 to 7.7.5", report) + save_run(args.out, config=vars(args), report=report) # Written only with --out. + + +if __name__ == "__main__": + main() diff --git a/lecture_7/dsb.py b/lecture_7/dsb.py new file mode 100644 index 0000000000000000000000000000000000000000..e3bfa0bf59dfa050ccef4eab8a3fc45c8e6da163 --- /dev/null +++ b/lecture_7/dsb.py @@ -0,0 +1,204 @@ +"""Diffusion Schrodinger Bridge: iterative proportional fitting run on path laws. + +Each half-step of proportional fitting is a relative-entropy projection onto the set of path laws +with one marginal correct, and the chain rule solves it in closed form by keeping the conditional +law of the previous iterate and resetting that one marginal. The conditional law is the current +process run backward in time, so a half-step amounts to reversing the current process and +restarting it from the requested law at the other end. + +Chapter 7, Section 7.19.1 of the notes. Runs the half-steps exactly on the two-state chain with +transition matrix A, source law (0.5, 0.5) and requested terminal law (0.2, 0.8), and checks the +diffusion time reversal in closed form on a Gaussian. Prints the backward matrices, the law after +each half-step, the forward matrices of the first round and the error at the free end. + + python lecture_7/dsb.py # 4 rounds, the first one worked by hand + python lecture_7/dsb.py --steps 1 # the single round the chapter writes out +""" + +# %% 01. Imports and configuration +from __future__ import annotations + +import math + +import torch + +from cis6270.runner import common_parser, print_report, save_run, set_seed + +DTYPE = torch.float64 # The chapter quotes six figures of the backward and forward matrices. + +A = torch.tensor([[0.7, 0.3], [0.4, 0.6]], dtype=DTYPE) # The reference transition matrix. +PI_0 = torch.tensor([0.5, 0.5], dtype=DTYPE) # The source law. +PI_1 = torch.tensor([0.2, 0.8], dtype=DTYPE) # The requested terminal law. +TILTED_0 = torch.tensor([[0.599160, 0.400840], [0.299266, 0.700734]], dtype=DTYPE) # A^h_0. +TILTED_1 = torch.tensor([[0.307297, 0.692703], [0.112490, 0.887510]], dtype=DTYPE) # A^h_1. + + +# %% 02. A path law as a joint table +# +# With two states and two steps a path law is a table over the triple (x_0, x_1, x_2), so every +# projection below is an exact operation on eight numbers. Section 7.12.1. + + +def reference_path_law(start: torch.Tensor, matrix: torch.Tensor, steps: int) -> torch.Tensor: + """R(x_{0:N}) as a tensor with one axis per time, from the starting law and one matrix per step.""" + law = start.clone() # [K] the law at time 0. + for _ in range(steps): + law = law.unsqueeze(-1) * matrix # Multiply by the transition matrix on a new axis. + return law # [K] * (N + 1). + + +def marginal(law: torch.Tensor, axis: int) -> torch.Tensor: + """The law of the state at one time, summed over every other time.""" + other = [d for d in range(law.dim()) if d != axis] + return law.sum(dim=other) # [K]. + + +def one_step_conditional(law: torch.Tensor, step: int) -> torch.Tensor: + """P(x_{k+1} | x_k), the forward transition matrix the path law carries at that step.""" + pair = law.sum(dim=[d for d in range(law.dim()) if d not in (step, step + 1)]) # [K, K]. + return pair / pair.sum(dim=1, keepdim=True) # Normalize each row. + + +def reverse_step(law: torch.Tensor, step: int) -> torch.Tensor: + """P(x_k | x_{k+1}), Bayes' rule on one step, which is the reversal a network has to learn.""" + pair = law.sum(dim=[d for d in range(law.dim()) if d not in (step, step + 1)]) # [K, K]. + return (pair / pair.sum(dim=0, keepdim=True)).T # [K, K] indexed by the later state first. + + +# %% 03. The two half-steps +# +# The odd iterate keeps the conditional law of the previous iterate given the final state and +# replaces the marginal at time 1, and the even iterate does the same at time 0. Section 7.19.1. + + +def fix_terminal(law: torch.Tensor, end: torch.Tensor) -> torch.Tensor: + """P'(x_{0:N}) = pi_1(x_N) P(x_{0:N-1} | x_N), the projection onto the laws ending at pi_1.""" + last = marginal(law, law.dim() - 1) # [K] the terminal marginal of the current law. + return law * (end / last) # Reweight every path by the ratio at its final state. + + +def fix_initial(law: torch.Tensor, start: torch.Tensor) -> torch.Tensor: + """P''(x_{0:N}) = pi_0(x_0) P'(x_{1:N} | x_0), the projection onto the laws starting at pi_0.""" + first = marginal(law, 0) # [K] the initial marginal of the current law. + shape = [-1] + [1] * (law.dim() - 1) # Broadcast the ratio along the first axis. + return law * (start / first).reshape(shape) + + +def proportional_fitting(law: torch.Tensor, start: torch.Tensor, end: torch.Tensor, + rounds: int) -> tuple[torch.Tensor, list[float], list[torch.Tensor]]: + """Alternate the two half-steps, returning the law, the terminal error and the law each round.""" + errors: list[float] = [] + iterates: list[torch.Tensor] = [] + for _ in range(rounds): + law = fix_terminal(law, end) # One half-step: the terminal marginal is now exact. + law = fix_initial(law, start) # The other: the initial marginal is now exact. + errors.append(float((marginal(law, law.dim() - 1) - end).abs().sum())) + iterates.append(law.clone()) + return law, errors, iterates + + +# %% 04. The reversal that a network learns +# +# On a finite state space the reversal is Bayes' rule on a matrix. On R^d it is the time reversal of +# a diffusion, whose drift is minus the forward drift plus epsilon times the score of the current +# marginal, and the score of the current iterate has no closed form. Section 7.19.1. + + +def reverse_drift(x: torch.Tensor, t: torch.Tensor, epsilon: float, variance_0: float + ) -> torch.Tensor: + """beta(x, t) = -b + epsilon grad log p_t for a driftless diffusion started from N(0, v_0).""" + variance = variance_0 + epsilon * t # The marginal of sqrt(epsilon) dB started from N(0, v_0). + return epsilon * (-x / variance) # The score of a centered Gaussian is -x over its variance. + + +def reversal_check(epsilon: float, variance_0: float, steps: int, paths: int, + generator: torch.Generator) -> tuple[float, float]: + """Run the reversal backward from p_1 and compare its variance at time 0 with v_0.""" + dt = 1.0 / steps + variance_1 = variance_0 + epsilon * 1.0 # The forward marginal at time 1. + x = torch.randn(paths, dtype=DTYPE, generator=generator) * math.sqrt(variance_1) + for step in range(steps): + t = torch.tensor(1.0 - step * dt, dtype=DTYPE) # Time runs down from 1 to 0. + noise = torch.randn(paths, dtype=DTYPE, generator=generator) * math.sqrt(dt) + x = x + reverse_drift(x, t, epsilon, variance_0) * dt + math.sqrt(epsilon) * noise + return float(x.var()), variance_0 + + +# %% 05. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=4) # --steps is the number of proportional fitting rounds. + args = parser.parse_args() + set_seed(args.seed) + loud = not args.quiet + rounds = max(args.steps, 1) + generator = torch.Generator().manual_seed(args.seed) + + reference = reference_path_law(PI_0, A, 2) # [2, 2, 2] the law P^0 = R. + marginals = [marginal(reference, k) for k in range(3)] + B2 = reverse_step(reference, 1) # P(x_1 | x_2), the reversal of the last step. + B1 = reverse_step(reference, 0) # P(x_0 | x_1), the reversal of the earlier step. + if loud: + print("The reference started from pi_0, read at the three time points") + print(f" {[[round(float(z), 3) for z in m] for m in marginals]}") + print(f" B_2 = {[[round(float(z), 6) for z in row] for row in B2]}") + print(f" B_1(1,1) = {float(B1[0, 0]):.6f}") + + half = fix_terminal(reference, PI_1) # The first half-step, which resets the terminal marginal. + backward_laws = [marginal(half, 2), marginal(half, 1), marginal(half, 0)] + drift_at_zero = float((backward_laws[2] - PI_0).abs().sum()) + full = fix_initial(half, PI_0) # The second half-step, which resets the initial marginal. + F0 = one_step_conditional(full, 0) + F1 = one_step_conditional(full, 1) + terminal = marginal(full, 2) + terminal_error = float((terminal - PI_1).abs().sum()) + tilt_gap = max(float((F0 - TILTED_0).abs().max()), float((F1 - TILTED_1).abs().max())) + if loud: + print("\nThe backward chain started from the requested terminal law") + written = ["(" + ", ".join(f"{float(z):.6f}" for z in m) + ")" for m in backward_laws] + print(f" {' -> '.join(written)}") + print(f" total absolute error at step 0: {drift_at_zero:.4f}") + print(f"\nThe forward matrices of the second half-step") + print(f" F_0 = {[[round(float(z), 6) for z in row] for row in F0]}") + print(f" F_1 = {[[round(float(z), 6) for z in row] for row in F1]}") + print(f" terminal law {[round(float(z), 6) for z in terminal]}, " + f"off by {terminal_error:.4f}") + print(f" largest gap against the tilted matrices A^h_0 and A^h_1: {tilt_gap:.4f}") + + _, errors, _ = proportional_fitting(reference, PI_0, PI_1, rounds) + ratios = [errors[k] / errors[k + 1] for k in range(len(errors) - 1) if errors[k + 1] > 0] + if loud: + print(f"\nterminal error by round {['%.3e' % e for e in errors]}") + print(f" ratios {[round(r, 1) for r in ratios]}") + + simulated, exact = reversal_check(1.0, 1.0, 200, 8192, generator) + if loud: + print(f"\nThe diffusion reversal on a Gaussian, run backward from p_1") + print(f" variance at time 0: simulated {simulated:.3f}, exact {exact:.3f}") + + report = { + "reference marginal at step 1": [round(float(z), 3) for z in marginals[1]], + "reference marginal at step 2": [round(float(z), 3) for z in marginals[2]], + "B_2(1,1)": float(B2[0, 0]), + "B_2(2,1)": float(B2[1, 0]), + "B_1(1,1)": float(B1[0, 0]), + "backward law after one step": [round(float(z), 6) for z in backward_laws[1]], + "backward law at step 0": [round(float(z), 6) for z in backward_laws[2]], + "error at the free end, half-step": drift_at_zero, + "F_0(1,1)": float(F0[0, 0]), + "F_1(1,2)": float(F1[0, 1]), + "terminal law after one round": [round(float(z), 6) for z in terminal], + "terminal error after one round": terminal_error, + "gap against the tilted matrices": tilt_gap, + "terminal error, last round": errors[-1], + "error ratio per round": ratios[0] if ratios else float("nan"), + "rounds run": rounds, + "reversal variance, simulated": simulated, + "reversal variance, exact": exact, + } + print_report("Diffusion Schrodinger Bridge, Chapter 7 Section 7.19.1", report) + save_run(args.out, config=vars(args), report=report) # Written only with --out. + + +if __name__ == "__main__": + main() diff --git a/lecture_7/dsbm.py b/lecture_7/dsbm.py new file mode 100644 index 0000000000000000000000000000000000000000..2bd4cb207fd749225f87200b1fa55909e3addd61 --- /dev/null +++ b/lecture_7/dsbm.py @@ -0,0 +1,207 @@ +"""Diffusion Schrodinger Bridge Matching: the Markov projection and the reciprocal projection. + +A mixture of bridges draws a pair of ends from a coupling and runs the reference bridge between +them, so its drift depends on the endpoint and the law is in general not Markov. Averaging that +drift over the endpoints of all bridges through a state removes the extra information and leaves a +diffusion with the same marginals, and replacing the bridges of a Markov law by reference bridges +keeps its endpoint coupling. Alternating the two projections holds both marginals fixed throughout +and improves only the pairing. + +Chapter 7, Section 7.19.2 of the notes. Runs the two projections exactly on the two-state chain, +started from the independent coupling of (0.5, 0.5) and (0.2, 0.8). Prints the averaged drift at one +point, the two projected matrices of the first round, and the gap against the bridge coupling. + + python lecture_7/dsbm.py # 9 rounds, which reach 1e-9 on this problem + python lecture_7/dsbm.py --steps 1 # the single round the chapter writes out +""" + +# %% 01. Imports and configuration +from __future__ import annotations + +import torch + +from cis6270.runner import common_parser, print_report, save_run, set_seed + +DTYPE = torch.float64 # The chapter quotes six figures of the two projected matrices. + +A = torch.tensor([[0.7, 0.3], [0.4, 0.6]], dtype=DTYPE) # The reference transition matrix. +PI_0 = torch.tensor([0.5, 0.5], dtype=DTYPE) # The source law. +PI_1 = torch.tensor([0.2, 0.8], dtype=DTYPE) # The requested terminal law. + + +# %% 02. The averaged bridge drift +# +# Two bridges passing through the same state at the same time carry different drifts, and the +# Markov drift at that state is their average weighted by how often each one passes through. +# A least-squares fit of a single number against the two targets lands on the same value. +# Section 7.19.2. + + +def bridge_drift(x1: float, x: float, t: float) -> float: + """(x_1 - x) / (1 - t), the drift of the Brownian bridge ending at x_1.""" + return (x1 - x) / (1.0 - t) + + +def markov_drift(endpoints: torch.Tensor, weights: torch.Tensor, x: float, t: float) -> float: + """v*(x, t) = E[(X_1 - x) / (1 - t) | X_t = x], the average over the bridges through x.""" + drifts = (endpoints - x) / (1.0 - t) # [n] one conditional drift per endpoint. + return float((weights * drifts).sum()) # The posterior average of the conditional drifts. + + +def least_squares_derivative(a: float, endpoints: torch.Tensor, weights: torch.Tensor, + x: float, t: float) -> float: + """The derivative in a of the average squared error of a single prediction against the drifts.""" + drifts = (endpoints - x) / (1.0 - t) # [n]. + return float(2.0 * (weights * (a - drifts)).sum()) # Vanishes at the weighted average. + + +# %% 03. The reciprocal projection +# +# A member of the reciprocal class is an endpoint coupling filled in with reference bridges. The +# projection of a Markov law keeps that law's endpoint coupling and redraws the interior from the +# reference. Section 7.19.2. + + +def reference_bridge_table(matrix: torch.Tensor) -> torch.Tensor: + """R(x_1 | x_0, x_2) = A(x_0, x_1) A(x_1, x_2) / A^2(x_0, x_2) for every triple. [K, K, K].""" + two_step = matrix @ matrix # [K, K] the reference over both steps. + joint = matrix[:, :, None] * matrix[None, :, :] # [x_0, x_1, x_2] the unnormalized weight. + return joint / two_step[:, None, :] # Normalize over the middle state. + + +def reciprocal_projection(coupling: torch.Tensor, matrix: torch.Tensor) -> torch.Tensor: + """proj_R(M) = M_{0,1} R_{.|0,1}, the coupling with the reference bridges put back in.""" + return coupling[:, None, :] * reference_bridge_table(matrix) # [K, K, K] a full path law. + + +# %% 04. The Markov projection +# +# Retaining the one-step conditionals of a path law and multiplying them into a chain gives the +# Markov law with the same marginals, and the averaging in each conditional is the discrete form of +# the drift average. Section 7.19.2. + + +def one_step_conditional(law: torch.Tensor, step: int) -> torch.Tensor: + """P(x_{k+1} | x_k) read off a path law, which is what the network fits at that transition.""" + pair = law.sum(dim=[d for d in range(law.dim()) if d not in (step, step + 1)]) # [K, K]. + return pair / pair.sum(dim=1, keepdim=True) + + +def markov_projection(law: torch.Tensor) -> tuple[list[torch.Tensor], torch.Tensor]: + """The chain of one-step conditionals, and the endpoint coupling it produces from pi_0.""" + conditionals = [one_step_conditional(law, k) for k in range(law.dim() - 1)] + joint = torch.diag(marginal(law, 0)) # [K, K] the initial law on the diagonal. + for matrix in conditionals: + joint = joint @ matrix # Multiply the conditionals into a chain. + return conditionals, joint + + +def marginal(law: torch.Tensor, axis: int) -> torch.Tensor: + """The law of the state at one time, summed over every other time.""" + return law.sum(dim=[d for d in range(law.dim()) if d != axis]) + + +# %% 05. Alternating the two projections +# +# Under a Markov projection every marginal is unchanged and under a reciprocal projection the +# endpoint coupling is unchanged, so both pi_0 and pi_1 stay exact once the starting law has them, +# and the iterations improve only the coupling. Section 7.19.2. + + +def imf_rounds(coupling: torch.Tensor, matrix: torch.Tensor, rounds: int, + target: torch.Tensor) -> tuple[torch.Tensor, list[float], list]: + """Run iterative Markovian fitting, returning the coupling, the gaps and the first round.""" + gaps: list[float] = [] + first: list = [] + for index in range(rounds): + law = reciprocal_projection(coupling, matrix) # Reference bridges between the current ends. + conditionals, coupling = markov_projection(law) # The Markov law and its endpoint coupling. + gaps.append(float((coupling - target).abs().sum())) + if index == 0: + first = conditionals + return coupling, gaps, first + + +def bridge_coupling(matrix: torch.Tensor, start: torch.Tensor, end: torch.Tensor, + rounds: int = 200) -> torch.Tensor: + """The exact bridge coupling, by scaling diag(pi_0) A^2 to the two marginals.""" + R01 = torch.diag(start) @ torch.linalg.matrix_power(matrix, 2) + f = torch.ones_like(start) + for _ in range(rounds): + g = end / (R01.T @ f) + f = start / (R01 @ g) + return f[:, None] * R01 * g[None, :] + + +# %% 06. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=9) # --steps is the number of Markovian fitting rounds. + args = parser.parse_args() + set_seed(args.seed) + loud = not args.quiet + rounds = max(args.steps, 1) + + endpoints = torch.tensor([3.0, 0.0], dtype=DTYPE) # Two bridges through x = 1 at t = 0.5. + weights = torch.tensor([0.5, 0.5], dtype=DTYPE) # They pass with equal probability. + average = markov_drift(endpoints, weights, 1.0, 0.5) + if loud: + print("Two bridges through x = 1 at t = 0.5, ending at 3 and at 0") + print(f" their drifts {bridge_drift(3.0, 1.0, 0.5):.0f} and " + f"{bridge_drift(0.0, 1.0, 0.5):.0f}, average {average:.0f}") + print(f" derivative of the squared error at a = {average:.0f}: " + f"{least_squares_derivative(average, endpoints, weights, 1.0, 0.5):.0f}") + print(f" derivative at a = 0: " + f"{least_squares_derivative(0.0, endpoints, weights, 1.0, 0.5):+.0f}") + + independent = PI_0[:, None] * PI_1[None, :] # [[0.1, 0.4], [0.1, 0.4]]. + target = bridge_coupling(A, PI_0, PI_1) + bridges = reference_bridge_table(A) + starting_gap = float((independent - target).abs().sum()) + coupling, gaps, first = imf_rounds(independent, A, rounds, target) + projected = torch.diag(PI_0) @ first[0] @ first[1] + ratios = [gaps[k] / gaps[k + 1] for k in range(len(gaps) - 1) if gaps[k + 1] > 0] + if loud: + print(f"\nStarting from the independent coupling " + f"{[[round(float(z), 1) for z in row] for row in independent]}") + print(f" the middle state is 1 with probability {float(bridges[0, 0, 0]):.6f} " + f"from 1 to 1, and {float(bridges[0, 0, 1]):.6f} from 1 to 2") + print(f" M_0 = {[[round(float(z), 6) for z in row] for row in first[0]]}") + print(f" M_1 = {[[round(float(z), 6) for z in row] for row in first[1]]}") + print(f" M_0(1,1) = 0.2 x {float(bridges[0, 0, 0]):.6f} + 0.8 x " + f"{float(bridges[0, 0, 1]):.6f} = {float(first[0][0, 0]):.6f}") + print(f" diag(pi_0) M_0 M_1 = " + f"{[[round(float(z), 6) for z in row] for row in projected]}") + print(f" row sums {[round(float(z), 3) for z in projected.sum(1)]}, " + f"column sums {[round(float(z), 3) for z in projected.sum(0)]}") + print(f"\ngap against pi*: {starting_gap:.4f} at the start, " + f"then {['%.2e' % g for g in gaps[:5]]}") + print(f" ratios {[round(r, 2) for r in ratios[:4]]}") + + reached = next((k + 1 for k, g in enumerate(gaps) if g < 1e-9), None) + report = { + "drift of the bridge ending at 3": bridge_drift(3.0, 1.0, 0.5), + "drift of the bridge ending at 0": bridge_drift(0.0, 1.0, 0.5), + "averaged Markov drift": average, + "least squares derivative at a = 1": least_squares_derivative( + average, endpoints, weights, 1.0, 0.5), + "R(mid = 1 | 1, 1)": float(bridges[0, 0, 0]), + "R(mid = 1 | 1, 2)": float(bridges[0, 0, 1]), + "M_0(1,1)": float(first[0][0, 0]), + "M_1(1,2)": float(first[1][0, 1]), + "endpoint coupling, entry 11": float(projected[0, 0]), + "endpoint coupling, entry 12": float(projected[0, 1]), + "row sums after the projection": [round(float(z), 6) for z in projected.sum(dim=1)], + "column sums after the projection": [round(float(z), 6) for z in projected.sum(dim=0)], + "gap, independent coupling": starting_gap, + "gap after one round": gaps[0], + "gap after the last round": gaps[-1], + "gap ratio per round": ratios[-1] if ratios else float("nan"), + "rounds to 1e-9": reached if reached else rounds, + } + print_report("Diffusion Schrodinger Bridge Matching, Chapter 7 Section 7.19.2", report) + save_run(args.out, config=vars(args), report=report) # Written only with --out. + + +if __name__ == "__main__": + main() diff --git a/lecture_7/dynamic_bridge.py b/lecture_7/dynamic_bridge.py new file mode 100644 index 0000000000000000000000000000000000000000..ad2be000c14fd74e32dc837b34224c5d25ed087e --- /dev/null +++ b/lecture_7/dynamic_bridge.py @@ -0,0 +1,257 @@ +"""A diffusion reference: Girsanov, the control problem, Hopf-Cole, and the Brownian bridge. + +Changing the drift of a diffusion and keeping its noise is the only change under which the two path +laws stay comparable, and over one short step the two increments differ only in their mean. Summing +those steps gives Girsanov's formula, averaging it under the controlled law turns relative entropy +on paths into an integral of a squared drift, and the Hopf-Cole change of variables turns the +resulting Hamilton-Jacobi-Bellman equation into the linear backward equation of the reference. + +Chapter 7, Sections 7.16 to 7.18 of the notes. Runs on a driftless diffusion in one dimension. +Prints the two computations of the relative entropy of a constant control, the residual of the +backward equation and of the Hamilton-Jacobi-Bellman equation, and the Brownian bridge between two +fixed points. + + python lecture_7/dynamic_bridge.py # 200 Euler-Maruyama steps + python lecture_7/dynamic_bridge.py --steps 20 # a coarse time grid, for a first look +""" + +# %% 01. Imports and configuration +from __future__ import annotations + +import math + +import torch + +from cis6270.runner import common_parser, print_report, save_run, set_seed + +DTYPE = torch.float64 # Closed forms and a Monte Carlo check, both in double precision. + +EPSILON = 1.0 # The reference noise variance per unit time. +START, END = 0.0, 2.0 # The two points the bridge is pinned at. +PATHS = 8192 # Trajectories in the Monte Carlo checks. + + +# %% 02. One short step +# +# Over a step of length dt both increments are Gaussian with the covariance epsilon dt, and they +# differ only in the mean, by u dt. The logarithm of the ratio of the two densities is the inner +# product of w, the increment beyond the reference drift, with u over epsilon, minus the squared +# norm of u times dt over 2 epsilon. Section 7.16.1. + + +def log_density_ratio_step(w: torch.Tensor, u: torch.Tensor, dt: float, + epsilon: float = EPSILON) -> torch.Tensor: + """log p^u / p^R over one step, with w = Delta X - b dt the part beyond the reference drift.""" + return (w * u).sum(dim=-1) / epsilon - (u ** 2).sum(dim=-1) * dt / (2.0 * epsilon) + + +# %% 03. Summing over the steps, and the average under the controlled law +# +# Adding the per-step contributions and letting dt go to zero gives Girsanov's formula, a martingale +# term built from the noise and a drift term built from the size of the control. Averaging under +# P^u leaves the two energy terms with coefficients 1/epsilon and -1/(2 epsilon). +# Sections 7.16.2 and 7.16.3. + + +def simulate_controlled(u: float, steps: int, paths: int, epsilon: float = EPSILON, + generator: torch.Generator | None = None) -> tuple[torch.Tensor, torch.Tensor]: + """Euler-Maruyama for dX = u dt + sqrt(epsilon) dB, returning the endpoints and the log ratio. + + The log ratio is accumulated step by step from the same Brownian increments the path uses, which + is how Section 7.27.3 accumulates it during training. + """ + dt = 1.0 / steps + x = torch.zeros(paths, 1, dtype=DTYPE) # [P, 1] every path starts at the origin. + log_ratio = torch.zeros(paths, dtype=DTYPE) # [P] the Girsanov integral along each path. + control = torch.full((paths, 1), u, dtype=DTYPE) # [P, 1] a control constant in space and time. + for _ in range(steps): + noise = torch.randn(paths, 1, dtype=DTYPE, generator=generator) * math.sqrt(dt) # dB. + w = control * dt + math.sqrt(epsilon) * noise # Delta X with no reference drift. + log_ratio = log_ratio + log_density_ratio_step(w, control, dt, epsilon) + x = x + w # The controlled state after this step. + return x.squeeze(1), log_ratio + + +def control_energy_entropy(u: float, epsilon: float = EPSILON) -> float: + """KL(P^u || R) = (1 / 2 epsilon) E int_0^1 ||u||^2 dt, for a control constant in time.""" + return u ** 2 / (2.0 * epsilon) # The integral of a constant over the unit interval. + + +def gaussian_entropy_gap(shift: float, variance: float) -> float: + """KL between two Gaussians with a common variance, the squared difference of the means over 2v.""" + return shift ** 2 / (2.0 * variance) + + +# %% 04. The value function and the optimal control +# +# Bellman's principle and Ito's formula give the Hamilton-Jacobi-Bellman equation, in which only two +# terms contain the control. Completing the square in those two isolates u* = -grad V, and +# substituting it back leaves a closed equation for V with a squared gradient in it. +# Sections 7.17 to 7.17.2. + + +def optimal_control(value, x: float, t: float) -> float: + """u*(x, t) = -grad V_t(x), the control pointing downhill on the cost-to-go.""" + xs = torch.tensor(x, dtype=DTYPE, requires_grad=True) + v = value(xs, torch.tensor(t, dtype=DTYPE)) + return float(-torch.autograd.grad(v, xs)[0]) + + +def hjb_residual(value, x: float, t: float, epsilon: float = EPSILON) -> float: + """d_t V + (epsilon / 2) grad^2 V - (1/2) ||grad V||^2, with no reference drift.""" + xs = torch.tensor(x, dtype=DTYPE, requires_grad=True) + ts = torch.tensor(t, dtype=DTYPE, requires_grad=True) + v = value(xs, ts) + dv_dt = torch.autograd.grad(v, ts, create_graph=True)[0] + dv_dx = torch.autograd.grad(v, xs, create_graph=True)[0] + d2v_dx2 = torch.autograd.grad(dv_dx, xs, create_graph=True)[0] + return float((dv_dt + 0.5 * epsilon * d2v_dx2 - 0.5 * dv_dx ** 2).detach()) + + +# %% 05. Removing the squared gradient +# +# Writing the value function as minus epsilon times a logarithm cancels the squared gradient against +# the extra term the Laplacian produces, and what is left is the Kolmogorov backward equation of the +# reference, which is linear. In the new variable the control is epsilon times the gradient of the +# logarithm of phi. Sections 7.17.3 and 7.17.4. + + +def hopf_cole(potential, epsilon: float = EPSILON): + """V_t = -epsilon log phi_t, the change of variables that linearizes the equation.""" + return lambda x, t: -epsilon * torch.log(potential(x, t)) + + +def backward_residual(potential, x: float, t: float, epsilon: float = EPSILON) -> float: + """d_t phi + (epsilon / 2) grad^2 phi, which vanishes for a reference transition density.""" + xs = torch.tensor(x, dtype=DTYPE, requires_grad=True) + ts = torch.tensor(t, dtype=DTYPE, requires_grad=True) + phi = potential(xs, ts) + dphi_dt = torch.autograd.grad(phi, ts, create_graph=True)[0] + dphi_dx = torch.autograd.grad(phi, xs, create_graph=True)[0] + d2phi_dx2 = torch.autograd.grad(dphi_dx, xs, create_graph=True)[0] + return float((dphi_dt + 0.5 * epsilon * d2phi_dx2).detach()) + + +# %% 06. The Brownian bridge +# +# Conditioning on a single final point y makes the potential the reference density of arriving at y +# from the current state, and epsilon times the gradient of its logarithm is the remaining +# displacement over the remaining time. The noise level cancels from the drift and survives in the +# variance. Section 7.18. + + +def bridge_potential(z: torch.Tensor, t: torch.Tensor, y: float = END, + epsilon: float = EPSILON) -> torch.Tensor: + """h(z, t) = R_{1|t}(y | z), the full Gaussian density with variance epsilon (1 - t).""" + remaining = 1.0 - t # The time left, which is the variance divided by epsilon. + return torch.exp(-((y - z) ** 2) / (2.0 * epsilon * remaining)) / torch.sqrt( + 2.0 * math.pi * epsilon * remaining) + + +def bridge_drift(z: torch.Tensor, t: torch.Tensor, y: float = END) -> torch.Tensor: + """(y - z) / (1 - t), the velocity that reaches y on time if nothing further perturbs the state.""" + return (y - z) / (1.0 - t) + + +def bridge_law(t: float, x: float = START, y: float = END, + epsilon: float = EPSILON) -> tuple[float, float]: + """The mean (1 - t) x + t y and the variance epsilon t (1 - t) of the bridge at time t.""" + return (1.0 - t) * x + t * y, epsilon * t * (1.0 - t) + + +def simulate_bridge(steps: int, paths: int, epsilon: float = EPSILON, + generator: torch.Generator | None = None) -> torch.Tensor: + """Euler-Maruyama for the bridge SDE, returning the state at every time. [steps + 1, P].""" + dt = 1.0 / steps + x = torch.full((paths,), START, dtype=DTYPE) + trajectory = [x.clone()] + for step in range(steps): + t = torch.tensor(step * dt, dtype=DTYPE) + noise = torch.randn(paths, dtype=DTYPE, generator=generator) * math.sqrt(dt) + x = x + bridge_drift(x, t) * dt + math.sqrt(epsilon) * noise # One Euler-Maruyama step. + trajectory.append(x.clone()) + return torch.stack(trajectory) # [steps + 1, P]. + + +# %% 07. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=200) # --steps is the number of Euler-Maruyama steps. + args = parser.parse_args() + set_seed(args.seed) + loud = not args.quiet + steps = max(args.steps, 2) + generator = torch.Generator().manual_seed(args.seed) + + endpoints, log_ratio = simulate_controlled(2.0, steps, PATHS, EPSILON, generator) + girsanov = float(log_ratio.mean()) # The average of the log density under P^u. + closed_form = control_energy_entropy(2.0) + endpoint_form = gaussian_entropy_gap(2.0, EPSILON) + if loud: + print(f"A constant control u = 2 at epsilon = {EPSILON:g}, over {steps} steps") + print(f" (1 / 2 eps) int ||u||^2 dt = {closed_form:.3f}") + print(f" KL(N(x + 2, 1) || N(x, 1)) = {endpoint_form:.3f}") + print(f" averaged Girsanov log ratio over {PATHS} paths = {girsanov:.3f}") + print(f" endpoint mean {float(endpoints.mean()):.3f}, " + f"variance {float(endpoints.var()):.3f}") + + value = hopf_cole(bridge_potential) # V_t = -epsilon log h, with h the pinned potential. + test_points = [(-0.5, 0.2), (0.5, 0.5), (1.5, 0.5), (2.5, 0.8)] + backward = max(abs(backward_residual(bridge_potential, x, t)) for x, t in test_points) + hjb = max(abs(hjb_residual(value, x, t)) for x, t in test_points) + control_gap = max( + abs(optimal_control(value, x, t) + - float(bridge_drift(torch.tensor(x, dtype=DTYPE), torch.tensor(t, dtype=DTYPE)))) + for x, t in test_points) + if loud: + print(f"\nThe pinned potential h(z, t) and its two equations") + print(f" largest backward residual d_t h + (eps/2) h'' = {backward:.2e}") + print(f" largest Hamilton-Jacobi-Bellman residual = {hjb:.2e}") + print(f" largest gap between -grad V and (y - z) / (1 - t) = {control_gap:.2e}") + + mean_half, variance_half = bridge_law(0.5) + sample = simulate_bridge(steps, PATHS, EPSILON, generator) + half = sample[steps // 2] # The state at t = 1/2 on the grid. + state = 0.6732 # The particular state Section 7.18 evaluates the drift at. + drift_here = float(bridge_drift(torch.tensor(state, dtype=DTYPE), + torch.tensor(0.25, dtype=DTYPE))) + mean_quarter, variance_quarter = bridge_law(0.25) + spreads = {e: math.sqrt(e * 0.25) for e in (0.25, 1.0, 4.0)} # Standard deviations at t = 1/2. + if loud: + print(f"\nThe bridge from {START:g} to {END:g} at epsilon = {EPSILON:g}") + print(f" at t = 0.25: mean {mean_quarter:.4f}, variance {variance_quarter:.4f}, " + f"standard deviation {math.sqrt(variance_quarter):.4f}") + print(f" at X_0.25 = {state}: drift = ({END:g} - {state}) / 0.75 = {drift_here:.4f}, " + f"against the straight-line velocity {END - START:.0f}") + print(f" at t = 1/2: mean {mean_half:.4f}, variance {variance_half:.4f}") + print(f" simulated mean {float(half.mean()):.4f}, variance {float(half.var()):.4f}") + print(f" standard deviation at t = 1/2 for eps = 0.25, 1, 4: " + f"{[round(s, 2) for s in spreads.values()]}") + + report = { + "KL from the control energy": closed_form, + "KL from the two endpoint laws": endpoint_form, + "KL from the simulated paths": girsanov, + "simulated endpoint mean": float(endpoints.mean()), + "backward equation residual": backward, + "Hamilton-Jacobi-Bellman residual": hjb, + "-grad V against the bridge drift": control_gap, + "bridge mean at t = 0.25": mean_quarter, + "bridge variance at t = 0.25": variance_quarter, + "bridge drift at X = 0.6732": drift_here, + "straight-line velocity": float(END - START), + "bridge mean at t = 1/2": mean_half, + "bridge variance at t = 1/2": variance_half, + "simulated mean at t = 1/2": float(half.mean()), + "simulated variance at t = 1/2": float(half.var()), + "simulated endpoint, mean": float(sample[-1].mean()), + "spread at t = 1/2, eps = 0.25": spreads[0.25], + "spread at t = 1/2, eps = 1": spreads[1.0], + "spread at t = 1/2, eps = 4": spreads[4.0], + } + print_report("A diffusion reference, Chapter 7 Sections 7.16 to 7.18", report) + save_run(args.out, config=vars(args), report=report) # Written only with --out. + + +if __name__ == "__main__": + main() diff --git a/lecture_7/optimal_transport.py b/lecture_7/optimal_transport.py new file mode 100644 index 0000000000000000000000000000000000000000..f97711e0fce33eb78f54f732b4e72edac053a30c --- /dev/null +++ b/lecture_7/optimal_transport.py @@ -0,0 +1,323 @@ +"""Optimal transport, from a pairing of four points to the Gaussian closed form. + +Two distributions fix how often each point appears and leave the assignment between them free. This +file prices every assignment with the squared cost, writes a pairing as a matrix, relaxes Monge's +map to Kantorovich's coupling, certifies the optimum with prices, and ends with the two Gaussians +whose optimal map is an affine function. + +Chapter 7, Sections 7.1 to 7.6 of the notes. Runs on the chapter's own points, with no training and +no data. Prints the cost of each pairing, the plan and the prices of the two-by-two example, the +sorted pairing on three points, and the transport of N(0,1) to N(3,4). + + python lecture_7/optimal_transport.py # every worked number of Sections 7.1 to 7.6 + python lecture_7/optimal_transport.py --steps 20 # a coarse quantile grid, for a first look +""" + +# %% 01. Imports and configuration +from __future__ import annotations + +import itertools +import math + +import torch + +from cis6270.runner import common_parser, print_report, save_run, set_seed + +DTYPE = torch.float64 # The chapter quotes six figures, so every number here is double precision. + + +# %% 02. The cost of moving one unit of mass +# +# The price of carrying one unit from x to y is the squared Euclidean distance, under which a move +# twice as long costs four times as much. Section 7.1.1. + + +def cost_matrix(x: torch.Tensor, y: torch.Tensor, p: float = 2.0) -> torch.Tensor: + """C_ij = |x_i - y_j|^p on the line. x: [n], y: [m] -> [n, m].""" + return (x[:, None] - y[None, :]).abs() ** p # [n, m] one price per route. + + +def plan_cost(C: torch.Tensor, pi: torch.Tensor) -> float: + """The Frobenius inner product , the mass on each route times the price of that route.""" + return float((C * pi).sum()) + + +# %% 03. Four points and two pairings +# +# Sources at 0 and 2, targets at 1 and 3, each point with mass 1/2. Section 7.1.2 counts the two +# arrows of a pairing without weighting them by mass, and Section 7.1.3 weights them. + + +def two_pairings() -> dict: + """The ordered and crossed pairings of the four points, by arrow count and by mass.""" + x = torch.tensor([0.0, 2.0], dtype=DTYPE) # Source points. + y = torch.tensor([1.0, 3.0], dtype=DTYPE) # Target points. + C = cost_matrix(x, y) # [[1, 9], [1, 1]] the four route prices. + ordered = torch.tensor([[0.5, 0.0], [0.0, 0.5]], dtype=DTYPE) # 0 -> 1 and 2 -> 3. + crossed = torch.tensor([[0.0, 0.5], [0.5, 0.0]], dtype=DTYPE) # 0 -> 3 and 2 -> 1. + independent = torch.full((2, 2), 0.25, dtype=DTYPE) # a_i b_j with a = b = (1/2, 1/2). + return { + "C": C, + "arrows_ordered": float(C[0, 0] + C[1, 1]), # 1 + 1, the two arrow lengths squared. + "arrows_crossed": float(C[0, 1] + C[1, 0]), # 9 + 1. + "mass_ordered": plan_cost(C, ordered), # Each arrow weighted by the 1/2 it carries. + "mass_crossed": plan_cost(C, crossed), + "mass_independent": plan_cost(C, independent), # 1/4 on every route. + } + + +# %% 04. A pairing written as a matrix +# +# A coupling is an n by m matrix with row sums a and column sums b, so pi 1_m = a and pi^T 1_n = b. +# The independent coupling a_i b_j always lies in that set. Section 7.1.3. + + +def is_coupling(pi: torch.Tensor, a: torch.Tensor, b: torch.Tensor, tol: float = 1e-9) -> bool: + """True when the row sums equal a, the column sums equal b, and no entry is negative.""" + return bool( + (pi >= -tol).all() + and (pi.sum(dim=1) - a).abs().max() < tol # pi 1_m = a. + and (pi.sum(dim=0) - b).abs().max() < tol # pi^T 1_n = b. + ) + + +# %% 05. Monge's formulation, and the source that has to be split +# +# A map sends each source point to one destination, so for n equal masses it is a permutation and +# the minimization runs over n! of them. A point mass pushes forward to a point mass, which is why +# no map carries delta_0 to a law with mass at two places. Sections 7.2 and 7.2.1. + + +def monge_permutations(x: torch.Tensor, y: torch.Tensor) -> list[tuple[float, tuple[int, ...]]]: + """Every permutation of n equal point masses and its average cost, cheapest first.""" + n = len(x) + C = cost_matrix(x, y) # [n, n]. + costs = [(float(sum(C[i, p[i]] for i in range(n)) / n), p) + for p in itertools.permutations(range(n))] # Mass 1/n on each of the n routes. + return sorted(costs) + + +def pushforward_of_a_point_mass(destination: float) -> dict[float, float]: + """T_# delta_0 is the point mass at T(0), whatever the map T does elsewhere.""" + return {destination: 1.0} # One destination, all of the mass, so nu is out of reach. + + +NU_ATOMS = {1.0: 0.5, 2.0: 0.5} # nu = (1/2) delta_1 + (1/2) delta_2 of Section 7.2.1. + + +# %% 06. Kantorovich's relaxation, solved by hand on the two-by-two +# +# The feasible set of the chapter's example is the one-parameter family pi(r), and the cost along it +# is 5.8 - 8r, decreasing, so the optimum sits at the largest feasible r. Sections 7.3 to 7.3.2. + +SOURCE_X = torch.tensor([0.0, 2.0], dtype=DTYPE) # Source points of the running example. +SOURCE_A = torch.tensor([0.6, 0.4], dtype=DTYPE) # Their masses. +TARGET_Y = torch.tensor([1.0, 3.0], dtype=DTYPE) # Target points. +TARGET_B = torch.tensor([0.3, 0.7], dtype=DTYPE) # Their required masses. + + +def feasible_plan(r: float) -> torch.Tensor: + """The coupling with mass r from 0 to 1, the one degree of freedom the margins leave.""" + return torch.tensor([[r, 0.6 - r], [0.3 - r, 0.1 + r]], dtype=DTYPE) # [2, 2]. + + +def feasible_interval() -> tuple[float, float]: + """Nonnegativity of the four entries of pi(r) bounds r, and the binding pair is 0 <= r <= 0.3.""" + lower = max(0.0, -0.1) # From r >= 0 and 0.1 + r >= 0. + upper = min(0.6, 0.3) # From 0.6 - r >= 0 and 0.3 - r >= 0. + return lower, upper + + +def cost_along_the_segment(r: float, C: torch.Tensor) -> float: + """ as a function of r alone, which collects into 5.8 - 8r.""" + return plan_cost(C, feasible_plan(r)) # Linear in r with slope C11 - C12 - C21 + C22 = -8. + + +# %% 07. Prices and the dual +# +# A price f_i at each source point and g_j at each target point is feasible when f_i + g_j <= C_ij on +# every route. Complementary slackness fixes the prices from the loaded routes of an optimal plan, +# one substitution at a time, after one price is pinned to remove the shift. Sections 7.4 to 7.4.4. + + +def prices_from_loaded_routes(C: torch.Tensor, pi: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Solve f_i + g_j = C_ij on the routes carrying mass, with f_1 = 0 fixing the shift.""" + n, m = C.shape + f = torch.full((n,), float("nan"), dtype=DTYPE) + g = torch.full((m,), float("nan"), dtype=DTYPE) + f[0] = 0.0 # Adding a constant to f and subtracting it from g changes no sum f_i + g_j. + for _ in range(n + m): # Each pass propagates along one more loaded route. + for i in range(n): + for j in range(m): + if pi[i, j] <= 0: + continue # An empty route imposes no equation. + if not torch.isnan(f[i]) and torch.isnan(g[j]): + g[j] = C[i, j] - f[i] # The route prices the target point. + elif torch.isnan(f[i]) and not torch.isnan(g[j]): + f[i] = C[i, j] - g[j] # The route prices the source point. + return f, g + + +def slack(C: torch.Tensor, f: torch.Tensor, g: torch.Tensor) -> torch.Tensor: + """The matrix C_ij - f_i - g_j, nonnegative exactly when the prices are feasible.""" + return C - f[:, None] - g[None, :] # [n, m]. + + +def dual_value(a: torch.Tensor, f: torch.Tensor, b: torch.Tensor, g: torch.Tensor) -> float: + """ + , the value of a pair of prices.""" + return float(a @ f + b @ g) + + +# %% 08. Powers of a distance, and non-crossing on the line +# +# With the squared cost the difference between the ordered and crossed pairings is -2 times the +# product of the two gaps, so an optimal plan on the line contains no crossing. The quantile form +# then computes W_p by sorting. Sections 7.5 and 7.5.1. + + +def crossing_difference(x1: float, x2: float, y1: float, y2: float) -> tuple[float, float]: + """The ordered cost minus the crossed cost, expanded and in the collected form -2 dx dy.""" + expanded = (x1 - y1) ** 2 + (x2 - y2) ** 2 - (x1 - y2) ** 2 - (x2 - y1) ** 2 + collected = -2.0 * (x2 - x1) * (y2 - y1) # Every squared term cancels, leaving the cross terms. + return expanded, collected + + +def quantile_levels(levels: int) -> torch.Tensor: + """Midpoints (i + 1/2) / levels of a uniform grid on [0, 1], the levels q the integral uses.""" + return (torch.arange(levels, dtype=DTYPE) + 0.5) / levels # [levels]. + + +def empirical_quantile(points: torch.Tensor, q: torch.Tensor) -> torch.Tensor: + """F^{-1}(q) = inf{x : F(x) >= q} for n equally weighted points, a step function of q.""" + ordered = points.sort().values # [n] the sorted points, which are the steps of F^{-1}. + index = torch.clamp((q * len(points)).floor().long(), max=len(points) - 1) # [levels]. + return ordered[index] # [levels] the quantile at each level. + + +def wasserstein_on_the_line(x: torch.Tensor, y: torch.Tensor, p: float, levels: int) -> float: + """W_p^p as the integral over quantile levels q of |F_mu^{-1}(q) - F_nu^{-1}(q)|^p. + + With n equally weighted points the quantile function is a step function whose values are the + sorted points, so the level-by-level pairing is the sorted pairing and one sort on each side + computes the whole integral. + """ + q = quantile_levels(levels) # [levels]. + difference = empirical_quantile(x, q) - empirical_quantile(y, q) # [levels]. + return float((difference.abs() ** p).mean()) # The integral of a step function over [0, 1]. + + +# %% 09. Gaussians, and the case where the plan is a map +# +# For two Gaussians on the line the squared 2-Wasserstein distance is the squared difference of the +# means plus the squared difference of the standard deviations, and the optimal map is the increasing +# affine function matching quantiles. Sections 7.6 to 7.6.2. + + +def gaussian_w2_squared(m0: float, s0: float, m1: float, s1: float) -> float: + """W_2^2 between N(m0, s0^2) and N(m1, s1^2).""" + return (m0 - m1) ** 2 + (s0 - s1) ** 2 + + +def brenier_map(x: torch.Tensor, m0: float, s0: float, m1: float, s1: float) -> torch.Tensor: + """T(x) = m1 + (s1 / s0) (x - m0), the gradient of a convex function in one dimension.""" + return m1 + (s1 / s0) * (x - m0) # Same shape as x. + + +# %% 10. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=3072) # --steps is the number of quantile levels in Section 7.5.1. + args = parser.parse_args() + set_seed(args.seed) + loud = not args.quiet + + pairings = two_pairings() + if loud: + print("Four points: sources at 0 and 2, targets at 1 and 3, mass 1/2 each") + print(f" arrows ordered {pairings['arrows_ordered']:.0f} " + f"crossed {pairings['arrows_crossed']:.0f}") + print(f" by mass ordered {pairings['mass_ordered']:.0f} " + f"crossed {pairings['mass_crossed']:.0f} " + f"independent {pairings['mass_independent']:.0f}") + + C = cost_matrix(SOURCE_X, TARGET_Y) # [[1, 9], [1, 1]] for the running example. + low, high = feasible_interval() + grid = torch.linspace(low, high, 7, dtype=DTYPE) # Seven points along the feasible segment. + best_r = float(grid[int(torch.tensor([cost_along_the_segment(float(r), C) + for r in grid]).argmin())]) + pi_star = feasible_plan(high) # The cost decreases in r, so the optimum is the right endpoint. + cost_star = plan_cost(C, pi_star) + if loud: + print(f"\nTwo-by-two: a = {SOURCE_A.tolist()}, b = {TARGET_B.tolist()}, C = {C.tolist()}") + print(f" feasible r in [{low:.1f}, {high:.1f}], cost(r) = " + f"{cost_along_the_segment(0.0, C):.1f} - " + f"{(cost_along_the_segment(0.0, C) - cost_along_the_segment(1.0, C)):.0f} r") + print(f" pi* = {pi_star.tolist()} = {cost_star:.1f}") + + f, g = prices_from_loaded_routes(C, pi_star) + S = slack(C, f, g) + dual = dual_value(SOURCE_A, f, TARGET_B, g) + gap = cost_star - dual # Equals sum_ij pi_ij (C_ij - f_i - g_j), term by term nonnegative. + if loud: + print(f"\nPrices f = {f.tolist()}, g = {g.tolist()}") + print(f" slack = {S.tolist()} dual value {dual:.1f} duality gap {gap:.1e}") + print(f" slackness: mass on a route of positive slack = " + f"{float((pi_star * (S > 1e-12)).sum()):.1f}") + + expanded, collected = crossing_difference(0.0, 2.0, 1.0, 3.0) + three_x = torch.tensor([4.0, 1.0, 6.0], dtype=DTYPE) # Section 7.5.2. + three_y = torch.tensor([9.0, 2.0, 3.0], dtype=DTYPE) + permutations = monge_permutations(three_x.sort().values, three_y.sort().values) + quantile_w22 = wasserstein_on_the_line(three_x, three_y, 2.0, args.steps) + C3 = cost_matrix(three_x.sort().values, three_y.sort().values) + f3 = torch.tensor([0.0, 0.0, 8.0], dtype=DTYPE) # Three equations in six unknowns, one solution. + g3 = torch.tensor([1.0, 1.0, 1.0], dtype=DTYPE) + dual3 = dual_value(torch.full((3,), 1 / 3, dtype=DTYPE), f3, + torch.full((3,), 1 / 3, dtype=DTYPE), g3) + if loud: + print(f"\nNo crossing: expanded {expanded:+.0f} = collected {collected:+.0f}") + print(f" three points, six permutations: {[round(c, 3) for c, _ in permutations]}") + print(f" quantile integral on {args.steps} levels: {quantile_w22:.4f}") + print(f" slack of f = (0, 0, 8) and g = (1, 1, 1): {slack(C3, f3, g3).tolist()}") + print(f" dual value {dual3:.4f}") + + source = torch.special.ndtri(quantile_levels(args.steps)) # N(0,1) read at the same levels. + mapped = brenier_map(source, 0.0, 1.0, 3.0, 2.0) # T(x) = 3 + 2x. + displacement_square = float(((mapped - source) ** 2).mean()) # E[(3 + X)^2] = 9 + 1. + w22_gauss = gaussian_w2_squared(0.0, 1.0, 3.0, 2.0) + if loud: + probe = brenier_map(torch.tensor([-1.0, 0.0, 1.0], dtype=DTYPE), 0.0, 1.0, 3.0, 2.0) + print(f"\nN(0,1) to N(3,4): T(x) = {float(probe[1]):.0f} + " + f"{float(probe[2] - probe[1]):.0f}x, T(-1) = {float(probe[0]):.0f}, " + f"T(0) = {float(probe[1]):.0f}, T(1) = {float(probe[2]):.0f}") + print(f" sampled mean square displacement {displacement_square:.3f}") + + report = { + "ordered pairing, by arrows": pairings["arrows_ordered"], + "crossed pairing, by arrows": pairings["arrows_crossed"], + "ordered pairing, by mass": pairings["mass_ordered"], + "crossed pairing, by mass": pairings["mass_crossed"], + "independent coupling, by mass": pairings["mass_independent"], + "atoms of T_# delta_0, of nu": f"{len(pushforward_of_a_point_mass(1.0))}, {len(NU_ATOMS)}", + "permutations at n = 2, 3, 256": f"2, 6, a number with " + f"{len(str(math.factorial(256)))} digits", + "argmin of the cost over r": best_r, + "": cost_star, + "dual value + ": dual, + "duality gap": gap, + "pi* is a coupling": is_coupling(pi_star, SOURCE_A, TARGET_B), + "W_2 of the two-by-two": cost_star ** 0.5, + "sorted pairing on three points": permutations[0][0], + "worst pairing on three points": permutations[-1][0], + "W_2^2 from the quantile integral": quantile_w22, + "dual value on three points": dual3, + "W_2^2 of N(0,1) and N(3,4)": float(w22_gauss), + "W_2 of N(0,1) and N(3,4)": float(w22_gauss) ** 0.5, + "E[(T(X) - X)^2] on a sample": displacement_square, + } + print_report("Optimal transport, Chapter 7 Sections 7.1 to 7.6", report) + save_run(args.out, config=vars(args), report=report) # Written only with --out. + + +if __name__ == "__main__": + main() diff --git a/lecture_7/ot_flow_matching.py b/lecture_7/ot_flow_matching.py new file mode 100644 index 0000000000000000000000000000000000000000..04f19b45a42724ee29747bc4ccf7771b2c528bc8 --- /dev/null +++ b/lecture_7/ot_flow_matching.py @@ -0,0 +1,215 @@ +"""Optimal transport inside flow matching: the same loss, trained on a different coupling. + +The flow matching loss of Chapter 2 is defined for any joint law whose two marginals are p_0 and +p_1, and the independent draw of a source sample and a data sample is one such law. Replacing it +with the optimal coupling of the batch changes the training pairs and leaves the loss, the network +and the sampler as they are. Under the optimal coupling no two paths cross, so the marginal +velocity at a point is the velocity of the single path through it and the trajectories come out +straight. + +Chapter 7, Sections 7.8 and 7.8.1 of the notes. Trains two velocity networks on the four-Gaussian +target of `cis6270.data`, one on the independent coupling and one on the exact coupling of each +minibatch. Prints the coupling cost, the straightness of the learned paths and the one-step error. + + python lecture_7/ot_flow_matching.py # 300 steps per model, under a minute + python lecture_7/ot_flow_matching.py --steps 20 # a few seconds, for a first look +""" + +# %% 01. Imports and configuration +from __future__ import annotations + +import torch + +from cis6270.data import gaussian_mixture_2d +from cis6270.metrics import mode_report +from cis6270.nets import VelocityMLP, count_parameters +from cis6270.runner import ( + common_parser, pick_device, print_report, save_run, set_seed, train, +) + + +# %% 02. The coupling of a batch +# +# A batch holds n source samples and m data samples with uniform masses 1/n and 1/m, and the +# transport problem on that batch has nm unknowns against n + m equality constraints. With n = m the +# vertex solution is a permutation, which the Hungarian algorithm returns in order n^3 operations. +# Section 7.8.1. + + +def batch_cost_matrix(x0: torch.Tensor, x1: torch.Tensor) -> torch.Tensor: + """C_ij = ||x0_i - x1_j||^2 between every source sample and every data sample. [n, m].""" + return torch.cdist(x0, x1) ** 2 # [n, m] squared Euclidean cost. + + +def optimal_pairing(x0: torch.Tensor, x1: torch.Tensor) -> torch.Tensor: + """The permutation of the data samples that minimizes the total squared cost of the batch.""" + from scipy.optimize import linear_sum_assignment # The assignment problem of Section 7.2. + + C = batch_cost_matrix(x0, x1).detach().cpu().numpy() # [n, n]. + _, columns = linear_sum_assignment(C) # The cheapest permutation, in order n^3 operations. + return torch.as_tensor(columns, device=x0.device, dtype=torch.long) # [n]. + + +def couple(x0: torch.Tensor, x1: torch.Tensor, kind: str) -> torch.Tensor: + """Reorder the data samples so that row i of the batch is the pair the coupling assigns to x0_i.""" + if kind == "independent": + return x1 # Two independent draws, which is the coupling of Chapters 2 through 6. + return x1[optimal_pairing(x0, x1)] # The optimal coupling of this batch. + + +def vertex_count(n: int, m: int) -> dict[str, int]: + """The size of the batch problem: nm unknowns, n + m constraints, at most n + m - 1 loaded.""" + return {"unknowns": n * m, "constraints": n + m, "loaded_routes": n + m - 1} + + +# %% 03. The loss, with pi the only symbol that changes +# +# The network still predicts a velocity, the target is still the straight-line displacement, and the +# time is still uniform on [0, 1]. Section 7.8. + + +def flow_matching_loss(model: VelocityMLP, x0: torch.Tensor, x1: torch.Tensor) -> torch.Tensor: + """E_t ||v_theta((1-t) x_0 + t x_1) - (x_1 - x_0)||^2 over the pairs the coupling supplies.""" + t = torch.rand(x0.shape[0], device=x0.device) # [B] one time per pair. + xt = (1.0 - t[:, None]) * x0 + t[:, None] * x1 # [B, 2] the point on the straight path. + ut = x1 - x0 # [B, 2] the conditional velocity, constant along the path. + return ((model(xt, t) - ut) ** 2).sum(dim=1).mean() # Sum coordinates, average pairs. + + +def conditional_velocity_energy(x0: torch.Tensor, x1: torch.Tensor) -> float: + """The average squared conditional velocity, which Proposition 7.5 makes the kinetic action.""" + return float(((x1 - x0) ** 2).sum(dim=1).mean()) + + +# %% 04. Sampling, and two measurements of the learned field +# +# A straight trajectory is integrated exactly by one Euler step, so the gap between a one-step +# sample and a many-step sample measures how far the learned paths are from straight. The second +# measurement compares the velocity along a trajectory with the net displacement of that +# trajectory. Section 7.8. + + +@torch.no_grad() +def integrate(model: VelocityMLP, x0: torch.Tensor, steps: int) -> tuple[torch.Tensor, float]: + """Euler integration from t = 0 to t = 1, returning the samples and the path straightness. + + The straightness is the average of ||v_t(x_t) - (x_1 - x_0)||^2 along the path, which is zero + exactly when the velocity never changes along a trajectory. + """ + x = x0.clone() # [B, 2]. + dt = 1.0 / steps + velocities = [] + for step in range(steps): + t = torch.full((x.shape[0],), step * dt, device=x.device) # [B] shared time. + v = model(x, t) # [B, 2] the learned velocity at the current state. + velocities.append(v) + x = x + dt * v # One Euler step. + net = (x - x0)[None, :, :] # [1, B, 2] the net displacement of each trajectory. + deviation = ((torch.stack(velocities) - net) ** 2).sum(dim=2).mean() # Scalar. + return x, float(deviation) + + +@torch.no_grad() +def one_step_error(model: VelocityMLP, x0: torch.Tensor, steps: int) -> float: + """The distance between the one-step sample x_0 + v_0(x_0) and the many-step sample.""" + fine, _ = integrate(model, x0, steps) # [B, 2] the reference trajectory endpoint. + t0 = torch.zeros(x0.shape[0], device=x0.device) + coarse = x0 + model(x0, t0) # [B, 2] a single Euler step of size 1. + return float((coarse - fine).norm(dim=1).mean()) + + +# %% 05. Training one model +def fit(kind: str, args, device: torch.device, data: torch.Tensor) -> dict: + """Train a velocity network on the coupling named by ``kind`` and measure what it learned.""" + set_seed(args.seed) # The same initialization and the same batches for both couplings. + model = VelocityMLP(dim=2, width=args.width).to(device) + costs: list[float] = [] + + def loss_fn(step: int) -> torch.Tensor: + x1 = data[torch.randint(0, len(data), (args.batch_size,), device=device)] # [B, 2]. + x0 = torch.randn_like(x1) # [B, 2] the source sample. + paired = couple(x0, x1, kind) # [B, 2] reordered by the coupling. + costs.append(conditional_velocity_energy(x0, paired)) + return flow_matching_loss(model, x0, paired) + + losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=True) + torch.manual_seed(args.seed + 1) # The same starting noise for both samplers. + x0 = torch.randn(args.samples, 2, device=device) + samples, straightness = integrate(model, x0, args.sample_steps) + report = mode_report(samples.cpu()) + return { + "kind": kind, + "loss": sum(losses[-50:]) / len(losses[-50:]), + "coupling_cost": sum(costs) / len(costs), + "straightness": straightness, + "one_step_error": one_step_error(model, x0, args.sample_steps), + "modes_found": report["modes_found"], + "nearest_distance": report["nearest_distance"], + } + + +# %% 06. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=300, batch_size=128, width=64, samples=512, sample_steps=50) + args = parser.parse_args() + set_seed(args.seed) + device = pick_device(args.device) + loud = not args.quiet + data = gaussian_mixture_2d(num_points=4096, seed=args.seed).to(device) # [4096, 2]. + + # The two-point batch of Section 7.8.1, where both couplings can be written out. + x = torch.tensor([[0.0], [2.0]]) # Two source samples on the line. + y = torch.tensor([[1.0], [3.0]]) # Two data samples. + C = batch_cost_matrix(x, y) # [[1, 9], [1, 1]]. + independent_energy = float(C.mean()) # Weight 1/4 on each of the four pairs. + optimal_energy = float(C[torch.arange(2), optimal_pairing(x, y)].mean()) # Weight 1/2 each. + sizes = vertex_count(256, 256) # The counting the chapter does, at n = m = 256. + if loud: + print("Two source samples at 0 and 2, two data samples at 1 and 3") + print(f" independent coupling, average squared velocity {independent_energy:.0f}") + print(f" optimal coupling, average squared velocity {optimal_energy:.0f}") + print(f"\nA batch of 256 and 256: {sizes['unknowns']:,} unknowns, " + f"{sizes['constraints']} constraints, at most {sizes['loaded_routes']} loaded routes") + print(f" an exact solve costs about {256 ** 3:.3g} operations, repeated every batch") + print(f" training batch {args.batch_size}, device={device}, parameters=" + f"{count_parameters(VelocityMLP(dim=2, width=args.width)):,}") + + results = {} + for kind in ("independent", "optimal"): + if loud: + print(f"\ntraining on the {kind} coupling") + results[kind] = fit(kind, args, device, data) + if loud: + r = results[kind] + print(f" loss {r['loss']:.4f} batch coupling cost {r['coupling_cost']:.3f}") + print(f" straightness {r['straightness']:.4f} " + f"one-step error {r['one_step_error']:.4f}") + + ratio = results["independent"]["coupling_cost"] / max(results["optimal"]["coupling_cost"], 1e-12) + report = { + "two-point batch, independent": independent_energy, + "two-point batch, optimal": optimal_energy, + "batch plan unknowns": sizes["unknowns"], + "batch plan constraints": sizes["constraints"], + "loaded routes at a vertex": sizes["loaded_routes"], + "coupling cost, independent": results["independent"]["coupling_cost"], + "coupling cost, optimal": results["optimal"]["coupling_cost"], + "cost ratio": ratio, + "training loss, independent": results["independent"]["loss"], + "training loss, optimal": results["optimal"]["loss"], + "straightness, independent": results["independent"]["straightness"], + "straightness, optimal": results["optimal"]["straightness"], + "one-step error, independent": results["independent"]["one_step_error"], + "one-step error, optimal": results["optimal"]["one_step_error"], + "modes found, independent": results["independent"]["modes_found"], + "modes found, optimal": results["optimal"]["modes_found"], + "nearest mode, independent": results["independent"]["nearest_distance"], + "nearest mode, optimal": results["optimal"]["nearest_distance"], + } + print_report("Optimal transport inside flow matching, Chapter 7 Section 7.8", report) + save_run(args.out, config=vars(args), report=report) # Written only with --out. + + +if __name__ == "__main__": + main() diff --git a/lecture_7/sf2m.py b/lecture_7/sf2m.py new file mode 100644 index 0000000000000000000000000000000000000000..da9a8f9b830a64994cf013ee79a3564403e39c1f --- /dev/null +++ b/lecture_7/sf2m.py @@ -0,0 +1,248 @@ +"""Simulation-Free Score and Flow Matching: two closed-form targets at one bridge sample. + +With the entropic plan computed on a minibatch, the bridge is a mixture of Brownian bridges between +pairs drawn from that plan, and a Brownian bridge at a time t is Gaussian. One normal draw gives a +sample, and at that sample both the bridge drift and the score of the bridge marginal are available +without simulating a path. Subtracting half the noise variance times the score turns the drift into +the velocity of the probability flow, whose dependence on epsilon cancels. + +Chapter 7, Section 7.19.3 of the notes. Runs on the endpoint pair x_0 = 0, x_1 = 2 at epsilon = 1, +and on the two-point batch of Section 7.8.1. Prints the two regression targets at t = 0.25 and at +t = 0.75, the scaled score at every time on the grid, and the entropic plan of the batch. + + python lecture_7/sf2m.py # 101 times on the grid + python lecture_7/sf2m.py --steps 20 # a coarse time grid, for a first look +""" + +# %% 01. Imports and configuration +from __future__ import annotations + +import torch + +from cis6270.runner import common_parser, print_report, save_run, set_seed + +DTYPE = torch.float64 # The chapter quotes four figures of both regression targets. + +X0, X1 = 0.0, 2.0 # The endpoint pair the chapter works through. +EPSILON = 1.0 # The reference noise variance per unit time. +DRAW = 0.4 # The normal draw z the chapter uses at both times. + + +# %% 02. One sample from the bridge +# +# The Brownian bridge at time t is Gaussian with mean the straight interpolation and variance +# epsilon t (1 - t), so a sample costs one normal draw and no intermediate step. +# Section 7.19.3. + + +def bridge_mean(x0: torch.Tensor, x1: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + """m_t = (1 - t) x_0 + t x_1, the mean of the bridge between the two ends.""" + return (1.0 - t) * x0 + t * x1 + + +def bridge_std(t: torch.Tensor, epsilon: float = EPSILON) -> torch.Tensor: + """sigma_t = sqrt(epsilon t (1 - t)), which vanishes at both ends and peaks at t = 1/2.""" + return torch.sqrt(epsilon * t * (1.0 - t)) + + +def bridge_sample(x0: torch.Tensor, x1: torch.Tensor, t: torch.Tensor, z: torch.Tensor, + epsilon: float = EPSILON) -> torch.Tensor: + """x = m_t + sigma_t z, one draw from the bridge at time t.""" + return bridge_mean(x0, x1, t) + bridge_std(t, epsilon) * z + + +# %% 03. The two conditional quantities +# +# Both are conditional on the drawn pair, and neither depends on a quantity a network has to +# estimate first. Section 7.19.3. + + +def conditional_drift(x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + """u_t(x | x_0, x_1) = (x_1 - x) / (1 - t), the remaining displacement over the remaining time.""" + return (x1 - x) / (1.0 - t) + + +def conditional_score(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor, + epsilon: float = EPSILON) -> torch.Tensor: + """grad log p_t(x | x_0, x_1) = (m_t - x) / (epsilon t (1 - t)), displacement over variance.""" + return (bridge_mean(x0, x1, t) - x) / (epsilon * t * (1.0 - t)) + + +# %% 04. From the bridge drift to a probability flow velocity +# +# Subtracting half the noise variance times the score converts a diffusion drift into the velocity +# of the probability flow, which has the same marginals with no noise at all. The second form shows +# the straight-line velocity plus a coefficient times the displacement of the sample from the mean, +# with the dependence on epsilon cancelled. Section 7.19.3. + + +def flow_coefficient(t: torch.Tensor) -> torch.Tensor: + """(1 - 2t) / (2 t (1 - t)), positive before t = 1/2 and negative after.""" + return (1.0 - 2.0 * t) / (2.0 * t * (1.0 - t)) + + +def flow_velocity(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor, + epsilon: float = EPSILON) -> torch.Tensor: + """u_t - (epsilon / 2) grad log p_t, the velocity the first network regresses onto.""" + return (conditional_drift(x1, x, t) + - 0.5 * epsilon * conditional_score(x0, x1, x, t, epsilon)) + + +def flow_velocity_collected(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor, + t: torch.Tensor) -> torch.Tensor: + """(x_1 - x_0) + (1 - 2t) / (2 t (1 - t)) (x - m_t), the same velocity with delta collected.""" + return (x1 - x0) + flow_coefficient(t) * (x - bridge_mean(x0, x1, t)) + + +def scaled_score(x0: torch.Tensor, x1: torch.Tensor, x: torch.Tensor, t: torch.Tensor, + epsilon: float = EPSILON) -> torch.Tensor: + """sigma_t grad log p_t = (m_t - x) / sigma_t = -z, which has unit scale at every time.""" + return bridge_std(t, epsilon) * conditional_score(x0, x1, x, t, epsilon) + + +# %% 05. The loss and the sampler +# +# With lambda(t) = sigma_t the score target is -z, so no part of the schedule dominates the +# gradient. At sampling time a noise level g(t) is fixed and the two networks are integrated +# together, with g^2 = epsilon recovering the reference noise and g = 0 the probability flow. +# Section 7.19.3. + + +def sf2m_loss(velocity: torch.Tensor, score: torch.Tensor, velocity_target: torch.Tensor, + score_target: torch.Tensor, t: torch.Tensor, epsilon: float = EPSILON + ) -> torch.Tensor: + """||v_theta - u^o_t||^2 + lambda(t)^2 ||s_theta - grad log p_t||^2, with lambda(t) = sigma_t.""" + weight = bridge_std(t, epsilon) ** 2 # lambda(t)^2, which turns the score target into -z. + return ((velocity - velocity_target) ** 2 + weight * (score - score_target) ** 2).mean() + + +def sampler_drift(velocity: torch.Tensor, score: torch.Tensor, noise_level: float + ) -> torch.Tensor: + """v_theta + (1/2) g(t)^2 s_theta, the drift of the family of samplers between the two ends.""" + return velocity + 0.5 * noise_level ** 2 * score + + +# %% 06. The role of the coupling +# +# Pairs from the entropic plan at strength 2 epsilon make the learned equation at g^2 = epsilon the +# Schrodinger bridge, and pairs from any other coupling make it the Markov projection of that +# mixture of bridges. The plan is computed on a minibatch. Section 7.19.3. + + +def entropic_plan(x: torch.Tensor, y: torch.Tensor, a: torch.Tensor, b: torch.Tensor, + epsilon: float = EPSILON, sweeps: int = 200) -> torch.Tensor: + """The entropic plan for the squared cost at strength 2 epsilon, by Sinkhorn on exp(-C / 2 eps).""" + C = (x[:, None] - y[None, :]) ** 2 # [n, m] the squared cost of the batch. + K = torch.exp(-C / (2.0 * epsilon)) # The Gibbs kernel a Brownian reference produces. + u, v = torch.ones_like(a), torch.ones_like(b) + for _ in range(sweeps): + u = a / (K @ v) + v = b / (K.T @ u) + return u[:, None] * K * v[None, :] # [n, m]. + + +# %% 07. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=101) # --steps is the number of times on the grid. + args = parser.parse_args() + set_seed(args.seed) + loud = not args.quiet + x0 = torch.tensor(X0, dtype=DTYPE) + x1 = torch.tensor(X1, dtype=DTYPE) + z = torch.tensor(DRAW, dtype=DTYPE) + + worked = {} + for value in (0.25, 0.75): + t = torch.tensor(value, dtype=DTYPE) + x = bridge_sample(x0, x1, t, z) + worked[value] = { + "mean": float(bridge_mean(x0, x1, t)), + "std": float(bridge_std(t)), + "sample": float(x), + "drift": float(conditional_drift(x1, x, t)), + "score": float(conditional_score(x0, x1, x, t)), + "flow": float(flow_velocity(x0, x1, x, t)), + "collected": float(flow_velocity_collected(x0, x1, x, t)), + "scaled": float(scaled_score(x0, x1, x, t)), + "coefficient": float(flow_coefficient(t)), + } + if loud: + w = worked[value] + print(f"t = {value}, z = {DRAW}") + print(f" m_t = {w['mean']:.4f} sigma_t = {w['std']:.4f} x = {w['sample']:.4f}") + print(f" u_t = {w['drift']:.4f} grad log p_t = {w['score']:.4f}") + print(f" u^o_t = {w['flow']:.4f} collected form {w['collected']:.4f} " + f"coefficient {w['coefficient']:+.4f}") + print(f" sigma_t grad log p_t = {w['scaled']:.4f}, which is -z") + + grid = torch.linspace(0.0, 1.0, max(args.steps, 3), dtype=DTYPE)[1:-1] # Interior times. + samples = bridge_sample(x0, x1, grid, z) # [levels] one bridge sample per time. + scaled = scaled_score(x0, x1, samples, grid) + scaled_gap = float((scaled + z).abs().max()) # The target is -z at every time. + flow_targets = flow_velocity(x0, x1, samples, grid) + score_targets = conditional_score(x0, x1, samples, grid) + loss_exact = float(sf2m_loss(flow_targets, score_targets, flow_targets, score_targets, grid)) + loss_zero = float(sf2m_loss(torch.zeros_like(flow_targets), torch.zeros_like(score_targets), + flow_targets, score_targets, grid)) # An untrained pair of networks. + stratified = torch.special.ndtri( + (torch.arange(max(args.steps, 3), dtype=DTYPE) + 0.5) / max(args.steps, 3)) + half = torch.tensor(0.5, dtype=DTYPE) + averaged = float(flow_velocity(x0, x1, bridge_sample(x0, x1, half, stratified), half).mean()) + symmetry = worked[0.25]["flow"] + worked[0.75]["flow"] - 2.0 * (X1 - X0) + if loud: + print(f"\nOn {len(grid)} interior times") + print(f" largest gap between sigma_t grad log p_t and -z: {scaled_gap:.2e}") + print(f" u^o averaged over the draw at t = 1/2: {averaged:.4f}, " + f"against the straight-line velocity {X1 - X0:.0f}") + print(f" the two flow targets about the straight-line velocity: {symmetry:.2e}") + print(f" loss at the exact targets {loss_exact:.2e}, at a zero network {loss_zero:.4f}") + + x = torch.tensor([0.0, 2.0], dtype=DTYPE) # The batch of Section 7.8.1. + y = torch.tensor([1.0, 3.0], dtype=DTYPE) + uniform = torch.full((2,), 0.5, dtype=DTYPE) + C = (x[:, None] - y[None, :]) ** 2 + plan = entropic_plan(x, y, uniform, uniform) + ordered_cost = float(C[0, 0] + C[1, 1]) # 1 + 1, the arrows of the exact pairing. + crossed_cost = float(C[0, 1] + C[1, 0]) # 9 + 1, the arrows of the crossed pairing. + if loud: + print(f"\nThe two-point batch at epsilon = {EPSILON:g}") + print(f" exact transport pairs 0 with 1 and 2 with 3 at cost {ordered_cost:.0f}, " + f"against {crossed_cost:.0f} for the crossed pairing") + print(f" entropic plan at strength {2 * EPSILON:g}: " + f"{[[round(float(z_), 6) for z_ in row] for row in plan]}") + print(f" mass on the crossed pairing {float(plan[0, 1] + plan[1, 0]):.6f}") + + report = { + "m_t at t = 0.25": worked[0.25]["mean"], + "sigma_t at t = 0.25": worked[0.25]["std"], + "bridge sample at t = 0.25": worked[0.25]["sample"], + "u_t at t = 0.25": worked[0.25]["drift"], + "grad log p_t at t = 0.25": worked[0.25]["score"], + "u^o_t at t = 0.25": worked[0.25]["flow"], + "coefficient at t = 0.25": worked[0.25]["coefficient"], + "sigma_t grad log p_t at t = 0.25": worked[0.25]["scaled"], + "bridge sample at t = 0.75": worked[0.75]["sample"], + "u_t at t = 0.75": worked[0.75]["drift"], + "u^o_t at t = 0.75": worked[0.75]["flow"], + "coefficient at t = 0.75": worked[0.75]["coefficient"], + "the two targets, sum minus 2 u": symmetry, + "largest gap of the scaled score": scaled_gap, + "u^o averaged over the draw": averaged, + "loss at the exact targets": loss_exact, + "loss at a zero network": loss_zero, + "sampler drift at g^2 = epsilon": float(sampler_drift( + torch.tensor(worked[0.25]["flow"], dtype=DTYPE), + torch.tensor(worked[0.25]["score"], dtype=DTYPE), EPSILON ** 0.5)), + "sampler drift at g = 0": worked[0.25]["flow"], + "ordered pairing of the batch": ordered_cost, + "crossed pairing of the batch": crossed_cost, + "entropic plan, ordered mass": float(plan[0, 0] + plan[1, 1]), + "entropic plan, crossed mass": float(plan[0, 1] + plan[1, 0]), + } + print_report("Simulation-Free Score and Flow Matching, Chapter 7 Section 7.19.3", report) + save_run(args.out, config=vars(args), report=report) # Written only with --out. + + +if __name__ == "__main__": + main() diff --git a/lecture_7/sinkhorn.py b/lecture_7/sinkhorn.py new file mode 100644 index 0000000000000000000000000000000000000000..5734e3fd6e69a4f0c07cbf959908ba8ffdb85bfb --- /dev/null +++ b/lecture_7/sinkhorn.py @@ -0,0 +1,234 @@ +"""Entropic transport and Sinkhorn's algorithm on the chapter's two-by-two example. + +Subtracting epsilon times the entropy from the transport cost makes the objective strictly convex, +and its minimizer is the Gibbs kernel exp(-C/epsilon) with one scaling per row and one per column. +The two marginal constraints determine the scalings by alternate rescaling, which is one relative +entropy projection per half-step. + +Chapter 7, Sections 7.9 and 7.10 of the notes. Runs on a = (0.6, 0.4), b = (0.3, 0.7) and the cost +matrix with entries 1, 9, 1, 1. Prints the kernel, the first sweep by hand, the marginal error after +each sweep, the converged plan and its multipliers, and the plan at four values of epsilon. + + python lecture_7/sinkhorn.py # 40 sweeps, every worked number of Section 7.10 + python lecture_7/sinkhorn.py --steps 1 # the first sweep alone, computed by hand +""" + +# %% 01. Imports and configuration +from __future__ import annotations + +import torch + +from cis6270.runner import common_parser, print_report, save_run, set_seed + +DTYPE = torch.float64 # The chapter quotes six figures of the kernel and of the converged plan. + +COST = torch.tensor([[1.0, 9.0], [1.0, 1.0]], dtype=DTYPE) # C of Section 7.3.1. +SOURCE_A = torch.tensor([0.6, 0.4], dtype=DTYPE) # Row sums. +TARGET_B = torch.tensor([0.3, 0.7], dtype=DTYPE) # Column sums. +EXACT_PLAN = torch.tensor([[0.3, 0.3], [0.0, 0.4]], dtype=DTYPE) # pi* of Section 7.3.2. +EPSILON = 2.0 # The strength the chapter runs by hand. + + +# %% 02. Entropy of a plan, and the entropic objective +# +# Entropy is larger when the mass spreads over many routes and smaller when it concentrates, with +# the convention 0 log 0 = 0. Subtracting epsilon H(pi) from the cost makes the objective strictly +# convex, so the minimizer is unique and every entry of it is positive. Section 7.9. + + +def entropy(pi: torch.Tensor) -> float: + """H(pi) = -sum_ij pi_ij log pi_ij, in nats, with an empty route contributing nothing.""" + positive = pi[pi > 0] # Drop the zero entries, which contribute 0 log 0 = 0. + return float(-(positive * positive.log()).sum()) + + +def entropic_objective(pi: torch.Tensor, C: torch.Tensor, epsilon: float) -> float: + """ - epsilon H(pi), the quantity the entropic problem minimizes over Pi(a, b).""" + return float((C * pi).sum()) - epsilon * entropy(pi) + + +# %% 03. Relative entropy from a positive matrix +# +# The definition keeps its meaning when xi is any strictly positive matrix. Taking xi to be the +# independent coupling splits the logarithm into three pieces and gives -H(pi) + H(a) + H(b), so +# the entropy term and a relative entropy from independence differ by a constant on Pi(a, b). +# Section 7.9.1. + + +def relative_entropy(pi: torch.Tensor, xi: torch.Tensor) -> float: + """KL(pi || xi) = sum_ij pi_ij log(pi_ij / xi_ij) for a strictly positive xi.""" + mask = pi > 0 # Routes with no mass contribute nothing. + return float((pi[mask] * (pi[mask] / xi[mask]).log()).sum()) + + +def independent_coupling(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """a_i b_j, the coupling of largest entropy and the one Chapters 2 through 6 used.""" + return a[:, None] * b[None, :] # [n, m]. + + +# %% 04. The Lagrangian and the scaling form +# +# Setting the derivative of the Lagrangian in a single entry to zero gives +# C_ij + epsilon (log pi_ij + 1) - f_i - g_j = 0, and solving for the logarithm splits the +# minimizer into a row factor, a kernel entry and a column factor. Section 7.9.2. + + +def gibbs_kernel(C: torch.Tensor, epsilon: float) -> torch.Tensor: + """K = exp(-C / epsilon), taken entrywise, with larger entries on the cheaper routes.""" + return torch.exp(-C / epsilon) # [n, m], every entry strictly positive. + + +def scaling_form(u: torch.Tensor, K: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + """diag(u) K diag(v), the plan with row i scaled by u_i and column j scaled by v_j.""" + return u[:, None] * K * v[None, :] # [n, m]. + + +def multipliers_from_scalings(u: torch.Tensor, epsilon: float) -> torch.Tensor: + """Invert u_i = exp(f_i / epsilon - 1/2) for f_i = epsilon (log u_i + 1/2).""" + return epsilon * (u.log() + 0.5) # [n] the additive price behind a multiplicative scaling. + + +# %% 05. The two updates +# +# The row sums of diag(u) K diag(v) are u * (K v) and the column sums are v * (K^T u), and solving +# each equation for its own scaling with the other held fixed gives the two updates. One sweep is +# the pair of them, at two matrix-vector products. Section 7.10. + + +def sinkhorn_sweeps( + K: torch.Tensor, a: torch.Tensor, b: torch.Tensor, sweeps: int, +) -> tuple[torch.Tensor, torch.Tensor, list[float]]: + """Alternate u <- a / (K v) and v <- b / (K^T u), returning the scalings and the sweep errors.""" + u = torch.ones_like(a) + v = torch.ones_like(b) + errors: list[float] = [] + for _ in range(sweeps): + u = a / (K @ v) # Row sums become exactly a. + v = b / (K.T @ u) # Column sums become exactly b, and the row sums move away again. + errors.append(marginal_error(scaling_form(u, K, v), a, b)) + return u, v, errors + + +def marginal_error(pi: torch.Tensor, a: torch.Tensor, b: torch.Tensor) -> float: + """Total absolute error of the row sums against a plus that of the column sums against b.""" + return float((pi.sum(dim=1) - a).abs().sum() + (pi.sum(dim=0) - b).abs().sum()) + + +def sweeps_to_tolerance(K: torch.Tensor, a: torch.Tensor, b: torch.Tensor, + tolerance: float = 1e-9, limit: int = 2000) -> int: + """The number of sweeps that brings the total marginal error below the tolerance.""" + _, _, errors = sinkhorn_sweeps(K, a, b, limit) + for index, error in enumerate(errors): + if error < tolerance: + return index + 1 + return limit + + +# %% 06. Each update as a relative entropy projection +# +# Minimizing KL(pi || xi) over the matrices with row sums a, with no condition on the columns, gives +# pi_ij = xi_ij exp(lambda_i - 1), which is xi with each row rescaled, and the row sums fix the +# factor at a_i / sum_j xi_ij. The u update performs exactly that rescaling. Section 7.10.1. + + +def project_rows(xi: torch.Tensor, a: torch.Tensor) -> torch.Tensor: + """The matrix with row sums a closest to xi in relative entropy.""" + return xi * (a / xi.sum(dim=1))[:, None] # One factor per row. + + +def project_columns(xi: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """The matrix with column sums b closest to xi in relative entropy.""" + return xi * (b / xi.sum(dim=0))[None, :] # One factor per column. + + +# %% 07. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.add_argument("--epsilon", type=float, default=EPSILON, help="regularization strength") + parser.set_defaults(steps=40) # --steps is the number of Sinkhorn sweeps. + args = parser.parse_args() + set_seed(args.seed) + loud = not args.quiet + epsilon = args.epsilon + + K = gibbs_kernel(COST, epsilon) + independent = independent_coupling(SOURCE_A, TARGET_B) + if loud: + print(f"H(pi*) = {entropy(EXACT_PLAN):.3f} H(a x b) = {entropy(independent):.3f} " + f"H(a) + H(b) = {entropy(SOURCE_A) + entropy(TARGET_B):.3f}") + print(f"KL(pi* || a x b) = {relative_entropy(EXACT_PLAN, independent):.3f} nats") + print(f"\nK = exp(-C/{epsilon:g}) = {[[round(float(k), 6) for k in row] for row in K]}") + print(f" kernel ratio between a cheap route and the expensive one " + f"{float(K[0, 0] / K[0, 1]):.1f}") + + # The first sweep, written out the way Section 7.10.2 writes it. + v = torch.ones_like(TARGET_B) + Kv = K @ v # Row sums of K, since v = (1, 1). + u = SOURCE_A / Kv + Ktu = K.T @ u + v = TARGET_B / Ktu + first = scaling_form(u, K, v) + if loud: + print(f"\nsweep 1 K v = {[round(float(z), 6) for z in Kv]}" + f" u = {[round(float(z), 6) for z in u]}") + print(f" K^T u = {[round(float(z), 6) for z in Ktu]}" + f" v = {[round(float(z), 5) for z in v]}") + print(f" diag(u) K diag(v) = {[[round(float(z), 4) for z in row] for row in first]}") + print(f" column sums {[round(float(z), 4) for z in first.sum(0)]} " + f"row sums {[round(float(z), 4) for z in first.sum(1)]}") + after_u = project_rows(first, SOURCE_A) # The next update of u, as a row projection. + print(f" after the next u update: rows {[round(float(z), 3) for z in after_u.sum(1)]} " + f"columns {[round(float(z), 3) for z in after_u.sum(0)]}") + + u, v, errors = sinkhorn_sweeps(K, SOURCE_A, TARGET_B, max(args.steps, 1)) + plan = scaling_form(u, K, v) + f = multipliers_from_scalings(u, epsilon) + g = multipliers_from_scalings(v, epsilon) + sums = f[:, None] + g[None, :] # The entropic counterparts of the dual prices. + if loud: + print(f"\nmarginal error by sweep {[round(e, 3) for e in errors[:5]]}") + print(f" successive ratios " + f"{[round(errors[i + 1] / errors[i], 2) for i in range(min(4, len(errors) - 1))]}") + print(f" pi^eps = {[[round(float(z), 6) for z in row] for row in plan]}") + print(f" u = {[round(float(z), 4) for z in u]} v = {[round(float(z), 4) for z in v]}") + print(f" f = {[round(float(z), 3) for z in f]} g = {[round(float(z), 3) for z in g]}") + print(f" f_i + g_j = {[round(float(z), 3) for z in sums.flatten()]} " + f"against C = {[round(float(z), 0) for z in COST.flatten()]}") + print(f" the route (2,2) carries {float(plan[1, 1]):.3f} and overshoots its cost by " + f"{float(sums[1, 1] - COST[1, 1]):.3f}") + + print("\neps pi^eps_21 sweeps to 1e-9") + for value in (10.0, 2.0, 0.5, 0.1): + kernel = gibbs_kernel(COST, value) + su, sv, _ = sinkhorn_sweeps(kernel, SOURCE_A, TARGET_B, 400) + p = scaling_form(su, kernel, sv) + print(f"{value:<7g} {float((COST * p).sum()):<14.7f} {float(p[1, 0]):<16.3e} " + f"{sweeps_to_tolerance(kernel, SOURCE_A, TARGET_B)}") + + report = { + "H(pi*)": entropy(EXACT_PLAN), + "H(a x b) = H(a) + H(b)": entropy(independent), + "KL(pi* || a x b)": relative_entropy(EXACT_PLAN, independent), + "K_11": float(K[0, 0]), + "K_12": float(K[0, 1]), + "sweep 1 marginal error": errors[0], + "sweep 5 marginal error": errors[min(4, len(errors) - 1)], + "sweeps to 1e-9": sweeps_to_tolerance(K, SOURCE_A, TARGET_B), + "pi^eps_11": float(plan[0, 0]), + "pi^eps_21": float(plan[1, 0]), + "": float((COST * plan).sum()), + "": float((COST * EXACT_PLAN).sum()), + "H(pi^eps)": entropy(plan), + "objective at pi^eps": entropic_objective(plan, COST, epsilon), + "objective at pi*": entropic_objective(EXACT_PLAN, COST, epsilon), + "objective at a x b": entropic_objective(independent, COST, epsilon), + "KL(pi^eps || a x b)": relative_entropy(plan, independent), + "f": [round(float(z), 3) for z in f], + "g": [round(float(z), 3) for z in g], + } + print_report("Entropic transport and Sinkhorn, Chapter 7 Sections 7.9 and 7.10", report) + save_run(args.out, config=vars(args), report=report) # Written only with --out. + + +if __name__ == "__main__": + main() diff --git a/lecture_7/static_bridge.py b/lecture_7/static_bridge.py new file mode 100644 index 0000000000000000000000000000000000000000..c56a67c4bc12af15bb8abb820b975a9e569b7c08 --- /dev/null +++ b/lecture_7/static_bridge.py @@ -0,0 +1,295 @@ +"""The static Schrodinger bridge: a path law, the chain rule, and the scaling that solves it. + +A path law over a finite time grid is an ordinary probability vector over paths, and its relative +entropy from a reference splits by the chain rule into a term for the two endpoints and a term for +the interior. Only the first term is constrained, so the bridge keeps the reference bridges exactly +and all of its freedom goes into the endpoint coupling, which iterative proportional fitting +computes by alternate rescaling. + +Chapter 7, Sections 7.12 to 7.15.3 of the notes. Runs on the two-state chain with transition matrix +A, source law (0.5, 0.5) and requested terminal law (0.2, 0.8). Prints the eight path probabilities, +the chain rule on two tables, the fitted potentials, the tilted matrices and the bridge path law. + + python lecture_7/static_bridge.py # 4 rounds of proportional fitting + python lecture_7/static_bridge.py --steps 1 # one round, the one worked by hand +""" + +# %% 01. Imports and configuration +from __future__ import annotations + +import itertools + +import torch + +from cis6270.runner import common_parser, print_report, save_run, set_seed + +DTYPE = torch.float64 # The chapter quotes six figures of the potentials and of the path law. + +A = torch.tensor([[0.7, 0.3], [0.4, 0.6]], dtype=DTYPE) # The reference transition matrix. +PI_0 = torch.tensor([0.5, 0.5], dtype=DTYPE) # The source law. +PI_1 = torch.tensor([0.2, 0.8], dtype=DTYPE) # The requested terminal law after two steps. +STEPS = 2 # N, the number of transitions, so a path is a triple and there are eight of them. + + +# %% 02. The probability of a whole path +# +# For a Markov process the probability of a list of states is the starting mass times the product of +# the transition probabilities along it. Section 7.12.1. + + +def path_law(start: torch.Tensor, matrices: list[torch.Tensor]) -> dict[tuple[int, ...], float]: + """R(x_{0:N}) over every path of the chain, as a dictionary keyed by the list of states.""" + size = len(start) + law = {} + for path in itertools.product(range(size), repeat=len(matrices) + 1): + weight = float(start[path[0]]) # The starting mass R_0(x_0). + for k, matrix in enumerate(matrices): + weight *= float(matrix[path[k], path[k + 1]]) # One transition probability per step. + law[path] = weight + return law # K^(N+1) entries, summing to one. + + +def relative_entropy(p: dict, q: dict) -> float: + """KL(P || R) = sum over paths of P log(P / R), the definition of Chapter 1 indexed by paths.""" + return sum(v * torch.log(torch.tensor(v / q[k], dtype=DTYPE)).item() + for k, v in p.items() if v > 0) + + +# %% 03. The chain rule for relative entropy +# +# Factoring each joint law into a marginal and a conditional and averaging the logarithm of the +# ratio splits the relative entropy into a marginal term and an averaged conditional term. +# Sections 7.13 and 7.13.1. + + +def kl_vectors(p: torch.Tensor, q: torch.Tensor) -> float: + """Relative entropy between two probability vectors over the same states.""" + mask = p > 0 + return float((p[mask] * (p[mask] / q[mask]).log()).sum()) + + +def chain_rule_terms(pi: torch.Tensor, other: torch.Tensor) -> tuple[float, float, float]: + """The total, the marginal term and the averaged conditional term of KL(pi || other).""" + total = kl_vectors(pi.flatten(), other.flatten()) # The four cells, computed directly. + marginal = kl_vectors(pi.sum(dim=1), other.sum(dim=1)) # The disagreement about the row. + rows = [kl_vectors(pi[i] / pi[i].sum(), other[i] / other[i].sum()) for i in range(len(pi))] + conditional = float((pi.sum(dim=1) * torch.tensor(rows, dtype=DTYPE)).sum()) + return total, marginal, conditional + + +# %% 04. Splitting a path law at its two endpoints +# +# With the pair (x_0, x_N) in the place of the first coordinate and the interior in the place of the +# second, the chain rule separates the endpoint coupling from the bridges. The second term is free +# and nonnegative, so the bridge takes the reference bridges and the problem collapses onto the +# coupling. Section 7.13.2. + + +def endpoint_coupling(law: dict) -> torch.Tensor: + """The law of the pair (x_0, x_N) under a path law, summed over the interior states.""" + coupling = torch.zeros(2, 2, dtype=DTYPE) + for path, weight in law.items(): + coupling[path[0], path[-1]] += weight # Collect every path with the same two ends. + return coupling + + +def reference_bridge(x0: int, mid: int, xN: int, matrix: torch.Tensor, + two_step: torch.Tensor) -> float: + """R(x_1 | x_0, x_2) = A(x_0, x_1) A(x_1, x_2) / A^2(x_0, x_2), the interior given the ends.""" + return float(matrix[x0, mid] * matrix[mid, xN] / two_step[x0, xN]) + + +def fill_with_reference_bridges(coupling: torch.Tensor, matrix: torch.Tensor) -> dict: + """P(x_{0:2}) = pi(x_0, x_2) R(x_1 | x_0, x_2), a coupling filled in with reference bridges.""" + two_step = matrix @ matrix + return {(x0, mid, xN): float(coupling[x0, xN]) * reference_bridge(x0, mid, xN, matrix, two_step) + for x0, mid, xN in itertools.product(range(2), repeat=3)} + + +# %% 05. The static problem and the two potentials +# +# The bridge coupling minimizes KL(pi || R_01) over the couplings of pi_0 and pi_1, which is the +# entropic problem with R_01 in place of the Gibbs kernel, so the minimizer is R_01 with one scaling +# per row and one per column and the two marginals determine them by alternation. +# Sections 7.14 and 7.14.2. + + +def reference_endpoint_law(start: torch.Tensor, matrix: torch.Tensor, steps: int) -> torch.Tensor: + """R_01 = diag(pi_0) A^N, the reference probability of starting at x and ending at y.""" + return torch.diag(start) @ torch.linalg.matrix_power(matrix, steps) # [K, K]. + + +def proportional_fitting( + R01: torch.Tensor, start: torch.Tensor, end: torch.Tensor, rounds: int, +) -> tuple[torch.Tensor, torch.Tensor, list[float]]: + """Alternate g <- pi_1 / (R_01^T f) and f <- pi_0 / (R_01 g), returning the terminal errors. + + One round is a g update followed by an f update, started from f = (1, 1), so the row sums are + exact at the end of a round and the error is measured on the columns. + """ + f = torch.ones_like(start) + g = torch.ones_like(end) + errors: list[float] = [] + for _ in range(rounds): + g = end / (R01.T @ f) # Column sums become exactly pi_1. + f = start / (R01 @ g) # Row sums become exactly pi_0, and the columns move again. + plan = f[:, None] * R01 * g[None, :] + errors.append(float((plan.sum(dim=0) - end).abs().sum())) + return f, g, errors + + +# %% 06. The squared cost from a Brownian reference +# +# For a Brownian motion of variance epsilon the transition density is a Gibbs kernel for the squared +# distance at strength 2 epsilon, so minimizing the relative entropy from its endpoint law is the +# entropic transport problem with the noise variance in the place of the regularization strength. +# Section 7.14.1. + + +def brownian_reference_plan(epsilon: float = 1.0) -> torch.Tensor: + """The bridge coupling of the four points of Section 7.3.1 under a Brownian reference.""" + x = torch.tensor([0.0, 2.0], dtype=DTYPE) # Source points. + y = torch.tensor([1.0, 3.0], dtype=DTYPE) # Target points. + a = torch.tensor([0.6, 0.4], dtype=DTYPE) # Their masses. + b = torch.tensor([0.3, 0.7], dtype=DTYPE) + endpoint_law = torch.exp(-(x[:, None] - y[None, :]) ** 2 / (2.0 * epsilon)) # exp(-C/2). + f, g, _ = proportional_fitting(endpoint_law, a, b, 200) + return f[:, None] * endpoint_law * g[None, :] # [2, 2]. + + +# %% 07. Making the bridge a Markov process +# +# The backward potential averages g over the reference, one matrix-vector product per step, and +# reweighting each transition by the ratio of the potential at the destination to the potential at +# the current state produces a transition matrix whose rows sum to one by that same recursion. +# Sections 7.15 and 7.15.1. + + +def backward_potential(matrix: torch.Tensor, g: torch.Tensor, steps: int) -> list[torch.Tensor]: + """h_N = g and h_k = A h_{k+1}, returned as the list h_0, ..., h_N.""" + potentials = [g] + for _ in range(steps): + potentials.insert(0, matrix @ potentials[0]) # One step back, into the front of the list. + return potentials + + +def tilted_matrix(matrix: torch.Tensor, h_now: torch.Tensor, h_next: torch.Tensor) -> torch.Tensor: + """A^h_k(x, y) = A(x, y) h_{k+1}(y) / h_k(x), Doob's h-transform of one step.""" + return matrix * h_next[None, :] / h_now[:, None] # [K, K], each row summing to one. + + +# %% 08. The report +def main() -> None: + parser = common_parser(__doc__.split("\n")[0]) + parser.set_defaults(steps=4) # --steps is the number of proportional fitting rounds. + args = parser.parse_args() + set_seed(args.seed) + loud = not args.quiet + rounds = max(args.steps, 1) + + reference = path_law(PI_0, [A] * STEPS) + if loud: + print("The eight paths of the two-state chain, with their reference probabilities") + for path, weight in reference.items(): + print(f" {','.join(str(s + 1) for s in path)} R = {weight:.6f}") + print(f" they sum to {sum(reference.values()):.6f}") + + table = torch.tensor([[0.5, 0.1], [0.1, 0.3]], dtype=DTYPE) # pi of Section 7.13.1. + uniform = torch.full((2, 2), 0.25, dtype=DTYPE) # pi', the uniform law on the four cells. + total, marginal, conditional = chain_rule_terms(table, uniform) + if loud: + print(f"\nChain rule on two tables: total {total:.4f} = marginal {marginal:.4f} " + f"+ conditional {conditional:.4f}") + + R01 = reference_endpoint_law(PI_0, A, STEPS) + reference_error = float((R01.sum(dim=0) - PI_1).abs().sum()) + f, g, errors = proportional_fitting(R01, PI_0, PI_1, rounds) + plan = f[:, None] * R01 * g[None, :] + independent = PI_0[:, None] * PI_1[None, :] + if loud: + print(f"\nA^2 = {[[round(float(z), 2) for z in row] for row in A @ A]}") + print(f"R_01 = {[[round(float(z), 3) for z in row] for row in R01]} " + f"column sums {[round(float(z), 3) for z in R01.sum(0)]} error {reference_error:.2f}") + first_g = PI_1 / R01.sum(dim=0) # The first update, with f = (1, 1). + print(f" first g = {[round(float(z), 4) for z in first_g]} " + f"R_01 g = {[round(float(z), 4) for z in R01 @ first_g]}") + print(f" first f = {[round(float(z), 4) for z in PI_0 / (R01 @ first_g)]}") + print(" round f g terminal error ratio") + fr = torch.ones_like(PI_0) + for index in range(rounds): + gr = PI_1 / (R01.T @ fr) + fr = PI_0 / (R01 @ gr) + previous = reference_error if index == 0 else errors[index - 1] + ratio = f"{previous / errors[index]:.1f}" if errors[index] > 0 else "exact" + print(f" {index + 1:<7d} ({float(fr[0]):.6f}, {float(fr[1]):.6f}) " + f"({float(gr[0]):.6f}, {float(gr[1]):.6f}) {errors[index]:.3e} {ratio}") + print(f" pi* = {[[round(float(z), 6) for z in row] for row in plan]}") + print(f" KL(pi* || R_01) = {kl_vectors(plan.flatten(), R01.flatten()):.4f} " + f"KL(independent || R_01) = {kl_vectors(independent.flatten(), R01.flatten()):.4f}") + + brownian = brownian_reference_plan() + if loud: + print("\nBrownian reference at epsilon = 1 on the four points of Section 7.3.1") + print(f" pi* = {[[round(float(z), 4) for z in row] for row in brownian]}, the plan the " + f"Gibbs kernel at strength 2 produced") + + potentials = backward_potential(A, g, STEPS) + tilted = [tilted_matrix(A, potentials[k], potentials[k + 1]) for k in range(STEPS)] + middle = PI_0 @ tilted[0] + terminal = middle @ tilted[1] + bridge_law = fill_with_reference_bridges(plan, A) + bridge_kl = relative_entropy(bridge_law, reference) + middle_from_paths = sum(v for k, v in bridge_law.items() if k[1] == 0) + if loud: + print() + for k, h in enumerate(potentials): + print(f"h_{k} = ({float(h[0]):.6f}, {float(h[1]):.6f}) " + f"ratio {float(h[1] / h[0]):.2f}") + for k, matrix in enumerate(tilted): + print(f"A^h_{k} = {[[round(float(z), 6) for z in row] for row in matrix]}") + print(f" (0.5, 0.5) -> {[round(float(z), 6) for z in middle]} -> " + f"{[round(float(z), 6) for z in terminal]}") + print(f" 1 / h_0 = {[round(float(z), 6) for z in 1.0 / potentials[0]]}, " + f"the f the scaling produced") + product = torch.diag(PI_0) @ tilted[0] @ tilted[1] # The endpoint law of the tilted chain. + print(f" diag(pi_0) A^h_0 A^h_1 = " + f"{[[round(float(z), 6) for z in row] for row in product]}") + print("\npath R P* P*/R") + for path in reference: + name = ",".join(str(s + 1) for s in path) + print(f" {name} {reference[path]:.6f} {bridge_law[path]:.6f} " + f"{bridge_law[path] / reference[path]:.6f}") + + report = { + "R(1,1,2)": reference[(0, 0, 1)], + "the eight paths sum to": sum(reference.values()), + "chain rule, total": total, + "chain rule, marginal term": marginal, + "chain rule, conditional term": conditional, + "reference terminal error": reference_error, + "terminal error, round 1": errors[0], + "terminal error, last round": errors[-1], + "error ratio per round": reference_error / errors[0], + "f": [round(float(z), 6) for z in f], + "g": [round(float(z), 6) for z in g], + "pi*_11": float(plan[0, 0]), + "pi*_12": float(plan[0, 1]), + "KL(pi* || R_01)": kl_vectors(plan.flatten(), R01.flatten()), + "KL(pi_0 x pi_1 || R_01)": kl_vectors(independent.flatten(), R01.flatten()), + "h_0": [round(float(z), 6) for z in potentials[0]], + "h_2 = g": [round(float(z), 6) for z in potentials[-1]], + "A^h_0(1,1)": float(tilted[0][0, 0]), + "A^h_1(1,2)": float(tilted[1][0, 1]), + "rows of A^h_0 sum to": float(tilted[0].sum(dim=1).max()), + "marginal at the middle step": float(middle[0]), + "marginal at the last step": [round(float(z), 6) for z in terminal], + "f against 1 / h_0": float((f - 1.0 / potentials[0]).abs().max()), + "KL(P* || R) over the paths": bridge_kl, + "middle state from the paths": middle_from_paths, + } + print_report("The static Schrodinger bridge, Chapter 7 Sections 7.12 to 7.15.3", report) + save_run(args.out, config=vars(args), report=report) # Written only with --out. + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..132643e5e13b21b01998e7d5f081985fcb4437de --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,27 @@ +[project] +name = "cis6270" +version = "1.0.0" +description = "Course code for CIS 6270, Discrete Generative Models, University of Pennsylvania" +requires-python = ">=3.11" +dependencies = ["torch>=2.4", "numpy>=1.26", "scipy>=1.11"] + +[project.optional-dependencies] +images = ["torchvision>=0.19"] +proteins = ["transformers>=4.40"] +dev = ["pytest>=8.0"] + +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[tool.setuptools] +packages = ["cis6270"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +filterwarnings = [ + "ignore::UserWarning", + # torch 2.14 warns from inside its own forward-mode autograd the first time + # torch.func.jvp runs, which Chapter 6 uses. Nothing here is deprecated. + "ignore:`torch.jit.script` is deprecated:FutureWarning", +] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..f852da68256e8a58911138b957bf8e15b82d9d32 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,10 @@ +# Python 3.11 or later. +# +# Lectures 1, 4, 5, 6 and 7 need only torch and numpy, and run on a laptop CPU. +# Lectures 2 and 3 need torchvision for MNIST, and transformers for the ESM-2 examples. +torch>=2.4 +numpy>=1.26 +scipy>=1.11 +torchvision>=0.19 # Lectures 2 and 3: MNIST and image grids. +transformers>=4.40 # Lecture 3 only, and only with --dataset esm2. +pytest>=8.0 # The test suite under tests/. diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000000000000000000000000000000000000..6378c0c3d5d403bc3181523494b69b7352c1dd5e --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,23 @@ +"""Shared setup for the test suite. + +Several lecture files print five decimal places of an arithmetic result, which float32 cannot +separate, so those files raise the default precision when they are imported. A test may import a +file that does not, so the suite raises the precision here for every test alike. + +This runs at module scope, not in a fixture, because pytest imports a conftest before it +imports or collects anything else, and a module-scoped fixture in a test file would otherwise +build its values first, in the lower precision. +""" + +import pytest +import torch + +torch.set_default_dtype(torch.float64) # Match what every script does in main(). +torch.set_num_threads(1) # These networks are small; one thread is the fastest here. + + +@pytest.fixture(scope="session", autouse=True) +def _restore_defaults(): + """Put the process back the way it was once the session ends.""" + yield + torch.set_default_dtype(torch.float32) diff --git a/tests/test_lecture_1.py b/tests/test_lecture_1.py new file mode 100644 index 0000000000000000000000000000000000000000..f9f17927d138ff99179517f081ec969ffce36953 --- /dev/null +++ b/tests/test_lecture_1.py @@ -0,0 +1,554 @@ +"""Chapter 1's hand-computed numbers, asserted against the six Lecture 1 scripts. + +Every row below is a value Chapter 1 works out by hand, paired with the label under which one of +the Lecture 1 scripts prints its own computation of it. The tolerance on each row is the rounding the +chapter itself prints: a value given to three decimals is checked to 5e-4, one given to five +decimals to 5e-6, and a value the chapter states exactly is checked to machine precision. A value +the chapter predicts and a script measures by sampling has a test of its own, with a band wide +enough for the Monte Carlo error at the sample size the test uses. + + cd cis6270 && PYTHONPATH=. python -m pytest tests/test_lecture_1.py -q +""" + +import math +import subprocess +import sys +from pathlib import Path + +import pytest + +from cis6270.runner import set_seed +from lecture_1 import calculus, fields, information, linear_algebra, optimization, probability + +SEED = 0 # The seed every fixture sets before running a script's sections. +SAMPLES = 4000 # Monte Carlo draws and random competitors, fixed so the printed labels are fixed. +GRID = 1000 # Points in the numerical checks of calculus.py and fields.py. +STEPS = 600 # Gradient-descent steps, chain steps, Euler steps and estimator repeats. +EXACT = 1e-12 # The tolerance for a value the chapter states exactly. +ROOT = Path(__file__).resolve().parents[1] # The repository root, for the command-line check. + + +# %% 01. Probability, Sections 1.3 to 1.11 +@pytest.fixture(scope="module") +def probability_worked() -> dict: + """Run every worked example of probability.py and return the values it computed.""" + set_seed(SEED) + probability.WORKED.clear() + probability.two_token_distribution(quiet=True) + probability.continuous_densities(quiet=True) + probability.expected_adenine(quiet=True) + probability.monte_carlo(SAMPLES, STEPS, quiet=True) + probability.worked_variance(quiet=True) + probability.peptide_covariance(quiet=True) + probability.noising_covariance(quiet=True) + probability.joint_marginal_conditional(quiet=True) + probability.bayes(quiet=True) + probability.independence(quiet=True) + return dict(probability.WORKED) + + +PROBABILITY_ROWS = [ + # Section 1.3.1, the two-token distribution and the event that the first token is A. + ("sum_x p_X(x)", 1.0, EXACT), + ("P(x_1 = A) = p(AA) + p(AC)", 0.5, EXACT), + # Section 1.4, densities whose height exceeds one. + ("uniform[0, 0.1] density height", 10.0, EXACT), + ("P(0 <= X <= 0.02)", 0.2, EXACT), + ("peak of N(0, 0.1^2)", 3.989, 5e-4), + # Sections 1.5.1 and 1.5.3, the expected adenine count and the Monte Carlo error scale. + ("E[h(X)], expected A count", 1.1, EXACT), + ("expected A count in 1000 draws", 1100.0, EXACT), + ("E[h(X)^2]", 1.9, EXACT), + ("Var(h(X))", 0.69, EXACT), + ("sd of a 10-sample average", 0.263, 5e-4), + ("sd of a 1000-sample average", 0.0263, 5e-5), + # Section 1.6.2, the variance of the values 0, 1, 2 by both routes. + ("E[X]", 1.1, EXACT), + ("E[X^2]", 1.7, EXACT), + ("Var(X) from the two moments", 0.49, EXACT), + ("Var(X) from the definition", 0.49, EXACT), + ("sigma_X", 0.7, EXACT), + # Sections 1.7.1 and 1.7.2, the three peptides and their covariance matrix. + ("mu_X", 2.0, EXACT), + ("mu_Y", 4.0, EXACT), + ("Cov(X, Y)", 5.0 / 3.0, EXACT), + ("Var(X) = Sigma_11", 2.0 / 3.0, EXACT), + ("Var(Y) = Sigma_22", 14.0 / 3.0, EXACT), + ("sigma_X across the peptides", 0.8165, 5e-5), + ("sigma_Y across the peptides", 2.160, 5e-4), + ("rho_XY", 0.9449, 5e-5), + ("Sigma symmetry, Sigma_21 - Sigma_12", 0.0, EXACT), + # Section 1.8, the noising equation at alpha_t = 0.6 and sigma_t = 0.8. + ("alpha_t^2 + sigma_t^2", 1.0, EXACT), + ("Cov(X_t)_11 with Cov(X_0) = I", 1.0, EXACT), + ("Cov(X_t)_11 with diag(4, 1)", 2.08, EXACT), + ("Cov(X_t)_22 with diag(4, 1)", 1.00, EXACT), + # Sections 1.9.1 to 1.9.3, the motif-and-binding table with its marginals and conditionals. + ("sum over the four cells", 1.0, EXACT), + ("p(X = 1), motif present", 0.60, EXACT), + ("p(X = 0), motif absent", 0.40, EXACT), + ("p(Y = 1), high binding", 0.44, EXACT), + ("p(Y = 0), low binding", 0.56, EXACT), + ("p(Y = 1 | X = 1)", 0.6, EXACT), + ("p(Y = 0 | X = 1)", 0.4, EXACT), + ("p(Y = 1 | X = 0)", 0.2, EXACT), + ("p(Y = 0 | X = 0)", 0.8, EXACT), + # Section 1.10, Bayes' rule and the law of total probability. + ("evidence p(Y = 1)", 0.44, EXACT), + ("posterior p(X = 1 | Y = 1)", 0.818, 5e-4), + ("likelihood ratio 0.6 / 0.2", 3.0, EXACT), + # Section 1.11, the failure of independence on the same table. + ("p(X = 1) p(Y = 1)", 0.264, EXACT), + ("p(X = 1, Y = 1)", 0.36, EXACT), + ("association missed by factorizing", 0.096, EXACT), +] + + +@pytest.mark.parametrize("label, chapter, tolerance", PROBABILITY_ROWS) +def test_probability(probability_worked, label, chapter, tolerance): + assert probability_worked[label] == pytest.approx(chapter, abs=tolerance) + + +def test_monte_carlo_lands_inside_its_own_error_bar(): + """Section 1.5.3: the sample average lands within four standard deviations of E[h(X)] = 1.1.""" + set_seed(SEED) + estimates = probability.monte_carlo(SAMPLES, STEPS, quiet=True) + size, value, error = estimates["ladder"][-1] # The last rung, at the largest N the ladder reached. + assert 10 * size > SAMPLES >= size + assert abs(value - estimates["exact"]) < 4.0 * error + + +def test_the_measured_spread_of_ten_sample_averages_matches_0_263(probability_worked): + """Section 1.5.3: repeating the ten-sample average measures the 0.263 the chapter predicts. + + The spread of a sample standard deviation over 600 rounds is a few percent, so the band is + ten percent of the predicted value, which is looser than the rounding of the printed digits. + """ + predicted = math.sqrt(0.69 / 10.0) + assert probability_worked[f"sd of {STEPS} ten-sample averages"] == pytest.approx( + predicted, rel=0.10) + + +def test_the_same_seed_gives_the_same_monte_carlo_ladder(): + """Rule 8 of the conventions: a run is reproducible from its seed alone.""" + set_seed(SEED) + first = probability.monte_carlo(1000, 20, quiet=True)["ladder"] + set_seed(SEED) + second = probability.monte_carlo(1000, 20, quiet=True)["ladder"] + assert first == second + + +# %% 02. Linear algebra, Sections 1.12 to 1.18 +@pytest.fixture(scope="module") +def linear_algebra_worked() -> dict: + """Run every worked example of linear_algebra.py and return the values it computed.""" + set_seed(SEED) + linear_algebra.WORKED.clear() + linear_algebra.simplex(quiet=True) + linear_algebra.softmax(quiet=True) + linear_algebra.one_transition(quiet=True) + spectrum = linear_algebra.eigen(quiet=True) + linear_algebra.powers(spectrum["values"], spectrum["vectors"], quiet=True) + linear_algebra.diagonalization(spectrum["values"], spectrum["vectors"], quiet=True) + linear_algebra.stationary(10, quiet=True) + linear_algebra.matrix_exponential(spectrum["values"], spectrum["vectors"], quiet=True) + return dict(linear_algebra.WORKED) + + +LINEAR_ALGEBRA_ROWS = [ + # Sections 1.12.2 and 1.12.3, expectation as a dot product and a mixture of two points. + ("E[h(X)] = p^T h", 1.3, EXACT), + ("mixture r_1 at lambda = 0.5", 0.30, EXACT), + ("mixture r_2 at lambda = 0.5", 0.20, EXACT), + ("mixture r_3 at lambda = 0.5", 0.50, EXACT), + ("sum_i r_i, still a distribution", 1.0, EXACT), + # Section 1.13.1, the softmax of z = (2, 1, 0). + ("e^{z_1}", 7.389, 5e-4), + ("e^{z_2}", 2.718, 5e-4), + ("e^{z_3}", 1.000, EXACT), + ("partition function sum_j e^{z_j}", 11.107, 5e-4), + ("softmax(z)_1", 0.665, 5e-4), + ("softmax(z)_2", 0.245, 5e-4), + ("softmax(z)_3", 0.090, 5e-4), + ("sum_i softmax(z)_i", 1.0, EXACT), + ("shift invariance, largest change", 0.0, 1e-12), + # Section 1.15.1, one step of the two-state chain from an even split. + ("p_{t+1}(1)", 0.55, EXACT), + ("p_{t+1}(2)", 0.45, EXACT), + # Sections 1.16 and 1.16.1, the eigenvalues and eigenvectors of A = [[2, 1], [1, 2]]. + ("trace A, lambda_1 + lambda_2", 4.0, EXACT), + ("det A, lambda_1 lambda_2", 3.0, EXACT), + ("lambda_1", 3.0, 1e-12), + ("lambda_2", 1.0, 1e-12), + ("v_1 second entry over first", 1.0, 1e-12), + ("v_2 second entry over first", -1.0, 1e-12), + ("largest |A v - lambda v|", 0.0, 1e-12), + # Section 1.16.2, powers along an eigendirection from x_0 = (3, 1). + ("c_1", 2.0, 1e-12), + ("c_2", 1.0, 1e-12), + ("lambda_1^10", 59049.0, 1e-6), + ("x_2 first entry, eigen route", 19.0, 1e-10), + ("x_2 second entry, eigen route", 17.0, 1e-10), + ("x_10 first entry", 118099.0, 1e-6), + ("x_10 second entry", 118097.0, 1e-6), + # Sections 1.16.3 and 1.16.4, diagonalization and the dominant eigenvalue. + ("V^{-1}_11", 0.5, 1e-12), + ("largest |V Lambda V^{-1} - A|", 0.0, 1e-12), + ("A^3_11 through the eigenvalues", 14.0, 1e-10), + ("A^3_12 through the eigenvalues", 13.0, 1e-10), + ("dominant coefficient 2^10", 1024.0, EXACT), + ("decaying coefficient 0.5^10", 0.000977, 5e-7), + ("ratio of the two, in millions", 1.05, 5e-3), + # Section 1.17, the stationary distribution and the rate of approach. + ("p*_1", 0.6, 1e-10), + ("p*_2", 0.4, 1e-10), + ("largest |P^T p* - p*|", 0.0, 1e-10), + ("trace P", 1.5, EXACT), + ("det P", 0.5, EXACT), + ("second eigenvalue, the rate", 0.5, 1e-12), + ("p_1(1)", 0.55, EXACT), + ("p_2(1)", 0.575, EXACT), + ("p_3(1)", 0.5875, EXACT), + ("p_10(1)", 0.59990, 5e-6), + # Section 1.18.3, the four-term series against the eigenvalue formula at t = 0.1. + ("four-term series, entry 11", 1.22733, 5e-6), + ("four-term series, entry 12", 0.12217, 5e-6), + ("eigenvalue route, entry 11", 1.22751, 5e-6), + ("eigenvalue route, entry 12", 0.12234, 5e-6), + ("largest gap between the routes", 1.8e-4, 5e-6), + ("first discarded term, entry 11", 1.71e-4, 5e-7), + ("first discarded term, entry 12", 1.67e-4, 5e-7), + ("x_1 first entry, 2e^3 + e", 42.889, 5e-4), + ("x_1 second entry, 2e^3 - e", 37.453, 5e-4), + # Sections 1.18.4 and 1.18.5, stability and the covariance read through its eigenvalues. + ("decay e^{-2t} at t = 1", 0.1353, 5e-5), + ("decay e^{-2t} at t = 3", 0.002479, 5e-7), + ("sqrt(lambda_1), Gaussian long axis", 1.732, 5e-4), +] + + +@pytest.mark.parametrize("label, chapter, tolerance", LINEAR_ALGEBRA_ROWS) +def test_linear_algebra(linear_algebra_worked, label, chapter, tolerance): + assert linear_algebra_worked[label] == pytest.approx(chapter, abs=tolerance) + + +def test_the_chain_halves_its_deviation_at_every_step(): + """Section 1.17: the deviation from p* lies along the second eigenvector, so d_n = 0.5^n d_0.""" + chain = linear_algebra.stationary(12, quiet=True) + for step, point in enumerate(chain["history"], start=1): + deviation = float((point - chain["p_star"]).abs().max()) + assert deviation == pytest.approx(0.1 * 0.5 ** step, abs=1e-12) + + +# %% 03. Calculus, Sections 1.22 to 1.30 +@pytest.fixture(scope="module") +def calculus_worked() -> dict: + """Run every worked example of calculus.py and return the values it computed.""" + set_seed(SEED) + calculus.WORKED.clear() + calculus.jacobian(quiet=True) + calculus.determinant(quiet=True) + calculus.change_of_variables(GRID, quiet=True) + calculus.composition(quiet=True) + calculus.branches(GRID, quiet=True) + calculus.gradient(quiet=True) + calculus.divergences(quiet=True) + calculus.continuity(GRID, STEPS, quiet=True) + return dict(calculus.WORKED) + + +CALCULUS_ROWS = [ + # Sections 1.22.2 and 1.22.3, the Jacobian at (1, 0) and the displacement it predicts. + ("J_phi(1, 0)_11", 2.0, EXACT), + ("J_phi(1, 0)_21 = cos(1)", 0.5403, 5e-5), + ("largest |autograd - closed form|", 0.0, EXACT), + ("predicted change, first output", 0.0, 1e-12), + ("predicted change, second output", -0.0546, 5e-5), + ("exact change, first output", 0.0001, 5e-6), + ("exact change, second output", -0.05464, 5e-6), + # Section 1.23, the determinant as the local volume factor. + ("det A for diag(2, 3)", 6.0, EXACT), + ("area of the image rectangle", 6.0, EXACT), + ("det of the collapsing map", 0.0, EXACT), + ("det J_phi at x_1 = 2", 4.0, 1e-12), + ("det J_phi at x_1 = 0.25", 0.5, 1e-12), + # Sections 1.24.1 and 1.24.2, change of variables in one and in two dimensions. + ("preimage of [1.4, 2.0], lower end", 0.2, EXACT), + ("preimage of [1.4, 2.0], upper end", 0.5, EXACT), + ("P(X in the preimage)", 0.3, EXACT), + ("p_Y on [1, 3] after Y = 2X + 1", 0.5, EXACT), + ("integral of p_Y over [1, 3]", 1.0, EXACT), + ("p_Y inside the image of A", 1.0 / 6.0, EXACT), + ("density times area for A", 1.0, EXACT), + ("p_Y(y_1 = 1) for the x_1^2 map", 0.5, EXACT), + ("largest |p_Y |det J| - p_X|", 0.0, 1e-12), + ("total mass after the x_1^2 map", 1.0, 1e-12), + # Section 1.25, the composition and its log-determinant. + ("det J_psi", 6.0, 1e-12), + ("chained Jacobian, entry 11", 8.0, 1e-12), + ("det of the composition", 24.0, 1e-10), + ("log 4", 1.3863, 5e-5), + ("log 6", 1.7918, 5e-5), + ("log 4 + log 6 = log 24", 3.1781, 5e-5), + ("log p_Y from a unit base density", -3.1781, 5e-5), + # Section 1.26, the two inverse branches of the squaring map. + ("p_Y(0.25), two branches at ±0.5", 1.0, 1e-12), + ("p_Y(0.04), two branches at ±0.2", 2.5, 1e-12), + ("integral of p_Y over [0.01, 1]", 0.9, 5e-5), + # Sections 1.27.1 and 1.27.2, the gradient and the directional derivative. + ("grad f(1, 2), first entry", 2.0, EXACT), + ("grad f(1, 2), second entry", 8.0, EXACT), + ("|grad f(1, 2)| = sqrt(68)", math.sqrt(68.0), 1e-12), + ("steepest direction, first entry", 0.243, 5e-4), + ("steepest direction, second entry", 0.970, 5e-4), + ("D_u f along (1, 0)", 2.0, EXACT), + ("D_u f along (0, 1)", 8.0, EXACT), + # Section 1.29, the divergence of the three fields. + ("div of (-x_1, -2 x_2)", -3.0, EXACT), + ("div of (2 x_1, -x_2)", 1.0, EXACT), + ("div of (-x_2, x_1)", 0.0, EXACT), + ("area factor e^{div t} at t = 0.5", 1.649, 5e-4), + # Section 1.30.2, the worked flow v(x) = x on a uniform source. + ("largest continuity residual", 0.0, 1e-6), + ("log p along the trajectory at t = 1", -1.0, 1e-12), + ("p_t(x_t) at t = 1", 0.3679, 5e-5), + ("density at t = log 2", 0.5, EXACT), + ("mass at t = log 2", 1.0, EXACT), + ("mass at t = 1", 1.000, 5e-4), +] + + +@pytest.mark.parametrize("label, chapter, tolerance", CALCULUS_ROWS) +def test_calculus(calculus_worked, label, chapter, tolerance): + assert calculus_worked[label] == pytest.approx(chapter, abs=tolerance) + + +def test_the_euler_walk_approaches_the_exact_trajectory(calculus_worked): + """Section 1.30.2: x_t = x_0 e^t, so a first-order walk closes on 0.5e as the step shrinks.""" + coarse = calculus.continuity(GRID, 50, quiet=True)["x_t"] + fine = calculus.continuity(GRID, 4 * 50, quiet=True)["x_t"] + assert abs(fine - 0.5 * math.e) < abs(coarse - 0.5 * math.e) / 3.0 + + +# %% 04. Information, Sections 1.33 to 1.35 +@pytest.fixture(scope="module") +def information_worked() -> dict: + """Run every worked example of information.py and return the values it computed.""" + set_seed(SEED) + information.WORKED.clear() + information.information(quiet=True) + information.entropies(quiet=True) + information.cross_entropies(quiet=True) + information.kl(SAMPLES, quiet=True) + information.monte_carlo(SAMPLES, STEPS, quiet=True) + information.maximum_likelihood(quiet=True) + return dict(information.WORKED) + + +INFORMATION_ROWS = [ + # Section 1.33.1, the information content of a common and of a rare outcome. + ("I(x) at p(x) = 0.5", 0.693, 5e-4), + ("I(x) at p(x) = 0.01", 4.605, 5e-4), + ("I(x) at p(x) = 1", 0.0, EXACT), + ("ratio of the two surprises", 6.6, 5e-2), + # Section 1.33.2, entropy on a skewed, a balanced and a uniform distribution. + ("H((0.9, 0.1))", 0.325, 5e-4), + ("H((0.5, 0.5)) = log 2", 0.693, 5e-4), + ("H of a point mass", 0.0, EXACT), + ("H of the uniform on 4 states", 1.386, 5e-4), + # Section 1.33.3, the cross-entropy of p = (0.8, 0.2) scored by q = (0.6, 0.4). + ("H(p, q)", 0.592, 5e-4), + ("H(p)", 0.500, 5e-4), + ("excess cost H(p, q) - H(p)", 0.092, 5e-4), + ("H(p, p) - H(p)", 0.0, EXACT), + # Section 1.34, the divergence, its split, its nonnegativity and its asymmetry. + ("KL(p || q)", 0.0915, 5e-5), + ("KL(q || p)", 0.1046, 5e-5), + ("H(p) + KL(p || q) = H(p, q)", 0.592, 5e-4), + ("KL(p || p)", 0.0, EXACT), + ("reverse over forward, percent", 14.0, 0.5), + ("largest log t - (t - 1)", 0.0, 1e-12), + (f"largest KL breach over {SAMPLES} pairs", 0.0, EXACT), + # Section 1.35, the empirical loss on ten DNA bases and its gap to the uniform model. + ("term for A, 0.4 log(1/0.4)", 0.36652, 5e-6), + ("term for C, 0.3 log(1/0.3)", 0.36119, 5e-6), + ("term for G, 0.2 log(1/0.2)", 0.32189, 5e-6), + ("term for T, 0.1 log(1/0.1)", 0.23026, 5e-6), + ("loss at the frequencies", 1.280, 5e-4), + ("loss under the uniform model", 1.386, 5e-4), + ("gap between the two losses", 0.1064, 5e-5), + ("KL(p_hat || uniform)", 0.1064, 5e-5), +] + + +@pytest.mark.parametrize("label, chapter, tolerance", INFORMATION_ROWS) +def test_information(information_worked, label, chapter, tolerance): + assert information_worked[label] == pytest.approx(chapter, abs=tolerance) + + +def test_cross_entropy_by_sampling_lands_inside_its_error_bar(): + """Section 1.33.3 estimated by the Monte Carlo of Section 1.5.3, to within four error bars.""" + set_seed(SEED) + estimate = information.monte_carlo(SAMPLES, STEPS, quiet=True) + assert abs(estimate["estimate"] - estimate["exact"]) < 4.0 * estimate["error"] + assert estimate["measured"] == pytest.approx(estimate["predicted"], rel=0.15) + + +# %% 05. Optimization, Sections 1.36 to 1.38 +@pytest.fixture(scope="module") +def optimization_worked() -> dict: + """Run every worked example of optimization.py and return the values it computed.""" + set_seed(SEED) + optimization.WORKED.clear() + optimization.unconstrained(STEPS, quiet=True) + optimization.equality_constraint(quiet=True) + optimization.normalization(quiet=True) + optimization.maximum_entropy(SAMPLES, quiet=True) + optimization.kkt(quiet=True) + optimization.over_the_simplex(STEPS, quiet=True) + return dict(optimization.WORKED) + + +OPTIMIZATION_ROWS = [ + # Section 1.36, three steps of gradient descent on (theta - 3)^2 at eta = 0.1. + ("theta_1", 0.600, 5e-4), + ("theta_2", 1.080, 5e-4), + ("theta_3", 1.464, 5e-4), + (f"theta_{STEPS} from the closed form", 3.0, 1e-12), + ("gradient at the minimizer theta = 3", 0.0, EXACT), + # Section 1.37.1, the squared distance to the origin along x_1 + x_2 = 1. + ("x_1 at the constrained minimum", 0.5, 1e-12), + ("x_2 at the constrained minimum", 0.5, 1e-12), + ("lambda, the equality multiplier", -1.0, 1e-12), + ("f at the minimum", 0.5, 1e-12), + ("f at the endpoint (1, 0)", 1.0, EXACT), + ("largest |grad f + lambda grad g|", 0.0, 1e-12), + # Section 1.37.2, normalization alone, and the likelihood minimizer it produces. + ("uniform p_i at K = 4", 0.25, EXACT), + ("max entropy under normalization", 1.386, 5e-4), + ("p_hat(A) = n(A) / N", 0.4, EXACT), + ("p_hat(T) = n(T) / N", 0.1, EXACT), + ("minimum empirical loss", 1.280, 5e-4), + ("loss drop from moving off p_hat", 0.0, EXACT), + # Section 1.37.3, maximum entropy under a GC-content target of 0.6. + ("e^lambda", 1.5, 1e-12), + ("lambda = log 1.5", 0.4055, 5e-5), + ("partition function Z", 5.0, 1e-12), + ("p*(A)", 0.2, 1e-12), + ("p*(C)", 0.3, 1e-12), + ("GC content of p*", 0.6, 1e-12), + ("H(p*)", 1.366, 5e-4), + ("H of the unconstrained uniform", 1.386, 5e-4), + ("entropy cost of the constraint", 0.0201, 5e-5), + (f"H above H(p*) over {SAMPLES} draws", 0.0, EXACT), + # Section 1.38, the bounded quadratic and its four KKT conditions. + ("x* at the bound", 1.0, EXACT), + ("lambda, the inequality multiplier", 4.0, EXACT), + ("g(x*) = x* - 1", 0.0, EXACT), + ("complementary slackness lambda g", 0.0, EXACT), + ("stationarity 2(x* - 3) + lambda", 0.0, EXACT), + ("f at the unconstrained minimum", 0.0, EXACT), + ("f under the bound", 4.0, EXACT), + ("value drop for epsilon = 0.001", 0.004, 5e-6), + # Section 1.38.1, the same likelihood problem solved on logits through the softmax. + ("sum_i p_i at the end", 1.0, 1e-12), + ("entries at or below zero", 0.0, EXACT), + ("p(A) after the descent", 0.4, 5e-4), + ("p(T) after the descent", 0.1, 5e-4), + ("loss after the descent", 1.280, 5e-4), + ("largest |p - p_hat|", 0.0, 5e-4), +] + + +@pytest.mark.parametrize("label, chapter, tolerance", OPTIMIZATION_ROWS) +def test_optimization(optimization_worked, label, chapter, tolerance): + assert optimization_worked[label] == pytest.approx(chapter, abs=tolerance) + + +def test_the_softmax_descent_reaches_the_multiplier_solution(): + """Section 1.38.1: the unconstrained route on logits reaches what Section 1.37.2 derived.""" + set_seed(SEED) + fitted = optimization.over_the_simplex(2000, quiet=True) + frequencies = optimization.BASE_COUNTS / optimization.BASE_COUNTS.sum() + assert fitted["distance"] < 1e-5 + assert fitted["losses"][-1] == pytest.approx( + float(-(frequencies * frequencies.log()).sum()), abs=1e-9) + assert fitted["losses"][0] > fitted["losses"][-1] + + +# %% 06. Fields, Sections 1.31 and 1.32 +@pytest.fixture(scope="module") +def fields_worked() -> dict: + """Run every worked example of fields.py and return the values it computed.""" + set_seed(SEED) + fields.WORKED.clear() + fields.scores(quiet=True) + fields.curvatures(quiet=True) + fields.heat(GRID, quiet=True) + fields.fokker_planck(GRID, quiet=True) + fields.euler_maruyama(SAMPLES, STEPS, quiet=True) + return dict(fields.WORKED) + + +FIELDS_ROWS = [ + # Section 1.31, the score of the standard Gaussian and of N(1, 4). + ("s(x) at x = 3", -3.0, EXACT), + ("s(x) at x = -2", 2.0, EXACT), + ("s(x) at the mode x = 0", 0.0, EXACT), + ("largest |s(x) + x|", 0.0, EXACT), + ("largest change from dropping Z", 0.0, EXACT), + ("s(5) under N(1, 4)", -1.0, EXACT), + # Section 1.32, the two curvatures of the bowl and the cancelling pair of the saddle. + ("d^2 f / d x_1^2", 2.0, EXACT), + ("d^2 f / d x_2^2", 6.0, EXACT), + ("Laplacian of x_1^2 + 3 x_2^2", 8.0, EXACT), + ("Laplacian from the four neighbors", 8.0, 5e-6), + ("Laplacian of the saddle", 0.0, EXACT), + # Section 1.32.1, the heat equation on the widening Gaussian p_t = N(0, t) at D = 1/2. + ("p_1(0) = 1 / sqrt(2 pi)", 0.39894, 5e-6), + ("d p_t(0) / d t at t = 1", -0.19947, 5e-6), + ("Laplacian of p_1 at 0", -0.39894, 5e-6), + ("-p_1(0) / sigma^2 at sigma^2 = 1", -0.39894, 5e-6), + ("D times the Laplacian of p_1 at 0", -0.19947, 5e-6), + ("largest heat-equation residual", 0.0, 1e-6), + # Section 1.32.2, the Fokker-Planck equation and the two equations it reduces to. + ("center of p_t at t = 1", 0.5, EXACT), + ("largest residual, drift and noise", 0.0, 1e-6), + ("largest residual at sigma_t = 0", 0.0, 1e-6), +] + + +@pytest.mark.parametrize("label, chapter, tolerance", FIELDS_ROWS) +def test_fields(fields_worked, label, chapter, tolerance): + assert fields_worked[label] == pytest.approx(chapter, abs=tolerance) + + +def test_the_heat_equation_matches_the_time_derivative_at_the_center(fields_worked): + """Section 1.32.1: D times the Laplacian at the origin equals d p_t(0)/dt to every digit.""" + assert fields_worked["D times the Laplacian of p_1 at 0"] == pytest.approx( + fields_worked["d p_t(0) / d t at t = 1"], abs=EXACT) + + +def test_the_simulated_walk_tracks_the_density_it_solves(): + """Section 1.32.2: dX_t = b dt + sigma dW_t has density N(bt, t), so the variance tracks t.""" + paths = 20000 # Enough paths that four standard errors is a tight band on both moments. + set_seed(SEED) + walk = fields.euler_maruyama(paths, 200, quiet=True) + assert abs(walk["mean"] - 0.5) < 4.0 * math.sqrt(1.0 / paths) + assert abs(walk["variance"] - 1.0) < 4.0 * math.sqrt(2.0 / paths) + assert walk["variance_gap"] < 0.05 + + +# %% 07. Every script runs from the command line +@pytest.mark.parametrize("script", ["probability", "linear_algebra", "calculus", "information", + "optimization", "fields"]) +def test_the_script_runs_clean(script): + """Rule 7 of the conventions: --steps 20 --quiet finishes in a few seconds on a CPU.""" + finished = subprocess.run( + [sys.executable, f"lecture_1/{script}.py", "--steps", "20", "--quiet"], + cwd=ROOT, env={"PYTHONPATH": str(ROOT), "PATH": "/usr/bin:/bin"}, + capture_output=True, text=True, timeout=120, + ) + assert finished.returncode == 0, finished.stderr + assert finished.stderr == "" + assert "Sections" in finished.stdout diff --git a/tests/test_lecture_2.py b/tests/test_lecture_2.py new file mode 100644 index 0000000000000000000000000000000000000000..b19643a14bc9fd91560d47ff03e11cf45555271c --- /dev/null +++ b/tests/test_lecture_2.py @@ -0,0 +1,235 @@ +"""The worked numbers of Chapter 2, asserted against the Lecture 2 scripts. + +Every figure here is quoted in the chapter: the Euler landings at one, two and four steps, the flux +balance on [-1, 1], both sides of the continuity equation at x = 2, the single endpoint pair, the +average at the crossing, the nine-unit gap between the two objectives, the table of three +interpolants at t = 0.25, and the parameter count of the MNIST U-Net. Nothing in this file trains. +""" + +import math +import sys +from pathlib import Path + +import pytest +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # Run from a fresh clone. + +from lecture_2 import flow_matching_2d as fm2d # noqa: E402 +from lecture_2 import interpolants as interp # noqa: E402 +from lecture_2 import velocity_fields as vf # noqa: E402 + +LOG2 = math.log(2.0) + + +# %% 01. Velocity fields, trajectories and the flow map (Sections 2.2 to 2.4) + + +def test_trajectories_match_the_chapter(): + report = vf.trajectories() + assert report["constant_field_trajectory"] == [1.0, 1.5, 2.0, 2.5, 3.0] + assert report["expansion_x_half_from_two"] == pytest.approx(2.828, abs=1e-3) + assert report["expansion_endpoints"] == [-2.0, 0.0, 2.0] + assert report["expansion_outer_spacing"] == pytest.approx(4.0) + assert report["contraction_endpoints"] == [-0.5, 0.0, 0.5] + assert report["contraction_outer_spacing"] == pytest.approx(1.0) + + +def test_rotation_preserves_the_radius(): + report = vf.trajectories() + assert report["rotation_quarter_turn"] == [0.0, 1.0] + assert report["rotation_half_turn"] == [-1.0, 0.0] + assert report["rotation_radius"] == pytest.approx(1.0, abs=1e-6) + assert report["rotation_squared_norm_rate"] == pytest.approx(0.0, abs=1e-6) + + +def test_flow_map_is_the_identity_at_zero_and_doubles_at_one(): + report = vf.flow_map_report() + assert report["flow_map_identity_error"] == pytest.approx(0.0, abs=1e-6) + assert report["flow_map_jacobian_det"] == pytest.approx(2.0, abs=1e-6) + assert report["trajectory_gap_minimum"] == pytest.approx(2.0, abs=1e-6) + + +# %% 02. Euler's method (Section 2.5) + + +def test_four_euler_steps_land_at_the_chapter_value(): + landing = vf.euler(lambda x, t: x, torch.tensor(1.0), 4) # dx/dt = x from x_0 = 1. + assert float(landing) == pytest.approx(2.44140625, abs=1e-6) + + +def test_euler_ladder_errors_and_ratios(): + report = vf.euler_ladder(4) + assert report["euler_step_counts"] == [1, 2, 4] + assert report["euler_landings"] == [2.0, 2.25, pytest.approx(2.4414, abs=1e-4)] + assert report["euler_errors"] == [pytest.approx(0.7183, abs=1e-4), + pytest.approx(0.4683, abs=1e-4), + pytest.approx(0.2769, abs=1e-4)] + assert report["euler_error_halving_ratios"] == [pytest.approx(1.53, abs=5e-3), + pytest.approx(1.69, abs=5e-3)] + assert report["euler_relative_error_four_steps"] == pytest.approx(0.1019, abs=1e-4) + + +def test_the_euler_error_closes_like_one_over_the_step_count(): + report = vf.euler_ladder(256) # The error times the step count settles at e/2 = 1.359. + scaled = report["euler_error_times_steps"] + assert scaled[-1] == pytest.approx(math.e / 2.0, abs=0.02) + assert report["euler_error_halving_ratios"][-1] == pytest.approx(2.0, abs=0.02) + + +# %% 03. The pushforward, the flux and the continuity equation (Sections 2.6 and 2.7) + + +def test_the_expansion_field_doubles_the_standard_deviation(): + report = vf.pushforward_moments(4096, seed=0) + assert report["pushforward_mean"] == pytest.approx(0.0, abs=0.05) + assert report["pushforward_variance"] == pytest.approx(4.0, abs=0.2) + + +def test_flux_balance_on_the_unit_interval(): + report = vf.flux_balance() + assert report["density_at_one"] == pytest.approx(0.2420, abs=1e-4) + assert report["flux_at_minus_one"] == pytest.approx(-0.1677, abs=1e-4) + assert report["flux_at_one"] == pytest.approx(0.1677, abs=1e-4) + assert report["flux_balance"] == pytest.approx(-0.3354, abs=1e-4) + assert report["enclosed_mass"] == pytest.approx(0.6827, abs=1e-4) + # The balance and the derivative of the enclosed mass agree to four figures. + assert report["enclosed_mass_rate"] == pytest.approx(report["flux_balance"], abs=1e-4) + + +def test_both_sides_of_the_continuity_equation_at_x_equals_two(): + report = vf.continuity_check() + assert report["density_at_two"] == pytest.approx(0.05399, abs=1e-5) + assert report["density_gradient_at_two"] == pytest.approx(-0.10798, abs=1e-5) + assert report["velocity_divergence"] == pytest.approx(LOG2, abs=1e-6) + assert report["continuity_time_derivative"] == pytest.approx(0.11227, abs=1e-5) + assert report["continuity_minus_flux_divergence"] == pytest.approx(0.11227, abs=1e-5) + + +def test_log_density_along_a_trajectory_falls_by_log_two(): + report = vf.log_density_along_trajectory() + assert report["log_density_rate"] == pytest.approx(-LOG2) + assert report["log_density_change"] == pytest.approx(-LOG2, abs=1e-6) + assert report["log_density_accumulated"] == pytest.approx(-LOG2) + assert report["prior_density_at_origin"] == pytest.approx(0.3989, abs=1e-4) + assert report["final_density_at_origin"] == pytest.approx(0.1995, abs=1e-4) + assert report["density_ratio_at_origin"] == pytest.approx(2.0, abs=1e-6) + + +# %% 04. The conditional path and the crossing (Sections 2.12 and 2.13) + + +def test_the_single_endpoint_pair(): + report = fm2d.worked_pair() + assert report["worked_pair_path"] == [-2.0, -0.75, 0.5, 1.75, 3.0] + assert report["worked_pair_velocity"] == pytest.approx(5.0) + + +def test_the_average_at_the_crossing(): + report = fm2d.crossing_average() + assert report["crossing_state"] == pytest.approx(2.0) + assert report["crossing_velocities"] == [4.0, -2.0] + assert report["crossing_average_velocity"] == pytest.approx(1.0) + + +def test_the_two_losses_differ_by_nine_at_the_crossing(): + report = fm2d.crossing_gap() + assert report["crossing_conditional_loss"] == [10.0, 9.0, 10.0] + assert report["crossing_marginal_loss"] == [1.0, 0.0, 1.0] + assert report["crossing_constant_gap"] == [9.0, 9.0, 9.0] + assert report["crossing_target_variance"] == pytest.approx(9.0) + + +# %% 05. The marginal velocity and the equivalence (Sections 2.13 and 2.15) + + +def test_the_marginal_velocity_at_the_two_ends(): + x = torch.randn(64, 2) + # v_0(x) = E[X_1] - x = -x, and v_1(x) = x - E[X_0] = x, since the coupling is independent. + assert torch.allclose(fm2d.marginal_velocity(x, torch.zeros(64)), -x, atol=1e-5) + assert torch.allclose(fm2d.marginal_velocity(x, torch.ones(64)), x, atol=1e-5) + + +def test_the_marginal_velocity_transports_the_prior_onto_the_data(): + from cis6270.metrics import mean_log_density, mode_report + + torch.manual_seed(0) + x_t = torch.randn(512, 2) # X_0 ~ p_0. + for step in range(50): # Euler on the exact field, with no network anywhere. + t = torch.full((512,), step / 50.0) + x_t = x_t + fm2d.marginal_velocity(x_t, t) / 50.0 + assert mode_report(x_t)["modes_found"] == 4 + assert mean_log_density(x_t) > -3.0 # The data themselves measure -2.14. + + +def test_the_losses_decompose_exactly_on_any_batch(): + from cis6270.nets import VelocityMLP + + torch.manual_seed(0) + model = VelocityMLP(dim=2, width=32) + x_0, x_1 = torch.randn(256, 2), torch.randn(256, 2) + 2.0 + t = torch.rand(256) + loss_cfm, loss_fm, gap, cross = fm2d.objectives(model, x_0, x_1, t) + # Expanding the square gives L_CFM = L_FM + cross + gap term by term, on every batch. + assert float(loss_cfm) == pytest.approx(float(loss_fm + cross + gap), rel=1e-5) + + +def test_the_two_objectives_share_a_gradient(): + from cis6270.data import gaussian_mixture_2d + from cis6270.nets import VelocityMLP + + torch.manual_seed(0) + model = VelocityMLP(dim=2, width=64) + data = gaussian_mixture_2d(4096, seed=0) + report = fm2d.equivalence_check(model, data, 8192) + assert report["loss_cfm"] == pytest.approx(report["loss_fm_plus_gap"], abs=0.3) + assert abs(report["loss_cross_term"]) < 0.3 # Zero in expectation, small on a large batch. + assert report["gradient_cosine"] > 0.99 # The residual is the batch's own sampling error. + + +# %% 06. The three interpolants (Section 2.18) + + +def test_every_schedule_reaches_the_right_endpoints(): + for name in interp.SCHEDULES: + assert interp.endpoint_and_table(name)[f"{name}_endpoint_error"] == pytest.approx(0.0, + abs=1e-6) + + +def test_the_table_of_three_interpolants_at_one_quarter(): + expected = { # Table 2.1, at X_0 = (2, 0), X_1 = (0, 2) and t = 0.25. + "linear": ([1.50, 0.50], [-2.0, 2.0]), + "quadratic": ([1.875, 0.125], [-1.0, 1.0]), + "trigonometric": ([1.848, 0.765], [-1.202, 2.902]), + } + for name, (position, velocity) in expected.items(): + report = interp.endpoint_and_table(name) + assert report[f"{name}_position_at_quarter"] == pytest.approx(position, abs=1e-3) + assert report[f"{name}_velocity_at_quarter"] == pytest.approx(velocity, abs=1e-3) + + +def test_only_the_trigonometric_schedule_keeps_identity_covariance(): + trig = interp.covariance_check("trigonometric", 8192, seed=0) + assert trig["trigonometric_coefficient_squares"] == [1.0, 1.0, 1.0, 1.0, 1.0] + assert trig["trigonometric_covariance_deviation"] < 0.05 + linear = interp.covariance_check("linear", 8192, seed=0) + assert linear["linear_coefficient_squares"] == [1.0, 0.625, 0.5, 0.625, 1.0] + assert linear["linear_covariance_deviation"] > 0.4 # Cov(X_{1/2}) is half the identity. + + +def test_the_quadratic_schedule_has_covered_a_sixteenth_at_one_quarter(): + # beta_{1/4} = (1/4)^2 = 0.0625, so the state has covered 6.25 percent of the segment. + _, beta, _, _ = interp.quadratic(torch.tensor([0.25])) + assert float(beta) == pytest.approx(0.0625) + _, beta_half, _, _ = interp.quadratic(torch.tensor([0.5])) + assert float(beta_half) == pytest.approx(0.25) # A quarter of the distance at the midpoint. + + +# %% 07. The MNIST U-Net (Section 2.20) + + +def test_the_unet_has_the_parameter_count_the_chapter_quotes(): + from lecture_2 import flow_matching_mnist as mnist + + model = mnist.FlowUNet() # The default base width of 32, so the widths are 32, 64 and 128. + assert sum(p.numel() for p in model.parameters()) == 471_265 diff --git a/tests/test_lecture_3.py b/tests/test_lecture_3.py new file mode 100644 index 0000000000000000000000000000000000000000..192b59272ae91565d0452ab5f78dc1bf2838cfcd --- /dev/null +++ b/tests/test_lecture_3.py @@ -0,0 +1,445 @@ +"""Lecture 3: the numbers Chapter 3 works by hand, and the invariants the code must keep. + +Every assertion here is either a quantity the chapter computes in the text or a property the +derivation guarantees. Nothing trains, so the whole file runs in a few seconds. +""" + +import math + +import pytest +import torch + +from cis6270.metrics import mode_report +from lecture_3 import ( + ddpm, forward_diffusion, guidance, latent_diffusion, samplers, score_matching, sde, +) + + +# %% The forward process, Sections 3.2 through 3.10 + + +def test_euler_maruyama_step_matches_the_worked_numbers(): + landed, drift, noise = sde.euler_maruyama_step(2.0, -1.0, 0.5, 0.04, 1.5) + assert landed == pytest.approx(2.11) + assert (drift, noise) == pytest.approx((-0.04, 0.15)) + + +def test_three_draws_average_above_the_drift_prediction(): + endpoints = [sde.euler_maruyama_step(2.0, -1.0, 0.5, 0.04, e)[0] for e in (1.5, -0.5, 0.2)] + assert endpoints == pytest.approx([2.11, 1.91, 1.98]) + assert sum(endpoints) / 3.0 == pytest.approx(2.00, abs=1e-6) # The drift alone reaches 1.96. + + +def test_increment_scales_as_the_square_root_of_the_step(): + assert sde.increment_scale(0.04)["std"] == pytest.approx(0.2) + assert sde.increment_scale(0.02)["std"] == pytest.approx(0.1414, abs=1e-4) + assert sde.increment_scale(0.04, -1.0, 0.5)["noise_over_drift"] == pytest.approx(2.5) + + +def test_ornstein_uhlenbeck_law_at_log_two(): + law = sde.ornstein_uhlenbeck_law(2.0, 1.0, math.sqrt(2.0), math.log(2.0)) + assert law["mean"] == pytest.approx(1.0) + assert law["variance"] == pytest.approx(0.75) + assert law["std"] == pytest.approx(0.8660, abs=1e-4) + assert law["stationary_variance"] == pytest.approx(1.0) + + +def test_ito_recovers_the_variance_and_the_chain_rule_alone_does_not(): + generator = torch.Generator().manual_seed(0) + check = sde.check_ito_on_square(8192, 200, 0.25, generator) + assert check["ito_mean"] == pytest.approx(0.25, abs=0.02) + assert check["realized_mean_square"] == pytest.approx(0.25, abs=0.02) + assert abs(check["no_curvature_mean"]) < 0.02 # The curvature term carries all the growth. + + +def test_fokker_planck_holds_on_the_spreading_gaussian(): + grid = torch.linspace(-3.0, 3.0, 41) + for t in (0.0, 0.5, 2.0): + residual = sde.fokker_planck_residual(grid, torch.full_like(grid, t)) + assert float(residual.abs().max()) < 1e-6 + + +def test_peak_of_the_gaussian_falls_at_the_rate_the_equation_gives(): + peak = sde.peak_extrapolation(0.1) + assert peak["p_0"] == pytest.approx(0.3989, abs=1e-4) + assert peak["rate"] == pytest.approx(-0.1995, abs=1e-4) + assert peak["linear"] == pytest.approx(0.3790, abs=1e-4) + assert peak["exact"] == pytest.approx(0.3804, abs=1e-4) + assert peak["gap"] == pytest.approx(0.0014, abs=1e-4) + + +# %% Score matching, Sections 3.14 through 3.17 + + +def test_posterior_average_at_the_running_observation(): + report = score_matching.posterior_report(7.4, (10.0, 9.0), (0.3, 0.7), epsilon_theta=-0.8) + assert report["weights"] == pytest.approx([0.2156, 0.7844], abs=1e-4) + assert report["clean_mean"] == pytest.approx(9.2156, abs=1e-4) + assert report["conditional_scores"] == pytest.approx([5 / 3, -5 / 9], abs=1e-4) + assert report["marginal"] == pytest.approx([-0.0765, -0.0765], abs=1e-4) + assert report["residual"] == pytest.approx(0.835, abs=1e-3) + assert report["score_hat"] == pytest.approx(4 / 3, abs=1e-4) + assert report["x0"] == pytest.approx([9.85, 9.85, 9.2156], abs=1e-4) + + +def test_one_parameter_fit_is_minimized_at_the_inverse_variance(): + assert score_matching.one_parameter_loss(0.5, 1.0) == pytest.approx(-0.375) + assert score_matching.one_parameter_loss(1.0, 1.0) == pytest.approx(-0.5) + assert score_matching.one_parameter_loss(0.25, 4.0) == pytest.approx(-0.125) + grid = [0.05 * i for i in range(1, 41)] + best = min(grid, key=lambda a: score_matching.one_parameter_loss(a, 4.0)) + assert best == pytest.approx(0.25) # a = 1 / v with v = 4. + + +def test_the_three_objectives_share_one_minimizer(): + generator = torch.Generator().manual_seed(1) + clean = score_matching.sample_clean(32768, generator) + epsilon = torch.randn(clean.shape, generator=generator) + found = score_matching.minimizers(clean, epsilon) + assert found["explicit"] == pytest.approx(1.0) + assert found["implicit"] == pytest.approx(1.0, abs=0.03) + assert found["denoising"] == pytest.approx(1.0, abs=0.03) + + +def test_integration_by_parts_holds_on_the_exact_field(): + generator = torch.Generator().manual_seed(2) + clean = score_matching.sample_clean(32768, generator) + points = score_matching.A_T * clean + score_matching.SIGMA_T * torch.randn( + clean.shape, generator=generator, + ) + identity = score_matching.integration_by_parts_gap(score_matching.exact_score, points) + assert identity["gap"] < 0.05 * abs(identity["left"]) + + +def test_hutchinson_estimates_the_trace_the_exact_loop_computes(): + generator = torch.Generator().manual_seed(3) + means = torch.randn(2, score_matching.TRACE_DIM, generator=generator) * 1.5 + points = torch.randn(8, score_matching.TRACE_DIM, generator=generator) + + def field(z): + return score_matching.high_dimensional_score(z, means, 1.0) + + exact, passes = score_matching.exact_trace(field, points) + assert passes == score_matching.TRACE_DIM # One backward pass per coordinate. + estimate, draws = score_matching.hutchinson_trace(field, points, 256, generator) + error = float(draws.std(dim=0).mean()) / math.sqrt(256) + assert float((estimate - exact).mean().abs()) < 4.0 * error + + +# %% The schedule, the corruption and DDPM, Sections 3.11, 3.12, 3.19 and 3.24 through 3.28 + + +def test_schedule_matches_the_continuous_curve(): + assert float(forward_diffusion.alphas_cumprod[207]) == pytest.approx(0.6402, abs=1e-4) + assert float(forward_diffusion.a[207]) == pytest.approx(0.8001, abs=1e-4) + assert float(forward_diffusion.sigma[207]) == pytest.approx(0.5999, abs=1e-4) + assert forward_diffusion.continuous_signal(0.207) == pytest.approx(0.7997, abs=1e-4) + assert float(forward_diffusion.alphas_cumprod[1000]) == pytest.approx(4.036e-5, rel=1e-3) + assert float(forward_diffusion.a[1000]) == pytest.approx(0.0064, abs=1e-4) + assert forward_diffusion.continuous_signal(1.0) == pytest.approx(0.0066, abs=1e-4) + assert float(forward_diffusion.a[500]) == pytest.approx(0.2803, abs=1e-4) + + +def test_every_copy_of_the_schedule_agrees(): + """No lecture file imports another, so each rebuilds the schedule. The copies must match.""" + for name in ("betas", "alphas", "alphas_cumprod"): + for other in (ddpm, latent_diffusion): + assert torch.equal(getattr(forward_diffusion, name), getattr(other, name)), name + assert torch.equal(ddpm.posterior_variance[1:], latent_diffusion.posterior_variance[1:]) + # The samplers read the same path off the continuous curve, which agrees at a fine grid. + for grid in (samplers.make_schedule(1000), guidance.schedule_tensors(1000)): + assert float(grid["alphas_cumprod"][207].sqrt()) == pytest.approx(0.8, abs=5e-3) + assert torch.allclose(grid["alphas_cumprod"], ddpm.alphas_cumprod, atol=1e-3) + + +def test_noise_levels_of_the_continuous_schedule(): + for signal, time, rate in ((0.98, 0.0589, 1.2721), (0.8, 0.2068, 4.2157), + (0.35, 0.4544, 9.1420), (0.04, 0.7994, 16.0073)): + located = forward_diffusion.level_time(signal) + assert located["t"] == pytest.approx(time, abs=1e-4) + assert located["beta_t"] == pytest.approx(rate, abs=1e-3) + assert located["sigma_t"] == pytest.approx(math.sqrt(1.0 - signal ** 2), abs=1e-9) + assert forward_diffusion.level_time(0.8)["sigma_t"] == pytest.approx(0.6, abs=1e-4) + + +def test_dropped_elbo_weight_at_step_five_hundred(): + weight = ddpm.elbo_weight( + ddpm.betas[500], ddpm.posterior_variance[500], + ddpm.alphas[500], ddpm.alphas_cumprod[500], + ) + assert float(ddpm.posterior_variance[500]) == pytest.approx(0.01003, abs=1e-5) + assert float(weight) == pytest.approx(0.005508, abs=1e-6) + assert ddpm.elbo_weight(0.2, 1 / 9, 0.8, 0.64) == pytest.approx(0.625, abs=1e-6) + + +def test_one_pixel_through_the_forward_code(): + a_k = float(forward_diffusion.a[207]) + sigma_k = float(forward_diffusion.sigma[207]) + assert a_k * 0.5 + sigma_k * (-1.0) == pytest.approx(-0.1998, abs=1e-4) + + +def test_q_sample_uses_the_supplied_noise_exactly(): + x_0 = torch.randn(4, 1, 28, 28) + noise = torch.randn_like(x_0) + k = torch.tensor([1, 207, 500, 1000]) + for module in (forward_diffusion, ddpm): + x_k, returned = module.q_sample(x_0, k, noise) + abar = module.alphas_cumprod[k].view(-1, 1, 1, 1) + assert torch.allclose(x_k, abar.sqrt() * x_0 + (1 - abar).sqrt() * noise, atol=1e-6) + assert returned is noise + + +def test_corruption_law_matches_the_gaussian_it_is_drawn_from(): + generator = torch.Generator().manual_seed(0) + x_0 = torch.randn(1, 1, 28, 28, generator=generator) + levels = [1, 207, 500, 1000] + law = forward_diffusion.measure_law(x_0, levels, 1024, generator) + assert law["mean_gap"] < 0.05 # Centered on a_k x_0, to one over the square root of 1024. + assert law["variance_gap"] < 0.05 # Spread sigma_k^2. + # Four times the draws halves the gap, which a systematic error in a coefficient would not. + finer = forward_diffusion.measure_law(x_0, levels, 4096, generator) + assert finer["mean_gap"] < 0.6 * law["mean_gap"] + + +def test_worked_training_batch_at_abar_point_eight(): + batch = forward_diffusion.worked_batch( + 0.8, (2.0, -1.0, 0.5), (0.5, -0.2, 1.0), (0.4, -0.1, 0.8), + ) + assert batch["a_k"] == pytest.approx(0.8944, abs=1e-4) + assert batch["sigma_k"] == pytest.approx(0.4472, abs=1e-4) + assert batch["x_k"] == pytest.approx([2.0125, -0.9839, 0.8944], abs=1e-3) + assert batch["loss"] == pytest.approx(0.02, abs=1e-6) + assert batch["score_targets"] == pytest.approx([-1.118, 0.447, -2.236], abs=1e-3) + assert batch["x_0_hat"] == pytest.approx([2.05, -1.05, 0.60], abs=1e-2) + + +def test_two_step_reverse_transition_agrees_in_both_forms(): + step = ddpm.two_step_transition(0.2, 7.4, -1.0, 0.3) + assert step["x_0_hat"] == pytest.approx(10.0, abs=1e-6) + assert step["mean_from_noise"] == pytest.approx(8.646, abs=1e-3) + assert step["mean_from_clean"] == pytest.approx(step["mean_from_noise"], abs=1e-6) + assert step["weights"] == pytest.approx([0.4969, 0.4969], abs=1e-4) + assert step["variance"] == pytest.approx(1 / 9) + assert step["x_{k-1}"] == pytest.approx(8.746, abs=1e-3) + assert step["elbo_weight"] == pytest.approx(0.625, abs=1e-6) + + +def test_unet_keeps_the_image_shape_and_the_last_reverse_step_adds_no_noise(): + torch.manual_seed(0) + model = ddpm.UNet(in_channels=1, base_channels=8, time_embed_dim=32) + x = torch.randn(2, 1, 28, 28) + k = torch.tensor([1, 1]) + assert model(x, k).shape == x.shape + first, second = ddpm.p_sample(model, x, k), ddpm.p_sample(model, x, k) + assert torch.allclose(first, second) # keep = 0 at k = 1, so the step is the mean. + + +# %% The three samplers, Sections 3.22, 3.26, 3.32, 3.34 and 3.35 + + +def test_one_backward_step_under_both_dynamics(): + step = samplers.one_step_comparison(1.2, 2.0, 0.5, 0.8, 0.1) + assert step["score"] == pytest.approx(-1.6) + assert step["velocity"] == pytest.approx(0.4) + assert step["ode_next"] == pytest.approx(1.16) + assert step["sde_drift"] == pytest.approx(2.0) + assert step["sde_mean"] == pytest.approx(1.0) + assert step["sde_noise_std"] == pytest.approx(0.4472, abs=1e-4) + + +def test_langevin_step_from_the_worked_state(): + assert samplers.langevin_step(2.0, -2.0, 0.01, 0.5) == pytest.approx(2.0507, abs=1e-4) + + +def test_schedule_stays_a_valid_discretization_at_any_grid_size(): + for num_steps in (20, 200, 1000): + schedule = samplers.make_schedule(num_steps) + assert float(schedule["betas"][1:].max()) < 1.0 + assert float(schedule["alphas_cumprod"][0]) == pytest.approx(1.0) + assert float(schedule["alphas_cumprod"][num_steps]) < 1e-4 # Terminal law is Gaussian. + thousand = samplers.make_schedule(1000) + assert float(thousand["alphas_cumprod"][207].sqrt()) == pytest.approx(0.8, abs=5e-3) + + +def test_every_sampler_finds_all_four_modes(): + budget = 60 + for name, points in ( + ("ancestral", samplers.ancestral_sample( + samplers.ExactScore(), samplers.make_schedule(budget), 512, + torch.Generator().manual_seed(0))), + ("flow", samplers.probability_flow_sample( + samplers.ExactScore(), 512, budget, torch.Generator().manual_seed(0))), + ("langevin", samplers.annealed_langevin_sample( + samplers.ExactScore(), 512, [0.7 - 0.69 * i / 19 for i in range(20)], 3, 0.15, + torch.Generator().manual_seed(0))), + ): + measured = mode_report(points) + assert measured["modes_found"] == 4, name + assert measured["mode_entropy"] > 1.30, name # log 4 = 1.386 when the modes are even. + + +def test_each_sampler_reports_the_evaluations_it_spent(): + score = samplers.ExactScore() + samplers.probability_flow_sample(score, 64, 37, torch.Generator().manual_seed(0)) + assert score.calls == 37 + score = samplers.ExactScore() + samplers.annealed_langevin_sample(score, 64, [0.5, 0.2, 0.05], 4, 0.1, + torch.Generator().manual_seed(0)) + assert score.calls == 12 # Three levels, four iterations each. + + +# %% Guidance, Sections 3.37 through 3.45 + + +def test_worked_guided_score(): + score, gradient = torch.tensor([[-1.0, -0.1]]), torch.tensor([[1.3, 1.1]]) + at_one = guidance.guided_score(score, gradient, 1.0)[0] + at_two = guidance.guided_score(score, gradient, 2.0)[0] + assert at_one.tolist() == pytest.approx([0.3, 1.0], abs=1e-6) + assert float(at_one.norm()) == pytest.approx(1.0440, abs=1e-4) + assert at_two.tolist() == pytest.approx([1.6, 2.1], abs=1e-6) + assert float(at_two.norm()) == pytest.approx(2.6401, abs=1e-4) + + +def test_classifier_free_step_and_the_score_it_implies(): + eps = guidance.classifier_free_field(torch.tensor(0.50), torch.tensor(0.10), 2.0) + assert float(eps) == pytest.approx(-0.30, abs=1e-6) + assert float(-eps / 0.5) == pytest.approx(0.6, abs=1e-6) + + +def test_guided_backward_step_agrees_in_both_parameterizations(): + guided_eps = float(guidance.guided_noise(torch.tensor(0.8), torch.tensor(0.6), 2.0, 0.5)) + assert guided_eps == pytest.approx(0.2, abs=1e-6) + velocity = -0.5 * 2.0 * (1.2 + (-guided_eps / 0.5)) # The collected form of the velocity. + assert velocity == pytest.approx(-0.8, abs=1e-6) + assert 1.2 - 0.1 * velocity == pytest.approx(1.28, abs=1e-6) + + +def test_classifier_gradient_equals_the_difference_of_the_two_fields(): + points = torch.randn(64, 2, generator=torch.Generator().manual_seed(0)) + gradient = guidance.classifier_gradient(points, 0.8, 0) + difference = (guidance.conditional_score(points, 0.8, 0) + - guidance.unconditional_score(points, 0.8)) + assert torch.allclose(gradient, difference, atol=1e-6) + + +def test_weights_are_projected_onto_the_simplex(): + weights = guidance.normalize_weights((-1.0, 3.0, 1.0), torch.device("cpu")) + assert float(weights.sum()) == pytest.approx(1.0) + assert float(weights[0]) == 0.0 # A negative weight is clamped away. + + +def test_scalarized_rewards_reverse_the_ranking_with_the_weights(): + first, second = (1.2, 0.2), (0.5, 1.1) + assert guidance.scalarize(first, (0.7, 0.3)) == pytest.approx(0.90, abs=1e-4) + assert guidance.scalarize(second, (0.7, 0.3)) == pytest.approx(0.68, abs=1e-4) + assert guidance.scalarize(first, (0.3, 0.7)) == pytest.approx(0.50, abs=1e-4) + assert guidance.scalarize(second, (0.3, 0.7)) == pytest.approx(0.92, abs=1e-4) + + +def test_competing_gradients_cancel_at_the_stated_weights(): + combined = 0.6 * torch.tensor([1.0, 0.8]) + 0.4 * torch.tensor([0.5, -1.2]) + assert combined.tolist() == pytest.approx([0.80, 0.0], abs=1e-6) + + +def test_selection_weights_and_the_guidance_schedule(): + weights = (2.0 * torch.tensor([1.2, 0.8, 0.3])).softmax(0) + assert weights.tolist() == pytest.approx([0.6193, 0.2783, 0.1024], abs=1e-4) + assert guidance.schedule_value(0.6) == pytest.approx(0.96) + assert guidance.schedule_value(0.0) == 0.0 and guidance.schedule_value(1.0) == 0.0 + + +def test_guidance_moves_mass_onto_the_target_mode(): + schedule = guidance.schedule_tensors(60) + fractions = [] + for w in (0.0, 0.25, 1.0): + points = guidance.sample_diffusion( + schedule, 512, guidance.TARGET_LABEL, w=w, + generator=torch.Generator().manual_seed(0), + ) + fractions.append(mode_report(points)["mode_fractions"][guidance.TARGET_LABEL]) + assert fractions[0] < 0.35 # Unconditional, so about a quarter. + assert fractions[0] < fractions[1] < fractions[2] + assert fractions[2] > 0.9 + + +def test_reward_guidance_prefers_the_right_hand_modes(): + schedule = guidance.schedule_tensors(60) + points = guidance.sample_diffusion( + schedule, 512, guidance.TARGET_LABEL, eta=2.0, lambdas=(1.0, 0.0), + generator=torch.Generator().manual_seed(0), + ) + fractions = mode_report(points)["mode_fractions"] + assert fractions[0] + fractions[1] > 0.9 # The two modes at x = +2. + + +# %% Latent diffusion, Sections 3.21, 3.29 and 3.46 + + +def test_latent_coordinate_corrupted_and_carried_back(): + abar = float(latent_diffusion.alphas_cumprod[500]) + a_500, sigma_500 = math.sqrt(abar), math.sqrt(1.0 - abar) + assert (a_500, sigma_500) == pytest.approx((0.2803, 0.9599), abs=1e-4) + z_500 = a_500 * 1.2 + sigma_500 * 0.4 + assert z_500 == pytest.approx(0.7204, abs=1e-4) + eps_hat = (z_500 - a_500 * 1.05) / sigma_500 + assert eps_hat == pytest.approx(0.4438, abs=1e-4) + beta_k, alpha_k = float(latent_diffusion.betas[500]), float(latent_diffusion.alphas[500]) + mean = (z_500 - beta_k / sigma_500 * eps_hat) / math.sqrt(alpha_k) + assert mean == pytest.approx(0.7193, abs=1e-4) + variance = float(latent_diffusion.posterior_variance[500]) + assert math.sqrt(variance) == pytest.approx(0.1002, abs=1e-4) + + +def test_masked_corruption_leaves_padding_at_zero(): + z0 = torch.randn(3, 8, 16) + mask = torch.ones(3, 8) + mask[:, 5:] = 0.0 # The last three positions are padding. + z_k, noise = latent_diffusion.q_sample_latent(z0, torch.tensor([1, 500, 1000]), mask) + assert torch.count_nonzero(z_k[:, 5:, :]) == 0 + assert noise.shape == z0.shape + + +def test_latent_loss_counts_valid_residues_only(): + z0 = torch.zeros(2, 4, 3) + z0_hat = torch.ones(2, 4, 3) # Unit error at every coordinate. + mask = torch.tensor([[1.0, 1.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0]]) + assert float(latent_diffusion.latent_loss(z0_hat, z0, mask)) == pytest.approx(1.0) + z0_hat[:, 2:, :] = 100.0 # A large error, but only on padding. + assert float(latent_diffusion.latent_loss(z0_hat, z0, mask)) == pytest.approx(1.0) + + +def test_denoiser_returns_a_latent_of_the_same_shape(): + torch.manual_seed(0) + model = latent_diffusion.LatentDenoiser(embed_dim=16, num_layers=1, num_heads=2, max_len=8) + z_k = torch.randn(2, 8, 16) + mask = torch.ones(2, 8) + assert model(z_k, torch.tensor([7, 900]), mask).shape == z_k.shape + + +def test_composition_proxies_on_a_known_peptide(): + charge, polar = latent_diffusion.composition_proxies("KRDE") + assert charge == pytest.approx(0.0) # Two positive and two negative over four residues. + assert polar == pytest.approx(1.0) # All four are in the polar set. + assert latent_diffusion.composition_proxies("AAAA") == pytest.approx((0.0, 0.0)) + + +def test_enforce_polar_reaches_the_minimum_at_the_cheapest_positions(): + torch.manual_seed(0) + logits = torch.randn(2, 12, 20) + logits[:, :, latent_diffusion.POLAR_COLUMNS] -= 5.0 # Make the argmax non-polar everywhere. + choices = logits.argmax(dim=-1) + sequences = latent_diffusion.enforce_polar(logits, choices, 5) + for sequence in sequences: + assert len(sequence) == 12 + assert sum(a in latent_diffusion.POLAR_RESIDUES for a in sequence) >= 5 + unchanged = latent_diffusion.enforce_polar(logits, choices, 0) # Nothing is missing. + assert unchanged == ["".join(latent_diffusion.AMINO_ACIDS[int(i)] for i in row) + for row in choices] + + +def test_the_cheapest_two_substitutions_cost_what_the_notes_says(): + costs = torch.tensor([0.9, 1.4, 2.6, 3.1, 3.3, 4.0, 4.2, 5.5, 6.1]) + cheapest, _ = costs.topk(2, largest=False) + assert cheapest.tolist() == pytest.approx([0.9, 1.4]) + assert float(cheapest.sum()) == pytest.approx(2.3, abs=1e-5) diff --git a/tests/test_lecture_4.py b/tests/test_lecture_4.py new file mode 100644 index 0000000000000000000000000000000000000000..9605dda9cfa7b9ea6c6e3c4ad41597db3cbeb2ad --- /dev/null +++ b/tests/test_lecture_4.py @@ -0,0 +1,535 @@ +"""Lecture 4: the worked numbers of Chapter 4, asserted against the files that print them. + +Every assertion here is a number the chapter computes by hand. The tests import the lecture files +directly, so they check the same functions the scripts run, and none of them trains anything. +""" + +from __future__ import annotations + +import importlib.util +import math +import sys +from pathlib import Path + +import pytest +import torch + +LECTURE = Path(__file__).resolve().parents[1] / "lecture_4" + + +def load(name: str): + """Import one lecture file by path, since lecture_4 is a folder of scripts and not a package.""" + spec = importlib.util.spec_from_file_location(f"lecture_4_{name}", LECTURE / f"{name}.py") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +ctmc = load("ctmc") +mdlm = load("mdlm") +uniform = load("uniform") +d3pm = load("d3pm") +sedd = load("sedd") +block = load("block") +guidance = load("guidance") +planner = load("planner") +search = load("search") +simplex = load("simplex") + + +# %% 01. ctmc: rates, the master equation, the exponential and the reversal + + +def test_two_rules_and_the_dna_row(): + """A jumps to C at rate 2 and to G and T at rate 1, so its diagonal entry is -4.""" + assert ctmc.obeys_the_two_rules(ctmc.DNA_RATES) + assert float(ctmc.DNA_RATES[0, 0]) == pytest.approx(-4.0) + assert float(ctmc.exit_rates(ctmc.DNA_RATES)[0]) == pytest.approx(4.0) + + +def test_short_interval_row_and_the_step_bound(): + """P_{t,t+0.1}(A, .) = (0.6, 0.2, 0.1, 0.1), and the step must satisfy h <= 0.25.""" + row = ctmc.short_interval_row(ctmc.DNA_RATES, 0, 0.1) + assert [round(float(v), 4) for v in row] == [0.6, 0.2, 0.1, 0.1] + assert math.exp(-0.4) == pytest.approx(0.6703, abs=5e-5) + assert ctmc.largest_valid_step(ctmc.DNA_RATES) == pytest.approx(0.25) + assert float(ctmc.short_interval_row(ctmc.DNA_RATES, 0, 0.3)[0]) == pytest.approx(-0.2) + + +def test_waiting_time_from_the_chapters_uniform(): + """Inverting S(s) = exp(-4s) at u = 0.42 gives a wait of 0.2169.""" + assert -math.log(0.42) / 4.0 == pytest.approx(0.2169, abs=5e-5) + assert math.log(2.0) / 4.0 == pytest.approx(0.1733, abs=5e-5) + + +def test_master_equation_on_four_dna_probabilities(): + """p_dot = (0.6, 0.2, -0.2, -0.6) and one Euler step gives (0.16, 0.22, 0.28, 0.34).""" + derivative = ctmc.master_equation(ctmc.START, ctmc.UNIT_RATES) + assert [round(float(v), 4) for v in derivative] == [0.6, 0.2, -0.2, -0.6] + arriving, leaving = ctmc.arrivals_and_departures(ctmc.START, ctmc.UNIT_RATES) + assert float(arriving[0]) == pytest.approx(0.9) + assert float(leaving[0]) == pytest.approx(0.3) + stepped = ctmc.START + 0.1 * derivative + assert [round(float(v), 4) for v in stepped] == [0.16, 0.22, 0.28, 0.34] + assert float(derivative.sum()) == pytest.approx(0.0, abs=1e-12) + + +def test_exact_evolution_matches_the_matrix_exponential(): + """1/4 + (p_0 - 1/4) exp(-4t) agrees with exp(tR) at t = 0.1 and at t = 1.""" + for t, expected in ((0.1, [0.1495, 0.2165, 0.2835, 0.3505]), + (1.0, [0.2473, 0.2491, 0.2509, 0.2527])): + exact = ctmc.exact_evolution(ctmc.START, t) + assert [round(float(v), 4) for v in exact] == expected + through_matrix = ctmc.START @ ctmc.transition_matrix(ctmc.UNIT_RATES, t) + assert float((exact - through_matrix).abs().max()) == pytest.approx(0.0, abs=1e-12) + + +def test_the_matrix_exponential_worked_for_dna(): + """At tau = log 2 the uniform kernel is 0.625 on the diagonal and 0.125 elsewhere.""" + pi = torch.full((4,), 0.25, dtype=torch.float64) + tau = math.log(2.0) + kernel = ctmc.transition_matrix(ctmc.r_pi(pi), tau) + assert float(kernel[0, 0]) == pytest.approx(0.625, abs=1e-12) + assert float(kernel[0, 1]) == pytest.approx(0.125, abs=1e-12) + assert float(kernel.sum(dim=1).min()) == pytest.approx(1.0) + closed = ctmc.pi_kernel(pi, math.exp(-tau)) + assert float((kernel - closed).abs().max()) == pytest.approx(0.0, abs=1e-12) + euler = {n: float(ctmc.euler_product(ctmc.r_pi(pi), tau, n)[0, 0]) for n in (1, 2, 10)} + assert [round(euler[n], 3) for n in (1, 2, 10)] == [0.480, 0.570, 0.616] + + +def test_masking_kernel_is_absorbing(): + """At alpha = 1/2 the masking kernel keeps each base with 0.5 and sends 0.5 to the mask.""" + mask_pi = torch.zeros(5, dtype=torch.float64) + mask_pi[4] = 1.0 + kernel = ctmc.pi_kernel(mask_pi, 0.5) + assert float(kernel[0, 0]) == pytest.approx(0.5) + assert float(kernel[0, 4]) == pytest.approx(0.5) + assert float(kernel[4, 4]) == pytest.approx(1.0) + + +def test_composition_eigenvalues_and_detailed_balance(): + """Intervals compose, R_pi has eigenvalues 0 and -1, and the cycle breaks detailed balance.""" + whole = ctmc.transition_matrix(ctmc.UNIT_RATES, 0.7) + halves = ctmc.transition_matrix(ctmc.UNIT_RATES, 0.3) @ ctmc.transition_matrix( + ctmc.UNIT_RATES, 0.4) + assert float((whole - halves).abs().max()) == pytest.approx(0.0, abs=1e-12) + pi = torch.full((4,), 0.25, dtype=torch.float64) + eigenvalues = sorted(round(float(v.real), 6) for v in torch.linalg.eigvals(ctmc.r_pi(pi))) + assert eigenvalues == [-1.0, -1.0, -1.0, 0.0] + assert float((pi @ ctmc.r_pi(pi)).abs().max()) == pytest.approx(0.0, abs=1e-12) + assert ctmc.detailed_balance_gap(ctmc.r_pi(pi), pi) == pytest.approx(0.0, abs=1e-12) + cycle_pi = torch.full((3,), 1.0 / 3.0, dtype=torch.float64) + assert float((cycle_pi @ ctmc.CYCLE).abs().max()) == pytest.approx(0.0, abs=1e-12) + assert ctmc.detailed_balance_gap(ctmc.CYCLE, cycle_pi) == pytest.approx(1.0 / 3.0) + + +def test_time_reversal_retraces_the_marginals(): + """With p_t = (0.5, 0.25, 0.125, 0.125) the reverse rate out of C toward A is 2.""" + marginal = torch.tensor([0.5, 0.25, 0.125, 0.125], dtype=torch.float64) + backward = ctmc.reverse_rates(ctmc.UNIT_RATES, marginal) + assert float(backward[1, 0]) == pytest.approx(2.0) + assert float(backward[1, 2]) == pytest.approx(0.5) + assert float(ctmc.exit_rates(backward)[1]) == pytest.approx(3.0) + net = marginal @ backward + marginal @ ctmc.UNIT_RATES + assert float(net.abs().max()) == pytest.approx(0.0, abs=1e-12) + + +# %% 02. mdlm: the corruption, the loss and the sampler + + +def test_reverse_step_worked(): + """At alpha_s = 0.7 and alpha_t = 0.4 the reveal probability is 0.5.""" + assert mdlm.reveal_probability(0.7, 0.4) == pytest.approx(0.5) + law = mdlm.reverse_step_distribution(0.7, 0.4, [0.6, 0.4, 0.0, 0.0]) + assert [round(v, 4) for v in law] == [0.3, 0.2, 0.0, 0.0, 0.5] + assert sum(law) == pytest.approx(1.0) + + +def test_one_term_of_the_loss_worked(): + """w(0.5) = 2 and the masked cross-entropy of 1.04982 give a weighted term of 2.09964.""" + assert mdlm.reveal_rate(0.5) == pytest.approx(2.0) + masked = -math.log(0.7) - math.log(0.5) + assert masked == pytest.approx(1.04982, abs=5e-6) + assert 2.0 * masked == pytest.approx(2.09964, abs=5e-6) + + +def test_the_rate_error_collapses_to_the_cross_entropy(): + """At a masked position the rate errors sum to w(t) times the negative log probability.""" + rates = [2.0 * d for d in (0.2, 0.6, 0.1, 0.1)] + assert [round(v, 4) for v in rates] == [0.4, 1.2, 0.2, 0.2] + total = mdlm.rate_error(2.0, rates[1]) + sum( + mdlm.rate_error(0.0, r) for i, r in enumerate(rates) if i != 1) + assert total == pytest.approx(2.0 * -math.log(0.6), abs=1e-9) + assert total == pytest.approx(1.0217, abs=5e-5) + assert 2.0 * -math.log(0.4) == pytest.approx(1.8326, abs=5e-5) + + +def test_corruption_and_sampling_shapes(): + """The corruption keeps or masks, and the sampler returns clean DNA with no mask left.""" + torch.manual_seed(0) + x0 = torch.randint(0, 4, (8, 6)) + zt = mdlm.corrupt(x0, torch.full((8, 1), 0.5)) + assert bool(((zt == x0) | (zt == mdlm.MASK)).all()) + kept = mdlm.corrupt(x0, torch.zeros(8, 1)).eq(x0).to(torch.get_default_dtype()) + assert float(kept.mean()) == 1.0 + assert int(mdlm.corrupt(x0, torch.ones(8, 1)).eq(mdlm.MASK).all()) + model = _tiny_denoiser(5, 6) + drawn = mdlm.sample(model, 4, 6, steps=4) + assert drawn.shape == (4, 6) and int(drawn.max()) < 4 + + +def _tiny_denoiser(vocab: int, length: int, num_classes: int = 0): + """A very small SequenceDenoiser, so the sampler tests cost nothing.""" + from cis6270.nets import SequenceDenoiser + + torch.manual_seed(0) + return SequenceDenoiser(vocab, length, width=16, depth=1, heads=2, num_classes=num_classes) + + +# %% 03. uniform: the posterior average of 2, and the rate error that keeps its staying term + + +def test_the_uniform_kernel_and_the_clean_factor(): + """At alpha_t = 1/2 the kernel is 0.625 at the clean base and 0.125 elsewhere.""" + assert float(uniform.kernel(0.5)[0, 0]) == pytest.approx(0.625) + assert float(uniform.kernel(0.5)[0, 1]) == pytest.approx(0.125) + assert uniform.clean_factor(0.5) == pytest.approx(5.0) + + +def test_the_worked_posterior_average_is_two(): + """The chapter's posterior average of the conditional reverse rates on one base comes out 2.""" + prior = torch.tensor([0.75, 0.25, 0.0, 0.0], dtype=torch.float64) + posterior = uniform.posterior_over_clean(0.5, 1, prior) + assert [round(float(v), 4) for v in posterior] == [0.375, 0.625, 0.0, 0.0] + assert float(uniform.conditional_rates(0.5, current=1, clean=0)[0]) == pytest.approx(5.0) + assert float(uniform.conditional_rates(0.5, current=1, clean=1)[0]) == pytest.approx(0.2) + assert uniform.marginal_reverse_rate(0.5, 1, 0, prior) == pytest.approx(2.0) + assert uniform.marginal_reverse_rate(0.5, 1, 2, prior) == pytest.approx(0.5) + marginal = prior @ uniform.kernel(0.5) + assert [round(float(v), 4) for v in marginal] == [0.5, 0.25, 0.125, 0.125] + assert float(marginal[0] / marginal[1]) == pytest.approx(2.0) + + +def test_the_rate_error_on_one_base(): + """A denoiser of (0.5, 0.3, 0.1, 0.1) gives rates 2.76 and 1.16 and a rate error of 0.754.""" + denoiser = torch.tensor([[[0.5, 0.3, 0.1, 0.1]]], dtype=torch.float64) + current = torch.tensor([[1]]) + alpha = torch.tensor([[0.5]], dtype=torch.float64) + rates = uniform.model_rates(denoiser, uniform.ratio_table(current, alpha))[0, 0] + rates = rates / (uniform.BETA / 4) # Quote the rates in the chapter's unit forward rate. + assert [round(float(v), 4) for v in rates] == [2.76, 1.0, 1.16, 1.16] + target = uniform.conditional_rates(0.5, current=1, clean=0) + error = float(uniform.rate_error(target[[0, 2, 3]], rates[[0, 2, 3]]).sum()) + assert error == pytest.approx(0.754, abs=5e-4) + assert float(uniform.rate_error(torch.tensor(5.0), torch.tensor(2.76))) == pytest.approx( + 0.731, abs=5e-4) + + +def test_the_rate_error_is_minimized_at_the_posterior_average(): + """0.375 l(5, B) + 0.625 l(0.2, B) is smallest at B = 2, where it equals 1.430.""" + def mixture(b: float) -> float: + return float(0.375 * uniform.rate_error(torch.tensor(5.0), torch.tensor(b)) + + 0.625 * uniform.rate_error(torch.tensor(0.2), torch.tensor(b))) + + assert mixture(2.0) == pytest.approx(1.430, abs=5e-4) + assert mixture(5.0) == pytest.approx(2.598, abs=5e-4) + assert mixture(0.2) == pytest.approx(4.235, abs=5e-4) + assert mixture(2.0) < min(mixture(1.9), mixture(2.1)) + + +# %% 04. d3pm: the finite step + + +def test_cumulative_transition_matrix(): + """With beta_1 = 0.2 and beta_2 = 0.3 the product is 0.67 on the diagonal and 0.11 off it.""" + product = d3pm.one_step(0.2) @ d3pm.one_step(0.3) + assert float(product[0, 0]) == pytest.approx(0.67, abs=5e-3) + assert float(product[0, 1]) == pytest.approx(0.11, abs=5e-3) + assert float((product - d3pm.cumulative(0.56)).abs().max()) == pytest.approx(0.0, abs=1e-6) + assert (1.0 - 1.0 / 1000) ** 1000 == pytest.approx(0.3677, abs=5e-5) + + +def test_reversing_one_finite_step(): + """Observing C under a clean A gives (0.5795, 0.3523, 0.0341, 0.0341) at the earlier step.""" + first, second = d3pm.one_step(0.2), d3pm.one_step(0.3) + clean_a = torch.tensor([[[1.0, 0.0, 0.0, 0.0]]]) + posterior = d3pm.exact_posterior(torch.tensor([[1]]), clean_a, second[None], first[None])[0, 0] + assert [round(float(v), 4) for v in posterior] == [0.5795, 0.3523, 0.0341, 0.0341] + + +def test_rate_matrices_chosen_for_the_data(): + """The Gaussian row is 0.4088 to a neighbour, and exp(0.3 R) has first row (0.6747, ...).""" + row = d3pm.gaussian_row() + assert float(row[1]) == pytest.approx(0.4088, abs=5e-5) + assert float(row[0]) == pytest.approx(0.0912, abs=5e-5) + adjacency = torch.tensor([[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 = d3pm.graph_rates(adjacency) + assert [round(float(v), 4) for v in R[0]] == [-1.5, 1.0, 0.5, 0.0] + kernel = torch.matrix_exp(0.3 * R) + assert [round(float(v), 4) for v in kernel[0]] == [0.6747, 0.1964, 0.1122, 0.0167] + + +# %% 05. sedd: score entropy + + +def test_score_entropy_at_three_predictions(): + """With the target ratio 2, a prediction of 1.8 costs 0.0107 nats and 1.0 costs 0.3863.""" + target = torch.tensor(2.0) + assert float(sedd.score_entropy(target, torch.tensor(1.8))) == pytest.approx(0.0107, abs=5e-5) + assert float(sedd.score_entropy(target, torch.tensor(1.0))) == pytest.approx(0.3863, abs=5e-5) + assert float(sedd.score_entropy(target, torch.tensor(3.0))) == pytest.approx(0.1891, abs=5e-5) + assert float(sedd.score_entropy(target, target)) == pytest.approx(0.0, abs=1e-6) + + +def test_the_target_ratio_from_the_denoiser_route(): + """0.375(5) + 0.625(0.2) is 2, and the kernel ratio reproduces it at a clean A.""" + assert 0.375 * 5.0 + 0.625 * 0.2 == pytest.approx(2.0) + x0 = torch.tensor([[0]]) # A clean A. + z = torch.tensor([[1]]) # Observed as a C. + ratio = sedd.kernel_ratio(x0, z, torch.tensor([[0.5]])) + assert float(ratio[0, 0, 0]) == pytest.approx(5.0) + assert float(ratio[0, 0, 2]) == pytest.approx(1.0) + + +# %% 06. block: the dial between diffusion and autoregression + + +def test_the_worked_block_bound(): + """ACGTAC in blocks of two at t = 0.5 gives 1.8326, 0.7134 and 0.2107, summing to 2.7567.""" + weight = 2.0 + contributions = [weight * (-math.log(0.5) - math.log(0.8)), + weight * -math.log(0.7), weight * -math.log(0.9)] + assert [round(v, 4) for v in contributions] == [1.8326, 0.7133, 0.2107] + assert sum(contributions) == pytest.approx(2.7567, abs=5e-5) + + +def test_the_one_position_identity(): + """At blocks of one position the time integral of w(t)(1 - alpha_t) is exactly one.""" + assert block.one_position_integral() == pytest.approx(1.0, abs=1e-9) + generator = torch.Generator().manual_seed(0) + estimate = block.one_position_estimate(0.7, 40000, generator) + assert estimate == pytest.approx(-math.log(0.7), abs=0.05) + + +def test_block_input_keeps_the_prefix_and_hides_the_future(): + """The prefix stays clean, the current block is corrupted, and the suffix is all mask.""" + x0 = torch.arange(8).remainder(4)[None, :].repeat(2, 1) # [2, 8]. + z = block.block_input(x0, torch.tensor([1, 1]), 2, torch.full((2, 1), 1.0)) + assert bool((z[:, :2] == x0[:, :2]).all()) # The prefix is clean. + assert bool((z[:, 2:4] == block.MASK).all()) # The block is fully masked at t = 1. + assert bool((z[:, 4:] == block.MASK).all()) # The future is hidden. + + +# %% 07. guidance: the h-transform and the worked guided jump + + +def test_h_transform_keeps_a_rate_matrix(): + """Reweighting by h(y)/h(x) leaves the row sums at zero, and a constant h changes nothing.""" + R = torch.ones(4, 4) - torch.diag(torch.full((4,), 4.0)) + transformed = guidance.h_transform(R, torch.tensor([2.0, 1.0, 0.5, 0.5])) + assert float(transformed.sum(dim=1).abs().max()) == pytest.approx(0.0, abs=1e-6) + constant = guidance.h_transform(R, torch.full((4,), 3.0)) + assert float((constant - R).abs().max()) == pytest.approx(0.0, abs=1e-6) + + +def test_the_reverse_process_as_an_h_transform(): + """h_t = p_t/pi is (2, 1, 0.5, 0.5), and the rates out of C come back as 2, 0.5 and 0.5.""" + pi = torch.full((4,), 0.25) + values = guidance.value_from_marginal(torch.tensor([0.5, 0.25, 0.125, 0.125]), pi) + assert [round(float(v), 4) for v in values] == [2.0, 1.0, 0.5, 0.5] + R = torch.ones(4, 4) - torch.diag(torch.full((4,), 4.0)) + reverse = guidance.h_transform(R, values) + assert float(reverse[1, 0]) == pytest.approx(2.0, abs=1e-5) + assert float(reverse[1, 2]) == pytest.approx(0.5, abs=1e-5) + clean_a = guidance.value_from_marginal(torch.tensor([0.625, 0.125, 0.125, 0.125]), pi) + assert float(clean_a[0] / clean_a[1]) == pytest.approx(5.0) + assert float(clean_a[1] / clean_a[0]) == pytest.approx(0.2) + + +def test_the_worked_guided_jump(): + """The classifier turns the rates 2, 0.5, 0.5 into 1, 1.5, 0.5, and gamma = 2 sharpens them.""" + unguided = torch.tensor([2.0, 0.5, 0.5]) + ratio = torch.tensor([0.1, 0.6, 0.2]) / 0.2 + guided = unguided * ratio + assert [round(float(v), 4) for v in guided] == [1.0, 1.5, 0.5] + assert float(guided.sum()) == pytest.approx(3.0) + shares = guided / guided.sum() + assert [float(v) for v in shares] == pytest.approx([1 / 3, 1 / 2, 1 / 6], abs=1e-6) + sharper = unguided * ratio ** 2 + assert [round(float(v), 4) for v in sharper] == [0.5, 4.5, 0.5] + assert float(sharper.sum()) == pytest.approx(5.5) + assert float((sharper / sharper.sum())[1]) == pytest.approx(0.8182, abs=5e-5) + assert 1.0 / float(sharper.sum()) == pytest.approx(0.1818, abs=5e-5) + + +def test_exact_and_gradient_ratios_agree_at_first_order(): + """Both estimates of p(c | z^{i <- a})/p(c | z) have the right shape and a positive value.""" + torch.manual_seed(0) + classifier = guidance.NoisyClassifier(6, width=16) + z = torch.tensor([[0, guidance.MASK, 2, guidance.MASK, 1, 3]]) + t = torch.tensor([0.5]) + with torch.no_grad(): + exact = guidance.exact_ratio(classifier, z, t) + approximate = guidance.gradient_ratio(classifier, z, t) + assert exact.shape == approximate.shape == (1, 6, 4) + assert bool((exact > 0).all()) and bool((approximate > 0).all()) + # The ratio at the token already in place is one under both estimates. + assert float(exact[0, 0, 0]) == pytest.approx(1.0, abs=1e-5) + assert float(approximate[0, 0, 0]) == pytest.approx(1.0, abs=1e-5) + + +# %% 08. planner: the order of the reveals and the weights it puts in the loss + + +def test_the_worked_planner_selection(): + """Scores (0.42, 0.91, 0.55, 0.78) at the masked positions reveal 3 and 6, and remask 5.""" + scores = torch.tensor([[0.42, 0.88, 0.91, 0.55, 0.31, 0.78]]) + masked = torch.tensor([[True, False, True, True, False, True]]) + revealed = sorted(int(i) + 1 for i in planner.topk_where(scores, masked, 2).nonzero()[:, 1]) + returned = sorted(int(i) + 1 for i in planner.topk_where(-scores, ~masked, 1).nonzero()[:, 1]) + assert revealed == [3, 6] + assert returned == [5] + assert int(planner.topk_where(scores, masked, 0).sum()) == 0 + + +def test_the_corrector_fluxes_match(): + """alpha_t beta(t) and (1 - alpha_t) w(t) both equal -alpha_dot, which is one here.""" + for t in (0.5, 0.25, 0.75): + forward, reverse = planner.corrector_fluxes(t) + assert forward == pytest.approx(1.0) + assert reverse == pytest.approx(1.0) + + +def test_the_worked_papl_weights(): + """A planner distribution of (0.10, 0.50, 0.15, 0.25) turns 2.9 nats into 2.12.""" + rho = torch.tensor([[0.10, 0.50, 0.15, 0.25]]) + cross_entropies = torch.tensor([[1.2, 0.3, 0.9, 0.5]]) + lam = 4.0 * rho + assert [round(float(v), 4) for v in lam[0]] == [0.4, 2.0, 0.6, 1.0] + assert float(lam.sum()) == pytest.approx(4.0) + assert float(cross_entropies.sum()) == pytest.approx(2.9) + assert float((lam * cross_entropies).sum()) == pytest.approx(2.12) + + +def test_planner_weights_average_to_one_and_reduce_to_mdlm(): + """lambda_i averages to one over the masked positions, and strength zero gives every weight one.""" + score = torch.tensor([[0.9, 0.1, 0.5, 0.2]]) + masked = torch.tensor([[True, True, False, True]]) + weights = planner.planner_weights(score, masked, 1.0) + assert float((weights * masked).sum()) == pytest.approx(float(masked.sum()), abs=1e-5) + plain = planner.planner_weights(score, masked, 0.0) + assert float((plain - 1.0).abs().max()) == pytest.approx(0.0, abs=1e-6) + + +# %% 09. simplex: the Dirichlet belief and the thinning representation + + +def test_the_worked_dirichlet_parameters(): + """At alpha_t = 1/2 the parameters are (5, 1, 1, 1) and at alpha_s = 0.8 they are (17, 1, 1, 1).""" + assert simplex.concentration(0.5) == pytest.approx(8.0) + assert simplex.concentration(0.8) == pytest.approx(20.0) + mean_t = simplex.belief_mean(0.5, 0) + assert [round(float(v), 4) for v in mean_t] == [0.625, 0.125, 0.125, 0.125] + assert [round(float(v), 4) for v in simplex.belief_mean(0.8, 0)] == [0.85, 0.05, 0.05, 0.05] + at_t = simplex.dirichlet_parameters(0.5, 0) + at_s = simplex.dirichlet_parameters(0.8, 0) + assert [round(float(v), 4) for v in at_t] == [5.0, 1.0, 1.0, 1.0] + assert [round(float(v), 4) for v in at_s] == [17.0, 1.0, 1.0, 1.0] + total = at_t.sum() + deviation = (at_t * (total - at_t) / (total ** 2 * (total + 1))).sqrt() + assert float(deviation[0]) == pytest.approx(0.161, abs=5e-4) + + +def test_the_increment_sits_on_the_clean_coordinate(): + """The increment is (12, 0, 0, 0), so the innovation is the one-hot vector of the clean base.""" + step = simplex.increment(0.5, 0.8, 0) + assert [round(float(v), 4) for v in step] == [12.0, 0.0, 0.0, 0.0] + first, second = simplex.thinning_law(0.5, 0.8) + assert (round(first, 4), round(second, 4)) == (8.0, 12.0) + weight = first / (first + second) + assert weight == pytest.approx(0.4) + mixed = weight * simplex.belief_mean(0.5, 0) + (1.0 - weight) * torch.eye(4)[0] + assert [round(float(v), 4) for v in mixed] == [0.85, 0.05, 0.05, 0.05] + + +def test_the_thinning_representation_reproduces_the_dirichlet(): + """W V + (1 - W) onehot(A) has the mean and the spread of Dir(17, 1, 1, 1).""" + torch.manual_seed(0) + at_t, at_s = simplex.dirichlet_parameters(0.5, 0), simplex.dirichlet_parameters(0.8, 0) + first, second = simplex.thinning_law(0.5, 0.8) + draws = 200_000 + thinned = torch.distributions.Dirichlet(at_t).sample((draws,)) + weight = torch.distributions.Beta(torch.tensor(first), torch.tensor(second)).sample((draws,)) + drawn = weight[:, None] * thinned + (1.0 - weight[:, None]) * torch.eye(4)[0] + total = at_s.sum() + exact = (at_s * (total - at_s) / (total ** 2 * (total + 1))).sqrt() + assert [round(float(v), 4) for v in exact] == [0.0779, 0.0476, 0.0476, 0.0476] + assert float(drawn.mean(0)[0]) == pytest.approx(0.85, abs=2e-3) + assert float(drawn.std(0)[0]) == pytest.approx(float(exact[0]), abs=2e-3) + + +def test_the_parameter_field_and_the_reverse_step(): + """The batched parameters match the per-position ones, and one step stays on the simplex.""" + x0 = torch.tensor([[0, 1], [2, 3]]) + field = simplex.parameter_field(x0, torch.full((2, 1), 0.5)) + assert [round(float(v), 4) for v in field[0, 0]] == [5.0, 1.0, 1.0, 1.0] + assert [round(float(v), 4) for v in field[1, 1]] == [1.0, 1.0, 1.0, 5.0] + torch.manual_seed(0) + belief = simplex.corrupt(x0, torch.full((2, 1), 0.5)) + assert belief.shape == (2, 2, 4) + assert float(belief.sum(-1).sub(1.0).abs().max()) == pytest.approx(0.0, abs=1e-6) + moved = simplex.reverse_step(belief, x0, 0.5, 0.8) + assert float(moved.sum(-1).sub(1.0).abs().max()) == pytest.approx(0.0, abs=1e-6) + assert bool((moved >= 0).all()) + + +def test_a_one_hot_belief_recovers_the_token_denoiser(): + """The only architectural change is at the input, so one-hot beliefs give identical logits.""" + torch.manual_seed(0) + model = simplex.BeliefDenoiser(6, width=16) + tokens = torch.tensor([[0, 1, 2, 3, 0, 1]]) + t = torch.tensor([0.4]) + one_hot = torch.nn.functional.one_hot(tokens, 4).to(torch.get_default_dtype()) + with torch.no_grad(): + through_belief = model(one_hot, t) + through_tokens = model.inner(tokens, t) + assert float((through_belief - through_tokens).abs().max()) == pytest.approx(0.0, abs=1e-6) + drawn = simplex.sample(model, 3, 6, steps=4) + assert drawn.shape == (3, 6) and int(drawn.max()) < 4 + + +# %% 10. search: rollouts, the archive and the two objectives + + +def test_the_worked_value_estimate(): + """Rollouts averaging 0.35 and 0.75 reweight a rate of 1.2 to 2.571.""" + parent = float(torch.tensor([0.1, 0.8, 0.3, 0.2]).mean()) + child = float(torch.tensor([0.7, 0.9, 0.6, 0.8]).mean()) + assert parent == pytest.approx(0.35) + assert child == pytest.approx(0.75) + assert child / parent == pytest.approx(2.143, abs=5e-4) + assert 1.2 * child / parent == pytest.approx(2.571, abs=5e-4) + assert 0.2 / parent == pytest.approx(0.571, abs=5e-4) + + +def test_the_worked_pareto_archive(): + """P3 at (0.4, 0.2) is dominated by P1 at (0.8, 0.3), and P1 and P2 both stay.""" + archive = search.ParetoArchive() + for name, scores in (((1,), (0.8, 0.3)), ((2,), (0.5, 0.7)), ((3,), (0.4, 0.2))): + archive.add(name, scores) + assert sorted(q for q, _ in archive.entries) == [(1,), (2,)] + archive.add((4,), (0.9, 0.9)) # A candidate that dominates both of them. + assert [q for q, _ in archive.entries] == [(4,)] + + +def test_the_two_objectives(): + """A GC-free sequence scores zero on the first axis, and the motif scores one on the second.""" + tokens = torch.tensor([[0, 1, 2, 3, 0, 0], [0, 0, 3, 3, 0, 3]]) + scores = search.objectives(tokens) + assert float(scores[0, 0]) == pytest.approx(2 / 6) + assert float(scores[0, 1]) == pytest.approx(1.0) # ACGT sits at the start. + assert float(scores[1, 0]) == pytest.approx(0.0) + assert float(scores[1, 1]) < 1.0 diff --git a/tests/test_lecture_5.py b/tests/test_lecture_5.py new file mode 100644 index 0000000000000000000000000000000000000000..46c8ca5e4113fc67fbe990afe9c754b2c1523dcf --- /dev/null +++ b/tests/test_lecture_5.py @@ -0,0 +1,550 @@ +"""The worked numbers of Chapter 5, asserted against the files of ``lecture_5/``. + +Every number checked here is printed in the notes. Where the chapter runs a simulation, the test +asserts the analytic value the simulation estimates, with a tolerance wide enough for the sample +size the file uses by default. +""" + +from __future__ import annotations + +import importlib.util +import math +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + +from cis6270.nets import SequenceDenoiser + +ROOT = Path(__file__).resolve().parents[1] + + +def load(name): + """Import one lecture file by path, so that ``lecture_5`` needs no package marker.""" + spec = importlib.util.spec_from_file_location(f"lecture5_{name}", ROOT / "lecture_5" / f"{name}.py") + module = importlib.util.module_from_spec(spec) + sys.path.insert(0, str(ROOT)) + spec.loader.exec_module(module) + return module + + +dfm = load("dfm") +velocity = load("velocity") +corrector = load("corrector") +kinetic = load("kinetic_optimal") +edits = load("edit_flows") +dirichlet = load("dirichlet") +fisher = load("fisher") +gumbel = load("gumbel") +guidance = load("guidance") +objectives = load("multi_objective") +rectify = load("rectify") + + +def close(values, expected, tolerance=5e-4): + """Compare a list of numbers against the chapter's, entry by entry.""" + assert len(values) == len(expected) + for got, want in zip(values, expected): + assert abs(float(got) - want) <= tolerance, (values, expected) + + +# %% Section 5.2 and 5.3: a velocity, its two rules, and the master equation + + +def test_worked_master_equation(): + u, p, p_dot, arriving, leaving = velocity.worked_master_equation() + assert velocity.obeys_the_two_rules(u) + assert torch.allclose(u, velocity.set_diagonal(u)) + close(p_dot.tolist(), [-0.1, -0.2, 0.7, -0.4]) + assert abs(float(p_dot.sum())) < 1e-9 + assert abs(arriving - 0.7) < 1e-9 and abs(leaving) < 1e-9 + close((p + 0.1 * p_dot).tolist(), [0.09, 0.18, 0.37, 0.36]) + + +def test_hazard_and_survival(): + assert abs(velocity.hazard(0.5) - 2.0) < 1e-9 + assert abs(velocity.hazard(0.9) - 10.0) < 1e-9 + assert abs(velocity.survival(4000) - 0.1) < 2e-3 # S(t) = 1 - kappa_t. + + +def test_continuous_and_discrete_side_by_side(): + p1 = torch.tensor([0.1, 0.6, 0.2, 0.1]) + moved, drawn, lam = velocity.side_by_side(2, p1, t=0.5, h=0.1) + close(moved.tolist(), [0.02, 0.12, 0.84, 0.02]) + close(drawn.tolist(), [0.02, 0.12, 0.84, 0.02]) + assert abs(lam - 0.2) < 1e-9 + + +# %% Sections 5.5 to 5.11: the mixture path, the velocity, the loss and the sampler + + +def test_dfm_worked_numbers(): + report = dfm.worked_numbers() + assert report["mixture_draw"] == "ATGA" + close(report["velocity_at_G"], [0.2, 1.2, -1.6, 0.2]) + close(report["euler_step_at_G"], [0.02, 0.12, 0.84, 0.02]) + assert abs(report["redraw_probability"] - 0.2) < 1e-9 + assert abs(report["example_loss_nats"] - 4.780) < 1e-3 + close(report["mask_posterior_C_G"], [0.625, 0.375]) + assert report["mask_posterior_unchanged"] # Proposition 6 of Discrete Flow Matching. + assert abs(report["posterior_entropy_nats"] - 0.6616) < 5e-4 + assert abs(report["kl_gap_to_uniform"] - 0.0316) < 5e-4 + + +def test_conditional_velocity_generates_the_conditional_path(): + """Section 5.6.2: the two terms of the master equation on the mixture path.""" + kappa, x0, x1 = 0.4, 0, 1 + path = torch.zeros(4, dtype=torch.double) + path[x0], path[x1] = 1.0 - kappa, kappa + total = sum(path[z] * dfm.conditional_velocity(z, x1, kappa) for z in (x0, x1)) + derivative = torch.zeros(4, dtype=torch.double) + derivative[x1], derivative[x0] = 1.0, -1.0 # kappa_dot_t times the two indicators. + assert torch.allclose(total, derivative, atol=1e-9) + + +def test_corrupt_reaches_both_endpoints(): + torch.manual_seed(0) + process = dfm.DiscreteFlowMatching(4) + x1 = torch.randint(0, 4, (64, 16)) + assert torch.equal(process.corrupt(x1, torch.ones(64)), x1) # kappa_t = 1 gives the data. + source = process.corrupt(x1, torch.zeros(64)) # kappa_t = 0 gives the source. + assert float((source == x1).double().mean()) < 0.4 # A uniform draw agrees about a quarter of the time. + halfway = process.corrupt(x1, torch.full((64,), 0.5)) + assert 0.55 < float((halfway == x1).double().mean()) < 0.75 # Heads with probability one half. + + +def test_two_sources(): + assert torch.equal(dfm.DiscreteFlowMatching(4, mask=4).source((3, 5)), + torch.full((3, 5), 4, dtype=torch.long)) # The all-mask source. + uniform = dfm.DiscreteFlowMatching(4).source((200, 20)) + assert int(uniform.min()) == 0 and int(uniform.max()) == 3 # Every base appears. + + +# %% Sections 5.13 and 5.14: zero-flux velocities and the corrector + + +def test_zero_flux_velocity_on_three_states(): + p, u, v, with_u, with_both = corrector.three_state_example() + assert torch.allclose(corrector.net_flux(v, p), torch.zeros(3), atol=1e-12) + close(with_u.tolist(), [-0.2, -0.3, 0.5]) + close(with_both.tolist(), [-0.2, -0.3, 0.5]) + assert abs(float(-u[0, 0]) - 1.0) < 1e-12 and abs(float(-(u + v)[0, 0]) - 5.0) < 1e-12 + + +def test_corrector_step_table(): + """Section 5.14.3 at t = 0.4 with h = 0.01 and the denoiser (0.1, 0.6, 0.2, 0.1).""" + p1 = torch.tensor([0.1, 0.6, 0.2, 0.1]) + expected = { + 0.0: ([0.00167, 0.01000, 0.00333, 0.00167], 0.98333), + 2.0: ([0.00300, 0.01800, 0.00600, 0.00300], 0.97000), + 10.0: ([0.00833, 0.05000, 0.01667, 0.00833], 0.91667), + } + for eta, (probabilities, stay) in expected.items(): + bar_u = corrector.corrector_velocity(p1, corrector.MASK, 0.4, eta) + close((0.01 * bar_u[:4]).tolist(), probabilities, tolerance=1e-5) + assert abs(float(1.0 + 0.01 * bar_u[corrector.MASK]) - stay) < 1e-5 + assert abs(corrector.masked_fraction_derivative(0.4, eta) + 1.0) < 1e-12 + + +def test_corrector_rates_and_admissible_steps(): + p1 = torch.tensor([0.1, 0.6, 0.2, 0.1]) + bar_u = corrector.corrector_velocity(p1, corrector.MASK, 0.4, 2.0) + assert abs(float(-bar_u[corrector.MASK]) - 3.0) < 1e-12 # (1 + 2 x 0.4)/(1 - 0.4). + visible = torch.zeros(4) + visible[1] = 1.0 + assert abs(float(corrector.corrector_velocity(visible, 1, 0.4, 2.0)[corrector.MASK]) - 2.0) < 1e-12 + close([corrector.admissible_step(0.9, eta) for eta in (0.0, 2.0, 10.0)], + [0.1, 0.0357, 0.01], tolerance=1e-4) + + +def test_jump_counts_are_one_plus_eta(): + """Section 5.14.3: the masked fraction stays at 1 - t and the jumps come to 1 + eta.""" + p1 = torch.tensor([0.25, 0.25, 0.25, 0.25]) + for eta in (0.0, 2.0, 10.0): + half, mean = corrector.simulate_jumps(p1, eta, positions=10_000, steps=300) + assert abs(half - 0.5) < 0.03 + assert abs(mean - (1.0 + eta)) < 0.25 + + +def test_parallel_step_error(): + assert abs(corrector.joint_reveal_rate(64) - 1.0 / 64) < 1e-12 + assert abs(corrector.pair_survival(2.0) - 0.1353) < 5e-4 + assert abs(corrector.pair_survival(10.0) - 4.54e-5) < 1e-7 + measured = corrector.simulate_pair(0.0, samples=40_000, steps=64) + assert abs(measured - 0.5 / 64) < 2e-3 # Half of the joint reveals land outside the data. + + +# %% Section 5.16: kinetic-optimal paths + + +def test_kinetic_optimal_recovers_the_mixture_velocity(): + q = torch.tensor([0.1, 0.6, 0.2, 0.1]) + p, p_dot = kinetic.mask_path(q, kappa_t=0.5) + u = kinetic.kinetic_optimal(p, p_dot) + assert abs(float(p_dot[kinetic.MASK] / p[kinetic.MASK]) + 2.0) < 1e-12 + assert abs(float(p_dot[1] / p[1]) - 2.0) < 1e-12 + assert abs(float(u[1, kinetic.MASK]) - 1.2) < 1e-12 # 0.30 times 4. + assert abs(float(u[1, kinetic.MASK]) - kinetic.mixture_rate(0.5) * 0.6) < 1e-12 + assert abs(float(u[1, 2])) < 1e-12 # No probability moves between two visible tokens. + assert torch.allclose(u @ p, p_dot, atol=1e-12) + + +def test_worked_metric_path_for_dna(): + beta_t, beta_dot = kinetic.beta_schedule(0.5) + assert abs(beta_t - 1.0) < 1e-12 and abs(beta_dot - 4.0) < 1e-12 + p = kinetic.metric_path(2, beta_t) + close(p.tolist(), [0.2245, 0.0826, 0.6103, 0.0826]) + u = kinetic.metric_velocity(2, beta_t, beta_dot) + assert abs(float(u[0, 1]) - 0.898) < 5e-4 # C to A. + assert abs(float(u[2, 1]) - 4.882) < 5e-4 # C to G. + assert abs(float(u[2, 0]) - 2.441) < 5e-4 # A to G. + assert abs(float(-u[1, 1]) - 5.780) < 1e-3 # The total rate out of C. + close((torch.eye(4)[1] + 0.05 * u[:, 1]).tolist(), [0.0449, 0.7110, 0.2441, 0.0]) + close((u @ p).tolist(), [-0.3997, -0.4774, 1.3546, -0.4774]) + close(kinetic.path_derivative(2, 0.5).tolist(), [-0.3997, -0.4774, 1.3546, -0.4774], 1e-3) + + +def test_energy_of_two_schedules(): + assert abs(kinetic.schedule_energy(*kinetic.LINEAR, 600, 1e-6) - 27.63) < 5e-3 + assert abs(kinetic.schedule_energy(*kinetic.LINEAR, 600, 1e-9) - 41.45) < 5e-3 + assert abs(kinetic.schedule_energy(*kinetic.GEODESIC, 600, 1e-6) - math.pi ** 2) < 1e-3 + assert abs(kinetic.geodesic_schedule(0.25) - 0.1464) < 5e-4 + assert abs(kinetic.geodesic_schedule(0.75) - 0.8536) < 5e-4 + assert abs(kinetic.geodesic_schedule(0.5, 0.25) - 0.6667) < 5e-4 + + +# %% Section 5.17: Edit Flows + + +def test_alignment_and_column_wise_path(): + z0 = [0, 1, 2, 3, edits.BLANK] # (A, C, G, T, eps). + z1 = [0, edits.BLANK, 2, 0, 1] # (A, eps, G, A, C). + assert edits.column_kinds(z0, z1) == ["none", "del", "none", "sub", "ins"] + drawn = edits.draw_state(z0, z1, [True, False, True, True, True]) + assert drawn == [0, 1, 2, 0, 1] # ACGAC, the draw of probability 0.5 ** 3 = 0.125. + assert abs(edits.expected_length([0, 1, 2, 3], [0, 2, 0, 1], 0.5) - 4.0) < 1e-12 + recovered = edits.align([0, 1, 2, 3], [0, 2, 0, 1]) + assert len(recovered[0]) == len(recovered[1]) # Any alignment pads to a common width. + + +def test_alignment_posterior_and_averaged_rates(): + first = ([0, 1, edits.BLANK], [0, 0, 1]) # Substitute C with A, then insert C. + second = ([0, edits.BLANK, 1], [0, 0, 1]) # Insert A, and leave the last column alone. + for kappa_t, expected in ((0.25, 0.4286), (0.5, 1 / 3), (0.75, 0.2)): + weights = edits.alignment_posterior([first, second], [0, 1], kappa_t) + assert abs(weights[0] - expected) < 5e-4 + posterior, rates = edits.averaged_rates([first, second], [0, 1], 0.5) + assert abs(rates["sub(2,A)"] - 0.667) < 1e-3 + assert abs(rates["ins(1,A)"] - 1.333) < 1e-3 + assert abs(rates["ins(2,C)"] - 0.667) < 1e-3 + assert abs(sum(rates.values()) - 2.667) < 1e-3 + + +def test_edit_loss_and_sampler_step(): + assert abs(edits.edit_loss(3.0, [0.8], kappa_t=0.5) - 3.446) < 1e-3 + grid = [0.01 * i for i in range(1, 801)] + assert abs(edits.loss_minimizer(2.0, grid) - 2.0) < 1e-6 + intensities = {"ins": [0.2, 0.3, 0.4], "del": [0.1, 0.9, 0.2], "sub": [0.4, 0.3, 0.5]} + insert, change = edits.step_probabilities(intensities, 0.05) + close(insert, [0.010, 0.015, 0.020]) + close(change, [0.025, 0.060, 0.035]) + exact, first_order = edits.untouched_probability(intensities, 0.05) + assert abs(exact - 0.8452) < 5e-4 and abs(first_order - 0.835) < 5e-4 + assert abs(exact - first_order - 0.0102) < 5e-4 + assert abs(edits.expected_length_change(intensities, 0.05) + 0.015) < 1e-9 + assert edits.edit_rates(intensities)[0] == 11 # Three intensities and two token distributions. + + +# %% Section 5.20: Dirichlet Flow Matching + + +def test_dirichlet_path_and_speeds(): + assert abs(dirichlet.dirichlet_density(torch.full((4,), 0.25), 0, 0.0) - 6.0) < 1e-9 + assert dirichlet.straight_line_corners(torch.tensor([0.34, 0.33, 0.33, 0.0]), 0.3) == 3 + assert dirichlet.straight_line_corners(torch.tensor([0.70, 0.30, 0.0, 0.0]), 0.6) == 1 + closed = dirichlet.speed_two_letters(0.3, 1.0) + assert abs(closed - 0.2580) < 5e-5 + assert abs(float(dirichlet.speed(0.3, 1.0, K=2)) - closed) < 1e-6 + x = torch.tensor([0.40, 0.30, 0.20, 0.10]) + close(dirichlet.speed_vector(x, 1.0).tolist(), [0.232, 0.199, 0.159, 0.104], 1e-3) + + +def test_dirichlet_velocity_and_posterior_average(): + x = torch.tensor([0.40, 0.30, 0.20, 0.10]) + toward_a = dirichlet.conditional_velocity(x, 0, 1.0) + close(toward_a.tolist(), [0.139, -0.070, -0.046, -0.023], 1e-3) + assert abs(float(toward_a.sum())) < 1e-12 + p_hat = torch.tensor([0.60, 0.25, 0.10, 0.05]) + v = dirichlet.marginal_velocity(x, p_hat, 1.0) + close(v.tolist(), [0.0553, -0.0133, -0.0262, -0.0158], 2e-4) + close((x + 0.1 * v).tolist(), [0.4055, 0.2987, 0.1974, 0.0984]) + assert abs(float((x + 0.1 * v).sum()) - 1.0) < 1e-12 + + +def test_dirichlet_continuity(): + time_term, space_term = dirichlet.continuity_terms(torch.tensor([0.5, 0.3, 0.2]), 0, 0.8) + assert abs(time_term - 0.6355) < 5e-4 + assert abs(space_term + 0.6355) < 5e-4 + assert abs(time_term + space_term) < 1e-5 + + +# %% Section 5.21: Fisher-Flow + + +def test_two_moves_under_two_metrics(): + near = torch.tensor([0.5, 0.5]), torch.tensor([0.4, 0.6]) + rare = torch.tensor([0.01, 0.99]), torch.tensor([0.001, 0.999]) + assert abs(fisher.euclidean_distance(*near) - 0.1414) < 5e-4 + assert abs(fisher.euclidean_distance(*rare) - 0.0127) < 5e-4 + assert abs(fisher.fisher_rao_distance(*near) - 0.201) < 1e-3 + assert abs(fisher.fisher_rao_distance(*rare) - 0.137) < 1e-3 + orthogonal = fisher.fisher_rao_distance(torch.tensor([1.0, 0.0]), torch.tensor([0.0, 1.0])) + assert abs(orthogonal - math.pi) < 1e-6 # Any two letters are the farthest apart. + + +def test_sphere_and_geodesic(): + p0 = torch.full((4,), 0.25) + p1 = fisher.smooth(torch.eye(4)[0]) + close(p1.tolist(), [0.91, 0.03, 0.03, 0.03]) + x0, x1 = fisher.to_sphere(p0), fisher.to_sphere(p1) + close(x1.tolist(), [0.954, 0.173, 0.173, 0.173], 1e-3) + assert abs(float((x0 * x1).sum()) - 0.737) < 1e-3 + assert abs(math.acos(float((x0 * x1).sum())) - 0.743) < 1e-3 + assert abs(fisher.fisher_rao_distance(p0, p1) - 1.485) < 1e-3 + mid = fisher.geodesic(x0, x1, torch.tensor(0.5)) + close(mid.tolist(), [0.780, 0.361, 0.361, 0.361], 1e-3) + close(fisher.to_simplex(mid).tolist(), [0.609, 0.130, 0.130, 0.130], 1e-3) + remaining = fisher.log_map(mid, x1) + close(remaining.tolist(), [0.2323, -0.1672, -0.1672, -0.1672], 1e-3) + assert abs(float(remaining.norm()) - 0.3713) < 5e-4 + assert abs(float(fisher.conditional_velocity(mid, x1, 0.5).norm()) - 0.743) < 1e-3 + stepped = fisher.exp_map(mid, 0.1 * fisher.conditional_velocity(mid, x1, 0.5)) + close(stepped.tolist(), [0.824, 0.327, 0.327, 0.327], 1e-3) + close(fisher.to_simplex(stepped).tolist(), [0.680, 0.107, 0.107, 0.107], 1e-3) + assert abs(float(stepped.norm()) - 1.0) < 1e-10 + + +def test_tangent_projection(): + x0, p1 = fisher.to_sphere(torch.full((4,), 0.25)), fisher.smooth(torch.eye(4)[0]) + mid = fisher.geodesic(x0, fisher.to_sphere(p1), torch.tensor(0.5)) + raw = torch.tensor([0.3, -0.2, 0.5, 0.1]) + assert abs(float((mid * raw).sum()) - 0.378) < 1e-3 + tangent = fisher.project(mid, raw) + close(tangent.tolist(), [0.005, -0.337, 0.363, -0.037], 1e-3) + assert abs(float((mid * tangent).sum())) < 1e-12 + + +# %% Section 5.22: Gumbel-Softmax Flow Matching + + +def test_gumbel_max_and_temperature(): + pi = torch.tensor([0.5, 0.2, 0.2, 0.1]) + torch.manual_seed(0) + frequencies = gumbel.gumbel_max(pi, 200_000) + close(frequencies.tolist(), pi.tolist(), tolerance=5e-3) + close([gumbel.temperature(t) for t in (0.0, 0.5, 0.9, 1.0)], [10.0, 2.2313, 0.6721, 0.4979], 5e-4) + + +def test_annealed_path_and_replicator(): + logits = torch.tensor([0.0, -0.45, 1.40, -0.05]) + expected = { + 0.0: [0.2438, 0.2331, 0.2805, 0.2426], + 0.5: [0.2142, 0.1751, 0.4012, 0.2095], + 0.9: [0.0955, 0.0489, 0.7669, 0.0887], + 1.0: [0.0528, 0.0214, 0.8781, 0.0477], + } + for t, values in expected.items(): + close(gumbel.annealed_path(logits, t).tolist(), values) + x = gumbel.annealed_path(logits, 0.5) + assert abs(float((x * logits).sum()) - 0.4724) < 5e-4 + noisy = gumbel.replicator_velocity(x, logits, 0.5) + close(noisy.tolist(), [-0.136, -0.217, 0.500, -0.147], 1e-3) + assert abs(float(noisy.sum())) < 1e-12 + close(gumbel.inference_velocity(x, 2, 0.5).tolist(), [-0.116, -0.094, 0.323, -0.113], 1e-3) + + +def test_gumbel_marginal_velocity(): + late = torch.tensor([0.138, 0.084, 0.647, 0.131]) + p_hat = torch.tensor([0.05, 0.10, 0.70, 0.15]) + coefficients = (gumbel.LAMBDA / gumbel.temperature(0.8)) * late * p_hat + close(coefficients.tolist(), [0.0228, 0.0278, 1.4976, 0.0650], 5e-4) + v = gumbel.marginal_velocity(late, p_hat, 0.8) + close(v.tolist(), [-0.200, -0.108, 0.454, -0.146], 1e-3) + close((late + 0.01 * v).tolist(), [0.1360, 0.0829, 0.6515, 0.1295]) + + +# %% Sections 5.24 to 5.29 and 5.33: guidance + + +def test_guided_base(): + p1 = torch.tensor([0.1, 0.2, 0.3, 0.4]) + rates = p1 / (1.0 - 0.6) + scores = torch.tensor([0.10, 0.35, 0.55, 0.25]) + close(rates.tolist(), [0.25, 0.50, 0.75, 1.00]) + assert abs(float(rates.sum()) - 2.50) < 1e-12 + close((scores / 0.20).tolist(), [0.50, 1.75, 2.75, 1.25]) + guided = guidance.tilt(rates, scores, 0.20) + close(guided.tolist(), [0.125, 0.875, 2.0625, 1.250]) + assert abs(float(guided.sum()) - 4.3125) < 1e-12 + close((guided / guided.sum()).tolist(), [0.029, 0.203, 0.478, 0.290], 1e-3) + close((p1 * scores).tolist(), [0.010, 0.070, 0.165, 0.100]) + h = guidance.h_value(p1, scores) + assert abs(h - 0.345) < 1e-12 + assert abs(float(guidance.tilt(rates, scores, h).sum()) - 2.50) < 1e-12 + assert guidance.reachable(100, 21) == 2001 + + +def test_guidance_strength_and_classifier_free(): + p1 = torch.tensor([0.1, 0.2, 0.3, 0.4]) + rates, scores = p1 / 0.4, torch.tensor([0.10, 0.35, 0.55, 0.25]) + harder = guidance.tilt(rates, scores, 0.20, gamma=2.0) + close(harder.tolist(), [0.0625, 1.5312, 5.6719, 1.5625], 1e-3) + assert abs(float(harder.sum()) - 8.828) < 1e-3 + close((harder / harder.sum()).tolist(), [0.007, 0.173, 0.642, 0.177], 1e-3) + mixed = guidance.classifier_free(guidance.tilt(rates, scores, 0.20), rates, gamma=2.0) + assert torch.allclose(mixed, harder, atol=1e-12) + + +def test_taylor_expansion(): + gradient = torch.tensor([-2.10, -1.15, -0.75, -1.30, -1.61]) + approximate = guidance.taylor_log_ratio(gradient, guidance.MASK)[:4] + close(approximate.tolist(), [-0.49, 0.46, 0.86, 0.31]) + exact = (torch.tensor([0.10, 0.35, 0.55, 0.25]) / 0.20).log() + close(exact.tolist(), [-0.693, 0.560, 1.012, 0.223], 1e-3) + rates = torch.tensor([0.25, 0.50, 0.75, 1.00]) * approximate.exp() + close((rates / rates.sum()).tolist(), [0.038, 0.194, 0.434, 0.334], 1e-3) + + +def test_guided_step_never_remasks_a_visible_position(): + """The denoiser in the sampler is a distribution over the data tokens, which exclude the mask.""" + torch.manual_seed(0) + model = SequenceDenoiser(guidance.V, 8, width=32, depth=1) + prop = guidance.PropertyModel(8, width=32) + xt = torch.randint(0, 4, (16, 8)) # Every position already holds a base. + stepped = guidance.guided_step(model, prop, xt, torch.full((16,), 0.9), 0.5, gamma=1.0) + assert int((stepped == guidance.MASK).sum()) == 0 # A revealed token never returns to the mask. + + +# %% Sections 5.30 to 5.32: several objectives at once + + +def test_pareto_front(): + scores = torch.tensor([[0.8, 0.3], [0.6, 0.7], [0.5, 0.2], [0.6, 0.3]]) + assert objectives.pareto_front(scores) == [0, 1] # The third and fourth are dominated. + close((scores @ torch.tensor([0.6, 0.4]))[:2].tolist(), [0.60, 0.64]) + close((scores @ torch.tensor([0.8, 0.2]))[:2].tolist(), [0.70, 0.62]) # The order reverses. + + +def test_mogdfm_step(): + changes = torch.tensor([[0.0, 0.0], [0.20, -0.05], [0.05, 0.15], [-0.10, 0.08]]) + delta_s, ranks, direction = objectives.combined_score( + changes, torch.tensor([1.2, 0.8]), torch.tensor([0.6, 0.4])) + close(ranks.tolist(), [0.500, 0.700, 0.850, 0.450]) + close(direction.tolist(), [0.0, 0.100, 0.090, -0.028]) + close(delta_s.tolist(), [-1.511, 1.541, 2.297, -2.327], 1e-3) + close(delta_s.exp().tolist(), [0.221, 4.667, 9.948, 0.098], 1e-2) + base = torch.tensor([0.3, 0.3, 0.1]) / 0.3 # The pretrained rates out of A toward C, G and T. + tilted = base * delta_s[1:].exp() + close(tilted.tolist(), [4.667, 9.948, 0.033], 1e-2) + angles, cosines = objectives.cone_angles(changes[1:], torch.tensor([0.6, 0.4])) + close(cosines.tolist(), [0.6727, 0.7894, -0.3032], 5e-4) + close(angles.tolist(), [47.7, 37.9, 107.7], 0.05) + survivors = tilted[angles <= 60.0] + close((survivors / survivors.sum()).tolist(), [0.319, 0.681], 1e-3) + assert abs(objectives.update_cone(60.0, 1 / 3) - 64.1) < 0.05 + + +def test_pcomole_rollout(): + omega, reference = torch.tensor([0.5, 0.5]), torch.zeros(2) + first = [torch.tensor(s) for s in ((0.8, 0.4), (0.6, 0.6), (0.9, 0.1), (0.5, 0.5))] + value, rewards = objectives.rollout_value(first, [1, 1, 0, 1], omega, reference, 2.0) + close(rewards, [1.682, 2.054, 0.0, 1.822], 1e-3) + assert abs(value - 1.390) < 1e-3 + second = [torch.tensor(s) for s in ((1.0, 0.0), (0.9, 0.1), (0.95, 0.05), (1.0, 0.1))] + other, _ = objectives.rollout_value(second, [1, 1, 1, 1], omega, reference, 2.0) + assert abs(other - 1.181) < 1e-3 + products = torch.tensor([0.6 * value, 1.2 * other, 0.0]) + close(products.tolist(), [0.834, 1.417, 0.0], 1e-3) + close((products / products.sum()).tolist(), [0.371, 0.629, 0.0], 1e-3) + + +# %% Sections 5.36 to 5.41: rectification and AReUReDi + + +def test_total_correlation_on_two_bases(): + source = torch.full((rectify.STATES,), 1.0 / rectify.STATES) + q, support = rectify.complementary_pairs() + coupling = q[None, :].repeat(rectify.STATES, 1) + assert abs(rectify.total_correlation(coupling, source) - math.log(4)) < 1e-9 + assert abs(rectify.hamming(coupling, source) - 1.5) < 1e-9 + assert abs(rectify.one_step_success(coupling, source, support) - 0.25) < 1e-9 + + +def test_redi_rounds(): + source = torch.full((rectify.STATES,), 1.0 / rectify.STATES) + q, support = rectify.complementary_pairs() + coupling = q[None, :].repeat(rectify.STATES, 1) + coupling = rectify.recouple(coupling, source, steps=300) + assert abs(rectify.total_correlation(coupling, source) - 1.101) < 2e-3 + assert abs(rectify.hamming(coupling, source) - 1.022) < 2e-3 + assert abs(rectify.one_step_success(coupling, source, support) - 0.403) < 2e-3 + columns = [a * 4 + b for a, b in support] + close([float(coupling[0, c]) for c in columns], [0.400, 0.400, 0.100, 0.100], 2e-3) + close([float(coupling[3, c]) for c in columns], [0.755, 0.082, 0.082, 0.082], 2e-3) + coupling = rectify.recouple(coupling, source, steps=300) + assert abs(rectify.total_correlation(coupling, source) - 0.843) < 2e-3 + assert abs(rectify.one_step_success(coupling, source, support) - 0.525) < 2e-3 + + +def test_areuredi_proposal_and_acceptance(): + 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]) + assert abs(rectify.anneal(11, 20) - 12.0) < 1e-12 + tchebycheff = rectify.tchebycheff(scores, weights) + close(tchebycheff.tolist(), [0.300, 0.350, 0.250, 0.310]) + tilts = (12.0 * tchebycheff).exp() + close(tilts.tolist(), [36.60, 66.69, 20.09, 41.26], 5e-2) + model = torch.tensor([0.40, 0.25, 0.20, 0.15]) + q_current, z_current, ratios = rectify.proposal(model, tilts, 0) + close(ratios.tolist(), [1.000, 1.822, 0.549, 1.127], 1e-3) + close(rectify.barker(ratios).tolist(), [0.500, 0.646, 0.354, 0.530], 1e-3) + assert abs(z_current - 0.5118) < 5e-4 + close(q_current.tolist(), [0.391, 0.315, 0.138, 0.155], 1e-3) + _, z_proposed, _ = rectify.proposal(model, tilts, 1) + assert abs(z_proposed - 0.3704) < 5e-4 + assert abs(rectify.acceptance(z_current, z_proposed) - 1.0) < 1e-12 + assert abs(rectify.acceptance(z_proposed, z_current) - 0.724) < 1e-3 + + +def test_areuredi_concentration(): + base = torch.tensor([0.50, 0.25, 0.15, 0.10]) + objective = torch.tensor([0.20, 0.30, 0.45, 0.45]) + tilted = rectify.tilted_distribution(base, objective, 20.0) + close(tilted.tolist(), [0.0127, 0.0468, 0.5643, 0.3762], 1e-3) + assert abs(float(tilted[2:].sum()) - 0.9405) < 1e-3 + bound, gap = rectify.concentration_bound(base, objective, 20.0) + assert abs(gap - 0.15) < 1e-9 + assert abs(bound * float(base[1]) - 0.0498) < 5e-4 + assert bound * float(base[1]) >= float(tilted[1]) + sharper = rectify.tilted_distribution(base, objective, 60.0) + assert abs(float(sharper[2:].sum()) - 0.9999) < 1e-3 + assert abs(float(sharper[1]) - 1.234e-4) < 1e-6 + + +# %% The command line of every file + + +@pytest.mark.parametrize("script", [ + "velocity", "kinetic_optimal", "rectify", "edit_flows", "multi_objective", +]) +def test_script_runs(script): + """The five files that need no training run end to end at the smallest grid.""" + result = subprocess.run( + [sys.executable, f"lecture_5/{script}.py", "--steps", "20", "--quiet"], + cwd=ROOT, capture_output=True, text=True, env={**os.environ, "PYTHONPATH": str(ROOT)}, + ) + assert result.returncode == 0, result.stderr[-2000:] + assert "nan" not in result.stdout.lower() diff --git a/tests/test_lecture_6.py b/tests/test_lecture_6.py new file mode 100644 index 0000000000000000000000000000000000000000..da104feab783fe4ce7e7035cce7cc9601fba3ee0 --- /dev/null +++ b/tests/test_lecture_6.py @@ -0,0 +1,398 @@ +"""The worked numbers of Chapter 6, asserted against the Lecture 6 files. + +Every number checked here appears in the chapter text, and every one is computed by the file under +test and not stored in it. The training loops are not run; the tests call the closed-form and +numerical functions the lectures build their reports from. +""" + +from __future__ import annotations + +import importlib.util +import math +import sys +from pathlib import Path + +import pytest +import torch + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + + +def load(name: str): + """Import one lecture file by path, so that ``lecture_6`` needs no package marker.""" + spec = importlib.util.spec_from_file_location(f"lecture_6_{name}", + ROOT / "lecture_6" / f"{name}.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def close(values, expected, tolerance=5e-4): + """Compare a list of numbers against the chapter's, coordinate by coordinate.""" + assert len(values) == len(expected) + for got, want in zip(values, expected): + assert got == pytest.approx(want, abs=tolerance) + + +# %% 01. flow_map.py, Sections 6.3 to 6.12 +def test_flow_map_listing(): + flow_map = load("flow_map") + x = torch.tensor(0.7, dtype=torch.float64) + assert float(flow_map.flow(x, 0.0, 1.0)) == pytest.approx(4.1, abs=1e-9) + assert float(flow_map.flow(flow_map.flow(x, 0.0, 0.5), 0.5, 1.0)) == pytest.approx(4.1, abs=1e-9) + assert float(flow_map.flow(flow_map.flow(x, 0.0, 1.0), 1.0, 0.0)) == pytest.approx(0.7, abs=1e-9) + from torch.func import jacrev + + assert float(jacrev(lambda z: flow_map.flow(z, 0.0, 1.0))(x)) == pytest.approx(3.0, abs=1e-9) + + +def test_flow_map_axioms_and_liouville(): + flow_map = load("flow_map") + x = torch.tensor(0.7, dtype=torch.float64) + identity, cocycle, tangent = flow_map.check_axioms(x, 200) + assert identity < 1e-12 and cocycle < 1e-10 and tangent < 1e-4 + assert flow_map.liouville_log_det(0.0, 1.0, 4000) == pytest.approx(math.log(3.0), abs=1e-6) + assert float(flow_map.jacobian_by_ode(x, 0.0, 1.0, 4000)) == pytest.approx(3.0, abs=1e-3) + + +def test_flow_map_transport_and_likelihood(): + flow_map = load("flow_map") + source = flow_map.normal_density(0.7, 0.0, 1.0) + data = flow_map.normal_density(4.1, 2.0, 3.0) + assert source == pytest.approx(0.3123, abs=1e-4) + assert data == pytest.approx(0.1041, abs=1e-4) + assert source / data == pytest.approx(3.0, abs=1e-6) + assert math.log(source) - math.log(3.0) == pytest.approx(-2.2625, abs=1e-4) + assert math.log(data) == pytest.approx(-2.2625, abs=1e-4) + + +# %% 02. residuals.py, Sections 6.16 and 6.17 +def test_stay_still_residuals(): + residuals = load("residuals") + report = residuals.stay_still_residuals() + assert report["stay-still Lagrangian"] == pytest.approx(1.0, abs=1e-9) + assert report["stay-still Eulerian"] == pytest.approx(1.0, abs=1e-9) + assert report["stay-still composition"] == pytest.approx(0.0, abs=1e-12) + assert report["stay-still diagonal gap"] == pytest.approx(0.0, abs=1e-12) + + +def test_identity_axiom_is_free(): + residuals = load("residuals") + torch.manual_seed(0) + model = residuals.Map(dim=2, width=16) + x, t = torch.randn(8, 2), torch.rand(8) + assert torch.allclose(model.X(x, t, t), x) + + +# %% 03. meanflow.py, Sections 6.13 and 6.18 +def test_meanflow_scalar_identity(): + meanflow = load("meanflow") + report = meanflow.scalar_check() + assert report["v_0,1(1)"] == pytest.approx(math.e - 1.0, abs=1e-5) + assert report["b(1, 0)"] == pytest.approx(1.0, abs=1e-6) + assert report["partial_s v"] == pytest.approx(-1.0, abs=1e-5) + assert report["grad_x v_0,1"] == pytest.approx(math.e - 1.0, abs=1e-5) + assert report["d/ds v_0,1"] == pytest.approx(math.e - 2.0, abs=1e-5) + assert report["identity gap"] == pytest.approx(0.0, abs=1e-5) + assert report["one map evaluation"] == pytest.approx(math.e, abs=1e-5) + assert report["one Euler step"] == pytest.approx(2.0, abs=1e-6) + + +def test_meanflow_correction_vanishes_on_the_gaussian_flow(): + meanflow = load("meanflow") + report = meanflow.gaussian_check() + assert report["v_0,t(0.7) spread"] == pytest.approx(0.0, abs=1e-6) + assert report["v_0,1(0.7)"] == pytest.approx(3.4, abs=1e-5) + assert report["b(0.7, 0)"] == pytest.approx(3.4, abs=1e-5) + + +# %% 04. shortcut.py, Section 6.20 +def test_shortcut_doubling(): + shortcut = load("shortcut") + report = shortcut.scalar_doubling() + assert report["s(1, 0, d)"] == pytest.approx(1.136102, abs=1e-6) + assert report["state after the first step"] == pytest.approx(1.284025, abs=1e-6) + assert report["s(x', d, d)"] == pytest.approx(1.458783, abs=1e-6) + assert report["average of the two"] == pytest.approx(1.297443, abs=1e-6) + assert report["s(1, 0, 2d) direct"] == pytest.approx(1.297443, abs=1e-6) + assert report["doubling gap"] == pytest.approx(0.0, abs=1e-12) + assert report["state after the doubled step"] == pytest.approx(math.exp(0.5), abs=1e-9) + assert shortcut.ladder()[0] == 1.0 / 128.0 and shortcut.ladder()[-1] == 1.0 + + +# %% 05. consistency.py, Section 6.20 +def test_consistency_boundary_parameterization(): + consistency = load("consistency") + report = consistency.scalar_consistency() + assert report["f(0.5)"] == pytest.approx(math.exp(0.5), abs=1e-6) + assert report["f(0.9)"] == pytest.approx(math.exp(0.1), abs=1e-6) + assert report["f(1.0)"] == pytest.approx(1.0, abs=1e-9) + assert report["c_out(1.0)"] == pytest.approx(0.0, abs=1e-12) + assert report["v_0.5,1(1)"] == pytest.approx(1.297443, abs=1e-6) + + +def test_consistency_teacher_floor(): + consistency = load("consistency") + report = consistency.teacher_floor() + assert report["f(1, 0.5)"] == pytest.approx(1.648721, abs=1e-6) + assert report["f(1.01, 0.51)"] == pytest.approx(1.648639, abs=1e-6) + assert report["distillation residual"] == pytest.approx(8.189e-5, rel=1e-3) + assert report["local Euler error"] == pytest.approx(5.0167e-5, rel=1e-3) + assert report["carried to the data end"] == pytest.approx( + report["distillation residual"], rel=1e-6) + + +def test_consistency_trajectory_head(): + consistency = load("consistency") + report = consistency.scalar_ctm() + assert report["a_0.0,0.5"] == pytest.approx(0.5, abs=1e-12) + assert report["g_0.0,0.5"] == pytest.approx(2.297443, abs=1e-6) + assert report["a_0.0,1.0"] == pytest.approx(0.0, abs=1e-12) + assert report["g_0.0,1.0"] == pytest.approx(math.e, abs=1e-6) + assert report["two routes apart"] == pytest.approx(0.0, abs=1e-12) + + +# %% 06. fmlm.py, Sections 6.26 and 6.27 +def test_fmlm_denoiser_cocycle(): + fmlm = load("fmlm") + report = fmlm.scalar_cocycle() + assert report["gamma"] == pytest.approx(0.2, abs=1e-12) + assert report["delta_s,t"] == pytest.approx(2.096158, abs=1e-6) + assert report["delta_s,u"] == pytest.approx(1.932957, abs=1e-6) + assert report["delta_u,t"] == pytest.approx(2.136959, abs=1e-6) + assert report["two legs combined"] == pytest.approx(2.096158, abs=1e-6) + assert report["convex weights"] == (0.25, 0.75) + assert report["convex form"] == pytest.approx(math.exp(0.6), abs=1e-6) + assert report["X_s,t(x)"] == pytest.approx(math.exp(0.6), abs=1e-9) + + +def test_fmlm_worked_kl(): + fmlm = load("fmlm") + report = fmlm.worked_kl() + assert report["worked gamma"] == pytest.approx(0.2, abs=1e-12) + close(report["composed target"], [0.06, 0.12, 0.76, 0.06], 1e-6) + assert report["target mass"] == pytest.approx(1.0, abs=1e-6) + close(report["the four terms"], [-0.030650, -0.026777, 0.118823, -0.030650], 1e-6) + assert report["worked KL"] == pytest.approx(0.030747, abs=1e-6) + + +def test_fmlm_gradient_comparison(): + """Section 6.27.4, the two gradients in the logits at the chapter's target and student.""" + fmlm = load("fmlm") + report = fmlm.gradient_comparison() + close(report["KL gradient"], [0.04, 0.03, -0.11, 0.04], 1e-6) + close(report["squared-error gradient"], [0.0099, 0.01335, -0.03315, 0.0099], 1e-6) + close(report["squared over KL"], [0.2475, 0.445, 0.3014, 0.2475], 1e-4) # A quarter to a half. + close(report["saturated student"], [1e-6, 0.1666665, 0.7222215, 0.111111], 1e-6) + assert report["saturated KL gradient"] == pytest.approx(-0.0600, abs=1e-4) + # The chapter quotes -6.10e-08 here. Differentiating its own expression + # d/dz_i (1/2) sum_j (q_j - p_j)^2 = q_i (q_i - p_i) - q_i sum_j q_j (q_j - p_j) + # at q_1 = 1e-06 gives -4.6171e-08. The first term on its own is -6.00e-08, so the quoted + # number is the first term with the second dropped, and this file prints the full expression. + assert report["saturated squared-error gradient, in 1e-08"] == pytest.approx(-4.6171, abs=1e-3) + assert report["orders of magnitude between them"] == pytest.approx(6.11, abs=0.01) + + +def test_fmlm_decoding_clock(): + fmlm = load("fmlm") + tau, start = fmlm.clock(fmlm.BOOK_VOCAB) + assert start == pytest.approx(0.999967, abs=1e-5) + assert tau(0.60) == pytest.approx(0.006, abs=5e-4) + assert tau(0.75) == pytest.approx(0.141, abs=5e-4) + assert tau(0.80) == pytest.approx(0.456, abs=1e-3) + assert tau(0.85) == pytest.approx(0.931, abs=1e-3) + grid = fmlm.even_in_tau(fmlm.BOOK_VOCAB) + close(grid, [0.0, 0.773490, 0.804481, 0.828095, 1.0], 5e-5) + close([tau(t) for t in (0.25, 0.5, 0.75)], [0.00009, 0.00124, 0.14109], 1e-5) + + +def test_fmlm_clock_weights(): + fmlm = load("fmlm") + rounded = [0.0, 0.773, 0.804, 0.828, 1.0] # The chapter quotes the grid to three decimals. + pairs = [((1 - b) / (1 - a), (b - a) / (1 - a)) for a, b in zip(rounded[:-1], rounded[1:])] + close([p[0] for p in pairs], [0.2270, 0.8634, 0.8776, 0.0], 1e-4) + close([p[1] for p in pairs], [0.7730, 0.1366, 0.1224, 1.0], 1e-4) + times = [torch.tensor([v]) for v in rounded[1:4]] + assert float(fmlm.composition_weight(*times)) == pytest.approx(0.494620, abs=1e-5) + early = [torch.tensor([v]) for v in (0.0, 0.773, 0.804)] + assert float(fmlm.composition_weight(*early)) == pytest.approx(0.830144, abs=1e-5) + + +def test_fmlm_endpoint_slice_is_trained(): + """Section 6.26 asks for t = 1 on purpose, since a continuous draw never lands there.""" + fmlm = load("fmlm") + assert 0.0 < fmlm.ENDPOINT_FRACTION < 1.0 + tokens, _ = fmlm.synthetic_dna(num_sequences=64, length=16, seed=0) + model = fmlm.TwoTimeMLP(dim=16 * 4, width=16) + torch.manual_seed(0) + loss_fn = fmlm.make_loss(model, tokens, 16, 4, 256, torch.device("cpu")) + assert torch.isfinite(loss_fn(0)) # The rows at t = 1 go through the loss like any other. + # The map at t = 1 is the prediction itself, so the one-step sampler reads that slice alone. + x = torch.randn(2, 16, 4) + s, t = torch.zeros(2), torch.ones(2) + prediction = fmlm.predict(model, x, s, t, 16, 4) + assert torch.allclose(fmlm.step_map(x, prediction, s, t), prediction) + + +# %% 07. categorical_map.py, Section 6.28 +def test_categorical_worked_cancellation(): + categorical = load("categorical_map") + report = categorical.worked_cancellation() + close(report["state x"], [-0.2591, -0.6675, -0.0467], 1e-4) + close(report["X_s,t(x)"], [0.1611, -0.2420, 0.0943], 1e-4) + close(report["pi_s,t(x)"], [0.5813, 0.1835, 0.2352], 1e-4) + close(report["pi_t,t(X_s,t(x))"], [0.6927, 0.0917, 0.2157], 1e-4) + assert report["gamma_s,t"] == pytest.approx(0.2, abs=1e-12) + close(report["d pi / dt"], [0.5570, -0.4592, -0.0979], 1e-4) + close(report["central difference"], [0.5570, -0.4592, -0.0979], 1e-4) + close(report["gap a"], [-0.111, 0.092, 0.020], 1e-3) + close(report["drift b"], [0.111, -0.092, -0.020], 1e-3) + + +def test_categorical_bound_is_positive_where_the_residual_is_zero(): + categorical = load("categorical_map") + report = categorical.worked_cancellation() + assert report["residual |a + b|^2"] < 1e-20 + assert report["4 KL"] == pytest.approx(0.156, abs=1e-3) + assert report["2 |b|^2"] == pytest.approx(0.042, abs=1e-3) + assert report["the bound"] == pytest.approx(0.199, abs=2e-3) + assert report["teacher entropy"] == pytest.approx(0.8043, abs=1e-3) + assert report["cross-entropy"] == pytest.approx(0.8434, abs=1e-3) + assert report["cross-entropy"] == pytest.approx( + 0.25 * report["4 KL"] + report["teacher entropy"], abs=1e-6) + + +def test_categorical_averaged_bound(): + categorical = load("categorical_map") + report = categorical.averaged_bound(1000, seed=0) + assert report["averaged residual"] < 1e-18 + assert report["averaged 4 KL"] == pytest.approx(0.39, abs=0.06) + assert report["averaged 2 |b|^2"] == pytest.approx(0.095, abs=0.03) + + +# %% 08. discrete_map.py, Sections 6.26 and 6.29 +def test_discrete_derivative_target_leaves_the_simplex(): + discrete = load("discrete_map") + report = discrete.worked_derivative() + assert report["derivative prefactor"] == pytest.approx(6.6667, abs=1e-3) + close(report["derivative target"], [-0.3333, -0.6667, 1.3333, -0.3333], 1e-4) + assert report["target sum"] == pytest.approx(0.0, abs=1e-6) + assert report["negative entries"] == 3 + + +def test_discrete_logit_repair(): + discrete = load("discrete_map") + report = discrete.worked_repair() + assert report["weighted average of dz/dt"] == pytest.approx(0.0, abs=1e-6) + close(report["d psi/dt"], [1.92, -1.14, -0.78], 1e-4) + close(report["additive target"], [-0.284, 0.828, 0.456], 1e-4) + assert report["additive target sum"] == pytest.approx(1.0, abs=1e-6) + assert report["additive negative entries"] == 1 + close(report["multiplicative factor"], [2.920, 0.544, 0.480], 1e-4) + close(report["quotient"], [0.0342, 1.1029, 0.6250], 1e-4) + assert report["quotient mass"] == pytest.approx(1.7622, abs=1e-4) + close(report["repaired target"], [0.0194, 0.6259, 0.3547], 1e-4) + close(report["log of the factor"], [1.0716, -0.6088, -0.7340], 1e-4) + assert report["gamma_s,t"] == pytest.approx(0.2, abs=1e-12) + assert report["gamma^E_s,t"] == pytest.approx(0.8, abs=1e-12) + + +# %% 09. stochastic_map.py, Section 6.31 +def test_stochastic_weak_floor(): + stochastic = load("stochastic_map") + report = stochastic.weak_floor(draws=200_000, seed=0) + assert report["K_1"] == pytest.approx(3.194528, abs=1e-5) + assert report["sqrt K_1"] == pytest.approx(1.7873, abs=1e-4) + assert report["sqrt K_0.5"] == pytest.approx(0.9269, abs=1e-4) + assert report["one full step"] == pytest.approx(2.2453, abs=1e-3) + assert report["two half steps"] == pytest.approx(0.4673, abs=1e-3) + assert report["gap"] == pytest.approx(1.778, abs=1e-3) + assert report["squared gap"] == pytest.approx(3.161, abs=1e-3) + assert report["floor 2 K_1"] == pytest.approx(math.e ** 2 - 1.0, abs=1e-9) + assert report["carried variance equals K_1"] == pytest.approx(report["K_1"], abs=1e-9) + assert report["the draw that would close it"] == pytest.approx(-0.195, abs=1e-3) + assert report["Monte Carlo floor"] == pytest.approx(6.389, abs=0.05) + + +def test_stochastic_chen_relations(): + stochastic = load("stochastic_map") + report = stochastic.worked_chen() + assert report["I0 over [0,1]"] == pytest.approx(0.50, abs=1e-9) + assert report["I1 over [0,1]"] == pytest.approx(-0.20, abs=1e-9) + assert report["fused variance at h=1"] == pytest.approx(1.0 / 3.0, abs=1e-9) + + +def test_stochastic_truncation(): + stochastic = load("stochastic_map") + report = stochastic.truncation_error(paths=512, grid=2048, seed=0) + for n, expected in enumerate((1.0, 1.0 / 3.0, 1.0 / 5.0)): + assert report[f"coefficient variance I{n}"] == pytest.approx(expected, rel=0.15) + errors = [report[f"pathwise error with {k}"] for k in (1, 2, 3)] + assert errors[0] > 5.0 * errors[1] > 25.0 * errors[2] + + +# %% 10. expanding.py, Section 6.32 +def test_expanding_local_clocks(): + expanding = load("expanding") + report = expanding.worked_clocks() + close(report["local clocks at t=0.8"], [0.800, 0.800, 0.692, 0.333], 1e-3) + close(report["clock speeds"], [1.0, 1.0, 1.538, 3.333], 1e-3) + close(report["source share"], [0.200, 0.200, 0.308, 0.667], 1e-3) + assert report["last clock from its speed"] == pytest.approx(0.667, abs=1e-3) + assert report["live coordinates at 0.8"] == 4 and report["live coordinates at 0.6"] == 3 + + +def test_expanding_composition_weight(): + expanding = load("expanding") + report = expanding.worked_omega() + assert report["omega at (0.3, 0.6, 0.9)"] == pytest.approx(0.125, abs=1e-6) + assert report["omega for a birth at 0.0"] == pytest.approx(0.125, abs=1e-6) + assert report["omega for a birth at 0.7"] == pytest.approx(0.0, abs=1e-9) + assert 0.0 < report["omega for a birth at 0.35"] < report["omega for a birth at 0.0"] + + +# %% 11. posterior_map.py, Sections 6.33 and 6.34 +def test_posterior_tilt_and_jensen_gap(): + posterior = load("posterior_map") + report = posterior.two_point() + close(report["tilted target"], [0.1192, 0.8808], 1e-4) + close(report["posterior at x=0.7"], [0.9820, 0.0180], 1e-4) + close(report["tilted posterior"], [0.8808, 0.1192], 1e-4) + assert report["V_t(0.7)"] == pytest.approx(0.1088, abs=1e-4) + assert report["V over the data"] == pytest.approx(1.4338, abs=1e-4) + assert report["single-endpoint estimate"] == pytest.approx(1.0, abs=1e-9) + assert report["Jensen gap"] == pytest.approx(0.434, abs=1e-3) + + +def test_posterior_value_bias_and_gradient(): + posterior = load("posterior_map") + report = posterior.value_bias() + close([report[f"E[V_{k}]"] for k in (1, 2, 4, 8)], [1.000, 1.217, 1.340, 1.393], 1e-3) + close([report[f"bias at K={k}"] for k in (1, 2, 4, 8)], [0.434, 0.217, 0.093, 0.041], 1e-3) + weights, weighted, plain = posterior.value_gradient( + torch.tensor([0.2, 1.5, 0.9]), torch.tensor([0.4, -1.2, 0.5])) + close(weights.tolist(), [0.1496, 0.5491, 0.3013], 1e-4) + assert weighted == pytest.approx(-0.4483, abs=1e-4) + assert plain == pytest.approx(-0.1000, abs=1e-6) + + +def test_posterior_renoising_and_meta_fixed_point(): + posterior = load("posterior_map") + report = posterior.renoise_example() + assert report["renoise scale"] == pytest.approx(0.75, abs=1e-9) + assert report["renoise variance"] == pytest.approx(0.1375, abs=1e-9) + assert report["renoise standard deviation"] == pytest.approx(0.3708, abs=1e-4) + assert report["proposal weight on the far mode"] == pytest.approx(0.420, abs=1e-3) + assert report["penalty at z=2"] == pytest.approx(10.125, abs=1e-6) + assert report["penalty at z=0"] == pytest.approx(6.125, abs=1e-6) + assert report["corrected weight on the far mode"] == pytest.approx(0.0131, abs=1e-4) + meta = posterior.stop_gradient_fixed_point(torch.tensor([1.0, 2.0, 5.0]), + torch.tensor([0.4, 1.0, 3.6])) + assert meta["E[w]"] == pytest.approx(8.0 / 3.0, abs=1e-6) + assert meta["E[G]"] == pytest.approx(5.0 / 3.0, abs=1e-6) + assert meta["stop-gradient fixed point"] == pytest.approx(0.625, abs=1e-6) + assert meta["per-sample ratio average"] == pytest.approx(0.540, abs=1e-6) + assert meta["gap"] == pytest.approx(-0.085, abs=1e-3) + assert meta["covariance of G/w and w"] == pytest.approx(0.2267, abs=1e-4) + assert meta["gap from the covariance"] == pytest.approx(meta["gap"], abs=1e-6) diff --git a/tests/test_lecture_7.py b/tests/test_lecture_7.py new file mode 100644 index 0000000000000000000000000000000000000000..05a30b28cb2a60ad91bf23f58e26789fc4f4a19c --- /dev/null +++ b/tests/test_lecture_7.py @@ -0,0 +1,709 @@ +"""Lecture 7: every number the chapter works by hand, asserted against the file that prints it. + +Chapter 7 computes almost all of its results on paper, so these tests are the chapter itself read +back. The tolerances are the number of digits the chapter quotes. +""" + +from __future__ import annotations + +import importlib.util +import math +import sys +from pathlib import Path + +import pytest +import torch + +LECTURE = Path(__file__).resolve().parent.parent / "lecture_7" + + +def load(name: str): + """Import one lecture file by path, since ``lecture_7`` is a folder of scripts, not a package.""" + spec = importlib.util.spec_from_file_location(f"lecture_7_{name}", LECTURE / f"{name}.py") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def tensor(*values): + return torch.tensor(values, dtype=torch.float64) + + +# %% Sections 7.1 to 7.6, optimal_transport.py + + +def test_four_points_and_two_pairings(): + ot = load("optimal_transport") + pairings = ot.two_pairings() + assert pairings["arrows_ordered"] == pytest.approx(2.0) + assert pairings["arrows_crossed"] == pytest.approx(10.0) + assert pairings["mass_ordered"] == pytest.approx(1.0) + assert pairings["mass_crossed"] == pytest.approx(5.0) + assert pairings["mass_independent"] == pytest.approx(3.0) + assert pairings["C"].tolist() == [[1.0, 9.0], [1.0, 1.0]] + + +def test_two_by_two_plan_and_prices(): + ot = load("optimal_transport") + assert ot.feasible_interval() == (0.0, 0.3) + C = ot.cost_matrix(ot.SOURCE_X, ot.TARGET_Y) + assert ot.cost_along_the_segment(0.0, C) == pytest.approx(5.8) + assert ot.cost_along_the_segment(0.3, C) == pytest.approx(3.4) + plan = ot.feasible_plan(0.3) + assert plan.tolist() == [[0.3, 0.3], [0.0, 0.4]] + assert ot.is_coupling(plan, ot.SOURCE_A, ot.TARGET_B) + f, g = ot.prices_from_loaded_routes(C, plan) + assert f.tolist() == [0.0, -8.0] and g.tolist() == [1.0, 9.0] + assert ot.slack(C, f, g).tolist() == [[0.0, 0.0], [8.0, 0.0]] + assert ot.dual_value(ot.SOURCE_A, f, ot.TARGET_B, g) == pytest.approx(3.4) + assert ot.plan_cost(C, plan) ** 0.5 == pytest.approx(1.844, abs=5e-4) + + +def test_no_crossing_and_three_points(): + ot = load("optimal_transport") + expanded, collected = ot.crossing_difference(0.0, 2.0, 1.0, 3.0) + assert expanded == pytest.approx(collected) and collected == pytest.approx(-8.0) + x = tensor(4.0, 1.0, 6.0) + y = tensor(9.0, 2.0, 3.0) + costs = [c for c, _ in ot.monge_permutations(x.sort().values, y.sort().values)] + assert costs == pytest.approx([11 / 3, 17 / 3, 35 / 3, 45 / 3, 77 / 3, 81 / 3]) + assert ot.wasserstein_on_the_line(x, y, 2.0, 3072) == pytest.approx(11 / 3) + + +def test_gaussian_closed_form(): + ot = load("optimal_transport") + assert ot.gaussian_w2_squared(0.0, 1.0, 3.0, 2.0) == pytest.approx(10.0) + assert math.sqrt(10.0) == pytest.approx(3.162, abs=5e-4) + points = tensor(-1.0, 0.0, 1.0) + assert ot.brenier_map(points, 0.0, 1.0, 3.0, 2.0).tolist() == [1.0, 3.0, 5.0] + source = torch.special.ndtri(ot.quantile_levels(3072)) + mapped = ot.brenier_map(source, 0.0, 1.0, 3.0, 2.0) + assert float(((mapped - source) ** 2).mean()) == pytest.approx(10.0, abs=1e-3) + + +# %% Sections 7.9 and 7.10, sinkhorn.py + + +def test_entropies_of_the_three_couplings(): + sk = load("sinkhorn") + independent = sk.independent_coupling(sk.SOURCE_A, sk.TARGET_B) + assert independent.flatten().tolist() == pytest.approx([0.18, 0.42, 0.12, 0.28]) + assert sk.entropy(sk.EXACT_PLAN) == pytest.approx(1.089, abs=5e-4) + assert sk.entropy(independent) == pytest.approx(1.284, abs=5e-4) + assert sk.entropy(sk.SOURCE_A) == pytest.approx(0.673, abs=5e-4) + assert sk.entropy(sk.TARGET_B) == pytest.approx(0.611, abs=5e-4) + assert sk.relative_entropy(sk.EXACT_PLAN, independent) == pytest.approx(0.195, abs=5e-4) + + +def test_sinkhorn_first_sweep_by_hand(): + sk = load("sinkhorn") + K = sk.gibbs_kernel(sk.COST, 2.0) + assert K.flatten().tolist() == pytest.approx([0.606531, 0.011109, 0.606531, 0.606531], + abs=5e-7) + v = torch.ones(2, dtype=torch.float64) + assert (K @ v).tolist() == pytest.approx([0.617640, 1.213061], abs=5e-7) + u = sk.SOURCE_A / (K @ v) + assert u.tolist() == pytest.approx([0.971440, 0.329744], abs=5e-7) + assert (K.T @ u).tolist() == pytest.approx([0.789208, 0.210792], abs=5e-7) + v = sk.TARGET_B / (K.T @ u) + assert v.tolist() == pytest.approx([0.38013, 3.32081], abs=5e-6) + plan = sk.scaling_form(u, K, v) + assert plan.flatten().tolist() == pytest.approx([0.2240, 0.0358, 0.0760, 0.6642], abs=5e-5) + assert plan.sum(dim=0).tolist() == pytest.approx([0.3, 0.7]) + assert plan.sum(dim=1).tolist() == pytest.approx([0.2598, 0.7402], abs=5e-5) + after = sk.project_rows(plan, sk.SOURCE_A) + assert after.sum(dim=0).tolist() == pytest.approx([0.558, 0.442], abs=5e-4) + + +def test_sinkhorn_convergence_and_epsilon(): + sk = load("sinkhorn") + K = sk.gibbs_kernel(sk.COST, 2.0) + u, v, errors = sk.sinkhorn_sweeps(K, sk.SOURCE_A, sk.TARGET_B, 40) + assert errors[:5] == pytest.approx([0.680, 0.382, 0.155, 0.047, 0.012], abs=5e-4) + assert sk.sweeps_to_tolerance(K, sk.SOURCE_A, sk.TARGET_B) == 17 + plan = sk.scaling_form(u, K, v) + assert plan.flatten().tolist() == pytest.approx( + [0.293122, 0.306878, 0.006878, 0.393122], abs=5e-7) + assert float((sk.COST * plan).sum()) == pytest.approx(3.45502, abs=5e-6) + assert sk.entropy(plan) == pytest.approx(1.124, abs=5e-4) + independent = sk.independent_coupling(sk.SOURCE_A, sk.TARGET_B) + assert sk.entropic_objective(plan, sk.COST, 2.0) == pytest.approx(1.208, abs=5e-4) + assert sk.entropic_objective(sk.EXACT_PLAN, sk.COST, 2.0) == pytest.approx(1.222, abs=5e-4) + assert sk.entropic_objective(independent, sk.COST, 2.0) == pytest.approx(1.792, abs=5e-4) + assert sk.relative_entropy(plan, independent) == pytest.approx(0.160, abs=5e-4) + f = sk.multipliers_from_scalings(u, 2.0) + g = sk.multipliers_from_scalings(v, 2.0) + assert f.tolist() == pytest.approx([3.766, -3.739], abs=5e-4) + assert g.tolist() == pytest.approx([-3.220, 4.872], abs=5e-4) + expected = {10.0: (4.0540165, 7), 2.0: (3.4550204, 17), 0.5: (3.4000004, 28), + 0.1: (3.4000000, 80)} + for epsilon, (cost, sweeps) in expected.items(): + kernel = sk.gibbs_kernel(sk.COST, epsilon) + su, sv, _ = sk.sinkhorn_sweeps(kernel, sk.SOURCE_A, sk.TARGET_B, 400) + assert float((sk.COST * sk.scaling_form(su, kernel, sv)).sum()) == pytest.approx( + cost, abs=5e-8) + assert sk.sweeps_to_tolerance(kernel, sk.SOURCE_A, sk.TARGET_B) == sweeps + + +# %% Sections 7.7 to 7.7.5, displacement.py + + +def test_one_particle_and_the_averaged_bound(): + dp = load("displacement") + times = torch.linspace(0.0, 1.0, 1025, dtype=torch.float64) + assert dp.action_of_a_path(dp.steady_particle(times), times) == pytest.approx(4.0) + assert dp.action_of_a_path(dp.dawdling_particle(times), times) == pytest.approx(8.0) + x, y = tensor(0.0, 2.0), tensor(1.0, 3.0) + ordered = torch.tensor([[0.5, 0.0], [0.0, 0.5]], dtype=torch.float64) + crossed = torch.tensor([[0.0, 0.5], [0.5, 0.0]], dtype=torch.float64) + independent = torch.full((2, 2), 0.25, dtype=torch.float64) + assert dp.average_squared_displacement(x, y, ordered) == pytest.approx(1.0) + assert dp.average_squared_displacement(x, y, independent) == pytest.approx(3.0) + assert dp.average_squared_displacement(x, y, crossed) == pytest.approx(5.0) + + +def test_displacement_interpolation_and_its_velocity(): + dp = load("displacement") + mean, std = dp.displacement_law(0.5) + assert (mean, std) == pytest.approx((1.5, 1.5)) + assert dp.gaussian_w2(0.0, 1.0, mean, std) == pytest.approx(math.sqrt(2.5)) + assert math.sqrt(2.5) == pytest.approx(1.581, abs=5e-4) + assert dp.gaussian_w2(0.0, 1.0, 3.0, 2.0) == pytest.approx(3.162, abs=5e-4) + grid = torch.linspace(-6.0, 9.0, 15001, dtype=torch.float64) + peak = float(grid[int(dp.mixture_density(grid).argmax())]) + assert peak == pytest.approx(0.13, abs=5e-3) + half = torch.tensor(0.5, dtype=torch.float64) + assert float(dp.velocity_field(torch.tensor(1.5, dtype=torch.float64), half)) == pytest.approx(3.0) + assert float(dp.velocity_field(torch.tensor(3.0, dtype=torch.float64), half)) == pytest.approx(4.0) + for t in (0.0, 0.25, 0.5, 0.75, 1.0): + assert dp.action_density(t, 4096) == pytest.approx(10.0, abs=1e-3) + assert abs(dp.continuity_residual(1.5, 0.5)) < 1e-12 + + +# %% Sections 7.8 and 7.8.1, ot_flow_matching.py + + +def test_batch_coupling_costs(): + fm = load("ot_flow_matching") + x = torch.tensor([[0.0], [2.0]]) + y = torch.tensor([[1.0], [3.0]]) + C = fm.batch_cost_matrix(x, y) + assert C.flatten().tolist() == pytest.approx([1.0, 9.0, 1.0, 1.0]) + assert float(C.mean()) == pytest.approx(3.0) + pairing = fm.optimal_pairing(x, y) + assert pairing.tolist() == [0, 1] + assert float(C[torch.arange(2), pairing].mean()) == pytest.approx(1.0) + assert fm.vertex_count(256, 256) == {"unknowns": 65536, "constraints": 512, + "loaded_routes": 511} + assert fm.conditional_velocity_energy(x, fm.couple(x, y, "optimal")) == pytest.approx(1.0) + + +def test_flow_matching_pipeline_runs(): + fm = load("ot_flow_matching") + from cis6270.nets import VelocityMLP + + torch.manual_seed(0) + model = VelocityMLP(dim=2, width=16) + x0 = torch.randn(8, 2) + x1 = torch.randn(8, 2) + assert float(fm.flow_matching_loss(model, x0, fm.couple(x0, x1, "optimal"))) > 0.0 + samples, straightness = fm.integrate(model, x0, 5) + assert samples.shape == (8, 2) and straightness >= 0.0 + + +# %% Sections 7.12 to 7.15.3, static_bridge.py + + +def test_eight_paths_and_the_chain_rule(): + sb = load("static_bridge") + law = sb.path_law(sb.PI_0, [sb.A, sb.A]) + assert law[(0, 0, 1)] == pytest.approx(0.105) + assert law[(0, 0, 0)] == pytest.approx(0.245) + assert law[(1, 1, 1)] == pytest.approx(0.180) + assert sum(law.values()) == pytest.approx(1.0) + table = torch.tensor([[0.5, 0.1], [0.1, 0.3]], dtype=torch.float64) + uniform = torch.full((2, 2), 0.25, dtype=torch.float64) + total, marginal, conditional = sb.chain_rule_terms(table, uniform) + assert total == pytest.approx(0.2180, abs=5e-5) + assert marginal == pytest.approx(0.0201, abs=5e-5) + assert conditional == pytest.approx(0.1979, abs=5e-5) + assert marginal + conditional == pytest.approx(total) + + +def test_proportional_fitting_on_the_chain(): + sb = load("static_bridge") + R01 = sb.reference_endpoint_law(sb.PI_0, sb.A, 2) + assert R01.flatten().tolist() == pytest.approx([0.305, 0.195, 0.26, 0.24]) + assert float((R01.sum(dim=0) - sb.PI_1).abs().sum()) == pytest.approx(0.73) + f, g, errors = sb.proportional_fitting(R01, sb.PI_0, sb.PI_1, 4) + assert errors == pytest.approx([3.933e-3, 2.098e-5, 1.119e-7, 5.966e-10], rel=1e-3) + assert f.tolist() == pytest.approx([1.072002, 0.937019], abs=5e-7) + assert g.tolist() == pytest.approx([0.350517, 1.843637], abs=5e-7) + plan = f[:, None] * R01 * g[None, :] + assert plan.flatten().tolist() == pytest.approx( + [0.114605, 0.385395, 0.085395, 0.414605], abs=5e-7) + assert sb.kl_vectors(plan.flatten(), R01.flatten()) == pytest.approx(0.2820, abs=5e-5) + independent = sb.PI_0[:, None] * sb.PI_1[None, :] + assert sb.kl_vectors(independent.flatten(), R01.flatten()) == pytest.approx(0.2847, abs=5e-5) + + +def test_h_transform_and_the_bridge_path_law(): + sb = load("static_bridge") + R01 = sb.reference_endpoint_law(sb.PI_0, sb.A, 2) + f, g, _ = sb.proportional_fitting(R01, sb.PI_0, sb.PI_1, 40) + potentials = sb.backward_potential(sb.A, g, 2) + assert potentials[2].tolist() == pytest.approx([0.350517, 1.843637], abs=5e-7) + assert potentials[1].tolist() == pytest.approx([0.798453, 1.246389], abs=5e-7) + assert potentials[0].tolist() == pytest.approx([0.932834, 1.067215], abs=5e-7) + tilted = [sb.tilted_matrix(sb.A, potentials[k], potentials[k + 1]) for k in range(2)] + assert tilted[0].flatten().tolist() == pytest.approx( + [0.599160, 0.400840, 0.299266, 0.700734], abs=5e-7) + assert tilted[1].flatten().tolist() == pytest.approx( + [0.307297, 0.692703, 0.112490, 0.887510], abs=5e-7) + assert tilted[0].sum(dim=1).tolist() == pytest.approx([1.0, 1.0]) + middle = sb.PI_0 @ tilted[0] + assert middle.tolist() == pytest.approx([0.449213, 0.550787], abs=5e-7) + assert (middle @ tilted[1]).tolist() == pytest.approx([0.2, 0.8]) + assert (1.0 / potentials[0]).tolist() == pytest.approx(f.tolist(), abs=1e-9) + plan = f[:, None] * R01 * g[None, :] + bridge = sb.fill_with_reference_bridges(plan, sb.A) + reference = sb.path_law(sb.PI_0, [sb.A, sb.A]) + assert bridge[(0, 0, 0)] == pytest.approx(0.092060, abs=5e-7) + assert bridge[(0, 0, 1)] == pytest.approx(0.207520, abs=5e-7) + assert bridge[(1, 1, 1)] == pytest.approx(0.310954, abs=5e-7) + ratios = {round(bridge[p] / reference[p], 6) for p in reference} + assert ratios == {0.375755, 1.976383, 0.328441, 1.727522} + assert sb.relative_entropy(bridge, reference) == pytest.approx(0.281961, abs=5e-7) + assert sum(v for k, v in bridge.items() if k[1] == 0) == pytest.approx(0.449213, abs=5e-7) + + +def test_brownian_reference_recovers_the_entropic_plan(): + sb = load("static_bridge") + plan = sb.brownian_reference_plan(1.0) + assert plan.flatten().tolist() == pytest.approx([0.2931, 0.3069, 0.0069, 0.3931], abs=5e-5) + + +# %% Sections 7.16 to 7.18, dynamic_bridge.py + + +def test_constant_control_two_ways(): + db = load("dynamic_bridge") + assert db.control_energy_entropy(2.0, 1.0) == pytest.approx(2.0) + assert db.gaussian_entropy_gap(2.0, 1.0) == pytest.approx(2.0) + generator = torch.Generator().manual_seed(0) + _, log_ratio = db.simulate_controlled(2.0, 200, 8192, 1.0, generator) + assert float(log_ratio.mean()) == pytest.approx(2.0, abs=0.1) + + +def test_the_bridge_potential_solves_both_equations(): + db = load("dynamic_bridge") + value = db.hopf_cole(db.bridge_potential) + for x, t in ((-0.5, 0.2), (0.5, 0.5), (2.5, 0.8)): + assert abs(db.backward_residual(db.bridge_potential, x, t)) < 1e-12 + assert abs(db.hjb_residual(value, x, t)) < 1e-12 + drift = float(db.bridge_drift(torch.tensor(x, dtype=torch.float64), + torch.tensor(t, dtype=torch.float64))) + assert db.optimal_control(value, x, t) == pytest.approx(drift, abs=1e-10) + + +def test_brownian_bridge_numbers(): + db = load("dynamic_bridge") + mean, variance = db.bridge_law(0.25) + assert mean == pytest.approx(0.5) + assert variance == pytest.approx(0.1875) + assert math.sqrt(variance) == pytest.approx(0.4330, abs=5e-5) + drift = float(db.bridge_drift(torch.tensor(0.6732, dtype=torch.float64), + torch.tensor(0.25, dtype=torch.float64))) + assert drift == pytest.approx(1.7691, abs=5e-5) + assert db.bridge_law(0.5)[1] == pytest.approx(0.25) + for epsilon, spread in ((0.25, 0.25), (1.0, 0.5), (4.0, 1.0)): + assert math.sqrt(db.bridge_law(0.5, epsilon=epsilon)[1]) == pytest.approx(spread) + + +# %% Section 7.19.1, dsb.py + + +def test_dsb_half_steps_on_the_chain(): + dsb = load("dsb") + reference = dsb.reference_path_law(dsb.PI_0, dsb.A, 2) + assert dsb.marginal(reference, 1).tolist() == pytest.approx([0.55, 0.45]) + assert dsb.marginal(reference, 2).tolist() == pytest.approx([0.565, 0.435]) + B2 = dsb.reverse_step(reference, 1) + assert B2.flatten().tolist() == pytest.approx( + [0.681416, 0.318584, 0.379310, 0.620690], abs=5e-7) + assert float(dsb.reverse_step(reference, 0)[0, 0]) == pytest.approx(0.636364, abs=5e-7) + half = dsb.fix_terminal(reference, dsb.PI_1) + assert dsb.marginal(half, 2).tolist() == pytest.approx([0.2, 0.8]) + assert dsb.marginal(half, 1).tolist() == pytest.approx([0.439731, 0.560269], abs=5e-7) + assert dsb.marginal(half, 0).tolist() == pytest.approx([0.466585, 0.533415], abs=5e-7) + assert float((dsb.marginal(half, 0) - dsb.PI_0).abs().sum()) == pytest.approx(0.0668, abs=5e-5) + full = dsb.fix_initial(half, dsb.PI_0) + assert dsb.one_step_conditional(full, 0).flatten().tolist() == pytest.approx( + [0.599738, 0.400262, 0.299771, 0.700229], abs=5e-7) + assert dsb.one_step_conditional(full, 1).flatten().tolist() == pytest.approx( + [0.309924, 0.690076, 0.113725, 0.886275], abs=5e-7) + assert dsb.marginal(full, 2).tolist() == pytest.approx([0.201967, 0.798033], abs=5e-7) + gap = max(float((dsb.one_step_conditional(full, 0) - dsb.TILTED_0).abs().max()), + float((dsb.one_step_conditional(full, 1) - dsb.TILTED_1).abs().max())) + assert gap < 0.003 + _, errors, _ = dsb.proportional_fitting(reference, dsb.PI_0, dsb.PI_1, 4) + assert errors == pytest.approx([3.933e-3, 2.098e-5, 1.119e-7, 5.966e-10], rel=1e-3) + assert errors[0] / errors[1] == pytest.approx(187.5, abs=0.1) + + +# %% Section 7.19.2, dsbm.py + + +def test_markov_drift_at_one_point(): + dm = load("dsbm") + assert dm.bridge_drift(3.0, 1.0, 0.5) == pytest.approx(4.0) + assert dm.bridge_drift(0.0, 1.0, 0.5) == pytest.approx(-2.0) + endpoints = tensor(3.0, 0.0) + weights = tensor(0.5, 0.5) + assert dm.markov_drift(endpoints, weights, 1.0, 0.5) == pytest.approx(1.0) + assert dm.least_squares_derivative(1.0, endpoints, weights, 1.0, 0.5) == pytest.approx(0.0) + assert dm.least_squares_derivative(0.0, endpoints, weights, 1.0, 0.5) == pytest.approx(-2.0) + + +def test_one_imf_round_on_the_chain(): + dm = load("dsbm") + bridges = dm.reference_bridge_table(dm.A) + assert float(bridges[0, 0, 0]) == pytest.approx(0.803279, abs=5e-7) + assert float(bridges[0, 0, 1]) == pytest.approx(0.538462, abs=5e-7) + target = dm.bridge_coupling(dm.A, dm.PI_0, dm.PI_1) + independent = dm.PI_0[:, None] * dm.PI_1[None, :] + assert independent.flatten().tolist() == pytest.approx([0.1, 0.4, 0.1, 0.4]) + assert float((independent - target).abs().sum()) == pytest.approx(0.0584, abs=5e-5) + coupling, gaps, first = dm.imf_rounds(independent, dm.A, 9, target) + assert first[0].flatten().tolist() == pytest.approx( + [0.591425, 0.408575, 0.307692, 0.692308], abs=5e-7) + assert first[1].flatten().tolist() == pytest.approx( + [0.298457, 0.701543, 0.119588, 0.880412], abs=5e-7) + projected = torch.diag(dm.PI_0) @ first[0] @ first[1] + assert projected.flatten().tolist() == pytest.approx( + [0.112688, 0.387312, 0.087312, 0.412688], abs=5e-7) + assert projected.sum(dim=0).tolist() == pytest.approx([0.2, 0.8]) + assert projected.sum(dim=1).tolist() == pytest.approx([0.5, 0.5]) + assert gaps[0] == pytest.approx(0.0077, abs=5e-5) + assert gaps[3] / gaps[4] == pytest.approx(7.37, abs=5e-3) + assert gaps[8] < 1e-9 + + +# %% Section 7.19.3, sf2m.py + + +def test_the_two_regression_targets(): + sm = load("sf2m") + x0, x1 = tensor(0.0)[0], tensor(2.0)[0] + z = tensor(0.4)[0] + for t_value, mean, sample, drift, flow, coefficient in ( + (0.25, 0.5, 0.6732, 1.7691, 2.2309, 4 / 3), + (0.75, 1.5, 1.6732, 1.3072, 1.7691, -4 / 3), + ): + t = torch.tensor(t_value, dtype=torch.float64) + assert float(sm.bridge_mean(x0, x1, t)) == pytest.approx(mean) + assert float(sm.bridge_std(t)) == pytest.approx(0.4330, abs=5e-5) + x = sm.bridge_sample(x0, x1, t, z) + assert float(x) == pytest.approx(sample, abs=5e-5) + assert float(sm.conditional_drift(x1, x, t)) == pytest.approx(drift, abs=5e-5) + assert float(sm.conditional_score(x0, x1, x, t)) == pytest.approx(-0.9238, abs=5e-5) + assert float(sm.flow_velocity(x0, x1, x, t)) == pytest.approx(flow, abs=5e-5) + assert float(sm.flow_velocity_collected(x0, x1, x, t)) == pytest.approx(flow, abs=5e-5) + assert float(sm.flow_coefficient(t)) == pytest.approx(coefficient) + assert float(sm.scaled_score(x0, x1, x, t)) == pytest.approx(-0.4, abs=5e-7) + + +def test_the_batch_plan_of_sf2m(): + sm = load("sf2m") + x, y = tensor(0.0, 2.0), tensor(1.0, 3.0) + uniform = tensor(0.5, 0.5) + plan = sm.entropic_plan(x, y, uniform, uniform, 1.0) + assert plan.sum(dim=0).tolist() == pytest.approx([0.5, 0.5]) + assert float(plan[0, 0] + plan[1, 1]) > float(plan[0, 1] + plan[1, 0]) + + +# %% Sections 7.21 to 7.23.3, discrete_bridge.py + + +def test_rate_error_is_not_symmetric(): + db = load("discrete_bridge") + assert db.scaling_cost(0.5) == pytest.approx(0.1534, abs=5e-5) + assert db.scaling_cost(2.0) == pytest.approx(0.3863, abs=5e-5) + assert float(db.rate_error(tensor(1.0)[0], tensor(1.0)[0])) == pytest.approx(0.0) + + +def test_two_state_continuous_time_bridge(): + db = load("discrete_bridge") + assert torch.matrix_exp(db.Q).flatten().tolist() == pytest.approx( + [0.6328, 0.3672, 0.5507, 0.4493], abs=5e-5) + R01 = db.endpoint_law(db.Q, db.PI_0) + assert R01.flatten().tolist() == pytest.approx([0.3164, 0.1836, 0.2754, 0.2246], abs=5e-5) + assert R01.sum(dim=0).tolist() == pytest.approx([0.5918, 0.4082], abs=5e-5) + f, g, plan = db.sinkhorn(R01, db.PI_0, db.PI_1) + assert f.tolist() == pytest.approx([1.0716, 0.9373], abs=5e-5) + assert g.tolist() == pytest.approx([0.3349, 1.9643], abs=5e-5) + assert float(g[1] / g[0]) == pytest.approx(5.865, abs=5e-4) + assert plan.flatten().tolist() == pytest.approx([0.1136, 0.3864, 0.0864, 0.4136], abs=5e-5) + assert torch.matrix_exp(db.Q / 2).flatten().tolist() == pytest.approx( + [0.7146, 0.2854, 0.4281, 0.5719], abs=5e-5) + phi = db.backward_potential(db.Q, g, 0.5) + assert phi.tolist() == pytest.approx([0.7999, 1.2667], abs=5e-5) + tilted = db.tilted_generator(db.Q, phi) + assert tilted.flatten().tolist() == pytest.approx( + [-1.5836, 1.5836, 0.9472, -0.9472], abs=5e-5) + assert db.backward_potential(db.Q, g, 0.0).tolist() == pytest.approx([0.9331, 1.0669], abs=5e-5) + at_zero = db.tilted_generator(db.Q, db.backward_potential(db.Q, g, 0.0)) + assert float(at_zero[0, 1]) == pytest.approx(1.1433, abs=5e-5) + assert float(at_zero[1, 0]) == pytest.approx(1.3120, abs=5e-5) + at_one = db.tilted_generator(db.Q, g) + assert float(at_one[0, 1]) == pytest.approx(5.8652, abs=5e-5) + assert float(at_one[1, 0]) == pytest.approx(0.2557, abs=5e-5) + marginal = db.bridge_marginal(db.Q, db.PI_0, f, g, 0.5) + assert marginal.tolist() == pytest.approx([0.4668, 0.5332], abs=5e-5) + assert db.bridge_marginal(db.Q, db.PI_0, f, g, 1.0).tolist() == pytest.approx([0.2, 0.8]) + reverse = db.reverse_rates(tilted, marginal) + assert float(reverse[1, 0]) == pytest.approx(1.3862, abs=5e-5) + assert float(reverse[0, 1]) == pytest.approx(1.0821, abs=5e-5) + assert float(marginal[0] * tilted[0, 1]) == pytest.approx( + float(marginal[1] * reverse[1, 0]), abs=5e-5) + + +def test_euler_sampler_and_the_cost_computed_twice(): + db = load("discrete_bridge") + R01 = db.endpoint_law(db.Q, db.PI_0) + f, g, plan = db.sinkhorn(R01, db.PI_0, db.PI_1) + errors = {} + for steps, expected in ((10, 0.0223), (20, 0.0106), (40, 0.0051), (100, 0.0021)): + terminal, leaving = db.euler_marginal(db.Q, db.PI_0, f, g, steps) + errors[steps] = float((terminal - db.PI_1).abs().sum() / 2) + assert errors[steps] == pytest.approx(expected, abs=1e-4) + if steps == 10: + assert leaving[:3] == pytest.approx([0.1143, 0.1187, 0.1245], abs=5e-5) + assert leaving[-1] == pytest.approx(0.3649, abs=5e-5) + assert db.relative_entropy(plan, R01) == pytest.approx(0.3235, abs=5e-5) + integral, at_half = db.path_cost(db.Q, db.PI_0, f, g) + assert at_half == pytest.approx(0.1300, abs=5e-5) + assert integral == pytest.approx(0.3235, abs=5e-5) + expected = {0.5: 0.3148, 1.0: 0.3235, 2.0: 0.3337, 4.0: 0.3348} + for lam, value in expected.items(): + law = db.endpoint_law(lam * db.Q, db.PI_0) + _, _, scaled = db.sinkhorn(law, db.PI_0, db.PI_1) + assert db.relative_entropy(scaled, law) == pytest.approx(value, abs=5e-5) + assert db.relative_entropy(db.PI_1, db.STATIONARY) == pytest.approx(0.3348, abs=5e-5) + + +# %% Section 7.24, ddsbm.py + + +def test_ddsbm_single_bond(): + dd = load("ddsbm") + kernel = dd.replacement_kernel(0.4, dd.PRIOR) + assert kernel.flatten().tolist() == pytest.approx( + [0.6, 0.2, 0.2, 0.2, 0.6, 0.2, 0.2, 0.2, 0.6]) + reference = torch.diag(dd.BOND_0) @ kernel + assert reference.sum(dim=0).tolist() == pytest.approx([0.48, 0.28, 0.24]) + assert (dd.BOND_1 / reference.sum(dim=0)).tolist() == pytest.approx( + [0.2083, 0.7143, 2.9167], abs=5e-5) + f, g, plan, errors = dd.sinkhorn(reference, dd.BOND_0, dd.BOND_1, 200) + assert f.tolist() == pytest.approx([1.1779, 0.9506, 0.5062], abs=5e-5) + assert g.tolist() == pytest.approx([0.1842, 0.6918, 3.0004], abs=5e-5) + assert errors[-1] < 1e-12 + assert float(plan[0, 2]) == pytest.approx(0.495, abs=5e-4) + assert dd.held_mass(plan) == pytest.approx(0.261, abs=5e-4) + assert float((dd.BOND_0 * dd.BOND_1).sum()) == pytest.approx(0.18) + assert float(torch.minimum(dd.BOND_0, dd.BOND_1).sum()) == pytest.approx(0.4) + expected = {0.2: 0.2250, 0.4: 0.2612, 0.6: 0.2926, 0.8: 0.3248, 0.99: 0.3812} + for r, value in expected.items(): + held = dd.sinkhorn(torch.diag(dd.BOND_0) @ dd.replacement_kernel(r, dd.PRIOR), + dd.BOND_0, dd.BOND_1, 200)[2] + assert dd.held_mass(held) == pytest.approx(value, abs=5e-5) + + +def test_ddsbm_pinned_rates_and_the_posterior_average(): + dd = load("ddsbm") + beta = dd.resampling_rate(0.4) + assert beta == pytest.approx(0.9163, abs=5e-5) + rates = dd.replacement_generator(beta, dd.PRIOR) + assert float(rates[0, 1]) == pytest.approx(0.3054, abs=5e-5) + assert rates.sum(dim=1).abs().max() < 1e-12 + assert dd.half_survival_time(0.4) == pytest.approx(math.log(0.8) / math.log(0.4)) + assert dd.half_survival_time(0.4) == pytest.approx(0.2435, abs=5e-5) + pinned = dd.replacement_kernel(0.5, dd.PRIOR) + assert float(pinned[0, 0]) == pytest.approx(2 / 3, abs=5e-5) + assert float(pinned[0, 1]) == pytest.approx(1 / 6, abs=5e-5) + toward = dd.pinned_rate(rates, pinned, 0, 2, 2) + away = dd.pinned_rate(rates, pinned, 2, 0, 2) + between = dd.pinned_rate(rates, pinned, 0, 1, 2) + assert toward == pytest.approx(1.2217, abs=5e-5) + assert away == pytest.approx(0.0764, abs=5e-5) + assert between == pytest.approx(0.3054, abs=5e-5) + assert toward / away == pytest.approx(16.0, abs=5e-4) + assert dd.rate_error(toward, 1.0) == pytest.approx(0.0229, abs=5e-5) + posterior = 0.6 * toward + 0.4 * between + assert posterior == pytest.approx(0.8552, abs=5e-5) + rates_two = [toward, between] + weights = [0.6, 0.4] + assert dd.posterior_loss(rates_two, weights, posterior) == pytest.approx(0.1357, abs=5e-5) + assert dd.posterior_loss(rates_two, weights, 1.0) == pytest.approx(0.1467, abs=5e-5) + assert dd.posterior_loss(rates_two, weights, 0.5) == pytest.approx(0.2395, abs=5e-5) + + +def test_ddsbm_factorization_counts(): + dd = load("ddsbm") + assert dd.feature_count(9) == (9, 36, 45) + assert dd.state_space_size(9, 5, 4) == pytest.approx(9.2e27, rel=5e-3) + assert dd.total_jump_rate(36, dd.resampling_rate(0.4), 4) == pytest.approx(24.74, abs=5e-3) + + +# %% Section 7.25, csbm.py + + +def test_csbm_references(): + cs = load("csbm") + A = cs.uniform_reference(0.35, 2) + assert A.flatten().tolist() == pytest.approx([0.65, 0.35, 0.35, 0.65]) + assert torch.linalg.matrix_power(A, 3).flatten().tolist() == pytest.approx( + [0.5135, 0.4865, 0.4865, 0.5135], abs=5e-5) + assert cs.uniform_reference(0.35, 5).sum(dim=1).tolist() == pytest.approx([1.0] * 5) + ordered = cs.ordered_reference(1.0, 5) + assert ordered[2].tolist() == pytest.approx( + [0.0545, 0.2442, 0.4026, 0.2442, 0.0545], abs=5e-5) + assert cs.transition_loss(tensor(0.3160, 0.6840), tensor(0.25, 0.75)) == pytest.approx( + 0.0110, abs=5e-5) + + +def test_csbm_dimf_on_sixteen_paths(): + cs = load("csbm") + A = cs.uniform_reference(cs.ALPHA, 2) + plan = cs.sinkhorn(torch.diag(cs.PI_0) @ torch.linalg.matrix_power(A, 3), cs.PI_0, cs.PI_1) + assert plan.flatten().tolist() == pytest.approx( + [0.10432, 0.39568, 0.09568, 0.40432], abs=5e-6) + assert float(cs.bridge_table(A, 3)[0, 0, 0, 0]) == pytest.approx(0.5348, abs=5e-5) + gap, start, target = cs.starting_gap(A, 3, cs.PI_0, cs.PI_1) + assert float(target[0, 0, 0, 0]) == pytest.approx(0.0558, abs=5e-5) + assert float(target[1, 1, 1, 1]) == pytest.approx(0.2162, abs=5e-5) + assert float(start[0, 0, 0, 0]) == pytest.approx(0.0535, abs=5e-5) + assert float(start[1, 1, 1, 1]) == pytest.approx(0.2139, abs=5e-5) + assert gap == pytest.approx(0.0173, abs=5e-5) + for axis, value in ((1, [0.4754, 0.5246]), (2, [0.4107, 0.5893])): + marginal = start.sum(dim=[d for d in range(4) if d != axis]) + assert marginal.tolist() == pytest.approx(value, abs=5e-5) + gaps, conditionals = cs.d_imf(cs.PI_0[:, None] * cs.PI_1[None, :], A, 3, 5, target) + assert conditionals[0].flatten().tolist() == pytest.approx( + [0.6243, 0.3757, 0.3265, 0.6735], abs=5e-5) + assert conditionals[1].flatten().tolist() == pytest.approx( + [0.5626, 0.4374, 0.2729, 0.7271], abs=5e-5) + assert conditionals[2].flatten().tolist() == pytest.approx( + [0.3160, 0.6840, 0.1192, 0.8808], abs=5e-5) + assert gaps == pytest.approx([2.9e-4, 5.0e-6, 8.4e-8, 1.4e-9, 2.4e-11], rel=0.05) + assert gaps[0] / gaps[1] == pytest.approx(59.0, abs=1.0) + + +def test_csbm_mixing_and_length(): + cs = load("csbm") + independent = cs.PI_0[:, None] * cs.PI_1[None, :] + expected = {0.2: (0.13589, 5.0), 0.4: (0.00512, 284.3)} + for alpha, (start, ratio) in expected.items(): + A = cs.uniform_reference(alpha, 2) + gap, _, target = cs.starting_gap(A, 3, cs.PI_0, cs.PI_1) + assert gap == pytest.approx(start, abs=5e-5) + later, _ = cs.d_imf(independent, A, 3, 2, target) + assert later[0] / later[1] == pytest.approx(ratio, rel=0.01) + A = cs.uniform_reference(cs.ALPHA, 2) + for steps, (start, ratio) in {2: (0.0574, 7.3), 5: (0.00156, 4674.0)}.items(): + gap, _, target = cs.starting_gap(A, steps, cs.PI_0, cs.PI_1) + assert gap == pytest.approx(start, abs=5e-5) + later, _ = cs.d_imf(independent, A, steps, 2, target) + assert later[0] / later[1] == pytest.approx(ratio, rel=0.01) + + +# %% Sections 7.26 and 7.27, branched.py + + +def test_two_branches_worked_out(): + br = load("branched") + linear = lambda t: 0.3 * t + late = lambda t: 0.3 * t ** 2 + assert float(br.branch_weights(linear, torch.tensor(0.5))) == pytest.approx(0.15) + assert float(br.branch_weights(linear, torch.tensor(1.0))) == pytest.approx(0.3) + near_zero = br.weighted_energy(2.0, lambda t: 1.0 - linear(t)) + near_one = br.weighted_energy(-2.0, linear) + assert near_zero == pytest.approx(1.7, abs=5e-5) + assert near_one == pytest.approx(0.3, abs=5e-5) + assert near_zero + near_one == pytest.approx(2.0, abs=5e-5) + far_linear = near_zero + br.weighted_energy(-4.0, linear) + far_late = br.weighted_energy(2.0, lambda t: 1.0 - late(t)) + br.weighted_energy(-4.0, late) + assert far_linear == pytest.approx(2.9, abs=5e-5) + assert far_late == pytest.approx(2.6, abs=5e-5) + + +def test_the_interpolant(): + br = load("branched") + grid = torch.tensor([0.0, 0.5, 1.0]) + assert br.interpolant(0.0, 2.0, 1.0, grid).tolist() == pytest.approx([0.0, 1.25, 2.0]) + assert br.interpolant_velocity(0.0, 2.0, 1.0, grid).tolist() == pytest.approx([3.0, 2.0, 1.0]) + assert br.interpolant_energy(0.0, 2.0, 1.0) == pytest.approx(2.1667, abs=5e-5) + assert br.interpolant_energy(0.0, 2.0, 0.0) == pytest.approx(2.0, abs=5e-5) + + +def test_entangled_tilt_and_parameterization(): + br = load("branched") + normalizer = 0.1 * 1.0 + 0.9 * 0.01 + assert normalizer == pytest.approx(0.109) + tilted = 0.1 / normalizer + assert tilted == pytest.approx(0.917, abs=5e-4) + assert tilted / 0.1 == pytest.approx(9.17, abs=5e-3) + reward = (1.0 - tilted) * math.log(0.01) + assert reward == pytest.approx(-0.380, abs=5e-4) + assert reward - math.log(normalizer) == pytest.approx(1.836, abs=5e-4) + assert -math.log(normalizer) == pytest.approx(2.216, abs=5e-4) + assert br.NOISE == pytest.approx(1.4142, abs=5e-5) + step = (-0.3 + 0.5) * 0.01 - 1.0 * 0.8 * 0.01 + br.NOISE * math.sqrt(0.01) * 0.4 + assert step == pytest.approx(0.0506, abs=5e-5) + assert 0.5 / br.NOISE == pytest.approx(0.3536, abs=5e-5) + assert 0.5 * (0.5 / br.NOISE) ** 2 * 0.01 == pytest.approx(0.000625, abs=5e-7) + force = br.compose_bias(torch.tensor([[0.5]]), torch.tensor([[[0.3, -0.4]]]), + torch.tensor([[[1.0, 0.0]]]))[0, 0] + assert force.tolist() == pytest.approx([0.5, -0.4]) + assert float(force.norm()) == pytest.approx(0.640, abs=5e-4) + weights = torch.softmax(torch.tensor([-1.0, 0.0, 1.0]), dim=0) + assert weights.tolist() == pytest.approx([0.090, 0.245, 0.665], abs=5e-4) + + +def test_entangled_training_lowers_the_loss(): + br = load("branched") + torch.manual_seed(0) + generator = torch.Generator().manual_seed(0) + force = br.BiasForce(width=16) + rollout = br.roll_out(force, 16, 8, torch.device("cpu"), generator) + weights = br.importance_weights(rollout) + assert float(weights.sum()) == pytest.approx(1.0) + assert br.force_check(rollout["positions"][-1]) < 1e-5 + loss = br.cross_entropy_loss(force, rollout, weights) + assert torch.isfinite(loss) + + +# %% Every file runs end to end + + +@pytest.mark.parametrize("name,flags", [ + ("optimal_transport", ["--steps", "300"]), + ("sinkhorn", ["--steps", "20"]), + ("displacement", ["--steps", "64"]), + ("ot_flow_matching", ["--steps", "3", "--samples", "64", "--sample-steps", "5"]), + ("static_bridge", ["--steps", "3"]), + ("dynamic_bridge", ["--steps", "20"]), + ("dsb", ["--steps", "2"]), + ("dsbm", ["--steps", "3"]), + ("sf2m", ["--steps", "21"]), + ("discrete_bridge", ["--steps", "4"]), + ("ddsbm", ["--steps", "20"]), + ("csbm", ["--steps", "3"]), + ("branched", ["--steps", "2", "--samples", "16"]), +]) +def test_the_script_runs(name, flags, monkeypatch, capsys): + module = load(name) + monkeypatch.setattr(sys, "argv", [name] + flags + ["--quiet"]) + module.main() + assert capsys.readouterr().out.strip()