SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\left. q \left(G _ {t - 1} \mid G _ {t}, x _ {0}\right) = q \left(x _ {t - 1} \mid x _ {0}\right) \right| _ {x _ {t - 1} = \mathcal {T} ^ {- 1} \left(A _ {t - 1} ^ {- \top} \left(G _ {t - 1} - G _ {t}\right)\right)} \cdot \left| \det  \left[ \frac {d}{d G _ {t - 1}} \mathcal {T} ^ {- 1} \left(A _ {t - 1} ^ {- \top} \left(G _ {t - 1} - G _ {t}\right)\right) \right] \right| \tag {94}
$$
This also allows us to define the reverse model like it was defined in DDPM:
$$
p _ {\theta} \left(G _ {t - 1} \mid G _ {t}\right) = q \left(G _ {t - 1} \mid G _ {t}, x _ {0}\right) \big | _ {x _ {0} = x _ {\theta} \left(G _ {t}, t\right)} \tag {95}
$$
This definition is consistent with the reverse process of SS-DDPM  $p_{\theta}(x_{t - 1}|x_{t:T}) = q(x_{t - 1}|x_0)|_{x_0 = x_\theta (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:
$$
\begin{array}{l} \mathcal {L} ^ {\text {D u a l}} (\theta) = \mathbb {E} _ {q \left(x _ {0}, G _ {1: T}\right)} \left[ \log p _ {\theta} \left(x _ {0} \mid G _ {1}\right) - \sum_ {t = 2} ^ {T} D _ {K L} \left(q \left(G _ {t - 1} \mid G _ {t}, x _ {0}\right) \| p _ {\theta} \left(G _ {t - 1} \mid G _ {t}\right)\right) - D _ {K L} \left(q \left(G _ {T} \mid x _ {0}\right) \| p _ {\theta} \left(G _ {T}\right)\right) \right] = (96) \\ = \mathbb {E} _ {q \left(x _ {0}, G _ {1: T}\right)} \left[ \log p _ {\theta} \left(x _ {0} \mid G _ {1}\right) - \sum_ {t = 2} ^ {T} \mathbb {E} _ {q \left(G _ {t - 1} \mid G _ {t}, x _ {0}\right)} \log \frac {q \left(G _ {t - 1} \mid G _ {t} , x _ {0}\right)}{p _ {\theta} \left(G _ {t - 1} \mid G _ {t}\right)} - D _ {K L} \left(q \left(x _ {T} \mid x _ {0}\right) \| p _ {\theta} \left(x _ {T}\right)\right) \right] = (97) \\ = \mathbb {E} _ {q \left(x _ {0: T}\right)} \left[ \log p _ {\theta} \left(x _ {0} \mid G _ {1}\right) - \sum_ {t = 2} ^ {T} \mathbb {E} _ {q \left(x _ {t - 1} \mid x _ {0}\right)} \log \frac {q \left(x _ {t - 1} \mid x _ {0}\right)}{p _ {\theta} \left(x _ {t - 1} \mid G _ {t}\right)} - D _ {K L} \left(q \left(x _ {T} \mid x _ {0}\right) \| p _ {\theta} \left(x _ {T}\right)\right) \right] = (98) \\ = \mathbb {E} _ {q \left(x _ {0: T}\right)} \left[ \log p _ {\theta} \left(x _ {0} \mid G _ {1}\right) - \sum_ {t = 2} ^ {T} D _ {K L} \left(q \left(x _ {t - 1} \mid x _ {0}\right) \| p _ {\theta} \left(x _ {t - 1} \mid G _ {t}\right)\right) - D _ {K L} \left(q \left(x _ {T} \mid x _ {0}\right) \| p _ {\theta} \left(x _ {T}\right)\right) \right] = \mathcal {L} ^ {\mathrm {s s}} (\theta) (99) \\ \end{array}
$$
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$ :
$$
q ^ {\mathrm {S S}} \left(x _ {t} \mid x _ {0}\right) \xrightarrow [ t \rightarrow 0 ]{} \delta \left(x _ {t} - x _ {0}\right) \tag {100}
$$
When  $t$  goes to  $T$ , we wish to obtain some standard distribution that doesn't depend on  $x_0$ :
$$
q ^ {\mathrm {S S}} \left(x _ {t} \mid x _ {0}\right) \xrightarrow [ t \rightarrow T ]{} q _ {T} ^ {\mathrm {S S}} \left(x _ {T}\right) \tag {101}
$$
In exponential families with linear parameterization of the natural parameter  $\eta_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  $\mathcal{G}_t(x_{t:T})$  using the sufficient statistic  $\mathcal{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
$$
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}, \left(1 - \bar {\alpha} _ {t} ^ {\mathrm {S S}}\right) I\right) \tag {102}
$$
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  $\overline{\alpha}_t^{\mathrm{DDPM}}$ , the following parameterization will produce the same process in the space of tail statistics:
$$
\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 {103}
$$
$$
\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}}} \tag {104}
$$
The KL divergence is computed as follows:
$$
D _ {K L} \left(q ^ {\mathrm {S S}} \left(x _ {t} \mid x _ {0}\right) \| p _ {\theta} ^ {\mathrm {S S}} \left(x _ {t} \mid G _ {t}\right)\right) = \frac {\bar {\alpha} _ {t} ^ {\mathrm {S S}} \left(x _ {0} - x _ {\theta}\right) ^ {2}}{2 \left(1 - \bar {\alpha} _ {t} ^ {\mathrm {S S}}\right)} \tag {105}
$$
F.2 Beta
$$
q \left(x _ {t} \mid x _ {0}\right) = \operatorname {B e t a} \left(x _ {t}; \alpha_ {t}, \beta_ {t}\right) \tag {106}
$$
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  $\nu_t$ , parameterizing  $\alpha_t = 1 + \nu_t x_0$  and  $\beta_t = 1 + \nu_t (1 - x_0)$ . By setting  $\nu_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  $\mathcal{T}(x) = \begin{pmatrix} \log x \\ \log (1 - x) \end{pmatrix}$ . However, under this parameterization, we can derive a one-dimensional tail statistic:
$$
q \left(x _ {t} \mid x _ {0}\right) \propto \exp \left\{\left(\alpha_ {t} - 1\right) \log x _ {t} + \left(\beta_ {t} - 1\right) \log \left(1 - x _ {t}\right) \right\} = \exp \left\{\nu_ {t} x _ {0} \log x _ {t} + \nu_ {t} \left(1 - x _ {0}\right) \log \left(1 - x _ {t}\right) \right\} = \tag {107}
$$
$$
= \exp \left\{\nu_ {t} x _ {0} \log \frac {x _ {t}}{1 - x _ {t}} + \nu_ {t} \log \left(1 - x _ {t}\right) \right\} = \exp \left\{\eta_ {t} \left(x _ {0}\right) \mathcal {T} \left(x _ {t}\right) + \log h _ {t} \left(x _ {t}\right) \right\} \tag {108}
$$
Therefore we can use  $\eta_t(x_0) = \nu_t x_0$  and  $\mathcal{T}(x) = \log \frac{x}{1 - x}$  to define the tail statistic:
$$
\mathcal {G} _ {t} \left(x _ {t: T}\right) = \sum_ {s = t} ^ {T} \nu_ {s} \log \frac {x _ {s}}{1 - x _ {s}} \tag {109}
$$
The KL divergence can then be calculated as follows:
$$
D _ {K L} \left(q ^ {\mathrm {s s}} \left(x _ {t} \mid x _ {0}\right) \| p _ {\theta} ^ {\mathrm {s s}} \left(x _ {t} \mid G _ {t}\right)\right) = \log \frac {\operatorname {B e t a} \left(\alpha_ {t} \left(x _ {\theta}\right) , \beta_ {t} \left(x _ {\theta}\right)\right)}{\operatorname {B e t a} \left(\alpha_ {t} \left(x _ {0}\right) , \beta_ {t} \left(x _ {0}\right)\right)} + \nu_ {t} \left(x _ {0} - x _ {\theta}\right) \left(\psi \left(\alpha_ {t} \left(x _ {0}\right)\right) - \psi \left(\beta_ {t} \left(x _ {0}\right)\right)\right) \tag {110}
$$
t = 0
t=7
t = 14
$\mathbf{t} = 21$
$t = 28$
$t = 35$
$t = 42$
t = 49
Figure 8: Visualization of the forward process in Dirichlet SS-DDPM on a three-dimensional probabilistic simplex.
t = 56
t = 63
F.3 Dirichlet
$$
q \left(x _ {t} \mid x _ {0}\right) = \operatorname {D i r i c h l e t} \left(x _ {t}; \alpha_ {t} ^ {1}, \dots , \alpha_ {t} ^ {K}\right) \tag {111}
$$
Similarly to the Beta distribution, we choose to fix the mode of the distribution at  $x_0$  and introduce a concentration parameter  $\nu_t$ , parameterizing  $\alpha_t^k = 1 + \nu_t x_0^k$ . This corresponds to using the natural parameter  $\eta_t(x_0) = \nu_t x_0$ . By setting  $\nu_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  $\mathcal{T}(x) = (\log x^1, \dots, \log x^K)^\top$  to define the tail statistic:
$$
\mathcal {G} _ {t} \left(x _ {t: T}\right) = \sum_ {s = t} ^ {T} \nu_ {s} \left(\log x _ {s} ^ {1}, \dots , \log x _ {s} ^ {K}\right) ^ {\top} \tag {112}
$$
$$
D _ {K L} \left(q ^ {\mathrm {s s}} \left(x _ {t} \mid x _ {0}\right) \| p _ {\theta} ^ {\mathrm {s s}} \left(x _ {t} \mid G _ {t}\right)\right) = \sum_ {k = 1} ^ {K} \left[ \log \frac {\Gamma \left(\alpha_ {t} ^ {k} \left(x _ {\theta}\right)\right)}{\Gamma \left(\alpha_ {t} ^ {k} \left(x _ {0}\right)\right)} + \nu_ {t} \left(x _ {0} ^ {k} - x _ {\theta} ^ {k}\right) \psi \left(\alpha_ {t} ^ {k} \left(x _ {0}\right)\right) \right] \tag {113}
$$
F.4 Categorical
$$
q \left(x _ {t} \mid x _ {0}\right) = \operatorname {C a t} \left(x _ {t}; p _ {t}\right) \tag {114}
$$
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 \overline{Q}_t$ . By setting  $\overline{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):
$$
\eta_ {t} \left(x _ {0}\right) = \log \left(x _ {0} \bar {Q} _ {t}\right) = x _ {0} \log \bar {Q} _ {t}, \text {w h e r e} \tag {115}
$$
log is taken element-wise, and we assume  $0\log 0 = 0$
The sufficient statistic here is a vector  $x_{t}^{\top}$ . Therefore, the tail statistic can be defined as follows:
$$
\mathcal {G} _ {t} \left(x _ {t: T}\right) = \sum_ {s = t} ^ {T} \left(\log \bar {Q} _ {s} \cdot x _ {s} ^ {\top}\right) = \sum_ {s = t} ^ {T} \log \left(\bar {Q} _ {s} x _ {s} ^ {\top}\right) \tag {116}
$$
$$
D _ {K L} \left(q ^ {\mathrm {s s}} \left(x _ {t} \mid x _ {0}\right) \| p _ {\theta} ^ {\mathrm {s s}} \left(x _ {t} \mid G _ {t}\right)\right) = \sum_ {i = 1} ^ {D} \left(x _ {0} \bar {Q} _ {t}\right) _ {i} \log \frac {\left(x _ {0} \bar {Q} _ {t}\right) _ {i}}{\left(x _ {\theta} \bar {Q} _ {t}\right) _ {i}} \tag {117}
$$
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_{\theta}(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  $(\cdot)$  function to normalize the tail statistics without breaking sufficiency:
$$
\tilde {G} _ {t} = \operatorname {S o f t M a x} \left(\mathcal {G} _ {t} \left(x _ {t: T}\right)\right) \tag {118}
$$
To see this, we can retrace the steps of the proof of Theorem 1:
$$
q ^ {\mathrm {s s}} \left(x _ {0} \mid x _ {t: T}\right) \propto q \left(x _ {0}\right) \exp \left\{x _ {0} G _ {t} \right\} = q \left(x _ {0}\right) \exp \left\{x _ {0} \log \tilde {G} _ {t} + \log \sum_ {i = 1} ^ {d} \exp \left(G _ {t}\right) _ {i} \right\} \propto q \left(x _ {0}\right) \exp \left\{x _ {0} \log \tilde {G} _ {t} \right\} \tag {119}
$$
$$
q ^ {\mathrm {s s}} \left(x _ {0} \mid x _ {t: T}\right) = \frac {q \left(x _ {0}\right) \exp \left\{x _ {0} \log \tilde {G} _ {t} \right\}}{\sum_ {\tilde {x} _ {0}} q \left(\tilde {x} _ {0}\right) \exp \left\{\tilde {x} _ {0} \log \tilde {G} _ {t} \right\}} = q ^ {\mathrm {s s}} \left(x _ {0} \mid \tilde {G} _ {t}\right) \tag {120}
$$
Also, Categorical distributions admit a convenient way to compute fractions involving  $q(G_t|x_0)$ :
$$
\begin{array}{l} q \left(G _ {t} \mid x _ {0}\right) = \sum_ {x _ {t: T}} \left(\prod_ {s \geq t} q \left(x _ {s} \mid x _ {0}\right)\right) \mathbb {1} \left[ G _ {t} = \mathcal {G} _ {t} \left(x _ {t: T}\right) \right] = (121) \\ = \sum_ {x _ {t: T}} \exp \left\{\sum_ {s \geq t} x _ {0} \log \left(\bar {Q} _ {s}\right) x _ {s} ^ {\mathsf {T}} \right\} \mathbb {1} \left[ G _ {t} = \mathcal {G} _ {t} \left(x _ {t: T}\right) \right] = \sum_ {x _ {t: T}} \exp \left\{x _ {0} G _ {t} \right\} \mathbb {1} \left[ G _ {t} = \mathcal {G} _ {t} \left(x _ {t: T}\right) \right] = (122) \\ = \exp \left\{x _ {0} G _ {t} \right\} \sum_ {x _ {t: T}} \mathbb {1} \left[ G _ {t} = \mathcal {G} _ {t} \left(x _ {t: T}\right) \right] = \exp \left\{x _ {0} G _ {t} \right\} \cdot \# _ {G _ {t}} (123) \\ \end{array}
$$
This allows us to define the likelihood term  $p(x_0 | x_{1:T})$  as follows:
$$
p \left(x _ {0} \mid G _ {1}\right) = \frac {q \left(G _ {1} \mid x _ {0}\right) \tilde {p} \left(x _ {0}\right)}{\sum_ {\tilde {x} _ {0}} q \left(G _ {1} \mid \tilde {x} _ {0}\right) \tilde {p} \left(\tilde {x} _ {0}\right)} = \frac {\exp \left\{x _ {0} G _ {1} \right\} \tilde {p} \left(x _ {0}\right)}{\sum_ {\tilde {x} _ {0}} \exp \left\{\tilde {x} _ {0} G _ {1} \right\} \tilde {p} \left(\tilde {x} _ {0}\right)}, \tag {124}
$$
where  $\tilde{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$ :
$$
\begin{array}{l} I \left(x _ {0}; G _ {t}\right) = \mathbb {E} _ {q \left(x _ {0}\right) q \left(G _ {t} \mid x _ {0}\right)} \log \frac {q \left(G _ {t} \mid x _ {0}\right)}{q \left(G _ {t}\right)} = \mathbb {E} _ {q \left(x _ {0}\right) q \left(G _ {t} \mid x _ {0}\right)} \log \frac {q \left(G _ {t} \mid x _ {0}\right)}{\sum_ {\tilde {x} _ {0}} q \left(G _ {t} \mid \tilde {x} _ {0}\right) q \left(\tilde {x} _ {0}\right)} = (125) \\ = \mathbb {E} _ {q \left(x _ {0}\right) q \left(G _ {t} \mid x _ {0}\right)} \log \frac {\exp \left\{x _ {0} G _ {t} \right\} \cdot \# _ {G _ {t}}}{\sum_ {\tilde {x} _ {0}} \exp \left\{\tilde {x} _ {0} G _ {t} \right\} \cdot \# _ {G _ {t}} \cdot q \left(\tilde {x} _ {0}\right)} = \mathbb {E} _ {q \left(x _ {0}\right) q \left(G _ {t} \mid x _ {0}\right)} \log \frac {\exp \left\{x _ {0} G _ {t} \right\}}{\sum_ {\tilde {x} _ {0}} \exp \left\{\tilde {x} _ {0} G _ {t} \right\} q \left(\tilde {x} _ {0}\right)} (126) \\ \end{array}
$$
It can be estimated using Monte Carlo. We use it when defining the noising schedule for Categorical SS-DDPM.
F.5 von Mises
$$
q \left(x _ {t} \mid x _ {0}\right) = \operatorname {v o n M i s e s} \left(x _ {t}; x _ {0}, \kappa_ {t}\right) \tag {127}
$$
The von Mises distribution has two parameters, the mode  $x$  and the concentration  $\kappa$ . It is natural to set the mode of the noising distribution to  $x_0$  and vary the concentration parameter  $\kappa_t$ . When  $\kappa_t$  goes to infinity, the von Mises distribution approaches the Dirac delta, centered at  $x_0$ . When  $\kappa_t$  goes to 0, it approaches a uniform distribution on a unit circle. The sufficient statistic is  $\mathcal{T}(x_t) = \begin{pmatrix} \cos x_t \\ \sin x_t \end{pmatrix}$ , and the corresponding natural parameter is  $\eta_t(x_0) = \kappa_t \begin{pmatrix} \cos x_0 \\ \sin x_0 \end{pmatrix}$ . The tail statistic  $\mathcal{G}_t(x_{t:T})$  is therefore defined as follows:
$$
\mathcal {G} _ {t} \left(x _ {t: T}\right) = \sum_ {s = t} ^ {T} \kappa_ {s} \binom {\cos x _ {s}} {\sin x _ {s}} \tag {128}
$$
The KL divergence term can be calculated as follows:
$$
D _ {K L} \left(q ^ {\mathrm {s s}} \left(x _ {t} \mid x _ {0}\right) \| p _ {\theta} ^ {\mathrm {s s}} \left(x _ {t} \mid G _ {t}\right)\right) = \kappa_ {t} \frac {I _ {1} \left(\kappa_ {t}\right)}{I _ {0} \left(\kappa_ {t}\right)} \left(1 - \cos \left(x _ {0} - x _ {\theta}\right)\right) \tag {129}
$$
$t = 0$
$\mathfrak{t} = 11$
$t = 22$
$t = 33$
$\mathfrak{t} = 44$
$t = 55$
$\mathfrak{t} = 66$
$t = 77$
Figure 9: Visualization of the forward process of the von Mises-Fisher SS-DDPM on the three-dimensional unit sphere.
$t = 88$
$t = 99$
F.6 von Mises-Fisher
$$
q \left(x _ {t} \mid x _ {0}\right) = \mathrm {v M F} \left(x _ {t}; x _ {0}, \kappa_ {t}\right) \tag {130}
$$
Similar to the one-dimensional case, we set the mode of the distribution to  $x_0$ , and define the schedule using the concentration parameter  $\kappa_t$ . When  $\kappa_t$  goes to infinity, the von Mises-Fisher distribution approaches the Dirac delta, centered at  $x_0$ . When  $\kappa_t$  goes to 0, it approaches a uniform distribution on a unit sphere. The sufficient statistic is  $\mathcal{T}(x) = x$ , and the corresponding natural parameter is  $\eta_t(x_0) = \kappa_t x_0$ . The tail statistic  $\mathcal{G}_t(x_{t:T})$  is therefore defined as follows:
$$
\mathcal {G} _ {t} \left(x _ {t: T}\right) = \sum_ {s = t} ^ {T} \kappa_ {s} x _ {s} \tag {131}
$$
$$
D _ {K L} \left(q ^ {\mathrm {S S}} \left(x _ {t} \mid x _ {0}\right) \| p _ {\theta} ^ {\mathrm {S S}} \left(x _ {t} \mid G _ {t}\right)\right) = \kappa_ {t} \frac {I _ {K / 2} \left(\kappa_ {t}\right)}{I _ {K / 2 - 1} \left(\kappa_ {t}\right)} x _ {0} ^ {\top} \left(x _ {0} - x _ {\theta}\right) \tag {132}
$$
F.7 Gamma
$$
q \left(x _ {t} \mid x _ {0}\right) = \Gamma \left(x _ {t}; \alpha_ {t}, \beta_ {t}\right) \tag {133}
$$
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:
$$
\beta_ {t} \left(x _ {0}\right) = \alpha_ {t} \left(\xi_ {t} + \left(1 - \xi_ {t}\right) x _ {0} ^ {- 1}\right) \tag {134}
$$
The mean of the distribution is  $\frac{\alpha_t}{\beta_t}$ , and the variance is  $\frac{\alpha_t}{\beta_t^2}$ . Therefore, we recover the Dirac delta, centered at  $x_0$ , when we set  $\xi_t$  to 0 and  $\alpha_t$  to infinity. To achieve some standard distribution that doesn't depend on  $x_0$ , we can set  $\xi_t$  to 1 and  $\alpha_t$  to some fixed value  $\alpha_T$ .
In this parameterization, the natural parameters are  $\alpha_{t} - 1$  and  $-\beta_{t}(x_{0})$ , and the corresponding sufficient statistics are  $\log x_{t}$  and  $x_{t}$ . Since the parameter  $\alpha_{t}$  doesn't depend on  $x_{0}$ , we only need the sufficient statistic  $\mathcal{T}(x_t) = x_t$  to define the tail statistic:
$$
\mathcal {G} \left(x _ {t: T}\right) = \sum_ {s = t} ^ {T} \alpha_ {s} \left(1 - \xi_ {s}\right) x _ {s} \tag {135}
$$
The KL divergence can be computed as follows:
$$
D _ {K L} \left(q ^ {\mathrm {s s}} \left(x _ {t} \mid x _ {0}\right) \| p _ {\theta} ^ {\mathrm {s s}} \left(x _ {t} \mid G _ {t}\right)\right) = \alpha_ {t} \left[ \log \frac {\beta_ {t} \left(x _ {0}\right)}{\beta_ {t} \left(x _ {\theta}\right)} + \frac {\beta_ {t} \left(x _ {\theta}\right)}{\beta_ {t} \left(x _ {0}\right)} - 1 \right] \tag {136}
$$
$\mathbf{t} = \mathbf{0}$
t = 7
xt
$\mathfrak{t} = 14$
t = 35
Figure 10: Visualization of the forward process in Wishart SS-DDPM on positive definite matrices of size  $2 \times 2$ .
$\mathfrak{t} = 49$
F.8 Wishart
$$
q \left(X _ {t} \mid X _ {0}\right) = \mathcal {W} _ {p} \left(X _ {t}; V _ {t}, n _ {t}\right) \tag {137}
$$
The natural parameters for the Wishart distribution are  $-\frac{1}{2} V_t^{-1}$  and  $\frac{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:
$$
\mu_ {t} \left(X _ {0}\right) = \xi_ {t} I + \left(1 - \xi_ {t}\right) X _ {0} ^ {- 1} \tag {138}
$$
$$
V _ {t} \left(X _ {0}\right) = n _ {t} ^ {- 1} \mu_ {t} ^ {- 1} \left(X _ {0}\right) \tag {139}
$$
We recover the Dirac delta, centered at  $X_0$ , when we set  $\xi_t$  to 0 and  $n_t$  to infinity. To achieve some standard distribution that doesn't depend on  $X_0$ , we set  $\xi_t$  to 1 and  $n_t$  to some fixed value  $n_T$ .
The sufficient statistic for this distribution is  $\mathcal{T}(X_t) = \binom{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  $\mathcal{T}(X_t) = X_t$  to define the tail statistic:
$$
\mathcal {G} \left(X _ {t: T}, t\right) = \sum_ {s = t} ^ {T} n _ {s} \left(1 - \xi_ {s}\right) X _ {s} \tag {140}
$$
$$
D _ {K L} \left(q ^ {\mathrm {s s}} \left(x _ {t} \mid x _ {0}\right) \| p _ {\theta} ^ {\mathrm {s s}} \left(x _ {t} \mid G _ {t}\right)\right) = - \frac {n _ {t}}{2} \left[ \log \left| V _ {t} ^ {- 1} \left(X _ {\theta}\right) V _ {t} \left(X _ {0}\right) \right| - \operatorname {t r} \left(V _ {t} ^ {- 1} \left(X _ {\theta}\right) V _ {t} \left(X _ {0}\right)\right) + p \right] \tag {141}
$$
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^{\mathrm{SS}}(x_0; G_t)$ , and for DDPM it would be  $I^{\mathrm{DDPM}}(x_0; x_t)$ . We would like to start from a well-performing DDPM model and define a similar schedule by matching  $I^{\mathrm{SS}}(x_0; G_t) \approx I^{\mathrm{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_{\mathcal{N}}^{\mathrm{SS}}(x_0;G_t) = I^{\mathrm{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:
$$
I ^ {\mathrm {s s}} \left(x _ {0}; x _ {t}\right) \approx I _ {\mathcal {N}} ^ {\mathrm {s s}} \left(x _ {0}; x _ {t}\right) \stackrel {?} {\Rightarrow} I ^ {\mathrm {s s}} \left(x _ {0}; G _ {t}\right) \approx I _ {\mathcal {N}} ^ {\mathrm {s s}} \left(x _ {0}; G _ {t}\right) \tag {142}
$$
Assuming the schedule is parameterized by a single parameter  $\nu$ , we can build a look-up table  $I^{\mathrm{SS}}(x_0; x_\nu)$  for a range of parameters  $\nu$ . Then we can use binary search to build a schedule to match the mutual information  $I^{\mathrm{SS}}(x_0; x_t)$  to the mutual information schedule  $I_{\mathcal{N}}^{\mathrm{SS}}(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.