Gaussian SDEs and masked CTMCs share one ELBO in this general-state-space primer

Foundations of Diffusion Models in General State Spaces: A Self-Contained Introduction

Vincent Pauline, Tobias Höppe, Kirill Neklyudov, Alexander Tong, Stefan Bauer, Andrea Dittadi

stat.ML, cs.LG

2025-12-05

A TUM/Mila primer puts Gaussian and categorical diffusion in one Markov frame: fix the forward noising, learn the reverse, train with an ELBO; limits are SDEs and CTMCs.

What problem this solves

Introductory notes on diffusion almost always start in Euclidean space: pixels are reals, noise is Gaussian. Discrete objects (tokens, residues, categories) live in a second literature, with different kernels, losses, and samplers. Hierarchical VAEs already stacked latent layers, but training required simulating long inference chains, intermediate marginals were non-identifiable, and it was unclear whether simple parametric conditionals could carry the load. Sohl-Dickstein et al. (2015) fixed the forward noising process. This TUM / Helmholtz / MCML / Mila primer takes that move and writes continuous and discrete state spaces in parallel, from discrete time through continuous-time limits, then closes with the infinitesimal generator.

Method

The recipe has three steps.

The forward process noises each coordinate independently, so the analysis can stay one-dimensional. The reverse process cannot: correlations in the data have to be rebuilt, so the network sees the full noisy state and emits per-coordinate reverse parameters.

On the continuous side, a one-step kernel is a convex mix of signal coefficient \(\alpha\) and noise \(\sigma\). As \(T\) goes to infinity the discrete update converges to an SDE whose drift is \((\mathrm{d}\log\alphat/\mathrm{d}t)\,xt\) and whose diffusion satisfies \(gt^2=\alphat^2\,\mathrm{d}/\mathrm{d}t(\sigmat^2/\alphat^2)\). The reverse SDE adds a \(g^2\nabla\log q\) term, i.e. the score. Absorbing the diffusion into the drift yields the probability-flow ODE.

On the discrete side, the transition matrix is the same convex mix: keep the token, or resample from a reference. The reference can be uniform, an absorbing [MASK] state, or a mixture. The continuous-time limit is a CTMC with rate matrix \(Rt\) (nonnegative off-diagonal, columns sum to zero). The reverse CTMC tilts those rates by the instantaneous marginal.

Masked diffusion is stricter: the only forward jump is token to [MASK]. After dropping parameter-free terms, the ELBO collapses to a time-weighted masked language modeling loss, weight \(\alpha't/(1-\alphat)\), only on masked positions.

Section 7 writes both sides as an infinitesimal generator \(\mathscr{L}t\). The operator describes how expectations of test functions move locally; its adjoint gives the Kolmogorov forward equation, covering Fokker-Planck for SDEs and the master equation for CTMCs. Jump-diffusions sit in the same slot. Section 8 moves diffusion into a learned latent, reviews embedding discrete data into continuous diffusion, and points at hybrids (Duo, CADD, CCDD).

Results

This is a tutorial, not a bake-off. The checkable claims are identities.

SettingForward limitTraining objective
Continuous GaussianSDE\(\frac12\int gt^2\\nabla\log q(xt\mid x0)-s\theta\^2\mathrm{d}t\)
Discrete absorbing / maskCTMCweighted MLM, weight \(\alpha't/(1-\alphat)\)
Bothgenerator \(\mathscr{L}t\)one ELBO on path space

Gaussian and categorical forwards both give closed-form \(q(xt\mid x0)\), so training never simulates the full noising chain. In continuous time the reconstruction term \(\mathcal{L}0\) vanishes.

Why it matters

If you already know DDPM or score SDEs, this is the shortest contrastive route into discrete diffusion: Gaussian kernel versus categorical kernel, score versus reverse rate. If you train masked diffusion LMs, it says why the loss looks like BERT MLM and where the time weight comes from.

The generator view is for later. Jump-diffusions and piecewise-deterministic processes do not fit SDE notation; they fit this operator.

It is not a methods paper. No new backbone, no FID or perplexity table. Reading it will not speed up a sampler.

Limitations

Engineering is out of scope by design: architectures, sampler tricks, schedule tuning. The reverse process is factorized per coordinate, which is an approximation. Variable length (insert/delete, Levenshtein-style) is cited, not developed at the same depth as substitution. Duo, CADD, and CCDD are 2025 pointers, not experiments in this manuscript. Generator derivations assume regularity conditions that implementations typically take as given on high-dimensional discrete sequences.

Terms

Source

What people are saying

Related papers

All paper explainers