Score-based Continuous-time Discrete Diffusion Models
Authors: Haoran Sun
Lijun Yu
Bo Dai
Dale Schuurmans
Hanjun Dai
Venue: ICLR 2023 poster (2023)
Keywords: discrete space diffusion
discrete score matching
continuous-time diffusion

Abstract: Score-based modeling through stochastic differential equations (SDEs) has provided a new perspective on diffusion models, and demonstrated superior performance on continuous data. However, the gradient of the log-likelihood function, \ie, the score function, is not properly defined for discrete spaces. This makes it non-trivial to adapt SDE with score functions to categorical data. In this paper, we extend diffusion models to discrete variables by introducing a stochastic jump process where the reverse process denoises via a continuous-time Markov chain. This formulation admits an analytical simulation during backward sampling. To learn the reverse process, we extend score matching to general categorical data, and show that an unbiased estimator can be obtained via simple matching of the conditional marginal distributions. We demonstrate the effectiveness of the proposed method on a set of synthetic and real-world music and image benchmarks.

Published as a conference paper at ICLR 2023
SCORE-BASED CONTINUOUS-TIME DISCRETE DIFFU- SION MODELS
Haoran Sun"
Lijun Yu
Bo Dai
Georgia Tech
Carnegie Mellon University lijun @ emu.edu
Google Research; Georgia Tech bodai @ google..com
hsun349@ gateehh.uu
Dale Schuurmans
Hanjun Dai
Google Research; University of Alberta schuurmans @ google.com
Google Research
hadai @ google.com
ABSTRACT
Score-based modeling through stochastic differential equations (SDEs) has pro- vided a new perspective on diffusion models, and demonstrated superior perfor- mance on continuous data. However, the gradient of the log-likelihood function, i.e., the score function, is not properly defined for discrete spaces. This makes it con-trivial to adapt the score-based modeling to categorical data. In this paper, we extend diffusion models to discrete variables by introducing stochastic jump process where the reverse process denoises via continuous-time Markov chain. This formulation admits an analytical simulation during backward sampling. To learn the reverse process, we extend score matching to general categorical data, and show that an unbiased estimator can be obtained via simple matching of the conditional marginal distributions. We demonstrate the effectiveness of the pro- posed method on set of synthetic and real-world music and image benchmarks.
INTRODUCTION
Diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020) have emerged as an important tech- nique for data distribution modeling, where data-coruppinng forward process coupled with denoising reverse process to simulate a diffusion relationship between the data distribution and an uninformed source. Such models admit stable learning procedures and have demonstrated superior performance on continuous data modeling in challenging scenarios (Dhariwal & Nichol, 2021), lead- ing to rapidly increasing popularity. Song et al. (2020) established a stochastic differential equation view of diffusion models by forming the limit of finer corruption and denoising steps in the forward and backward processes, rendering a continuum of distributions. This perspective has provided a unified framework under a new score-bas score-based learning objective, and inspired a variety of simulation methods for efficient sampling and inference (De Bortoli et al., 2021; Zhang & Chen, 2022).
Given the advantages of diffuuson models in terms of flexibility, learning tractability, and sampling, there have been several attempts to extend the approach to discrete data. Recent attempts have inves- tigated alternative corruption operations for discrete data in the forward process, yielding promising results (Hoogeboom et al., 202 1b;a; Austin et al., 2021). However, these extensions still execute a finite sequence of corruption and restoration steps, and remain restricted a fixed reverse sampling strategy that can be sub-optimal To overcome this limitation, we investigate whether continuous- time discrete diffusion formulation might admit more effective estimation and generation.
Such an extension is highly non-trivial however. The continuous-time diffusion framework is based on stochastic differential equation (SDE) with respect the score function, which itself is the gradient of the log- ligelioel with respect to a continuous variable. Although this can be used to characterize a continuum of infinitesimally evolving distributions over a continuous space, such formulation no longer exists for discrete variables, since the gradient of the log- likeliiel does not exist with respect to a discrete variable. Recently, Campbell et al. (2022) made significant progress in
Work done during an intersship Gongle.

---

Published as a conference paper at ICLR 2023
closing this gap by proposing discrete data distribution that evolves in a continuous- Markov chain. With this formulation they were able to approximate maximum likelihood training with an evidence lower bound (ELBO) surrogate and using predictor-correet generalize some of the benefits of continuous-time modeling to a discrete space.
sampling scheme to
However, limitation of this previous work is the reliance an ELBO approximation to the MLE when is known that score- based learning yields superior estimation quality (when it can be ap- plied) due to its unbiasedness (Song et al., 2020). Of course, developing an analog of the score matching for discrete spaces is non-trivial due to non-differentiabiliyy Nevertheless, a score func- tion for a discrete random variable cannot be arbitrary: whenever two score functions match over a space, their corresponding distributions should also match; the score function should characterize the direction of infinitesimal evolution of a discrete distribution; and finally the score should enable tractable score-based estimation. In this paper, we investigate the design of such score functions for discrete spaces achieving these desired properties, and provide corresponding learning and sampling methods. In addition, we complete the agenda for continuous time diffusion in discrete spaces by formulating a coherent SDE in terms of stochastic jump processes. The main contributions are:
We extend the definition of score functions to generic categorical discrete variables, and derive a continuous-tim discrete diffusion model via continuous time Markov chain in Section 3.1;
We derive a score-based objective called categorical ratio matching for estimating the proposed model in Section 3.2, which can be tractably optimized, showing that a previous proposal for binary discrete data (Hyvärinen, 2007) can be obtained as a special case;
In Section 4, we develop numerical simulation technique for reverse sampling, then provide an analytical sampling method based on implicit modeling of the conditional marginals:
We discuss three architectural choices and present a novel "hollow" neural model in Section 5;
We evaluate the proposed SDDM on set of synthetic and real-world music and image bench- marks, achieving promising results in Section 6.
BACKGROUND
Diffusion models (Sohl-Dickstein et al., 2015) are characterized by forward Markov process that transforms an observation Zo Adata (zo) to reference distribution 2T ~ qr and and backward process, which also Markovian, that recovers the data distribution from Tr- Specifically, the forward process is defined through simple corruption operations, 4++1jt- For continuous data the corruption kernel usually adds Gaussian noise (Ho et al., 2020); whereas for discrete data, where 3o E X with finite cardinality |.\, the corruption kernel can be uniform, discrete Gaussian, or some other choice (Austin et al., 2021; Johnson et al., 2021). Given the corruption kernel, after Ttepps, the forward process forms joint distribution,
90:T( daate (zo) 1TTt0 9--1je (1++1
The backward process can be derived from the joint distribution via Bayes' rule,
(o:T(F)))) = (xT)) IT==0 lt 8+1 (xe (e))), 2c/c+1 (2x(20441)
( (wn))
(2)
(==1(
where q: (x+) denotes the marginal distribution and the prior gr is usually a simple distribution. Typically the backward kernel 9oja-+-1 (#4)) is intractable; thus is usually parameterized by a
neural network, denoted Pftez and learned via ELBO (Sohl-Dickstein et al., 2015; Ho et al., 2020;
Austin et al., 2021) or score-matching (Song & Ermon, 2019). Due to the structure of the joint distribution (1) and (2), the ELBO admits a particularly simple formulation as
[DKL. (911+1,0l pi(++)] Equio [Log p (colez1)
L.b -E DKL. (erioller)
Excro
(3)
which is applicable for both continuous and discrete diffusion model learning. For continuous vari- ables with Gaussian corruptions ( (xz) = N' (œ141i /1- Bit1 Bt+14), and a backward 1 kernel pa (#x( = -V (I: V1-8.+1 (ret1 + B4417 (#++1)) . such that r is learned as a neural network and Be is predefined variance schedule, the ELBO (3) can be rewritten as T 1
Lab =
(1 ar)) Brzsce Epn, (x/()) [1:2 (x') - - log (a' [1]=))]]]]]
(4)
0=0
2

---

Published as a conference paper at ICLR 2023
where de IT (1 Be). The Equation 4 is highly related to score- more-matching (Hyvärinen & Dayan, 2005) with a score function defined as Ve log Pr+ (++1)((F2-1((2)))
Reverse sampling in discrete-ting models is somewhat restricted by the forward process. Therefore, continuous- tinnuuous-tmeeeeeo diffusion models have been constructed with indexed from [0. T]. Song et al. (2020) develop SDE perspective and define diffusion process with respect the score
dr =f (z, dt + (())dw,
forward SDE,
(5)
dr = [ (z, of (1) (ogd) + g(()du.
reverse SDE,
(6)
where and w standard Wiener processes, f (xx t) is vector-valud drift function, and g a scalar-valued diffusion coefficient. A score-based learning method can be easily derived as an extension of (4). However, the score function v. log Pr(2:) is not defined for discrete variable, and the SDE above no longer suffices to characterize a continuous-time extension for discrete diffusion.
3 CONTINUOUS TIME DISCRETE SCORE MATCHING
Although Campbell et al. (2022) bypass the score function a continuous-ting time extension, by lever- aging stochastic process view of the ELBO approximation, we will instead focus score-based extension for continuous-time discrete diffusion.
3.1 CONTINUOUS TIME MODELING
Consider the finite discrete state space = CD, where C= (1,,,..,, is code book. To generalize score matching from continuous space R" to discrete space X' we first model the forward process as a continuous time Markov chain {Xe]rejorry whose transition probability is characterized by rate matrices Q: E x In particular, if we let q denote the distribution for the forward process X,, the transition probability will satisfy the Kolmogorov forward equation:
d de 4t/s (1 xr,) ZEX 9t/. x, (z,
<t
If the forward process starts at the target distribution qu Tdata, the marginal at time t has the form:
(<(xt)) = Jx Taata (2o) (:jo)))j((( (zg|=o))do
(8)
By properly choosing the rate matrices Qt, we can achieve a final distribution close known tractable distribution 9r x Teet Then, the reverse time process X, X can be expressed as generative process from the reference distribution Tea to the target distribution Tdata (Anderson & Rhodes, 1983; Van Handel, 2007).
Proposition 3.1. The reverse time process X, of the continuous time Markov chain X₃ is also Markov process, whose transition probabilities 4s for satisfy:
(0)e(2,((1e)) (x(x)) (m)))) 4 (1e(2rs) <t
Since considering finite space X. the reverse time process X is uniquely determined by its rate matrices, which we denote as R. Using the transition Equation 9, we have the following.
Proposition 3.2. For a continuous time Markov chain {Xtheeoo. with distribution g and rate matrices Qt, the rate matrices Re for the reverse process satisfy:
9c(y) qe (x)
Ra(2, 3)
Q2(
(10)
Equation 10 provides closed form expression for the reverse time rate R. Therefore, once we know the ratio qi (x)), we can obtain the generative flow towards Waat-
3.2 CATEGORICAL RATIO MATCHING
In general, the ratio qe (y)) qq (x) in Equation 10 is intractable. This ratio behaves analogously to the score function V log (xr) in Equation 6. Such connection inspires learning the ratio ge (v) ////// (x:) in
3

---

Published as a conference paper at ICLR 2023
similar manner to learning the score function for diffusion models in continuous spaces. For binary discrete variable, Hyvärinen (2007) proposed to match the ratio π (X\, X = -///////// X = 1- c) as an extension for score matching. In this work, we consider the more general categorical case, where neither the score function V log 7t (I) nor the binary ratio is well defined. The generalization relies on the singleton conditional distribution:
a (X c/a: ( = T = D CEc x =d)
(11)
yielding a score function categorical space we seek to match. The sufficiency of matching Equa- tion 11 is guaranteed by the property that the joint distribution is completely determined by its singleton conditional distributions (Brook, 1964; Lyu, 2012).
Proposition 3.3. Consider random variables X = =X1,...,XD)) E X, and two probability dis- tributions n1, nT. We have Tι T2, if and only if their conditional distributions equal Ti(X = x² x = (x x (d) ), for any x E X and d = l........ D.
Going back the ratio q: (v(y/q (I) in reverse time rate Equation 10, Proposition 3.3 tells can recover the probability ratio by employing a time dependent neural network pr(50) to match the conditional distribution.
(1 (x(x = (1) (((dd
Pc(X, = -
Pr (x (, (x(
9x(X; = ==d/add = (12)
To train (e(:; f), we minimize the expected cross entropy with respect to the conditional distributions along the forward process:
D Σ
0" arg min
Σ CEC
(x(x; = c/o, log =cjx (00 (13)
ZIEX
where the loss is minimized if pt (X;" |x) (x = qe |x] This loss function matches the ratio in Equation 11, hence name categorical ratio matching generalization of Hyvärinen (2007). Since the spirit of this loss function is the same as the score matching in continuous space, we also interchangeably call this method as categorical score matching. Equation 13, the expectation over qe (-) can be efficiently estimated via Monte Carlo sampling. However, the conditional distribution qe (Xx = < is intractable, which brings difficulty to the training. Fortunately, using the property of conditional distribution, we can avoid computing this intractable term in the loss function.
Proposition 3.4. The categorical ratio matching loss function in Equation 13 can be simplified as:
arg min /' EmEx ge(ee)
D
Σ
log pe (x = at 10 ;) dt.
(14)
Using this surprisingly neat loss function in Equation 14, we can efficiently learn the conditional >d: ()) determines a reverse process, and we use p(:; 0) to denote its joint distribution in order to distinguish with the true reverse process. We will sometimes drop the 0 if it does not create ambiguity.
distribution. The learned Pe |x
CONTINUOUS TIME DISCRETE SAMPLING
We next develop and efficient sampling scheme for the proposed diffusion model.
4.1 CONTINUOUS TIME SIMULATION FOR FORWARD PROCESS
In continuous time Markov chain, the transition matrix (1)((-1)) from time time t in the forward process can be obtained by solving the ODE in Equation 7. For general rate matrices
Q E || solving the ODE is intractable. Therefore, we follow standard practice in diffusion models (Austin et al., 2021; Campbell et al., 2022) and approximate the forward process by factor- izing X, = (X! (x,,,..,? where each sub- process X propagates independently. Furthermore, we also define the sub-rate matrix Qf Q,((t) in terms of fixed base rate Q PAP-1 RCx
4

---

Published as a conference paper at ICLR 2023
and time schedule function B(t); see design of B(t) in Appendix C.1. In this way, the sub- tobbtrannition matrix in each dimension can then be easily computed (Campbell et al., 2022) as:
f P exp A SA ((r)dr) P-1
(15)
In particular, use uniform stationary base rate Q 11" CI, which is a natural choice for categorical discrete distributions as it admits a uniform distribution as its stationary distribution (Hoogeboom et al., 2021b; Austin al., 2021). For simplicity of comparison we the same schedule function as in Camphell et al. (2022) for all the experiments, and in Appendix C.1 we discuss other possible choices.
4.2 DISCRETE TIME SAMPLING FOR REVERSE PROCESS
Even given the learned conditional distribution (e 0), simulating the reverse time process is much harder than simulating the forward process, as the reverse process cannot be factorized. For example, consider that are only different at the d-th site. The rate for the reversed process to jump to from at time is:
1= Qr(y,
Re y; 0) = px
(16)
Such jump rate depends on both the time and the value in the other dimensions zy \ Hence, we can not simulate each dimension in parallel, making exact simulation of the reverse time process X less efficient. Considering Ri (zz, 0) is already an approximation of Ri (Zt, y). we employ Euler's s method for parallel simulation. Specifically, given Ia at time t, we fix the rate Equation 16, then determine the transition probabilities for dimension d at time t - according to:
- - (
R2 (I,, X = c; A), 11- Ex (2,,, - cc;),
c x c c=
P (X = <lx)
(17)
We clip the quantities to ensure all probabilities are non-negative. Then, we collect new value from each dimension to obtain a new state yt-c, which has the factorized probability:
ye-elar;; 0) I2_ (X 1=)
(18)
In this way, we can update all dimensions of æ, in parallel, making the simulation of the reverse time process much more efficient.
4.3 ANALYTICAL SAMPLING FOR REVERSE PROCESS
The Euler's method mentioned above assumes the rate matrix R2 is fixed during T E [t e, t]. Such an approximation makes the samples x, deviate from the correct time slice marginal (-) especially when the time step is large. To mitigate this approximation error, diffusion models usually intro- duce a corrector to improve sample quality (Song et al., 2020; Campbell et al., 2022), however this increases computational cost. Instead, we propose alternative approach that leverages implicit modeling of pt {x) 0). Specifically, for distribution q, we have:
q. (x,4 /ol)
E.g -ojo (0)((2))12 Val (X;])
(19)
Since Q:jo is tractable, computing Equation 19 is efficient once know 4o (X. 1d). Hence, can replace the explicit approximation in Equation 12 by the following implicit approximation:
Poje (x6 I 0) = 20j: (x,1(²)
Pe (x 0) = Σnd Poje(ay x Vd
(x)4 (20)
Equation 20 provides tractable transformation from Pite (X 0) to p. (X +4;0)
0). Hence, we Vol 8). To
can continue using the categorical ratio matching loss in Equation 14 to train (x conduct backward sampling via this new parameterization, we consider the true reverse process:
gt q (Xd (Xt_ 1zx) = Σ
2oje (aff(f))))))) (x- xx)
(21)
mg
(aa ang. /dl Dg(agla \d g( log xz_))q(
(22)
x
Σ
20je (= ) (x
(23)
5

---

Published as a conference paper at ICLR 2023
By substituting pol (X i Vd 0), we have:
Pt-ejo (X 1
Σng (olo - ejjo (x
(24)
Thus, we obtain an analytical expression for the reverse process sampling that avoids the simulation error in Euler's method. Hence, we refer to this method as analytical sampling.
5 PARAMETERIZATION
The key the simplicity of the objective function in Equation 14 is the special structure of the conditional marginal distributions pr (x(;(lxz" [= 8) and ( (x v th dimension do not depend on the current value æ at time t. However, this elegance in the objective brings additional challenges to design an efficient and flexible neural network parameterization. A key design constraint is that the predictions at dimension d must not depend on æt otherwise the information leak would make Equation 14 trivial solve. The prediction can depend on the other coordinates in however. We offer three alternative architectures that are able to satisfy this constraint but incur different trade-offs between flexibility and computational cost. With out loss of generality we consider the parameterization of PepP (X,]2) [1; ;)), however the same designs can be Ad 0).
such that the marginals d-
directly applied to the parameterization of Poje (X8 my
5.1 ENERGY BASED MODELS
The most general and flexible form of parameterization an Energy based model (EBM), where an arbitrary neural network fo (z, f) CD x R + R can be used to specify the energy of any sample. In this case the conditional marginals can be modeled as:
Px(X,' = cla) : exp (- (-
c,a,
E exp (c, 2.))) (25)
where we overload the notation to use [c, ( to denote le?.x!...... X = C, xt ...] such that only the d-th value is replaced with c. This modeling flexibility however comes at a high computa- tional cost. Specifically, to evaluate I[- P+ (x,(
one needs O(D x C.) rounds of evaluation of fo. which is computationally prohibitive when modeling high dimensional data with deep neural parameterization of fa. Nevertheless, this connection between EBMs and diffusion provides another way to learn time-dependent EBMs, which might be of independent interest.
5.2 MASKED MODELS
To alleviate the computation overhead of EBMs while preserving flexibility, masked model is natural choice. Specifically, let a masking function mz(z)- = [x ie....., ....,."' zdd1 .MASK, xd+1 ...., .."..."! x!
replace the d-th dimension of given special mask token MASK. Then one can formulate the following conditional parameterization:
Pe (e(X m Ad 0) = Softmax (m(d).))) where fo(æ, f) {C (CUMAN MASK) x R
(26)
Here f can still be a general neural network with only a minor requirement on handling then masked token. Overall this approaches requires (DD) rounds of evaluation of fg. Since discrete data, one can further reduce D at the cost of increasing vocabulary size C, to further reduce the rounds of feed-forward evaluations of So-
dealing with
5.3 HOLLOW TRANSFORMERS
Even though the two parameterizations above allow flexibility in the neural network design, they require a number of feed-forward evaluations that scales in the dimensionality and/or the vocabu- lary size of discrete data. As an alternative, we propose Transformer variant that requires only ((1) feed-forward evaluations. The key constraint is that the diagonals of the Jacobian matrix of (..X/ 0) be zero for any input Tt.
Many techniques have been developed to guarantee this property, including autoregressive mask- ing (Germain et al., 2015; Vaswani al., 2017), which creates triangular Jacobian for multi-layer networks, or hollow masking (Chen & Duvenaud, 2019), which considers the full context but only permits single layer of dense interaction between dimensions.

---

Published as a conference paper at ICLR 2023
P( P(
P(xabo)))
Jacobian matrix (X xx)
Here, capture the full context for each dimension with expressive deep neural architectures,
0xr
introduce the hollow Transformer, a model that runs two autoregressive Transformers, in each direction, for L-layers; see Fig- The hollow Jacobian matrix is obtained via the summation of up- per and lower triangular Jacobians the full context for each position is ob- tained. We add one additional Trans- former layer at the top, where the query vector comes from the two embedding directions of the corresponding dimen- sion, and attention on the key and value vectors are conducted jointly with the two directions. For clarity we omit de- tails of each Transformer layer we use the standard architecture (Vaswani al.. 2017). Overall, this specially de- signed architecture avoids any dependence on dimensionality or vocabulary size in terms of the number of feed-forward evaluations required, while being able to leverage the expressiveness of multi-layer Transformers.
K2 - V.
X1
34
V
Kg
Transformer
Transformer Layer
Laper
33
Transformet
Transformet Layer
Loyer
X
X2
Figure 1: Diagram illustration of the Hollow Transformer.
EXPERIMENTS
We present empirical evaluation of the proposed diffusion approach on synthetic data, CIFAR1O dataset and the monophonic music dataset. The primary goal is to compare the effectiveness of the proposed categorical ratio matching with the alternative parameterizations presented in Section 5. We use the time duration 7 = and the uniform forward process in all cases. Experiments are conducted on machines equipped with TPU-v4 chips. For more details please refer Appendix C.
6.1 DEEP EBMS ON SYNTHETIC DATA
As discussed in Section 5.1, one can leverage the EBMs for modeling conditional marginals. Here we verify this approach using synthetic dataset from the discrete EBM learning community. Fol- lowing prior works (Dai et al., 2020; Zhang et al., 2022), use seven different distributions of 32-dimensional binary discrete data to evaluate different approaches. Each discrete distribution is converted from 2-D continuous point (Ir, y) E B2 by quantizing both and y with 16-bit repre- sentations using a Gray code (Gray, 1953).
We parameterize the energy function fo(2, t) using the same 3-layer MLP as used in prior works (Dai et al., 2020; Zhang et al., 2022), with one minor change to directly add sinusoidal embedding of into each hidden layer before activation. The uniform rate constant is set to 1.0 and we use time resolution of le for simulation. To measure sample quality, we follow (Zhang al., 2022) and compare 4,000 generated samples to true data samples using exponential Hamming MMD (Gretton et al., 2012), repeated 10 times to report average results in Table 1. As baselines we include PCD (Tieleman, 2008), ALOE (Dai al., 2020) with larger network for dual parameteriza- tion, and GF-GFN (Zhang et al., 2022). Here we find that the proposed continuous-ting ratio matching approach is able to successfully learn good EBM, yiclding an MMD value that is consistently lower than the baseline methods. We also visualize the distribution obtained via SDDM in Figure 2 via reverse mapping of discrete Gray codes into 2-D points in continuous space. These plots show similar distribution to the ground truth data (see Appendix C.2 for more information).
categorical
Ablation study: we further provide ablation studies on three different parameterizations proposed in Section and two different sampling methods. For the full results please see Appendix C.2.1. Overall when 3-layer transformers are used in these three parameterizations, the performance are comparable each other, which shows that the masked model and hollow model can achieve better quality-speed trade-offss than EBMs. We will revisit the comparison among samplers in next section.
6.2 IMAGE MODELING ON CIFARIO

---

Published as a conference paper at ICLR 2023
O #
2spirals
8gaussians
circles
moons
pinwheel
swissroll
chekcerboard
Figure
Visualization of sampled discrete binary data in 2D space via decoding of Gray codes.
Table 1: Quality of generated binary samples from the learned EBMs, in terms of MMD with
exponential Hamming kernel using bandwidth=0.1 (in units of 1 x 10-4, the lower the better).
2spirals
Sgaussians
circles
moons
pinwheel
swissroll
checkerbourd
PCD (Tieleman, 2008)
2.160
0.954
0.188
0.962
0.505
1.382
2.831
ALOE+ (Dai et al., 2020)
0.149
0.078
0.636
0.516
1.746
0.718
12.138
EB-GEN (/hang ct al., 2022)
0.583
0.531
0.305
0.121
0.492
0.274
1,2016
SDDM (this paper)
0.096
0.040
0.086
0.150
0.243
0.225
0.162
Next we evaluate image generation quality using the CIFAR10 dataset. In this case, each raw image
Analytical
has shape (32, 32, 3) with 256 choices per each pixel value. We compare sev- eral diffusion models from the literature, with the main
Fwd Fuler
goal of understanding per- formance in the different spaces with either continu- ous or discrete time formulations. We primarily evaluate the performance of the masked model for- mulation presented in Section 5.2, and report the commonly used Inception Score (IS) and Fréchet Inception Distance (FID) against the training set. As the image pixels are naturally discretized continuous values, instead evaluate the proposed method in the quantization hash-code space via pretrained VQGAN (Esser et al., 2021). The main purpose of doing this is to 1) evaluate score matching complex categorical discrete data; and 2) alleviate the burden of the Transformers mod- eling long sequences. The resulting data lies in 64-dimensional categorical space with vocabulary size [C]=512. Note that VQGAN is a lossy compressor where the reconstructed images from the de- coder obtain IS=9,67 and FID=9.05 9.0,, which is the upper- bound of categorical diffusion built on top of this vector quantizer (VQ). Here we simply use 12-layer BERT-base architecture parameterize the fo t). See Appendix C.3 for more details about the VQ and architecture used.
Figure 3: Visualization of reverse sampling with different samplers.
From Table 2 we can see that, dealing with images harder than an ordinal discrete space, as it loses ordinal structure as prior. The proposed approach able to approach the performance limit of VQ based categorical discrete diffusion, and improves the performance of D3PM and rLDR in the same VQ space with the same parameterization. In Figure? we visualize the reverse sampling procedure (unevenly sampled time steps) using two proposed samplers, where the analytical sampler achieves reasonable quality in fewer steps.
categorical discrete space generally much
To quantitatively show why the continuous time modeling would achieve better flcxibility, we re- port the image quality with different number of sampling steps in Table 3. We can see as D3PM was trained with 1,000 steps, the performance drops quickly with fewer steps. Our continuous-time ver- sion can be more robust
with much fewer steps. Also the analytical sampler can be more robust when fewer steps are given, where the forward Euler may need corrector for better performance.
6.3 MONOPIONNIC MUSIC MODELING
Finally, we conduct a study on discrete music generation following the same settings used in Camp- bell et al. (2022). This benchmark is originally from the Lakh pianoroll dataset (Raffel, 2016; Dong et al., 2018), cleaned by removing some trivial music sequences. In the end there are 6,000 music sequences for training and 973 sequences for evaluation. Each sequence is of length 256, where the
8

---

Published as a conference paper at ICLR 2023
Table 2: Metrics on sample quality for different diffusion models CIFAR1O dataset. Here Incep- tion Score (IS) and Fréchet Inception Distance (FID) are compared. We follow the common practice to compare the 50,000 unconditionally sampled images from the model and the images from train- ing dataset. Approaches with representations in different state spaces are listed in separate sections. Yorr re-implementation in VQ space, where we tuned configurations and report the best result.
State space
Methods
IS 1
FID
DDPM (Ho et al., 2020)
9.46
3.17
Continuous state
NCSN (Song et al., 2020)
9.89
2.20
D3PM Gauss (Austin et al., 2021)
8.56
7.34
Ordinal discrete state
rLDR-0 (Campbell et al., 2022)
8.74
8.10
+LDR-10 (Campbell et al., 2022)
9.49
3.74
D3PM Uniform (Austin et al., 2021)
5.99
51.27
D3PM Absorbing (Austin et al., 2021)
6.78
30.97
Categorical discrete state
VQGAN (Esser et al., 2021) reconstruction
9.67
9.05
D3PM-VQ (Austin et al., 2021
8.85
16.47
7LDR- VQ (Campbell et al., 2022)*
5.71
40.06
SDDM-VQ (this paper)
8.98
12.23
Table FID/IS on CIFARIO with different sampling steps, using continoons/diserrce--iiee methods.
FID with
250
200
100
50
40
20
IS 1 with # steps
250
200
100
50
40
20
D3PM
32.61
38.56
62.15
84.72
90.1
102.3
D3PM
7.71
7.16
4.68
4.48
SDDM (Faler)
12.5
12.64
13.61
16.51
17.73
28
SDDM (Euler)
8.8
8.69
8.51
8.16
7.96
6.96
SDDM (Analytical)
13.39
13.29
15.96
SDDM (Analytical)
8.82
8.78
8.31
8.14
vocabulary size |x| 129 consists of 128 notes and rest. The vocabulary is scrambled there is no ordering information preserved. This creates a challenging categorical discrete data modeling task. Here we primarily compare against existing diffusion models in discrete spaces.
We use the simulation time e= for all methods. Since the di- mensionality is high, the conditional marginals Poje (X4(x) are parame- terized by the proposed hollow Trans- former (as discussed in Section 5.3) with 6 Transformer layers. Ap- pendix C.4 provides more details about the training and parameterizations where for each sequence, a prefix of length 32 is given as conditioning and the models are asked to generate the remaining 224 tokens. We follow the same evaluation protocol in Campbell et al. (2022) and report the Hellinger Distance and Proportion Outliers after 5 runs per each test case in Table 4. Overall we can see the proposed SDDM is able to consistently improve the two metrics.
Table Conditional music generation.
Methods
Hellinger Distance
Proportion of Outliers
+L.R-- Birth/Death
0.3928 - 0.0010
0.1316 + -0.0122
-LDR-0 Uniform
0,3765 = 0.0013
0.1106 = 0,0010
7LDR22 Uniform
0.3762 = 0.0015
0.1091 = 0.0014
D3PM Uniform
0.3839 = 0.0002
0.1137 + 0.0010
SDDM (this paper)
0.3736 0.0024
0.1093 0.0016
The evaluation was conducted on the held-out test set,
LIMITATION AND CONCLUSION
In this paper we presented new learning and sampling paradigm for continuous-ting time diffusion models in categorical discrete spaces. We developed an extension of score matching to categorical discrete variables, showing that the corresponding continuous-time score matching properly aligns the reverse process with the posterior of the forward processes. The new learning paradigm also nat- urally introduces new sampling algorithms. Despite the promising preliminary results, there are still limitations in the current treatment. The main bottleneck is the design of the conditional marginal parameterization, which requires non-trivial trade-offs between computational cost and flexibility of the architectures; score matching for general categorical discrete variables does not benefit from prior knowledge about ordinal discrete data; and finally unifying score matching between continu- ous and discrete spaces would be needed to handle data in mixed spaces. We believe this initiative sheds new light on score matching for discrete-spen space diffusion.

---

Published as a conference paper at ICLR 2023
REFERENCES
Brian DO Anderson and lan B Rhodes. Smoothing algorithms for nonlinear finiteddimensional systems, Stochastics: An International Journal of Probability and Stochastic Processes, 9(1-2): 139-165, 1983.
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34: 1998111993, 2021.
D Brook. On the distinction between the conditional probability and the joint probability approaches in the specification of neurest-neighhour systems. Biometrika, 51334):481-433 1964.
Andrew Campbell, Joe Benton, Valentin De Bortoli, Tom Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models. arXiv preprint arXiv:2205. 14987, 2022.
Huiwen Chang, Han Zhang. Lu Jiang, Ce Liu, and William T Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11315-1 11325, 2022.
Ricky TQ Chen and David K Duvenaud. Neural networks with cheap differential operators. Ad- vances in Neural Information Processing Systems, 32, 2019.
Max Cohen, Guillaume Quispe, Sylvain Le Corff, Charles Ollion, and Eric Moulines. Diffusion bridges vector quantized variational autoencoders. arXiv preprint 20w:2202204895, 2022.
Hanjun Dai, Rishabh Singh, Dai, Charles Sutton, and Dale Schuurmans. Learning discrete energy-based models via auxiliary-varian variable local exploration. Advances in Neural Information Processing Systems, 13:4034-11 2020.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34, 2021.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. 248-255. Icee, 2009.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780-877 2021.
Hao-Wen Dong, Wen-Yi Yi Hsiao, Li-Chia Yang, and Yi-Hsuan Yang. Musegan: Multi-track quential generative adversarial networks for symbolic music generation and accompaniment. In Proceedings of the AAAI Conference Artificial Intelligence, volume 32, 2018
Patrick Esser, Robin Rombach, and Bjom Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recogni- fion, pp.
2021.
Mathieu Germain, Karol Gregor, lain Murray, and Hugo Larochelle. Made: Masked autoencoder for distribution estimation. In International conference on machine learning, pp. 881-889. PMLR, 2015.
Frank Gray. Pulse code communication. United States Patent Number 2632058, 1953.
Arthur Gretton, Karsten M Borgwardt, Malte Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-77) 2012.
Shuyang Gu, Dong Chen, Jianmin Bao, Fang Wen, Bo Zhang, Dongdong Chen, Lu Yuan, and Baining Guo. Vector quantized diffusion model for text-to-image synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10696-10706, 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840 6851, 2020.
10

---

Published as a conference paper at ICLR 2023
Emiel Hoogeboom, Alexey A Gritsenko, Jasmijn Bastings, Ben Poole, Rianne van den Berg, and Tim Salimans. Autoregressive diffusion models. arXiv preprint arXiv:2110.02037, 2021a.
Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling, Argmax flows and multinomial diffusion: Learning categorical distributions. Advances in Neural Information Processing Systems, 34:12454-12465, 2021b.
Aapo Hyvärinen. Some extensions of score matching. Computational statistics & data analysis, 51 (5):2499 2007.
Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score match- ing. Journal of Machine Learning Research, 6(4), 2005.
Daniel D Johnson, Jacob Austin, Rianne den Berg, and Daniel Tarlow. Beyond in-place comuu-- tion: Insertion and deletion in denoising probabilistic models. arXiv preprint arXi::2107.07675, 2021.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization, arXiv preprint arXiv:1412.6900 2014.
Siwei Lyu. Interpretation and generalization of score matching. arXiv preprint arXiv:1205.2629, 2012.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pp. 8162-8171, PMLR, 2021.
Colin Raffel. Learning-base baseed methods for comparing sequences, with applications to cudio-to-mide aligument and matching. Columbia University, 2016.
Jascha Sohl-Dicksteen, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learn- ing. pp. 2265-2265. PMLR, 2015.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint MrXiv:2011.13455 2020.
Tijmen Tieleman. Training restricted boltzmann machines using approximations to the likelihood gradient. In Proceedings of the 25th international conference on Machine learning, pp. 1064 1071, 2008.
Aaron Van Den Oord, Oriol Vinyals, al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Ramon Van Handel. Filtering, stability, and robustness. PhD thesis, California Institute of Technol- ogy, 2007.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Dinghuai Zhang, Nikolay Malkin, Zhen Liu, Alexandra Volokhova, Aaron Courville, and Yoshua Bengio. Generative flow networks for discrete probabilistic modeling. arXiv preprint 2rXiv:2202.01361, 2022.
Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. arXiv preprint urXiv:2204. /3902, 2022.
11

---

Published as a conference paper at ICLR 2023
A RELATED WORK
The discrete diffusion model has actually been described in the original paper by (Sohl-Dickstein al., 2015), in which binomial diffusion process is considered for binary variable modeling. After the resurgence of the diffusion models with successful applications on continuous variable distribution modeling (Ho et al., 2020; Song et al., 2020), the recent discrete diffusion extensions mainly focus on the design of forward corruption kernels. The multinomial diffusion extension is proposed (Hoogeboom et al., 2021b), which inspires more structured perturbation operations in the forward process in Austin et al. (2021) beyond uniform noise injection, including discretized Gaus- sian perturbation, transition proportion to token embedding distance, and absorbing state corrup- tion. Furthermore, Hoogeboom et al. (2021a) introduce the standard masking operation in language model into forward process, and Johnson et al. (2021) consider the injection and deletion operation for corruption that is beyond the in-place perturbation. Discrete diffusion models demonstrate bet- ter flexibility and have been used as alternative of autoregressive models, egg., in Gu et al. (2022); Cohen et al. (2022), discrete diffusion models are used for latent codes quantization and composed with decoder/encoder in VQ-VAE. These models still follow the finite-step corruption in forward processes and exploit MLE-related objectives for learning, which is orthogonal to our focus.
The most related work is Campbell et. (2022), which generalize the discrete diffusion processes by characterizing the continuum of discrete variable evolving through continuous- time Markov chain. Our work also considers the continuous-time discrete diffusion. The major differences lie in two- fold: i), we designed score-based learning, while the learning in Campbell et al. (2022) relies on ELBO ii), our derivation admits analytical sampling strategy during backward sampling, in addi- tion to the commonly used numerical simulations in Campbell et al. (2022). Together with (Camp- bell et al., 2022), we complete the missing puzzle of discrete diffusion models in continuous-time framework.
B DEFFERED PROOF
B.1 PROOF FOR PROPOSITION 3.1
Proof. For a time u E (0, T), denote the filtration F,, = ofX. : v E [0, which is the sigma algebra of the events, Xu, from time 0 to time u. Also, denote the filtration for the reverse time process as F,, = or (X, : v E [0, u]) = o(x, o E T], which is the sigma algebra of events, Xu. from time to time T. By the Markov property, Fu, and F(u) are conditionally independent given X,, Now, consider any Fr-t, and any < t. Then we have
P(XT-. = F. X - = xt, xt,A)
(27)
=P(X, = x, = It, A) = P(X, = Rs. X, = Nt, A)
(28)
P(X, I1, A)
IAA|X, Zo. X, 2, PPX, 2z 21|X, = I. )P = x.a
(29)
P(A|X, = 23 PX,, =xx==00 =xa)
=P(X, = =, = z.) P(X, Is))
P(X, = xx)) qt (z)( (o((1e) (xe)
(30)
is independent of A, which proves the reverse time process X, is continuous time Markov chain.
Also, from the derivation above we have:
ge(Wa) IPX, = x,,x, =xx==xA xz, A) = P(X, =,|X, = xx) (x.(()) (x())))
gejs (ze
(31)
Hence, we have established the proposition.
12

---

Published as a conference paper at ICLR 2023
B.2 PROOF FOR PROPOSITION 3.2
Proof. Since the transition kernel for the time reversal process X, satisfies Equation 9, we consider its time derivative. For y,, we have
91 qr--(U)
di 2T -1-- - dt (9 (x)) -
(32)
d gr (1))
9T (V- d 9T (T-
( dt (
(33)
6r (x))
For the first term in Equation 33, we use Kolmogorov forward equation:
d
d
Σ
Teduaa (20))r))(((())
(34)
dt
dt
d Tetata (20) 1 (100000000))) dt
Σ
(35)
ToE
Σ meEX
2r --ounuulloo000r--c03 Tedaaa (20)
(36)
Σ -
(37)
Since X' is finite space, the summation over is finite, hence obtain:
lim de (((
d
lim E. (= (/- 00 (x)
(-t (r- 9T -
(38)
For the second term in Equation 33, we use Kolmogorov backward equation to obtain the derivative:
d
di -r--
Σ -
(39)
Then we have:
lim GT ( d
t+t -r- (x) dt QT ( = 9T- (v)
Qr (V,
(40)
(x)
By combining these two results, and using the property that:
d
9r- (x) QT - (I,
a 3) lim Q - 0
(41)
t0 di
Qr (T-
then relabelling T - s by yields:
Re(x,y) 9:(u)
(x.
(42)
(z)
B.3 PROOF FOR PROPOSITION 3.3
Proof. For one direction, if π; Te. it is trivial that their conditional distributions match. For the other direction, consider z, X' and an arbitrary probability distribution x. We have:
(1)
TΦ=² 22 y ..., y --1,yy)
(43)
(zx
((x x2 x3 .... x n- P)
Tyy y2 3 ..., y D-1 y0 D) t (x 22 3² ..... y DD (DD) (a x2 ..... x D-1 yD
<(x], y2 y² .... D- 11 D) (x x² ya .... y D- yD) a(x1, x² ... D-1 1)
(44)
D
x (1
lid-1 , D
(45)
11 I:D)
d=1
Thus the probability ratio can be decomposed as a product completely determined by the singleton conditional distributions. Hence, if nt and n. have the same singleton conditional distributions, then ti (y)/ (z) = Tz (B)/TT(()) Vz, y. As they are both distributions, we can easily see 1/ Σ, (y)/m/) E, T2())/ (T2) 1/ ma (x)), Va. Hence π2.
13

---

Published as a conference paper at ICLR 2023
B.4 PROOF FOR PROPOSITION 3.4
The key idea for the proof is that the conditional distribution on d-th dimension does not rely on the value zd. That is to say, the value of (xx clz²²) po = clz}", 0) does not depend on zę. Hence, we have:
Σ (x)
ΣΣ
2x = (d 1o =
(46)
ZEEN'
d=C
qla0/" Vi xx==)) =) log log =c c/2²
qc(Te)
(47)
EdeEc _dd=)
d=1 CEC TEE
a10 /d X
, =c)
ΣEE E d=1cEC Ail e'El
l Ad ,d")
log Dogpr(22 (X = cl20
(48)
Exec Vd xd=
uexx /d X _x==) =c)
log logpt (x* $
cl22";;0 Σ (0(2)( .c")
(49)
Evec /dd X xd=
d=1
e'EC
2l Ad x = logn = cz)
(50)
dm 1 CEC
D
(1x)) logp: x/
(51)
d=1 11EE
Σ
Σ
x :0)
(52)
Z.EX' d=1
By substituting this result into original score matching loss function:
0 arg min /' ".
'T
D
(z.)
Σ
Σ
qi(x = =
log logpi(X"
de
(53)
ZEEX
CEC
obtain simplified and tractable loss function:
8" = arg min S
D
Σ (x(1))
Σ
log p. (x" = =0 dt
(54)
x1EX
1=1
and prove the proposition 3.4.
C EXPERIMENTAL DETAILS
C.1 NOISE SCHEDULE OF ((t)
Empirically we find the following time schedule (()) to be effective in some situations:
1/4
= cos "
2 +1
(55)
This provides cosine style noise schedule so that the forward process has reasonable noise level to contribute to sample quality (Nichol & Dhariwal, 2021).
C.2 SYNTHETIC DATA
To parameterize fe t), we use the same 3-layer MLP as used in Zhang et al. (2022). Each hidden layer has dimensionality of 256, with elu activations. We use a constant learning rate of 1e-4 and train the model using 300k steps, where the per-step batch size is 128. The data is generated on the fly, using the data generator provided by Dai et al. (2020).
14

---

Published as a conference paper at ICLR 2023
Table 5: Conditional music generation compared against ground truth completions, with different
noise schedules.
Methods
Hellinger Distance
Proportion of Outliers
SDDM
0.3736 + 0.0024 0.3735 + 0.0025
0.1093 0.0016 0.1071 + 0.0017
SDDM with (55)
2spirals
8gaussians
circles
moons
pinwheel
swissroll
chekcerboard
Figure 4: Visualization of the true data in 2D space via decoding of Gray codes.
C.2.1 ABLATION STUDY
In this section we provide the ablation study three parameterization methods proposed in sec- tion 5, as well as the two proposed sampling methods (namely the forward Euler's method, denoted Fwd Euler, and the analytical sampling method, denoted analytical). Overall we can see all the combinations of parameterizaaionntsampp different samplers and models are mixed, which indicates that more efficient parameterizations like masked or hollow models can achieve better quality-sper trade-offs than EBM in this synthetic data.
achieve reasonable results. The comparison among
C.3 EXPERIMENTS ON CIFAR10
Vector Quantization We model the images learned categorical discrete latent space with a Vector Quantized Variational Autoencoder (VQ-VAB) (Van Den Oord al. (2017)). The VQVAE encoder maps an image in H x W x to I x 4 training. we use the VQGAN (Esser al. (2021)) variant which uses GAN loss and perceptual loss in addition to the ELBO objective.
W tokens with vocabulary size of 512. During
We follow the implementation of VQGAN in MaskGIT (Chang al. (2022)) with adaptations for a smaller down sample rate and fewer parameters. In particular, we use three convolution blocks in the encoder, with average pooling for down sampling between them. The three blocks use 64, 128, and 256 filters respectively, where each block consists of two standard residual blocks. We use nearest neighbor lookup to map the encoder outputs to a token index, according to codebook of 512 x 256. The decoder mirrors the encoder architecture. The overall parameters for the VQVAE is 12.35M.
We use the straight-through gradient estimator (Van Den Oord et al. (2017)) for the quantization part during training. To acquire a general VQ-VAE model as the bridge between pixel space and latent space, train on the ImageNet dataset (Deng al. (2009)) at 64 x 64 resolution for 90 epochs. The model is trained with the Adam optimizer (31 = 0, B2 = 0.99) (Kingma & Ba (2014)), using a peak learning rate of le 4 with a schedule of linear warm up and cosine decay. The GAN loss and perceptual loss are added with weight 0.1.
For the 32 32 images in the CIFAR10 dataset, 8 8 latent codes are produced. The reconstruction achieves FID of 9.05 and IS of 9.67.
Paramterization and training We parameterize fo t) using the masked modeling (see sec- tion 5.2). The backbone of the neural network is the same as BERT-base, which consists of 12 layers of Transformers, where each layer has attention heads, embedding size of 768 and hid- den layer size of 3072 for MLPs. We simply concatenate the time embedding representation of t together with the other tokens. After obtaining the final embedding of each token, we feed that into 2-block ResNet (using the same parameterization used in Campbell al. (2022) for their music
15

---

Published as a conference paper at ICLR 2023
Table 6: Ablation study of SDDM on synthetic dataset, using the same experimetal protocol in
Methods
2spirals
8gaussians
circles
moons
pinwheel
swissroll
chekcerboard
Analyyyltta
0.250
0.075
0.267
0.037
0.323
0.305
0,383
MaaicclAnalltical
0.143
0.080
0.126
0.064
0.242
0.116
0.345
Hollow+Analytical
0.227
0.052
0.229
0.117
0.049
0.201
0.557
EBM+Fwd Eulcr
0,199
0.022
0.173
0.095
0.167
0.287
0.211
Masked+Fwd Euler
0.120
0,028
0.204
0.195
0.064
0.086
0.248
Hollow+Fwo Euler
0.179
0.088
0.077
0.034
0.288
0.281
0.148
generation experiments) and predict logits for the masked position. We use a constant uniform rate of 0.007 for the forward process.
We train the model with 4x4x4 TPU-v4 chips, with batch size of 128. The training is done after 700k steps in about 20h. The learning rate warmed up from to le- during the first 3% steps of training. and then decays to 0 in a linear schedule. The final evaluation is done on the exponential moving average of the model parameters, with the decay rate of 0.999.
C.4 MUSIC DATASET
We parameterize the fo(x,t) t) using the proposed hollow Transformer (see section 5.3). Each Trans- former component has 6 layers with embedding size of 256, attention heads and hidden dimension of 2048. We simply concatenate the time embedding together with other tokens for the We experiment it using 2x2 TPU-v4 chips, and it took around 12 hours to get to convergence, or roughly 2m steps with batch size of 64.
16


Reviews:
[{"review_id": "letWdofeIW", "reviewer": "Reviewer_tzeK", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The work examines the problem of continuous-time discrete diffusion models, and the paper mainly focuses on 1) extending the existing discrete-time discrete diffusion models to the continuous-time via stochastic jump processes, 2) extending the score function to generic categorical variables and deriving a score-matching based learning for the model, 3) different sampling methods for the reverse process. Experimental validations on synthetic data, image and music generation have shown promising results comparing to existing discrete-time and continuous-time discrete diffusion models.", "summary_of_the_review": "I think the score-based continuous-time extension of the discrete diffusion models is interesting and the paper makes a novel contribution to it. However, I am leaning towards the negative side due to the weak presentation and not-so-convincing experiments.", "strength_and_weaknesses": "*Strengths*\n\nI think this work studies an interesting problem and it is nice to see the authors take the effort to extend the score functions to the categorical data and manage to leverage it to derive a score-matching based objective for estimating the continuous-time diffusion models. The main idea of using score-based objectives instead of ELBO surrogates is well-motivated, and the related works are properly referenced.\n\nThe authors have tried many different parametrizations and a new design for the conditional marginals, and it is also nice to see the proposed score-based learning naturally introduces a new analytic sampling algorithm. Experiments with toy 2-D datasets, CIFAR-10, and monophonic music dataset show the soundness of the proposed method, and the results are competitive with previous state-of-the-art methods. \n\n\n*Weaknesses*\n\nThe presentation is on the weaker side, and I think the paper can be significantly improved by having better clarity.\n- The authors extend the ratio test proposed in Hyvarinen (2007) to the categorical variables but point to equation (14) as the \"score-matching\" objective. However, it does not resemble those used in Hyvarinen (2007) or Song et al. (2020). It would be better if the authors could discuss the relationship of the proposed objective to the classic ones.\n- It would be nicer if the authors could discuss more regarding the similarities and differences from Campbell et al. (2022). For example, is there any subtle difference between the derivation of the continuous time modeling in Section 3.1 and those in Section 3.1 in Campbell et al. (2022)? Also, comparisons between the learning objectives would also make it easier to understand the benefits of the proposed method.\n\nThe parametrization of the conditional marginal distributions. Even though the authors have tried to design an efficient structure through different architecture choices including the hollow transformers, the current design is more computationally demanding comparing to existing methods and prevents it from directly modeling the CIFAR10 data with resolution as small as 32x32, not to mention other datasets with higher resolution. So that it is not easy to perform an apple-to-apple comparison with existing work.\n\nThe experimental validations are not very convincing. Two benefits of the proposed method are discussed in the introduction, including 1) the continuous-time discrete diffusion formulation is more flexible comparing to the discrete-time counterparts and may admit more effective estimation and generation; 2) the score-based learning yields superior estimation quality comparing to the ELBO surrogate used in Campbell et al. (2022) for the continuous-time discrete diffusion models. However, neither of them is convincingly validated in the experiment section. First, there is no formal comparison of the computations used during training and generation between the proposed model and the discrete-time ones. Second, the noise schedule adopted here is different from what used in Campbell at al. (2022). It has been shown that noise schedule can have noticeable impact on the final performance. Thus, this difference may question the validity of the conclusion that the better performance of the proposed model can be solely attributed to the score-based learning process. Third, it would be desirable to perform some ablation studies to validate the designs. For example, comparisons among all three parametrizations in the synthetic experiments and comparisons between the sampling methods discussed in Section 4.2 and 4.3.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "l1_FTSCI2l", "reviewer": "Reviewer_grv3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The submission considers score-based generative models with SDEs, but for discrete variables. The main challenge here is that the existing score-based methods do not apply to discrete variables since the log likelihood is not well defined. To deal with this issue, the authors use continuous-time discrete-state Markov process to model the forward process, with the according backward process also being a continuous-time discrete-state Markov process. Under this setting, a score-based estimation method is proposed for the diffusion model, with a sampling scheme for the reverse process. \n\nThe performance of the proposed method is thoroughly studied empirically, demonstrating an improved performance. ", "summary_of_the_review": "The submission is well written with lots of technical details. A new method is proposed with empirical results supporting the performance improvement. To the best of my knowledge, the submission extends the existing score-based method to discrete variables, which is a contribution to this area. ", "strength_and_weaknesses": "Strength\n\n1. The submission is well written. The notations are well defined. The used concepts are background details are clearly provided in the submission, making the paper self-contained. \n\n2. The empirical analysis is thorough and clear. The protocol, competing methods, expected results and the achieved results are clearly provided.  \n\n3. The methodology is novel. While the whole framework follows the existing framework of score-based generative models with SDEs, the proposed score function for discrete state variables is an extension to existing works. \n\nWeakness\n\n1. I am confused by the stochastic jump process. To the best of my knowledge, this terminology is more often used for jump processes like Poisson jump process. Specifically, it emphasizes more the discreteness of the path, instead of the discreteness of the variables: many stochastic jump processes have continuous-time variables. The model used in this submission is more a continuous-time discrete-state Markov process. Would it be possible for authors to clarify the rigorous definition of this terminology? \n\n2. I am wondering whether we can \"naively\" use a score-based generative models with SDEs for discrete variables, by simply doing some discretization when conducting sampling? This might be a bad idea when the variable has only two categories, but what if it has more than 200 categories like the image example in the experiments? What is the problem of this simple solution? Answering this question can further motivate the submission. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "VbcPcRDnQy", "reviewer": "Reviewer_eZKv", "summary": "", "questions": "", "limitations": "", "rating": 10, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper presents a diffusion-type generative model for categorical data. For this, the authors derive a variational learning scheme to approximate the discrete state space analog of the score of the marginal likelihood. The authors present a forward noise and backward denoising model based on a continuous-time Markov chain model. For learning the authors present three different neural-network-based parametrizations for the variational approximation. The algorithm is extensively evaluated for synthetic and real-world data and shows very good results.", "summary_of_the_review": "The paper presents a new algorithm for generative modeling with categorical data. It is very nicely written and it shows very competitive results.", "strength_and_weaknesses": "Overall, I do not have much to say about this paper, besides that, I really really enjoyed reading it. It is written in an easy-to-follow way, even though continuous-time stochastic process modeling usually requires immense technical machinery. The approximate inference algorithm is derived in a principled way and the empirical results of the paper are very promising.\n\nSome sugestions I have are:\n- Including the result of equation (47) in the main text, would have made it easier for me to understand how the learning of the conditionals is related to the ratio in equation (10)\n- In Sec 4.1. the transition matrix is given as an aside as $q_{t \\mid s}(\\cdot \\mid \\cdot)=\\int_s^t \\exp(Q_{\\tau})\\, \\mathrm d \\tau$. This is not correct as it is the solution to the linear ODE given in equation (7), for which the solution is given by the state-transition function, for which there is generally no decomposition as a matrix exponential, see, e.g., \"Peano–Baker series\", or, \"ordered exponential\".\n- The results for reversing a CTMC are not new and well known in the context of continuous-time filtering and smoothing. The authors should consider citing the relevant work, see, e.g., [1-3], and the references therein.\n\n[1] Anderson, Brian DO, and Ian B. Rhodes. \"Smoothing algorithms for nonlinear finite-dimensional systems.\" Stochastics: An International Journal of Probability and Stochastic Processes 9.1-2 (1983): 139-165.\n[2] Elliott, R. \"Reverse-time Markov processes (Corresp.).\" IEEE transactions on information theory 32.2 (1986): 290-292.\n[3] Van Handel, Ramon. Filtering, stability, and robustness. Diss. California Institute of Technology, 2007.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "IuFND8G6Ts-", "reviewer": "Reviewer_4YQP", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "There is a recent emergence of interests in adapting diffusion models to handle discrete data and this work adds to that literature. To understand the main difference from prior work, it is important to note that there are two (almost) equivalent frameworks for learning diffusion models: the variational lower bound framework and the denoising score matching framework. Prior work on discrete diffusion models have mostly focused on the variational lower bound framework because score functions are not defined in discrete domains. This paper replaces the continuous score with a discrete analogy which were studied in previous work on ratio matching/generalized score matching. The authors then design a transformer-based architecture to model this \"discrete score\" (which is essentially all complete conditionals at time t). Experiments show that the proposed discrete score matching objective is effective in fitting discrete energy-based models and diffusion models, with comparable or worse results on image data but better results on music data. ", "summary_of_the_review": "The work investigates an under-explored direction for developing discrete diffusion models. However, the paper is not ready yet due to lack of necessary theoretical developments to address the major problems and insufficient empirical evidence for supporting their claims. ", "strength_and_weaknesses": "## Strengths\n\n* This paper is well-motivated. A well-known fact in the diffusion model community is that score-based objectives often lead to better likelihoods than the ELBO. It is natural to ask if this holds for discrete data. However, the score-based interpretation was overlooked in the prior literation of discrete diffusion models. They mostly mentioned that scores are undefined in this case but ignored the fact that Hyvärinen had written about discrete generalization of score matching (called ratio matching) in his own work. \n\n* The proposed learning objectives are technically sound. As far as I checked, the derivations are correct. Modeling the complete conditionals also seems a very natural choice to me. \n\n## Weaknesses\n\n* The paper claims they extend score matching to general categorical data in several places. However, this is not true given the prior work on ratio matching and generalized score matching. The discrete analogy of score has also appeared in these works. From this perspective, the technical novelty seems to be quite limited. \n\n* Although the motivation suggests that using score-based loss function should work better than ELBOs, the results on image data indicates the opposite. On simple datasets like CIFAR10, the proposed method was outperformed by prior discrete diffusion models, including D3PM, Campbell et al. \n\n* One reason the paper uses for explaining the superiority of continuous diffusion models over discrete ones is that score matching for discrete variables does not benefit from prior knowledge about ordinal discrete data. However, I believe this is not the case for the \"discrete score\" definition in (11), which leverages the \"sparse neighborhood\" knowledge by only looking at states that differ by +-1. \n\n* Section 4 is confusing. I do not get the reasoning \"Such a jump rate depends on both the time t and ... Hence, it is inefficient to simulate the reverse time process X_t exactly\". ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster