AISTATS_VESDE: The Diffusion Repo That Already Knows the Answer

By swapping a learned score network for a closed-form Gaussian mixture, this tiny research codebase turns sampler design into a test of discretization error, correction steps, and variance schedule choice.

8 min read • View on GitHub • More from wanshuiyin

A wide editorial scene shows a laboratory bench where two glass vessels feed a single analytic machine. One vessel carries a tight inward spiral, while the other releases a widening cloud that is held in check by a correction ring. The image explains that the repo compares sampler geometry, not image quality. — This is not a model zoo. It is a controlled lab where the target distribution is known first, so the sampler becomes the thing under inspection.

This folder contains two .py files, DPUM.py is the DPUM algorithm implemented in [1], and PFODE_Corrector_VESDE.py is the VE-based algorithm we improved on it. Readers only need to click run, the file will automatically generate a closed-form score function and run the algorithm to calculate the KL divergence.

wanshuiyin, Project Creator and Sole Contributor · Repository: wanshuiyin/AISTATS_VESDE

Key Takeaways

AISTATS_VESDE turns diffusion sampling into a controlled experiment by making the target score analytic instead of learned.
That setup removes model approximation from the story and exposes the sampler, schedule, and corrector as the real variables.
The repo's main contrast is geometric: a DPUM baseline on one side, a VE-style predictor-corrector path on the other.
Its flat two-script structure is a feature because it keeps the experiment reproducible and easy to inspect.

The setup: a diffusion experiment with the answer sheet visible

Most diffusion repos try to make the score network stronger. AISTATS_VESDE takes a sharper route. It builds a Gaussian mixture with a closed-form score, then asks a narrower question: when the answer is known, which sampler gets closest with the least numerical damage?

That framing explains why the codebase feels small but serious. Two Python scripts are enough because the experiment is synthetic, reproducible, and hostile to hand-waving. Every run compares sampling behavior against the same analytic target.

Why closed-form scores change the entire experiment

In a learned diffusion model, the score estimate is part of what you are trying to improve. Here, the score comes from the target density itself. That removes approximation error from the equation and leaves the sampler, the noise schedule, and the correction step exposed.

The repository's loop stays analytic from end to end, so the interesting part is how the sampler and corrector behave when the target is already known.

A close-up of overlapping Gaussian hills drawn like contour ridges on tracing paper, except the background is pure white. Fine score vectors point inward and outward around the density, and a caliper measures the distance between a sample path and the target contour. The image explains why exact score computation changes the experiment. — The closed-form target is the point. Once the score is exact, the sampler no longer hides behind model error.

DPUM versus VE-corrector: same target, different geometry

This is the repo's main technical contrast. Both scripts work against the same analytic mixture, but they push the sampler through different geometry. One follows the DPUM baseline from Chen et al. through a variance preserving path. The other shifts into a variance exploding setup and adds a predictor-corrector step to pull samples back into alignment.

Approach	Score source	Noise geometry	Corrector	What it isolates	Practical takeaway
DPUM baseline	Closed-form score on the synthetic mixture	Variance preserving	Underdamped correction	Solver behavior under the original PFODE setup	A strong baseline for measuring discretization error
VE-corrector variant	Closed-form score on the same mixture	Variance exploding	Predictor-corrector with a Langevin-style nudge	Whether VE geometry stabilizes the path	The repo's main experimental twist
Typical learned-score diffusion repo	A neural network approximation	Usually fixed by the model family	Optional or bundled	Training error mixed with sampling error	Harder to tell what actually improved

A split editorial scene shows two paths toward the same basin. On the left, a narrow trajectory collapses inward through a tight corridor. On the right, a wider cloud expands first, then a thin correction spring guides it back toward the basin. The image explains the geometric difference between the two sampling strategies. — The table gives the facts. The illustration gives the shape of the trade-off.

The interesting part is not that one path looks fancier. It is that the repo keeps the target fixed while changing the sampler geometry, which makes the comparison legible. That is a rare luxury in diffusion work, where training, architecture, and sampler quality usually blur together.

Inside the code: two scripts, one controlled experiment

DPUM.py and PFODE_Corrector_VE.py share the same basic experiment loop. A mixture generator creates random component means and covariances, the analytic score is computed from that synthetic density, and the sampler steps are evaluated against the same target.

That flat structure matters. There is no framework layer, no training pipeline, and no hidden abstraction stack. The repo reads like research notebook code that was cleaned just enough to make the argument reproducible.

Where this sits in the diffusion landscape

This repository is not competing with image generators or large training libraries. It sits closer to a microscope. Its job is to strip away the expensive, fuzzy parts of diffusion research and leave one question on the table: how much of the result is the solver, and how much is everything else?

That makes the project small in scope but sharp in intent. If you care about sample quality, solver efficiency, or the geometry of diffusion dynamics, this is useful precisely because it refuses to be broad.