A Continuous Time Framework for Discrete Denoising Models
Abstract
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.
Introduction
Diffusion/score-based/denoising models [1, 2, 3, 4] are a popular class of generative models that achieve state-of-the-art sample quality with good coverage of the data distribution [5] all whilst using a stable, non-adversarial, simple to implement training objective.
The general framework is to define a forward noising process that takes in data and gradually corrupts it until the data distribution is transformed into a simple distribution that is easy to sample.
The model then learns to reverse this process by learning the logarithmic gradient of the noised marginal distributions known as the score.
Most previous work on denoising models operates on a continuous state space.
However, there are many problems for which the data we would like to model is discrete.
This occurs, for example, in text, segmentation maps, categorical features, discrete latent spaces and the direct 8-bit representation of images.
Previous work has tried to realize the benefits of the denoising framework on discrete data problems, with promising initial results [6, 7, 8, 9, 10, 11, 12, 13].
All of these previous approaches train and sample the model in discrete time .
Unfortunately, working in discrete time has notable drawbacks.
It generally forces the user to pick a partition of the process at training time and the model only learns to denoise at these fixed time points.
Due to the fixed partition, we are then limited to a simple ancestral sampling strategy.
In continuous time, the model instead learns to denoise for any arbitrary time point in the process.
This complete specification of the reverse process enables much greater flexibility in defining the reverse sampling scheme.
For example, in continuous state spaces, continuous time samplers that greatly reduce the sampling time have been devised [14, 15, 16, 17] as well as ones that improve sample quality [4, 18].
The continuous time interpretation has also enabled the derivation of interesting theoretical properties such as error bounds [19] in continuous state spaces.
{A,A,A,B,B,B,C,C,C} {D,E,W,W,F,X,K,A,N} {D,P,W,W,F,X,K,A,N}{D,E,W,W,F,X,K,A,N}{D,E,W,W,F,X,K,A,N} {D,E,W,W,F,X,K,Y,N}{D,E,W,W,F,X,K,Y,N} {A,A,A,B,B,B,C,C,C}{A,A,A,M,B,B,C,C,C}{A,A,A,M,B,B,C,C,C}{A,A,A,M,B,B,X,C,C} t= t=T =state change can occur at any t ∈ ℝ t= t=T
Noising process
Generative process
Noise
Discrete Data
Denoising Network { }
[Figure]
[Figure]
To allow these benefits to be exploited for discrete state spaces as well, we formulate a continuous time framework for discrete denoising models.
Specifically, our contributions are as follows.
We formulate the forward noising process as a Continuous Time Markov Chain (CTMC) and identify the generative CTMC that is the time-reversal of this process.
We then bound the log likelihood of the generated data distribution, giving a continuous time equivalent of the ELBO that can be used for efficient training of a parametric approximation to the true generative reverse process.
To efficiently simulate the parametric reverse process, we leverage tau-leaping [20] and propose a novel predictor-corrector type scheme that can be used to improve simulation accuracy.
The continuous time framework allows us to derive a bound on the error between the true data distribution and the samples generated from the approximate reverse process simulated with tau-leaping.
Finally, we demonstrate our proposed method on the generative modeling of images from the CIFARdataset and monophonic music sequences.
Notably, we find our tau-leaping with predictor-corrector sampler can provide higher quality CIFAR10 samples than previous discrete time discrete state approaches, further closing the performance gap between when images are modeled as discrete data or as continuous data.
Proofs for all propositions and theorems are given in the Appendix.
Background on Discrete Denoising Models
In the discrete time, discrete state space case, we aim to model discrete data x 0 ∈ X with finite cardinality S = |X| .
We assume x 0 ∼ p data ( x 0 ) for some discrete data distribution p data ( x 0 ) .
We define a forward noising process that transforms p data ( x 0 ) to some distribution q K ( x K ) that closely approximates an easy to sample distribution p ref ( x K ) .
This is done by defining forward kernels q k +1 | k
( x k +1 | x k ) that all admit p ref as a stationary distribution and mix reasonably quickly.
For example, one can use a simple uniform kernel [6, 8], q k +1 | k
( x k +1 | x k )
=
δ x k +1 ,x k (1 − β )
+ (1 − δ x k +1 ,x k ) β/
( S − 1) where δ is a Kronecker delta.
The corresponding p ref is the uniform distribution over all states.
Other choices include: an absorbing state kernel—where for each state there is a small probability that it transitions to some absorbing state—or a discretized Gaussian kernel—where only transitions to nearby states have significant probability (valid for spaces with ordinal structure) [8].
After defining q k +1 | k , we have a forward joint decomposition as follows
[Equation]
The joint distribution q 0: K ( x 0: K ) also admits a reverse decomposition: q 0: K ( x 0: K ) =
q K ( x K ) (cid:81) K − 1 k =
q k | k +1 ( x k | x k +1 ) where q k | k +1 ( x k | x k +1 )
=
q k +1 | k
( x k +1 | x k ) q k
( x k ) q k +
x k +1 ) .
Here q k ( x k ) denotes the marginal of q 0: K ( x 0: K ) at time k .
If one had access to q k | k +1 and could sample q K ( x K ) exactly, then samples from p data ( x 0 ) could be produced by first sampling q K ( x K ) and then ancestrally sampling the q k | k +1 reverse kernels.
However, in practice, q k | k +1 is intractable and needs to be approximated with a parametric reverse kernel, p θk |
k +1 .
This kernel is commonly defined through the analytic q k | k +1 , 0 distribution and a parametric ‘denoising’ model p θ 0 | k +1 [6, 8],
[Equation]
Though q K ( x K ) is also intractable, for large K we can reliably approximate it with p ref ( x K ) .
Note that the faster the transitions mix, the more accurate this approximation becomes.
Approximate samples from p data ( x 0 ) can then be obtained by sampling the generative joint distribution
[Equation]
where θ is trained through minimizing the negative discrete time (DT) ELBO which is an upper bound on the negative model log-likelihood
[Equation]
It was shown in [1] that L DT can be re-written as
[Equation]
where KL is the Kullback-Leibler divergence.
The forward kernels q k +1 | k are chosen such that q k | 0 ( x k | x 0 ) can be computed efficiently in a time independent of k .
With this, θ can be efficiently trained by taking a random selection of terms from L DT in each minibatch and performing a stochastic gradient step.
Continuous Time Framework 3.1 Forward process and its time reversal
Our method is built upon a continuous time process from t = 0 to t = T .
State transitions can occur at any time during this process as opposed to the discrete time case where transitions only occur when one of the finite number of transition kernels is applied (see Figure 1).
This process is known as a
Continuous Time Markov Chain (CTMC), we provide a short overview of CTMCs in Appendix A for completeness.
Giving an intuitive introduction here, we can define a CTMC through an initial distribution q 0 and a transition rate matrix R t ∈ R S × S .
If the current state is ˜ x , then the transition rate matrix entry R t (˜ x, x ) is the instantaneous rate (occurrences per unit time) at which state ˜ x transitions to state x .
Loosely speaking, the next state in the process will likely be one for which R t (˜ x, x ) is high, and furthermore, the higher the rate is, the less time it will take for this transition to occur.
It turns out that the transition rate, R t , also defines the infinitesimal transition probability for the process between the two time points t − ∆ t and t
[Equation]
where o (∆ t ) represents terms that tend to zero at a faster rate than ∆ t .
Comparing to the discrete time case, we see that R t assumes an analogous role to the discrete time forward kernel q k +1 | k in how we define the forward process.
Therefore, just as in discrete time, we design R t such that: i) the forward process mixes quickly towards an easy to sample (stationary) distribution, p ref , (e.g. uniform), ii) we can analytically obtain q t |
( x t | x 0 ) distributions to enable efficient training (see Section 4.1 for how this is done).
We initialize the forward CTMC at q 0 ( x 0 ) =
p data ( x 0 ) at time t = 0 .
We denote the marginal at time t =
T as q T ( x T ) , which should be close to p ref ( x T ) .
We now consider the time reversal of the forward process, which will take us from the marginal q T ( x T ) back to the data distribution p data ( x 0 ) through a reverse transition rate matrix, ˆ R t ∈ R S × S :
[Equation]
In discrete time, one uses Bayes rule to go from q k +1 | k to q k | k +1 .
We can use similar ideas to calculate ˆ R t from R t as per the following result.
Proposition 1.
For a forward in time CTMC, { x t } t ∈ [0 ,T ] , with rate matrix R t , initial distribution p data ( x 0 ) and terminal distribution q T ( x T ) , there exists a CTMC with initial distribution q T ( x T ) at t = T , terminal distribution p data ( x 0 ) at t = 0 and transition rate matrix ˆ R t that runs backwards in time and is almost everywhere equivalent to the time reversal of the forward CTMC, { x t } t ∈ [ T, 0] .
Furthermore, ˆ R t is related to R t by the following expression
[Equation]
where q t | 0 ( x | x 0 ) are the conditional marginals of the forward process and q
| t ( x 0 | x ) =
q t | 0 ( x | x 0 ) p data ( x 0 ) /q
t ( x ) with q t ( x ) being the marginal of the forward process at time t .
When x = ˜ x , ˆ R t ( x, x ) = − (cid:80) x ′ ̸
= x ˆ R t ( x, x ′ ) because the rows must sum to zero (see Appendix A).
Unfortunately, ˆ R t is intractable due to the intractability of q t ( x ) and thus of q
| t ( x 0 | x ) .
Therefore, we consider an approximation ˆ R θt of ˆ R t by approximating q 0 | t ( x 0 | x ) with a parametric denoising model, p θ
| t ( x 0 | x ) :
[Equation]
ˆ R θt ( x, x )
= − (cid:80) x
′ ̸
= x ˆ R θt ( x, x ′ ) as before.
As a further analogy to the discrete time case, notice that when x ̸ =
˜
x , ˆ R θt has the same form as the discrete time parametric reverse kernel, p θk | k +1 defined in eq (1) but with the forward kernel, q k +1 | k , replaced by the forward rate, R t .
3.2 Continuous Time ELBO
In discrete time, θ is trained by minimizing the discrete time negative ELBO, L DT , formed from the forward and reverse processes.
We mirror this approach in continuous time by minimizing the corresponding continuous time (CT) negative ELBO, L CT , as derived below.
Proposition 2.
For the reverse in time CTMC with initial distribution p ref ( x T ) , terminal distribution p θ 0 ( x 0 ) , and reverse rate ˆ R θt , an upper bound on the negative model log-likelihood, E p data ( x 0 ) [ − log p θ 0 ( x 0 )] , is given by
[Equation]
where C is a constant independent of θ and
Z
[Equation]
r t (˜ x | x ) gives the probability of transitioning from x to ˜ x , given that we know a transition occurs at time t .
We can optimize this objective efficiently with stochastic gradient descent.
For a gradient update, we sample a batch of datapoints from p data ( x 0 ) , noise each datapoint using a random time, t ∼ U (0 , T ) , x ∼ q t | 0 ( x | x 0 ) and finally sample an auxiliary ˜ x from r t (˜ x | x ) for each x .
Intuitively, ( x , ˜ x ) are a pair of states following the forward in time noising process.
Minimizing the second term in L CT maximizes the reverse rate for this pair, but going in the backwards direction, ˜ x to x .
This is how ˆ R θt learns to reverse the noising process.
Intuition on the first term and a direct comparison to
L DT is given in Appendix C.1.
The first argument of ˆ R θt is input into p θ 0 | t
so we naively require two network forward passes on x and ˜ x to evaluate the objective.
We can avoid this by approximating the q t ( x ) sample in the first term with ˜ x
meaning we need only evaluate the network once on ˜ x .
The approximation is valid because, as we show in Appendix C.4, ˜ x is approximately distributed according to q t + δt for δt very small.
Efficient Forward and Backward Sampling 4.1 Choice of Forward Process
The transition rate matrix R t needs to be chosen such that the forward process: i) mixes quickly towards p ref , and ii) the q t | 0 ( x | x 0 ) distributions can be analytically obtained.
The Kolmogorov differential equation for the CTMC needs to be integrated to obtain q t | 0 ( x | x 0 ) .
This can be done analytically when R t and R t ′ commute for all t, t ′ , see
Appendix E. An easy way to meet this condition is to let R t = β ( t ) R b where R b ∈ R S × S is a user-specified time independent base rate matrix and β ( t ) ∈
R is a time dependent scalar.
We then obtain the analytic expression
[Equation]
where R b = Q Λ Q − 1 is the eigendecomposition of matrix R b and exp [ · ] the element-wise exponential.
Our choice of β schedule is guided by [3, 4], β ( t ) =
ab t log( b ) .
The hyperparameters a and b are selected such that q T ( x T )
≈ p ref ( x T ) at the terminal time t = T while having a steady speed of ‘information corruption’ which ensures that ˆ R t does not vary quickly in a short span of time.
We experiment with a variety of R b matrices, for example, a uniform rate, R b = 11 T − S Id , where 11 T is a matrix of ones and Id is the identity.
For problems with a heavy spatial bias, e.g. images, we can instead use a forward rate that only encourages transitions to nearby states; details and the links to the corresponding discrete time processes can be found in Appendix E.
4.2 Factorizing Over Dimensions
Our aim is to model data that is D dimensional, with each dimension taking one value from S possibilities.
We now slightly redefine notation and say x 1: D ∈ X D , |X| = S .
In this setting, calculating transition probabilities naively would require calculating S D rate values corresponding to each of the possible next states.
This is intractable for any reasonably sized S and D .
We avoid this problem simply by factorizing the forward process such that each dimension propagates independently.
Since this is a continuous time process and each dimension’s forward process is independent of the others, the probability two or more dimensions transition at exactly the same time is zero.
Therefore, overall in the full dimensional forward CTMC, each transition only ever involves a change in exactly one dimension.
For the time reversal CTMC, it will also be true that exactly one dimension changes in each transition.
This makes computation tractable because of the S D rate values, only D × ( S − 1) + 1 are non-zero - those corresponding to transitions where exactly one dimension changes plus the no change transition.
Finally, we note that even though dimensions propagate independently in the forward direction, they are not independent in the reverse direction because the starting points for each dimension’s forward process are not independent for non factorized p data .
The following proposition shows the exact forms for the forward and reverse rates in this case.
Proposition 3.
If the forward process factorizes as q t | s ( x 1: D t
[Equation]
then the forward and reverse rates are of the form
[Equation]
where R dt ∈ R S × S and δ x 1: D \ d , ˜ x 1: D \ d is 1 when all dimensions except for d are equal.
To find ˆ R θ 1: D t we simply replace q 0 | t ( x d 0 | x 1: D ) with p θ 0 | t ( x d 0 | x 1: D ) which is easily modeled with a neural network that outputs conditionally independent state probabilities in each dimension.
In
Appendix C.3 we derive the form of L CT when we use this factorized form for R 1: D t and ˆ R θ 1: D t .
4.3 Simulating the Generative Reverse Process with Tau-Leaping
The parametric generative reverse process is a CTMC with rate matrix ˆ R θ 1: D t .
Simulating this process from distribution p ref ( x 1: D T ) at time t
= T back to t = 0 will produce approximate samples from p data ( x 1: D 0 ) .
The process could be simulated exactly using Gillespie’s Algorithm [21, 22, 23] which alternates between i) sampling a holding time to remain in the current state and ii) sampling a new state according to the current rate matrix, ˆ R θ 1: D t (see Appendix F).
This is inefficient for large D because we would need to step through each transition individually and so only one dimension would change for each simulation step.
Instead, we use tau-leaping [20, 23], a very popular approximate simulation method developed in chemical physics.
Rather than step back through time one transition to the next, tau-leaping leaps from t to t − τ and applies all transitions that occurred in [ t, t − τ ] simultaneously.
To make a leap, we assume ˆ R θ 1: D t and x 1: D t remain constant in [ t, t − τ ] .
As we propagate from t to t − τ , we count all of the transitions that occur, but hold off on actually applying them until we reach t − τ , such that x 1: D t remains constant in [ t, t − τ ] .
Assuming ˆ R θ 1: D t and x 1: D t remain constant, the number of times a transition from x 1: D t to ˜ x 1: D occurs in [ t, t − τ ] is Poisson distributed with mean τ ˆ R θ 1: D t ( x 1: D t , ˜ x 1: D ) .
Once we reach t − τ , we apply all transitions that occurred simultaneously i.e. x 1: D
[Equation]
=
[Equation]
) where P i is a Poisson random variable with mean τ ˆ R θ 1: D t ( x 1: D t , ˜ x 1: D i ) .
Note the sum assumes a mapping from X to Z .
[Figure]
Here,
, |X| = 5 , P 12 = 1 , P
2 , all
P
= 0 .
Using our knowledge of ˆ R θ 1: D t , we can further unpack this update.
Namely, ˆ
[Equation]
D t ( x 1: D t , ˜ x 1: D ) can only be non-zero when ˜ x 1: D has a different value to x 1: D t in exactly one dimension (rates for multidimensional changes are zero).
Explicitly summing over these options we get x 1: D t − τ = x 1: D t +
[Equation]
=
[Equation]
\ x
[Equation]
P ds
[Equation]
s − x dt ) e d where e d is a one-hot vector with a 1 at dimension d and P ds is a Poisson random variable with mean τ ˆ R θ 1: D t ( x 1: D t , x 1: D t +( s − x dt ) e d ) .
Since multiple
P ds can be non-zero, we see that tau-leaping allows x 1: D t to change in multiple dimensions in a single step.
Figure 2 visualizes this idea.
During the [ t, t − τ ] interval, one jump occurs in dimension 1 and two jumps occur in dimension 2 .
These are all applied simultaneously once we reach t − τ .
When our discrete data has ordinal structure (e.g.
Section 6.2) our mapping to Z is not arbitrary and making multiple jumps within the same dimension ( (cid:80) Ss =1 \ x dt P ds > 1 ) is meaningful.
In the non-ordinal/categorical case (e.g. Section 6.3) the mapping to Z is arbitrary and so, although taking simultaneous jumps in different dimensions is meaningful, taking multiple jumps within the same dimension is not.
For this type of data, we reject changes to x dt for any d for which (cid:80) Ss =1 \ x dt P ds > 1 .
In practice, the rejection rate is very small when R 1: D t is suitable for categorical data (e.g. uniform), see Appendix H.3.
In Section 4.5, our error bound accounts for this low probability of rejection and also the low probability of an out of bounds jump that we observe in practice in the ordinal case.
The tau-leaping approximation improves with smaller τ , recovering exact simulation in the limit as τ → 0 .
Exact simulation is similar to an autoregressive model in that only one dimension changes per step.
Increasing τ and thus the average number of dimensions changing per step gives us a natural way to modulate the ‘autoregressiveness’ of the model and trade sample quality with compute (Figure 4 right).
We refer to our method of using tau-leaping to simulate the reverse CTMC as τ LDR (tau-leaping denoising reversal) which we formalize in Algorithm 1 in Appendix F.
We note that theoretically, one could approximate ˆ
R
θ 1: D t as constant in the interval [ t, t − τ ] , and construct a transition probability matrix by solving the forward Kolmogorov equation with the matrix exponential P t − τ | t ≈ exp ( τ ˆ R θ 1: D t ) .
However, for the learned ˆ R θ 1: D t ∈ R S D × S D matrix, it is intractable to compute this matrix exponential so we use tau-leaping for sampling instead.
4.4 Predictor-Corrector
During approximate reverse sampling, we aim for the marginal distribution of samples at time t to be close to q t ( x t ) (the marginal at time t of the true CTMC).
The continuous time framework allows us to exploit additional information to more accurately follow the reverse progression of marginals, { q t ( x t ) } t ∈ [ T, 0] and improve sample quality.
Namely, after a tau-leaping ‘predictor’ step using rate ˆ R θ t , we can apply ‘corrector’ steps with rate R c t which has q t ( x t ) as its stationary distribution.
The corrector steps bring the distribution of samples at time t closer to the desired q t ( x t ) marginal.
R ct is easy to calculate as stated below
Proposition 4.
For a forward CTMC with marginals { q t ( x t ) } t ∈ [0 ,T ] , forward rate, R t , and corresponding reverse CTMC with rate ˆ R t , the rate R ct =
R t + ˆ R t has q t ( x t ) as its stationary distribution.
In practice, we approximate R ct by replacing ˆ R t with ˆ R θt .
This is directly analogous to PredictorCorrector sampling in continuous state spaces [4] that predict by integrating the reverse SDE and correct with score-based Markov chain Monte Carlo steps, see Appendix F.2 for further discussion.
4.5 Error Bound
Our continuous time framework also allows us to provide a novel theoretical bound on the error between the true data distribution and the sample distribution generated via tau-leaping (without predictor-corrector steps), in terms of the error in our approximation of the reverse rate and the mixing of the forward noising process.
We assume we have a time-homogeneous rate matrix R t on X , from which we construct the factorized rate matrix R 1: D t on X D by setting R dt = R t for each d .
Note that by rescaling time by a factor of β ( t ) we can transform our choice of rate from Section 4.1 to be time-homogeneous.
We will denote | R | = sup t ∈ [0 ,T ] ,x ∈X | R t ( x, x ) | , and let t mix be the (1/4)-mixing time of the CTMC with rate R t (see [24, Chapter 4.5]).
Theorem 1.
For any D ≥ 1 and distribution p data on X D , let { x t } t ∈ [0 ,T ] be a CTMC starting in p data with rate matrix R 1: D t as above.
Suppose that ˆ R θ 1: D t is an approximation to the reverse rate matrix and let ( y k ) k =0 , 1 ,...,N be a tau-leaping approximation to the reverse dynamics with maximum step size τ .
Suppose further that there is some constant M > 0 independent of D such that
[Equation]
for all t ∈ [0 , T ] .
Then under the assumptions in Appendix B.5, there are constants C 1 , C 2 >depending on X and R t but not D such that, if L ( y 0 ) denotes the law of y 0 , we have the total variation bound
[Equation]
The first term of the above bound captures the error introduced by our approximation of the reverse rate ˆ R 1: D t with ˆ R θ 1: D t .
The second term reflects the error introduced by the tau-leaping approximation, and is linear in both T and τ , showing that as we take our tau-leaping steps to be arbitrarily small, the error introduced by tau-leaping goes to zero.
The final term describes the mixing of the forward chain, and captures the error introduced since p ref and q T are not exactly equal.
We choose to make the dependence of the bound on the dimension D explicit, since we are specifically interested in applying tau-leaping to high dimensional problems where we make transitions in different dimensions simultaneously in a single time step.
The bound grows at worst quadratically in the dimension, versus e.g. exponentially.
The bound is therefore useful in showing us that we do not need to make τ impractically small in high dimensions.
Other than gaining these intuitions, we do not expect the bound to be particularly tight in practice and further it would not be practical to compute because of the difficulty in finding M , C 1 and C 2 .
The assumptions listed in Appendix B.5 hold approximately for tau-leaping in practice when we use spatially biased rates for ordinal data such that jump sizes are small or uniform rates for non-ordinal data such that the dimensional rejection rate is small.
These assumptions could be weakened, however,
Theorem 1 would become much more involved, obscuring the intuition and structure of the problem.
Related Work
The application of denoising models to discrete data was first described in [1] using a binomial diffusion process for a binary dataset.
Each reverse kernel p θk | k +1 was directly parameterized without using a denoising model p θ 0 | k .
In [25] an approach for discrete categorical data was suggested using a uniform forward noising kernel, q k +1 | k , and a reverse kernel parameterized through a denoising model, though no experiments were performed with the approach.
Experiments on text and segmentation maps were then performed with a similar model in [6].
Other forward kernels were introduced in [8] that are more appropriate for certain data types such as the spatially biased Gaussian kernel.
[9, 13] apply the approach to discrete latent space modeling using uniform and absorbing state forward
[Figure]
kernels.
Whilst a link to continuous time for the forward process is mentioned in [8], all of these approaches train and sample in discrete time.
We show in Appendix G that this involves making an implicit approximation for multi-dimensional data.
We extend this line of work by training and sampling in continuous time.
Other works also operate in discrete space but less rigidly follow the diffusion framework.
A corruption process tailored to text is proposed in [12], whereby token deletion and insertion is also incorporated.
[26] also focus on text, creating a generative reverse chain that repeatedly applies the same denoising kernel.
The corruption distribution is also defined through the same denoising kernel to reduce distribution shift between training and sampling.
In [7], a more standard masking based forward process is used but the reversal is interpreted from an order agnostic autoregressive perspective.
They also describe how their model can be interpreted as the reversal of a continuous time absorbing state diffusion but do not utilize this perspective in training or sampling.
[27] propose a denoising type framework that can be used on binary data where the forward and reverse process share the same transition kernel.
Finally, in [11], the discrete latent space of a VQVAE is modeled by quantizing an underlying continuous state space diffusion with probabilistic quantization functions.
Experiments 6.1 Demonstrative Example
We first verify the method can accurately produce samples from the entire support of the data distribution and that tau-leaping can accurately simulate the reverse CTMC.
To do this, we create a dataset formed of 2d samples of a state space of 32 arranged such that the histogram of the training dataset forms a ‘ τ ’ shape.
We train a denoising model using the L CT objective with p θ 0 | t parameterized through a residual MLP (full details in Appendix H.1).
We then sample the parameterized reverse process using an exact method (up to needing to numerically integrate the reverse rate) and tauleaping.
Figure 3 top-right shows the marginals during reverse simulation with τ = 0 .
and we indeed produce samples from the entire support of p data .
Furthermore, we find that with sufficiently small τ , we can match the fidelity of exact simulation of the reverse CTMC (Figure 3 left).
The value of τ dictates the number of network evaluations in the reverse process according to NFE = T/τ .
In all experiments we use T = 1 .
Exact simulation results in a non zero Hellinger distance between the generated and training distributions because of imperfections in the learned ˆ R θt model.
6.2 Image Modeling
We now demonstrate that our continuous time framework gives us improved generative modeling performance versus operating in discrete time.
We show this on the CIFAR-10 image dataset.
Images are typically stored as discrete data, each pixel channel taking one value from 256 possibilities.
Continuous state space methods have to somehow get around this fact by, for example, adding a discretization function at the end of the generative process [3] or adding uniform noise to the data.
[Table]
[Figure]
[Table]
[Figure]
Here, we model the images directly in discrete space.
We parameterize p θ 0 | t using the standard
U-net architecture [3] with the modifications for discrete state space suggested by [8].
We use a spatially biased rate matrix and train with an augmented L CT loss including direct p θ 0 | t supervision, full experimental details are in Appendix H.2.
Figure 4 left shows randomly generated unconditional CIFAR10 samples from the model and we report sample quality metrics in Table 1.
We see that our method ( τ LDR- 0 ) with 0 corrector steps has better Inception Score but worse FID than the D3PM discrete time method.
However, our τ LDRmethod with 10 corrector steps per predictor step at the end of the reverse sampling process ( t < 0 .
1 T ) greatly improves sample quality, beating the discrete time method in both metrics and further closes the performance gap with methods modeling images as continuous data.
The derivation of the corrector rate which gave us this improved performance required our continuous time framework.
D3PM achieves the highest ELBO but we note that this does not correlate well with sample quality.
In Table 1, τ was adjusted such that both τ LDR- 0 and τ LDR- 10 used 1000 p θ 0 | t evaluations in the reverse sampling procedure.
We show how FID score varies with number of p θ 0 | t evaluations for τ LDR- { 0 , 3 , 10 } in Figure 4 right.
The optimum number of corrector steps depends on the sampling budget, with lower numbers of corrector steps being optimal for tighter budgets.
This is due to the increased τ required to maintain a fixed budget when we use a larger number of corrector steps.
6.3 Monophonic Music
In this experiment, we demonstrate our continuous time model improves generation quality on non-ordinal/categorical discrete data.
We model songs from the Lakh pianoroll dataset [29, 30].
We select all monophonic sequences from the dataset such that at each of the 256 time steps either one from 128 notes is played or it is a rest.
Therefore, our data has state space size S = 129 and dimension D = 256 .
We scramble the ordering of the state space when mapping to Z to destroy any ordinal structure.
We parameterize p θ 0 | t with a transformer architecture [31] and train using a conditional form of L CT targeting the conditional distribution of the final 14 bars (224 time steps) given the first 2 bars of the song.
We use a uniform forward rate matrix, R t , full experimental details
[Table]
[Figure]
[Table]
[Figure]
We quantify sample quality in Table 2.
We use two metrics: the Hellinger distance between the histograms of generated and ground truth notes and the proportion of outlier notes in the generations but not in the ground truth.
Using our method, we compare between a birth/death and uniform forward rate matrix R t .
The birth/death rate is only non-zero for adjacent states whereas the uniform rate allows transitions between arbitrary states which is more appropriate for the categorical case thus giving improved sample quality.
Adding 2 corrector steps per predictor step further improves sample quality.
We also compare to the discrete time method D3PM [8] with its most suitable corruption process for categorical data.
We find it performs worse than our continuous time method.
Discussion
We have presented a continuous time framework for discrete denoising models.
We showed how to efficiently sample the generative process with tau-leaping and provided a bound on the error of the generated samples.
On discrete data problems, we found our predictor-corrector sampler improved sample quality versus discrete time methods.
Regarding limitations, our model requires many model evaluations to produce a sample.
Our work has opened the door to applying the work improving sampling speed on continuous data [14, 15, 16, 17, 32] to discrete data problems too.
Modeling performance on images is also slightly behind continuous state space models, we hope this gap is further closed with bespoke discrete state architectures and corruption process tuning.
Finally, we note that the ELBO values for the discrete time model on CIFAR10 are better than for our method.
In this work, we focused on sample quality rather than using our model to give data likelihoods e.g. for compression downstream tasks.
Acknowledgements
Andrew Campbell and Joe Benton acknowledge support from the EPSRC CDT in Modern Statistics and Statistical Machine Learning (EP/S023151/1).
Arnaud Doucet is partly supported by the EPSRC grant EP/R034710/1.
He also acknowledges support of the UK Defence Science and Technology
Laboratory (DSTL) and EPSRC under grant EP/R013616/1.
This is part of the collaboration between
US DOD, UK MOD and UK EPSRC under the Multidisciplinary University Research Initiative.
This project made use of time on Tier 2 HPC facility JADE2, funded by EPSRC (EP/T022205/1).
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] We detail the continuous time framework in Section 3, show our theoretical result in Section 4.5 and demonstrate improved performance in Section(b)
Did you describe the limitations of your work?
[Yes] See Section(c)
Did you discuss any potential negative societal impacts of your work?
[Yes] See Appendix I (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] See Appendix B. (b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix B 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] The code is available at https://github.com/andrew-cr/tauLDR (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] See Appendix H (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] See Table(d)
Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix H 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] We cite the authors of the models and datasets we use in Section(b)
Did you mention the license of the assets?
[Yes] See Appendix H (c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] We include our curated data in the supplement (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
Appendix
Contents
A
Primer on Continuous Time Markov Chains
B Proofs
B.1 Proof of Proposition 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.2 Proof of Proposition 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.3 Proof of Proposition 3 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.4 Proof of Proposition 4 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.5 Proof of Theorem 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C Continuous Time ELBO Details
C.1 Comparison with the Discrete Time ELBO . . . . . . . . . . . . . . . . . . . . . .
C.2 Conditional Form . .
C.3 Continuous Time ELBO with Factorization Assumptions . . . . . . . . . . . . . .
C.4 One Forward Pass . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
D Direct Denoising Model Supervision
E Choice of Forward Process
F CTMC
Simulation
F.
Exact CTMC and Tau-Leaping . .
F.
Predictor-Corrector Discussion . .
G Implicit Dimensional Assumptions Made
in Discrete Time
H Experimental Details
H.1 Demonstrative Example
H.2 Image Modeling
H.3 Monophonic Music
I Ethical Considerations
The appendix is organized as follows.
In Section A, we provide a short introduction to Continuous
Time Markov Chains, including the relevant results we use in this work.
Proofs for all the Propositions and Theorems from the main text are in Section
B.
We then describe in Section C some additional intuitions and forms of our proposed objective, L CT .
In Section D, we describe how an additional direct denoising model supervision term can be added to the objective to improve empirical performance.
Details for how we define the forward process in our model can be found in Section
E. Section F describes in more detail how CTMCs can be simulated and includes the algorithmic description of tau-leaping.
We argue in Section G that operating in discrete time forces an implicit assumption when using a factorized forward process on multi-dimensional data.
Full experimental details for all investigations can be found in Section H as well as additional plots and results from our models.
Finally, in Section I, we consider the social impacts of our research.
[Equation]
t
[Equation]
x
A Primer on Continuous Time Markov Chains
A Continuous Time Markov Chain (CTMC) is a right continuous stochastic process { x t } t ∈ [0 ,T ] satisfying the Markov property, with x t taking values in a discrete state space X .
Since the CTMC is
Markov, future behaviour of the process depends only on the current state and not the history.
A schematic representation of a CTMC path is shown in Figure 6.
The process repeatedly transitions from one state to another after having waited in the previous state for a randomly determined amount of time.
A CTMC can be completely characterised by its jumps and holding times.
Specifically, the time between each jump or holding time is exponentially distributed with mean ν ( x ) where x is the state in which the process is holding.
The next state that is jumped to is drawn from a jump probability distribution r (˜ x | x ) .
The holding and jumping procedure is then repeated.
There is an equivalent definition involving the transition rate matrix, R ∈ R S × S , that we use in the main paper.
The transition rate matrix is defined as
[Equation]
[Equation]
where R (˜ x, x ) is the (˜ x, x ) element of the transition rate matrix and q t | t − ∆ t ( x | ˜ x ) is the infinitesimal transition probability of being in state x at time t given that the process was in state ˜ x
at time t − ∆ t .
Conversely, the CTMC can itself be defined through this infinitesimal transition probability
[Equation]
(4) where o (∆ t ) represents terms that tend to zero at a faster rate than ∆ t .
From this definition of the transition rate matrix, we can infer the following properties:
[Equation]
R (˜ x, x ) is the rate at which probability mass moves from state ˜ x to x .
R ( x, x ) is the total rate at which probability mass moves out of state x and is thus negative.
In the time-homogeneous case, R has simple relations to the jump and holding time definitions.
[Equation]
In the time-inhomogeneous case, our transition rate matrix will now depend on time, R t , and these simple relations to the jump and holding time definition do not hold.
However, R t will still follow equations (3), (4) and (5).
The CTMC transition probabilities satisfy the Kolmogorov forward and backward equations.
For t > s ,
[Equation]
The Kolmogorov forward equation also gives us a differential equation for the marginals of the
CTMC.
[Equation]
Exponential and Poisson Random Variables
In the time homogeneous case, holding times are exponentially distributed with mean ν ( x ) = − 1 /R ( x, x ) .
The tau-leaping algorithm relies on the fact that the number of events in interval [0 , t ] is Poisson distributed with mean 1 ν t when the inter-event times are exponentially distributed with mean ν .
B Proofs B.1 Proof of Proposition
Proof.
We recall that a process { x t } t ∈ [0 ,T ] taking values in X is called a CTMC if it is rightcontinuous and satisfies the Markov property.
Denote { y t } t ∈ [0 ,T ] = { x T − t } t ∈ [0 ,T ] except at the jump times of the forward process τ n with n ∈ N , where y T − τ n = x − τ n = lim t ≤ τ n ,t → τ n x t .
Hence, { y t } t ∈ [0 ,T ] is almost surely equal to { x T − t } t ∈ [0 ,T ] and is right-continuous.
Since the Markov property is symmetric, we get that { y t } t ∈ [0 ,T ] is a CTMC.
We now compute its transition matrix.
Let x, ˜ x ∈ X with x ̸ =
˜ x , using the Kolmogorov forward equation, we have
[Equation]
t ( y, ˜ x ) , where { p t | s , s, t ∈ [0 , T ] , t > s } is the transition probability system associated with { y t } t ∈ [0 ,T ] and { ˆ R T − t } t ∈ [0 ,T ] is the transition rate matrix associated with { y t } t ∈ [0 ,T ] .
Note that
[Equation]
where { q t | s , s, t ∈
[0 , T ] , t > s } is the transition probability system associated with { x t } t ∈ [0 ,T ] and { q t , t ∈ [0 , T ] } are the marginals of { x t } t ∈ [0 ,T ] .
Now, writing the backward Kolmogorov equation for { x t } t ∈ [0 ,T ]
[Equation]
Re-labeling the time indices we obtain,
[Equation]
Letting s → t and using that lim s → t q T − s | T − t ( x | ˜ x ) = 0 , we get that
[Equation]
Re-labeling the time-indices on the rate matrices, we obtain
[Equation]
Now we write the marginal ratio q t (˜ x )
q t ( x ) in a different form
[Equation]
Substituting in this form for the marginal ratio concludes the proof.
B.2 Proof of Proposition
In this section, we detail two proofs for Proposition 2.
The first is a formal proof using results from stochastic processes.
We then provide a second informal proof for the same result to gain intuition into the L CT objective that only relies on elementary results from CTMCs.
Proof 1 - Stochastic Processes
Proof.
Let us write Q for the path measure of the forward CTMC with rate matrix R t , ˆ Q for the path measure of its exact time reversal and P θ for the path measure of the approximate reverse process with rate matrix ˆ R θt .
Also, we use superscripts to notate conditioning on the starting point, for example Q x 0 denotes the path measure of the forward process conditioned to start in x 0 .
With this notation, we have
[Equation]
where P θ , ˆ Q run in the reverse time direction.
Writing ˆ W s for a reverse path and integrating wrt p data (d x 0 ) we have
[Equation]
where ˆ R θt ( x ) is shorthand for − ˆ R θt ( x, x ) .
When x 0 ∼ p data , x T ∼ q T | 0 ( ·| x 0 ) , ˆ W ∼ ˆ Q (d W | ˆ W 0 = x T , ˆ W T = x 0 ) , the reverse path is distributed according to p data (d x 0 ) Q x 0 (d W ) and therefore ( ˆ W s − , ˆ W s ) is distributed like ( W T − s , W
( T − s ) − ) and thus we have (cid:90)
[Equation]
Using Dynkin’s lemma and the fact that P θt ( x | y )
ˆ R θt ( y )
=
ˆ R θt ( y, x ) we can re-expresss this final line as
[Equation]
which rearranges to give the continuous time ELBO in the form of Proposition 2.
Proof 2 - Limit of Discrete Time ELBO
Proof.
Consider a partitioning of [0 , T ] ,
= t 0 < t 1 < · · · < t k − 1 < t k < t k +1 < · · · < t K − 1 < t K = T .
Let t k − t k − 1 = ∆ t for all k .
In subscripts we use k as a shorthand for t k when this does not cause confusion.
Considering a CTMC with this time partitioning converts the problem into a discrete time Markov Chain with forward transition kernel, q k +1 | k
( x k +1 | x k ) and parameterized reverse kernel, p θk |
k +1 ( x k | x k +1 ) .
Therefore, we can write the negative ELBO in its discrete time form, L DT
[Equation]
In the following
[Equation]
we will write the transition kernels in terms of the CTMC rate matrices and take the limit as ∆ t → 0 to obtain a continuous time negative ELBO
[Equation]
First, consider one item from the inner sum of L DT
[Equation]
where we have absorbed terms that do not depend on θ into C .
We now write p θk | k +
( x k | x k +1 ) in terms of ˆ R θk .
[Equation]
where on the last line we have used the series expansion for log(1 + z )
=
z − z 2 2 + o ( z 2 ) valid for | z | ≤ 1 , z ̸ = − 1 .
For any finite R θk ( x k , x k ) , ∆ t can be taken small enough such that the series expansion holds.
We now substitute this form for log p θ k | k +1 into L k and further write the expectation over q k +1 | k
( x k +1 | x k )
=
δ x k ,x k +
+ R k
( x k , x k +1 )
∆ t + o (∆ t ) as an explicit sum.
[Equation]
We can isolate ˆ R θk within the log through the following re-arrangement
[Equation]
where the first two terms are independent of θ and tend to 0 as ∆ t → 0 .
Note that we assume ˆ
R θk ( x k +1 , x k ) > 0 for x k +1 ̸ = x k pairs which have R k ( x k , x k +1 ) > 0 .
This assumption is valid because, for x k +1 ̸ = x k , we have
[Equation]
and we assume p θ 0 | k
( x 0 | x k +1 ) > 0 which is valid when we parameterize p θ 0 | k with a softmax output.
We assume an irreducible Markov chain, hence q k | 0 > 0 for t k > 0 .
With this re-arrangement, and absorbing constant terms into C , we obtain
[Equation]
The second term can be re-written so that it is more efficient to approximate with Monte Carlo.
Currently the denoising model p θ 0 | k has to be evaluated for each term in the sum (cid:80) x
k +1 ̸ = x k which would require multiple forward passes of the neural network.
We can instead create a new probability distribution to sample from as follows.
Define
[Equation]
where
[Equation]
So we now have
[Equation]
Examining the other terms in L DT we have E
[Equation]
ref ( x K )) (cid:3) which does not depend on θ and E q 1 | 0 ( x 1 | x 0 ) (cid:104) log p {+θ 
+}
| 1 ( x 0 | x 1 ) (cid:105) which we expand here
[Equation]
where on the final line we have used eq 6.
In summary,
[Equation]
We now take the limit of L DT as ∆ t → 0 and K → ∞ .
[Equation]
We can estimate the integral with Monte Carlo if we consider it to be an expectation with respect to a uniform distribution over times (0 , T ) .
We also write ˆ R θt ( x, x ) explicitly as the negative off diagonal row sum.
[Equation]
B.3 Proof of Proposition
Proof.
We assume q t | s ( x 1: D t | x 1: D s ) factorizes as (cid:81)
Dd =1 q t | s ( x dt | x ds ) where q t | s ( x dt | x ds ) , d = 1 , . . .
, D are the transition probabilities for independent singular dimensional CTMCs each with forward rate R dt (˜ x d , x d ) .
In the following, we will drop time subscripts on x arguments.
To find the correspondence between R 1: D t and R dt , we use the Kolmogorov forward equation
[Equation]
Substitute in our factorized form for q t | s into the LHS
[Equation]
d
[Equation]
y
[Equation]
: D
[Equation]
We therefore obtain (cid:88)
[Equation]
This must be true for all possible factorizable forward process transitions, q t | s , including q t | s ( y 1: D | ˜ x 1: D ) =
δ y 1: D , ˜ x 1: D .
This choice gives us our forward rate relation
[Equation]
Substituting this into our expression for the reverse rate from Proposition 1 we obtain ˆ R 1: D t ( x 1: D , ˜ x 1: D ) =
(cid:88)
[Equation] [Equation]
B.4 Proof of Proposition
Proof.
By the Kolmogorov forward equation applied to the forwards process, we have
[Equation]
In addition, applying the Kolmogorov forward equation to the reverse process, which has the same marginals as the forward but time-reversed, we get
[Equation]
Summing these two equations gives
[Equation]
Therefore, by comparison with the Kolmogorov equation, R t + ˆ R t is the rate matrix of a CTMC with invariant distribution q t .
B.5 Proof of Theorem
In this section, we derive a bound on the error of our tau-leaping diffusion model.
Because the tauleaping approximation is only interesting in the case where multiple jumps are made along different dimensions in a single step, we choose to make the dependence of our bound on the dimension of our model explicit, rather than simply considering the case of fixed D and τ → 0 .
Recall from the main text that we have a time-homogeneous rate matrix R t on X , from which we construct the factorised rate matrix R 1: D t on X D by setting R dt
=
R t for each d , and will denote | R | = sup t ∈ [0 ,T ] ,x ∈X | R t ( x, x ) | , and let t mix be the (1/4)-mixing time of the CTMC with rate R t .
We also define addition on the state space X D using a mapping from X to Z as in Section 4.3 and component-wise addition.
Theorem 1.
For any D ≥ 1 and distribution p data on X D , let { x t } t ∈ [0 ,T ] be a CTMC starting in p data with rate matrix R 1: D t as above.
Suppose that ˆ R θ 1: D t is an approximation to the reverse rate matrix and let ( y k ) k =0 , 1 ,...,N be a tau-leaping approximation to the reverse dynamics with maximum step size τ .
Suppose further that there is some constant M > 0 independent of D such that
[Equation]
for all t ∈ [0 , T ] .
Then under the assumptions listed below, there are constants C 1 , C 2 > 0 depending on X and R t but not D such that, if L ( y 0 ) denotes the law of y 0 , we have the total variation bound |
[Equation]
L
[Equation]
The above theorem holds under the following assumptions, where we write x ∼ y for x, y ∈ S D if they differ in at most one coordinate.
Assumption 1.
The data distribution p data is strictly positive.
Assumption 2.
There exists a constant C 1 > 0 , depending on S and R t but not D , such that for all t ∈ [0 , T ] and x, y ∈ S D such that x ∼ y , we have
[Equation]
Assumption 3.
There exists a constant C 2 > 0 , depending on S and R t but not D , such that for all t ∈ [0 , T ] and all x, y ∈ S D such that x ∼ y , we have
[Equation]
If instead we were to allow C 1 and C 2 to depend on the dimension D , then Assumptions 2 and 3 follow trivially from Assumption 1 and the finiteness of the state space.
However, we choose the stronger formulation above in order to make explicit the dependence of the error bound on the dimension, as previously explained.
As remarked in the main text, in most cases of practical interest (including the two examples explored in Section 6), Assumption 3 holds only approximately.
However, we still expect the bound in
Assumption 3 to hold whenever x, y are in addition chosen such that the tau-leaping approximation of the reverse process makes a jump between them with reasonably high probability.
For example, in the case where our data is ordinal, we expect that for any x ∼ y jumps from x to y are only common when x is close to y , and thus ˆ R t ( x, x + z ) and ˆ R t ( y, y + z ) should be reasonably close whenever a jump from x to y occurs.
Under a weaker assumption of this form, the proof of Theorem 1 can be adapted to work along similar lines, at the cost of a significant increase in technicality.
We therefore choose to focus on the simpler case where Assumption 3 holds as it illustrates the key ideas.
In order to prove Theorem 1 we will require the following lemmas.
Proposition 5.
Let ( x t ) t ∈ [0 ,T ] and ( y t ) t ∈ [0 ,T ] be continuous time
Markov chains on a finite state space S with generators G t and H t respectively which are both bounded and continuous in t .
Let the
Markov kernels associated to X and Y be K and L respectively.
Then for any probability distribution ν on S we have
[Equation]
Proof.
We define a coupling of ( x t ) t ∈ [0 ,T ] and ( y t ) t ∈ [0 ,T ] as follows, based on the construction in Chapter 20.1 of [24].
First take Z ∼ ν and set x 0 = y
= Z .
Also define the variables ˜ x 0 = ˜
y
= Z .
Next, fix λ such that | G t ( x, x ) | , | H t ( x, x ) | ≤ λ for all x ∈ S , t ∈ [0 , T ] , let ( N s ) 1 ≤ s ≤ T be a Poisson process on [0 , T ] of rate λ , and set N 0 = 0 .
We write N = N T , and S 1 , S 2 , . . .
, S N for the arrival times and set S n +1 = T .
We construct x t and y t for t > 0 inductively as follows.
For t ∈ [0 , S 1 ) let x t =
y t
= x 0 .
Let 1 ≤ j ≤ N .
Given ( x r : r < S j ) , ( y r : r < S j ) , and ˜ x j , ˜ y
j , define the probability measures
[Equation]
Sample (˜ x j +1 , ˜ y j +1 ) from a maximal coupling of ( ρ j , ρ ′ j ) and for t ∈ [ S j , S j +1 ) set x t = ˜ x
j +1 , y t = ˜
y j +1 .
Finally set x T = x S N and y T =
y S N .
Now, observe that ( x t , y t ) t ∈ [0 ,T ] defined in this way is a coupling of the given Markov chains.
Moreover,
[Equation]
and using the fact that jumps are coupled maximally
[Equation]
as required.
Proposition 6.
For all t ∈ [0 , T ] and x, y ∈
X D such that x ∼ y , we have
[Equation]
Moreover, it follows that ˆ R t is bounded and continuous in t .
Proof.
Omitting the superscripts for brevity where the notation is clear, we have
[Equation]
where the second line follows from Kolmogorov’s forward equation and the final inequality follows from Assumption 2 plus the fact that R t ( z, x ) (resp.
R t ( z, y ) ) is only non-zero when x ∼ z (resp.
y ∼ z ), and there are at most | S
|| D
| values of x (resp.
y ) for which this holds.
We now give the proof of Theorem 1.
Proof of Theorem 1.
Let us label the time steps used in tau-leaping by
= t 0 < t 1 < · · · < t N = T , denote τ k = t k − t k − 1 , and denote the target stationary distribution by π D ( x 1: D )
=
(cid:81)
Dd =1 π ( x d ) , where π is the invariant distribution of the single-dimensional transition matrix R 1 t .
Also, let R θ, ( τ ) k be the Markov kernel corresponding to applying the tau-leaping approximation with rate matrix ˆ R θt k to move from t k to t k − 1 , and denote R θ, ( τ ) =
R θ, ( τ ) N
R θ, ( τ ) N − 1 . . .
R θ, ( τ )so that R θ, ( τ ) expresses the full dynamics of the tau-leaping process and we have L (ˆ y 0 )
=
π D R θ, ( τ ) .
Then, as in [19] we can decompose
[Equation]
where P R is the path measure of the exact reverse process.
We deal with the second term first.
Let t mix be the (1/4)-mixing time of the single-dimension CTMC with rate matrix R 1 t , i.e.
[Equation]
It then follows from
[Equation]
that t D mix , the (1/4)-mixing time of the full CTMC with rate matrix R 1: D t , satisfies the inequality t Dmix ≤ { 1+ ⌈ log 2 D ⌉} t mix .
If we view ( x mt Dmix ) m ∈ N as a discrete-time Markov chain, then standard results on Markov chain mixing (see, for example, Chapter 4.5 of [24]) show that
[Equation]
It then follows that for any T ≥
we have
[Equation]
completing the bound on the second term.
To bound the first term
[Equation]
we define P k
=
( P R ) T − t k − 1 | T − t k and decompose it as
[Equation]
by proceeding inductively.
So it suffices to find bounds on the total variation distance accumulated on each interval [ t k − 1 , t k ]
[Equation]
Let R θk be the Markov kernel corresponding to running the chain from t k to t k − 1 with constant rate matrix ˆ R θt k .
Since by Proposition 6 the reverse rate matrix ˆ R t is bounded and continuous in t
[Equation]
using
Proposition 5 we made deduce that for any distribution ν on S we have
[Equation]
The first half of this expression can be bounded using the Mean Value Theorem, according to
[Equation]
where in the first line we have used that the summand is only non-zero when y ∼ x , and there are at most | S
|| D
| values of y for which this holds.
The second term can be bounded using condition (2), to get
[Equation]
Combining these two expressions, we get a bound on || ν
P k − ν R θk || TV .
[Equation]
It remains to bound ||
ν R θk − ν R θ, ( τ ) k || TV .
Note that performing tau-leaping with rate matrix ˆ R θt k starting in x t k is equivalent to running a continuous time Markov chain from time t k to t k − 1 with constant rate matrix ˆ R θ, ( τ ) t k given by
[Equation]
τ
[Equation]
t
[Equation]
(followed potentially by a clamping operation to keep us within X D ).
By an analogous argument to the proof of Proposition 5,
[Equation]
where the expectation is taken over ( x t ) t ∈ [ t k − 1 ,t k ] distributed according to the exact CTMC with rate matrix ˆ R θ t k .
(Note we have disregarded the clamping operation, since this can only decrease the resulting total variation distance.)
We may rewrite this bound in terms of the exact reverse process using condition (2) to get
[Equation]
Let J t be the number of jumps that ( x t ) makes between t k and t , and label the times of these jumps as s 1 , . . .
, s j where t ≤ s 1 ≤ · · · ≤ s j ≤ t k and j =
J t for convenience.
Then by Assumption 3, we have
[Equation]
where we have made the substitution
z
= y − x t k .
We conclude that
[Equation]
where to bound E [ J t ] we have observed that jumps of ( x t ) occur at a rate bounded above by sup x
| ˆ R θt k ( x, x ) |
[Equation]
and in the last line we have used the condition (2) and Assumption 2.
Since the above holds for any choice of x t k
[Equation]
it follows that
[Equation]
Summing over k and putting all our bounds together, we get ||L ( y 0 ) − p data || TV ≤
[Equation]
[Table]
as required
[Equation]
C Continuous Time ELBO Details C.1 Comparison with the Discrete Time ELBO
It is easiest to gain intuition on the L CT objective by comparing it to its discrete time counterpart,
L DT , and examining the way in which L DT in the limit becomes L CT when we take the time step size to be very small.
We repeat the definition of L CT here for convenience
[Equation]
Recall that a single term from the KL sum in L DT up to an additive constant independent of θ is
[Equation]
Minimizing this term is to sample ( x k , x k +1 ) from the forward dynamics and then maximize the assigned model probability for the pairing in the reverse direction.
A similar idea can be used to understand L CT .
First, we write log p θk | k +
( x k | x k +1 ) in terms of ˆ R θk as [Equation] k ,
[Equation]
where we have separated the cases when x k = x k +
and when x k ̸ = x k +1 (see the proof of L CT for the full details).
The first term will become the (cid:80) x ′ ̸
= x ˆ R θt ( x, x ′ ) term in L CT whilst the second term will become the Z t ( x ) log (cid:0)
ˆ R θt (˜ x, x ) (cid:1) term.
Now, when we minimize L CT , we are sampling ( x, ˜ x ) from the forward process and then maximizing the assigned model probability for the pairing in the reverse direction, just as in L DT .
The slight extra complexity comes from the fact we are considering the case when x k = x k +
and the case when x k ̸ = x k +1 separately.
When x k = x k +1 , this corresponds to the first term in L CT which we can see is minimizing the reverse rate out of x which is exactly maximizing the model probability for no transition to occur.
When x k ̸ = x k +1 , this corresponds to the second term in L CT , which is maximizing the reverse rate from ˜ x to x which in turn maximizes the model probability for the ˜ x to x transition to occur.
C.2 Conditional Form
For the conditional form of L CT , denoted as ¯ L CT , we instead upper bound the negative conditional model log-likelihood, E p data ( x 0 ,y )
[ − log p θ
( x 0 | y )] where y is our conditioner.
¯ L CT has the following form
[Equation]
where
[Equation]
This follows easily from considering the conditional form of the discret time ELBO, ¯ L DT and using the same arguments as before to go from discrete time to continuous time.
[Equation]
C.3 Continuous Time ELBO with Factorization Assumptions
In the following Proposition, we show the form of L CT when we use a factorized forward process.
We note that in the proof we rearrange the sampling distribution from p data ( x 1: D 0 )
q t | 0 ( x 1: D | x 1: D 0 ) r t (˜ x 1: D | x 1: D ) to p data ( x 1: D 0 ) ψ t (˜ x 1: D | x 1: D 0 ) ϕ
t ( x 1: D | ˜ x 1: D , x 1: D 0 ) .
This is not strictly necessary but it allows us to analytically sum over the intermediate x 1: D variable which greatly reduces the variance of the resulting objective.
Proposition 7.
The L CT objective when we substitute in the factorized forms for the forward and reverse process given in Proposition 3 is
[Equation]
+ C
[Equation]
) where ◦ represents the concatenation of a D − 1 dimensional vector, x 1: D \ d with a scalar x d , such that the resultant D dimensional vector has x d at its d th dimension.
ψ t (˜ x 1: D | x 1: D 0 ) is defined as the marginal of the forward noising process joint, (cid:82) q t | 0 ( x 1: D | x 1: D 0 ) r t (˜ x 1: D | x 1: D ) d x 1: D .
Proof.
We first re-write the general form of L CT here
[Equation]
where
[Equation]
With a factorized forward process, ˆ R θt becomes
[Equation]
where ˆ
[Equation]
Substituting this form for ˆ R θ 1: D t into the first term in L CT we get (cid:88)
[Equation]
Now we tackle the second term in L CT .
We first re-arrange the distribution over which we take the expectation:
[Equation]
We have,
[Equation]
To find the normalization constant
[Equation]
we can sum the proportional term over x 1: D (cid:88)
[Equation]
Therefore
[Equation]
where ϕ
[Equation]
| ˜ x
[Equation]
=
[Equation] [Equation]
D 0 )
Now we write the second term as
T
[Equation]
(cid:35)
C.4 One Forward Pass
To evaluate the L CT objective, we naively need to perform two forward passes of the denoising network: p θ 0 | t ( x 0 | x ) to calculate ˆ R θt ( x, x ′ ) and p θ
| t ( x 0 | ˜ x ) to calculate ˆ R θt (˜ x, x ) .
This is wasteful because ˜ x is created from x by applying a single forward transition which on multi-dimensional problems means ˜ x differs from x in only a single dimension.
To exploit the fact that ˜ x and x are very similar, we approximate the sample x ∼ q t ( x ) with the sample ˜ x ∼ (cid:80) x
q t ( x ) r t (˜ x | x ) .
This gives the more efficient objective,
[Equation]
[Table]
The approximation is valid because q t ( x ) and (cid:80) x
q t ( x ) r t (˜ x | x ) are very similar distributions, now show
[Equation]
where on the third line we have used the Kolmogorov forward equation and defined δ t
= 1 / (cid:80) x ′ ̸
= x R t (˜ x, x ′ ) .
The distribution (cid:80) x q t ( x ) r t (˜ x | x ) is therefore q t + δt (˜ x ) approximated using a first-order Taylor expansion around q t (˜ x ) .
We notice that δt is the average time to the next transition at time t .
δt can be calculated for the practical settings we consider, its varies between 2 × 10 − 6 T and 2 × 10 − 8 T in the image modelling task and is 1 × 10 − 3 T in the monophonic music task.
We perform an ablation experiment comparing between training with L eCT and L CT on the monophonic music dataset, the results are shown in Table 3.
We find that we gain a small boost in performance when using the more efficient L eCT objective alongside the improved efficiency.
We hypothesize that this is because of a slight reduction in variance for the L eCT objective due to increased negative correlation between the two terms in the objective when ˜ x is shared between them.
D Direct Denoising Model Supervision
Following [8], we can introduce direct p θ 0 | t supervision into the optimization objective which has been found empirically to improve performance.
We first contextualize the change by expressing L CT with the dependence on p θ
| t made explicit.
[Equation]
The signal for p θ 0 | t ( x 0 | x ) comes through a sum over x 0 weighted by the ratio q t | 0 ( x | x 0 ) q t |
(˜ x | x 0 ) .
We can also provide a direct denoising signal by predicting the clean datapoint x 0 from the corrupted version x and using the negative log-likelihood loss.
[Equation]
Proposition 8.
The true denoising distribution, q 0 | t , minimizes L ll
Proof.
[Equation]
where C is a constant independent of θ .
Therefore, minimizing L ll is equivalent to minimizing the
KL divergence between q 0 | t and p θ 0 | t , which is minimized when p θ 0 | t = q 0 | t .
If we obtain the true denoising distribution, p θ 0 | t = q 0 | t , then we will have the true reverse rate, ˆ R t .
[8] find that optimizing with an objective combining L ll and L DT performs best, which we can also do in continuous time
[Equation]
where λ is a hyperparameter.
In [8], it was found that training with L ll alone resulted in poorer performance than when the ELBO was included in the loss.
We provide a theoretical hypothesis as to why this may be the case here.
We show that minimizing L ll is equivalent to minimizing an upper bound on the negative ELBO in discrete time and thus by training with L ll we are simply minimizing a looser bound on the negative model log-likelihood than if we were to use the negative
ELBO directly.
Proposition 9.
Minimizing the sum of negative log-likelihoods
[Equation]
is equivalent to minimizing an upper bound on the negative ELBO.
Proof.
[Equation]
Consider one term from the sum
[Equation]
Now,
[Equation]
Therefore,
[Equation]
We can see that only the first term depends on θ .
E Choice of Forward Process
We need to choose the structure of R t such that we can analytically obtain q t | 0 marginals to enable efficient training.
Proposition 10.
If R t and R t ′ commute for all t , t ′ then q t | 0 ( x = j | x 0 = i ) = (cid:0) exp (cid:2) (cid:82) t
R s ds (cid:3)(cid:1) ij where exp here is understood to be the matrix exponential function.
Proof.
Let ( P t ) ij =
q t | 0 ( x = j | x 0 = i ) .
We show that P t = exp (cid:16)(cid:82) t
R s ds (cid:17) is a solution to the
Kolmogorov forward equation, which in matrix form reads, ∂ t P t = P t R t .
Writing the matrix exponential in sum form
[Equation]
Now, differentiating and using the fact that R t , R t ′ commute.
[Equation]
As stated in the main text, we achieve the commutative property by selecting R t = β ( t ) R b where β ( t ) is a time dependent scalar and R b is a constant base matrix.
We can utilize the eigendecomposition of R b = Q Λ Q − 1 to efficiently calculate P t ,
P t = exp
[Equation]
Since Λ is a diagonal matrix, the matrix exponential coincides with the element wise exponential making the final expression tractable to compute.
We choose β ( t ) =
ab t log b because this makes the integral which dictates the variance of q t | 0 have a simple form (cid:82) t 0 β ( s )
ds =
ab t − a .
For categorical problems, we found a uniform rate matrix works well, R t
= β 11 T − βS Id .
This is directly analogous to the discrete time uniform transition matrix: P = α 11 T + (1 − Sα )Id with α depending on the time discretization used.
Indeed, if one calculates the corresponding discrete transition matrix for the uniform R t rate through the matrix exponential, the uniform transition matrix is obtained.
Another categorical corruption process is the absorbing state process.
In discrete time, the transition matrix is given by P = α 1 e T ∗ + (1 − α )
Id where e ∗ is the one-hot encoding of the absorbing state.
The corresponding absorbing state continuous time process has transition rate matrix:
R t = β 1e T ∗ − β Id .
The correspondence for more complex transition matrices e.g. the Discretized
Gaussian matrix in [8] is much harder to find analytically especially if the time inhomogeneous case is considered.
For datasets with an ordinal structure, we construct a new rate matrix that maintains a bias towards nearby states using a similar approach as that taken by [8] to construct the Discretized
Gaussian matrix.
We construct this matrix by first picking a desired stationary distribution, p ref , and then filling in matrix entries such that we encourage transitions to nearby states whilst keeping p ref as our stationary distribution.
Specifically, we let p ref be a discretized Gaussian over the state space, i.e.
[Equation]
To find a condition on the rate such that this is the case, recall the Kolmogorov differential equation for the marginals
[Equation]
Now, consider a rate that is in detailed balance with p ref
[Equation]
Substituting this rate into the Kolmogorov equation, we see that p ref is the stationary distribution
[Equation]
where the last line follows from the fact that the row sum of a rate matrix is zero.
Note that any
R t = β ( t ) R b will also have this stationary distribution as the multiplication by β ( t ) can be seen as just a scaling of the time axis.
From the detailed balance equation, we gain a condition on R b such that our desired p ref is the stationary distribution
[Equation]
This gives constraints on diagonal elements within R b but does not fully define the entire matrix.
To do this, we first make the assumption that µ is selected to be at the center of the state space.
Then we set off diagonal terms to the right of the diagonal in the top half of the rate matrix and off diagonal terms to the left of the diagonal in the bottom half to be 1.
Finally, progressing in from the top and bottom of the rate matrix we make definitions of rate matrix values that have not already been defined by the detailed balance condition.
For clarity, we provide a pictorial representation of this scheme for an 8 × 8 rate matrix below    · □ □ □ □ □ □ △ · □ □ □ □ △ △ △ · □ □ △ △ △ △ △ · △ △ △ △ △ △ △ · △ △ △ △ △ △ □ · △ △ △ △ □ □ □ · △ △ □ □ □ □ □ · 
  where □ represents a value we will define, △ represents a value that is defined relative to another entry through the detailed balance condition and · is a diagonal entry that is equal to the negative off diagonal row sum.
We could define □ values to be 0 to gain a sparse rate matrix, however, we found in early experiments that allowing transitions to further away states greatly reduces the mixing time and gives better performance.
We define □ in each row similarly, by setting it equal to exp [ − i 2 /σ
2 r ] where i is the distance away from the ‘ 1 ’ value in that row and σ r is a hyperparameter defining the length scale in state space of a typical transition.
This biases our forward process to make transitions between nearby states, at a length scale of σ r .
F CTMC Simulation F.1 Exact CTMC and Tau-Leaping
In this section, we first describe exact CTMC simulation before giving an algorithmic description of tau-leaping.
When a CTMC has a time-homogeneous rate matrix, we can use Gillespie’s Algorithm [21, 22, 23] to exactly simulate it.
This algorithm is based on the jump chain/holding time definition of the CTMC.
It repeats the following two steps:
Draw a holding time from an exponential distribution with mean − 1 /R ( x, x ) and wait in the current state x for that amount of time.
• Sample the next state from r (˜ x | x ) =
(1 − δ x, ˜ x ) R ( x, ˜ x ) (cid:80) x ′̸ = x R ( x,x ′ )
This Algorithm can be adjusted for the case when we have a time-inhomogeneous rate matrix using the modified next reaction method [33].
However, both algorithms still step through each transition in the CTMC individually and are thus unsuitable in our case because only one dimension would change for each simulation step making it very computationally expensive to produce a sample.
Instead we use tau-leaping that allows multiple dimensions to change in a single simulation step.
We detail this method in Algorithm 1.
F.2 Predictor-Corrector Discussion
In this section we compare predictor-corrector sampling schemes as applied to continuous state spaces and discrete state spaces.
t ← T x 1: D t ∼ p ref ( x 1: D T ) while t > 0 do
Compute p θ
| t ( x d 0 | x 1: D t ) , d = 1 , . . .
, D with one forward pass of the denoising network for d = 1 , . . .
, D do for s = 1 , . . .
, S \ x dt do ˆ R θ d t ( x 1: D t , s ) ← R dt ( s, x dt ) (cid:80) x d 0 p θ
| t ( x d 0 | x 1: D t ) q t | 0 ( s | x d 0 ) q t | 0 ( x dt | x d 0 )
P ds ← Poisson (cid:16) τ
ˆ R θ d t ( x 1: D t , s ) (cid:17) end endfor d = 1 , . . .
, D do if data is categorical AND (cid:80) Ss =1 P ds > 1 then x dt − τ ← x dt // reject change else x
[Equation]
end end x 1: D t − τ ← Clamp ( x 1: D t − τ , min = 1 , max = S ) t ← t − τ end
The predictor-corrector scheme in continuous state spaces was introduced in [4].
It consists of alternating between a predictor step and a corrector step:
Predictor
[Equation]
Corrector
[Equation]
where x i is the state at sampling step i , s θ is the learned score model approximating ∇ x log q t ( x ) and γ i , ϵ
i are the step sizes for the predictor and corrector respectively.
We see that both take similar forms, except the corrector adds in a factor √ 2 more Gaussian noise during the update step.
In discrete state spaces, rather than sampling using gradient guided stochastic steps as in the continuous state space case, we sample by simulating CTMCs with defined rates.
When we take a predictor step, we simulate using ˆ R θt and when we take a corrector step we simulate using R cθt =
ˆ R θt + R t .
If we simulate the CTMC exactly, we have seen in the previous section that this amounts to sampling next states from the categorical distribution defined by normalizing the row of the rate matrix corresponding to the current state.
Therefore, corrector sampling can be seen as sampling from a slightly noisier categorical distribution defined through R cθ t as compared to the predictor categorical distribution defined through ˆ R θt .
This is analogous to the increased Gaussian noise applied during a corrector step in continuous state spaces.
Adding corrector steps brings the marginal of the samples closer to q t ( x ) and continued application of the corrector will further explore the domain of q t ( x ) .
In previous work on continuous state predictor-corrector methods, the number of corrector steps has been small (e.g. 1 or 2 corrector steps per predictor step) or indeed the corrector steps have been removed altogether.
In this work we have found that using up to 10 corrector steps per predictor steps can be beneficial during certain regions of the reverse generative process.
Additionally, in continuous state spaces, it has been observed that too many corrector steps can result in unwanted noise in the generated data [34].
We hypothesize that corrector steps are better utilized in discrete state spaces to explore the domain of q t ( x ) than in continuous state spaces.
This is because, the corrector update is defined largely through the reverse rate itself, ˆ R θt , just with the categorical probabilities being annealed slightly more towards uniform through the addition of the forward rate R t .
This may be a more effective update than simply adding extra Gaussian noise in the continuous state space case.
Furthermore, the denoising model in continuous state spaces can be seen as outputting a point estimate of x 0 of dimension D .
However, in discrete state spaces, the denoising model outputs a categorical distribution over every dimension (output dimension D × S ) allowing it to express some uncertainty information in the x 0 prediction, albeit with conditional independence between the dimensions.
Adding corrector steps in discrete state spaces would then allow information to mix between dimensions for the current time step, exploring modes of q t ( x ) .
We explore this idea on the image modelling task in Figure 7.
We run the reverse generative process until time t = 0 .
4 at which point we hold the time constant and apply 1000 corrector steps.
We see that the resulting progression of x t states explores potential local modes of q t ( x ) in the local region of image space.
G Implicit Dimensional Assumptions Made in Discrete Time
In discrete time, the parametric reverse kernel, p θ k | k +1 , is commonly defined through a denoising model p θ
| k +1 .
Here, we examine this definition in the multi-dimensional case where the forward process factorizes, as in Appendix C.3 and previous discrete time work [8].
We begin by writing the true full dimensional reverse kernel, q k | k +1 , in terms of the true denoising distribution, q
|
k +1 .
[Equation]
=
[Equation]
=
[Equation]
(cid:88)
[Equation] [Equation]
[Equation]
( x d 0 | x 1:
[Equation]
k , x
[Equation]
: D k +1 )
q k | 0 ,k +1 ( x
[Equation]
| x d 0 , x dk +1 ) where on the final line we have used the fact that the forward process is independent across dimensions.
To create our approximate reverse kernel, p θk | k +1 , we approximate q 0 | k +
( x d 0 | x 1: d − 1 k , x 1: D k +1 ) with p θ 0 | k +1 ( x d 0 | x 1: D k +1 ) , p θk | k +1 ( x 1: D k | x 1: D k +1 ) =
D (cid:89) d
= (cid:88) x dp θ 0 | k +1 ( x d 0 | x 1: D k +1 )
q k | 0 ,k +1 ( x dk | x d 0 , x dk +1 )
We throw away the extra x 1: d − 1 k conditioning because we use a non-autoregressive model that takes in x 1: D k +1 and in a single forward pass gives conditionally independent probabilities over x d 0 , d = 1 , . . .
, D .
For finite K , this approximation can never match the true kernel because we are not conditioning on all relevant information.
Of course, as K gets larger, this approximation becomes more accurate.
Since we operate in the continuous regime, we do not have to make this approximation because the conditionally independent denoising model, q
| t ( x d 0 | x 1: D ) , appears directly in our reverse rate, ˆ R 1: D t , when we factorize the forward process (see Proposition 3).
H Experimental Details
In this section, we provide additional details for the experiments we performed applying our method to practical problems.
The code for our models is available at https://github.com/andrew-cr/tauLDR .
Before describing the specifics for each experiment, we first explain the implementation details common to all.
When we evaluate the objective L CT on each minibatch of training datapoints, we must sample a time for each from t ∼ U (0 , T ) which represents the point in the forward process which we will noise to.
Training instabilities can be found if t is sampled very close to 0 because the reverse rate, ˆ
R t , becomes ill-conditioned in this region.
This phenomenon is also observed in continuous state space models because the score, ∇ x log q t ( x ) , becomes ill-conditioned close to t = 0 .
The reverse rate and score become ill conditioned close to the start of the forward process because the marginal probability, q t ( x ) , will be highly peaked around the data manifold and log q t ( x ) will explode in regions that are not close to the data.
To avoid these issues, a common trick is to set a minimum time such that t ∼ U ( ϵ, T ) .
ϵ is set such that the level of noising at t =
ϵ is very small and reverse sampling to this point will produce samples very close to p data .
In our experiments, we set ϵ = 0 .
01 T .
During reverse sampling, we use tau-leaping to simulate the reverse process from t =
T until t =
ϵ because the reverse rate is not trained for t <
ϵ .
This produces a sample close to p data .
We found improved performance in metrics such as FID if we then complete a final step to remove the small amount of noise that may still be present in the sample.
Specifically, we pass the sample through the denoising model p θ 0 | t ( x 0 | x t ) with t =
ϵ to obtain an output of shape D × S where D is the dimensionality of the problem.
This is a probability distribution over the states for each of the dimensions.
We set the value of each dimension to the state with the highest probability.
This then produces a sample which has all of the noise removed.
The specific value of T within our model is arbitrary because the forward process can be scaled in the time axis to provide the same noising process for any T .
Therefore, we simply set T = 1 .
H.1 Demonstrative Example
Our 2d dataset is created by sampling 1M 2d points from a 32 × 32 state space with probability proportional to the pixel values of a 32 × 32 grayscale image of a τ character.
For our forward process, we use a Gaussian rate (see Appendix E) with stationary distribution standard deviation σ 0 = 8 and rate length scale σ r = 1 .
We use a rate schedule of β ( t ) =
5 × 5 t log(5) .
To represent p θ 0 | t we use a residual MLP.
The architecture consists of an input linear layer to lift the input dimension of 2 to the internal network dimension of 16 .
Then, there are 2 residual blocks each consisting of: a single hidden layer MLP of hidden dimension 32 , a residual connection to the input of the MLP, a layer norm, and finally a FiLM layer [35] modulated by the time embedding.
At the output, there is a single linear layer with output size of 2 × 32 = 64 representing state probabilities in each of the 2 dimensions.
The time is embedded using the Transformer sinusoidal position embedding [36] creating an embedding of size 32 .
Then, the embedding is further processed by a single hidden layer MLP with hidden layer size 32 and output size 128 .
To create the FiLM parameters in each residual block, the time embedding is passed through a linear layer with output of size 32 to provide a multiplicative and additive modulation to the state dimension of 16 .
We minimize the L CT objective using Adam with a learning rate of 0 . 0001 and batch size of 32 for 1M steps.
For the exact simulation we use the next reaction method with modifications for time dependent transition rates [33].
This method steps through each transition in the exact simulation path individually by calculating the time to the next occurrence of each transition type and applying the transition that occurs soonest.
Exact algorithmic details can be found in [33].
To calculate the time to the next occurrence for a transition, we need to integrate the reverse rate matrix (eq (13) in [33]).
We do this with euler integration with a step size of 0.001.
H.2 Image Modeling
We train on the CIFAR10 training dataset that contains 50000 images of dimension 3 × 32 × 32 .
We evaluate the test ELBO on the CIFAR10 test dataset which consists of 10000 images.
For the forward noising process, we use the the Gaussian rate (see Appendix E) with stationary distribution standard deviation of σ 0 = 512 and rate length scale σ r = 6 .
This effectively defines a uniform stationary distribution since the state space is of size 256 .
We found this performs better than a more concentrated Gaussian.
Our β schedule is β ( t ) =
× 100 t log 100 .
This was selected in accordance with σ r such that the overall shape of progression of the q t | 0 variances approximately matches that of the schedule proposed in [3].
Our p θ 0 | t model is parameterized with the standard U-net [37] architecture introduced in [3].
The network follows the PixelCNN++ backbone [38] with group normalization layers.
There are four feature map resolutions ( 32 × 32 to 4 × 4 ) in the downsampling/upsampling stacks.
At each resolution there are two convolutional residual blocks.
There is a self-attention block between the residual blocks at the 16 × 16 resolution level [39].
The time is input into the network by first embedding with the Transformer sinusoidal position embedding [36].
This time embedding is passed into each residual block by passing it through a SiLU activation [40] and then a linear layer before adding it onto the hidden state within the residual block between the two convolution operations.
The original architecture of [3] has an output of dimension 3 × 32 × 32 as it makes a point prediction of x 0 given x t .
In order for the model to output probabilities over x 0 (i.e. an output dimension of 3 × 32 × 32 × 256 ) we make the adjustments suggested in [8].
Specifically, we use their truncated logistic distribution parameterization where the model outputs the mean and log scale of a logistic distribution i.e. an output dimension of 3 × 32 × 32 × 2 .
The probability for a state is then the integral of this continuous distribution between this state and the next when mapped onto the real line.
To impart a residual inductive bias on the output, the mean of the logistic distribution is taken to be tanh ( x t + µ ′ ) where x t is the normalized input into the model and µ ′ is mean outputted from the network.
The normalization operation takes the input in the range 0 , . .
, 255 and maps it to [ − 1 , 1] .
In total, our network has approximately 35 .
7 million parameters.
We optimize with the auxiliary objective described in Appendix D with λ = 0 .
Within the auxiliary objective, we use the one-forward pass version of the continuous time ELBO, L eCT .
We optimize with Adam for 2M steps with a learning rate of 0.0002 and batch size of 128.
We use the standard set of training tricks to improve optimization [3, 4].
Throughout training we maintain an exponential moving average of the parameters with decay factor 0.9999.
These average parameters are used during testing.
At the start of optimization we use a linear learning rate warm-up for the first 5000 steps.
We clip the gradient norm at a norm value of 1.0.
We set the dropout rate for the network at 0.1.
The skip connections for each residual block are rescaled by a factor of 1 √ 2 .
The input images have random horizontal flips applied to them during training.
For sampling in Table 1 we set τ = 0 .
001 for τ LDR- 0 and set τ = 0 .
002 for τ LDR- 10 .
The 10 corrector steps per predictor steps for τ LDR- 10 are introduced after t < 0 .
1 T .
We found that introducing the corrector steps near the end of the reverse sampling process had the best improvement in sample quality for the smallest increase in computational cost.
When performing tau-leaping with the corrector rate, R ct , we have control over what τ we use since we are sampling a different
CTMC (with q t as its stationary distribution) to the original reverse CTMC.
We found that setting the corrector rate τ to be 1 .
5 times the original τ for the reverse CTMC achieves the best performance in this example.
We train using 4 V100 GPUs on an academic research cluster.
To calculate Inception and FID values, we use pytorch-fid [41] and a further development 1 .
We verified this library produced comparable values to previous work by calculating the Inception and FID scores for the published images from the DDPM [3] method.
We show a large array of unconditional samples from the τ LDR- 10 model in Figure 8.
We now also present statistics from the reverse sampling process with standard tau-leaping with τ = 0 .
Figure 9 shows the proportion of dimensions that transition during a single step of tau-leaping.
We see that during the initial stages, every dimensions changes during every tau-leaping step, but nearer the end of the process, more dimensions will have settled in their final positions and the proportion is less.
In
Figure 10 we show the proportion of dimensions that are clipped due to proposing an out of bounds jump.
Overall, the proportion is small.
It is largest at the start of the process when we have initially sampled from the approximately uniform p ref and there will be dimensions close to the boundary.
As pixel values settle to their final values, the proportion reduces.
Figure 11 shows the progression of a selection of dimensions during the reverse sampling process.
A similar picture emerges where dimensions eventually settle in a region of the state space.
We also note that larger jumps are made in a single tau-leaping step nearer the start of the reverse process and smaller jumps are made nearer the end.
H.3 Monophonic Music
We generate our training dataset from the Lakh pianoroll dataset [29, 30] (license CC By 4.0).
This dataset consists of 174,154 multitrack pianorolls.
We go through all songs and all tracks within each song and select sequences that match the following criteria: they are monophonic (only one note played at a time), there is not a period longer than one bar in which no note is played, there is more than one type of note played in the sequence and finally there is no one note played for more than 50 time steps out of the total 256 time steps.
This removes the uninteresting and trivial sequences present within the dataset.
We then remove any duplicates in the result.
This leaves us withtraining examples and 950 testing examples.
Each song consists of 256 time steps (16 per bar) and each time step takes one from 129 values i.e. we have D = 256 and S = 129 .
This state value represents either a note from 128 options or a rest.
We scramble the ordering of this state space when mapping to the integers from 0 to 128.
When we input into the denoising network, we input as one-hot 129 dimensional vectors.
For the forward noising process, we use a uniform rate matrix, R b = 11 T − S Id and set β ( t ) = 0 .
We found a constant in time β ( t ) was sufficient for this dataset.
In our comparison, we used a birth/death rate matrix defined as   − λ λ  . . .
λ − 2 λ λ . . . λ − 2 λ λ . . .
λ − λ 
 this is the rate matrix for a birth/death process.
We set λ = 1 and β ( t ) = 12 × 10000 t log 10000 .
These hyperparameters were selected such that the forward process has a steady rate of noising whilst still having q T very close to p ref .
We chose to compare these types of rate matrix because the birth/death rate is inappropriate for this categorical data as adjacent states have no meaning
[Figure]
r o p o r t i o n o f
d
i m e n s i o n
[Figure]
[Figure]
p r o p o r t i o n o f
d i m e n s i o n s
[Figure]
since the mapping to the integers was arbitrary.
The uniform rate is suitable for this categorical data because, during a time interval, it has a uniform probability to transition to any other state.
The
D3PM baseline was implemented also with a time homogeneous uniform forward kernel set such that the rate of noising is matched in the discrete and continuous time cases.
We define our conditional denoising network, p θ
| t ( x 0 | x, y ) using a transformer architecture inspired by [31].
It takes an input of shape ( B, D, S ) where B is the batch size, D is the dimensionality ( 256 ) and S is the state size ( 129 ).
This final dimension contains the one-hot vectors.
The conditioning on the initial bars is achieved by concatenating the conditioning information y with the noisy input x .
At the start of the network, there is an input embedding linear layer with output of size 128 which is our model dimension for the transformer.
Then a transformer positional embedding is added to the hidden state.
Next a stack of 6 transformer encoder layers are applied which consist of a self attention block and a one hidden layer MLP.
The self attention block uses 8 heads and the MLP has a hidden layer size of 2048 .
At the output of each internal block, we apply dropout with rate 0 .
Finally, there is a stack of 2 residual MLP layers.
Each consists of a one hidden layer MLP with a hidden dimension of 2048 .
There is a residual connection between the input and output of the MLP.
A layer norm is applied to the output of the block.
To create the output of the network, there is an output linear layer with an output shape of ( B, D, S ) where now the S dimension has logit probabilities.
To instill a residual bias into the network, we add the one-hot input to the output logits.
All activations are ReLU.
The time is input into the network through FiLM layers [35].
First, the time is embedded using the sinusoidal transformer position embedding as in the U-net architecture used for image modeling to create an embedding size of 128 .
This is then passed into a single hidden layer MLP with hidden size
[Figure]
time
[Figure]
r e j e c t i o n r a t e
2048 and output size 512 .
Within each encoder and residual MLP block, there is a FiLM linear layer which takes in the 512 time embedding and outputs two FiLM parameters each of size 128 .
These are the scale and offset applied to the hidden state.
In the encoder blocks, this FiLM transform is applied after the self attention block and again after the fully connected block.
In the residual MLP blocks, it is applied after the layer norm operation.
Our network has approximately 7 million parameters in total.
We optimize using Adam for 1M steps with a batch size of 64 and learning rate of 0.0002.
We use the conditional ¯ L CT objective with additional direct p θ 0 | t supervision as described in Appendix D with weight λ = 0 . 001 .
We also make the same one forward pass approximation as explained in Appendix
C.4.
We use the standard set of training tricks to improve optimization [3, 4].
Throughout training we maintain an exponential moving average of the parameters with decay factor 0.9999.
These average parameters are used during testing.
At the start of optimization we use a linear learning rate warm-up for the first 5000 steps.
We clip the gradient norm at a norm value of 1.0.
We train on a single V
GPU on an academic cluster.
For sampling with τ LDR- 0 we use τ = 0 .
001 and for sampling with τ LDR- 2 we include 2 corrector steps per predictor step after t < 0 .
9 T .
The corrector rate is simulated with τ = 0 . 0001 which we found to perform best.
We reject any dimension in which 2 or more jumps are proposed as this is categorical data.
We plot the rejection rate in Figure 12.
Most of the time, the rejection rate is zero and there are few steps for which it increases slightly.
We show a large batch of samples from the first 10 songs in the test dataset in Figure 13.
We see that there is variation between the sampled completions and they consistently follow the style of the conditioning first two bars of the song.
Audio samples from the model are available at https://github.com/andrew-cr/tauLDR .
Finally, we examine the progression of a random selection of dimensions during reverse sampling for the uniform and birth/death rate matrix cases.
Figure 14 shows the progression for the uniform case, we see that large jumps through the state space are made throughout the reverse process.
Figureshows the progression for the birth/death case.
At the start of reverse sampling, no dimensions move as the rejection rate is high in this case because the rate matrix is not suitable for categorical data.
Nearer the end, small jumps are made between adjacent states but since large jumps between any category do not occur for this rate matrix, the performance will overall be worse.
I Ethical Considerations
Our work increases our theoretical understanding of denoising generative models and also improves generation capabilities within some discrete datasets.
Deep generative models are generic methods
[Figure]
Sample
[Figure] [Figure]
[Figure]
for learning from unstructured data and can have negative social impacts when misused.
For example, they can be used to spread misinformation by reducing the resources required to create realistic fake content.
Furthermore, generative models will produce samples that accurately reflect the statistics of their training dataset.
Therefore, if samples from these models are interpreted as an objective truth without fully considering the biases present in the original data, then they can perpetuate discrimination against minority groups.
In this work, we train on datasets that contain less sensitive data such as pictures of objects and music samples.
The methods we presented, however, could be used to model images of people or text from the internet which will contain biases and potentially harmful content that the model will then learn from and reproduce.
Great care must be taken when training these models on real world datasets and when deploying them so as to mitigate and prevent the harms that they can cause.