Bridging discrete and continuous state spaces: Exploring the Ehrenfest process in time-continuous diffusion models

TL;DR

The paper addresses bridging discrete-state Markov jump processes with continuous-state diffusion models by focusing on the Ehrenfest process as a discrete analogue of the Ornstein–Uhlenbeck process. It derives time-reversed rates via conditional expectations and introduces a loss function to learn the reversed dynamics, elucidating a direct connection to score-based generative modeling in the OU limit. Computational strategies, including dimensional factorization and tau-leaping, enable scalable training and sampling, while numerical experiments on MNIST and CIFAR-10 demonstrate competitive performance and reveal a pathway for transferring insights between discrete and continuous diffusion frameworks. Overall, the work provides a principled theoretical link and a practical algorithm for leveraging discrete-state diffusion in contexts traditionally addressed by continuous-score methods, with potential benefits for discrete data modalities and future theoretical generalization.

Abstract

Generative modeling via stochastic processes has led to remarkable empirical results as well as to recent advances in their theoretical understanding. In principle, both space and time of the processes can be discrete or continuous. In this work, we study time-continuous Markov jump processes on discrete state spaces and investigate their correspondence to state-continuous diffusion processes given by SDEs. In particular, we revisit the $\textit{Ehrenfest process}$, which converges to an Ornstein-Uhlenbeck process in the infinite state space limit. Likewise, we can show that the time-reversal of the Ehrenfest process converges to the time-reversed Ornstein-Uhlenbeck process. This observation bridges discrete and continuous state spaces and allows to carry over methods from one to the respective other setting. Additionally, we suggest an algorithm for training the time-reversal of Markov jump processes which relies on conditional expectations and can thus be directly related to denoising score matching. We demonstrate our methods in multiple convincing numerical experiments.

Figures (9)

• Figure 1: We display two time-reversed processes from $t = 2$ to $t = 0$ that transport a standard Gaussian (left panels, in green) to a multimodal Gaussian mixture model (left panels, in orange), or a binomial distribution to a binomial mixture, respectively, once using a diffusion process in continuous space (upper panel) and once a time-reversed (scaled) Ehrenfest process in discrete space with $S=100$ (lower panel). Crucially, in both cases we use the (state-continuous) score function to employ the time-reversal, which for this problem is known analytically, see \ref{['app: Gaussian mixture score']}. The plots demonstrate that the distributions of the processes seem indeed very close one another, implying that the approximation \ref{['eq: relation conditional expectation and score']} is quite accurate even for a moderate state space size $S + 1$.
• Figure 2: We plot histograms of $500.000$ samples from the time-reversed scaled Ehrenfest process at different times. The processes have been trained with three different losses.
• Figure 3: MNIST samples obtained with the time-reversed scaled Ehrenfest process which was trained with $\mathcal{L}_\mathrm{OU}$.
• Figure 4: CIFAR-10 samples from the Ehrenfest process with a pretrained model, further finetuned with $\mathcal{L}_\mathrm{OU}$.
• Figure 5: CIFAR-10 samples from the Ehrenfest process with a pretrained model, further finetuned with $\mathcal{L}_\mathrm{Taylor}$.

Theorems & Definitions (21)

• Lemma 2.1
• proof
• Remark 2.2: Conditional expectation
• Lemma 3.1
• proof
• Proposition 3.2: State space limit of Ehrenfest process
• Remark 3.3: Learning of conditional expectation
• Proposition 3.4
• proof
• Remark 3.5: Convergence of the time-reversed Ehrenfest process