A Continuous Time Framework for Discrete Denoising Models
Anonymous Author(s)
AffiliationAddress email
Abstract
We provide the first complete continuous time framework for denoising diffusionmodels of discrete data.
This is achieved by formulating the forward noising processand corresponding reverse time generative process as Continuous Time Markov
Chains (CTMCs).
The model can be efficiently trained using a continuous timeversion of the ELBO.
We simulate the high dimensional CTMC using techniquesdeveloped in chemical physics and exploit our continuous time framework to derivehigh performance samplers that we show can outperform discrete time methodsfor discrete data.
The continuous time treatment also enables us to derive a noveltheoretical result bounding the error between the generated sample distribution andthe true data distribution.
Introduction
Diffusion/score-based/denoising models [1, 2, 3, 4] are a popular class of generative models thatachieve state-of-the-art sample quality with good coverage of the data distribution [5] all whilst usinga stable, non-adversarial, simple to implement training objective.
The general framework is to definea forward noising process that takes in data and gradually corrupts it until the data distribution istransformed into a simple distribution that is easy to sample.
The model then learns to reverse thisprocess 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 aremany problems for which the data we would like to model is discrete.
This occurs, for example, intext, segmentation maps, categorical features, discrete latent spaces, and the direct 8-bit representationof images.
Previous work has tried to realize the benefits of the denoising framework on discrete dataproblems, 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, workingin discrete time has notable drawbacks.
It generally forces the user to pick a partition of the process attraining 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 insteadlearns to denoise for any arbitrary time point in the process.
This complete specification of the reverseprocess enables much greater flexibility in defining the reverse sampling scheme.
For example, incontinuous state spaces, continuous time samplers that greatly reduce the sampling time have beendevised [14, 15, 16, 17] as well as ones that improve sample quality [4, 18].
The continuous timeinterpretation has also enabled the derivation of interesting theoretical properties such as error bounds[19] in continuous state spaces.
To allow these benefits to be exploited for discrete state spaces as well, we formulate a continuoustime 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 beused for efficient training of a parametric approximation to the true generative reverse process.
Toefficiently simulate the parametric reverse process, we leverage tau-leaping [20] and propose a novelpredictor-corrector type scheme that can be used to improve simulation accuracy.
The continuoustime framework allows us to derive a bound on the error between the true data distribution andthe 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 CIFAR-10dataset and monophonic music sequences.
Notably, we find our tau-leaping with predictor-correctorsampler can provide higher quality CIFAR10 samples than previous discrete time discrete stateapproaches, further closing the performance gap between when images are modeled as discrete dataor 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 finitecardinality S = |X| .
We assume x 0 ∼ p data ( x 0 ) for some discrete data distribution p data ( x 0 ) .
Wedefine a forward noising process that transforms p data ( x 0 ) to some distribution q K ( x K ) that closelyapproximates an easy to sample distribution p ref ( x K ) .
This is done by defining forward kernelsq k +1 | k ( x k +1 | x k ) that all admit p ref as a stationary distribution and mix reasonably quickly.
Forexample, 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 distributionover all states.
Other choices include: an absorbing state kernel—where for each state there is a smallprobability that it transitions to some absorbing state—or a discretized Gaussian kernel—where onlytransitions 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:
[Equation]
1 ) where q k | k +
[Equation]
x k
[Equation] [Equation]
+
[Equation]
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 couldsample 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 reversekernel, p θk |
k +1 .
This kernel is commonly defined through the analytic q k | k +1 , 0 distribution and aparametric ‘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 ) .
Notethat the faster the transitions mix, the more accurate this approximation becomes.
Approximatesamples 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 upperbound 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 thatq k | 0 ( x k | x 0 ) can be computed efficiently in a time independent of k .
With this, θ can be efficientlytrained by taking a random selection of terms from L DT in each minibatch and performing a stochasticgradient step.
Continuous Time Framework
3.1 Forward process and its time reversal
In continuous time, we imagine x k , x k +1 to be samples from a right continuous process separated byan infinitesimal time step ∆ t , i.e. samples at time t − ∆ t and t .
The process { x t } t ∈ [0 ,T ] is known asa Continuous Time
Markov Chain (CTMC).
We provide a short introduction to CTMCs in Appendix
A for completeness.
A CTMC can be defined through an initial distribution q 0 and a transition ratematrix R t ∈ R S × S .
This transition rate matrix is the continuous time analogue of the discrete timeforward kernel, q k +1 | k , because it defines the infinitesimal transition probability for the process
[Equation]
where R t (˜ x, x ) is the (˜ x, x ) entry in the transition rate matrix at time t .
Just as in discrete time, we create a forward noising process by defining R t such that: i) the forwardprocess mixes quickly towards an easy to sample (stationary) distribution, p ref , (e.g. uniform), ii) wecan analytically obtain q t |
( x t | x 0 ) distributions to enable efficient training (see Section 4.1 for howthis is done).
For the forward process starting at p data ( x 0 ) at time t = 0 , we denote the marginal attime 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 marginalq 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 tocalculate ˆ 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 distributionp data ( x 0 ) and terminal distribution q T ( x T ) , there exists a CTMC with initial distribution q T ( x T ) att
= T , terminal distribution p data ( x 0 ) at t = 0 and transition rate matrix ˆ R t that runs backwardsin 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]
x ̸ = ˜ x, where q t | 0 ( x | x 0 ) are the conditional jump 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 .
Whenx =
˜ 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 denoisingmodel, p θ
| t ( x 0 | x ) :
[Equation]
ˆ R θt ( x, x )
= − (cid:80) x
′ ̸
= x ˆ R θt ( x, x ′ ) as before.
When x ̸ = ˜ x , ˆ R θt has the same form as the discrete timeparametric reverse kernel, p θk | k +1 defined in eq (1) but with the forward kernel, q k +1 | k , replaced bythe forward rate, R t .
3.2 Continuous Time ELBO
In discrete time, θ is trained by minimizing the discrete time negative ELBO, L DT , formed fromthe forward and reverse processes.
We mirror this approach in continuous time by minimizing thecorresponding 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, 115 E p data
[Equation]
x 0 ) [ − log p θ 0 ( x 0 )] , is given by
[Equation]
where C is a constant independent of θ and
[Equation]
We can optimize this objective efficiently with stochastic gradient descent.
For a gradient update, wesample 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 .
r t (˜ x | x ) is a properdistribution, giving the probability of transitioning from x to ˜ x , given that we know a transition occursat time t .
Evaluating this objective naively requires two forward passes of p θ 0 | t
but, in practice, weneed only one by making a small approximation, see Appendix C.4.
Intuitively, ( x , ˜ x ) are a pair of states following the forward in time noising process.
Minimizing thesecond 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 directcomparison to L DT is given in Appendix C.1.
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 quicklytowards p ref , and ii) the q t | 0 ( x | x 0 ) distributions can be analytically obtained.
The Kolmogorovdifferential equation for the CTMC needs to be integrated to obtain q t | 0 ( x | x 0 ) .
This can be doneanalytically when R t and R t ′ commute for all t, t ′ , see
Appendix E. An easy way to meet thiscondition is to let R t
= β ( t ) R b where R b ∈ R S × S is a user-specified time independent base ratematrix 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.
The β schedule is guided by [3, 4], β ( t ) =
ab t log( b ) .
The hyperparameters a and b are selectedsuch that q T ( x T )
≈ p ref ( x T ) at the terminal time t = T while having a steady speed of ‘informationcorruption’ 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 140 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 thelinks 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 Spossibilities.
We now slightly redefine notation and say x 1: D ∈ X D , |X| =
S .
This results inrate matrices, R t and ˆ R θt , of size S D × S D .
It would be intractable to consider them in theirentirety for any reasonably sized S , D .
We avoid this problem simply by factorizing the forwardprocess such that each dimension propagates independently.
Since this is a continuous time processand each dimension’s forward process is independent of the others, the probability two or moredimensions transition at exactly the same time is zero.
Therefore, overall in the full dimensionalforward CTMC, each transition only ever involves a change in exactly one dimension.
For the timereversal CTMC, it will also be true that exactly one dimension changes in each transition.
Thismakes computation tractable as to simulate the next transition, instead of needing to calculate S Drate values, we need only calculate D × ( S − 1) values (for each dimension, the rates of transition to each of the other states).
Finally, we note that even though dimensions propagate independently inthe forward direction, they are not independent in the reverse direction because the starting pointsfor each dimension’s forward process are not independent for non factorized p data .
The followingproposition 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 | x 1: D s ) = (cid:81)
Dd =1 q t | s ( x dt | x ds ) , t > s ,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 witha 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 processfrom distribution p ref ( x 1: D T ) at time t
= T back to t = 0 will produce approximate samples fromp data ( x 1: D 0 ) .
Simulating this process exactly (e.g. by using Gillespie’s Algorithm [21, 22, 23]) wouldbe inefficient for large D , because we would need to step through each transition individually and soonly one dimension would change for each simulation step.
Instead, we use tau-leaping [20, 23], a very popular approximate simulation method developed inchemical physics.
Rather than step back through time one transition to the next, tau-leaping leapsfrom t to t − τ and applies all transitions that occurred in [ t, t − τ ] simultaneously.
To make aleap, 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 reacht − τ , such that x 1: D t remains constant in [ t, t − τ ] .
Assuming ˆ R θ 1: D t and x 1: D t remain constant, thenumber 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 simultaneouslyi.e.
x 1: D t − τ = x 1: D t + (cid:80) i P i (˜ x 1: D i − x 1: D t ) 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 .
Using our knowledge of ˆ R θ 1: D t , we can further unpack this update.
Namely, ˆ R θ 1: 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 (ratesfor multi-dimensional changes are zero).
Explicitly summing over these options we get x 1: D t − τ =x 1: D t + (cid:80) Dd =1 (cid:80) Ss =1 \ x dt P ds ( 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 dscan be non-zero, we see that tau-leaping allows x 1: D t to change in multiple dimensions in a singlestep.
When our discrete data has ordinal structure (e.g. Section 6.2) our mapping to Z is not arbitraryand making multiple jumps within the same dimension ( (cid:80) Ss =1 \ x dt P ds > 1 ) is meaningful.
In thenon-ordinal/categorical case (e.g. Section 6.3) the mapping to Z is arbitrary and so, although takingsimultaneous jumps in different dimensions is meaningful, taking multiple jumps within the samedimension 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 andalso 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 changesper step.
Increasing τ and thus the average number of dimensions changing per step gives us anatural way to modulate the ‘autoregressiveness’ of the model and trade sample quality with compute(Figure 2 right).
We refer to our method of using tau-leaping to simulate the reverse
CTMC asτ LDR
(tau-leaping denoising reversal) and give a precise algorithmic description in Algorithm 1 in
Appendix F.
4.4 Predictor-Corrector
During approximate reverse sampling, we aim for the marginal distribution of samples at time t to beclose to q t ( x t ) (the marginal at time t of the true CTMC).
The continuous time framework allowsus 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 ct which has q t ( x t ) as its stationary distribution.
Thecorrector steps bring the distribution of samples at time t closer to the desired q t ( x t ) marginal.
R ct iseasy to calculate as stated below
Proposition 4.
For a forward CTMC with marginals { q t ( x t ) } t ∈ [0 ,T ] , forward rate matrix, R t , andcorresponding reverse CTMC with rate matrix ˆ R t , the rate R ct =
R t + ˆ R t has q t ( x t ) as its stationarydistribution.
In practice, we approximate R ct by replacing ˆ R t with ˆ R θt .
This is directly analogous to PredictorCorrector sampling strategies in continuous state spaces [4] that predict by integrating the reverse
SDE and correct with score-based Markov chain Monte Carlo steps.
4.5 Error Bound
Our continuous time framework also allows us to provide a novel theoretical bound on the errorbetween the true data distribution and the sample distribution generated via tau-leaping (withoutpredictor-corrector steps), in terms of the error in our approximation of the reverse rate and the mixingof the forward noising process.
We assume we have a time-homogeneous rate matrix R t on X , from which we construct the factorizedrate 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 inp data with rate matrix R 1: D t as above.
Suppose that ˆ R θ 1: D t is an approximation to the reverse ratematrix and let ( y k ) k =0 , 1 ,...,N be a tau-leaping approximation to the reverse dynamics with maximum 229 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 > 0depending on X and R t but not D such that, if L ( y 0 ) denotes the law of y 0 , we have the totalvariation bound
[Equation]
The first term of the above bound captures the error introduced by our approximation of the reverserate ˆ 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 forwardchain, 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 specificallyinterested in applying tau-leaping to high dimensional problems where we make transitions in differentdimensions simultaneously in a single time step.
The bound grows at worst quadratically in thedimension, versus e.g. exponentially.
This shows we do not need to make τ impractically small inhigh dimensions.
The assumptions listed in Appendix B.5 hold approximately for tau-leaping in practice when we usespatially biased rates for ordinal data such that jump sizes are small or uniform rates for non-ordinaldata 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.
[Figure]
Related Work
[Figure] The application of denoising models to discrete data was first described in [1] using a binomialdiffusion process for a binary dataset.
Each reverse kernel p θk | k +1 was directly parameterized withoutusing a denoising model p θ 0 | k .
In [25] an approach for discrete categorical data was suggested using auniform 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 segmentationmaps 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 forwardkernels.
Whilst a link to continuous time for the forward process is mentioned in [8], all of theseapproaches train and sample in discrete time.
We show in Appendix G that this involves makingan implicit approximation for multi-dimensional data.
We extend this line of work by training andsampling in continuous time.
Other works also operate in discrete space but less rigidly follow the diffusion framework.
Acorruption process tailored to text is proposed in [12], whereby token deletion and insertion is alsoincorporated.
[26] also focus on text, creating a generative reverse chain that repeatedly appliesthe same denoising kernel.
The corruption distribution is also defined through the same denoisingkernel to reduce distribution shift between training and sampling.
In [7], a more standard maskingbased forward process is used but the reversal is interpreted from an order agnostic autoregressiveperspective.
Finally, in [11], the discrete latent space of a VQVAE is modeled by quantizing anunderlying 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 a target datadistribution and that tau-leaping can accurately simulate the reverse CTMC.
To do this, we create adataset formed of 2d samples of a state space of 32 arranged such that the histogram of the trainingdataset forms a ‘ τ ’ shape.
We train a denoising model using the L CT objective with p θ 0 | t parameterizedthrough a residual MLP (full details in Appendix H.1).
We then sample the reverse process defined bythis denoising model using an exact method (up to needing to numerically integrate the reverse rate)and tau-leaping.
Figure 1 top-right shows the marginals during reverse simulation with τ = 0 .
004and
we indeed produce samples from the entire support of p data .
Furthermore, we find that withsufficiently small τ , we can match the fidelity of exact simulation of the reverse CTMC (Figure 1left).
The value of τ dictates the number of network evaluations in the reverse process according to NFE = T/τ . In all experiments we use T = 1 .
[Table] [Table]
NFE
[Figure] [Figure]
[Figure] [Figure]
6.2 Image Modeling
We now demonstrate that our continuous time framework gives us improved generative modelingperformance versus operating in discrete time.
We show this on the CIFAR-10 image dataset.
Imagesare 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 adiscretization function at the end of the generative process [3] or adding uniform noise to the data.
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 aspatially biased rate matrix and train with an augmented L CT loss including direct p θ
| t supervision,full experimental details are in Appendix H.2.
Figure 2 left shows randomly generated unconditional CIFAR10 samples from the model and wereport sample quality metrics in Table 1.
We see that our method ( τ LDR- 0 ) with 0 corrector steps hasbetter Inception Score but worse FID than the D3PM discrete time method.
However, our τ LDR- 10method 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 closesthe performance gap with methods modeling images as continuous data.
The derivation of thecorrector 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 thereverse sampling procedure.
We show how FID score varies with number of p θ 0 | t evaluations forτ LDR- { 0 , 3 , 10 } in Figure 2 right.
The optimum number of corrector steps depends on the samplingbudget, with lower numbers of corrector steps being optimal for tighter budgets.
This is due to theincreased τ required to maintain a fixed budget when we use a larger number of corrector steps.
[Table]
[Figure]
[Table]
[Figure]
6.3 Monophonic Music
[Figure]
In this experiment, we demonstrate our continuous time model improves generation quality onnon-ordinal/categorical discrete data.
We model songs from the Lakh pianoroll dataset [28, 29].
We select all monophonic sequences from the dataset such that at each of the 256 time steps eitherone from 128 notes is played or it is a rest.
Therefore, our data has state space size S = 129 anddimension D = 256 .
We scramble the ordering of the state space when mapping to Z to destroyany ordinal structure.
We parameterize p θ 0 | t with a transformer architecture [30] and train using aconditional 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 detailsare given in Appendix H.3.
Conditional completions of unseen test songs are shown in Figure 3.
We quantify sample quality in Table 2.
We use two metrics: the Hellinger distance between thegenerated and ground truth notes and the proportion of outlier notes in the generations but not inthe ground truth.
Using our method, we compare between a birth/death and uniform forward ratematrix R t .
The birth/death rate is only non-zero for adjacent states whereas the uniform rate allowstransitions between arbitrary states which is more appropriate for the categorical case thus givingimproved 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 forcategorical 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 toefficiently sample the generative process with tau-leaping and provided a bound on the error of thegenerated samples.
On discrete data problems, we found our predictor-corrector sampler improvedsample quality versus discrete time methods.
Regarding limitations, our model requires many modelevaluations to produce a sample.
Our work has opened the door to applying the work improvingsampling speed on continuous data [14, 15, 16, 17, 31] to discrete data problems too.
Modelingperformance on images also is slightly behind continuous state space models, we hope this gap isfurther closed with bespoke discrete state architectures and corruption process tuning.
Finally, wenote that the ELBO values for the discrete time model on CIFAR10 are better than for our method.
Inthis work, we focused on sample quality rather than using our model to give data likelihoods e.g. forcompression downstream tasks.