SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mingyuan Zhou*, Tianqi Chen, Zhendong Wang, and Huangjie Zheng
Austin, TX 78712
We introduce beta diffusion, a novel generative modeling method that integrates demasking and denoising to generate data within bounded ranges. Using scaled and shifted beta distributions, beta diffusion utilizes multiplicative transitions over time to create both forward and reverse diffusion processes, maintaining beta distributions in both the forward marginals and the reverse conditionals, given the data at any point in time. Unlike traditional diffusion-based generative models relying on additive Gaussian noise and reweighted evidence lower bounds (ELBOs), beta diffusion is multiplicative and optimized with KL-divergence upper bounds (KLUBs) derived from the convexity of the KL divergence. We demonstrate that the proposed KLUBs are more effective for optimizing beta diffusion compared to negative ELBOs, which can also be derived as the KLUBs of the same KL divergence with its two arguments swapped. The loss function of beta diffusion, expressed in terms of Bregman divergence, further supports the efficacy of KLUBs for optimization. Experimental results on both synthetic data and natural images demonstrate the unique capabilities of beta diffusion in generative modeling of range-bounded data and validate the effectiveness of KLUBs in optimizing diffusion models, thereby making them valuable additions to the family of diffusion-based generative models and the optimization techniques used to train them.
Diffusion-based deep generative models have been gaining traction recently. One representative example is Gaussian diffusion [57, 59, 23, 60, 35] that uses a Gaussian Markov chain to gradually diffuse images into Gaussian noise for training. The learned reverse diffusion process, defined by a Gaussian Markov chain in reverse order, iteratively refines noisy inputs towards clean photorealistic images. Gaussian diffusion can also be viewed from the lens of denoising score matching [28, 64, 59, 60] and stochastic differential equations [61]. They have shown remarkable success across a wide range of tasks, including but not limited to generating, restoring, and editing images [14, 24, 50, 53, 54, 70, 67], transforming 2D to 3D [49, 1], synthesizing audio [9, 37, 73], reinforcement learning [31, 66, 48], quantifying uncertainty [19], and designing drugs and proteins [56, 44, 32].
Constructing a diffusion-based generative model often follows a general recipe [57, 23, 35, 34]. The recipe involves three basic steps: First, defining a forward diffusion process that introduces noise into the data and corrupts it with decreasing signal-to-noise ratio (SNR) as time progresses from 0 to 1. Second, defining a reverse diffusion process that denoises the corrupted data as time reverses from 1 to 0. Third, discretizing the time interval from 0 to 1 into a finite number of intervals, and viewing the discretized forward and reverse processes as a fixed inference network and a learnable generator, respectively. Auto-encoding variational inference [36, 51] is then applied to optimize the parameters of the generator by minimizing a weighted negative ELBO that includes a Kullback-Leibler (KL) divergence-based loss term for each discretized reverse step.
*Corresponding to: mingyuan.zhou@mccombs.utexas.edu  
PyTorch code is available at: https://github.com/mingyuanzhou/Beta-Diffusion
Figure 1: Illustration of the beta forward diffusion process for two example images. The first column displays the original images, while the other 21 columns display the images noised and masked by beta diffusion at time  $t = 0, 0.05, \dots, 1$ , using  $\eta = 10000$  and the sigmoid diffusion schedule with  $c_{0} = 10$  and  $c_{1} = -13$ .
Figure 2: Illustration of reverse beta diffusion for two example generations. The time  $t$  decreases from 1 to 0 when moving from left to right. In each image generation, the top row shows the trajectory of  $z_{t}$  (rescaled for visualization), which has been demasked and denoised using reverse diffusion, whereas the bottom row shows  $\hat{x}_0 = f_\theta (z_t,t)$ , whose theoretical optimal solution is equal to  $\mathbb{E}[x_0\mid z_t]$ . See Appendix D for more details.
Although the general diffusion-modeling recipe is simple in concept, it requires access to the corrupted data at any time during the forward diffusion process given a data observation, as well as the analytic form of the conditional posterior for any earlier time given both a data observation and its corrupted version at the present time. The latter requirement, according to Bayes' rule, implies access to the analytical form of the distribution of a corrupted data observation at the present time given its value at any previous time. Linear operations of the Gaussian distributions naturally satisfy these requirements since they are conjugate to themselves with respect to their mean parameters. This means that the marginal form and the conditional distribution of the mean remain Gaussian when two Gaussian distributions are mixed. Similarly, the requirements can be met under the categorical distribution [26, 2, 18, 27] and Poisson distribution [10]. However, few additional distributions are known to meet these requirements and it remains uncertain whether negative ELBO would be the preferred loss.
While previous works have primarily used Gaussian, categorical, or Poisson distribution-based diffusion processes, this paper introduces beta diffusion as a novel addition to the family of diffusion-based generative models. Beta diffusion is specifically designed to generate data within bounded ranges. Its forward diffusion process is defined by the application of beta distributions in a multiplicative manner, whereas its reverse diffusion process is characterized by the use of scaled and shifted beta distributions. Notably, the distribution at any point in time of the forward diffusion given a data observation remains a beta distribution. We illustrate the forward beta diffusion process in Figure 1, which simultaneously adds noise to and masks the data, and the reverse one in Figure 2, which iteratively performs demasking and denoising for data generation. We provide the details on how these images are obtained in Appendix D.
Since the KL divergence between two beta distributions is analytic, one can follow the general recipe to define a negative ELBO to optimize beta diffusion. However, our experiments show that minimizing the negative ELBO of beta diffusion can fail to optimally estimate the parameters of the reverse diffusion process. For each individual time-dependent KL loss term of the negative ELBO, examining it in terms of Bregman divergence [4] reveals that the model parameters and corrupted data are placed in its first and second arguments, respectively. However, to ensure that the optimal solution under the Bregman divergence agrees with the expectation of the clean data given the corrupted data, the order of the two arguments must be swapped.
By swapping the Bregman divergence's two arguments, we obtain an upper bound on the KL divergence from the joint distribution of corrupted observations in the reverse chain to that in the forward chain. This bound arises from the convexity of the KL divergence. In addition, there exists another Bregman divergence that upper bounds the KL divergence from the univariate marginal of a corrupted observation in the reverse chain to that in the forward chain. These two Bregman divergences, which can be derived either through KL divergence upper bounds (KLUBs) or the logarithmically convex beta function, share the same optimal solution but have distinct roles in targeting reverse accuracy at each step or counteracting accumulation errors over the course of reverse
diffusion. We further demonstrate that combining these two KLUBs presents a computationally viable substitution for a KLUB derived at the chain level, which upper bounds the KL divergence from the joint distribution of all latent variables in a forward diffusion chain to that of its reverse. Either KLUB works on its own, which is not unexpected as they both share the same optimal solutions in theory, but combining them could lead to the best overall performance. In beta diffusion, the KL divergence is asymmetric, which enables us to derive an alternative set of KLUBs by swapping its two arguments. We will demonstrate that these augment-swapped KLUBs, which will be referred to as AS-KLUBs, essentially reduce to negative ELBOs. In Gaussian diffusion, the KL divergence is often made symmetric, resulting in KLUBs that are equivalent to (weighted) negative ELBOs.
Our main contributions are the introduction of beta diffusion as a novel diffusion-based multiplicative generative model for range-bounded data, as well as the proposal of KLUBs as effective loss objectives for optimizing diffusion models, in place of (weighted) negative ELBOs. Additionally, we introduce the log-beta divergence, a Bregman divergence corresponding to the differentiable and strictly convex log-beta function, as a useful tool for analyzing KLUBs. These contributions enhance the existing family of diffusion-based generative models and provide a new perspective on optimizing them.
2 Beta Diffusion and Optimization via KLUBs
We begin by specifying the general requirements for constructing a diffusion-based generative model and establish the notation accordingly [57, 23, 35]. Let  $x_0$  denote the observed data, and let  $z_s$  and  $z_t$  represent their corrupted versions at time  $s$  and time  $t$ , respectively, where  $0 < s < t < 1$ . In the forward diffusion process, we require access to random samples from the marginal distribution  $q(z_t \mid x_0)$  at any time  $t$ , as well as an analytical expression of the probability density function (PDF) of the conditional distribution  $q(z_s \mid z_t, x_0)$  for any  $s < t$ .
The forward beta diffusion chain uses diffusion scheduling parameters  $\alpha_{t}$  to control the decay of its expected value over the course of forward diffusion, given by  $\mathbb{E}[z_t\mid x_0] = \alpha_t x_0$ , and a positive concentration parameter  $\eta$  to control the tightness of the diffusion process around its expected value. We typically set  $\alpha_{t}$  to approach 1 and 0 as  $t$  approaches 0 and 1, respectively, and satisfy the condition
$$
1 \geq \alpha_ {0} > \alpha_ {s} > \alpha_ {t} > \alpha_ {1} \geq 0 \text {f o r a l l} s \in (0, t), t \in (0, 1).
$$
Let  $\Gamma(\cdot)$  denote the gamma function and  $B(\cdot, \cdot)$  denote the beta function. The beta distribution  $\mathrm{Beta}(x; a, b) = B(a, b)^{-1} x^{a - 1} (1 - x)^{b - 1}$  is a member of the exponential family [6, 65, 7]. Its log partition function is a log-beta function as  $\ln B(a, b) = \ln \Gamma(a) + \ln \Gamma(b) - \ln \Gamma(a + b)$ , which is differentiable, and strictly convex on  $(0, \infty)^2$  as a function of two variables [15]. As a result, the KL divergence between two beta distributions can be expressed as the Bregman divergence associated with the log-beta function. Specifically, as in Appendix A, one can show by their definitions that
$$
\begin{array}{l} \operatorname {K L} (\operatorname {B e t a} (\alpha_ {p}, \beta_ {p}) | | \operatorname {B e t a} (\alpha_ {q}, \beta_ {q})) = \ln \frac {B (\alpha_ {q} , \beta_ {q})}{B (\alpha_ {p} , \beta_ {p})} - (\alpha_ {q} - \alpha_ {p}, \beta_ {q} - \beta_ {p}) \left( \begin{array}{c} \nabla_ {\alpha} \ln B (\alpha_ {p}, \beta_ {p}) \\ \nabla_ {\beta} \ln B (\alpha_ {p}, \beta_ {p}) \end{array} \right) \\ = D _ {\ln B (a, b)} \left(\left(\alpha_ {q}, \beta_ {q}\right), \left(\alpha_ {p}, \beta_ {p}\right)\right). \tag {1} \\ \end{array}
$$
We refer to the above Bregman divergence as the log-beta divergence. Moreover, if  $(\alpha_{q},\beta_{q})$  are random variables, applying Proposition 1 of Banerjee et al. [4], we can conclude that the optimal value of  $(\alpha_{p},\beta_{p})$  that minimizes this log-beta divergence is  $(\alpha_{p}^{*},\beta_{p}^{*}) = \mathbb{E}[(\alpha_{q},\beta_{q})]$ .
Next, we introduce a conditional bivariate beta distribution, which given a data observation has (scaled and shifted) beta distributions for not only its two marginals but also two conditionals. These properties are important for developing the proposed diffusion model with multiplicative transitions.
2.1 Conditional Bivariate Beta Distribution
We first present the conditional bivariate beta distribution in the following Lemma, which generalizes previous results on the distribution of the product of independent beta random variables [30, 38, 33].
Lemma 1 (Conditional Beta Bivariate Distribution). Denote  $(z_s, z_t)$  as variables over a pair of time points  $(s, t)$ , with  $0 < s < t < 1$ . Given a random sample  $x_0 \in (0, 1)$  from a probability-valued data distribution  $p_{data}(x_0)$ , we define a conditional bivariate beta distribution over  $(z_s, z_t)$  with PDF:
$$
q \left(z _ {s}, z _ {t} \mid x _ {0}\right) = \frac {\Gamma (\eta)}{\Gamma \left(\eta \alpha_ {t} x _ {0}\right) \Gamma \left(\eta \left(1 - \alpha_ {s} x _ {0}\right)\right) \Gamma \left(\eta \left(\alpha_ {s} - \alpha_ {t}\right) x _ {0}\right)} \frac {z _ {t} ^ {\eta \alpha_ {t} x _ {0} - 1} \left(1 - z _ {s}\right) ^ {\eta \left(1 - \alpha_ {s} x _ {0}\right) - 1}}{\left(z _ {s} - z _ {t}\right) ^ {1 - \eta \left(\alpha_ {s} - \alpha_ {t}\right) x _ {0}}}. \tag {2}
$$
Margins: Given  $x_0$ , the two univariate marginals of this distribution are both beta distributed as
$$
q \left(z _ {s} \mid x _ {0}\right) = \operatorname {B e t a} \left(\eta \alpha_ {s} x _ {0}, \eta \left(1 - \alpha_ {s} x _ {0}\right)\right), \tag {3}
$$
$$
q \left(z _ {t} \mid x _ {0}\right) = \operatorname {B e t a} \left(\eta \alpha_ {t} x _ {0}, \eta \left(1 - \alpha_ {t} x _ {0}\right)\right). \tag {4}
$$
**Conditionals:** Given  $x_0$ , a random sample  $(z_t, z_s)$  from this distribution can be either generated in forward order, by multiplying a beta variable from (3) with beta variable  $\pi_{s \to t}$ , as
$$
z _ {t} = z _ {s} \pi_ {s \rightarrow t}, \pi_ {s \rightarrow t} \sim \operatorname {B e t a} \left(\eta \alpha_ {t} x _ {0}, \eta \left(\alpha_ {s} - \alpha_ {t}\right) x _ {0}\right), z _ {s} \sim q \left(z _ {s} \mid x _ {0}\right), \tag {5}
$$
or generated in reverse order, by combining a beta variable from (4) with beta variable  $p_{s \leftarrow t}$ , as
$$
z _ {s} = z _ {t} + \left(1 - z _ {t}\right) p _ {s \leftarrow t}, \quad p _ {s \leftarrow t} \sim \operatorname {B e t a} \left(\eta \left(\alpha_ {s} - \alpha_ {t}\right) x _ {0}, \eta \left(1 - \alpha_ {s} x _ {0}\right)\right), \quad z _ {t} \sim q \left(z _ {t} \mid x _ {0}\right). \tag {6}
$$
The proof starts with applying change of variables to obtain two scaled and shifted beta distributions
$$
q \left(z _ {t} \mid z _ {s}, x _ {0}\right) = \frac {1}{z _ {s}} \operatorname {B e t a} \left(\frac {z _ {t}}{z _ {s}}; \eta \alpha_ {t} x _ {0}, \eta \left(\alpha_ {s} - \alpha_ {t}\right) x _ {0}\right), \tag {7}
$$
$$
q \left(z _ {s} \mid z _ {t}, x _ {0}\right) = \frac {1}{1 - z _ {t}} \operatorname {B e t a} \left(\frac {z _ {s} - z _ {t}}{1 - z _ {t}}; \eta \left(\alpha_ {s} - \alpha_ {t}\right) x _ {0}, \eta \left(1 - \alpha_ {s} x _ {0}\right)\right), \tag {8}
$$
and then takes the products of them with their corresponding marginals, given by (3) and (4), respectively, to show that the PDF of the joint distribution defined in (5) and that defined in (6) are both equal to the PDF of  $q(z_{s},z_{t} \mid x_{0})$  defined in (2). The detailed proof is provided in Appendix B. To ensure numerical accuracy, we will calculate  $z_{s} = z_{t} + (1 - z_{t})p_{s\leftarrow t}$  in (6) in the logit space as
$$
\operatorname {l o g i t} \left(z _ {s}\right) = \ln \left(e ^ {\log_ {1 0} (z _ {t})} + e ^ {\log_ {1 0} \left(p _ {s \leftarrow t}\right)} + e ^ {\log_ {1 0} \left(p _ {s \leftarrow t}\right)}\right). \tag {9}
$$
2.2 Continuous Beta Diffusion
Forward Beta Diffusion. We can use the conditional bivariate beta distribution to construct a forward beta diffusion chain, beginning with the beta distribution from (3) and proceeding with the scaled beta distribution from (7). The marginal at any time  $t$  for a given data observation  $x_0$ , as shown in (4), stays as beta-distributed in the forward chain. For the beta distribution given by (4), we have
$$
\mathbb {E} [ z _ {t} \mid x _ {0} ] = \alpha_ {t} x _ {0}, \mathrm {v a r} [ z _ {t} \mid x _ {0} ] = \frac {(\alpha_ {t} x _ {0}) (1 - \alpha_ {t} x _ {0})}{\eta + 1}, \mathrm {S N R} _ {t} = \left(\frac {\mathbb {E} [ z _ {t} \mid x _ {0} ]}{\mathrm {s t d} [ z _ {t} \mid x _ {0} ]}\right) ^ {2} = \frac {\alpha_ {t} x _ {0} (\eta + 1)}{1 - \alpha_ {t} x _ {0}}.
$$
Thus when  $\alpha_{t}$  approaches 0 (i.e.,  $t\to 1$ ), both  $z_{t}$  and  $\mathrm{SNR}_t$  are shrunk towards 0, and if  $\alpha_{1} = 0$ , we have  $z_{1}\sim \mathrm{Beta}(0,\eta)$ , a degenerate beta distribution that becomes a unit point mass at 0.
Infinite Divisibility. We consider beta diffusion as a form of continuous diffusion, as its forward chain is infinitely divisible given  $x_0$ . This means that for any time  $k \in (s,t)$ , we can perform forward diffusion from  $z_s$  to  $z_t$  by first setting  $z_k = z_s\pi_{s \to k}$ , where  $\pi_{s \to k} \sim \mathrm{Beta}(\eta \alpha_kx_0,\eta (\alpha_k - \alpha_s)x_0)$ , and then setting  $z_t = z_k\pi_{k \to t}$ , where  $\pi_{k \to t} \sim \mathrm{Beta}(\eta \alpha_tx_0,\eta (\alpha_t - \alpha_k)x_0)$ . The same approach can be used to show the infinite divisibility of reverse beta diffusion given  $x_0$ .
Reverse Beta Diffusion. We follow Gaussian diffusion to use  $q(z_{s} \mid z_{t}, x_{0})$  to help define  $p_{\theta}(z_{s} \mid z_{t})$  [58, 35, 72]. To construct a reverse beta diffusion chain, we will first need to learn how to reverse from  $z_{t}$  to  $z_{s}$ , where  $s < t$ . If we know the  $x_{0}$  used to sample  $z_{t}$  as in (3), then we can readily apply the conditional in (8) to sample  $z_{s}$ . Since this information is unavailable during inference, we make a weaker assumption that we can exactly sample from  $z_{t} \sim q(z_{t}) = \mathbb{E}_{x_{0} \sim p_{data}}[q(z_{t} \mid x_{0})]$ , which is the "  $x_{0}$ -free" univariate marginal at time  $t$ . It is straightforward to sample during training but must be approximated during inference. Under this weaker assumption on  $q(z_{t})$ , utilizing (8) but replacing its true  $x_{0}$  with an approximation  $\hat{x}_{0} = f_{\theta}(z_{t}, t)$ , where  $f_{\theta}$  denotes the learned generator parameterized by  $\theta$ , we introduce our "  $x_{0}$ -free" and hence "causal" time-reversal distribution as
$$
p _ {\theta} \left(z _ {s} \mid z _ {t}\right) = q \left(z _ {s} \mid z _ {t}, \hat {x} _ {0} = f _ {\theta} \left(z _ {t}, t\right)\right). \tag {10}
$$
2.3 Optimization via KLUBs and Log-Beta Divergence
KLUB for Time Reversal. The time-reversal distribution  $p_{\theta}(z_s \mid z_t)$  reaches its optimal when its product with  $q(z_t)$  becomes equivalent to  $q(z_s, z_t) = \mathbb{E}_{x_0 \sim p_{data}}[q(z_s, z_t \mid x_0)]$ , which is a marginal bivariate distribution that is "  $x_0$ -free." Thus we propose to optimize  $\theta$  by minimizing  $\mathrm{KL}(p_{\theta}(z_s \mid z_t) q(z_t) || q(z_s, z_t))$  in theory but introduce a surrogate loss in practice:
Lemma 2 (KLUB (conditional)). The KL divergence from  $q(z_{s},z_{t})$  to  $p_{\theta}(z_s|z_t)q(z_t)$ , two "  $x_0$ -free" joint distributions defined by forward and reverse diffusions, respectively, can be upper bounded:
$$
\operatorname {K L} \left(p _ {\theta} \left(z _ {s} \mid z _ {t}\right) q \left(z _ {t}\right) \mid \mid q \left(z _ {s}, z _ {t}\right)\right) \leq \operatorname {K L U B} _ {s, t} = \mathbb {E} _ {\left(z _ {t}, x _ {0}\right) \sim q \left(z _ {t} \mid x _ {0}\right) p _ {\text {d a t a}} \left(x _ {0}\right)} \left[ \operatorname {K L U B} \left(s, z _ {t}, x _ {0}\right) \right], \tag {11}
$$
$$
\operatorname {K L U B} \left(s, z _ {t}, x _ {0}\right) = \operatorname {K L} \left(q \left(z _ {s} \mid z _ {t}, \hat {x} _ {0} = f _ {\theta} \left(z _ {t}, t\right)\right) \mid \mid q \left(z _ {s} \mid z _ {t}, x _ {0}\right)\right). \tag {12}
$$
The proof in Appendix B utilizes the equation  $p_{\theta}(z_s \mid z_t) q(z_t) = \mathbb{E}_{x_0 \sim p_{data}}[p_{\theta}(z_s \mid z_t) q(z_t \mid x_0)]$  and then applies the convexity of the KL divergence [11, 74] and the definition in (10).
Log-Beta Divergence. To find out the optimal solution under KLUB, following (1), we can express  $\mathrm{KLUB}(s,z_t,x_0)$  given by (12) as a log-beta divergence as
$$
D _ {\ln B (a, b)} \left\{\left[ \eta \left(\alpha_ {s} - \alpha_ {t}\right) x _ {0}, \eta \left(1 - \alpha_ {s} x _ {0}\right) \right], \left[ \eta \left(\alpha_ {s} - \alpha_ {t}\right) f _ {\theta} \left(z _ {t}, t\right), \eta \left(1 - \alpha_ {s} f _ {\theta} \left(z _ {t}, t\right)\right) \right] \right\}. \tag {13}
$$
We note  $\mathrm{KLUB}_{s,t}$  defined in (11) can also be written as  $\mathbb{E}_{z_t\sim q(z_t)}\mathbb{E}_{x_0\sim q(x_0\mid z_t)}[\mathrm{KLUB}(s,z_t,x_0)]$  where the log-beta divergence for  $\mathrm{KLUB}(s,z_t,x_0)$ , defined as in (13), includes  $x_0\sim q(x_0\mid z_t)$  in its first argument and the generator  $f_{\theta}(z_t,t)$  in its second argument. Therefore, applying Proposition 1 of Banerjee et al. [4], we have the following Lemma.
Lemma 3. The objective  $\mathrm{KLUB}_{s,t}$  defined in (11) for any  $s < t$  is minimized when
$$
f _ {\theta^ {*}} \left(z _ {t}, t\right) = \mathbb {E} \left[ x _ {0} \mid z _ {t} \right] = \mathbb {E} _ {x _ {0} \sim q \left(x _ {0} \mid z _ {t}\right)} \left[ x _ {0} \right] \quad f o r a l l z _ {t} \sim q \left(z _ {t}\right). \tag {14}
$$
Thus under the  $\mathrm{KLUB}_{s,t}$ -optimized  $\theta^{*}$ , we have  $p_{\theta^{*}}(z_s|z_t) = q(z_s|z_t,\mathbb{E}[x_0|z_t])$ , which is different from the optimal solution of the original KL loss in (11), which is  $p_{\theta}^{*}(z_{s}|z_{t}) = q(z_{s},z_{t}) / q(z_{t}) = q(z_{s}|z_{t}) = \mathbb{E}_{x_{0}\sim q(x_{0}|z_{t})}[q(z_{s}|z_{t},x_{0})]$ . It is interesting to note that they only differ on whether the expectation is carried out inside or outside the conditional posterior.
In practice, we need to control the gap between  $p_{\theta^*}(z_s \mid z_t)$  and  $q(z_s \mid z_t)$  and hence  $s$  needs to be close to  $t$ . Furthermore, the assumption of having access to unbiased samples from the true marginal  $q(z_t)$  is also rarely met. Thus we need to discretize the time from 1 to  $t$  into sufficiently fine intervals and perform time-reversal sampling over these intervals. Specifically, we can start with  $z_1 = 0$  and iterate (10) over these intervals to obtain an approximate sample from  $z_t \sim q(z_t)$ . However, the error could accumulate along the way from  $z_1$  to  $z_t$ , to which we present a solution below.
KLUB for Error Accumulation Control. To counteract error accumulation during time reversal, we propose to approximate the true marginal  $q(z_{t}^{\prime})$  using a "distribution-cycle-consistency" approach. This involves feeding a random sample  $z_{t}$  from  $q(z_{t})$  into the generator  $f_{\theta}$ , followed by the forward marginal  $q(z_{t}^{\prime}|\hat{x}_{0} = f_{\theta}(z_{t},t))$ , with the aim of recovering the distribution  $q(z_{t}^{\prime})$  itself. Specifically, we propose to approximate  $q(z_{t}^{\prime})$  with  $p_{\theta}(z_{t}^{\prime})\coloneqq \mathbb{E}_{z_{t}\sim q(z_{t})}[q(z_{t}^{\prime}|\hat{x}_{0} = f_{\theta}(z_{t},t))]$  by minimizing  $\mathrm{KL}(p_{\theta}(z_t^{\prime})||q(z_t^{\prime}))$  in theory, but introducing a surrogate loss in practice:
Lemma 4 (KLUB (marginal)). The KL divergence  $\mathrm{KL}(p_{\theta}(z_t'))||q(z_t'))$  can be upper bounded:
$$
\mathrm {K L} \left(p _ {\theta} \left(z _ {t} ^ {\prime}\right) \| q \left(z _ {t} ^ {\prime}\right)\right) \leq \mathrm {K L U B} _ {t} = \mathbb {E} _ {\left(z _ {t}, x _ {0}\right) \sim q \left(z _ {t} \mid x _ {0}\right) p _ {\text {d a t a}} \left(x _ {0}\right)} \left[ \mathrm {K L U B} \left(z _ {t}, x _ {0}\right) \right], \tag {15}
$$