PaperlayerView exact revision on arXiv

Source-linked AI summary

Blackout Diffusion: Generative Diffusion Models in Discrete-State Spaces

Javier E Santos, Zachary R. Fox, Nicholas Lubbers, Yen Ting Lin

arXiv:2305.11089v1cs.LGcs.CV

TL;DR

Many data domains are discrete, while diffusion models have largely relied on Gaussian processes and existing discrete-state theories have lacked an exact, non-variational treatment. This paper develops an exact framework for arbitrary discrete-state Markov processes and demonstrates it with Blackout Diffusion, which generates samples from an empty-image prior.

  • Problem

    Many real-world data domains are discrete, while most diffusion models use Gaussian processes and prior discrete-state formulations have not provided an exact, non-variational analysis.

    Source: §1 · pp. 1–2

  • Method

    The paper derives an exact framework for arbitrary discrete-state Markov processes, including reverse-time processes and score functions, and instantiates it with Blackout Diffusion.

    Source: §1 · p. 2

  • Results

    Blackout Diffusion generates CIFAR-10, CelebA, and Binarized MNIST samples from a completely black-image prior without treating data as continuous.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The Blackout Diffusion experiments show that Gaussian noise is not essential to generative diffusion modeling, since discrete-state Poissonian noise can also be used.

    Source: §4 · p. 7

    The derivation first treats a subset of Markov processes whose transitions are restricted to neighboring states and whose generator is tridiagonal.

    Source: §2.2 · p. 3

Abstract

from arXiv · show

Typical generative diffusion models rely on a Gaussian diffusion process for training the backward transformations, which can then be used to generate samples from Gaussian noise. However, real world data often takes place in discrete-state spaces, including many scientific applications. Here, we develop a theoretical formulation for arbitrary discrete-state Markov processes in the forward diffusion process using exact (as opposed to variational) analysis. We relate the theory to the existing continuous-state Gaussian diffusion as well as other approaches to discrete diffusion, and identify the corresponding reverse-time stochastic process and score function in the continuous-time setting, and the reverse-time mapping in the discrete-time setting. As an example of this framework, we introduce ``Blackout Diffusion'', which learns to produce samples from an empty image instead of from noise. Numerical experiments on the CIFAR-10, Binarized MNIST, and CelebA datasets confirm the feasibility of our approach. Generalizing from specific (Gaussian) forward processes to discrete-state processes without a variational approximation sheds light on how to interpret diffusion models, which we discuss.

1. Introduction

Most diffusion models use Gaussian processes, but many data domains are inherently discrete and exact theoretical treatments for arbitrary discrete-state processes are lacking. The paper develops an exact framework for such processes and introduces Blackout Diffusion as an example.

1. Introduction

  • Discrete data, including small-count biological measurements, categorical fluid regions, graphs, DNA, and protein sequences, motivate diffusion models that do not rely on continuous representations.The introduction notes that quantization and de-quantization may not be ideal for these domains.Source: §1 · p. 1
  • Existing work on discrete-state diffusion includes specific processes and variational formulations, but an exact, non-variational analysis had not been presented.The cited approaches include multinomial diffusion, masked diffusion, and continuous-time discrete-state models.Source: §1 · p. 2
  • The paper gives an exact framework for arbitrary discrete-state Markov processes and derives the reverse-time generator, score function, and discrete-time reverse mapping.It addresses whether discrete analogues of Brownian-bridge mappings and Stein score functions can be used.Source: §1 · p. 2
  • Blackout Diffusion uses a pure-death process with a completely black image as its prior and generates CIFAR-10, CelebA, and Binarized MNIST samples without treating data as continuous.The paper distinguishes it from Mask Diffusion despite both processes converging almost surely to a singular state.Source: §1 · p. 2

2. Method

The section develops a discrete-state Markov diffusion framework, then specifies Blackout Diffusion and its training and sampling procedures. It also relates the discrete-state theory to continuous Gaussian diffusion.

2.1. Notations

  • For an N-dimensional state space, each data coordinate takes values in a finite discrete set Ω, with states labeled and ordered from 0 to M.The notation applies to discrete- and continuous-time processes, indexed as X_k and X_t, respectively.Source: §2.1 · p. 3

2.2. Forward and Backward equations

  • The forward process follows Chapman–Kolmogorov equations, while its continuous-time form is a Master Equation governed by a generator and its adjoint.For transitions between neighboring states, step operators express the forward equation using state-dependent transition rates.Source: §2.2 · p. 3

2.3. Reverse-time process

  • Reverse-time transition rates depend on the forward solution; the formulation extends to arbitrary transition matrices, unordered state spaces, and discrete-time processes.The reverse process is explicitly time-dependent and describes how conditional probabilities evolve backward in time.Source: §2.3 · pp. 3–4

2.4. Relationship to Gaussian diffusion

  • As the number of densely distributed discrete states grows, both forward and reverse processes converge to Gaussian diffusion, with reversal consistent with Anderson’s reverse process.The continuous limit also identifies a discrete-state score function analogous to the Stein score, with a learning transformation for each transition event into a state.Source: §2.4 · p. 4

2.5. Blackout Diffusion

  • Blackout Diffusion uses a pure-death process: pixel intensities decay through neighboring states toward an all-black, single-point prior.The forward distribution is binomial, and the reverse process admits a finite-time binomial bridge; the paper uses continuous time for the 256-state image space.Source: §2.5 · p. 4

2.6. Diffusion Schedule and Loss functions

  • Observation times are chosen to uniformize changes in Fisher information, and the network predicts the difference X_0 − X_tk−1 from X_tk rather than transition rates.The likelihood-based loss has instantaneous and finite-time forms; with uniform observation times and small Δ, the finite-time form converges asymptotically to the instantaneous one.Source: §2.6 · pp. 5–6

2.7. Generating sample images

  • Generation starts from an all-black image and repeatedly predicts reverse changes, sampling them with either τ-leaping and Poisson draws or the exact binomial bridge.The selected terminal time is t_T = 15; the binomial bridge is specific to Blackout Diffusion.Source: §2.7 · p. 6

3. Numerical Experiments

The experiments test whether Blackout Diffusion is feasible for image generation, rather than aiming to advance the state of the art. Results report feasibility across three datasets and compare loss and sampling choices.

Experimental scope

  • Blackout Diffusion was evaluated on CIFAR-10, Binarized MNIST, and CelebA to demonstrate feasibility, not state-of-the-art performance.For CIFAR-10, the authors used NCSN++ with positivity and inference-clipping modifications and fixed T = 1000.Source: §3 · pp. 6–7

Results

  • The experiments suggest sample quality is not sensitive to the loss function, while binomial-bridge sampling is slightly advantageous.The paper reports CIFAR-10 FID and Inception score in Table 1, and results for Binarized MNIST and CelebA in Table 2.Source: Table 1 · p. 7, §3 · p. 7

Experimental setup

  • In this feasibility experiment, with T = 1000 and both loss functions evaluated under binomial and Poisson sampling, each was trained once due to limited computational resources using two NVIDIA A100 GPUs for ∼72 hr.Source: §3 · p. 7

4. Discussion and Conclusion

The paper develops an exact framework for discrete-state Markov diffusion and uses Blackout Diffusion to show that generative diffusion need not rely on Gaussian noise. It also discusses how the framework informs interpretation and future study of diffusion models.

Theoretical contributions

  • The framework generalizes diffusion models from Wiener-process SDEs to discrete-state Markov processes and derives discrete-state score functions and an exact maximum-likelihood loss.The variational bound of Campbell et al. (2022) is tight in the continuum-time limit; discrete-state score-based modeling remains for future work.Source: §4 · p. 7

Interpretation of diffusion

  • Blackout Diffusion demonstrates that Poissonian noise can support generative diffusion, so Gaussian noise is not essential in this example.The paper reports this conclusion from its numerical experiment on Blackout Diffusion.Source: §4 · p. 7
  • Blackout Diffusion’s forward process adds noise for t < log 2 but removes it for t > log 2 as samples converge to a blank image.The paper says whether a process adds or removes noise depends on the stochastic process and how noise is defined; its forward process degrades samples and its reverse process reconstructs them.Source: §4 · p. 8

Implications and future work

  • The authors argue that a clear theoretical framework lets researchers examine forward evolution, reverse processes, and loss functions more closely than ad hoc performance tuning alone.They caution that ad hoc procedures may be less transferable across tasks or datasets.Source: §4 · p. 9
  • The authors hypothesize that discrete-state models may be more computationally efficient than Gaussian models because their state spaces are smaller.They cite reasonable results using NCSN++ out of the box as an empirical hint, not a demonstration of superior efficiency.Source: §4 · p. 9
  • The framework’s instantaneous loss needs no analytic bridge formula, and discrete-state forward evolution can be computed numerically using transition-matrix powers or exponentiation.The authors anticipate these properties could support a library of discrete-state processes matched to discrete data and systems.Source: §4 · p. 9

Reproducibility and code availability

The authors state that experiment code is available in a GitHub repository.

Reproducibility and code availability

  • The experiment code is deposited in the Blackout-Diffusion GitHub repository.The deposit has C-number C23047 approval from the Richard P. Feynman Center for Innovation at Los Alamos National Laboratory.Source: Reproducibility and code availability · p. 9

A. Deriving the Reverse-Time Evolutions

The appendix derives reverse-time evolution equations for discrete-time and continuous-time processes, then extends the derivation to arbitrary discrete-state Markov transitions.

A. Deriving the Reverse-Time Evolutions

  • The continuous-time derivation evolves the joint conditional distribution using the forward solution, then divides by a finite terminal probability to obtain the reverse-time equation.The discrete-time derivation follows the same general approach with additional technical details and yields the reverse-time stochastic process.Source: A. Deriving the Reverse-Time Evolutions · pp. 13–14
  • The reverse-time process permits transitions opposite to the forward transitions, but its rates or probabilities are not identical to the forward ones.For discrete time, the derivation uses conditional Bayes’ rule to prescribe reverse transition probabilities.Source: A. Deriving the Reverse-Time Evolutions · p. 15

B. Generalization of the Theoretical Formulation for Arbitrary Discrete-State Markov Processes

  • For a finite state space, the derivation represents the transition matrix with powers of shift operators for general n-step transitions, then applies the reverse-time procedure.The reverse-time equation uses the pre-image of the forward transition.Source: B. Generalization of the Theoretical… · p. 15

C. Kramers–Moyal expansion of the reverse-time evolution in discrete-state spaces

The Kramers–Moyal expansion connects discrete-state reversal to a continuous-time SDE and identifies a discrete-state score function. For the pure-death process, the conditional state at an intermediate time follows a binomial bridge distribution.

C. Kramers–Moyal expansion of the reverse-time evolution in discrete-state spaces

  • The reverse-time expansion needs asymptotic matching because its transition rates depend on non-local information at transition preimages.The analysis assumes a large system size V and expands the scaled test function.Source: C. Kramers–Moyal expansion of the… · pp. 15–16
  • In the dense-grid limit, reverse transition rates match forward rates up to a correction proportional to the density derivative, yielding the reverse-time Itô SDE.The correction is (σ/V) ∂_m̃ρ(m̃,s).Source: C. Kramers–Moyal expansion of the… · pp. 16–17
  • The discrete-state score depends on the forward solution at the preimage of a viable transition and is asymptotically proportional to the continuous-state Stein score.The proportionality constant is 1/V, proportional to the forward solution's variance.Source: C. Kramers–Moyal expansion of the… · p. 17
  • The paper identifies the discrete-state score but does not directly learn this function.For general Markov processes with multiple viable transitions into a state, the analysis gives distinct score functions for each transition.Source: C. Kramers–Moyal expansion of the… · p. 17

D. Deriving the binomial bridge formula

  • For the pure-death process, conditioning on the state at time t makes the intermediate state at time s a binomially distributed bridge.Its parameters use r = (e^-s − e^-t)/(1 − e^-t), with the remaining population at t setting the number of trials.Source: D. Deriving the binomial bridge formula · p. 18

E. Observation times

The section proposes choosing observation times according to the forward process's Fisher information. In the experiments reported, this heuristic performed significantly better than the other tested timing heuristics.

E. Observation times

  • The proposed schedule places more observation times where the forward process has higher Fisher information.The authors define a time density proportional to the time-dependent part of the Fisher information.Source: E. Observation times · p. 18
  • The time density is symmetric around t = log 2, when the images are maximally noisified.The schedule handles endpoint singularities by placing half the times on each side and choosing a finite final time.Source: E. Observation times · p. 18, Figure 4 · p. 19
  • Fisher-information-based observation times performed significantly better than the other tested heuristics.Comparisons included entropy, entropic production rate, variance, KL divergence between observation times, coefficients of variation, and a uniform temporal grid.Source: E. Observation times · p. 19

F. Likelihood functions

The section derives likelihood losses for learning reverse transition rates in continuous-time discrete-state systems. It gives an instantaneous Bernoulli-based loss and a continuum-time full likelihood, with a finite-time correction for the Blackout process.

F. Likelihood functions

  • The instantaneous loss is the KL divergence between Bernoulli distributions induced by the ground-truth and model-predicted transition processes.For a short interval, the model assigns transition and no-transition probabilities κθ(t)Δt and 1 − κθ(t)Δt.Source: F. Likelihood functions · p. 19
  • The continuum-time full negative log-likelihood integrates κθ(τ) − λ(τ) log κθ(τ) over time, up to a θ-independent constant.The full likelihood is formed by integrating the instantaneous expression over 0 ≤ t < ∞.Source: F. Likelihood functions · p. 19
  • In practice, training samples an observation time and reweights the loss by the Fisher-information-based density φ(t).The time is sampled in a Monte Carlo fashion, with one observation time drawn for each training image.Source: F. Likelihood functions · pp. 19–20
  • For finite observation intervals, the constant-rate τ-leaping model approximates the true binomial distribution with a Poisson distribution, and the section derives a corrected likelihood.The model-predicted distribution is Poisson(κθ), while the ground-truth probability mass function is the binomial bridge distribution.Source: F. Likelihood functions · p. 20

G. Algorithms for general continuous-time discrete-state Markov processes

The section extends training and inference procedures to general continuous-time discrete-state Markov processes with multiple transition types. The procedures require a supplied forward solution, and inference uses τ-leaping when the binomial bridge is unavailable.

G. Algorithms for general continuous-time discrete-state Markov processes

  • The general procedures assume a supplied forward solution and represent each state's possible transitions with indexed preimages and forward rates.The framework allows R potential transitions from each discrete state.Source: G. Algorithms for general… · p. 20
  • Training samples a transition type and uses a neural network to estimate its reverse rate, while inference applies sampled reverse events using τ-leaping.For general cases, the binomial bridge formula no longer holds, so inference falls back to τ-leaping.Source: G. Algorithms for general… · p. 21
  • The exact reverse-time process agrees with Campbell et al. (2022), indicating their variational bound is tight in the continuum-time limit.The authors report preliminary, not shown, feasibility for learning a birth-and-death process on CIFAR-10 and focus this paper on Blackout Diffusion.Source: G. Algorithms for general… · p. 21

H. Binarized MNIST dataset

Blackout Diffusion was trained on resized, threshold-binarized MNIST images and generated samples with an FID of 0.023 against the training dataset.

H. Binarized MNIST dataset

I. Celeb-A 64x64 dataset

Blackout Diffusion was trained on CelebA 64x64 with the infinitesimal loss; 50,000 generated samples had an FID of 3.22 against the training dataset, and Figure 13 shows generated images.

I. Celeb-A 64x64 dataset

  • The CelebA 64x64 model used the infinitesimal loss function 11.Source: §I · p. 21
  • FID was 3.22 for 50,000 generated samples compared with the CelebA training dataset.Source: §I · p. 21, Figure 13 · p. 29
  • Figure 13 shows 144 samples generated by Blackout Diffusion trained on Celeb-A using the instantaneous loss function (Eq. 11), with τ-leaping (Poisson random number) inference and FID=3.22.Source: Figure 13 · p. 29
Loading 2305.11089v1…