Delete lecture_3

lecture_3/README.md

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

lecture_3/ddpm.py

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

lecture_3/forward_diffusion.py

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

lecture_3/guidance.py

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

lecture_3/latent_diffusion.py

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

lecture_3/samplers.py

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

lecture_3/score_matching.py

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

lecture_3/sde.py

@@ -1,249 +0,0 @@
-"""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()