Title: Star-Shaped Denoising Diffusion Probabilistic Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: Generative models

Diffusion

Exponential Family
Abstract: Denoising Diffusion Probabilistic Models (DDPMs) provide the foundation for the recent breakthroughs in generative modeling.
Their Markovian structure makes it difficult to define DDPMs with distributions other than Gaussian or discrete.
In this paper, we introduce Star-Shaped DDPM (SS-DDPM).
Its *star-shaped diffusion process* allows us to bypass the need to define the transition probabilities or compute posteriors.
We establish duality between star-shaped and specific Markovian diffusions for the exponential family of distributions and derive efficient algorithms for training and sampling from SS-DDPMs.
In the case of Gaussian distributions, SS-DDPM is equivalent to DDPM.
However, SS-DDPMs provide a simple recipe for designing diffusion models with distributions such as Beta, von Mises–Fisher, Dirichlet, Wishart and others, which can be especially useful when data lies on a constrained manifold.
We evaluate the model in different settings and find it competitive even on image data, where Beta SS-DDPM achieves results comparable to a Gaussian DDPM.
Our implementation is available at https://github.com/andrey-okhotin/star-shaped

Star-Shaped Denoising Diffusion Probabilistic Models

Abstract

Denoising Diffusion Probabilistic Models (DDPMs) provide the foundation for the recent breakthroughs in generative modeling. Their Markovian structure makes it difficult to define DDPMs with distributions other than Gaussian or discrete. In this paper, we introduce Star-Shaped DDPM (SS-DDPM). Its star-shaped diffusion process allows us to bypass the need to define the transition probabilities or compute posteriors. We establish duality between star-shaped and specific Markovian diffusions for the exponential family of distributions and derive efficient algorithms for training and sampling from SS-DDPMs. In the case of Gaussian distributions, SS-DDPM is equivalent to DDPM. However, SS-DDPMs provide a simple recipe for designing diffusion models with distributions such as Beta, von Mises-Fisher, Dirichlet, Wishart and others, which can be especially useful when data lies on a constrained manifold. We evaluate the model in different settings and find it competitive even on image data, where Beta SS-DDPM achieves results comparable to a Gaussian DDPM. Our implementation is available at https://github.com/andrey-okhotin/star-shaped .

1 Introduction

Deep generative models have shown outstanding sample quality in a wide variety of modalities. Generative Adversarial Networks (GANs) (Goodfellow et al., 2014; Karras et al., 2021), autoregressive models (Ramesh et al., 2021), Variational Autoencoders (Kingma & Welling, 2013; Rezende et al., 2014), Normalizing Flows (Grathwohl et al., 2018; Chen et al., 2019) and energy-based models (Xiao et al., 2020) show impressive abilities to synthesize objects. However, GANs are not robust to the choice of architecture and optimization method (Arjovsky et al., 2017; Gulrajani et al., 2017; Karras et al., 2019; Brock et al., 2018), and they often fail to cover modes in data distribution (Zhao et al., 2018; Thanh-Tung & Tran, 2020). Likelihood-based models avoid mode collapse but may overestimate the probability in low-density regions (Zhang et al., 2021).

30

31

CT

Ct-1

3t

gss(xt[co)

b

DDPM

（ctct-1)

20

Ct-1

Ct

Recently, diffusion probabilistic models (Sohl-Dickstein et al., 2015; Ho et al., 2020) have received a lot of attention. These models generate samples using a trained Markov process that starts with white noise and iteratively removes noise from the sample. Recent works have shown that diffusion models can generate samples comparable in quality or even better than GANs (Song et al., 2020b; Dhariwal & Nichol, 2021), while they do not suffer from mode collapse by design, and also they have a log-likelihood comparable to autoregressive models (Kingma et al., 2021). Moreover, diffusion models show these results in various modalities such as images (Saharia et al., 2021), sound (Popov et al., 2021; Liu et al., 2022) and shapes (Luo & Hu, 2021; Zhou et al., 2021).

The main principle of diffusion models is to destroy information during the forward process and then restore it during the reverse process. In conventional diffusion models like denoising diffusion probabilistic models (DDPM) destruction of information occurs through the injection of Gaussian noise, which is reasonable for some types of data, such as images. However, for data distributed on manifolds, bounded volumes, or with other features, the injection of Gaussian noise can be unnatural, breaking the data structure. Unfortunately, it is not clear how to replace the noise distribution within traditional diffusion models. The problem is that we have to maintain a connection between the distributions defining the Markov noising process that gradually destroys information and its marginal distributions. While some papers explore other distributions, such as delta functions (Bansal et al., 2022) or Gamma distribution (Nachmani et al., 2021), they provide ad hoc solutions for special cases that are not easily generalized.

In this paper, we present Star-Shaped Denoising Diffusion Probabilistic Models (SS-DDPM), a new approach that generalizes Gaussian DDPM to an exponential family of noise distributions. In SS-DDPM, one only needs to define marginal distributions at each diffusion step (see Figure 1). We provide a derivation of SS-DDPM, design efficient sampling and training algorithms, and show its equivalence to DDPM (Ho et al., 2020) in the case of Gaussian noise. Then, we outline a number of practical considerations that aid in training and applying SS-DDPMs. In Section 5, we demonstrate the ability of SS-DDPM to work with distributions like von Mises-Fisher, Dirichlet and Wishart. Finally, we evaluate SS-DDPM on image and text generation. Categorical SS-DDPM matches the performance of Multinomial Text Diffusion (Hoogeboom et al., 2021) on the text8 dataset, while our Beta diffusion model achieves results, comparable to a Gaussian DDPM on CIFAR-10.

2 Theory

2.1 DDPMs

We start with a brief introduction of DDPMs. The Gaussian DDPM (Ho et al., 2020) is defined as a forward (diffusion) process q DDPM ( x 0: T ) and a corresponding reverse (denoising) process p DDPM θ ( x 0: T ) . The forward process is defined as a Markov chain with Gaussian conditionals:

where q ( x 0 ) is the data distribution. Parameters β t are typically chosen in advance and fixed, defining the noise schedule of the diffusion process. The noise schedule is chosen in such a way that the final x T no longer depends on x 0 and follows a standard Gaussian distribution q DDPM ( x T ) = N ( x T ; 0 , I ) .

30

21

32

CT-1

CT

DDPM(ct|Ct-1)

q

Ct-1

Ct

CT

pe(ct-1|ct)

G2

GT-1

GT

(a) Denoising Diffusion Probabilistic Models

(b) Star-Shaped Denoising Diffusion Probabilistic Models

The reverse process p DDPM θ ( x 0: T ) follows a similar structure and constitutes the generative part of the model:

The forward process q DDPM ( x 0: T ) of DDPM is typically fixed, and all the parameters of the model are contained in the generative part of the model p DDPM θ ( x 0: T ) . These parameters are tuned to maximize the variational lower bound (VLB) on the likelihood of the training data:

One of the main challenges in defining DDPMs is the computation of the posterior q DDPM ( x t -1 | x t , x 0 ) . Specifically, the transition probabilities q DDPM ( x t | x t -1 ) have to be defined in such a way that this posterior is tractable. Specific DDPM-like models are available for Gaussian (Ho et al., 2020), Categorical (Hoogeboom et al., 2021) and Gamma (Kawar et al., 2022) distributions. Defining such models remains challenging in more general cases.

2.2 Star-Shaped DDPMs

As previously discussed, extending the DDPMs to other distributions poses significant challenges. In light of these difficulties, we propose to construct a model that only relies on marginal distributions q ( x t | x 0 ) in its definition and the derivation of the loss function.

We define star-shaped diffusion as a non-Markovian forward process q SS ( x 0: T ) that has the following structure:

where q ( x 0 ) is the data distribution. We note that in contrast to DDPM all noisy variables x t are conditionally independent given x 0 instead of constituting a Markov chain. This structure of the forward process allows us to utilize other noise distributions, which we discuss in more detail later.

2.3 Defining the reverse model

In DDPMs the true reverse model q DDPM ( x 0: T ) has a Markovian structure (Ho et al., 2020), allowing for an efficient sequential generation algorithm:

For the star-shaped diffusion, however, the Markovian assumption breaks:

Consequently, we now need to approximate the true reverse process by a parametric model which is conditioned on the whole tail x t : T .

It is crucial to use the whole tail x t : T rather than just one variable x t when predicting x t -1 in a star-shaped model. As we show in Appendix B, if we try to approximate the true reverse process with a Markov model, we introduce a substantial irreducible gap into the variational lower bound. Such a sampling procedure fails to generate realistic samples, as can be seen in Figure 3.

Intuitively, in DDPMs the information about x 0 that is contained in x t +1 is nested into the information about x 0 that is contained in x t .

That is why knowing x t allows us to discard x t +1 . In star-shaped diffusion, however, all variables contain independent pieces of information about x 0 and should all be taken into account when making predictions.

We can write down the variational lower bound as follows:

With this VLB, we only need the marginal distributions q ( x t -1 | x 0 ) to define and train the model, which allows us to use a wider variety of noising distributions. Since conditioning the predictive model p θ ( x t -1 | x t : T ) on the whole tail x t : T is typically impractical, we propose a more efficient way to implement the reverse process next.

2.4 Efficient tail conditioning

Instead of using the full tail x t : T , we would like to define some statistic G t = G t ( x t : T ) that would extract all information about x 0 from the tail x t : T . Formally speaking, we call G t a sufficient tail statistic if the following equality holds:

One way to define G t is to concatenate all the variables x t : T into a single vector. This, however, is impractical, as its dimension would grow with the size of the tail T -t +1 .

The Pitman-Koopman-Darmois (Pitman, 1936) theorem (PKD) states that exponential families admit a sufficient statistic with constant dimensionality. It also states that no other distribution admits one: if such a statistic were to exist, the distribution has to be a member of the exponential family. Inspired by the PKD, we turn to the exponential family of distributions. In the case of star-shaped diffusion, we cannot apply the PKD directly, as it was formulated for i.i.d. samples and our samples are not identically distributed. However, we can still define a sufficient tail statistic G t for a specific subset of the exponential family, which we call an exponential family with linear parameterization :

Theorem 1. Assume the forward process of a star-shaped model takes the following form:

Let G t be a tail statistic, defined as follows:

Then, G t is a sufficient tail statistic:

Here definition (13) is the standard definition of the exponential family, where h t ( x t ) is the base measure , η t ( x 0 ) is the vector of natural parameters with corresponding sufficient statistics T ( x t ) , and Ω t ( x 0 ) is the log-partition function . The key assumption added is the linear parameterization of the natural parameters (14). We provide the proof in Appendix C. When A t is scalar, we denote it as a t instead.

For the most part, the premise of Theorem 1 restricts the parameterization of the distributions rather than the family of the distributions involved. As we discuss in Appendix F, we found it easy to come up with linear parameterization for a wide range of distributions in the exponential family. For example, we can obtain a linear parameterization for the Beta distribution q ( x t | x 0 ) = Beta( x t ; α t , β t ) using x 0 as the mode of the distribution and introducing a new concentration parameter ν t :

In this case, η t ( x 0 ) = ν t x 0 , T ( x t ) = log x t 1 -x t , and we can use equation (15) to define the sufficient tail statistic G t . We provide more examples in Appendix F. We also provide an implementation-ready reference sheet for a wide range of distributions in the exponential family in Table 6.

We suspect that, just like in PKD, this trick is only possible for a subset of the exponential family. In the general case, the dimensionality of the sufficient tail statistic G t would have to grow with the size of the tail x t : T . It is still possible to apply SS-DDPM in this case, however, crafting the (now only approximately) sufficient statistic G t would require more careful consideration and we leave it for future work.

2.5 Final model definition

To maximize the VLB (11), each step of the reverse process should approximate the true reverse distribution:

Similarly to DDPM (Ho et al., 2020), we choose to approximate q SS ( x 0 | x t : T ) with a delta function centered at the prediction of some model x θ ( G t ( x t : T ) , t ) . This results in the following definition of the reverse process of SS-DDPM:

The distribution p SS θ ( x 0 | x 1: T ) can be fixed to some small-variance distribution p SS θ ( x 0 | ˆ x 0 ) centered at the final prediction ˆ x 0 = x θ ( G 1 ( x 1: T ) , 1) , similar to the dequantization term, commonly used in DDPM. If this distribution has no trainable parameters, the corresponding term can be removed from the training objective. This dequantization distribution would then only be used for log-likelihood estimation and, optionally, for sampling.

Together with the forward process (7) and the VLB objective (11), this concludes the general definition of the SS-DDPM model. The model structure is illustrated in Figure 2. The corresponding training and sampling algorithms are provided in Algorithms 1 and 2.

The resulting model is similar to DDPM in spirit. We follow the same principles when designing the forward process: starting from a low-variance distribution, centered at x 0 at t = 1 , we gradually increase the entropy of the distribution q SS ( x t | x 0 ) until there is no information shared between x 0 and x t at t = T .

We provide concrete definitions for Beta, Gamma, Dirichlet, von Mises, von Mises-Fisher, Wishart, Gaussian and Categorical distributions in Appendix F.

Algorithm 1 SS-DDPM training

repeat

Algorithm 2 SS-DDPM sampling

2.6 Duality between star-shaped and Markovian diffusion

While the variables x 1: T follow a star-shaped diffusion process, the corresponding tail statistics G 1: T form a Markov chain:

since x t is conditionally independent from G t +2: T given G t +1 (see Appendix E for details). Moreover, we can rewrite the probabilistic model in terms of G t and see that variables ( x 0 , G 1: T ) form a (not necessarily Gaussian) DDPM.

In the case of Gaussian distributions, this duality makes SS-DDPM and DDPM equivalent. This equivalence can be shown explicitly:

Theorem 2. Let α DDPM t define the noising schedule for a DDPM model (1-2) via β t = ( α DDPM t -1 -α DDPM t ) /α DDPM t -1 . Let q SS ( x 0: T ) be a Gaussian SS-DDPM forward process with the following noising schedule and sufficient tail statistic:

Then the tail statistic G t follows a Gaussian DDPM noising process q DDPM ( x 0: T ) | x 1: T = G 1: T defined by the schedule α DDPM t . Moreover, the corresponding reverse processes and VLB objectives are also equivalent.

We show this equivalence in Appendix D. We make use of this connection when choosing the noising schedule for other distributions.

This equivalence means that SS-DDPM is a direct generalization of Gaussian DDPM. While admitting the Gaussian case, SS-DDPM can also be used to implicitly define a non-Gaussian DDPM in the space of sufficient tail statistics for a wide range of distributions.

3 Practical considerations

While the model is properly defined, there are several practical considerations that are important for the efficiency of star-shaped diffusion.

Choosing the right schedule It is important to choose the right noising schedule for a SS-DDPM model. It significantly depends on the number of diffusion steps T and behaves differently given different noising schedules, typical to DDPMs. This is illustrated in Figure 4, where we show the noising schedules for Gaussian SS-DDPMs that are equivalent to DDPMs with the same cosine schedule.

























t

ss

t

,

T

=50

ss

t

,

T

=100

ss

t

,

T

=250

ss

t

,

T

=1000

ss

t

,

T

=4000

DDPM

t

t

/

T

Since the variables G t follow a DDPM-like process, we would like to somehow reuse those DDPM noising schedules that are already known to work well. For Gaussian distributions, we can transform a DDPM noising schedule into the corresponding SS-DDPM noising schedule analytically by equating I ( x 0 ; G t ) = I ( x 0 ; x DDPM t ) . In general case, we look for schedules that have approximately the same level of mutual information I ( x 0 ; G t ) as the corresponding mutual information I ( x 0 ; x DDPM t ) for a DDPM model for all timesteps t . We estimate the mutual information using Kraskov (Kraskov et al., 2004) and DSIVI (Molchanov et al., 2019) estimators and build a look-up table to match the noising schedules. This procedure is described in more detail in Appendix G. The resulting schedule for the Beta SS-DDPM is illustrated in Figure 5. Note how with the right schedule appropriately normalized tail statistics G t look and function similarly to the samples x t from the corresponding Gaussian DDPM. We further discuss this in Appendix H.

Implementing the sampler During sampling, we can grow the tail statistic G t without any overhead, as described in Algorithm 2. However, during training, we need to sample the tail statistic for each object to estimate the loss function. For this we need to sample the full tail x t : T from the forward process q SS ( x t : T | x 0 ) , and then compute the tail statistic G t . In practice, this does not add a noticeable overhead and can be computed in parallel to the training process if needed.

Reducing the number of steps We can sample from DDPMs more efficiently by skipping some timestamps. This wouldn't work for SS-DDPM, because changing the number of steps would require changing the noising schedule and, consequently, retraining the model.

However, we can still use a similar trick to reduce the number of function evaluations. Instead of skipping some timestamps x t 1 +1: t 2 -1 , we can draw them from the forward process using the current prediction x θ ( G t 2 , t 2 ) , and then use these samples to obtain the tail statistic G t 1 . For Gaussian SS-DDPM this is equivalent to skipping these timestamps in the corresponding DDPM. In general case, it amounts to approximating the reverse process with a different reverse process:

We observe a similar dependence on the number of function evaluations for SS-DDPMs and DDPMs.

Time-dependent tail normalization As defined in Theorem 1, the tail statistics can have vastly different scales for different timestamps. The values of coefficients a t can range from thousandths when t approaches T to thousands when t approaches zero. To make the tail statistics suitable for use in neural networks, proper normalization is crucial. In most cases, we collect the time-dependent means and variances of the tail statistics across the training dataset and normalize the tail statistics to zero mean and unit variance. We further discuss this issue in Appendix H.

Architectural choices To make training the model easier, we make some minor adjustments to the neural network architecture and the loss function.

Our neural networks x θ ( G t , t ) take the tail statistic G t as an input and are expected to produce an estimate of x 0 as an output. In SS-DDPM the data x 0 might lie on some manifold, like the unit sphere or the space of positive definite matrices. Therefore, we need to map the neural network output to that manifold. We do that on a case-by-case basis, as described in Appendices I-L.

Different terms of the VLB can have drastically different scales. For this reason, it is common practice to train DDPMs with a modified loss function like L simple rather than the VLB to improve the stability of training (Ho et al., 2020). Similarly, we can optimize a reweighted variational lower bound when training SS-DDPMs.

4 Related works

Our work builds upon Denoising Diffusion Probabilistic Models (Ho et al., 2020). Interest in diffusion models has increased recently due to their impressive results in image (Ho et al., 2020; Song et al., 2020b; Dhariwal & Nichol, 2021) and audio (Popov et al., 2021; Liu et al., 2022) generation.

SS-DDPM is most closely related to DDPM. Like DDPM, we only rely on variational inference when defining and working with our model. SS-DDPM can be seen as a direct generalization of DDPM and essentially is a recipe for defining DDPMs with non-Gaussian distributions. The underlying non-Gaussian DDPM is constructed implicitly and can be seen as dual to the star-shaped formulation that we use throughout the paper.

Other ways to construct non-Gaussian DDPMs include Binomial diffusion (Sohl-Dickstein et al., 2015), Multinomial diffusion (Hoogeboom et al., 2021) and Gamma diffusion (Nachmani et al., 2021). Each of these works presents a separate derivation of the resulting objective, and extending them to other distributions is not straightforward. On the other hand, SS-DDPM provides a single recipe for a wide range of distributions.

DDPMs have several important extensions. Song et al. (2020a) provide a family of non-Markovian diffusions that all result in the same training objective as DDPM. One of them, denoted Denoising Diffusion Implicit Model (DDIM), results in an efficient deterministic sampling algorithm that requires a much lower number of function evaluations than conventional stochastic sampling. Their derivations also admit a star-shaped forward process (at σ 2 t = 1 -α t -1 ), however, the model is not studied in the star-shaped regime. Their reverse process remains Markovian, which we show to not be sufficient to invert a star-shaped forward process. Denoising Diffusion Restoration Models (Kawar et al., 2022) provide a way to solve general linear inverse problems using a trained DDPM. They can be used for image restoration, inpainting, colorization and other conditional generation tasks. Both DDIMs and DDRMs rely on the explicit form of the underlying DDPM model and are derived for Gaussian diffusion. Extending these models to SS-DDPMs is a promising direction for future work.

Song et al. (2020b) established the connection between DDPMs and models based on score matching. This connection gives rise to continuous-time variants of the models, deterministic solutions and more precise density estimation using ODEs. We suspect that a similar connection might hold for SS-DDPMs as well, and it can be investigated further in future works.

Other works that present diffusion-like models with other types of noise or applied to manifold data, generally stem from score matching rather than variational inference. Flow Matching (Lipman et al., 2022) is an alternative probabilistic framework that works with any differentiable degradation process. De Bortoli et al. (2022) and Huang et al. (2022) extended score matching to Riemannian manifolds, and Chen & Lipman (2023) proposed Riemannian Flow Matching. Bansal et al. (2022) proposed Cold Diffusion, a non-probabilistic approach to reversing general degradation processes.

To the best of our knowledge, for the first time, an approach for constructing diffusion without a consecutive process was proposed by Rissanen et al. (2022) (IHDM) and further expanded on by Daras et al. (2022) and Hoogeboom & Salimans (2022). IHDM uses a similar star-shaped structure that results in a similar variational lower bound. Adding a deterministic process based on the heat equation allows the authors to keep the reverse process Markovian without having to introduce the tail statistics. As IHDM heavily relies on blurring rather than adding noise, the resulting diffusion dynamics become very different. Conceptually our work is much closer to DDPM than IHDM.

5 Experiments

We evaluate SS-DDPM with different families of noising distributions. The experiment setup, training details and hyperparameters are listed in Appendices I-L.

Synthetic data We consider two examples of starshaped diffusion processes with Dirichlet and Wishart noise to generate data from the probabilistic simplex and from the manifold of p.d. matrices respectively. We compare them to DDPM, where the predictive network x θ ( x t , t ) is parameterized to always satisfy the manifold constraints. As seen in Table 1, using the appropriate distribution results in a better approximation. The data and

modeled distributions are illustrated in Table 3. This shows the ability of SS-DDPM to work with different distributions and generate data from exotic domains.

Geodesic data We apply SS-DDPM to a geodesic dataset of fires on the Earth's surface (EOSDIS, 2020) using a three-dimensional von Mises-Fisher distribution. The resulting samples and the source data are illustrated in Table 3. We find that SS-DDPM is not too sensitive to the distribution family and can fit data in different domains.

Discrete data Categorical SS-DDPM is similar to Multinomial Text Diffusion (MTD, (Hoogeboom et al., 2021)). However, unlike in the Gaussian case, these models are not strictly equivalent. We follow a similar setup to MTD and apply Categorical SS-DDPM to unconditional text generation on the text8 dataset (Mahoney, 2011). As shown in Table 2, SS-DDPM achieves similar results to MTD, allowing to use different distributions in a unified manner. While D3PM (Austin et al., 2021) provides some improvements, we follow a simpler setup from MTD. We expect the improvements from D3PM to directly apply to Categorical SS-DDPM.

Image data Finally, we evaluate SS-DDPM on CIFAR-10. Since the training data is constrained to a [0 , 1] segment, we use Beta distributions. We evaluate our model with various numbers of generation steps, as described in equation (25), and report the resulting Fréchet Inception Distance (FID, (Heusel et al., 2017)) in Figure 6. Beta SS-DDPM achieves comparable quality with the Improved DDPM (Nichol & Dhariwal, 2021) and is slightly better on lower numbers of steps. As expected, DDIM performs better when the number of diffusion steps is low, but both SS-DDPM and DDPM outperform DDIM on longer runs. The best FID score achieved by Beta SS-DDPM is 3 . 17 . Although the FID curves for DDPM and DDIM do not achieve this score in Figure 6, Ho et al. (2020) reported an FID score of 3 . 17 for 1000 DDPM steps, meaning that SS-DDPM performs similarly to DDPM in this setting.

10 2

10 3

sampling steps

2

4

6

8

10

12

14

16

FID

DDPM

DDIM

Beta SS-DDPM

Dirichlet

Wishart

Real

Generated

8

6 Conclusion

We propose an alternative view on diffusion-like probabilistic models. We reveal the duality between star-shaped and Markov diffusion processes that allows us to go beyond Gaussian noise by switching to a star-shaped formulation. It allows us to define diffusion-like models with arbitrary noising distributions and to establish diffusion processes on specific manifolds. We propose an efficient way to construct a reverse process for such models in the case when the noising process lies in a general subset of the exponential family and show that star-shaped diffusion models can be trained on a variety of domains with different noising distributions. On image data, star-shaped diffusion with Beta distributed noise attains comparable performance to Gaussian DDPM, challenging the optimality of Gaussian noise in this setting. The star-shaped formulation opens new applications of diffusion-like probabilistic models, especially for data from exotic domains where domain-specific non-Gaussian diffusion is more appropriate.

References

Arjovsky, M., Chintala, S., and Bottou, L. Wasserstein gan. arxiv 2017. arXiv preprint arXiv:1701.07875 , 30:4, 2017.

Austin, J., Johnson, D. D., Ho, J., Tarlow, D., and van den Berg, R. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems , 34:1798117993, 2021.

Bansal, A., Borgnia, E., Chu, H.-M., Li, J. S., Kazemi, H., Huang, F., Goldblum, M., Geiping, J., and Goldstein, T. Cold diffusion: Inverting arbitrary image transforms without noise. arXiv preprint arXiv:2208.09392 , 2022.

Brock, A., Donahue, J., and Simonyan, K. Large scale gan training for high fidelity natural image synthesis. arXiv preprint arXiv:1809.11096 , 2018.

Chen, R. T. and Lipman, Y. Riemannian flow matching on general geometries. arXiv preprint arXiv:2302.03660 , 2023.

Chen, R. T., Behrmann, J., Duvenaud, D. K., and Jacobsen, J.-H. Residual flows for invertible generative modeling. Advances in Neural Information Processing Systems , 32, 2019.

Daras, G., Delbracio, M., Talebi, H., Dimakis, A. G., and Milanfar, P. Soft diffusion: Score matching for general corruptions. arXiv preprint arXiv:2209.05442 , 2022.

De Bortoli, V., Mathieu, E., Hutchinson, M., Thornton, J., Teh, Y. W., and Doucet, A. Riemannian score-based generative modeling. arXiv preprint arXiv:2202.02763 , 2022.

Dhariwal, P. and Nichol, A. Diffusion models beat gans on image synthesis. arXiv preprint arXiv:2105.05233 , 2021.

EOSDIS. Land, atmosphere near real-time capability for eos (lance) system operated by nasa's earth science data and information system (esdis). https://earthdata.nasa.gov/earth-observation-data/nearreal-time/firms/active-fire-data, 2020.

Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial nets. Advances in neural information processing systems , 27, 2014.

Grathwohl, W., Chen, R. T., Bettencourt, J., Sutskever, I., and Duvenaud, D. Ffjord: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint arXiv:1810.01367 , 2018.

Gulrajani, I., Ahmed, F., Arjovsky, M., Dumoulin, V., and Courville, A. C. Improved training of wasserstein gans. Advances in neural information processing systems , 30, 2017.

He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. corr abs/1512.03385 (2015), 2015.

Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems , 30, 2017.

Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. arXiv preprint arXiv:2006.11239 , 2020.

Hoogeboom, E. and Salimans, T. Blurring diffusion models. arXiv preprint arXiv:2209.05557 , 2022.

Hoogeboom, E., Nielsen, D., Jaini, P., Forré, P., and Welling, M. Argmax flows and multinomial diffusion: Learning categorical distributions. Advances in Neural Information Processing Systems , 34:12454-12465, 2021.

Huang, C.-W., Aghajohari, M., Bose, J., Panangaden, P., and Courville, A. C. Riemannian diffusion models. Advances in Neural Information Processing Systems , 35:2750-2761, 2022.

Karras, T., Laine, S., and Aila, T. A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pp. 4401-4410, 2019.

Karras, T., Aittala, M., Laine, S., Härkönen, E., Hellsten, J., Lehtinen, J., and Aila, T. Alias-free generative adversarial networks. Advances in Neural Information Processing Systems , 34, 2021.

Kawar, B., Elad, M., Ermon, S., and Song, J. Denoising diffusion restoration models. In Advances in Neural Information Processing Systems , 2022.

Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In International Conference on Learning Representations , 2015.

Kingma, D. P. and Welling, M. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

Kingma, D. P., Salimans, T., Poole, B., and Ho, J. Variational diffusion models. arXiv preprint arXiv:2107.00630 , 2, 2021.

Kraskov, A., Stögbauer, H., and Grassberger, P. Estimating mutual information. Physical review E , 69(6):066138, 2004.

Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747 , 2022.

Liu, S., Su, D., and Yu, D. Diffgan-tts: High-fidelity and efficient text-to-speech with denoising diffusion gans. arXiv preprint arXiv:2201.11972 , 2022.

Luo, S. and Hu, W. Diffusion probabilistic models for 3d point cloud generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 2837-2845, 2021.

Mahoney, M. Large text compression benchmark, 2011. URL http://www.mattmahoney.net/ dc/text.html .

Molchanov, D., Kharitonov, V., Sobolev, A., and Vetrov, D. Doubly semi-implicit variational inference. In The 22nd International Conference on Artificial Intelligence and Statistics , pp. 2593-2602. PMLR, 2019.

Nachmani, E., Roman, R. S., and Wolf, L. Denoising diffusion gamma models. arXiv preprint arXiv:2110.05948 , 2021.

Nichol, A. and Dhariwal, P. Improved denoising diffusion probabilistic models. arXiv preprint arXiv:2102.09672 , 2021.

Pitman, E. J. G. Sufficient statistics and intrinsic accuracy. Mathematical Proceedings of the Cambridge Philosophical Society , 32(4):567-579, 1936. doi: 10.1017/S0305004100019307.

Popov, V., Vovk, I., Gogoryan, V., Sadekova, T., and Kudinov, M. Grad-tts: A diffusion probabilistic model for text-to-speech. In International Conference on Machine Learning , pp. 8599-8608. PMLR, 2021.

Ramesh, A., Pavlov, M., Goh, G., Gray, S., Voss, C., Radford, A., Chen, M., and Sutskever, I. Zeroshot text-to-image generation. In International Conference on Machine Learning , pp. 8821-8831. PMLR, 2021.

Rezende, D. J., Mohamed, S., and Wierstra, D. Stochastic backpropagation and approximate inference in deep generative models. In International conference on machine learning , pp. 1278-1286. PMLR, 2014.

Rissanen, S., Heinonen, M., and Solin, A. Generative modelling with inverse heat dissipation. arXiv preprint arXiv:2206.13397 , 2022.

Saharia, C., Ho, J., Chan, W., Salimans, T., Fleet, D. J., and Norouzi, M. Image super-resolution via iterative refinement. arXiv preprint arXiv:2104.07636 , 2021.

Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning , pp. 2256-2265. PMLR, 2015.

Song, J., Meng, C., and Ermon, S. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502 , 2020a.

Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020b.

Thanh-Tung, H. and Tran, T. Catastrophic forgetting and mode collapse in gans. In 2020 International Joint Conference on Neural Networks (IJCNN) , pp. 1-10. IEEE, 2020.

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Xiao, Z., Kreis, K., Kautz, J., and Vahdat, A. Vaebm: A symbiosis between variational autoencoders and energy-based models. arXiv preprint arXiv:2010.00654 , 2020.

Yin, M. and Zhou, M. Semi-implicit variational inference. In International Conference on Machine Learning , pp. 5660-5669. PMLR, 2018.

Zhang, L., Goldstein, M., and Ranganath, R. Understanding failures in out-of-distribution detection with deep generative models. In International Conference on Machine Learning , pp. 12427-12436. PMLR, 2021.

Zhao, S., Ren, H., Yuan, A., Song, J., Goodman, N., and Ermon, S. Bias and generalization in deep generative models: An empirical study. Advances in Neural Information Processing Systems , 31, 2018.

Zhou, L., Du, Y., and Wu, J. 3d shape generation and completion through point-voxel diffusion. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pp. 5826-5835, 2021.

A Variational lower bound for the SS-DDPM model

B True reverse process for Markovian and Star-Shaped DDPM

If the forward process is Markovian, the corresponding true reverse process is Markovian too:

For star-shaped models, however, the reverse process has a general structure that cannot be reduced further:

In this case, a Markovian reverse process can be a very poor approximation to the true reverse process. Choosing such an approximation adds an irreducible gap to the variational lower bound:

Intuitively, there is little information shared between x t -1 and x t , as they are conditionally independent given x 0 . Therefore, we would expect the distribution q SS ( x t -1 | x t ) to have a much higher entropy than the distribution q SS ( x t -1 | x t : T ) , making the irreducible gap (35) large. The dramatic effect of this gap is illustrated in Figure 3.

This gap can also be computed analytically for Gaussian DDPMs when the data is coming from a standard Gaussian distribution q ( x 0 ) = N ( x 0 ; 0 , 1) . According to equations (34-35), the best Markovian reverse process in this case is p θ ( x t -1 | x t ) = q SS ( x t -1 | x t ) . It results in the following value of the variational lower bound:

If the reverse process is matched exactly (and has a general structure), the variational lower bound reduces to the negative entropy of the data distribution:

As shown in Figure 7, the Markovian approximation adds a substantial irreducible gap to the variational lower bound.

C Proof of Theorem 1

























SS





We start by establishing the following lemma. It would allow us to change the variables in the condition of a conditional distribution. Essentially, we would like to show that if a variable y is only dependent on x through some function h ( x ) , we can write the distribution p ( y | x ) as p ( y | z ) | z = h ( x ) .

Lemma 1. Assume random variables x and y follow a joint distribution p ( x, y ) = p ( x ) p ( y | x ) , where p ( y | x ) is given by f ( y, h ( x )) . Then p ( y | x ) = p ( y | z ) | z = h ( x ) , where the variable z is defined as z = h ( x ) .

Proof. We can write down the joint distribution over all variables as follows:

By integrating x out, we get

Theorem 1. Assume the forward process of a star-shaped model takes the following form:

Let G t be a tail statistic, defined as follows:

Then, G t is a sufficient tail statistic:

Finally, we obtain

Proof.

D Gaussian SS-DDPM is equivalent to Gaussian DDPM

Theorem 2. Let α DDPM t define the noising schedule for a DDPM model (1-2) via β t = ( α DDPM t -1 -α DDPM t ) /α DDPM t -1 . Let q SS ( x 0: T ) be a Gaussian SS-DDPM forward process with the following noising schedule and sufficient tail statistic:

Then the tail statistic G t follows a Gaussian DDPM noising process q DDPM ( x 0: T ) | x 1: T = G 1: T defined by the schedule α DDPM t . Moreover, the corresponding reverse processes and VLB objectives are also equivalent.

Proof. We start by listing the necessary definitions and expressions used in the Gaussian DDPM model (Ho et al., 2020).

In this notation, the variational lower bound for DDPM can be written as follows (here we omit the term corresponding to x T because it's either zero or a small constant):

For convenience, we additionally define α DDPM T +1 = 0 . Since all the involved distributions are typically isotropic, we consider a one-dimensional case without the loss of generality.

First, we show that the forward processes q SS ( x 0 , G 1: T ) and q DDPM ( x 0: T ) are equivalent.

Since G t is a linear combination of independent (given x 0 ) Gaussian random variables, its distribution (given x 0 ) is also Gaussian. All the terms conveniently cancel out, and we recover the same form as the forward process of DDPM.

Since both processes share the same q ( x 0 ) , and have the same Markovian structure, matching these marginals sufficiently shows that the processes are equivalent:

Consequently, the posteriors are matching too:

Next, we show that the reverse processes p SS θ ( G t -1 | G t ) and p DDPM θ ( x t -1 | x t ) are the same. Since G t in SS-DDPM is the same as x t in DDPM, we can assume the predictive models x SS θ ( G t , t ) and x DDPM θ ( x t , t ) to coincide at x t = G t .

where x t -1 ∼ p SS θ ( x t -1 | G t ) = q SS ( x t -1 | x 0 ) | x 0 = x θ ( G t ,t ) . Therefore, p SS θ ( G t -1 | G t ) is also a Gaussian distribution. Let's take care of the mean first:

Therefore,

Finally, we show that the variational lower bounds are the same.

Since G 1 from the star-shaped model is the same as x 1 from the DDPM model, the first term log p SS θ ( x 0 | G 1 ) coincides with log p DDPM θ ( x 0 | x 1 ) | x 1 = G 1 .

Therefore, we finally obtain L SS ( θ ) = L DDPM ( θ ) .

Now, let's derive the variance:

E Duality between star-shaped and Markovian diffusion: general case

We assume that T ( x t ) and all matrices A t (except possibly A T ) are invertible. Then, there is a bijection between the set of tail statistics G t : T and the tail x t : T .

First, we show that the variables G 1: T form a Markov chain.

The last equation holds since the reverse of a Markov chain is also a Markov chain.

This means that a star-shaped diffusion process on x 1: T implicitly defines some Markovian diffusion process on the tail statistics G 1: T . Due to the definition of G t -1 = G t + A T t -1 T ( x t -1 ) , we can write down the following factorization of that process:

The posteriors q ( G t -1 | G t , x 0 ) can then be computed by a change of variables:

This also allows us to define the reverse model like it was defined in DDPM:

This definition is consistent with the reverse process of SS-DDPM p θ ( x t -1 | x t : T ) = q ( x t -1 | x 0 ) | x 0 = x θ ( G t ,t ) . Now that both the forward and reverse processes are defined, we can write down the corresponding variational lower bound. Because the model is structured exactly like a DDPM, the variational lower bound is going to look the same. However, we can show that it is equivalent to the variational lower bound of SS-DDPM:

As we can see, there are two equivalent ways to write down the model. One is SS-DDPM, a star-shaped diffusion model, where we only need to define the marginal transition probabilities q ( x t | x 0 ) . Another way is to rewrite the model in terms of the tail statistics G t : T . This way we obtain a non-Gaussian DDPM that is implicitly defined by the star-shaped model. Because of this equivalence, we see the star-shaped diffusion of x 1: T and the Markovian diffusion of G 1: T as a pair of dual processes.

F SS-DDPM in different families

The main principles for designing a SS-DDPM model are similar to designing a DDPM model. When t goes to zero, we wish to recover the Dirac delta, centered at x 0 :

When t goes to T , we wish to obtain some standard distribution that doesn't depend on x 0 :

In exponential families with linear parameterization of the natural parameter η t ( x 0 ) = a t f ( x 0 ) + b t , we can define the schedule by choosing the parameters a t and b t that satisfy conditions (100-101). After that, we can use Theorem 1 to define the tail statistic G t ( x t : T ) using the sufficient statistic T ( x t ) of the corresponding family. However, as shown in the following sections, in some cases linear parameterization admits a simpler sufficient tail statistic.

The following sections contain examples of defining the SS-DDPM model for different families. These results are summarized in Table 6.

F.1 Gaussian

F.2 Beta

Since Gaussian SS-DDPM is equivalent to a Markovian DDPM, it is natural to directly reuse the schedule from a Markovian DDPM. As we show in Theorem 2, given a Markovian DDPM defined by α DDPM t , the following parameterization will produce the same process in the space of tail statistics:

The KL divergence is computed as follows:

There are many ways to define a noising schedule. We choose to fix the mode of the distribution at x 0 and introduce a concentration parameter ν t , parameterizing α t = 1 + ν t x 0 and β t = 1 + ν t (1 -x 0 ) . By setting ν t to zero, we recover a uniform distribution, and by setting it to infinity we obtain the Dirac delta, centered at x 0 . Generally, the Beta distribution has a two-dimensional sufficient statistic T ( x ) = ( log x log(1 -x ) ) . However, under this parameterization, we can derive a one-dimensional tail statistic:

Therefore we can use η t ( x 0 ) = ν t x 0 and T ( x ) = log x 1 -x to define the tail statistic:

The KL divergence can then be calculated as follows:

t=0

t=7

t =14

t=21

t=28

t=35

t = 42

t =49

t=63

F.3 Dirichlet

Similarly to the Beta distribution, we choose to fix the mode of the distribution at x 0 and introduce a concentration parameter ν t , parameterizing α k t = 1 + ν t x k 0 . This corresponds to using the natural parameter η t ( x 0 ) = ν t x 0 . By setting ν t to zero, we recover a uniform distribution, and by setting it to infinity we obtain the Dirac delta, centered at x 0 . We can use the sufficient statistic T ( x ) = (log x 1 , . . . , log x K ) T to define the tail statistic:

The KL divergence can then be calculated as follows:

F.4 Categorical

In this case, we mimic the definition of the categorical diffusion model, used in D3PM. The noising process is parameterized by the probability vector p t = x 0 Q t . By setting Q 0 to identity, we recover the Dirac delta, centered at x 0 .

In this parameterization, the natural parameter admits linearization (after some notation abuse):

log is taken element-wise, and we assume 0 log 0 = 0 .

The sufficient statistic here is a vector x T t . Therefore, the tail statistic can be defined as follows:

The KL divergence can then be calculated as follows:

Note that, unlike in the Gaussian case, the categorical star-shaped diffusion is not equivalent to the categorical DDPM. The main difference here is that the input G t to the predictive model x θ ( G t ) is now a continuous vector instead of a one-hot vector.

As discussed in Section H, proper normalization is important for training SS-DDPMs. In the case of Categorical distributions, we can use the SoftMax( · ) function to normalize the tail statistics without breaking sufficiency:

To see this, we can retrace the steps of the proof of Theorem 1:

Therefore,

Also, Categorical distributions admit a convenient way to compute fractions involving q ( G t | x 0 ) :

This allows us to define the likelihood term p ( x 0 | x 1: T ) as follows:

where ˜ p ( x 0 ) can be defined as the frequency of the token x 0 in the dataset.

It also allows us to estimate the mutual information between x 0 and G t :

It can be estimated using Monte Carlo. We use it when defining the noising schedule for Categorical SS-DDPM.

F.5 von Mises

The von Mises distribution has two parameters, the mode x and the concentration κ . It is natural to set the mode of the noising distribution to x 0 and vary the concentration parameter κ t . When κ t goes to infinity, the von Mises distribution approaches the Dirac delta, centered at x 0 . When κ t goes to 0, it approaches a uniform distribution on a unit circle. The sufficient statistic is T ( x t ) = ( cos x t sin x t ) , and the corresponding natural parameter is η t ( x 0 ) = κ t ( cos x 0 sin x 0 ) . The tail statistic G t ( x t : T ) is therefore defined as follows:

The KL divergence term can be calculated as follows:

t=0

t=11

t=33

t = 44

t=55

t=66

t =77

t=88

t=99

F.6 von Mises-Fisher

F.7 Gamma

Similar to the one-dimensional case, we set the mode of the distribution to x 0 , and define the schedule using the concentration parameter κ t . When κ t goes to infinity, the von Mises-Fisher distribution approaches the Dirac delta, centered at x 0 . When κ t goes to 0, it approaches a uniform distribution on a unit sphere. The sufficient statistic is T ( x ) = x , and the corresponding natural parameter is η t ( x 0 ) = κ t x 0 . The tail statistic G t ( x t : T ) is therefore defined as follows:

The KL divergence term can be calculated as follows:

There are many ways to define a schedule. We choose to interpolate the mean of the distribution from x 0 at t = 0 to 1 at t = T . This can be achieved with the following parameterization:

The mean of the distribution is α t β t , and the variance is α t β 2 t . Therefore, we recover the Dirac delta, centered at x 0 , when we set ξ t to 0 and α t to infinity. To achieve some standard distribution that doesn't depend on x 0 , we can set ξ t to 1 and α t to some fixed value α T .

In this parameterization, the natural parameters are α t -1 and -β t ( x 0 ) , and the corresponding sufficient statistics are log x t and x t . Since the parameter α t doesn't depend on x 0 , we only need the sufficient statistic T ( x t ) = x t to define the tail statistic:

The KL divergence can be computed as follows:

xt

t = 0

t = 7

t = 14

t = 21

t = 28

t = 35

t = 42

t = 49

t = 56

t = 63

F.8 Wishart

The natural parameters for the Wishart distribution are -1 2 V -1 t and n t -p -1 2 . To achieve linear parameterization, we need to linearly parameterize the inverse of V t rather than V t directly. Similar to the Gamma distribution, we interpolate the mean of the distribution from X 0 at t = 0 to I at t = T . This can be achieved with the following parameterization:

We recover the Dirac delta, centered at X 0 , when we set ξ t to 0 and n t to infinity. To achieve some standard distribution that doesn't depend on X 0 , we set ξ t to 1 and n t to some fixed value n T .

The sufficient statistic for this distribution is T ( X t ) = ( X t log | X t | ) . Since the parameter n t doesn't depend on X 0 , we don't need the corresponding sufficient statistic log | X t | and can use T ( X t ) = X t to define the tail statistic:

The KL divergence can then be calculated as follows:

G Choosing the noise schedule

In DDPMs, we train a neural network to predict x 0 given the current noisy sample x t . In SS-DDPMs, we use the tail statistic G t as the neural network input instead. Therefore, it is natural to search for a schedule, where the 'level of noise' in G t is similar to the 'level of noise' in x t from some DDPM model. Similarly to D3PM, we formalize the 'level of noise' as the mutual information between the clean and noisy samples. For SS-DDPM it would be I SS ( x 0 ; G t ) , and for DDPM it would be I DDPM ( x 0 ; x t ) . We would like to start from a well-performing DDPM model and define a similar schedule by matching I SS ( x 0 ; G t ) ≈ I DDPM ( x 0 ; x t ) . Since Gaussian SS-DDPM is equivalent to DDPM, the desired schedule can be found in Theorem 2. In other cases, however, it is difficult to find a matching schedule.

In our experiments, we found the following heuristic to work well enough. First, we find a Gaussian SS-DDPM schedule that is equivalent to the DDPM with the desired schedule. We denote the corresponding mutual information as I SS N ( x 0 ; G t ) = I DDPM ( x 0 ; x t ) . Then, we can match it in the original space of x t and hope that the resulting mutual information in the space of tail statistics is close too:

Assuming the schedule is parameterized by a single parameter ν , we can build a look-up table I SS ( x 0 ; x ν ) for a range of parameters ν . Then we can use binary search to build a schedule to match the mutual information I SS ( x 0 ; x t ) to the mutual information schedule I SS N ( x 0 ; x t ) . While this procedure doesn't allow to match the target schedule exactly , it provides a good enough approximation and allows to obtain an adequate schedule. We used this procedure to find the schedule for the Beta diffusion in our experiments with image generation.

To build the look-up table, we need a robust way to estimate the mutual information. The target mutual information I SS N ( x 0 ; x t ) can be computed analytically when the data q ( x 0 ) follows a Gaussian distribution. When the data follows an arbitrary distribution, it can be approximated with a Gaussian mixture and the mutual information can be calculated using numerical integration. Estimating the mutual information for arbitrary noising distributions is more difficult. We find that the Kraskov estimator (Kraskov et al., 2004) works well when the mutual information is high ( I > 2 ). When the mutual information is lower, we build a different estimator using DSIVI bounds (Molchanov et al., 2019).

This conditional entropy is available in closed form for many distributions in the exponential family. Since the marginal distribution q SS ( x 0 , x t ) = ∫ q SS ( x t | x 0 ) q ( x 0 ) dx 0 is a semi-implicit distribution (Yin & Zhou, 2018), we can use the DSIVI sandwich (Molchanov et al., 2019) to obtain an upper and lower bound on the marginal entropy H [ x t ] :

These bounds are asymptotically exact and can be estimated using Monte Carlo. We use K = 1000 when the mutual information is high ( 1 2 ≤ I < 2 ), K = 100 when the mutual information is lower ( 0 . 002 ≤ I < 1 2 ), and estimate the expectations using M = 10 8 K -1 samples for each timestamp. For values I > 2 we use the Kraskov estimator with M = 10 5 samples and k = 10 neighbors. For values I < 0 . 002 we fit an exponential curve i ( t ) = e at + b to interpolate between the noisy SIVI estimates, obtained with K = 50 and M = 10 5 .

For evaluating the mutual information I ( x 0 ; G t ) between the clean data and the tail statistics, we use the Kraskov estimator with k = 10 and M = 10 5 .

The mutual information look-up table for the Beta star-shaped diffusion, as well as the used estimations of the mutual information, are presented in Figure 11. The resulting schedule for the beta diffusion is presented in Figure 12, and the comparison of the mutual information schedules for the tail statistics between Beta SS-DDPM and the referenced Gaussian SS-DDPM is presented in Figure 13.

For Categorical SS-DDPM we estimate the mutual information using Monte Carlo. We then choose the noising schedule to match the cosine schedule used by Hoogeboom et al. (2021) using a similar technique.

H Normalizing the tail statistics

We illustrate different strategies for normalizing the tail statistics in Figure 14. Normalizing by the sum of coefficients is not enough, therefore we resort to matching the mean and the variance empirically. We refer to this trick as time-dependent tail normalization.

Also, proper normalization allows us to visualize the tail statistics by projecting them back into the original domain using T -1 ( ˜ G t ) . The effect of normalization is illustrated in Figure 15.

































































I

(

x

0;

x

)























































t

I Synthetic data

We compare the performance of DDPM and SS-DDPM on synthetic tasks with exotic data domains. In these tasks, we train and sample from DDPM and SS-DDPM using T = 64 steps. We use an MLP with 3 hidden layers of size 512, swish activations and residual connections through hidden layers (He et al., 2015). We use sinusoidal positional time embeddings (Vaswani et al., 2017) of size 32 and concatenate them with the normalized tail statistics ˜ G t . We add a mapping to the corresponding domain on top of the network. During training, we use gradient clipping and EMA weights to improve stability. All models on synthetic data were trained for 350k iterations with batch size 128. In all our experiments with SS-DDPM on synthetic data we use time-dependent tail normalization. We use DDPM with linear schedule and L simple or L vlb as a loss function. We choose between L simple and L vlb objective based on the KL divergence between the data distribution and the model distribution D KL ( q ( x 0 ) ∥ p θ ( x 0 )) . To make an honest comparison, we precompute normalization statistics for DDPM in the same way that we do in time-dependent tail normalization. In DDPM, a neural network makes predictions using x t as an input, so we precompute normalization statistics for x t and fix them during the training and sampling stages.

Probabilistic simplex We evaluate Dirichlet SS-DDPM on a synthetic problem of generating objects on a threedimensional probabilistic simplex. We use a mixture of three Dirichlet distributions with different parameters as training







































I

(

x

0;

Gt

)





0













t





















|

Gt

|

x

0 =0.1

x

0 =0.5

x

0 =0.9

(a) No normalization

˜

G

t

=

G

t













t











Gt

x

0 =0.1

x

0 =0.5

x

0 =0.9

(b) Normalized by the sum of coefficients

˜

G

t

=

G

t

∑

T

s

=

t

a

s













t















Gt

x

0 =0.1

x

0 =0.5

x

0 =0.9

(c) Matched the mean and variance

˜

G

t

= (

G

t

-

E

G

t

)

std(

T

(

x

0

))

std(

G

t

)

+

E

T

(

x

)

Sum

Stats

0=↓

100

200

300

400

500

600

700

800

900

999

data. The Dirichlet SS-DDPM forward process is illustrated in Figure 8. To map the predictions to the domain, we put the Softmax function on the top of the MLP. We optimize Dirichlet SS-DDPM on the VLB objective without any modifications using Adam with a learning rate of 0 . 0004 . The DDPM was trained on L vlb using Adam with a learning rate of 0 . 0002 . An illustration of the samples is presented in Table 4.

Original

DDPM

Dirichlet SS-DDPM

Symmetric positive definite matrices We evaluate Wishart SS-DDPM on a synthetic problem of generating symmetric positive definite matrices of size 2 × 2 . We use a mixture of three Wishart distributions with different parameters as training data. The Wishart SS-DDPM forward processes are illustrated in Figure 10. For the case of symmetric positive definite matrices V , MLP predicts a lower triangular factor L θ from the Cholesky decomposition V θ = L θ L T θ . For stability of sampling from the Wishart distribution and estimation of the loss function, we add a scalar matrix 10 -4 I to the predicted symmetric positive definite matrix V θ . In Wishart SS-DDPM we use an analogue of L simple as a loss function. To improve the training stability, we divide each KL term corresponding to a timestamp t by n t , the de-facto concentration parameter

from the used noising schedule (see Table 6). Wishart SS-DDPM was trained using Adam with a learning rate of 0 . 0004 . The DDPM was trained on L simple using Adam with a learning rate of 0 . 0004 . Since positive definite 2 × 2 matrices can be interpreted as ellipses, we visualize the samples by drawing the corresponding ellipses. An illustration of the samples is presented in Table 5.

Original

DDPM

Wishart SS-DDPM

J Geodesic data

We evaluate von Mises-Fisher SS-DDPM on the problem of restoring the distribution on the sphere from empirical data of fires on the Earth's surface (EOSDIS, 2020). We illustrate points on the sphere using a 2D projection of the Earth's map. We train and sample from von Mises-Fisher SS-DDPM using T = 100 steps. The forward process is illustrated in Figure 9. We use the same MLP architecture as described in Appendix I and also use time-dependent tail normalization. To map the prediction onto the unit sphere, we normalize the three-dimensional output of the MLP. We optimize L vlb using the AdamW optimizer with a learning rate of 0.0002 and exponential decay with γ = 0 . 999997 . The model is trained for 2 , 000 , 000 iterations with batch size 100 . For inference, we also use EMA weights with a decay of 0 . 9999 .

K Discrete data

We evaluate Categorical SS-DDPM on the task of unconditional character-level text generation on the text8 dataset. We evaluate our model on sequences of length 256. Using the property of Categorical SS-DDPM, that we can directly compute mutual information between x 0 and G t (see Appendix F.4), we match the noising schedule of G t to the noising schedule of x t by Austin et al. (2021). We optimize Categorical SS-DDPM on the VLB objective. Following Austin et al. (2021), we use the default T5 encoder architecture with 12 layers, 12 heads, mlp dim 3072 and qkv dim 768 . We add positional embeddings to the sequence of tail statistics G t and also add the sinusoidal positional time embedding to the beginning. We use Adam (Kingma & Ba, 2015) with learning rate 5 × 10 -4 with a 10 , 000 -step learning rate warmup, but instead of inverse sqrt decay, we use exponential decay with γ = 0 . 999995 . We use a standard 90 , 000 , 000 / 5 , 000 , 000 / 500 , 000 train-test-validation split and train neural network for 512 epochs (time costs when using 3 NVIDIA A100 GPUs: training took approx. 112 hours and estimating NLL on the test set took approx. 2 . 5 hours). Since the Softmax function does not break the sufficiency of tail statistic in Categorical SS-DDPMs (see Appendix F.4), we can use the Softmax function as an efficient input normalization for the neural network. Categorical SS-DDPM also provides us with a convenient way to define log p ( x 0 | G 1 ) , and we make use of it to estimate the NLL score (see Appendix F.4).

L Image data

In our experiments with Beta SS-DDPM, we use the NCSN++ neural network architecture and train strategy by Song et al. (2020b) (time costs when using 4 NVIDIA 1080 GPUs: training took approx. 96 hours, sampling of 50 , 000 images took approx. 10 hours). We put a sigmoid function on the top of NCSN++ to map the predictions to the data domain. We use time-dependent tail normalization and train Beta SS-DDPM with T = 1000 . We matched the noise schedule in Beta SS-DDPM to the cosine noise schedule in Gaussian DDPM (Nichol & Dhariwal, 2021) in terms of mutual information, as described in Appendix G.

M Changing discretization

When sampling from DDPMs, we can skip some timestamps to trade off computations for the quality of the generated samples. For example, we can generate x t 1 from x t 2 in one step, without generating the intermediate variables x t 1 +1: t 2 -1 :

In the general case of SS-DDPM, we can't just skip the variables. If we skip the variables, the corresponding tail statistics will become atypical and the generative process will fail. To keep the tail statistics adequate, we can sample all the intermediate variables x t but do it in a way that doesn't use additional function evaluations:

In case of Gaussian SS-DDPM this trick is equivalent to skipping the variables in DDPM:

In general case, this trick can be formalized as the following approximation to the reverse process:

Peer reviews:
comments: SUMMARY:
This introduces the so-called Star-Shaped DDPM (SS-DDPM). Instead of using Gaussian forward diffusion step in a Markovian manner, the diffusion process is directly conditioned on the data, in order to construct the "true" reverse process (posterior), resulting in the star-shaped diffusion process. This opens the door for using other transition process as demonstrated in the paper.

STRENGTHS:
1. This is a solid paper that I have reviewed so far.  All the conclusions are backed by the theory.
2. The paper is well presented and easy to follow the main the idea with full materials in the supplementary
3. The idea presented in the paper is novel and may have wide applications.

WEAKNESSES:
Frankly speaking, I like to see this paper to be accepted.  I cannot form any weakness.  The only point I may raise is it would be much much better for the authors to share their implementation source code.

QUESTIONS:
When using other transition distributions in the forward process, what is the strategy to actually diffuse the data (into noises)?

LIMITATIONS:
Not much in this perspective. No comments
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-06-25

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a new probabilistic structure for denoising models that does away with a Markov forward process. The authors derive the reverse process in terms of a sufficient statistic that allows for efficient reverse sampling. The form of the model is derived for a variety of noising distributions and the model is evaluated on a range of different data modalities.\n\nSTRENGTHS:\nThis work can potentially help many people in the generative modelling community because there has been a lot of work trying to extend corruption process to different data types e.g. discrete, manifold data and this paper has a general recipe that appears to work for a large chunk of these and more which could be really helpful when encountering non-standard data types. The example of the Wishart distribution on positive definite matrices is very cool.\n\nI can foresee many extensions to this work, looking to bring more distributions into the framework, extending it to continuous time which could bring up a lot of interesting technical problems and broadening the class of possible distributions by introducing approximate sufficient statistics.\n\nThe paper has a good suite of experiments on a variety of real world practically relevant data such as text and images and the performance of the model seems good\n\nWEAKNESSES:\nI think the paper could benefit from more intuitions and visualizations for the role of G. Initially, the model seems strange as x_t can jump a lot between x_t and x_{t+1} due to going to x_0 and resampling q(x_t | x_0). However, it seems that G instead is the smoothly changing statistic that is more analogous to the state in previous diffusion models. For example, I think figure 9 in the supplementary is quite interesting and could be moved to the main as it shows that a correctly normalized G looks very similar to the state in norm diffusion models. It would be good to have more intuition as to the role of G and how it is more closely linked to the 'x_t' of normal diffusion models.\n\nFurther, there is not really a experiment that really motivates the use of these more interesting distributions in a scenario where normal diffusion models really don't work at all. I appreciate that this is somewhat an open question at the moment in the literature more generally and the community has already proven an interest in adapting the corruption process to the data modality of interest so this may not be a big requirement.\n\nEdit after rebuttal: I have read the authors rebuttal and appreciate the answers to my questions about the derivations and the FID scores. I intend to keep my score as it is.\n\nQUESTIONS:\nIn the proof of Theorem 1, how do you go from line (23) to line (24), as in, how do you know what q(G_t | x_0) is, for the numerator in (24)?\n\nFigure 6 is confusing to me because the figure makes it seem like the model is better than DDPM at all NFE but in the text the best FID is stated to be 3.17 which is the same as the best FID from DDPM and so shouldn't Figure 6 show the method and DDPM at least overlapping at 10^3 sampling steps?\n\nIn DDPM, the Gaussian form of the reverse distribution is justified by considering the limit as the noise in each step is taken to 0 and the number of steps is taken to infinity. Is there any such justification for the Gaussian form that you assume in your model or is this a heuristic?\n\nLIMITATIONS:\nThe limitations of the framework e.g. restriction of the distributions and ability to find the sufficient statistics was nicely discussed.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a star-shaped diffusion probabilistic model, which is non-Markovian, and more like an autoregressive way of predicting intermediate states xt. Specifically, the authors define a forward process q(xt|x0) that each intermediate state xt is directly related to the initial state x0. The reverse process is defined by p(xt|xt, xt+1, ..., xT). To efficiently solve p(xt|xt, xt+1, ..., xT), the authors propose a compressive representation Gt(xt, xt+1, ..., xT) to extract all the information of x0 contained in history states {xt+1, ..., xT}, so the modified distribution becomes p(xt|Gt). The training and inference go in a similar way as DDPM. Simple experiments show that the proposed method performs better than the baseline methods MTD and DDPM.\n\nSTRENGTHS:\n1. This paper provides a novel view of the definition of diffusion models and is very inspirational. It is worth thinking about whether the Markov process is the best design choice considering the fact that the whole history {xt, xt+1, ..., xT} contains more information than a single state xt. \n\n2. The author proposed a novel diffusion model with basic evaluations, which may inspire the community for further exploration.\n\nWEAKNESSES:\n1. The presentation needs to be improved: (1) T, Ωt, and ht in Eq.13 lack interpretation; (2) It is not well proved that Gt covers all the information of x0 in {xt, xt+1, ..., xT}. Although the authors mention that there are explanations in the supplementary material, you need to ensure the reader has a basic understanding by just reading the paper.\n\n2. The advantage and application value of the proposed method is not clear. For example, can you provide a practical case that DDPM can not address while SS-DDPM can?\n\n3. The experiments are too simple and not persuasive. To prove the advantage of SS-DDPM, it is encouraged to compare with state-of-the-art methods on diverse benchmarks, \n\n4. As the authors said in line 128, \"In general case the dimensionality of Gt would grow with the size of the tail\"\n\n4. There is no evidence that SS-DDPM can apply acceleration methods like DDIM.\n\nQUESTIONS:\nI doubt the reverse process of q(xt|x0) should be p(x0|xt), can you explain why the reverse process is p(xt|xt, xt+1, ..., xT)?\n\nLIMITATIONS:\nThe scope of this paper is interesting, but it needs more time to polish the theory, presentation, and experiments. Considering its current form, I think this paper does not meet the acceptance criteria of NeurIPS. However, if there is enough evidence to show the value of this paper, I may also change my score.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 1, "DATE": "2023-07-05"}

comments: SUMMARY:
This paper proposes a non-Markovian diffusion model named the star-shaped diffusion model that generates a sequence of noised image from the original image in the forward process.

This paper studies the theoretical foundation of such new type of model showing that if the forward process is based on a subset of the exponential family of distribution, a tail statistics is sufficient statistics for the reversed process. Based on such a foundation, the reversed process can be conducted.

STRENGTHS:
The overall structure of this new diffusion model is new, novel and interesting.

The application of PKD theory is sound.

WEAKNESSES:
Although the model is pretty novel and new, there is a lack of motivation from practical problems -- the experiment are weak and kind of toy. A showcase of the unique application of this model is appreciated.

Lack of comparison of some other non-markovian diffusion on constrained domains such as [1,2]

[1] Learning Diffusion Bridges on Constrained Domains
[2] First Hitting Diffusion Models for Generating Manifold, Graph and Categorical Data

QUESTIONS:
Is there any ODE-version of this SDE-based model?

LIMITATIONS:
See above
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

comments: SUMMARY:
The paper presents the Star-Shaped DDPM (SS-DDPM), a general recipe for designing a diffusion model with a noising process lying in a general subset of the exponential family. With a Gaussian noising process, SS-DDPM recovers the DDPM. Diverse experiments on synthetic and practical image and text datasets demonstrate the effectiveness of the proposed SS-DDPM.

STRENGTHS:
(1) A general recipe for designing a diffusion model, termed Star-Shaped DDPM (SS-DDPM), is proposed.

(2) The recipe is derived and analyzed in detail, with convincing statistical justifications.

(3) Diverse experiments are conducted to demonstrate the effectiveness of the SS-DDPM.

WEAKNESSES:
(1) Many technical details are given in the supplementary material.

(2) The advantages and disadvantages of the SS-DDPM over the DDPM are not explicitly discussed.

QUESTIONS:
What are the limitations of the proposed SS-DDPM, when compared with other diffusion models?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-09

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposed a star-shaped denoising diffusion probabilistic model (DDPM), which extends DDPM to non-Gaussian noises. As a result, the backward/generative process requires conditioning on tails. The authors then propose an efficient tail conditioning strategy which works when the forward process follows an exponential family with linear parameterization. Duality between star-shaped and Markov diffusion processes are also established which provides theoretical support for star-shaped DDPM for its ability to go beyond Gaussian noises. The effectiveness of star-shaped DDPM is demonstrated on several experiments.\n\nSTRENGTHS:\n1. This paper is written very clearly and well organized. \n2. It also provides a general framework for DDPM that can incoporate no-Gaussian noises, where previous efforts are designed for specific noises and not generalizable to other distributions.\n3. Duality between star-shaped and Markovian diffusion is established, which shows SS-DDPM and vanilla DDPM are equivalent in the Gaussian case.\n\nWEAKNESSES:\n1. While experiments on synthetic data demonstrate the ability of SS-DDPM to deal with non-Gaussian noises, the effectiveness of SS-DDPM on real data needs further verification. The current experiment on real data seems a bit inadequate.\n\n2. Star-shaped denoising diffusion models have been proposed before (Rissanen et al. (2022)), and there is no comparison to this previous method.\n\nQUESTIONS:\n1. How does DDPM perform on the Geodesic data in the experiments?\n2. How much improvement can SS-DDPM have over DDPM on real data? The advantage over DDPM seems a bit tiny based on current evaluation. Also, in which case would non-Gaussian noise be useful for DDPMs?\n3. Rissanen et al. (2022) also proposed a similar star-shaped DDPM, where the reverse process is Markovian based on the heat equation. How does your method compare to their's?\n\nLIMITATIONS:\nYes, they do.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-17"}