A Continuous Time Framework for Discrete Denoising Models

Andrew Campbell; Joe Benton; Valentin De Bortoli; Tom Rainforth; George Deligiannidis; Arnaud Doucet

Source-linked AI summary

A Continuous Time Framework for Discrete Denoising Models

Andrew Campbell, Joe Benton, Valentin De Bortoli, Tom Rainforth, George Deligiannidis, Arnaud Doucet

arXiv:2205.14987v2May 30, 2022stat.MLcs.LG

TL;DR

Discrete denoising models largely rely on discrete-time procedures, limiting when the model learns to denoise and how the reverse process is sampled. The paper formulates both directions as continuous-time Markov chains, trains the reverse process with a continuous-time ELBO, and simulates it with tau-leaping. Its predictor-corrector sampler improves sample quality over discrete-time methods on discrete data, while the framework also yields an error bound for tau-leaping samples.

Problem
Existing denoising approaches for discrete data operate in discrete time, restricting denoising to fixed time points and limiting reverse sampling flexibility.
Method
The paper formulates forward and reverse discrete denoising processes as Continuous Time Markov Chains, trains an approximate reverse process with a continuous-time ELBO, and simulates high-dimensional chains using tau-leaping.
Results
The predictor-corrector sampler improves sample quality over discrete-time methods on discrete data, and the framework derives a bound on error between true and generated sample distributions.
Takeaways & Limitations
Continuous-time modeling enables flexible reverse sampling schemes for discrete data and supports theoretical analysis of tau-leaping approximation error.
Takeaways & Limitations
The bound grows at worst quadratically with dimension, is not expected to be particularly tight, and is impractical to compute because key constants are difficult to obtain.

Abstract

from arXiv · show

We provide the first complete continuous time framework for denoising diffusion models of discrete data. This is achieved by formulating the forward noising process and corresponding reverse time generative process as Continuous Time Markov Chains (CTMCs). The model can be efficiently trained using a continuous time version of the ELBO. We simulate the high dimensional CTMC using techniques developed in chemical physics and exploit our continuous time framework to derive high performance samplers that we show can outperform discrete time methods for discrete data. The continuous time treatment also enables us to derive a novel theoretical result bounding the error between the generated sample distribution and the true data distribution.

1 Introduction

Discrete data applications motivate extending denoising models beyond their predominantly continuous-state setting. The paper proposes a continuous-time formulation intended to support flexible reverse sampling for such data.

Discrete data arise in text, segmentation maps, categorical features, discrete latent spaces, and 8-bit images.
Discrete-time denoising models restrict learning to fixed time points and generally limit sampling to ancestral strategies.
Continuous-time models learn to denoise at arbitrary times, enabling more flexible reverse sampling schemes.
The paper formulates discrete denoising as forward and reverse Continuous Time Markov Chains and trains with a continuous-time ELBO.
The framework uses tau-leaping and predictor-corrector sampling, derives a sample-distribution error bound, and evaluates image and monophonic-music generation.

2 Background on Discrete Denoising Models

Discrete denoising models corrupt data through a sequence of kernels toward an easy-to-sample distribution, then approximate the intractable reverse process with a learned model. Training uses a discrete-time ELBO.

The forward process transforms pdata into a terminal distribution that approximates an easy-to-sample stationary distribution.
Forward kernels are designed to mix reasonably quickly and admit the reference distribution as a stationary distribution.
If exact reverse kernels and terminal samples were available, ancestral reverse sampling would recover the data distribution.
The exact reverse kernel is intractable, so a parametric reverse kernel is used instead.
The denoising model is trained by minimizing the negative discrete-time ELBO, an upper bound on negative model log-likelihood.
Efficient training samples random ELBO terms from minibatches and performs stochastic gradient updates.

3 Continuous Time Framework

The paper replaces discrete transition kernels with a continuous-time Markov chain whose reverse-time process is defined by a time-reversed rate matrix. An approximate reverse rate is trained with a continuous-time ELBO and simulated efficiently.

3.1 Forward process and its time reversal: A CTMC allows state transitions at any time, specified by an initial distribution and a time-dependent transition rate matrix Rt.
3.1 Forward process and its time reversal: The forward rate Rt is designed to mix toward an easy-to-sample stationary distribution and support analytically tractable conditional marginals.
3.1 Forward process and its time reversal: Time reversal defines a reverse CTMC that starts from qT and ends at pdata, with reverse rates related to the forward rates.
3.1 Forward process and its time reversal: The exact reverse rate is intractable because it depends on the forward marginal and conditional distribution, so the model uses an approximation.
The continuous-time negative ELBO provides an upper bound on negative model log-likelihood for training the approximate reverse process.
Training samples data, a random time, a noised state, and an auxiliary transition pair to learn reversal of the noising process.

4 Efficient Forward and Backward Sampling

The framework factorizes high-dimensional CTMCs across dimensions and uses tau-leaping to simulate the reverse process efficiently. Its error bound separates reverse-rate, tau-leaping, and forward-chain mixing errors, with dimension dependence growing at worst quadratically.

Factorizing Over Dimensions: Factorizing the forward process makes only D × (S −1) + 1 rate values non-zero, enabling tractable simulation and preserving one-dimensional transitions in the exact forward and reverse CTMCs.Each dimension propagates independently, so simultaneous transitions have probability zero in the continuous-time process.
Simulating the Reverse Process: The reverse process can be sampled exactly with Gillespie’s Algorithm, but this is inefficient for large D because each step changes only one dimension.The method therefore uses tau-leaping instead of computing the intractable full matrix exponential.
Tau-Leaping: Tau-leaping assumes the reverse rate is constant over [t −τ, t], counts transitions as Poisson variables, and applies them simultaneously at t −τ.This permits multiple dimensions to change in one step; multiple jumps within one dimension are meaningful for ordinal data but rejected for categorical data.
Tau-Leaping: The tau-leaping approximation improves with smaller τ, recovering exact simulation as τ →0, while larger τ trades sample quality against compute by increasing the average number of dimensions changing per step.The paper names this reverse-CTMC sampler τLDR.
Error Bound: The error bound contains reverse-rate approximation, tau-leaping, and forward-chain mixing terms; the tau-leaping contribution is linear in both T and τ.As tau-leaping steps become arbitrarily small, its approximation error goes to zero.
Error Bound: The bound grows at worst quadratically in dimension D, but is not expected to be tight in practice and is impractical to compute because M, C1, and C2 are difficult to obtain.Its purpose is primarily to show that τ need not be impractically small in high dimensions.

5 Related Work

Earlier discrete denoising approaches include categorical, text, segmentation-map, and other corruption processes, but train and sample in discrete time. This work extends that line by using continuous time for both training and sampling.

Earlier work applied denoising models to binary, categorical, text, and segmentation-map data using several corruption and reverse-kernel designs.
Prior approaches used uniform or spatially biased forward kernels and denoising-model parameterizations for discrete data.
These approaches trained and sampled in discrete time, despite discussion of continuous-time forward processes.
The paper extends prior work by training and sampling discrete denoising models in continuous time.

6 Experiments

Experiments test reverse-CTMC simulation on a synthetic distribution and evaluate continuous-time discrete denoising on CIFAR-10 and monophonic music. The method reproduces support, improves image metrics with correctors, and outperforms the discrete-time baseline on categorical music data.

Synthetic distribution: On a synthetic 2D distribution, sufficiently small tau-leaping steps reproduce the full data support and match exact reverse-CTMC simulation fidelity.The experiment uses a 32-state space whose training histogram forms a ‘τ’ shape.
Image Modeling: For CIFAR-10, τLDR-0 has better Inception Score but worse FID than D3PM, while τLDR-10 beats D3PM on both metrics.τLDR-10 adds 10 corrector steps per predictor step near the end of reverse sampling.
Image Modeling: The continuous-time framework enables the corrector-rate derivation that substantially improves τLDR-10 sample quality.
Image Modeling: The optimal number of corrector steps depends on the sampling budget, with fewer correctors preferred under tighter budgets.More corrector steps require a larger tau to maintain a fixed number of network evaluations.
Monophonic Music: For monophonic music, the model conditions on two bars to complete 14 unseen bars and faithfully preserves the conditioning style.The task uses 256-step sequences with 128 notes or rest at each step.
Monophonic Music: On categorical music data, a uniform forward rate improves sample quality over birth/death rates, corrector steps improve it further, and D3PM performs worse.The comparison uses Hellinger distance and the proportion of outlier notes.

7 Discussion

The paper presents a continuous-time framework for discrete denoising models, using CTMCs, a continuous-time ELBO, and tau-leaping for reverse-process sampling. It concludes that predictor-corrector sampling improves sample quality over discrete-time methods, while noting evaluation and efficiency limitations.

Contributions: The framework formulates discrete denoising as a forward CTMC and identifies the reverse-time generative CTMC.Training uses a continuous-time ELBO to approximate the true reverse process.
Discussion: The paper reports improved sample quality from its predictor-corrector sampler compared with discrete-time methods for discrete data.The authors also note that the model requires many evaluations per sample and remains slightly behind continuous-state-space models on images.
CTMC background: A CTMC evolves through state jumps separated by randomly determined holding times, with transition behavior specified by a rate matrix.In the time-homogeneous case, holding times are exponential and jump destinations follow a jump-probability distribution.
Objective: The continuous-time ELBO is obtained by taking the small-step limit of the discrete-time ELBO after expressing transition kernels through CTMC rate matrices.The derivation uses expansions of transition probabilities and logarithms as the time step approaches zero.

B.4 Proof of Proposition 4

The proof combines the forward and reverse Kolmogorov equations to characterize the sum of their rate matrices.

Proof: The argument applies the Kolmogorov forward equation to both the forward process and its time-reversed process.Both processes share the same marginals, with the reverse process evolving in the opposite time direction.
Proof: Adding the forward and reverse Kolmogorov equations shows that R_t + ˆR_t is the rate matrix of a CTMC with invariant distribution q_t.The result follows by comparing the summed marginal equations with the Kolmogorov equation.
Proof: The resulting rate-matrix identity supplies the structural relation needed for the proposition.The proof concludes by identifying the summed generator through comparison with the marginal differential equation.

B.5 Proof of Theorem 1

The theorem bounds the total variation error of tau-leaping for a high-dimensional reverse CTMC by decomposing mixing and per-interval approximation errors. The bound is dimension-aware and depends on assumptions controlling local rates and their variation.

High-dimensional process: The bound applies to a factorized high-dimensional CTMC whose stationary distribution is the product of the single-dimension invariant distributions.The proof also relates full-chain mixing to the single-dimension mixing time.
Assumptions: The theorem assumes positive data probabilities and dimension-independent local bounds on forward and reverse rate behavior.These stronger assumptions make the dependence of the error bound on dimension explicit.
Error decomposition: The proof decomposes the final total variation error into the error from reverse-process approximation and the error from finite-time mixing.The decomposition reduces the analysis to bounding error accumulated over individual tau-leaping intervals.
Interval approximation: The tau-leaping approximation is analyzed by comparing each constant-rate interval with the exact reverse process.Coupling bounds control the discrepancy between the two Markov chains, including the effects of jumps within an interval.
Conclusion: The theorem is established by summing interval-wise bounds and combining them with the mixing estimate.The resulting expression gives the total variation bound for the law of the tau-leaping output.

C.4 One Forward Pass

The paper introduces a more efficient LCT evaluation that uses one denoising-network forward pass instead of two. The approximation exploits the similarity between neighboring states and provides a small empirical performance boost on monophonic music.

Motivation: Naively evaluating LCT requires two denoising-network forward passes, creating unnecessary computation.The redundancy arises because the two input states differ in only one dimension after a single forward transition.
One Forward Pass: The efficient objective LeCT replaces one sample distribution with a nearby distribution formed by applying a single forward transition.The approximation uses the fact that the resulting distributions are very similar.
Results: Table 3 compares LeCT and LCT metrics on the monophonic music test set using mean±std over five samples per test song.The caption specifies the dataset, comparison, evaluation split, and aggregation procedure.
Approximation: The nearby-state distribution approximates q_{t+δt}, where δt is the average time to the next transition.The paper reports δt values between 2×10^-6T and 2×10^-8T for image modeling and 1×10^-3T for monophonic music.
Results: Training with LeCT gives a small performance boost on monophonic music while improving computational efficiency.The authors hypothesize that sharing the perturbed state reduces objective variance through increased negative correlation.

D Direct Denoising Model Supervision

The section adds direct denoising supervision to the continuous-time objective and analyzes its relationship to the ELBO. This supervision targets the clean-data posterior and can be combined with the continuous-time ELBO for improved performance.

Direct denoising supervision predicts the clean datapoint x0 from corrupted input x using a negative log-likelihood loss.
Proposition 8 states that the true denoising distribution q0|t minimizes Lll.
Minimizing Lll is equivalent to minimizing the KL divergence between q0|t and pθ 0|t.
If pθ 0|t = q0|t, the model obtains the true reverse rate, and combining Lll with LDT can perform best.
Minimizing Lll alone minimizes a looser upper bound on the negative model log-likelihood than directly using the negative ELBO.
The continuous-time rate matrix is designed to provide analytically tractable qt|0 marginals and efficient training.

F.2 Predictor-Corrector Discussion

Predictor-corrector sampling for discrete states alternates CTMC simulations using learned reverse and corrector rates. Corrector steps can explore local modes by mixing information across dimensions.

The reverse process is implemented with tau-leaping, sampling Poisson jump counts for dimensions and rejecting unsuitable categorical proposals.
In discrete state spaces, predictor steps simulate the learned reverse CTMC, while corrector steps use the reverse rate plus the forward rate.
Corrector sampling corresponds to a slightly noisier categorical transition distribution than predictor sampling.
Adding corrector steps moves sample marginals closer to qt(x) and can further explore its domain.
Discrete denoising models output categorical distributions over every dimension, allowing uncertainty in x0 predictions and cross-dimensional mixing during correction.
Figure 7 applies 1000 corrector steps at t = 0.4 and shows exploration of potential local modes of qt(x).

H Experimental Details

The experiments use CTMC-specific training and sampling choices across synthetic, image, and music settings. Practical details include near-zero-time stabilization, final denoising, and task-specific corrector schedules.

Training samples a forward-process time t uniformly, but times near zero can cause ill-conditioned reverse rates and training instabilities.
Reverse sampling uses tau-leaping from T to ϵ, followed by argmax denoising to remove residual noise.
The time horizon is set to T = 1 because rescaling the time axis preserves the forward noising process.
The synthetic experiment uses 1M samples from a 32 × 32 state space derived from a grayscale τ image and a Gaussian rate process.
CIFAR10 experiments use 50000 training images and 10000 test images with a Gaussian rate whose stationary distribution is effectively uniform.
For CIFAR10, the denoising network is a modified U-net whose output represents probabilities over 256 states per image value.
Image sampling uses τ = 0.001 for τLDR-0 and τ = 0.002 for τLDR-10, with 10 corrector steps introduced after t < 0.1T.
For the corrector CTMC, a corrector rate τ set to 1.5 times the reverse-process τ performs best in the reported example.

I Ethical Considerations

The paper reports theoretical and generative advances while acknowledging risks from applying these models to sensitive real-world data. Training uses less sensitive images and music, but broader applications may reproduce harmful biases.

The work claims increased theoretical understanding of denoising generative models and improved generation capabilities on some discrete datasets.
Figures 13 and 14 depict conditional samples and reverse-sampling dimension trajectories, respectively.
Generative models can spread misinformation by reducing the resources needed to create realistic fake content.
Treating generated samples as objective truth without considering training-data biases can perpetuate discrimination against minority groups.
Although the experiments use less sensitive pictures of objects and music, the methods could model biased images of people or internet text.