SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $\alpha_{t} = \alpha(s_{t})$  with  $\alpha(s)$  defined in equation (12). In this example, the score function  $S_{t}(\boldsymbol{x}_{t}) = \nabla_{\boldsymbol{x}_{t}} \log p_{t}(\boldsymbol{x}_{t})$  has an analytical form
$$
\boldsymbol {S} _ {t} (\boldsymbol {x} _ {t}) = \boldsymbol {S} _ {0} \left(\boldsymbol {x} _ {t}; \sqrt {\alpha_ {t}} \boldsymbol {\mu}, \left(\alpha_ {t} \sigma^ {2} + 1 - \alpha_ {t}\right) \mathbf {I}\right), \quad 0 \leq t \leq T - 1.
$$
where  $S_0(\pmb{x}; \pmb{\mu}, \sigma^2\mathbf{I})$  is the score function of the Gaussian mixture model  $p(\pmb{x}; \pmb{\mu}, \sigma^2\mathbf{I}) = \sum_{i=1}^{8} 1/8 \cdot \mathcal{N}(\pmb{x}; \pmb{\mu}_i, \sigma^2\mathbf{I})$ . We set the number of layers  $T = 5$  and  $\alpha_t = 1 - t/5$  for  $t = 0, \dots, 4$ . Figure 8 shows the sample trajectories generated by HSIVI. We see clearly that semi-implicit distributions are guided toward the target distribution following the diffusion bridge.
D.2 High-dimensional conditioned diffusion
We also test SIVI-LB and HSIVI-LB for fitting the posterior in high-dimensional conditioned diffusion. The auxiliary bridge is formed using the same geometric interpolation as for HSIVI-SM,
Figure 10: Comparison of 10,000 samples generated by DDPM, DDIM, and HSIVI-SM.
Table 4: Frobenius distances between the estimated covariance matrices and that of the ground truth. For each method, we collect 100,000 samples to estimate the covariance matrix.
T=1	T=2	T=3	T=5
HSIVI-SM	0.0886	0.0813	0.0431	0.0333
HSIVI-LB	0.0883	0.0825	0.0722	0.0433
$$
p _ {\text {b a s e}} = \mathcal {N} (\boldsymbol {x}; \boldsymbol {y}, \sigma^ {2} \mathbf {I}), \quad \lambda_ {t} = 1 - \frac {t}{T - 1} \text {f o r} 0 \leq t \leq T - 1.
$$
From Figure 9, we see that SIVI-LB also underestimates the posterior variance and 5-layer HSIVI-LB fits the variance better. This phenomenon is also observed in the performances of SIVI-SM and HSIVI-SM in Figure 3. The quantitative comparison between different numbers of layers is reported in Table 4, where we see that for both HSIVI-SM and HSIVI-LB, the variational approximation gets more accurate with more layers. We also find that HSIVI-SM fits better than HSIVI-LB consistently.
D.3 Toy examples of diffusion model acceleration
We compare the samples from DDPM, DDIM, and our proposed HSIVI-SM with 5 and 10 steps in Figure 10. We find that DDIM and DDPM fail to converge to the target distribution with a small number of steps, while HSIVI-SM can provide noticeably better samples. Moreover, DDPM tends to underestimate the variance as evidenced by the narrower region occupied by the samples.
D.4 MNIST
Figure 11 shows the samples from DDPM, DDIM, and HSIVI-SM with  $T = 10$  steps. We see that the samples produced by HSIVI-SM is much cleaner and more recognizable than those produced by DDPM and DDIM.
DDPM  $(T = 10)$
DDIM  $(T = 10)$
Figure 11: Comparison of the quality of uncurated samples generated by DDPM, DDIM, and HSIVISM with 10 discrete time steps on MNIST.
HSIVI-SM (T=10)
Table 5: Number of parameters in the score model (or noise model) used by different methods in Table 2. 'M' refers to million.
CIFAR-10	CelebA
other methods	38.72M	78.66M
HSIVI-SM (ours)	38.72M	78.66M
D.5 CIFAR-10 & CelebA
Figure 12 shows the uncurated samples from our proposed HSIVI-SM method with different numbers of layers on CIFAR-10  $(28 \times 28)$ , CelebA  $(64 \times 64)$  and ImageNet  $(64 \times 64)$ . We also compare the sampling time of different methods when NFE = 5 in Figure 13.
One can observe that HSIVI-SM has almost the same running time as the simplest DDIM algorithm. Finally, we report the number of parameters in the score model (or noise model) used by different methods in Table 5, which corresponds to Table 2 and Figure 13. In our implementations of HSIVI-SM, the number of parameters in the noise model equals that in the conditional layer  $q_{t}(\pmb{x}_{t}|\pmb{x}_{t + 1};\phi)$ . We find that our model with the same parameters reaches better results in Table 2.
E.1 Target distribution approximation
In this part, we set the conditional layer to be  $q_{\phi}(\pmb{x}|\pmb{z}) = \mathcal{N}(\pmb{x};\pmb{\mu}(\pmb{z};\phi^{\mu}),\mathrm{diag}\{\exp (\phi^{\sigma})\})$  and the mixing layer to be  $\mathcal{N}(\mathbf{0},\mathbf{I})$  for SIVI. Here,  $\{\phi^{\mu},\phi^{\sigma}\} = \phi$  are the variational parameters. For  $T$ -layer hierarchical semi-implicit variational distribution with  $T \geq 2$ , the variational prior  $q_{T}(\pmb{x}_{T})$  is set to be  $\mathcal{N}(\mathbf{0},\mathbf{I})$ . Each conditional layer  $q_{t}(\pmb{x}_{t}|\pmb{x}_{t + 1};\phi_{t})$  for  $t = 0,\dots ,T - 1$  is a conditional Gaussian distribution
$$
q _ {t} \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {t + 1}; \phi_ {t}\right) = \mathcal {N} \left(\boldsymbol {x} _ {t}; \boldsymbol {\mu} \left(\boldsymbol {x} _ {t + 1}; \phi_ {t} ^ {\mu}\right), \operatorname {d i a g} \left\{\exp \left(\phi_ {t} ^ {\sigma}\right) \right\}\right).
$$
Note that the  $\phi^{\sigma}$  and  $\{\phi_t^\sigma\}_{t=0}^{T-1}$  above are all vectors with the same dimension as  $\pmb{x}$ . We use sequential training for HSIVI in the two experiments in this part. The parameters  $\{\phi_t\}_{t=0}^{T-1}$  are independent across different  $t$ . If not otherwise specified, we use the Adam optimizer (Kingma & Ba, 2015) with  $\beta = (0.9, 0.99)$  for training.
E.1.1 Gaussian mixture model
For the experiment on the Gaussian mixture model, we construct 5-layer hierarchical semi-implicit variational distributions. The mean of each conditional layer  $\pmb{\mu}(\pmb{z};\phi^{\mu})$  in SIVI or  $\pmb{\mu}(\pmb{x}_{t + 1};\phi_t^\mu)$  in HSIVI has a residual form, i.e.  $\pmb{\mu}(\pmb{z};\phi^{\mu}) = \pmb{z} + \bar{\pmb{\mu}} (\pmb{z};\phi^{\mu})$  and  $\pmb {\mu}(\pmb{x}_{t + 1};\phi_t^\mu) = \pmb{x}_{t + 1} + \bar{\pmb{\mu}} (\pmb{x}_{t + 1};\phi_t^\mu)$ , for  $t = 0,\dots ,T - 1$ .  $\bar{\pmb{\mu}} (\pmb {z};\phi^{\mu})$  in SIVI and  $\{\bar{\pmb{\mu}} (\pmb{x}_{t + 1};\phi_t^\mu)\}_{t = 0}^4$  in HSIVI all have the same structures of multi-layer perceptrons (MLPs) with layer widths [2, 50, 50, 2] and ReLU activation functions. For each  $t$ ,  $\pmb{f}_t(\pmb{x}_t;\psi_t)$  in HSIVI-SM and  $\pmb{f}(\pmb {x};\psi)$  in SIVI-SM are parameterized by MLPs with layer widths [2, 128, 128, 2] and ReLU activation functions.
(a) CIFAR-10  $(32\times 32)$
(b) CelebA  $(64\times 64)$
Figure 12: Uncurated samples generated by HSIVI-SM with different numbers of layers on CIFAR-10, CelebA and ImageNet.
(c) ImageNet  $(64\times 64)$
The noise levels in the diffusion bridge are  $1 - \alpha(s_{t}) = 1 - t/5$  for  $t \in \{0,1,\dots,4\}$ . We set the learning rate of variational parameters  $\phi_t$  (or  $\phi$ ) to 0.001 and the learning rate of  $\psi_t$  (or  $\psi$ ) to 0.002 in both SIVI and HSIVI. For HSIVI-LB and HSIVI-SM, we run 80000 variational parameter updates for every conditional layer; for SIVI-LB and SIVI-SM, we run  $5 \times 80000$  variational parameter updates. For HSIVI-SM and SIVI-SM, in each nested training loop of  $f_t(x_t;\psi_t)$  (or  $f(\pmb{x};\psi)$ ), we update  $\psi_t$  (or  $\psi$ ) one time after each update of  $\phi_t$  (or  $\phi$ ). All the algorithms are trained with a batch size of 64.
E.1.2 High-dimensional conditioned diffusion
For the experiment on high-dimensional conditioned diffusion, we examine the performances of SIVI and 5-layer HSIVI. The ground truth is formed by running 100,000 independent stochastic gradient Langevin dynamics (SGLD) chains with a step size of 0.0001 and collecting the results
Figure 13: Sampling time  $(\downarrow)$  of different methods when NFE = 5 on CIFAR-10 and CelebA. Results are averaged by 100 independent runs with a batch size of 128 on a single Nvidia 2080Ti GPU.
after 10,000 iterations. For  $t = 0,\dots ,T - 2$  the mean of each conditional layer  $\pmb {\mu}(\pmb{x}_{t + 1};\phi_t^\mu)$  in HSIVI has a residual form, i.e.  $\pmb {\mu}(\pmb {x}_{t + 1};\phi_t^\mu) = \pmb {x}_{t + 1} + \bar{\pmb{\mu}} (\pmb {x}_{t + 1};\phi_t^\mu)$ . For SIVI and  $t = T - 1$  in HSIVI, we assume  $\pmb {\mu}(\pmb {z};\phi^{\mu}) = \bar{\pmb{\mu}} (\pmb {z};\phi^{\mu})$  and  $\pmb {\mu}(\pmb {x}_{t + 1};\phi_t^\mu) = \bar{\pmb{\mu}} (\pmb {x}_{t + 1};\phi_t^\mu)$ . For each  $t$ ,  $\bar{\pmb{\mu}}_t(\pmb {x};\phi_t^\mu)$  in HSIVI and  $\bar{\pmb{\mu}} (z;\phi^{\mu})$  in SIVI are MLPs with layer widths [300, 512, 512, 300] and ReLU activation functions. For each  $t$ ,  $\pmb {f}_t(\pmb {x}_t;\psi_t)$  in HSIVI-SM and  $\pmb {f}(\pmb {x};\psi)$  in SIVI-SM are MLPs with layer widths [300, 512, 512, 300] and ReLU activation functions. For both SIVI and HSIVI, we train each conditional layer for 100,000 iterations with a batch size of 128. For HSIVI-SM and SIVI-SM, in each nested training loop of  $\pmb {f}_t(\pmb {x}_t;\psi_t)$  (or  $\pmb {f}(\pmb {x};\psi)$ ), we update  $\psi_{t}$  (or  $\psi$ ) one time after each update of  $\phi_t$  (or  $\phi$ ). We set the learning rate to be 0.0001 for  $\phi_t$  (or  $\phi$ ) and 0.0005 for  $\psi_t$  (or  $\psi$ ).
E.2 Diffusion model acceleration
In this part, we use the diffusion bridge to construct the auxiliary distributions and joint training as mentioned in Section 4.2. With a pre-trained score model or noise model, we consider the generative tasks as score-based variational inference problems. Therefore, we do not use any training data to train HSIVI-SM.
For HSIVI-SM, the variational prior  $q_{T}(\pmb{x}_{T})$  is set to be  $\mathcal{N}(\mathbf{0},\mathbf{I})$ . To avoid the large memory consumption, we use the joint training method where the parameters of the conditional layers  $q_{t}(\pmb{x}_{t}| \pmb{x}_{t + 1};\phi)$  and  $f_{t}(\pmb{x}_{t};\psi)$  are the same across different  $t$ . The  $t$ -th conditional layer is a conditional Gaussian distribution
$$
q _ {t} \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {t + 1}; \phi\right) = \mathcal {N} \left(\boldsymbol {x} _ {t}; \boldsymbol {\mu} _ {t} \left(\boldsymbol {x} _ {t + 1}; \phi^ {\mu}\right), \operatorname {d i a g} \left(\sigma_ {t} ^ {2} \exp \left(\phi^ {\sigma}\right)\right)\right),
$$
where  $\{\phi^{\mu},\phi^{\sigma}\} = \phi$  are the variational parameters,  $\phi^{\sigma}$  is a vector with the same dimension as  $\pmb{x}$ , and  $\sigma_{t}$  is a fixed scalar value. We use the generalized inference process in DDIM (Song et al., 2020a) with the noise level  $\eta >0$  to initialize  $\pmb {\mu}_t(\pmb {x}_{t + 1};\phi^{\mu})$  and determine the value of  $\sigma_t$  for each  $t$ . If not otherwise specified, we use the Adam optimizer (Kingma & Ba, 2015) with  $\beta = (0.9,0.99)$  for training.
For our implementation, we referenced the training code of diffusion model acceleration for our models in the repository from (Dockhorn et al., 2022).
E.2.1 Toy examples of diffusion model acceleration
For pre-training the score model  $S^{*}(\pmb{x}, s)$ , we consider quadratic noise levels  $1 - \alpha(s) = s^2$  for  $s \in [0,1]$ . We then train  $S^{*}(\pmb{x}, s)$  on 1000 fixed noise levels  $\{1 - \alpha(i/1000)\}_{i=1}^{1000}$  by optimizing the DDPM loss in equation (13) for 200,000 iterations with a learning rate of 0.0003 and a batch size of 100. For constructing the diffusion bridge, we choose  $T$  discrete time steps  $\{s_t\}_{t=0}^{T-1}$  so that  $1 - \alpha(s_t) = [0.01 + (\sqrt{0.8} - 0.1)t/T]^2$  for  $t = 0,1,\dots,T-1$ .
Model architecture The model architecture of  $S^{*}(\pmb{x}, s)$  is
$$
\boldsymbol {S} ^ {*} (\boldsymbol {x}, s) = \mathrm {M L P} ^ {\mathrm {d e c}} \left(\mathrm {M L P} ^ {\mathrm {e m b x}} (\boldsymbol {x}) + \mathrm {M L P} ^ {\mathrm {e m b t}} (1 - \alpha (s))\right),
$$
where  $\mathrm{MLP}^{\mathrm{dec}}$  is a decoder implemented as MLPs with layer widths [128, 128, 128, 2],  $\mathrm{MLP}^{\mathrm{embx}}$  is a data embedding block implemented as MLPs with layer widths [2, 128], and  $\mathrm{MLP}^{\mathrm{embt}}$  is a time
embedding block implemented as MLPs with layer widths [256, 128, 128]. We use the sinusoidal positional embedding (Vaswani et al., 2017) of  $1 - \alpha(s)$  as the input of  $\mathrm{MLP}^{\mathrm{embt}}$ . All these three MLPs use GELU as activation functions. We use the generalized inference process with noise level  $\eta = 1.0$  to initialize the conditional layers. The architecture of  $f_{t}(\boldsymbol{x}_{t}; \psi)$  is the same as that of  $\boldsymbol{S}^{*}(\boldsymbol{x}, s)$ . We initialize  $f_{t}(\boldsymbol{x}_{t}; \psi)$  with  $\boldsymbol{S}_{t}^{*}(\boldsymbol{x}_{t}) := \boldsymbol{S}^{*}(\boldsymbol{x}_{t}, s_{t})$ .
Training setting The learning rate is set to be 0.0002 for  $q_{t}(\pmb{x}_{t}|\pmb{x}_{t + 1};\phi)$  and 0.0005 for  $\pmb{f}_{t}(\pmb{x}_{t};\psi)$  on Swissroll, Circles, and Moons for both  $T = 5,10$ . On Checkerboard, the learning rate is set to be 0.00001 (0.00002) for  $q_{t}(\cdot |\pmb{x}_{t + 1};\phi)$  and 0.00005 (0.0001) for  $\pmb{f}_{t}(\pmb{x}_{t};\psi)$  when  $T = 5$  ( $T = 10$ ). We train HSIVI-SM for 25,000 iterations with a batch size of 64 in all cases. In each nested training loop of  $\pmb{f}_{t}(\pmb{x}_{t};\psi)$ , we update  $\psi$  3 times after each update of  $\phi$ .
E.2.2 MNIST
For the experiment on MNIST, we use the pre-trained noise model  $\epsilon^{*}(\pmb{x}, s)$  and train HSIVI-SM with  $\epsilon$ -training introduced in Section C.4. The following construction of noise schedule comes from Song et al. (2020a). Let  $\beta_{j} = \beta_{\min} + \frac{\beta_{\max} - \beta_{\min}}{999} j$  for  $j = 0, \dots, 999$ , where  $\beta_{\min} = 0.0001$ ,  $\beta_{\max} = 0.02$ . We pre-train the noise model on the 1000 fixed noise levels  $1 - \alpha(s) \coloneqq \prod_{j=0}^{s} \beta_{j}$  for  $s = 0, \dots, 999$  by equation (14). The noise model is trained for 100,000 iterations with a learning rate of 0.0001 and a batch size of 64. We then choose  $T$  discrete time steps  $s_{t} = \lfloor 800 \cdot \frac{t^{2}}{T^{2}} \rfloor$  for  $t = 0, \dots, T-1$  to construct the  $T$ -layer diffusion bridge.
Model architecture The pre-trained noise model  $\epsilon^{*}(\pmb{x}, s)$  follows the UNet structure employed by Ho et al. (2020) where the number of input channels and output channels is reduced to one. Additionally, we pad the image size to  $32 \times 32$  to fit  $\epsilon^{*}(\pmb{x}, s)$ . We use the generalized inference process with noise level  $\eta = 0.2$  to initialize the conditional layers. The architecture of  $\pmb{f}_t(\pmb{x}_t; \psi)$  is the same as that of  $\epsilon^{*}(\pmb{x}, s)$ . We initialize  $\pmb{f}_t(\pmb{x}_t; \psi)$  with  $-\epsilon^{*}(\pmb{x}_t, s_t) / \sqrt{1 - \alpha(s_t)}$ .
Training setting For both  $T = 5,10$ , the learning rate is set to be  $1.6\times 10^{-5}$  for  $\phi$  and  $6.4\times 10^{-5}$  for  $\psi$ . We train HSIVI-SM for 10,000 iterations with a batch size of 64 in all cases. In each nested training loop of  $f_{t}(x_{t};\psi)$ , we update  $\psi$  20 times after each update of  $\phi$ .
E.2.3 CIFAR-10, CelebA & ImageNet
For experiments on CIFAR-10 and CelebA, we use the pre-trained noise model  $\epsilon^{*}(\pmb{x},s)$  and train HSIVI-SM with  $\epsilon$ -training introduced in Section C.4. We use the same noise schedule as in the experiment on MNIST. Let  $\beta_{j} = \beta_{\min} + \frac{\beta_{\max} - \beta_{\min}}{999} j$  for  $j = 0,\ldots,999$ , where  $\beta_{\min} = 0.0001, \beta_{\max} = 0.02$ . We take the pretrained noise model for CIFAR10 and ImageNet separately from https://github.com/tqch/ddpm-torch/releases/download/checkpoints/cifar10_2040.pt and https://openaipublic.blob.core.windows.net/diffusion/march-2021/imagenet64_uncond_100M_1500K.pt. On CelebA, We pre-train the noise model on the 1000 fixed noise levels  $1 - \alpha(s) := \prod_{j=0}^{s} \beta_{j}$  for  $s = 0,\ldots,999$  by optimizing equation (14). The noise model is trained for 600 epochs, with a learning rate of 0.00002 and batch size of 128. We then choose  $T$  discrete time steps  $s_{t} = \left\lfloor 800 \cdot \frac{t^{2}}{T^{2}} \right\rfloor$  for  $t = 0,\ldots,T-1$  to construct the  $T$ -layer diffusion bridge.
Model architecture On CIFAR-10 and CelebA, the structure of  $\epsilon^{*}(\pmb{x}, s)$  is exactly the UNet<sup>3</sup> employed in Ho et al. (2020) without modification; on ImageNet, the structure of  $\epsilon^{*}(\pmb{x}, s)$  is exactly the UNet in Nichol & Dhariwal (2021). We use the generalized inference process with noise level  $\eta = 0.2$  to initialize the conditional layers. The architecture of  $f_{t}(\pmb{x}_{t}; \psi)$  is the same as that of  $\epsilon^{*}(\pmb{x}, s)$ . We initialize  $f_{t}(\pmb{x}_{t}; \psi)$  with  $-\epsilon^{*}(\pmb{x}_{t}, s_{t}) / \sqrt{1 - \alpha(s_{t})}$ .
Training setting The number of layers, which is also the number of function evaluations (NFE), is set to be  $T = 5, 10, 15$  in our test cases. On CIFAR-10, the learning rate is set to be  $1.6 \times 10^{-5}$  for  $q_{t}(\cdot | \boldsymbol{x}_{t+1}; \phi)$  and  $8 \times 10^{-5}$  for  $\boldsymbol{f}_{t}(\boldsymbol{x}_{t}; \psi)$ ; on CelebA, the learning rate is set to be  $1.2 \times 10^{-6}$  for  $q_{t}(\cdot | \boldsymbol{x}_{t+1}; \phi)$  and  $6 \times 10^{-6}$  for  $\boldsymbol{f}_{t}(\boldsymbol{x}_{t}; \psi)$ ; on ImageNet, the learning rate is set to be  $1 \times 10^{-5}$
<sup>3</sup>We use the Pytorch implementation of UNet structure in https://github.com/tqch/ddpm-torch.
for  $q_{t}(\cdot | \boldsymbol{x}_{t+1}; \phi)$  and  $5 \times 10^{-5}$  for  $f_{t}(\boldsymbol{x}_{t}; \psi)$ . We trained HSIVI-SM for 10,000 iterations with a batch size of 128. During each nested training loop of  $f_{t}(\boldsymbol{x}_{t}; \psi)$ , we update  $\psi$  20 times after each update of  $\phi$ , since we find  $f_{t}(\boldsymbol{x}_{t}; \psi)$  needs more training empirically to provide reliable guidance. For  $T = 10, 15$ , we use the above training settings; for  $T = 5$ , we find that further fine-tuning on the well-trained 15-layer HSIVI-SM for 1,000 iterations yields better results, and we utilize this strategy to optimize the 5-layer HSIVI-SM with a  $0.1 \times$  smaller learning rate. Experiments need about 1.5 days on CIFAR-10, need about 3 days on CelebA and 4 days on ImageNet using 8 Nvidia 2080 Ti GPUs. During the training, we find that HSIVI-SM converges in the first  $30\%$  iterations on CIFAR-10 and converges in the first  $50\%$  iterations on CelebA.
For the application of accelerating the sampling process of diffusion models, our HSIVI-SM training involves three models: the score model (or noise model), the conditional layers  $q_{t}(\pmb{x}_{t}|\pmb{x}_{t + 1};\phi)$ , and  $\pmb{f}_t(\pmb{x}_t;\psi)$ . As a result, HSIVI-SM requires higher memory consumption due to the involvement of multiple models. Additionally, since our HSIVI algorithm approximates the target distribution using the score function, it necessitates a pre-trained score model (or noise model) with high accuracy and additional training steps. Finally, we recognize that the alternative method HSIVI-LB remains unexplored for accelerating the diffusion model, and we defer this aspect to future research.