generative-model-lab

A working notebook of generative modelling, built in public: small, runnable modules that trace one idea from sampling to structure generation, and that report failures as carefully as successes.

The through-line is a single question:

When a conditional generative model is asked to produce something new, how do we know it is still trustworthy once the data moves off the training distribution?

That question is asked three times, at three different levels of structure:

Module	Object being generated	Dimensionality	Status
01_mcmc_toy	samples from a 2D double-well distribution	2D, sampling	✅ runnable
02_diffusion_toy	samples from 2D toy densities (two-moons, 8-gaussians)	2D, diffusion	✅ runnable
03_omics_synthesis	synthetic patients for small-sample omics cohorts	tabular, ~10²–10³	🚧 v0 demo
04_se3_diffusion	protein-backbone-like frames on SE(3), with property guidance	3D rigid-body	🚧 v0 demo
05_spectral_view	the operators behind 01–04: the noising step, the denoiser, the guidance term	linear algebra	✅ runnable
06_krylov_mcmc	the reaction coordinate a d-dimensional density implies, recovered from the chain itself	d-dimensional, sampling	✅ runnable

Nothing here is a research contribution, and nothing here is a wrapper around someone else's checkpoint. The modules are deliberately small enough to be read end to end — the point is to make the mechanics visible: the noise process, the score, the equivariance, the operating point, and the external validation.

Module 05 is the one that looks at the machinery instead of the output: it treats the forward noising step as a linear operator (and computes its spectrum), the MMSE denoiser as an orthogonal projection (and measures the residual's inner products), and guidance as the gradient of a functional (and shows what a linear versus a radial functional does). It is the linear-algebra reading of modules 01, 02 and 04.

Quickstart

pip install -r requirements.txt

# 01: sampling, and what the operator language actually looks like as a plot
python 01_mcmc_toy/run_demo.py

# 02: a from-scratch 2D diffusion model (MLP + DDPM/DDIM)
python 02_diffusion_toy/run_demo.py

# 04: SO(3)/SE(3) frames, an equivariance test, and property-guided sampling
python 04_se3_diffusion/run_demo.py

# 05: the linear-algebra reading of the modules above (operators, projections, functionals)
python 05_spectral_view/spectral_forward.py
python 05_spectral_view/projection_demo.py
python 05_spectral_view/guidance_functionals.py

# 06: the reaction coordinate derived from the dynamics rather than chosen by hand (Krylov + MCMC)
python 06_krylov_mcmc/krylov_basics.py
python 06_krylov_mcmc/cv_discovery.py
python 06_krylov_mcmc/dimension_scan.py

# 03: small-sample omics augmentation, demo mode (no download needed)
python 03_omics_synthesis/run_experiment.py --demo

Each script writes figures to figures/ and prints a short summary table, so the whole repo can be checked in a couple of minutes on a laptop CPU.

What each module is actually about

01 — MCMC: getting the operator language out of the head and onto the screen

Random-walk Metropolis, MALA and HMC on the same symmetric double-well potential: trace plots, autocorrelation, ESS, acceptance rate, and barrier-crossing rates rather than raw counts, with error bars over 10 seeds. This is the "spectral gap" story told with pictures — mixing time is visible as slow ACF decay, and mode hopping is what the gap is about. The headline result is deliberately anti-textbook: with a cheap 2D target, tuned RWM is the most cost-effective of the three (18.8 ± 1.0 ESS per 1k evaluations, against 3.6 and 3.5), while HMC crosses the barrier twice as often per step but pays 12 units of cost per iteration for the privilege. show_one_step.py prints one complete step of each sampler - every term of the acceptance ratio - and checks three invariants of the target, exiting non-zero if any of them fails.

02 — Diffusion: the same sampling problem, solved by learning a score

A tiny MLP trained to denoise 2D toy densities, sampled with DDPM and DDIM. The interesting part is not that it works, but what changes when it does not: mode collapse on the 8-gaussians target, over-smoothed two-moons, and the trade-off between fewer DDIM steps and sample quality.

03 — Omics synthesis: does augmenting a small cohort actually help the external validation?

The honest version of this experiment is the one that is usually not reported. Synthetic samples usually improve the internal metrics and do nothing — or worse — for a cohort collected elsewhere, because the generator is happy to learn the batch effect instead of the biology. This module is built around that failure mode: the same classifier is evaluated on held-out internal data and on a study collected separately, and both numbers are reported.

Data: MetaboLights MTBLS1 (small, public, tabular). run_experiment.py --demo runs the whole pipeline on a synthetic cohort with the same shape, so the repo is runnable without any download.

04 — SE(3) diffusion: the same idea, one symmetry group later

Protein backbones are not vectors, they are rigid frames: $T_i = (R_i, t_i) \in SE(3)$, one per residue. Adding Gaussian noise to the entries of a rotation matrix breaks the symmetry the problem has, which is why this module implements the noise in the tangent space and maps it back with the exponential map. The test that matters is written as an assertion, not a picture:

$$f(Q \cdot T) = Q \cdot f(T) \quad \text{for a random global rotation } Q$$

and a small property-guided sampler (push two residues together, avoid clashes) — which is the mechanism that conditional structure generation (motif scaffolding, binder design) is built on.

06 — Krylov + MCMC: let the chain pick the reaction coordinate

Enhanced sampling needs a collective variable, and it normally gets one by hand — a distance, an angle, an RMSD. This module derives it instead: estimate a transfer operator from the trajectory, take its leading eigenvector, and check it three ways against a target whose answer is known exactly. The operator is never formed — a Lanczos run only applies it — and the headline numbers are that the CV lands 3.45 degrees from the true direction (top principal component: 64.9) and recovers 96 percent of the barrier (PCA: 0 percent), in eight dimensions, with the barrier along a dense random direction that no coordinate axis can see. The honest half is the dimension scan: d = 2 … 20 does not break it, a shrinking spectral gap does, and the ratio to watch is readable in the Ritz spectrum before the answer is known. Lanczos and Arnoldi are transcribed from Golub & Van Loan and Trefethen & Bau, and the module README says so before it says anything else.

Evaluation discipline (the part I care most about)

Whatever the module, the same three rules apply:

1. Report the operating point, not just the AUC. A ranking metric can look fine while the threshold that a clinic would actually use is useless.
2. Split by source, not by row. If the external cohort exists, it is the only number that matters; internal cross-validation is a sanity check.
3. Say what the generator learned that it should not have. Batch, site, plate, sampling date — if the synthetic data carries them, the metric will too.

Roadmap

See ROADMAP.md. Next up: real MetaboLights data in module 03 (replacing the demo cohort), a proper SE(3) transformer denoiser in module 04, and a molecule-level module (property-guided generation on a MoleculeNet property) as the AIDD-facing counterpart.

License

MIT — see LICENSE. No proprietary or patient-identifiable data is used anywhere in this repository; module 03 works on public datasets or on synthetic cohorts generated by the code itself.