SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 2: Comparing FID scores for KLUB and negative ELBO-optimized Beta Diffusion on CIFAR-10 with varying NFE under  $\eta = {10000}$  and two different mini-batch sizes  $B$  .
Loss B	-ELBO 512	-ELBO 288	KLUB 512	KLUB 288
20	16.04	16.10	17.06	16.09
50	6.82	6.82	6.48	5.96
200	4.55	4.84	3.69	3.31
500	4.39	4.65	3.45	3.10
1000	4.41	4.61	3.38	3.08
2000	4.50	4.66	3.37	3.06
(a) -ELBO
Figure 4: Uncurated randomly-generated images by beta diffusion optimized with -ELBO or KLUB with  $\eta = 10000$  and  $B = 288$ . The generation with NFE = 1000 starts from the same random seed.
(b) KLUB
Table 2 presents a comparison between KLUB and negative ELBO-optimized beta diffusion across different NFE under two different mini-batch sizes  $B$ . Table 3 in the Appendix includes the results under several different combinations of  $\eta$  and  $B$ . We also include Figure 4 to visually compare generated images under KLUB and negative ELBO. The findings presented in Tables 2 and 3 and Figure 4 provide further validation of KLUB's efficacy in optimizing beta diffusion.
As each training run takes a long time and FID evaluation is also time-consuming, we have not yet optimized the combination of these hyperparameters given the limit of our current computation resources. Thus the results reported in this paper, while demonstrating that beta diffusion can provide competitive image generation performance, do not yet reflect the full potential of beta diffusion. These results are likely to be further improved given an optimized hyperparameter setting or a network architecture that is tailored to beta diffusion. We leave these further investigations to our future work.
We introduce beta diffusion characterized by the following properties: 1) Analytic Marginal: Given a probability-valued data observation  $x_0 \in (0,1)$ , the distribution at any time point  $t \in [0,1]$  of the forward beta diffusion chain, expressed as  $q(z_{t} \mid x_{0})$ , is a beta distribution. 2) Analytical Conditional: Conditioning on a data  $x_0$  and a forward-sampled latent variable  $z_{t} \sim q(z_{t} \mid x_{0})$ , the forward beta diffusion chain can be reversed from time  $t$  to the latent variable at any previous time  $s \in [0,t)$  by sampling from an analytic conditional posterior  $z_{s} \sim q(z_{s} \mid z_{t},x_{0})$  that follows a scaled and shifted beta distribution. 3) KLUBs: We introduce the combination of two different Kullback-Leibler Upper Bounds (KLUBs) for optimization and represent them under the log-beta Bregman divergence, showing that their optimal solutions of the generator are both achieved at  $f_{\theta^*}(z_t,t) = \mathbb{E}[x_0 \mid z_t]$ . We also establish the connection between augment-swapped KLUBs and (weighted) negative ELBOs for diffusion models. Our experimental results confirm the distinctive qualities of beta diffusion when applied to generative modeling of range-bounded data spanning disjoint regions or residing in high-dimensional spaces, as well as the effectiveness of KLUBs for optimizing beta diffusion.
The authors acknowledge the support of NSF-IIS 2212418, NIH-R37 CA271186, the Fall 2022 McCombs REG award, the Texas Advanced Computing Center (TACC), and the NSF AI Institute for Foundations of Machine Learning (IFML).
[66] Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=AHvFDPi-FA.
[67] Zhendong Wang, Yifan Jiang, Yadong Lu, Yelong Shen, Pengcheng He, Weizhu Chen, Zhangyang Wang, and Mingyuan Zhou. In-context learning unlocked for diffusion models. arXiv preprint arXiv:2305.01115, 2023.
[68] Zhendong Wang, Yifan Jiang, Huangjie Zheng, Peihao Wang, Pengcheng He, Zhangyang Wang, Weizhu Chen, and Mingyuan Zhou. Patch diffusion: Faster and more data-efficient training of diffusion models. arXiv preprint arXiv:2304.12526, 2023. URL https://arxiv.org/abs/2304.12526.
[69] Zhendong Wang, Huangjie Zheng, Pengcheng He, Weizhu Chen, and Mingyuan Zhou. Diffusion-GAN: Training GANs with diffusion. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=HZf7UbpWHuA.
[70] Zhixin Wang, Xiaoyun Zhang, Ziying Zhang, Huangjie Zheng, Mingyuan Zhou, Ya Zhang, and Yanfeng Wang. DR2: Diffusion-based robust degradation remover for blind face restoration. In CVPR 2023: Conference on Computer Vision and Pattern Recognition, 2023. URL https://openreview.net/forum?id=HmcDUPEGA0.
[71] Zike Wu, Pan Zhou, Kenji Kawaguchi, and Hanwang Zhang. Fast diffusion model. arXiv preprint arXiv:2306.06991, 2023.
[72] Zhisheng Xiao, Karsten Kreis, and Arash Vahdat. Tackling the generative learning trilemma with denoising diffusion GANs. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=JprM0p-qOCo.
[73] Dongchao Yang, Jianwei Yu, Helin Wang, Wen Wang, Chao Weng, Yuexian Zou, and Dong Yu. Diffsound: Discrete diffusion model for text-to-sound generation. IEEE/ACM Transactions on Audio, Speech, and Language Processing, 31:1720-1733, 2023. doi: 10.1109/TASLP.2023.3268730.
[74] Mingzhang Yin and Mingyuan Zhou. Semi-implicit variational inference. In International Conference on Machine Learning, pages 5660–5669, 2018.
[75] Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=Loek7hfb46P.
[76] Hongkai Zheng, Weili Nie, Arash Vahdat, Kamyar Azizzadenesheli, and Anima Anandkumar. Fast sampling of diffusion models via operator learning. In International Conference on Machine Learning, pages 42390-42402. PMLR, 2023.
[77] Huangjie Zheng and Mingyuan Zhou. Exploiting chain rule and Bayes' theorem to compare probability distributions. Advances in Neural Information Processing Systems, 34:14993-15006, 2021.
[78] Huangjie Zheng, Pengcheng He, Weizhu Chen, and Mingyuan Zhou. Truncated diffusion probabilistic models and diffusion-based adversarial auto-encoders. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=HDxgaKk9561.
Algorithm 1 Training of Beta Diffusion
Require: Dataset  $\mathcal{D}$  whose values are bounded from 0 to 1, Mini-batch size  $B$ , concentration parameter  $\eta = 10000$ , data shifting parameter  $S_{shift} = 0.6$ , data scaling parameter  $S_{calc} = 0.39$ , generator  $f_{\theta}$ , loss balance coefficient  $\omega = 0.99$ , time reversal coefficient  $\pi = 0.95$ , and a decreasing function that returns scheduling parameter  $\alpha_{t} \in (0,1)$  given  $t \in [0,1]$ , such as a beta linear schedule defined by  $\alpha_{t} = e^{-0.5(\beta_{\max} - \beta_{\min})t^{2} - \beta_{\min}t}$ , where  $\beta_{\max} = 20$  and  $\beta_{\min} = 0.1$  and a sigmoid schedule defined by  $\alpha_{t} = 1 / (1 + e^{-c_{0} - (c_{1} - c_{0})t})$ , where  $c_{0} = 10$  and  $c_{1} = -13$ .
1: repeat
2: Draw a mini-batch  $X_0 = \{x_0^{(i)}\}_{i=1}^B$  from  $\mathcal{D}$
3: for  $i = 1$  to  $B$  do
4:  $t_i\sim \mathrm{Unif}(1\mathrm{e} - 5,1)$
5:  $s_i = \pi t_i$
6: Compute  $\alpha_{s_i}$  and  $\alpha_{t_i}$
7:  $x_0^{(i)} = x_0^{(i)}*S_{\text{cal}} + S_{\text{hif}}$
8:  $z_{t_i}\sim \mathrm{Beta}(\eta \alpha_{t_i}x_0^{(i)},\eta (1 - \alpha_{t_i}x_0^{(i)}))$
9:  $\hat{x}_0^{(i)} = f_\theta (z_{t_i},t_i)*S_{calc} + S_{shift}$
0: Using (18) to compute the loss as
can be run in parallel
$$
\begin{array}{l} \mathcal {L} _ {i} = \omega \operatorname {K L U B} \left(s _ {i}, z _ {t _ {i}}, x _ {0} ^ {(i)}\right) + (1 - \omega) \operatorname {K L U B} \left(z _ {t _ {i}}, x _ {0} ^ {(i)}\right) \\ = \omega D _ {\ln B (a, b)} \left\{\left[ \eta \left(\alpha_ {s _ {i}} - \alpha_ {t _ {i}}\right) x _ {0} ^ {(i)}, \eta \left(1 - \alpha_ {s _ {i}} x _ {0} ^ {(i)}\right) \right], \left[ \eta \left(\alpha_ {s _ {i}} - \alpha_ {t _ {i}}\right) \hat {x} _ {0} ^ {(i)}, \eta \left(1 - \alpha_ {s _ {i}} \hat {x} _ {0} ^ {(i)}\right) \right] \right\} \\ + (1 - \omega) D _ {\ln B (a, b)} \left\{\left[ \eta \alpha_ {t _ {i}} x _ {0} ^ {(i)}, \eta \left(1 - \alpha_ {t _ {i}} x _ {0} ^ {(i)}\right) \right], \left[ \eta \alpha_ {t _ {i}} \hat {x} _ {0} ^ {(i)}, \eta \left(1 - \alpha_ {t _ {i}} \hat {x} _ {0} ^ {(i)}\right) \right] \right\} \tag {24} \\ \end{array}
$$
or swap the two arguments of both log-beta Bregman divergences in (24) if the loss is (weighted) -ELBO.
11: end for
12: Perform SGD with  $\frac{1}{B}\nabla_{\theta}\sum_{i = 1}^{B}\mathcal{L}_{i}$
13: until converge
Table 3: Comparing FID scores for KLUB and negative ELBO-optimized Beta Diffusion on the CIFAR-10 image dataset with varying NFE under several different combinations of concentration parameter  $\eta$  and mini-batch size  $B$ . We train the model with 200M images and use it to calculate FID. We compute FID one time in each experiment. The other model parameters are set as  $S_{\text{scale}} = 0.39$ ,  $S_{\text{shift}} = 0.60$ ,  $\pi = 0.95$ ,  $\omega = 0.99$ ,  $lr = 2\mathrm{e} - 4$ ,  $c_0 = 10$ , and  $c_1 = -13$ .
Loss	η × 10-4	B	NFE = 10	20	50	200	500	1000	2000
-ELBO	1	512	37.64	16.04	6.82	4.55	4.39	4.41	4.50
-ELBO	1	288	37.54	16.10	6.82	4.84	4.65	4.61	4.66
KLUB	1	512	39.05	17.06	6.48	3.69	3.45	3.38	3.37
KLUB	0.1	512	41.28	20.38	9.72	5.85	4.98	4.90	4.88
KLUB	0.3	512	36.70	16.47	7.03	4.10	3.65	3.67	3.66
KLUB	1	288	37.67	16.09	5.96	3.31	3.10	3.08	3.06
KLUB	0.8	288	36.46	15.58	5.98	3.49	3.22	3.21	3.23
KLUB	1.2	288	38.49	16.59	6.31	3.68	3.36	3.37	3.24
KLUB	1	128	38.16	16.47	6.29	3.74	3.44	3.40	3.47
KLUB	2	128	39.80	17.09	6.58	3.93	3.76	3.75	3.65
Beta Diffusion: Appendix
Algorithm 2 Sampling of Beta Diffusion
Require: Number of function evaluations (NFE)  $J = 200$ , generator  $f_{\theta}$ , and timesteps  $\{t_j\}_{j=0}^J$ :  $t_j = 1 - (1 - 1e-5) * (J - j) / (J - 1)$  for  $j = 1, \dots, J$  and  $t_0 = 0$
1: if NFE  $>350$  then
2:  $\alpha_{t_j} = 1 / (1 + e^{-c_0 - (c_1 - c_0)t_j})$
3: else
4:  $\alpha_{t_j} = (1 / (1 + e^{-c_1}))^{t_j}$
5: end if
6: Initialize  $\hat{x}_0 = \mathbb{E}[x_0] * S_{\text{scale}} + S_{\text{shift}}$
7:  $z_{t_J} \sim \mathrm{Beta}(\eta \alpha_{t_J} \hat{x}_0, \eta (1 - \alpha_{t_J} \hat{x}_0))$
8: for  $j = J$  to 1 do
9:  $\hat{x}_0 = f_\theta (z_{t_j},\alpha_{t_j})*S_{calc} + S_{shift}$
10:  $p_{(t_{j - 1}\leftarrow t_j)}\sim \mathrm{Beta}\left(\eta (\alpha_{t_{j - 1}} - \alpha_{t_j})\hat{x}_0,\eta (1 - \alpha_{t_{j - 1}}\hat{x}_0)\right),$
$\triangleright$  which is implemented in the logit space as
$$
\begin{array}{l} \operatorname {l o g i t} \left(p _ {\left(t _ {j - 1} \leftarrow t _ {j}\right)}\right) = \ln u - \ln v, \\ u \sim \operatorname {G a m m a} \left(\eta \left(\alpha_ {t _ {j - 1}} - \alpha_ {t _ {j}}\right) \hat {x} _ {0}, 1\right), \\ v \sim \operatorname {G a m m a} \left(\eta \left(1 - \alpha_ {t _ {j - 1}} \hat {x} _ {0}\right), 1\right) \\ \end{array}
$$
11:  $z_{t_{j-1}} = z_{t_j} + (1 - z_{t_j})p_{(t_{j-1}\leftarrow t_j)}$ ,
$$
\operatorname {l o g i t} \left(z _ {t _ {j - 1}}\right) = \ln \left(e ^ {\log \mathrm {i t} \left(z _ {t _ {j}}\right)} + e ^ {\log \mathrm {i t} \left(p _ {\left(t _ {j - 1} \leftarrow t _ {j}\right)}\right)} + e ^ {\log \mathrm {i t} \left(z _ {t _ {j}}\right) + \log \mathrm {i t} \left(p _ {\left(t _ {j - 1} \leftarrow t _ {j}\right)}\right)}\right)
$$
12: end for
13: return  $(\hat{x}_0 - S_{shift}) / S_{scale}$  or  $(z_{t_0} / \alpha_{t_0} - S_{shift}) / S_{scale}$
A Log-beta Divergence and KL Divergence between Beta Distributions
By the definition of Bregman divergence, the log-beta divergence corresponding to the log-beta function  $\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  $a$  and  $b$ , can be expressed as
$$
\begin{array}{l} D _ {\ln B (a, b)} \left(\left(\alpha_ {q}, \beta_ {q}\right), \left(\alpha_ {p}, \beta_ {p}\right)\right) \\ = \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_ {p}} \ln B (\alpha_ {p}, \beta_ {p}) \\ \nabla_ {\beta_ {p}} \ln B (\alpha_ {p}, \beta_ {p}) \end{array} \right) \\ = \ln \frac {B (\alpha_ {q} , \beta_ {q})}{B (\alpha_ {p} , \beta_ {p})} - (\alpha_ {q} - \alpha_ {p}, \beta_ {q} - \beta_ {p}) \left( \begin{array}{c} \psi (\alpha_ {p}) - \psi (\alpha_ {p} + \beta_ {p}) \\ \psi (\beta_ {p}) - \psi (\alpha_ {p} + \beta_ {p}) \end{array} \right) \\ = \ln \frac {B (\alpha_ {q} , \beta_ {q})}{B (\alpha_ {p} , \beta_ {p})} - (\alpha_ {q} - \alpha_ {p}) \psi (\alpha_ {p}) - (\beta_ {q} - \beta_ {p}) \psi (\beta_ {p}) + (\alpha_ {q} - \alpha_ {p} + \beta_ {q} - \beta_ {p}) \psi (\alpha_ {p} + \beta_ {p}), \\ \end{array}
$$
which is equivalent to the analytic expression of
$$
\operatorname {K L} \left(\operatorname {B e t a} \left(\alpha_ {p}, \beta_ {p}\right) \mid \mid \operatorname {B e t a} \left(\alpha_ {q}, \beta_ {q}\right)\right).
$$
B Proof
Proof of Lemma 1. The joint distribution of  $z_{t}$  and  $z_{s}$  in (5) can be expressed as
$$
\begin{array}{l} q \left(z _ {t}, z _ {s} \mid x _ {0}\right) = q \left(z _ {t} \mid z _ {s}\right) q \left(z _ {s} \mid 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) \operatorname {B e t a} \left(z _ {s}; \eta \alpha_ {s} x _ {0}, \eta \left(1 - \alpha_ {s} x _ {0}\right)\right) \\ = \frac {1}{z _ {s}} \frac {\Gamma (\eta \alpha_ {s} x _ {0})}{\Gamma (\eta \alpha_ {t} x _ {0}) \Gamma (\eta (\alpha_ {s} - \alpha_ {t}) x _ {0})} \left(\frac {z _ {t}}{z _ {s}}\right) ^ {\eta \alpha_ {t} x _ {0} - 1} \left(1 - \frac {z _ {t}}{z _ {s}}\right) ^ {\eta (\alpha_ {s} - \alpha_ {t}) x _ {0} - 1} \\ \times \frac {\Gamma (\eta)}{\Gamma (\eta \alpha_ {s} x _ {0}) \Gamma (\eta (1 - \alpha_ {s} x _ {0}))} z _ {s} ^ {\eta \alpha_ {s} x _ {0} - 1} (1 - z _ {s}) ^ {\eta (1 - \alpha_ {s} x _ {0}) - 1} \\ = \frac {\Gamma (\eta)}{\Gamma (\eta \alpha_ {t} x _ {0}) \Gamma (\eta (\alpha_ {s} - \alpha_ {t}) x _ {0}) \Gamma (\eta (1 - \alpha_ {s} x _ {0}))} \\ \times z _ {t} ^ {\eta \alpha_ {t} x _ {0} - 1} \left(z _ {s} - z _ {t}\right) ^ {\eta \left(\alpha_ {s} - \alpha_ {t}\right) x _ {0} - 1} \left(1 - z _ {s}\right) ^ {\eta \left(1 - \alpha_ {s} x _ {0}\right) - 1}. \tag {25} \\ \end{array}
$$
The joint distribution of  $z_{t}$  and  $z_{s}$  in (6) can be expressed as
$$
\begin{array}{l} q \left(z _ {t}, z _ {s} \mid x _ {0}\right) = q \left(z _ {s} \mid z _ {t}, x _ {0}\right) q \left(z _ {t} \mid 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) \operatorname {B e t a} \left(z _ {t}; \eta \alpha_ {t} x _ {0}, \eta \left(1 - \alpha_ {t} x _ {0}\right)\right) \\ = \frac {1}{1 - z _ {t}} \frac {\Gamma (\eta (1 - \alpha_ {t} x _ {0}))}{\Gamma (\eta (1 - \alpha_ {s} x _ {0})) \Gamma (\eta (\alpha_ {s} - \alpha_ {t}) x _ {0})} \\ \times \left(\frac {z _ {s} - z _ {t}}{1 - z _ {t}}\right) ^ {\eta (\alpha_ {s} - \alpha_ {t}) x _ {0} - 1} \left(1 - \frac {z _ {s} - z _ {t}}{1 - z _ {t}}\right) ^ {\eta (1 - \alpha_ {s} x _ {0}) - 1} \\ \times \frac {\Gamma (\eta)}{\Gamma (\eta \alpha_ {t} x _ {0}) \Gamma (\eta (1 - \alpha_ {t}) x _ {0})} z _ {t} ^ {\eta \alpha_ {t} x _ {0} - 1} (1 - z _ {t}) ^ {\eta (1 - \alpha_ {t} x _ {0}) - 1} \\ = \frac {\Gamma (\eta)}{\Gamma \left(\eta \alpha_ {t} x _ {0}\right) \Gamma \left(\eta \left(\alpha_ {s} - \alpha_ {t}\right) x _ {0}\right) \Gamma \left(\eta \left(1 - \alpha_ {s} x _ {0}\right)\right)} \\ \times \left(z _ {s} - z _ {t}\right) ^ {\eta \left(\alpha_ {s} - \alpha_ {t}\right) x _ {0} - 1} \left(1 - z _ {s}\right) ^ {\eta \left(1 - \alpha_ {s} x _ {0}\right) - 1} z _ {t} ^ {\eta \alpha_ {t} x _ {0} - 1}. \tag {26} \\ \end{array}
$$
The joint distribution shown in (25) is the same as that in (26).
Proof of Lemma 2. Since we can re-express  $q(z_{t})$  and  $q(z_{s},z_{t})$  as
$$
q \left(z _ {t}\right) = \mathbb {E} _ {x _ {0} \sim p _ {d a t a} \left(x _ {0}\right)} \left[ q \left(z _ {t} \mid x _ {0}\right) \right],
$$
$$
q \left(z _ {s}, z _ {t}\right) = \mathbb {E} _ {x _ {0} \sim p _ {\text {d a t a}} \left(x _ {0}\right)} \left[ q \left(z _ {s}, z _ {t} \mid x _ {0}\right) \right],
$$
and  $q(z_{s}, z_{t} \mid x_{0}) = q(z_{s} \mid z_{t}, x_{0})q(z_{t} \mid x_{0})$ , using the convexity of the KL divergence, we have
$$
\begin{array}{l} \mathrm {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) \\ = \operatorname {K L} \left(\mathbb {E} _ {x _ {0} \sim p _ {d a t a} (x _ {0})} \left[ p _ {\theta} \left(z _ {s} \mid z _ {t}\right) q \left(z _ {t} \mid x _ {0}\right) \right] \mid \mid \mathbb {E} _ {x _ {0} \sim p _ {d a t a} (x _ {0})} \left[ q \left(z _ {s}, z _ {t} \mid x _ {0}\right) \right]\right) \\ \leq \mathbb {E} _ {x _ {0} \sim p _ {\text {d a t a}} (x _ {0})} [ \mathrm {K L} (p _ {\theta} (z _ {s} \mid z _ {t}) q (z _ {t} \mid x _ {0}) | | q (z _ {s}, z _ {t} \mid x _ {0})) ] \\ = \mathbb {E} _ {x _ {0} \sim p _ {d a t a} (x _ {0})} \mathbb {E} _ {(z _ {s}, z _ {t}) \sim p _ {\theta} (z _ {s} \mid z _ {t}) q (z _ {t} \mid x _ {0})} \ln \frac {p _ {\theta} (z _ {s} \mid z _ {t}) q (z _ {t} \mid x _ {0})}{q (z _ {s} \mid z _ {t} , x _ {0}) q (z _ {t} \mid x _ {0})} \\ = \mathbb {E} _ {x _ {0} \sim p _ {\text {d a t a}} (x _ {0})} \mathbb {E} _ {z _ {t} \sim q (z _ {t} \mid x _ {0})} \mathbb {E} _ {z _ {s} \sim p _ {\theta} (z _ {s} \mid z _ {t})} \ln \frac {p _ {\theta} (z _ {s} \mid z _ {t})}{q (z _ {s} \mid z _ {t} , x _ {0})} \\ = \mathbb {E} _ {\left(z _ {t}, x _ {0}\right) \sim q \left(z _ {t} \mid x _ {0}\right) p _ {d a t a} \left(x _ {0}\right)} \left[ \mathrm {K L} \left(p _ {\theta} \left(z _ {s} \mid z _ {t}\right) | | q \left(z _ {s} \mid z _ {t}, x _ {0}\right)\right) \right] \\ = \mathbb {E} _ {\left(z _ {t}, x _ {0}\right) \sim q \left(z _ {t} \mid x _ {0}\right) p _ {d a t a} \left(x _ {0}\right)} [ \mathrm {K L} \left(q \left(z _ {s} \mid z _ {t}, \hat {x} _ {0} = f _ {\theta} \left(z _ {t}, t\right) \right| | q \left(z _ {s} \mid z _ {t}, x _ {0}\right)\right) ]. \\ \end{array}
$$
Proof of Lemma 4. Since we can re-express  $p_{\theta}(z_t')$  and  $q(z_t')$  as
$$
\begin{array}{l} p _ {\theta} (z _ {t} ^ {\prime}) := \mathbb {E} _ {z _ {t} \sim q (z _ {t})} [ q (z _ {t} ^ {\prime} | \hat {x} _ {0} = f _ {\theta} (z _ {t}, t)) ] = \mathbb {E} _ {(z _ {t}, x _ {0}) \sim q (z _ {t} | x _ {0}) p _ {d a t a} (x _ {0})} [ q (z _ {t} ^ {\prime} | \hat {x} _ {0} = f _ {\theta} (z _ {t}, t)) ], \\ q (z _ {t} ^ {\prime}) = \mathbb {E} _ {(z _ {t}, x _ {0}) \sim q (z _ {t} | x _ {0}) p _ {d a t a} (x _ {0})} [ q (z _ {t} ^ {\prime} | x _ {0}) ], \\ \end{array}
$$
using the convexity of the KL divergence, we have
$$
\begin{array}{l} \mathrm {K L} \left(p _ {\theta} \left(z _ {t} ^ {\prime}\right) | | q \left(z _ {t} ^ {\prime}\right)\right) \\ = \operatorname {K L} \left(\mathbb {E} _ {\left(z _ {t}, x _ {0}\right) \sim q \left(z _ {t} \mid x _ {0}\right) p _ {d a t a} \left(x _ {0}\right)} \left[ q \left(z _ {t} ^ {\prime} \mid \hat {x} _ {0} = f _ {\theta} \left(z _ {t}, t\right)\right) \right] \mid \mid \mathbb {E} _ {\left(z _ {t}, x _ {0}\right) \sim q \left(z _ {t} \mid x _ {0}\right) p _ {d a t a} \left(x _ {0}\right)} \left[ q \left(z _ {t} ^ {\prime} \mid x _ {0}\right) \right]\right) \\ \leq \mathbb {E} _ {\left(z _ {t}, x _ {0}\right) \sim q \left(z _ {t} \mid x _ {0}\right) p _ {d a t a} \left(x _ {0}\right)} \left[ \mathrm {K L} \left(q \left(z _ {t} ^ {\prime} \mid f _ {\theta} \left(z _ {t}, t\right)\right) | | q \left(z _ {t} ^ {\prime} \mid x _ {0}\right)\right) \right]. \\ \end{array}
$$
Proof of Lemma 6. The proof utilizes the convexity of the KL divergence to show that
$$
\begin{array}{l} \operatorname {K L} \left(p _ {\theta} \left(z _ {t _ {1: T}}\right) | | q \left(z _ {t _ {1: T}}\right)\right) = \operatorname {K L} \left(\mathbb {E} _ {x _ {0} \sim p _ {d a t a} (x _ {0})} \left[ p _ {\theta} \left(z _ {1: T} \mid x _ {0}\right) \right] | | \mathbb {E} _ {x _ {0} \sim p _ {d a t a} (x _ {0})} \left[ q \left(z _ {1: T} \mid z _ {0}\right) \right]\right) \\ \leq \mathbb {E} _ {x _ {0} \sim p _ {\text {d a t a}} (x _ {0})} \mathrm {K L} \left(p _ {\theta} \left(z _ {t _ {1: T}} \mid x _ {0}\right) | | q \left(z _ {t _ {1: T}} \mid x _ {0}\right)\right) \\ = \mathbb {E} _ {x _ {0} \sim p _ {d a t a} (x _ {0})} \mathbb {E} _ {z _ {t _ {1: T}} \sim p _ {\theta} (z _ {t _ {1: T}} \mid x _ {0})} \ln \frac {p _ {\theta} (z _ {t _ {1 : T}} \mid x _ {0})}{q (z _ {t _ {1 : T}} \mid x _ {0})} \\ = \mathbb {E} _ {x _ {0} \sim p _ {d a t a} (x _ {0})} \left[ \mathrm {K L} \left(p _ {p r i o r} \left(z _ {t _ {T}} \mid x _ {0}\right) | | q \left(z _ {t _ {T}} \mid x _ {0}\right)\right) \right. \\ + \sum_ {j = 2} ^ {T} \mathbb {E} _ {z _ {t _ {j}} \sim p _ {\theta} (z _ {t _ {j}} \mid x _ {0})} \mathbb {E} _ {z _ {t _ {j - 1}} \sim p _ {\theta} (z _ {t _ {j - 1}} \mid z _ {t _ {j}})} \ln \left. \frac {p _ {\theta} (z _ {t _ {j - 1}} \mid z _ {t _ {j}})}{q (z _ {t _ {j - 1}} \mid z _ {t _ {j}} , x _ {0})} \right]. \\ \end{array}
$$
We note  $p_{\theta}(z_{t_j} \mid x_0) \propto p(x_0 \mid z_{t_j}) p_{\theta}(z_{t_j})$ , which in general is intractable to sample from, motivating us to replace  $p_{\theta}(z_{t_j} \mid x_0)$  with  $q(z_{t_j} \mid x_0)$  for tractable computation.
Proof of Lemma 7. Utilizing the convexity of the KL divergence, we present an upper-bound of the augmented-swapped KL divergence  $\mathrm{KL}(q(z_{t_{1:T}}) || p_{\theta}(z_{t_{1:T}}))$ , referred to as AS-KLUB, as
$$
\begin{array}{l} \operatorname {K L} \left(q \left(z _ {t _ {1: T}}\right) \mid \mid p _ {\theta} \left(z _ {t _ {1: T}}\right)\right) = \operatorname {K L} \left(\mathbb {E} _ {x _ {0} \sim p _ {d a t a} (x _ {0})} [ q \left(z _ {t _ {1: T}} \mid x _ {0}\right) ] \mid \mid \mathbb {E} _ {x _ {0} \sim p _ {d a t a} (x _ {0})} \left[ p _ {\theta} \left(z _ {t _ {1: T}} \mid x _ {0}\right) \right]\right) \\ \leq \operatorname {A S - K L U B} = \mathbb {E} _ {x _ {0} \sim p _ {\text {d a t a}} (x _ {0})} \operatorname {K L} \left(q \left(z _ {t _ {1}: T} \mid x _ {0}\right) \| p _ {\theta} \left(z _ {t _ {1}: T} \mid x _ {0}\right)\right) \\ = \mathbb {E} _ {x _ {0} \sim p _ {d a t a} (x _ {0})} \mathbb {E} _ {z _ {t _ {1: T}} \sim q (z _ {t _ {1: T}} \mid x _ {0})} \ln \frac {q (z _ {t _ {1 : T}} \mid x _ {0})}{p _ {\theta} (z _ {t _ {1 : T}} \mid x _ {0})} \\ = \mathbb {E} _ {x _ {0} \sim p _ {d a t a} (x _ {0})} \left[ \mathrm {K L} \left(q \left(z _ {t _ {T}} \mid x _ {0}\right) | | p _ {p r i o r} \left(z _ {t _ {T}} \mid x _ {0}\right)\right) \right. \\ + \sum_ {j = 2} ^ {T} \mathbb {E} _ {z _ {t _ {j}} \sim q (z _ {t _ {j}} \mid x _ {0})} \mathbb {E} _ {z _ {t _ {j - 1}} \sim q (z _ {t _ {j - 1}} \mid z _ {t _ {j}}, x _ {0})} \ln \left. \frac {q (z _ {t _ {j - 1}} \mid z _ {t _ {j}} , x _ {0})}{p _ {\theta} (z _ {t _ {j - 1}} \mid z _ {t _ {j}})} \right]. \\ \end{array}
$$
Utilizing the convexity of the negative logarithmic function, we have