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

Source: arXiv - 2512.05092v1

Overview

Diffusion models have become the go‑to technique for generating images, audio, and even text, but most tutorials assume the data lives in a Euclidean space. This paper lifts that restriction and builds a single, self‑contained theory that works for both continuous domains (e.g., pixel values) and discrete structures (e.g., token sequences). By unifying stochastic differential equations (SDEs) with continuous‑time Markov chains (CTMCs), the authors give developers a clear roadmap for extending diffusion‑based generation to any kind of data.

Key Contributions

• Unified framework that treats diffusion on arbitrary state spaces—continuous ℝⁿ, finite alphabets, or hybrids.
• Discrete‑time and continuous‑time derivations side‑by‑side, showing how forward noising kernels translate into reverse‑time dynamics.
• General ELBO formulation that recovers the standard training loss for both Gaussian and categorical corruptions.
• Catalog of forward corruption kernels (Gaussian, uniform, masking/absorbing, etc.) and analysis of how each shapes the reverse process.
• Pedagogical layering: a gentle intro for newcomers, a synthesis for practitioners, and a deep‑theory bridge for experts in continuous diffusion.
• Reusable proof toolkit (Fokker–Planck, master equation, variational identities) that can be plugged into future diffusion research.

Methodology

1. Forward Process

   • Continuous: Apply a Gaussian Markov kernel at each timestep, which in the limit becomes an SDE of the form

     dx = f(x,t)dt + g(t)dW
     

   • Discrete: Use a Markov transition matrix (e.g., uniform mixing, token‑masking, or absorbing states) that defines a CTMC on a finite alphabet.

2. Reverse Process

   • Derive the time‑reversed dynamics using the Fokker–Planck equation for SDEs and the master equation for CTMCs.
   • Show that the reverse kernel can be expressed as a learned neural network that approximates the true reverse drift or transition probabilities.

3. Variational Objective

   • Start from the joint distribution of data and noisy latent variables.
   • Apply the standard ELBO trick to obtain a tractable loss that decomposes into a reconstruction term and a KL term, valid for any state space.

4. Bridging Discrete & Continuous

   • Map discrete transition kernels onto continuous‑time generators, highlighting the mathematical analogy (e.g., diffusion coefficient ↔ transition rate matrix).
   • Provide a “dictionary” that lets practitioners translate intuition from image diffusion to token diffusion (and vice‑versa).

The whole development stays at a level where a developer familiar with basic probability and neural nets can follow the derivations without needing deep stochastic calculus.

Results & Findings

• Theoretical equivalence: The ELBO derived for discrete CTMCs reduces exactly to the familiar diffusion loss when the state space is ℝⁿ and the forward kernel is Gaussian.
• Kernel impact: Different forward corruptions lead to markedly different reverse dynamics; for instance, masking kernels produce sparse gradients that are easier to learn for language models.
• Empirical sanity checks (illustrative experiments): training a simple categorical diffusion model on MNIST digits (treated as 10‑class labels) matches the performance of a continuous‑pixel diffusion model when the same ELBO is used.
• Proof reuse: The authors demonstrate that a handful of core identities (e.g., change‑of‑measure for Markov processes) suffice to re‑derive most existing diffusion results, confirming the unifying power of their framework.

Practical Implications

• Broader data modalities: Engineers can now design diffusion pipelines for graphs, molecules, or code tokens without reinventing the math from scratch.
• Custom corruption strategies: By picking a forward kernel that respects domain structure (e.g., masking only syntactically valid tokens), the reverse model learns more efficiently, potentially reducing training time and improving sample quality.
• Interoperable libraries: The paper’s modular view encourages the creation of diffusion libraries where the forward kernel is a plug‑in component, making it trivial to switch between Gaussian noise, uniform mixing, or task‑specific corruptions.
• Hybrid models: For multimodal tasks (image + caption), one can run a continuous SDE on pixel space and a CTMC on the caption simultaneously, using a shared latent schedule.
• Better debugging: Understanding the forward‑reverse relationship through the master/Fokker–Planck equations gives developers analytical tools to diagnose training instability (e.g., mismatched noise schedules).

Limitations & Future Work

• The paper focuses on theoretical unification and provides only minimal empirical validation; large‑scale benchmarks (e.g., ImageNet, large language models) are left for future studies.
• Scalability of discrete kernels: While the framework supports arbitrary transition matrices, constructing efficient, expressive kernels for very large vocabularies remains an open engineering challenge.
• Extending the theory to continuous‑discrete hybrid spaces (e.g., diffusion on manifolds with categorical attributes) is mentioned but not fully explored.
• The authors suggest investigating adaptive noise schedules that are jointly optimized across state‑space types, and formalizing privacy‑preserving diffusion where the forward kernel incorporates differential‑privacy noise.

Authors

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

Paper Information

• arXiv ID: 2512.05092v1
• Categories: stat.ML, cs.LG
• Published: December 4, 2025
• PDF: Download PDF