Original: S CORE - BASED C ONTINUOUS - TIME D ISCRETE D IFFUSION M ODELS
Revised: S CORE - BASED C ONTINUOUS - TIME D ISCRETE D IFFUSION M ODELS

Original: Anonymous authors Paper under double-blind review
Revised: Google Research; University of Alberta schuurmans

Original: A BSTRACT
Revised: A BSTRACT

Original: Score-based modeling through stochastic differential equations (SDEs) has provided a new perspective on diffusion models, and demonstrated superior performance on continuous data.
Revised: Score-based modeling through stochastic differential equations (SDEs) has provided a new perspective on diffusion models, and demonstrated superior performance on continuous data.

Original: However, the gradient of the log-likelihood function, i.e. , the score function, is not properly defined for discrete spaces.
Revised: However, the gradient of the log-likelihood function, i.e. , the score function, is not properly defined for discrete spaces.

Original: This makes it non-trivial to adapt the score-based modeling to categorical data.
Revised: This makes it non-trivial to adapt the score-based modeling to categorical data.

Original: 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.
Revised: 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.

Original: This formulation admits an analytical simulation during backward sampling.
Revised: This formulation admits an analytical simulation during backward sampling.

Original: 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.
Revised: 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.

Original: We demonstrate the effectiveness of the proposed method on a set of synthetic and real-world music and image benchmarks.
Revised: We demonstrate the effectiveness of the proposed method on a set of synthetic and real-world music and image benchmarks.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020) have emerged as an important technique for data distribution modeling, where a data-corrupting forward process is coupled with a denoising reverse process to simulate a diffusion relationship between the data distribution and an uninformed source.
Revised: Diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020) have emerged as an important technique for data distribution modeling, where a data-corrupting forward process is coupled with a denoising reverse process to simulate a diffusion relationship between the data distribution and an uninformed source.

Original: Such models admit stable learning procedures and have demonstrated superior performance on continuous data modeling in challenging scenarios (Dhariwal & Nichol, 2021), leading to rapidly increasing popularity.
Revised: Such models admit stable learning procedures and have demonstrated superior performance on continuous data modeling in challenging scenarios (Dhariwal & Nichol, 2021), leading to rapidly increasing popularity.

Original: Song et al.
Revised: Song et al.

Original: (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.
Revised: (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.

Original: This perspective has provided a unified framework under a new score-based learning objective, and inspired a variety of simulation methods for efficient sampling and inference (De Bortoli et al., 2021; Zhang & Chen, 2022).
Revised: This perspective has provided a unified framework under a new score-based learning objective, and inspired a variety of simulation methods for efficient sampling and inference (De Bortoli et al., 2021; Zhang & Chen, 2022).

Original: Given the advantages of diffusion models in terms of flexibility, learning tractability, and sampling, there have been several attempts to extend the approach to discrete data.
Revised: Given the advantages of diffusion models in terms of flexibility, learning tractability, and sampling, there have been several attempts to extend the approach to discrete data.

Original: Recent attempts have investigated alternative corruption operations for discrete data in the forward process, yielding promising results (Hoogeboom et al.
Revised: Recent attempts have investigated alternative corruption operations for discrete data in the forward process, yielding promising results (Hoogeboom et al.

Original: , 2021b;a; Austin et al., 2021).
Revised: , 2021b;a; Austin et al., 2021).

Original: However, these extensions still execute a finite sequence of corruption and restoration steps, and remain restricted to a fixed reverse sampling strategy that can be sub-optimal.
Revised: However, these extensions still execute a finite sequence of corruption and restoration steps, and remain restricted to a fixed reverse sampling strategy that can be sub-optimal.

Original: To overcome this limitation, we investigate whether a continuoustime discrete diffusion formulation might admit more effective estimation and generation.
Revised: To overcome this limitation, we investigate whether a continuoustime discrete diffusion formulation might admit more effective estimation and generation.

Original: Such an extension is highly non-trivial however.
Revised: Such an extension is highly non-trivial however.

Original: The continuous-time diffusion framework is based on a stochastic differential equation (SDE) with respect to the score function, which itself is the gradient of the log-likelihood with respect to a continuous variable.
Revised: The continuous-time diffusion framework is based on a stochastic differential equation (SDE) with respect to the score function, which itself is the gradient of the log-likelihood with respect to a continuous variable.

Original: Although this can be used to characterize a continuum of infinitesimally evolving distributions over a continuous space, such a formulation no longer exists for discrete variables, since the gradient of the log-likelihood does not exist with respect to a discrete variable.
Revised: Although this can be used to characterize a continuum of infinitesimally evolving distributions over a continuous space, such a formulation no longer exists for discrete variables, since the gradient of the log-likelihood does not exist with respect to a discrete variable.

Original: Recently, Campbell et al.
Revised: Recently, Campbell et al.

Original: (2022) made significant progress in closing this gap by proposing a discrete data distribution that evolves in a continuous-time Markov chain.
Revised: (2022) made significant progress in closing this gap by proposing a discrete data distribution that evolves in a continuous-time Markov chain.

Original: With this formulation they were able to approximate maximum likelihood training with an evidence lower bound (ELBO) surrogate and using a predictor-corrector sampling scheme to generalize some of the benefits of continuous-time modeling to a discrete space.
Revised: With this formulation they were able to approximate maximum likelihood training with an evidence lower bound (ELBO) surrogate and using a predictor-corrector sampling scheme to generalize some of the benefits of continuous-time modeling to a discrete space.

Original: However, a limitation of this previous work is the reliance on an ELBO approximation to the MLE when it is known that score-based learning yields superior estimation quality (when it can be applied) due to its unbiasedness (Song et al., 2020).
Revised: However, a limitation of this previous work is the reliance on an ELBO approximation to the MLE when it is known that score-based learning yields superior estimation quality (when it can be applied) due to its unbiasedness (Song et al., 2020).

Original: Of course, developing an analog of the score matching for discrete spaces is non-trivial due to non-differentiability.
Revised: Of course, developing an analog of the score matching for discrete spaces is non-trivial due to non-differentiability.

Original: Nevertheless, a score function 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.
Revised: Nevertheless, a score function 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.

Original: 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.
Revised: 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.

Original: In addition, we complete the agenda for continuous time diffusion in discrete spaces by formulating a coherent SDE in terms of stochastic jump processes.
Revised: In addition, we complete the agenda for continuous time diffusion in discrete spaces by formulating a coherent SDE in terms of stochastic jump processes.

Original: The main contributions are:
Revised: The main contributions are:

Original: • We extend the definition of score functions to generic categorical discrete variables, and derive a continuous-time 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¨arinen, 2007) can be obtained as a special case; •
Revised: • We extend the definition of score functions to generic categorical discrete variables, and derive a continuous-time 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 ¨ arinen, 2007) can be obtained as a special case; •

Original: In Section 4, we develop a 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 a set of synthetic and real-world music and image benchmarks, achieving promising results in Section 6.
Revised: In Section 4, we develop a 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 a set of synthetic and real-world music and image benchmarks, achieving promising results in Section 6.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: Diffusion models (Sohl-Dickstein et al., 2015) are characterized by a forward Markov process that transforms an observation x 0 ∼ π data ( x 0 ) to a reference distribution x T ∼ q T ( x T ) , and a backward process, which is also Markovian, that recovers the data distribution from x T .
Revised: Diffusion models (Sohl-Dickstein et al., 2015) are characterized by a forward Markov process that transforms an observation x 0 ∼ π data ( x 0 ) to a reference distribution x T ∼ q T ( x T ) , and a backward process, which is also Markovian, that recovers the data distribution from x T .

Original: Specifically, the forward process is defined through simple corruption operations, q t +1 | t .
Revised: Specifically, the forward process is defined through simple corruption operations, q t +1 | t .

Original: For continuous data the corruption kernel usually adds Gaussian noise (Ho et al., 2020); whereas for discrete data, where x 0 ∈ X with finite cardinality |X| , the corruption kernel can be uniform, discrete Gaussian, or some other choice (Austin et al., 2021;
Revised: For continuous data the corruption kernel usually adds Gaussian noise (Ho et al., 2020); whereas for discrete data, where x 0 ∈ X with finite cardinality |X| , the corruption kernel can be uniform, discrete Gaussian, or some other choice (Austin et al., 2021;

Original: Johnson et al., 2021).
Revised: Johnson et al., 2021).

Original: Given the corruption kernel, after T -steps, the forward process forms a joint distribution,
Revised: Given the corruption kernel, after T -steps, the forward process forms a joint distribution,

Original: [Equation]
Revised: [Equation]

Original: The backward process can be derived from the joint distribution via Bayes’ rule,
Revised: The backward process can be derived from the joint distribution via Bayes’ rule,

Original: [Equation]
Revised: [Equation]

Original: where q t ( x t ) denotes the marginal distribution and the prior q T is usually a simple distribution.
Revised: where q t ( x t ) denotes the marginal distribution and the prior q T is usually a simple distribution.

Original: Typically the backward kernel q [Equation] 1 ( x t | x k +1 ) is intractable; thus it is usually parameterized by a neural network, denoted p θt | t +1 , and learned via ELBO (Sohl-Dickstein et al., 2015;
Revised: Typically the backward kernel q t | t +1 ( x t | x k +1 ) is intractable; thus it is usually parameterized by a neural network, denoted p θt | t +1 , and learned via ELBO (Sohl-Dickstein et al., 2015;

Original: Ho et al., 2020;
Revised: Ho et al., 2020;

Original: Austin et al., 2021) or score-matching (Song & Ermon, 2019).
Revised: Austin et al., 2021) or score-matching (Song & Ermon, 2019).

Original: Due to the structure of the joint distribution (1) and (2), the ELBO admits a particularly simple formulation as
Revised: Due to the structure of the joint distribution (1) and (2), the ELBO admits a particularly simple formulation as

Original: [Equation]
Revised: [Equation]

Original: which is applicable for both continuous and discrete diffusion model learning.
Revised: which is applicable for both continuous and discrete diffusion model learning.

Original: For continuous variables with Gaussian corruptions [Equation]
Revised: For continuous variables with Gaussian corruptions q ( x t +1 | x t ) =

Original: = N (cid:0) x t +1 ; (cid:112)
Revised: N (cid:0) x t +1 ; (cid:112)

Original: 1 − β t
Revised: 1 − β t

Original: +1 x t , β t +1 I (cid:1) , and a backward kernel p θ ( x t | x t +1 ) =
Revised: +1 x t , β t +1 I (cid:1) , and a backward kernel p θ ( x t | x t +1 ) =

Original: N (cid:16) xt ;
Revised: N (cid:16) x

Original: [Equation] 
Revised: [Equation] [Equation]

Revised: , β t

Original: [Equation]
Revised: [Equation]

Original:  , such that r θt is learned as a neural network and β t is a predefined variance schedule, the ELBO (3) can be rewritten as
Revised: I (cid:17) , such that r θt is learned as a neural network and β t is a predefined variance schedule, the ELBO (3) can be rewritten as

Original: [Equation]
Revised: [Equation]

Original: where α t = (cid:81) t − 1 t =0 (1 − β t ) .
Revised: where α t = (cid:81) t − 1 t =0 (1 − β t ) .

Original: The Equation 4 is highly related to score-matching (Hyv¨arinen &
Revised: The Equation 4 is highly related to score-matching (Hyv¨arinen &

Original: Dayan, 2005) with a score function defined as ∇ x log p t +1 | t ( x t +1 | x t ) .
Revised: Dayan, 2005) with a score function defined as ∇ x log p t +1 | t ( x t +1 | x t ) .

Original: Reverse sampling in discrete-time models is somewhat restricted by the forward process.
Revised: Reverse sampling in discrete-time models is somewhat restricted by the forward process.

Original: Therefore, continuous-time diffusion models have been constructed with t indexed from [0 , T ] .
Revised: Therefore, continuous-time diffusion models have been constructed with t indexed from [0 , T ] .

Original: Song et al.
Revised: Song et al.

Original: (2020) develop an SDE perspective and define a diffusion process with respect to the score as
Revised: (2020) develop an SDE perspective and define a diffusion process with respect to the score as

Original: [Equation]
Revised: [Equation]

Original: where w and ¯ w are standard Wiener processes, f ( x, t ) is a vector-valued drift function, and g ( t ) is a scalar-valued diffusion coefficient.
Revised: where w and ¯ w are standard Wiener processes, f ( x, t ) is a vector-valued drift function, and g ( t ) is a scalar-valued diffusion coefficient.

Original: A score-based learning method can be easily derived as an extension of (4).
Revised: A score-based learning method can be easily derived as an extension of (4).

Original: However, the score function ∇ x log p t ( x ) is not defined for a discrete variable, and the SDE above no longer suffices to characterize a continuous-time extension for discrete diffusion.
Revised: However, the score function ∇ x log p t ( x ) is not defined for a discrete variable, and the SDE above no longer suffices to characterize a continuous-time extension for discrete diffusion.

Original:  C ONTINUOUS T IME D ISCRETE S CORE M ATCHING
Revised:  C ONTINUOUS T IME D ISCRETE S CORE M ATCHING

Original: Although Campbell et al.
Revised: Although Campbell et al.

Original: (2022) bypass the score function in a continuous-time extension, by leveraging a stochastic process view of the ELBO approximation, we will instead focus on a score-based extension for continuous-time discrete diffusion.
Revised: (2022) bypass the score function in a continuous-time extension, by leveraging a stochastic process view of the ELBO approximation, we will instead focus on a score-based extension for continuous-time discrete diffusion.

Original: 3. C ONTINUOUS T IME M ODELING
Revised: 3. C ONTINUOUS T IME M ODELING

Original: Consider the finite discrete state space X
Revised: Consider the finite discrete state space X

Original: =
Revised: =

Original: C D , where C = { 1 , 2 , ..., C } is a code book.
Revised: C D , where C = { 1 , 2 , ..., C } is a code book.

Original: To generalize score matching from a continuous space R n to discrete space X , we first model the forward process as a continuous time Markov chain { X t } t ∈ [0 ,T ] , whose transition probability is characterized by rate matrices Q t ∈ R |X|×|X| .
Revised: To generalize score matching from a continuous space R n to discrete space X , we first model the forward process as a continuous time Markov chain { X t } t ∈ [0 ,T ] , whose transition probability is characterized by rate matrices Q t ∈ R |X|×|X| .

Original: In particular, if we let q denote the distribution for the forward process X t , the transition probability will satisfy the Kolmogorov forward equation:
Revised: In particular, if we let q denote the distribution for the forward process X t , the transition probability will satisfy the Kolmogorov forward equation:

Original: [Equation]
Revised: [Equation]

Original: If the forward process starts at the target distribution q 0 = π data , the marginal at time t has the form:
Revised: If the forward process starts at the target distribution q 0 = π data , the marginal at time t has the form:

Original: [Equation]
Revised: [Equation]

Original: By properly choosing the rate matrices Q t , we can achieve a final distribution close to a known tractable distribution q T ≈ π ref .
Revised: By properly choosing the rate matrices Q t , we can achieve a final distribution close to a known tractable distribution q T ≈ π ref .

Original: Then, the reverse time process X t =
Revised: Then, the reverse time process X t =

Original: X T − t can be expressed as a generative process from the reference distribution π ref to the target distribution π data (Anderson &
Revised: X T − t can be expressed as a generative process from the reference distribution π ref to the target distribution π data (Anderson &

Original: Rhodes, 1983; Van Handel, 2007).
Revised: Rhodes, 1983; Van Handel, 2007).

Original: Proposition 3.1.
Revised: Proposition 3.1.

Original: The reverse time process X t of the continuous time Markov chain X t is also a
Revised: The reverse time process X t of the continuous time Markov chain X t is also a

Original: Markov process, whose transition probabilities q s | t ( ·|· ) for s < t satisfy:
Revised: Markov process, whose transition probabilities q s | t ( ·|· ) for s < t satisfy:

Original: [Equation]
Revised: [Equation]

Original: Since we are considering a finite space X , the reverse time process X is uniquely determined by its rate matrices, which we denote as R t .
Revised: Since we are considering a finite space X , the reverse time process X is uniquely determined by its rate matrices, which we denote as R t .

Original: Using the transition in Equation 9, we have the following.
Revised: Using the transition in Equation 9, we have the following.

Original: Proposition 3.2.
Revised: Proposition 3.2.

Original: For a continuous time Markov chain { X t } t ∈ [0 ,T ] with distribution q and rate matrices Q t , the rate matrices R t for the reverse process satisfy: 
Revised: For a continuous time Markov chain { X t } t ∈ [0 ,T ] with distribution q and rate matrices Q t , the rate matrices R t for the reverse process satisfy: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Equation 10 provides a closed form expression for the reverse time rate R t .
Revised: Equation 10 provides a closed form expression for the reverse time rate R t .

Original: Therefore, once we know the ratio q t ( y ) /q
Revised: Therefore, once we know the ratio q t ( y ) /q

Original: t ( x ) , we can obtain the generative flow towards π data .
Revised: t ( x ) , we can obtain the generative flow towards π data .

Original: 3. C ATEGORICAL R ATIO M ATCHING
Revised: 3. C ATEGORICAL R ATIO M ATCHING

Original: In general, the ratio q t ( y ) /q
Revised: In general, the ratio q t ( y ) /q

Original: t ( x ) in Equation 10 is intractable.
Revised: t ( x ) in Equation 10 is intractable.

Original: This ratio behaves analogously to the score function ∇ log π ( x ) in Equation 6.
Revised: This ratio behaves analogously to the score function ∇ log π ( x ) in Equation 6.

Original: Such a connection inspires learning the ratio q t ( y ) /q
Revised: Such a connection inspires learning the ratio q t ( y ) /q

Original: t ( x ) in a similar manner to learning the score function for diffusion models in continuous spaces.
Revised: t ( x ) in a similar manner to learning the score function for diffusion models in continuous spaces.

Original: For binary discrete variable, Hyv ¨ arinen (2007) proposed to match the ratio π ( X \ d , X d = c ) /π
Revised: For binary discrete variable, Hyv¨arinen (2007) proposed to match the ratio π ( X \ d , X d = c ) /π

Original: ( X \ d , X d
Revised: ( X \ d , X d

Original: = 1 − c ) as an extension for score matching.
Revised: = 1 − c ) as an extension for score matching.

Original: In this work, we consider the more general categorical case, where neither the score function ∇ log π ( x ) nor the binary ratio is well defined.
Revised: In this work, we consider the more general categorical case, where neither the score function ∇ log π ( x ) nor the binary ratio is well defined.

Original: The generalization relies on the singleton conditional distribution:
Revised: The generalization relies on the singleton conditional distribution:

Original: [Equation]
Revised: [Equation]

Original: yielding a score function in categorical space we seek to match.
Revised: yielding a score function in categorical space we seek to match.

Original: The sufficiency of matching Equation 11 is guaranteed by the property that the joint distribution is completely determined by its singleton conditional distributions (Brook, 1964; Lyu, 2012).
Revised: The sufficiency of matching Equation 11 is guaranteed by the property that the joint distribution is completely determined by its singleton conditional distributions (Brook, 1964; Lyu, 2012).

Original: Proposition 3.3.
Revised: Proposition 3.3.

Original: Consider random variables
Revised: Consider random variables

Original: X
Revised: X

Original: =
Revised: =

Original: ( X 1 , ..., X D ) ∈ X , and two probability distributions π 1 , π 2 .
Revised: ( X 1 , ..., X D ) ∈ X , and two probability distributions π 1 , π 2 .

Original: We have π 1 =
Revised: We have π 1 =

Original: π 2 , if and only if their conditional distributions are equal π 1 ( X d = x d | x \ d )
Revised: π 2 , if and only if their conditional distributions are equal π 1 ( X d = x d | x \ d )

Original: =
Revised: =

Original: π 2 ( X d = x d | x \ d ) , for any x ∈ X and d = 1 , ..., D .
Revised: π 2 ( X d = x d | x \ d ) , for any x ∈ X and d = 1 , ..., D .

Original: Going back to the ratio q t ( y ) /q
Revised: Going back to the ratio q t ( y ) /q

Original: t ( x ) in reverse time rate Equation 10, Proposition 3.3 tells us we can recover the probability ratio by employing a time dependent neural network p t ( · ; θ ) to match the conditional distribution,
Revised: t ( x ) in reverse time rate Equation 10, Proposition 3.3 tells us we can recover the probability ratio by employing a time dependent neural network p t ( · ; θ ) to match the conditional distribution,

Original: [Equation]
Revised: [Equation]

Original: To train p t ( · ; θ ) , we minimize the expected cross entropy with respect to the conditional distributions along the forward process:
Revised: To train p t ( · ; θ ) , we minimize the expected cross entropy with respect to the conditional distributions along the forward process:

Original: [Equation]
Revised: [Equation]

Original: where the loss is minimized if p t ( X dt | x \ dt ; θ ) ≡ q t ( X dt | x \ dt ) .
Revised: where the loss is minimized if p t ( X dt | x \ dt ; θ ) ≡ q t ( X dt | x \ dt ) .

Original: This loss function matches the ratio in
Revised: This loss function matches the ratio in

Original: Equation 11, hence we name it as categorical ratio matching as a generalization of Hyv ¨ arinen (2007).
Revised: Equation 11, hence we name it as categorical ratio matching as a generalization of Hyv ¨ arinen (2007).

Original: 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.
Revised: 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.

Original: In Equation 13, the expectation over q t ( · ) can be efficiently estimated via Monte Carlo sampling.
Revised: In Equation 13, the expectation over q t ( · ) can be efficiently estimated via Monte Carlo sampling.

Original: However, the conditional distribution q t ( X dt = c | x \ dt ) is intractable, which brings difficulty to the training.
Revised: However, the conditional distribution q t ( X dt = c | x \ dt ) is intractable, which brings difficulty to the training.

Original: Fortunately, using the property of conditional distribution, we can avoid computing this intractable term in the loss function.
Revised: Fortunately, using the property of conditional distribution, we can avoid computing this intractable term in the loss function.

Original: Proposition 3.4.
Revised: Proposition 3.4.

Original: The categorical ratio matching loss function in Equation 13 can be simplified as:
Revised: The categorical ratio matching loss function in Equation 13 can be simplified as:

Original: [Equation]
Revised: [Equation]

Original: Using this surprisingly neat loss function in Equation 14, we can efficiently learn the conditional distribution.
Revised: Using this surprisingly neat loss function in Equation 14, we can efficiently learn the conditional distribution.

Original: The learned p t ( X dt | x \ dt ; θ ) determines a reverse process, and we use p ( · ; θ ) to denote its joint distribution in order to distinguish with the true reverse process.
Revised: The learned p t ( X dt | x \ dt ; θ ) determines a reverse process, and we use p ( · ; θ ) to denote its joint distribution in order to distinguish with the true reverse process.

Original: We will sometimes drop the θ if it does not create ambiguity.
Revised: We will sometimes drop the θ if it does not create ambiguity.

Original:  C ONTINUOUS T IME D ISCRETE S AMPLING
Revised:  C ONTINUOUS T IME D ISCRETE S AMPLING

Original: We next develop and efficient sampling scheme for the proposed diffusion model.
Revised: We next develop and efficient sampling scheme for the proposed diffusion model.

Original: 4. C ONTINUOUS T IME S IMULATION FOR F ORWARD P ROCESS
Revised: 4. C ONTINUOUS T IME S IMULATION FOR F ORWARD P ROCESS

Original: In a continuous time Markov chain, the transition matrix q t | s ( ·|· ) from time s to time t in the forward process can be obtained by solving the ODE in Equation 7.
Revised: In a continuous time Markov chain, the transition matrix q t | s ( ·|· ) from time s to time t in the forward process can be obtained by solving the ODE in Equation 7.

Original: For general rate matrices
Revised: For general rate matrices

Original: Q t ∈ R |X|×|X| , solving the ODE is intractable.
Revised: Q t ∈ R |X|×|X| , solving the ODE is intractable.

Original: Therefore, we follow standard practice in diffusion models (Austin et al., 2021; Campbell et al., 2022) and approximate the forward process by factorizing X t
Revised: Therefore, we follow standard practice in diffusion models (Austin et al., 2021; Campbell et al., 2022) and approximate the forward process by factorizing X t

Original: =
Revised: =

Original: ( X 1 t , ..., X Dt ) where each sub-process X dt propagates independently.
Revised: ( X 1 t , ..., X Dt ) where each sub-process X dt propagates independently.

Original: Furthermore, we also define the sub-rate matrix Q dt = Qβ ( t ) in terms of a fixed base rate Q = P Λ P − 1 ∈ R C × C and time schedule function β ( t ) ; see design of β ( t ) in Appendix C.1.
Revised: Furthermore, we also define the sub-rate matrix Q d t = Qβ ( t ) in terms of a fixed base rate Q = P Λ P − 1 ∈ R C × C and time schedule function β ( t ) ; see design of β ( t ) in Appendix C.1.

Original: In this way, the sub-transition matrix in each dimension can then be easily computed (Campbell et al., 2022) as:
Revised: In this way, the sub-transition matrix in each dimension can then be easily computed (Campbell et al., 2022) as:

Original: [Equation]
Revised: [Equation]

Original: In particular, we use a uniform stationary base rate Q = 11 T − 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 et al., 2021).
Revised: In particular, we use a uniform stationary base rate Q = 11 T − 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 et al., 2021).

Original: For simplicity of comparison we use the same schedule function as in Campbell et al.
Revised: For simplicity of comparison we use the same schedule function as in Campbell et al.

Original: for all the experiments, and in Appendix C.

Revised: for all the experiments, and in Appendix C.

Original: we discuss other possible choices.
Revised: we discuss other possible choices.

Original: 4. D ISCRETE T IME S AMPLING FOR R EVERSE P ROCESS
Revised: 4. D ISCRETE T IME S AMPLING FOR R EVERSE P ROCESS

Original: Even given the learned conditional distribution p t ( · ; θ ) , simulating the reverse time process is much harder than simulating the forward process, as the reverse process cannot be factorized.
Revised: Even given the learned conditional distribution p t ( · ; θ ) , simulating the reverse time process is much harder than simulating the forward process, as the reverse process cannot be factorized.

Original: For example, consider x, y that are only different at the d -th site.
Revised: For example, consider x, y that are only different at the d -th site.

Original: The rate for the reversed process to jump to y from x at time t is:
Revised: The rate for the reversed process to jump to y from x at time t is:

Original: [Equation]
Revised: [Equation]

Original: Such a jump rate depends on both the time t and the value in the other dimensions x \ dt .
Revised: Such a jump rate depends on both the time t and the value in the other dimensions x \ dt .

Original: Hence, we can not simulate each dimension in parallel, making exact simulation of the reverse time process X t less efficient.
Revised: Hence, we can not simulate each dimension in parallel, making exact simulation of the reverse time process X t less efficient.

Original: Considering R d t ( x t , y ; θ ) is already an approximation of R d t ( x t , y ) , we employ Euler’s method for parallel simulation.
Revised: Considering R dt ( x t , y ; θ ) is already an approximation of R dt ( x t , y ) , we employ Euler’s method for parallel simulation.

Original: Specifically, given x t at time t , we fix the rate in Equation 16, then determine the transition probabilities for dimension d at time t − ϵ according to:
Revised: Specifically, given x t at time t , we fix the rate in Equation 16, then determine the transition probabilities for dimension d at time t − ϵ according to:

Original: [Equation]
Revised: [Equation]

Original: We clip the quantities to ensure all probabilities are non-negative.
Revised: We clip the quantities to ensure all probabilities are non-negative.

Original: Then, we collect a new value from each dimension to obtain a new state y t − ϵ , which has the factorized probability:
Revised: Then, we collect a new value from each dimension to obtain a new state y t − ϵ , which has the factorized probability:

Original: [Equation]
Revised: [Equation]

Original: In this way, we can update all dimensions of x t in parallel, making the simulation of the reverse time process much more efficient.
Revised: In this way, we can update all dimensions of x t in parallel, making the simulation of the reverse time process much more efficient.

Original: 4. A NALYTICAL S AMPLING FOR R EVERSE P ROCESS
Revised: 4. A NALYTICAL S AMPLING FOR R EVERSE P ROCESS

Original: The Euler’s method mentioned above assumes the rate matrix R τ is fixed during τ ∈ [ t − ϵ, t ] .
Revised: The Euler’s method mentioned above assumes the rate matrix R τ is fixed during τ ∈ [ t − ϵ, t ] .

Original: Such an approximation makes the samples x t deviate from the correct time slice marginal q t ( · ) , especially when the time step ϵ is large.
Revised: Such an approximation makes the samples x t deviate from the correct time slice marginal q t ( · ) , especially when the time step ϵ is large.

Original: To mitigate this approximation error, diffusion models usually introduce a corrector to improve sample quality (Song et al., 2020; Campbell et al., 2022), however this increases computational cost.
Revised: To mitigate this approximation error, diffusion models usually introduce a corrector to improve sample quality (Song et al., 2020; Campbell et al., 2022), however this increases computational cost.

Original: Instead, we propose an alternative approach that leverages implicit modeling of p t ( X d | x \ dt ; θ ) .
Revised: Instead, we propose an alternative approach that leverages implicit modeling of p t ( X d | x \ d t ; θ ) .

Original: Specifically, for distribution q , we have:
Revised: Specifically, for distribution q , we have:

Original: [Equation]
Revised: [Equation]

Original: Since q t | 0 is tractable, computing Equation 19 is efficient once we know q 

Revised: Since q t | 0 is tractable, computing Equation 19 is efficient once we know q

Original: | t ( X d 0 | x \ dt ) .
Revised: | t ( X d 0 | x \ dt ) .

Original: Hence, we can replace the explicit approximation in Equation 12 by the following implicit approximation:
Revised: Hence, we can replace the explicit approximation in Equation 12 by the following implicit approximation:

Original: [Equation]
Revised: [Equation]

Original: Equation 20 provides a tractable transformation from p 0 | t ( X d 0 | x \ dt ; θ ) to p t ( X dt | x \ dt ; θ ) .
Revised: Equation 20 provides a tractable transformation from p 0 | t ( X d 0 | x \ dt ; θ ) to p t ( X dt | x \ dt ; θ ) .

Original: Hence, we can continue using the categorical ratio matching loss in Equation 14 to train p 0 | t ( X d 0 | x \ dt ; θ ) .
Revised: Hence, we can continue using the categorical ratio matching loss in Equation 14 to train p 0 | t ( X d 0 | x \ dt ; θ ) .

Original: To conduct backward sampling via this new parameterization, we consider the true reverse process:
Revised: To conduct backward sampling via this new parameterization [Equation] we consider the true reverse process:

Original: [Equation]
Revised: [Equation]

Original: By substituting p 0 ( X d 0 | x \ dt ; θ ) , we have:
Revised: By substituting p 0 ( X d 0 | x \ dt ; θ ) , we have:

Original: [Equation]
Revised: [Equation]

Original: Thus, we obtain an analytical expression for the reverse process sampling that avoids the simulation error in Euler’s method.
Revised: Thus, we obtain an analytical expression for the reverse process sampling that avoids the simulation error in Euler’s method.

Original: Hence, we refer to this method as analytical sampling.
Revised: Hence, we refer to this method as analytical sampling.

Original:  P ARAMETERIZATION
Revised:  P ARAMETERIZATION

Original: The key to the simplicity of the objective function in Equation 14 is the special structure of the conditional marginal distributions p t ( X dt
Revised: The key to the simplicity of the objective function in Equation 14 is the special structure of the conditional marginal distributions p t ( X dt

Original: | x \ dt ; θ ) and p 

Revised: | x \ dt ; θ ) and p

Original: | t ( X d 0 | x \ dt ; θ ) such that the marginals of dth dimension do not depend on the current value x dt at time t .
Revised: | t ( X d 0 | x \ dt ; θ ) such that the marginals of dth dimension do not depend on the current value x dt at time t .

Original: However, this elegance in the objective brings additional challenges to design an efficient and flexible neural network parameterization.
Revised: However, this elegance in the objective brings additional challenges to design an efficient and flexible neural network parameterization.

Original: A key design constraint is that the predictions at a dimension d must not depend on x dt , otherwise the information leak would make Equation 14 trivial to solve.
Revised: A key design constraint is that the predictions at a dimension d must not depend on x dt , otherwise the information leak would make Equation 14 trivial to solve.

Original: The prediction can depend on the other coordinates in x t however.
Revised: The prediction can depend on the other coordinates in x t however.

Original: We offer three alternative architectures that are able to satisfy this constraint but incur different trade-offs between flexibility and computational cost.
Revised: We offer three alternative architectures that are able to satisfy this constraint but incur different trade-offs between flexibility and computational cost.

Original: With out loss of generality we consider the parameterization of p t ( X dt | x \ d t ; θ ) , however the same designs can be directly applied to the parameterization of p 0 | t ( X d 0 | x \ dt ; θ ) .
Revised: With out loss of generality we consider the parameterization of p t ( X dt | x \ d t ; θ ) , however the same designs can be directly applied to the parameterization of p 0 | t ( X d 0 | x \ dt ; θ ) .

Original: 5. E NERGY BASED MODELS
Revised: 5. E NERGY BASED MODELS

Original: The most general and flexible form of parameterization is an Energy based model (EBM), where an arbitrary neural network f θ ( x, t ) :
Revised: The most general and flexible form of parameterization is an Energy based model (EBM), where an arbitrary neural network f θ ( x, t ) :

Original: C D × R (cid:55)→ R can be used to specify the energy of any sample.
Revised: C D × R (cid:55)→ R can be used to specify the energy of any sample.

Original: In this case the conditional marginals can be modeled as:
Revised: In this case the conditional marginals can be modeled as:

Original: [Equation]
Revised: [Equation]

Original: where we overload the notation to use [ c, x \ dt ] to denote [ x 0 t , x 1 t , . .
Revised: where we overload the notation to use [ c, x \ dt ] to denote [ x 0 t , x 1 t , . .

Original: , X dt = c, x d +1 t , . .
Revised: , X dt = c, x d +1 t , . .

Original: such that only the d -th value is replaced with c .
Revised: such that only the d -th value is replaced with c .

Original: This modeling flexibility however comes at a high computational cost.
Revised: This modeling flexibility however comes at a high computational cost.

Original: Specifically, to evaluate (cid:81)
Revised: Specifically, to evaluate (cid:81)

Original: Dd =1 p t ( X dt | x \ dt ; θ ) one needs O ( D × C ) rounds of evaluation of f θ , which is computationally prohibitive when modeling high dimensional data with a deep neural parameterization of f θ .
Revised: Dd =1 p t ( X dt | x \ dt ; θ ) one needs O ( D × C ) rounds of evaluation of f θ , which is computationally prohibitive when modeling high dimensional data with a deep neural parameterization of f θ .

Original: Nevertheless, this connection between EBMs and diffusion provides another way to learn time-dependent EBMs, which might be of independent interest.
Revised: Nevertheless, this connection between EBMs and diffusion provides another way to learn time-dependent EBMs, which might be of independent interest.

Original: 5. M ASKED MODELS
Revised: 5. M ASKED MODELS

Original: To alleviate the computation overhead of EBMs while preserving flexibility, a masked model is a natural choice.
Revised: To alleviate the computation overhead of EBMs while preserving flexibility, a masked model is a natural choice.

Original: Specifically, let a masking function m d ( x ) =
Revised: Specifically, let a masking function m d ( x ) =

Original: [ x 1 , . . .
Revised: [ x 1 , . . .

Original: , x d − 1 , MASK , x d +1 , . . .
Revised: , x d − 1 , MASK , x d +1 , . . .

Original: , x D ] replace the d -th dimension of a given x to a special mask token MASK.
Revised: , x D ] replace the d -th dimension of a given x to a special mask token MASK.

Original: Then one can formulate the following conditional parameterization:
Revised: Then one can formulate the following conditional parameterization:

Original: [Equation]
Revised: [Equation]

Original: Here f can still be a general neural network with only a minor requirement on handling then masked token.
Revised: Here f can still be a general neural network with only a minor requirement on handling then masked token.

Original: Overall this approaches requires O ( D ) rounds of evaluation of f θ .
Revised: Overall this approaches requires O ( D ) rounds of evaluation of f θ .

Original: Since we are dealing with 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 f θ .
Revised: Since we are dealing with 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 f θ .

Original: 5. H OLLOW T RANSFORMERS
Revised: 5. H OLLOW T RANSFORMERS

Original: 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 vocabulary size of discrete data.
Revised: 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 vocabulary size of discrete data.

Original: As an alternative, we propose a Transformer variant that requires only
Revised: As an alternative, we propose a Transformer variant that requires only

Original: O (1) feed-forward evaluations.
Revised: O (1) feed-forward evaluations.

Original: The key constraint is that the diagonals of the Jacobian matrix of p t ( X t | x t ; θ ) be zero for any input x t .
Revised: The key constraint is that the diagonals of the Jacobian matrix of p t ( X t | x t ; θ ) be zero for any input x t .

Original: Many techniques have been developed to guarantee this property, including autoregressive masking (Germain et al., 2015; Vaswani et al., 2017), which creates a triangular Jacobian for multi-layer networks, or hollow masking (Chen & Duvenaud, 2019), which considers the full context but only permits a single layer of dense interaction between dimensions.
Revised: Many techniques have been developed to guarantee this property, including autoregressive masking (Germain et al., 2015; Vaswani et al., 2017), which creates a triangular Jacobian for multi-layer networks, or hollow masking (Chen & Duvenaud, 2019), which considers the full context but only permits a single layer of dense interaction between dimensions.

Original: 𝑄 𝑥 !
Revised: 𝑄 𝑥 !

Original: 𝑥 " 𝑥 # 𝑥 $ 𝑡 𝑡
Revised: 𝑥 " 𝑥 # 𝑥 $ 𝑡 𝑡

Original: TransformerLayer 𝐿×
Revised: TransformerLayer 𝐿×

Original: TransformerLayer TransformerLayer
Revised: TransformerLayer TransformerLayer

Original: TransformerLayer 𝑄 % 𝐾 & ∘ ,𝐾 ’ ∘ →
Revised: TransformerLayer 𝑄 % 𝐾 & ∘ ,𝐾 ’ ∘ →

Original: [𝑉 & ,𝑉 ’ ] 𝑃(𝑥 ! |𝑥 \! )
Revised: [𝑉 & ,𝑉 ’ ] 𝑃(𝑥 ! |𝑥 \! )

Original: 𝑃(𝑥 # |𝑥 \# ) 𝑃(𝑥 $ |𝑥 \$ ) 𝑃(𝑥 % |𝑥
Revised: 𝑃(𝑥 # |𝑥 \# ) 𝑃(𝑥 $ |𝑥 \$ ) 𝑃(𝑥 % |𝑥

Original: 𝑉 & K ’ 𝑉 ( 𝐾 ( 𝑥 !
Revised: 𝑉 & K ’ 𝑉 ( 𝐾 ( 𝑥 !

Original: 𝑥 " 𝑥 # 𝑥 $ 𝑡 𝑥 ! 𝑥 " 𝑥 # 𝑥 $ 𝑡 𝑥 !
Revised: 𝑥 " 𝑥 # 𝑥 $ 𝑡 𝑥 ! 𝑥 " 𝑥 # 𝑥 $ 𝑡 𝑥 !

Original: 𝑥 " 𝑥 # 𝑥 $ 𝑥 ! 𝑥 " 𝑥 # 𝑥 $ 𝑥 !
Revised: 𝑥 " 𝑥 # 𝑥 $ 𝑥 ! 𝑥 " 𝑥 # 𝑥 $ 𝑥 !

Original: 𝑥 " 𝑥 # 𝑥 $
Revised: 𝑥 " 𝑥 # 𝑥 $

Original: Jacobian matrix 𝜕𝑃 )
Revised: Jacobian matrix 𝜕𝑃 )

Original: (𝑋 ∣ 𝑥 ) ) 𝜕𝑥 )
Revised: (𝑋 ∣ 𝑥 ) ) 𝜕𝑥 )

Original: Here, to capture the full context for each dimension with expressive deep neural architectures, we introduce the hollow Transformer , a model that runs two autoregressive Transformers, one in each direction, for L -layers; see Figure 1.
Revised: Here, to capture the full context for each dimension with expressive deep neural architectures, we introduce the hollow Transformer , a model that runs two autoregressive Transformers, one in each direction, for L -layers; see Figure 1.

Original: The hollow Jacobian matrix is obtained via the summation of upper and lower triangular Jacobians so the full context for each position is obtained.
Revised: The hollow Jacobian matrix is obtained via the summation of upper and lower triangular Jacobians so the full context for each position is obtained.

Original: We add one additional Transformer layer at the top, where the query vector comes from the two embedding directions of the corresponding dimension, and attention on the key and value vectors are conducted jointly with the two directions.
Revised: We add one additional Transformer layer at the top, where the query vector comes from the two embedding directions of the corresponding dimension, and attention on the key and value vectors are conducted jointly with the two directions.

Original: For clarity we omit details of each Transformer layer as we use the standard architecture (Vaswani et al., 2017).
Revised: For clarity we omit details of each Transformer layer as we use the standard architecture (Vaswani et al., 2017).

Original: Overall, this specially designed 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.
Revised: Overall, this specially designed 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.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We present an empirical evaluation of the proposed diffusion approach on synthetic data, CIFAR dataset and the monophonic music dataset.
Revised: We present an empirical evaluation of the proposed diffusion approach on synthetic data, CIFAR dataset and the monophonic music dataset.

Original: The primary goal is to compare the effectiveness of the proposed categorical ratio matching with the alternative parameterizations presented in Section 5.
Revised: The primary goal is to compare the effectiveness of the proposed categorical ratio matching with the alternative parameterizations presented in Section 5.

Original: We use the time duration T = 1 and the uniform forward process in all cases.
Revised: We use the time duration T = 1 and the uniform forward process in all cases.

Original: Experiments are conducted on machines equipped with TPU-v4 chips.
Revised: Experiments are conducted on machines equipped with TPU-v4 chips.

Original: For more details please refer to Appendix C.
Revised: For more details please refer to Appendix C.

Original: 6.  D EEP EBM S ON S YNTHETIC D ATA
Revised:  D EEP EBM S ON S YNTHETIC D ATA

Original: As discussed in Section 5.1, one can leverage the EBMs for modeling conditional marginals.
Revised: As discussed in Section 5.1, one can leverage the EBMs for modeling conditional marginals.

Original: Here we verify this approach using a synthetic dataset from the discrete EBM learning community.
Revised: Here we verify this approach using a synthetic dataset from the discrete EBM learning community.

Original: Following prior works (Dai et al., 2020; Zhang et al., 2022), we use seven different distributions of 32-dimensional binary discrete data to evaluate different approaches.
Revised: Following prior works (Dai et al., 2020; Zhang et al., 2022), we use seven different distributions of 32-dimensional binary discrete data to evaluate different approaches.

Original: Each discrete distribution is converted from a 2-D continuous point ( x, y ) ∈ R 2 by quantizing both x and y with 16-bit representations using a Gray code (Gray, 1953).
Revised: Each discrete distribution is converted from a 2-D continuous point ( x, y ) ∈ R 2 by quantizing both x and y with 16-bit representations using a Gray code (Gray, 1953).

Original: We parameterize the energy function f θ ( x, 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 a sinusoidal embedding of t into each hidden layer before activation.
Revised: We parameterize the energy function f θ ( x, 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 a sinusoidal embedding of t into each hidden layer before activation.

Original: The uniform rate constant is set to 1.
Revised: The uniform rate constant is set to 1.

Original: and we use a time resolution of 1 e − 3 for simulation.
Revised: and we use a time resolution of 1 e − 3 for simulation.

Original: To measure sample quality, we follow (Zhang et al., 2022) and compare 4,000 generated samples to true data samples using exponential Hamming
Revised: To measure sample quality, we follow (Zhang et al., 2022) and compare 4,000 generated samples to true data samples using exponential Hamming

Original: MMD (Gretton et al., 2012), repeated 10 times to report average results in Table 1.
Revised: MMD (Gretton et al., 2012), repeated 10 times to report average results in Table 1.

Original: As baselines we include PCD (Tieleman, 2008), ALOE (Dai et al., 2020) with a larger network for dual parameterization, and EB-GFN (Zhang et al., 2022).
Revised: As baselines we include PCD (Tieleman, 2008), ALOE (Dai et al., 2020) with a larger network for dual parameterization, and EB-GFN (Zhang et al., 2022).

Original: Here we find that the proposed continuous-time categorical ratio matching approach is able to successfully learn a good EBM, yielding an MMD value that is consistently lower than the baseline methods.
Revised: Here we find that the proposed continuous-time categorical ratio matching approach is able to successfully learn a good EBM, yielding an MMD value that is consistently lower than the baseline methods.

Original: 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.
Revised: 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.

Original: These plots show a similar distribution to the ground truth data (see Appendix C.2 for more information).
Revised: These plots show a similar distribution to the ground truth data (see Appendix C.2 for more information).

Original: Ablation study: we further provide ablation studies on three different parameterizations proposed in Section 5 and two different sampling methods.
Revised: Ablation study: we further provide ablation studies on three different parameterizations proposed in Section 5 and two different sampling methods.

Original: For the full results please see Appendix C.2.1.
Revised: For the full results please see Appendix C.2.1.

Original: Overall when 3-layer transformers are used in these three parameterizations, the performance are comparable to each other, which shows that the masked model and hollow model can achieve better quality-speed trade-offs than EBMs.
Revised: Overall when 3-layer transformers are used in these three parameterizations, the performance are comparable to each other, which shows that the masked model and hollow model can achieve better quality-speed trade-offs than EBMs.

Original: We will revisit the comparison among samplers in next section.
Revised: We will revisit the comparison among samplers in next section.

Original: 6. I MAGE MODELING ON CIFAR 
Revised: 6. I MAGE MODELING ON CIFAR

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Figure]
Revised: [Figure]

Original: Next we evaluate image generation quality using the CIFAR10 dataset.
Revised: Next we evaluate image generation quality using the CIFAR10 dataset.

Original: In this case, each raw image has shape (32 , 32 , 3) with 256 choices per each pixel value.
Revised: In this case, each raw image has shape (32 , 32 , 3) with 256 choices per each pixel value.

Original: We compare several diffusion models from the literature, with the main goal of understanding performance in the different spaces with either continuous or discrete time formulations.
Revised: We compare several diffusion models from the literature, with the main goal of understanding performance in the different spaces with either continuous or discrete time formulations.

Original: We primarily evaluate the performance of the masked model formulation presented in Section 5.2, and report the commonly used Inception 
Revised: We primarily evaluate the performance of the masked model formulation presented in Section 5.2, and report the commonly used Inception Score (IS) and Fr´echet

Original: Score (IS) and Fr´echet Inception Distance (FID) against the training set.
Revised: Inception Distance (FID) against the training set.

Original: As the image pixels are naturally discretized continuous values
Revised: As the image pixels are naturally discretized continuous values

Original: , we instead evaluate the proposed method in the quantization hash-code space via pretrained VQGAN (Esser et al., 2021).
Revised: , we instead evaluate the proposed method in the quantization hash-code space via pretrained VQGAN (Esser et al., 2021).

Original: The main purpose of doing this is to 1) evaluate score matching on complex categorical discrete data; and 2) alleviate the burden of the Transformers modeling long sequences.
Revised: The main purpose of doing this is to 1) evaluate score matching on complex categorical discrete data; and 2) alleviate the burden of the Transformers modeling long sequences.

Original: The resulting data lies in a 64-dimensional categorical space with vocabulary size |C| =512.
Revised: The resulting data lies in a 64-dimensional categorical space with vocabulary size |C| =512.

Original: Note that VQGAN is a lossy compressor where the reconstructed images from the decoder obtain IS=9.67 and FID=9.05, which is the upper-bound of categorical diffusion built on top of this vector quantizer (VQ).
Revised: Note that VQGAN is a lossy compressor where the reconstructed images from the decoder obtain IS=9.67 and FID=9.05, which is the upper-bound of categorical diffusion built on top of this vector quantizer (VQ).

Original: Here we simply use a 12-layer BERT-base architecture to parameterize the f θ ( x, t ) .
Revised: Here we simply use a 12-layer BERT-base architecture to parameterize the f θ ( x, t ) .

Original: See Appendix C.3 for more details about the VQ and architecture used.
Revised: See Appendix C.3 for more details about the VQ and architecture used.

Original: From Table 2 we can see that, dealing with images in a categorical discrete space is generally much harder than an ordinal discrete space, as it loses ordinal structure as a prior.
Revised: From Table 2 we can see that, dealing with images in a categorical discrete space is generally much harder than an ordinal discrete space, as it loses ordinal structure as a prior.

Original: The proposed approach is able to approach the performance limit of VQ based categorical discrete diffusion, and improves the performance of D3PM and τ LDR in the same VQ space with the same parameterization.
Revised: The proposed approach is able to approach the performance limit of VQ based categorical discrete diffusion, and improves the performance of D3PM and τ LDR in the same VQ space with the same parameterization.

Original: In Figure 3 we visualize the reverse sampling procedure (unevenly sampled time steps) using two proposed samplers, where the analytical sampler achieves reasonable quality in fewer steps.
Revised: In Figure 3 we visualize the reverse sampling procedure (unevenly sampled time steps) using two proposed samplers, where the analytical sampler achieves reasonable quality in fewer steps.

Original: To quantitatively show why the continuous-time modeling would achieve better flexibility, we report the image quality with different number of sampling steps in Table 3.
Revised: To quantitatively show why the continuous-time modeling would achieve better flexibility, we report the image quality with different number of sampling steps in Table 3.

Original: We can see as D3PM was trained with 1,000 steps, the performance drops quickly with fewer steps.
Revised: We can see as D3PM was trained with 1,000 steps, the performance drops quickly with fewer steps.

Original: Our continuous-time version can be more robust even with much fewer steps.
Revised: Our continuous-time version can be more robust even with much fewer steps.

Original: Also the analytical sampler can be more robust when fewer steps are given, where the forward Euler may need corrector for better performance.
Revised: Also the analytical sampler can be more robust when fewer steps are given, where the forward Euler may need corrector for better performance.

Original: 6. M ONOPHONIC MUSIC MODELING
Revised: 6. M ONOPHONIC MUSIC MODELING

Original: Finally, we conduct a study on discrete music generation following the same settings used in Campbell et al.
Revised: Finally, we conduct a study on discrete music generation following the same settings used in Campbell et al.

Original: This benchmark is originally from the Lakh pianoroll dataset (Raffel, 2016;
Revised: This benchmark is originally from the Lakh pianoroll dataset (Raffel, 2016;

Original: Dong et al., 2018), cleaned by removing some trivial music sequences.
Revised: Dong et al., 2018), cleaned by removing some trivial music sequences.

Original: In the end there are 6,000 music sequences for training and 973 sequences for evaluation.
Revised: In the end there are 6,000 music sequences for training and 973 sequences for evaluation.

Original: Each sequence is of length 256, where the
Revised: Each sequence is of length 256, where the

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table] 
Revised: [Table] [Figure] [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: vocabulary size |X| =
Revised: vocabulary size |X| =

Original: 129 consists of 128 notes and a rest.
Revised: 129 consists of 128 notes and a rest.

Original: The vocabulary is scrambled so there is no ordering information preserved.
Revised: The vocabulary is scrambled so there is no ordering information preserved.

Original: This creates a challenging categorical discrete data modeling task.
Revised: This creates a challenging categorical discrete data modeling task.

Original: Here we primarily compare against existing diffusion models in discrete spaces.
Revised: Here we primarily compare against existing diffusion models in discrete spaces.

Original: [Table]
Revised: [Table]

Original: We use the simulation time ϵ = 1 e −for all methods.
Revised: We use the simulation time ϵ = 1 e −for all methods.

Original: Since the dimensionality is high, the conditional marginals p 0 | t ( X d | x \ dt ) are parameterized by the proposed hollow Transformer (as discussed in Section 5.3) with 6 Transformer layers.
Revised: Since the dimensionality is high, the conditional marginals p 0 | t ( X d | x \ dt ) are parameterized by the proposed hollow Transformer (as discussed in Section 5.3) with 6 Transformer layers.

Original: Appendix C.4 provides more details about the training and parameterizations.
Revised: Appendix C.4 provides more details about the training and parameterizations.

Original: The evaluation was conducted on the held-out test set, where for each sequence, a prefix of length 32 is given as conditioning and the models are asked to generate the remaining 224 tokens.
Revised: The evaluation was conducted on the held-out test set, where for each sequence, a prefix of length 32 is given as conditioning and the models are asked to generate the remaining 224 tokens.

Original: We follow the same evaluation protocol as in Campbell et al.
Revised: We follow the same evaluation protocol as in Campbell et al.

Original: and report the Hellinger Distance and Proportion of Outliers after 5 runs per each test case in
Revised: and report the Hellinger Distance and Proportion of Outliers after 5 runs per each test case in

Original: Table 4.
Revised: Table 4.

Original: Overall we can see the proposed SDDM is able to consistently improve the two metrics.
Revised: Overall we can see the proposed SDDM is able to consistently improve the two metrics.

Original:  L IMITATION AND C ONCLUSION
Revised:  L IMITATION AND C ONCLUSION

Original: In this paper we presented a new learning and sampling paradigm for continuous-time diffusion models in categorical discrete spaces.
Revised: In this paper we presented a new learning and sampling paradigm for continuous-time diffusion models in categorical discrete spaces.

Original: 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.
Revised: 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.

Original: The new learning paradigm also naturally introduces new sampling algorithms.
Revised: The new learning paradigm also naturally introduces new sampling algorithms.

Original: Despite the promising preliminary results, there are still limitations in the current treatment.
Revised: Despite the promising preliminary results, there are still limitations in the current treatment.

Original: 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 continuous and discrete spaces would be needed to handle data in mixed spaces.
Revised: 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 continuous and discrete spaces would be needed to handle data in mixed spaces.

Original: We believe this initiative sheds new light on score matching for discrete-space diffusion.
Revised: We believe this initiative sheds new light on score matching for discrete-space diffusion.

Original: A R ELATED W ORK
Revised: A R ELATED W ORK

Original: The discrete diffusion model has actually been described in the original paper by (Sohl-Dickstein et al., 2015), in which a binomial diffusion process is considered for binary variable modeling.
Revised: The discrete diffusion model has actually been described in the original paper by (Sohl-Dickstein et al., 2015), in which a binomial diffusion process is considered for binary variable modeling.

Original: 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.
Revised: 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.

Original: The multinomial diffusion extension is proposed (Hoogeboom et al., 2021b), which inspires more structured perturbation operations in the forward process in Austin et al.
Revised: The multinomial diffusion extension is proposed (Hoogeboom et al., 2021b), which inspires more structured perturbation operations in the forward process in Austin et al.

Original: (2021) beyond uniform noise injection, including discretized Gaussian perturbation, transition proportion to token embedding distance, and absorbing state corruption.
Revised: (2021) beyond uniform noise injection, including discretized Gaussian perturbation, transition proportion to token embedding distance, and absorbing state corruption.

Original: Furthermore, Hoogeboom et al.
Revised: Furthermore, Hoogeboom et al.

Original: (2021a) introduce the standard masking operation in language model into forward process, and Johnson et al.
Revised: (2021a) introduce the standard masking operation in language model into forward process, and Johnson et al.

Original: (2021) consider the injection and deletion operation for corruption that is beyond the in-place perturbation.
Revised: (2021) consider the injection and deletion operation for corruption that is beyond the in-place perturbation.

Original: Discrete diffusion models demonstrate better flexibility and have been used as alternative of autoregressive models, e.g. , in Gu et al.
Revised: Discrete diffusion models demonstrate better flexibility and have been used as alternative of autoregressive models, e.g. , in Gu et al.

Original: Cohen et al.
Revised: Cohen et al.

Original: (2022), discrete diffusion models are used for latent codes quantization and composed with decoder/encoder in VQ-VAE.
Revised: (2022), discrete diffusion models are used for latent codes quantization and composed with decoder/encoder in VQ-VAE.

Original: These models still follow the finite-step corruption in forward processes and exploit MLE-related objectives for learning, which is orthogonal to our focus.
Revised: These models still follow the finite-step corruption in forward processes and exploit MLE-related objectives for learning, which is orthogonal to our focus.

Original: The most related work is Campbell et al.
Revised: The most related work is Campbell et al.

Original: (2022), which generalize the discrete diffusion processes by characterizing the continuum of discrete variable evolving through continuous-time Markov chain.
Revised: (2022), which generalize the discrete diffusion processes by characterizing the continuum of discrete variable evolving through continuous-time Markov chain.

Original: Our work also considers the continuous-time discrete diffusion.
Revised: Our work also considers the continuous-time discrete diffusion.

Original: The major differences lie in twofold: i) , we designed score-based learning, while the learning in Campbell et al.
Revised: The major differences lie in twofold: i) , we designed score-based learning, while the learning in Campbell et al.

Original: (2022) relies on
Revised: (2022) relies on

Original: ELBO ii) , our derivation admits an analytical sampling strategy during backward sampling, in addition to the commonly used numerical simulations in Campbell et al.
Revised: ELBO ii) , our derivation admits an analytical sampling strategy during backward sampling, in addition to the commonly used numerical simulations in Campbell et al.

Original: Together with (Campbell et al., 2022), we complete the missing puzzle of discrete diffusion models in continuous-time framework.
Revised: Together with (Campbell et al., 2022), we complete the missing puzzle of discrete diffusion models in continuous-time framework.

Original: B D EFFERED P ROOF B. P ROOF FOR P ROPOSITION 3. 
Revised: B D EFFERED P ROOF B. P ROOF FOR P ROPOSITION 3.

Original: Proof.
Revised: Proof.

Original: For a time u ∈ (0 , T ) , denote the filtration F u = σ { X v : v ∈ [0 , u ] } , which is the sigma algebra of the events, X u , from time 0 to time u .
Revised: For a time u ∈ (0 , T ) , denote the filtration F u = σ { X v : v ∈ [0 , u ] } , which is the sigma algebra of the events, X u , from time 0 to time u .

Original: Also, denote the filtration for the reverse time process as F u = σ { X v : v ∈ [0 , u ] } =
Revised: Also, denote the filtration for the reverse time process as F u = σ { X v : v ∈ [0 , u ] } =

Original: σ { X v : v ∈ [ u, T ] } , which is the sigma algebra of events,
Revised: σ { X v : v ∈ [ u, T ] } , which is the sigma algebra of events,

Original: X u , from time u to time T .
Revised: X u , from time u to time T .

Original: By the Markov property, F u and F ( u ) are conditionally independent given X u .
Revised: By the Markov property, F u and F ( u ) are conditionally independent given X u .

Original: Now, consider any A ∈ F T − t , and any s < t .
Revised: Now, consider any A ∈ F T − t , and any s < t .

Original: Then we have
Revised: Then we have

Original: [Equation]
Revised: [Equation]

Original: ) is independent of A , which proves the reverse time process X t is a continuous time Markov chain.
Revised: 30) is independent of A , which proves the reverse time process X t is a continuous time Markov chain.

Original: Also, from the derivation above we have:
Revised: Also, from the derivation above we have:

Original: [Equation]
Revised: [Equation]

Original: Hence, we have established the proposition.
Revised: Hence, we have established the proposition.

Original: B. P ROOF FOR P ROPOSITION 3. 
Revised: B. P ROOF FOR P ROPOSITION 3.

Original: Proof.
Revised: Proof.

Original: Since the transition kernel for the time reversal process X t satisfies Equation
Revised: Since the transition kernel for the time reversal process X t satisfies Equation

Original: [Equation]
Revised: [Equation]

Original: we consider its time derivative.
Revised: we consider its time derivative.

Original: For x ̸ =
Revised: For x ̸ =

Original: y , we have
Revised: y , we have

Original: [Equation]
Revised: [Equation]

Original: For the first term in Equation
Revised: For the first term in Equation

Original: [Equation]
Revised: [Equation]

Original: we use Kolmogorov forward equation:
Revised: we use Kolmogorov forward equation:

Original: [Equation]
Revised: [Equation]

Original: Since X is a finite space
Revised: Since X is a finite space

Original: [Equation]
Revised: [Equation]

Original: the summation over z is finite
Revised: the summation over z is finite

Original: [Equation]
Revised: [Equation]

Original: hence we obtain:
Revised: hence we obtain:

Original: [Equation]
Revised: [Equation]

Original: For the second term in Equation
Revised: For the second term in Equation

Original: [Equation]
Revised: [Equation]

Original: we use Kolmogorov backward equation to obtain the derivative:
Revised: we use Kolmogorov backward equation to obtain the derivative:

Original: [Equation]
Revised: [Equation]

Original: Then we have:
Revised: Then we have:

Original: [Equation]
Revised: [Equation]

Original: By combining these two results
Revised: By combining these two results

Original: [Equation]
Revised: [Equation]

Original: and using the property that:
Revised: and using the property that:

Original: [Equation]
Revised: [Equation]

Original: then relabelling T − s by t yields:
Revised: then relabelling T − s by t yields:

Original: [Equation]
Revised: [Equation]

Original: B. P ROOF FOR P ROPOSITION 3. 
Revised: B. P ROOF FOR P ROPOSITION 3.

Original: Proof.
Revised: Proof.

Original: For one direction
Revised: For one direction

Original: [Equation]
Revised: [Equation]

Original: if π 1 = π 2 , it is trivial that their conditional distributions match.
Revised: if π 1 = π 2 , it is trivial that their conditional distributions match.

Original: For the other direction, consider x
Revised: For the other direction, consider x

Original: [Equation]
Revised: [Equation]

Original: y ∈ X and an arbitrary probability distribution π .
Revised: y ∈ X and an arbitrary probability distribution π .

Original: We have:
Revised: We have:

Revised: [Equation]

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Thus the probability ratio can be decomposed as a product completely determined by the singleton conditional distributions.
Revised: Thus the probability ratio can be decomposed as a product completely determined by the singleton conditional distributions.

Original: Hence, if π 1 and π 2 have the same singleton conditional distributions, then π 1 ( y )
Revised: Hence, if π 1 and π 2 have the same singleton conditional distributions, then π 1 ( y )

Original: /π 1 ( x )
Revised: /π 1 ( x )

Original: =
Revised: =

Original: π 2 ( y ) /π
Revised: π 2 ( y ) /π

Original: 2 ( x ) , ∀ x, y .
Revised: 2 ( x ) , ∀ x, y .

Original: As they are both distributions, we can easily see 1 /π 1 ( x ) = (cid:80)
Revised: As they are both distributions, we can easily see 1 /π 1 ( x ) = (cid:80)

Original: y π 1 ( y ) /π
Revised: y π 1 ( y ) /π

Original: 1 ( x ) = (cid:80)
Revised: 1 ( x ) = (cid:80)

Original: y π 2 ( y )
Revised: y π 2 ( y )

Original: /π
Revised: /π

Original: 2 ( x ) = 1 /π 2 ( x ) , ∀ x .
Revised: 2 ( x ) = 1 /π 2 ( x ) , ∀ x .

Original: Hence π 1 = π 2 .
Revised: Hence π 1 = π 2 .

Original: B. P ROOF FOR P ROPOSITION 3. 
Revised: B. P ROOF FOR P ROPOSITION 3.

Original: The key idea for the proof is that the conditional distribution on d -th dimension does not rely on the value x d .
Revised: The key idea for the proof is that the conditional distribution on d -th dimension does not rely on the value x d .

Revised: That is to say, the value of q t ( X d = c | x \ dt ) log p t ( X d = c

Revised: | x \ dt ; θ ) does not depend on x dt .

Original: Hence, we have:
Revised: Hence, we have:

Original: [Equation]
Revised: [Equation]

Original: By substituting this result into original score matching loss function:
Revised: By substituting this result into original score matching loss function:

Original: [Equation]
Revised: [Equation]

Original: (cid:34) D (cid:88)

Original: d = (cid:32) − (cid:88) c ∈C q t ( X d = c | x \ dt ) log p t ( X d = c

Original: | x \ dt ; θ ) (cid:33)(cid:35) dt (51) we obtain a simplified and tractable loss function:
Revised: we obtain a simplified and tractable loss function:

Original: [Equation]
Revised: [Equation]

Original: 52) and prove the proposition 3.4.
Revised: and prove the proposition 3. [Equation]

Original: C E XPERIMENTAL DETAILS C. N OISE SCHEDULE OF β ( t )
Revised: C E XPERIMENTAL DETAILS C. N OISE SCHEDULE OF β ( t )

Original: Empirically we find the following time schedule β ( t ) to be effective in some situations:
Revised: Empirically we find the following time schedule β ( t ) to be effective in some situations:

Original: [Equation]
Revised: [Equation]

Original: This provides a cosine style noise schedule so that the forward process has a reasonable noise level to contribute to sample quality (Nichol & Dhariwal, 2021).[Table]
Revised: This provides a cosine style noise schedule so that the forward process has a reasonable noise level to contribute to sample quality (Nichol & Dhariwal, 2021).

Original: C. S YNTHETIC DATA
Revised: C. S YNTHETIC DATA

Original: To parameterize f θ ( x, t ) , we use the same 3-layer MLP as used in Zhang et al.
Revised: To parameterize f θ ( x, t ) , we use the same 3-layer MLP as used in Zhang et al.

Original: Each hidden layer has dimensionality of 256, with elu activations.
Revised: Each hidden layer has dimensionality of 256, with elu activations.

Original: We use a constant learning rate of 1 e − 4 and train the model using 300k steps, where the per-step batch size is 128.
Revised: We use a constant learning rate of 1 e − 4 and train the model using 300k steps, where the per-step batch size is 128.

Original: The data is generated on the fly, using the data generator provided by Dai et al.
Revised: The data is generated on the fly, using the data generator provided by Dai et al.

Original: [Table]
Revised: [Table]

Revised: 2spirals 8gaussians circles

Original: [Table]
Revised: [Table]

Original:  A BLATION STUDY
Revised: C.2. A BLATION STUDY

Original: In this section we provide the ablation study on three parameterization methods proposed in section [Table] as well as the two proposed sampling methods (namely the forward Euler’s method, denoted as Fwd Euler [Table]
Revised: In this section we provide the ablation study on three parameterization methods proposed in section 5, as well as the two proposed sampling methods (namely the forward Euler’s method, denoted as Fwd Euler, and the analytical sampling method, denoted as analytical).

Original: and the analytical sampling method, denoted as analytical).

Original: Overall we can see all the combinations of parameterization+sampler achieve reasonable results.
Revised: Overall we can see all the combinations of parameterization+sampler achieve reasonable results.

Original: The comparison among different samplers and models are mixed [Table] which indicates that more efficient parameterizations like masked or hollow models can achieve better quality-speed trade-offs than EBM in this synthetic data.
Revised: The comparison among different samplers and models are mixed, which indicates that more efficient parameterizations like masked or hollow models can achieve better quality-speed trade-offs than EBM in this synthetic data.

Original: C. E XPERIMENTS ON CIFAR 
Revised: C. E XPERIMENTS ON CIFAR

Original: Vector Quantization We model the images in a learned categorical discrete latent space with a
Revised: Vector Quantization We model the images in a learned categorical discrete latent space with a

Original: Vector Quantized Variational Autoencoder (VQ-VAE) (Van Den Oord et al.
Revised: Vector Quantized Variational Autoencoder (VQ-VAE) (Van Den Oord et al.

Original: The VQVAE encoder maps an image in H × W × 3 to H 4 × W 4 tokens with a vocabulary size of 512.
Revised: The VQVAE encoder maps an image in H × W × 3 to H 4 × W 4 tokens with a vocabulary size of 512.

Original: During training, we use the VQGAN (Esser et al.
Revised: During training, we use the VQGAN (Esser et al.

Original: variant which uses a GAN loss and a perceptual loss in addition to the ELBO objective.
Revised: variant which uses a GAN loss and a perceptual loss in addition to the ELBO objective.

Original: We follow the implementation of VQGAN in MaskGIT (Chang et al.
Revised: We follow the implementation of VQGAN in MaskGIT (Chang et al.

Original: with adaptations for a smaller down sample rate and fewer parameters.
Revised: with adaptations for a smaller down sample rate and fewer parameters.

Original: In particular, we use three convolution blocks in the encoder, with average pooling for down sampling between them.
Revised: In particular, we use three convolution blocks in the encoder, with average pooling for down sampling between them.

Original: The three blocks use 64, 128, and 256 filters respectively, where each block consists of two standard residual blocks.
Revised: The three blocks use 64, 128, and 256 filters respectively, where each block consists of two standard residual blocks.

Original: We use nearest neighbor lookup to map the encoder outputs to a token index, according to a codebook of 512 × 256 .
Revised: We use nearest neighbor lookup to map the encoder outputs to a token index, according to a codebook of 512 × 256 .

Original: The decoder mirrors the encoder architecture.
Revised: The decoder mirrors the encoder architecture.

Original: The overall parameters for the VQVAE is 12.35M.
Revised: The overall parameters for the VQVAE is 12.35M.

Original: We use the straight-through gradient estimator (Van Den Oord et al.
Revised: We use the straight-through gradient estimator (Van Den Oord et al.

Original: for the quantization part during training.
Revised: for the quantization part during training.

Original: To acquire a general VQ-VAE model as the bridge between pixel space and latent space, we train it on the ImageNet dataset (Deng et al.
Revised: To acquire a general VQ-VAE model as the bridge between pixel space and latent space, we train it on the ImageNet dataset (Deng et al.

Original: at 64 × 64 resolution for 90 epochs.
Revised: at 64 × 64 resolution for 90 epochs.

Original: The model is trained with the Adam optimizer ( β 1 = 0 , β 2 = 0 . 99 ) (Kingma & Ba (2014)), using a peak learning rate of 1 e − 4 with a schedule of linear warm up and cosine decay.
Revised: The model is trained with the Adam optimizer ( β 1 = 0 , β 2 = 0 . 99 ) (Kingma & Ba (2014)), using a peak learning rate of 1 e − 4 with a schedule of linear warm up and cosine decay.

Original: The GAN loss and perceptual loss are added with weight 0.1.
Revised: The GAN loss and perceptual loss are added with weight 0.1.

Original: For the 32 × 32 images in the CIFAR10 dataset, 8 × 8 latent codes are produced.
Revised: For the 32 × 32 images in the CIFAR10 dataset, 8 × 8 latent codes are produced.

Original: The reconstruction achieves FID of 9.05 and IS of 9.67.
Revised: The reconstruction achieves FID of 9.05 and IS of 9.67.

Original: Paramterization and training We parameterize f θ ( x, t ) using the masked modeling (see section 5.2).
Revised: Paramterization and training We parameterize f θ ( x, t ) using the masked modeling (see section 5.2).

Original: The backbone of the neural network is the same as BERT-base, which consists oflayers of Transformers, where each layer has 12 attention heads, embedding size of 768 and hidden layer size of 3072 for MLPs.
Revised: The backbone of the neural network is the same as BERT-base, which consists oflayers of Transformers, where each layer has 12 attention heads, embedding size of 768 and hidden layer size of 3072 for MLPs.

Original: We simply concatenate the time embedding representation of t together with the other tokens.
Revised: We simply concatenate the time embedding representation of t together with the other tokens.

Original: After obtaining the final embedding of each token, we feed that into a 2-block ResNet (using the same parameterization as used in Campbell et al.
Revised: After obtaining the final embedding of each token, we feed that into a 2-block ResNet (using the same parameterization as used in Campbell et al.

Original: for their music  generation experiments) and predict logits for the masked position.
Revised: for their music [Table] generation experiments) and predict logits for the masked position.

Original: We use a constant uniform rate of 0.007 for the forward process.
Revised: We use a constant uniform rate of 0.007 for the forward process.

Original: We train the model with 4x4x4 TPU-v4 chips, with batch size of 128.
Revised: We train the model with 4x4x4 TPU-v4 chips, with batch size of 128.

Original: The training is done after 700k steps in about 20h.
Revised: The training is done after 700k steps in about 20h.

Original: The learning rate is warmed up from 0 to 1 e − 4 during the first 3% steps of training, and then decays to 0 in a linear schedule.
Revised: The learning rate is warmed up from 0 to 1 e − 4 during the first 3% steps of training [Table] and then decays to 0 in a linear schedule.

Original: The final evaluation is done on the exponential moving average of the model parameters, with the decay rate of 0.999.
Revised: The final evaluation is done on the exponential moving average of the model parameters, with the decay rate of 0.999.

Original: C. M USIC DATASET
Revised: C. M USIC DATASET

Original: We parameterize the f θ ( x, t ) using the proposed hollow Transformer (see section 5.3).
Revised: We parameterize the f θ ( x, t ) using the proposed hollow Transformer (see section 5.3).

Original: Each Transformer component has 6 layers with embedding size of 256, 8 attention heads and hidden dimension of 2048.
Revised: Each Transformer component has 6 layers with embedding size of 256, 8 attention heads and hidden dimension of 2048.

Original: We simply concatenate the time embedding together with other tokens for the feed-forward.
Revised: We simply concatenate the time embedding together with other tokens for the feed-forward.

Original: 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.
Revised: 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.