SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Andrey Okhotin*
HSE University, MSU University
andrey.okhotin@gmail.com
Dmitry Molchanov*
BAYESG
Budva, Montenegro
dmolch111@gmail.com
Vladimir Arkhipkin
Sber AI
arkhipkin.v98@gmail.com
Grigory Bartosh
AMLab, Informatics Institute
g.bartosh@uva.nl
Viktor Ohanesian
v.v.oganesyan@gmail.com
Aibek Alanov
AIRI, HSE University
alanov.aibek@gmail.com
Constructor University
dvetrov@constructor.university
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.
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).
Figure 1: Markovian forward processes of DDPM (left) and the star-shaped forward process of SS-DDPM (right).
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 noisig 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-DDDM), a new approach that generalizes Gaussian DDPM to an exponential family of noise distributions. In SS-DDDM, 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.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^{\mathrm{DDPM}}(x_{0:T})$  and a corresponding reverse (denoising) process  $p_{\theta}^{\mathrm{DDPM}}(x_{0:T})$ . The forward process is defined as a Markov chain with Gaussian conditionals:
$$
q ^ {\mathrm {D D P M}} \left(x _ {0: T}\right) = q \left(x _ {0}\right) \prod_ {t = 1} ^ {T} q ^ {\mathrm {D D P M}} \left(x _ {t} \mid x _ {t - 1}\right), \tag {1}
$$
$$
q ^ {\mathrm {D D P M}} \left(x _ {t} \mid x _ {t - 1}\right) = \mathcal {N} \left(x _ {t}; \sqrt {1 - \beta_ {t}} x _ {t - 1}, \beta_ {t} \mathbf {I}\right), \tag {2}
$$
where  $q(x_0)$  is the data distribution. Parameters  $\beta_{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^{\mathrm{DDPM}}(x_T) = \mathcal{N}(x_T;0,\mathbf{I})$ .
(a) Denoising Diffusion Probabilistic Models
(b) Star-Shaped Denoising Diffusion Probabilistic Models
Figure 2: Model structure of DDPM and SS-DDPM.
The reverse process  $p_{\theta}^{\mathrm{DDPM}}(x_{0:T})$  follows a similar structure and constitutes the generative part of the model:
$$
p _ {\theta} ^ {\mathrm {D D P M}} \left(x _ {0: T}\right) = q ^ {\mathrm {D D P M}} \left(x _ {T}\right) \prod_ {t = 1} ^ {T} p _ {\theta} ^ {\mathrm {D D P M}} \left(x _ {t - 1} \mid x _ {t}\right), \tag {3}
$$
$$
p _ {\theta} ^ {\mathrm {D D P M}} \left(x _ {t - 1} \mid x _ {t}\right) = \mathcal {N} \left(x _ {t - 1}; \mu_ {\theta} \left(x _ {t}, t\right), \Sigma_ {\theta} \left(x _ {t}, t\right)\right). \tag {4}
$$
The forward process  $q^{\mathrm{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_{\theta}^{\mathrm{DDPM}}(x_{0:T})$ . These parameters are tuned to maximize the variational lower bound (VLB) on the likelihood of the training data:
$$
\mathcal {L} ^ {\mathrm {D D P M}} (\theta) = \mathbb {E} _ {q ^ {\mathrm {D D P M}}} \left[ \log p _ {\theta} ^ {\mathrm {D D P M}} \left(x _ {0} \mid x _ {1}\right) - \sum_ {t = 2} ^ {T} D _ {K L} \left(q ^ {\mathrm {D D P M}} \left(x _ {t - 1} \mid x _ {t}, x _ {0}\right) \| p _ {\theta} ^ {\mathrm {D D P M}} \left(x _ {t - 1} \mid x _ {t}\right)\right) \right] \tag {5}
$$
$$
\mathcal {L} _ {\theta} ^ {\text {D D P M}} (\theta) \rightarrow \max  _ {\theta} \tag {6}
$$
One of the main challenges in defining DDPMs is the computation of the posterior  $q^{\mathrm{DDPM}}(x_{t - 1}|x_t,x_0)$ . Specifically, the transition probabilities  $q^{\mathrm{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^{\mathrm{ss}}(x_{0:T})$  that has the following structure:
$$
q ^ {\mathrm {s s}} \left(x _ {0: T}\right) = q \left(x _ {0}\right) \prod_ {t = 1} ^ {T} q ^ {\mathrm {s s}} \left(x _ {t} \mid x _ {0}\right), \tag {7}
$$
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^{\mathrm{DDPM}}(x_{0:T})$  has a Markovian structure (Ho et al., 2020), allowing for an efficient sequential generation algorithm:
$$
q ^ {\mathrm {D D P M}} \left(x _ {0: T}\right) = q ^ {\mathrm {D D P M}} \left(x _ {T}\right) \prod_ {t = 1} ^ {T} q ^ {\mathrm {D D P M}} \left(x _ {t - 1} \mid x _ {t}\right). \tag {8}
$$
For the star-shaped diffusion, however, the Markovian assumption breaks:
$$
q ^ {\mathrm {s s}} \left(x _ {0: T}\right) = q ^ {\mathrm {s s}} \left(x _ {T}\right) \prod_ {t = 1} ^ {T} q ^ {\mathrm {s s}} \left(x _ {t - 1} \mid x _ {t: T}\right). \tag {9}
$$
Consequently, we now need to approximate the true reverse process by a parametric model which is conditioned on the whole tail  $x_{t:T}$ .
$$
p _ {\theta} ^ {\mathrm {s s}} \left(x _ {0: T}\right) = p _ {\theta} ^ {\mathrm {s s}} \left(x _ {T}\right) \prod_ {t = 1} ^ {T} p _ {\theta} ^ {\mathrm {s s}} \left(x _ {t - 1} \mid x _ {t: T}\right). \tag {10}
$$
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$ .
Figure 3: Markov reverse process fails to recover realistic images from star-shaped diffusion, while a general reverse process produces realistic images. The top row is equivalent to DDIM at  $\sigma_t^2 = 1 - \alpha_{t-1}$ . A similar effect was also observed by Bansal et al. (2022).
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:
$$
\mathcal {L} ^ {\mathrm {s s}} (\theta) = \mathbb {E} _ {q ^ {\mathrm {s s}}} \left[ \log p _ {\theta} \left(x _ {0} \mid x _ {1: T}\right) - \sum_ {t = 2} ^ {T} D _ {K L} \left(q ^ {\mathrm {s s}} \left(x _ {t - 1} \mid x _ {0}\right) \| p _ {\theta} ^ {\mathrm {s s}} \left(x _ {t - 1} \mid x _ {t: T}\right) \right] \right. \tag {11}
$$
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 noisy distributions. Since conditioning the predictive model  $p_\theta(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} = \mathcal{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:
$$
q ^ {\mathrm {s s}} \left(x _ {t - 1} \mid x _ {t: T}\right) = q ^ {\mathrm {s s}} \left(x _ {t - 1} \mid G _ {t}\right). \tag {12}
$$
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:
$$
q ^ {\mathrm {s s}} \left(x _ {t} \mid x _ {0}\right) = h _ {t} \left(x _ {t}\right) \exp \left\{\eta_ {t} \left(x _ {0}\right) ^ {\mathsf {T}} \mathcal {T} \left(x _ {t}\right) - \Omega_ {t} \left(x _ {0}\right) \right\}, \tag {13}
$$
$$
\eta_ {t} \left(x _ {0}\right) = A _ {t} f \left(x _ {0}\right) + b _ {t}. \tag {14}
$$
Let  $G_{t}$  be a tail statistic, defined as follows:
$$
G _ {t} = \mathcal {G} _ {t} \left(x _ {t: T}\right) = \sum_ {s = t} ^ {T} A _ {s} ^ {\top} \mathcal {T} \left(x _ {s}\right). \tag {15}
$$
Then,  $G_{t}$  is a sufficient tail statistic:
$$
q ^ {\mathrm {s s}} \left(x _ {t - 1} \mid x _ {t: T}\right) = q ^ {\mathrm {s s}} \left(x _ {t - 1} \mid G _ {t}\right). \tag {16}
$$
Here definition (13) is the standard definition of the exponential family, where  $h_t(x_t)$  is the base measure,  $\eta_t(x_0)$  is the vector of natural parameters with corresponding sufficient statistics  $\mathcal{T}(x_t)$ , and  $\Omega_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}) = \mathrm{Beta}(x_{t};\alpha_{t},\beta_{t})$  using  $x_0$  as the mode of the distribution and introducing a new concentration parameter  $\nu_{t}$ :
$$
\alpha_ {t} = 1 + \nu_ {t} x _ {0}, \tag {17}
$$
$$
\beta_ {t} = 1 + \nu_ {t} \left(1 - x _ {0}\right). \tag {18}
$$
In this case,  $\eta_t(x_0) = \nu_t x_0$ ,  $\mathcal{T}(x_t) = \log \frac{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:
$$
p _ {\theta} ^ {\mathrm {s s}} \left(x _ {t - 1} \mid x _ {t: T}\right) \approx q ^ {\mathrm {s s}} \left(x _ {t - 1} \mid x _ {t: T}\right) = \int q ^ {\mathrm {s s}} \left(x _ {t - 1} \mid x _ {0}\right) q ^ {\mathrm {s s}} \left(x _ {0} \mid x _ {t: T}\right) d x _ {0}. \tag {19}
$$
Similarly to DDPM (Ho et al., 2020), we choose to approximate  $q^{\mathrm{SS}}(x_0|x_{t:T})$  with a delta function centered at the prediction of some model  $x_{\theta}(\mathcal{G}_t(x_{t:T}),t)$ . This results in the following definition of the reverse process of SS-DDPM:
$$
p _ {\theta} ^ {\mathrm {S S}} \left(x _ {t - 1} \mid x _ {t: T}\right) = q ^ {\mathrm {S S}} \left(x _ {t - 1} \mid x _ {0}\right) | _ {x _ {0} = x _ {\theta} \left(\mathcal {G} _ {t} \left(x _ {t: T}\right), t\right)}. \tag {20}
$$
The distribution  $p_{\theta}^{\mathrm{ss}}(x_0|x_{1:T})$  can be fixed to some small-variance distribution  $p_{\theta}^{\mathrm{ss}}(x_0|\hat{x}_0)$  centered at the final prediction  $\hat{x}_0 = x_\theta (\mathcal{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^{\mathrm{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
$$
x _ {0} \sim q (x _ {0})
$$
$$
t \sim \operatorname {U n i f o r m} (1, \dots , T)
$$
$$
x _ {t: T} \sim q ^ {\mathrm {s s}} (x _ {t: T} | x _ {0})
$$
$$
G _ {t} = \sum_ {s = t} ^ {T} A _ {s} ^ {\top} \mathcal {T} \left(x _ {s}\right)
$$
Move along  $\nabla_{\theta}\mathrm{KL}(q^{\mathrm{ss}}(x_{t - 1}|x_0)\| p_\theta^{\mathrm{ss}}(x_{t - 1}|G_t))$
until Convergence
Algorithm 2 SS-DDPM sampling
$$
x _ {T} \sim q ^ {\mathrm {s s}} (x _ {T})
$$
$$
G _ {T} = A _ {T} ^ {\top} \mathcal {T} \left(x _ {T}\right)
$$
$$
\mathbf {f o r} t = \bar {T} \text {t o 2 d o}
$$
$$
\tilde {x} _ {0} = x _ {\theta} \left(G _ {t}, t\right)
$$
$$
x _ {t - 1} \sim q ^ {\mathrm {s s}} \left(x _ {t - 1} \mid x _ {0}\right) | _ {x _ {0} = \tilde {x} _ {0}}
$$
$$
G _ {t - 1} = G _ {t} + A _ {t - 1} ^ {\mathsf {T}} \mathcal {T} (x _ {t - 1})
$$
$$
x _ {0} \sim p _ {\theta} ^ {\mathrm {s s}} (x _ {0} | G _ {1})
$$
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:
$$
G _ {t} = \sum_ {s = t} ^ {T} A _ {s} ^ {\top} \mathcal {T} \left(x _ {s}\right) = G _ {t + 1} + A _ {t} ^ {\top} \mathcal {T} \left(x _ {t}\right), \tag {21}
$$
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  $\overline{\alpha}_t^{\mathrm{DDPM}}$  define the noising schedule for a DDPM model (1-2) via  $\beta_{t} = (\overline{\alpha}_{t - 1}^{\mathrm{DDPM}} - \overline{\alpha}_{t}^{\mathrm{DDPM}}) / \overline{\alpha}_{t - 1}^{\mathrm{DDPM}}$ . Let  $q^{\mathrm{SS}}(x_{0:T})$  be a Gaussian SS-DDPM forward process with the following noising schedule and sufficient tail statistic:
$$
q ^ {\mathrm {S S}} \left(x _ {t} \mid x _ {0}\right) = \mathcal {N} \left(x _ {t}; \sqrt {\bar {\alpha} _ {t} ^ {\mathrm {S S}}} x _ {0}, 1 - \bar {\alpha} _ {t} ^ {\mathrm {S S}}\right), \tag {22}
$$
$$
\mathcal {G} _ {t} \left(x _ {t: T}\right) = \frac {1 - \bar {\alpha} _ {t} ^ {\mathrm {D D P M}}}{\sqrt {\bar {\alpha} _ {t} ^ {\mathrm {D D P M}}}} \sum_ {s = t} ^ {T} \frac {\sqrt {\bar {\alpha} _ {s} ^ {\mathrm {S S}}} x _ {s}}{1 - \bar {\alpha} _ {s} ^ {\mathrm {S S}}}, \text {w h e r e} \tag {23}
$$
$$
\frac {\bar {\alpha} _ {t} ^ {\mathrm {S S}}}{1 - \bar {\alpha} _ {t} ^ {\mathrm {S S}}} = \frac {\bar {\alpha} _ {t} ^ {\mathrm {D D P M}}}{1 - \bar {\alpha} _ {t} ^ {\mathrm {D D P M}}} - \frac {\bar {\alpha} _ {t + 1} ^ {\mathrm {D D P M}}}{1 - \bar {\alpha} _ {t + 1} ^ {\mathrm {D D P M}}}. \tag {24}
$$
Then the tail statistic  $G_{t}$  follows a Gaussian DDPM noising process  $q^{\mathrm{DDPM}}(x_{0:T})|_{x_{1:T} = G_{1:T}}$  defined by the schedule  $\overline{\alpha}_{t}^{\mathrm{DDPM}}$ . 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 noisng schedule for a SS-DDPM model. It significantly depends on the number of diffusion steps  $T$  and behaves differently given different noisng schedules, typical to DDPMs. This is illustrated in Figure 4, where we show the noisng schedules for Gaussian SS-DDPMs that are equivalent to DDPMs with the same cosine schedule.
Figure 4: The noising schedule  $\overline{\alpha}_t^{\mathrm{SS}}$  for Gaussian star-shaped diffusion, defined for different numbers of steps  $T$  using eq. (24). All the corresponding equivalent DDPMs have the same cosine schedule  $\overline{\alpha}_t^{\mathrm{DDPM}}$ .
Figure 5: Top: samples  $x_{t}$  from a Gaussian DDPM forward process with a cosine noise schedule. Bottom: samples  $G_{t}$  from a Beta SS-DDPM forward process with a noise schedule obtained by matching the mutual information. Middle: corresponding samples  $x_{t}$  from that Beta SS-DDPM forward process. The tail statistics have the same level of noise as  $x_{t}^{\mathrm{DDPM}}$ , while the samples  $x_{t}^{\mathrm{BetaSS}}$  are diffused much faster.
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_t^{\mathrm{DDPM}})$ . 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_t^{\mathrm{DDPM}})$  for a DDPM model for all timesteps  $t$ . We estimate the mutual information using Kraskov (Kriskov 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.