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.