lecture_3/README.md

# Lecture 3: Score Matching and Diffusion Models

Chapter 3 of the notes builds a generator out of corruption. A stochastic differential equation
pushes a clean sample toward a standard Gaussian, 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. To reverse the motion we need the
gradient of the log density, the score, and we learn it by regression on the noise that was drawn.
Discretize the reverse equation on a grid and you get DDPM; rewrite the Laplacian as a divergence
and you get 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. Chapter 2
runs its clock the other way, and we keep this one 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 we compare the three 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` we work with a process we already know, so
`--steps` sets the number of simulation steps, measured levels or reverse steps; the other three
train a network, so there `--steps` 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, and their ordering 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 here in the code. 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. We decode the generated latents and get 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 the width and the
  training budget are flags, so you can raise either one. 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. Each lecture file stands on its own, 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, so you can compare the three samplers and the guidance weights
  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))`,
  in place of 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.
  We write the Hutchinson estimator 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. Offline, a fixed random projection stands in for the
  language-model head; the decoding path and the composition constraint run unchanged, and the
  residues they read out are arbitrary. `--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 we read the two properties off its coordinates directly and keep the
  same signature, which leaves nothing to train.