Title: Hierarchical Semi-Implicit Variational Inference with Application to Diffusion Model Acceleration
Venue: NeurIPS 2023
Accepted: yes
Keywords: Hierarchical semi-implicit variational inference

Score based training

Diffusion model
Abstract: Semi-implicit variational inference (SIVI) has been introduced to expand the analytical variational families by defining expressive semi-implicit distributions in a hierarchical manner. However, the single-layer architecture commonly used in current SIVI methods can be insufficient when the target posterior has complicated structures. In this paper, we propose hierarchical semi-implicit variational inference, called HSIVI, which generalizes SIVI to allow more expressive multi-layer construction of semi-implicit distributions. By introducing auxiliary distributions that interpolate between a simple base distribution and the target distribution, the conditional layers can be trained by progressively matching these auxiliary distributions one layer after another. Moreover, given pre-trained score networks, HSIVI can be used to accelerate the sampling process of diffusion models with the score matching objective. We show that HSIVI significantly enhances the expressiveness of SIVI on several Bayesian inference problems with complicated target distributions. When used for diffusion model acceleration, we show that HSIVI can produce high quality samples comparable to or better than the existing fast diffusion model based samplers with a small number of function evaluations on various datasets.

Hierarchical Semi-Implicit Variational Inference with Application to Diffusion Model Acceleration

Abstract

Semi-implicit variational inference (SIVI) has been introduced to expand the analytical variational families by defining expressive semi-implicit distributions in a hierarchical manner. However, the single-layer architecture commonly used in current SIVI methods can be insufficient when the target posterior has complicated structures. In this paper, we propose hierarchical semi-implicit variational inference, called HSIVI, which generalizes SIVI to allow more expressive multi-layer construction of semi-implicit distributions. By introducing auxiliary distributions that interpolate between a simple base distribution and the target distribution, the conditional layers can be trained by progressively matching these auxiliary distributions one layer after another. Moreover, given pre-trained score networks, HSIVI can be used to accelerate the sampling process of diffusion models with the score matching objective. We show that HSIVI significantly enhances the expressiveness of SIVI on several Bayesian inference problems with complicated target distributions. When used for diffusion model acceleration, we show that HSIVI can produce high quality samples comparable to or better than the existing fast diffusion model based samplers with a small number of function evaluations on various datasets.

1 Introduction

Variational inference (VI) is an approximate Bayesian inference method that is gaining in popularity, where one tries to find an approximation to the target posterior distribution using an optimization approach (Jordan et al., 1999; Wainwright & Jordan, 2008; Blei et al., 2016). To do that, it first posits a family of variational distributions and then seeks the closest member from this family that minimizes some statistical distance to the target posterior, usually the Kullback-Leibler (KL) divergence. As the posterior is not analytically available, an equivalent formulation is often adopted in practice where one maximizes the evidence lower bound (ELBO) instead (Jordan et al., 1999).

One classical VI method is mean-field VI, which assumes a factorizable structure of the variational distributions over the parameters or latent variables (Bishop & Tipping, 2000). This often leads to closed-form coordinate-ascent update rules when certain conditional conjugacy conditions are satisfied. In practice, the conditional conjugacy may not hold and the true posterior could be much

more complicated than what a factorized variational distribution can accurately approximate. In recent years, several attempts have been made in VI that alleviate these constraints by designing more flexible variational families (Jaakkola & Jordan, 1998; Saul & Jordan, 1996; Giordano et al., 2015; Tran et al., 2015; Rezende & Mohamed, 2015; Dinh et al., 2017; Kingma et al., 2016; Papamakarios et al., 2019), together with generic training algorithms via Monte Carlo gradient estimators (Nott et al., 2012; Paisley et al., 2012; Ranganath et al., 2014; Rezende et al., 2014; Kingma & Welling, 2014). While successful, these approaches all assume tractable densities of variational distributions. To further expand the capacity of variational families, one approach is to incorporate the implicit models that have intractable densities but are easy to sample from (Huszár, 2017; Tran et al., 2017; Mescheder et al., 2017; Shi et al., 2018a,b; Song et al., 2019). However, as the densities are intractable for implicit models, one often resorts to density ratio estimation for ELBO evaluation during training, which is known to be difficult in high dimensional settings (Sugiyama et al., 2012). To avoid density ratio estimation, semi-implicit variational inference (SIVI) has been proposed where the variational distributions are formed through a semi-implicit hierarchical construction, and various training criteria have been employed (Yin & Zhou, 2018; Moens et al., 2021; Titsias & Ruiz, 2019; Yu & Zhang, 2023).

While striking a good balance between approximation flexibility and training efficiency, current SIVI methods often use a single conditional layer which can be insufficient when the target posterior possesses complicated structures (e.g., multimodality, see an example in Section 5.1). To enhance the expressiveness of single-layer models, an intuitive but effective approach is to extend them to multi-layer hierarchical models (Vahdat & Kautz, 2020; Ranganath et al., 2016; Sobolev & Vetrov, 2019). In this paper, we propose hierarchical semi-implicit variational inference (HSIVI), which is a generalization of SIVI that allows multiple conditional layers. Instead of training the hierarchical semi-implicit model end to end, we introduce auxiliary distributions that interpolate between a simple base distribution and the target distribution to guide the intermediate semi-implicit distributions toward the target distribution. The conditional layers are then trained sequentially to match these auxiliary bridging distributions given the fitted semi-implicit distributions from the previous layers (Figure 1), using different criteria from before. This way, HSIVI allows progressive learning of the target distribution that significantly reduces the burden of each conditional layer. Moreover, HSIVI with the score matching objective can also be used to accelerate the sampling process of diffusion models where the pre-trained score networks corresponding to different noise levels provide a natural sequence of bridging distributions. In experiments, we demonstrate the effectiveness of HSIVI on both Bayesian inference tasks with complicated target distributions and diffusion model acceleration.

2 Background on semi-implicit variational inference

The semi-implicit variational family (Yin & Zhou, 2018; Titsias & Ruiz, 2019) is defined as

where ϕ are the variational parameters, q ϕ ( x | z ) is called the conditional layer, and q ( z ) is called the mixing layer. This variational family is said to be semi-implicit as q ϕ ( x | z ) is required to be explicit and q ( z ) is often implicit. The semi-implicit variational family is capable of capturing more complicated dependencies between variables (Yin & Zhou, 2018; Titsias & Ruiz, 2019; Yu & Zhang, 2023) than explicit variational families without the hierarchical structure. Given the observed data D , the classical VI methods often use the evidence lower bound (ELBO) for training, which is defined as ELBO := E q ϕ ( x ) [log p ( D, x ) -log q ϕ ( x )] . However, as q ϕ ( x ) is no longer tractable in SIVI, alternative training objectives have been introduced.

ELBO related objectives Yin & Zhou (2018) considered a sequence of lower bounds of the ELBO

It is an asymptotically exact surrogate in the sense that lim K →∞ L SIVI-LB = ELBO. Titsias & Ruiz (2019) proposed unbiased implicit variational inference (UIVI) which uses samples from the inverse conditional distribution q ϕ ( z | x ) (from an MCMC run, e.g. Hamiltonian Monte Carlo (Neal, 2011)) to provide an unbiased gradient estimator of the exact ELBO. See more details of UIVI in Appendix B.

q

0

(

x

0

|

x

1

)

q

1

(

x

1

|

x

2

)

q

2

(

x

2

|

x

3

)

LSIVI

-

f

(

p

0

(

x

0

) ∥

q

0

(

x

0

) )

LSIVI

-

f

(

p

1

(

x

1

) ∥

q

1

(

x

1

) )

LSIVI

-

f

(

p

2

(

x

2

) ∥

q

2

(

x

2

) )

q

3

(

x

3

|

x

4

)

q

4

(

x

4

)

LSIVI

-

f

(

p

3

(

x

3

) ∥

q

3

(

x

3

) )

q

i

(

x

i

|

x

i

+1

)

Conditional

layer

Auxiliary

distribution

Base

distribution

p

0

(

x

0

)

p

1

(

x

1

)

p

2

(

x

2

)

p

3

(

x

3

)

p

i

(

x

i

)

Score matching objective Besides the ELBO, score based distance measures have also been used for variational inference where the score function S ( x ) := ∇ x log p ( x | D ) = ∇ x log p ( D, x ) is assumed to be tractable (Liu et al., 2016; Zhang et al., 2018; Hu et al., 2018). Yu & Zhang (2023) considered the following Fisher divergence between the target distribution and the semi-implicit variational distribution

By reformulating D Fisher as the maximum of the following optimization problem

and using a similar trick as in denoising score matching (Vincent, 2011; Song & Ermon, 2019), one can transform the minimization of D Fisher into the following minimax problem which is tractable

In practice, f ψ ( x ) is parametrized using neural networks. The above minimax optimization problem can be efficiently solved by optimizing ψ and ϕ alternately.

3 Hierarchical semi-implicit variational inference

The semi-implicit variational family q ϕ ( x ) in equation (1) is indeed a single-layer model in the sense that it contains only one conditional layer. Our main idea is to expand this single-layer semi-implicit variational family into its multi-layer variants and introduce a sequence of auxiliary distributions to guide the semi-implicit distributions toward the target distribution. This leads to a new SIVI method which we call hierarchical semi-implicit variational inference (HSIVI). We start with the following definition which is motivated by equation (1).

Definition 1 (Hierarchical Semi-Implicit Distribution) . Let x T ∼ q T ( x T ) for some T ∈ N ∗ , where q T ( x T ) is called the variational prior. Let q t ( x t | x t +1 ; ϕ t ) be the t -th conditional layer for 0 ≤ t ≤ T -1 . Denote { ϕ k } T -1 k = t by ϕ ≥ t . The t -th layer hierarchical semi-implicit distribution q t ( x t ; ϕ ≥ t ) is defined recursively from T -1 to 0 by

where q T ( x T ; ϕ ≥ T ) := q T ( x T ) . Here, the t -th conditional layer q t ( x t | x t +1 ; ϕ t ) is required to be explicit and reparametrizable with a tractable score function ∇ x t log q t ( x t | x t +1 ; ϕ t ) .

Compared to the single-layer semi-implicit variational family (1), the family of hierarchical semiimplicit distributions provides a principled way to construct more expressive mixing layers using multilayer architectures. Also, unlike the hierarchical variational models (Ranganath et al., 2016) which require an extra reverse model and explicit variational prior, hierarchical semi-implicit distributions inherit the advantage of SIVI that allows q t ( x t ; ϕ ≥ t ) to be implicit, and as shown next, they do not require a reverse model and can be progressively trained using the simple algorithms of SIVI for each conditional layer, from t = T -1 to t = 0 .

3.1 Progressive approximation with the auxiliary bridge

In this section, we introduce a bridging technique for progressively approximating the target distribution p ( x ) using hierarchical semi-implicit distributions. Rather than approximating p ( x ) with q 0 ( x ; ϕ ≥ 0 ) directly, we construct a sequence of intermediate auxiliary distributions { p t ( x ) } T -1 t =0 as a bridge between the target distribution p 0 ( x ) := p ( x ) and an easy-to-approximate distribution p T -1 ( x ) , to amortize the difficulty of one-pass fitting. A typical example of an auxiliary bridge is the geometric interpolation as described below.

Example 1 (Geometric Interpolation) . Let S ( x ) := ∇ log p ( x ) be the score function of target distribution p ( x ) and S base ( x ) := ∇ log p base ( x ) be the score function of a base distribution p base ( x ) . In geometric interpolation (Neal, 2001; Bernton et al., 2019), each auxiliary distribution p t ( x ) for 0 ≤ t ≤ T -1 has the following probability density function (pdf) and score function

t t =0 0

Intuitively, we expect the distance between two neighboring distributions p t ( x ) and p t +1 ( x ) to be not too large so that it would be easy to construct a conditional distribution q t ( x t | x t +1 ) such that p t ( x t ) ≈ ∫ q t ( x t | x t +1 ) p t +1 ( x t +1 )d x t +1 . Note that the auxiliary bridge { p t ( x ) } T -1 t =0 does not necessarily need to have analytical pdfs (up to a constant). In fact, it suffices if they have tractable score functions { S t ( x ) } T -1 t =0 which lead to another type of auxiliary bridge (Example 2 in Section 4).

3.2 Sequential training of HSIVI

Given the auxiliary distributions { p t ( x t ) } T -1 t =0 , a natural approach is to progressively train the hierarchical semi-implicit distribution q t ( x t ; ϕ ≥ t ) to match p t ( x t ) from t = T -1 to t = 0 . Let the parameters ϕ t in the t -th conditional layer be independent across different t s. We first train q T -1 ( x T -1 ; ϕ T -1 ) to match p T -1 ( x T -1 ) by optimizing ϕ T -1 w.r.t. the single-layer SIVI objective L SIVIf ( p T -1 ( x T -1 ) ∥ q T -1 ( x T -1 ; ϕ T -1 )) . For t = T -2 , . . . , 0 , given the trained semi-implicit distribution q t +1 ( x t +1 ; ϕ ≥ t +1 ) , we can fix it as the mixing layer and train the t -th conditional layer q t ( x t | x t +1 ; ϕ t ) by optimizing ϕ t w.r.t. the single-layer SIVI objective L SIVIf ( p t ( x t ) ∥ q t ( x t ; ϕ ≥ t )) as well. Note this is fine as the mixing layer can be implicit in SIVI. Here, f is some distance criterion, e.g. L SIVI-LB in equation (2) or L SIVI-SM in equation (4). In this article, we mainly focus on L SIVI-LB and L SIVI-SM, while other distance criteria can also be applied. We summarize this sequential training procedure in Algorithm 1.

Score based training In addition to the common assumption that p t ( x ) is known up to a constant, it is worth noting that L SIVI-LB is also applicable when only the score functions { S t ( x ) } T -1 t =0 are available which is important for the diffusion bridge construction of auxiliary distributions in Example 2. Concretely, assume q t ( x t | x t +1 ; ϕ t ) is induced by a parametrized transform x t = h t ( x t +1 , ϵ ; ϕ t )

where ϵ ∼ p ϵ ( ϵ ) is a random noise. The only term in L SIVI-LB ( p t ( x t ) ∥ q t ( x t ; ϕ ≥ t )) containing p t ( x t ) is E q t ( x t ; ϕ ≥ t ) log p t ( x t ) (see equation (2)) whose gradient takes the form

In the training of HSIVI-SM, each term L SIVI-SM ( p t ( x t ) ∥ q t ( x t ; ϕ ≥ t )) involves a nested optimization of f t ( x t ; ψ t ) . When the score functions are computationally expensive, we find that an alternative parametrization f t ( x t ; ψ t ) := S t ( x t ) -g t ( x t ; ψ t ) is useful to avoid the time-consuming evaluation of S t ( x t ) when optimizing ψ t in equation (4). The reason for this lies in Proposition 1. See Appendix C.2 for the proof of Proposition 1.

Proposition 1. Let q t ( x t , x t +1 ; ϕ ≥ t ) = q t ( x t | x t +1 ; ϕ t ) q t +1 ( x t +1 ; ϕ ≥ t +1 ) . The minimax optimization of L SIVI-SM ( p t ( x t ) ∥ q t ( x t ; ϕ ≥ t )) is equivalent to

Marginal approximation v.s. joint approximation Previous works (Bernton et al., 2019; Bao et al., 2022) often construct a joint distribution p ( x 0: T ) and minimize KL( p ( x 0: T -1 ) ∥ q ( x 0: T -1 )) where q ( x 0: T -1 ) is a variational distribution. In HSIVI, we directly approximate p t ( x t ) using the semi-implicit variational distributions. When p ( x 0: T -1 ) is complex and T is small, the variational distribution q ( x 0: T -1 ) may be insufficient to fully capture the joint distribution p ( x 0: T -1 ) . For example, the optimal fit of the joint distribution for diffusion models established by AnalyticDPM (Bao et al., 2022) does not guarantee that the marginal distributions would be approximated well (see Table 2 for comparison).

4 Application to diffusion model acceleration

4.1 Review of diffusion models

Recently, diffusion models have shown great success on many generative modeling benchmarks, including image generation (Ho et al., 2020; Song et al., 2020a,b), graph generation (Niu et al., 2020), and text generation (Austin et al., 2021). Diffusion models work by adding noise to the training data in the forward process and then removing the noise to recover the data in the backward process, which can be integrated into a general stochastic differential equation (SDE) framework. The forward process { u s } s ∈ [0 ,L ] is usually described by

where p 0 ( · ) is the data distribution, w s is a standard Brownian motion, and f ( u s , s ) and g ( s ) are the drift and diffusion coefficients respectively. To generate samples from the data distribution, one can run the following backward process

where p s ( · ) is the pdf of u s and ¯ w s is a standard Brownian motion when time flows from L to 0 . As the score function ∇ u s log p s ( u s ) is intractable, we need to estimate it by denoising score matching (Vincent, 2011; Song et al., 2020b). See more details of diffusion models and the training objectives in Appendix A.

4.2 Diffusion model acceleration via HSIVI

While diffusion models prove effective for generative modeling, it often takes a large number of discretization steps in the backward process (9) to produce high quality samples, which caps their potential for real time applications. Note that the forward process (8) naturally provides another type of auxiliary bridge, which combined with HSIVI, can be used to accelerate the sampling process of diffusion models.

Example 2 (Diffusion Bridge) . Consider the forward process { u s } s ∈ [0 ,L ] with L > 0 (defined in equation (8)) in diffusion models. We choose T discrete time steps 0 ≈ s 0 < · · · < s T -1 ≤ L and

let x t := u s t with probability density function p t ( · ) . Assume each auxiliary distributions p t ( · ) for 0 ≤ t ≤ T -1 admits a score function as

where S ∗ ( x , s ) is a pre-trained score model with the denoising score matching loss (equation (13) in Appendix A). Let us denote S ∗ ( x , s t ) by S ∗ t ( x ) for short. With sufficient samples from the data distribution p 0 ( x ) and model capacity, the approximation S ∗ t ( x ) can be reasonably accurate for almost all x and t (Song et al., 2020b).

As the pre-trained score model provides a diffusion bridge from the simple distribution p T -1 (e.g., standard Gaussian) to the data distribution, we can train the hierarchical semi-implicit distributions to approximate the diffusion bridge within the HSIVI framework. Given the expressiveness of hierarchical semi-implicit distributions, we may expect an accurate approximation of the data distribution with a small number T and hence acceleration can be achieved.

However, the memory usage during the sequential training process for HSIVI might be large because of the necessity for independent parameters. Therefore, we may employ a parameter sharing scheme which is commonly assumed in diffusion models (Song & Ermon, 2019; Ho et al., 2020) such that different conditional layers share the same parameters ϕ . Note that sequential training is not suitable in this setting. Therefore, we propose a joint training procedure that minimizes a weighted sum of the SIVI objectives

where β ( t ) : { 0 , . . . , T -1 } → R + is a positive weighting function and f is some distance criterion. See Algorithm 2 in Appendix C.3 for more details of joint training.

More specifically, in this work, we mainly focus on building the diffusion bridge with variance preserving SDE (VP-SDE) (Song et al., 2020b) such that u s | u 0 ∼ N ( √ α ( s ) u 0 , (1 -α ( s )) I ) with a decreasing function α ( s ) of s . We use L HSIVI-SM in equation (10) for training and set the weighting function β ( t ) = 1 -α ( s t ) as recommended in Song et al. (2020b), which tends to train layers that are far from t = 0 first during the training, resembling the sequential training. Another popular formulation of diffusion models is to fit a noise model ϵ ∗ ( x , s ) that predicts the noise added to a noisy sample x at time s (Ho et al., 2020). HSIVI-SM also generalizes to the case where a pre-trained noise model is available. The pre-trained noise model forms a (generalized) diffusion bridge by letting ϵ ∗ t ( x ) = ϵ ∗ ( x , s t ) , and we call the corresponding training method ' ϵ -training'. We provide a reparametrized objective function ˜ L HSIVI-SM for ϵ -training in Appendix C.4.

Several efforts have been made to accelerate the sampling process of diffusion models, including faster numerical ordinary differential equation (ODE) solvers (Song et al., 2020a; Zhang & Chen, 2022; Lu et al., 2022) and distillation techniques (Luhman & Luhman, 2021; Salimans & Ho, 2022; Zheng et al., 2022). Our approach is different from these previous efforts in that we accelerate the stochastic diffusion model directly (hence would provide more diverse samples (Figure 6)) and do not require sampling datasets from the diffusion models prior to distillation which is computationally expensive. From a Bayesian perspective, HSIVI is related to Song & Ermon (2019), where the authors used the annealed Langevin dynamics guided by a pre-trained score model to sample from the data distribution. By solving this problem using a variational inference approach, HSIVI enjoys faster sampling speed and scales better to high-dimensional data.

5 Experiments

In this section, we first compare HSIVI to its single-layer counterpart, SIVI, on two inference tasks. We use the sequential training method where each conditional layer in the hierarchical semiimplicit variational distributions has independent parameters. We then apply HSIVI-SM to diffusion model acceleration on various datasets. As the memory consumption for generative models is large, we use the joint training method where the conditional layers in hierarchical semi-implicit distributions have shared parameters across different t s. For all experiments, each conditional layer is modeled as a Gaussian distribution with parametrized mean and variance. More details of the model architectures and hyper-parameters are included in Appendix E. The code is available at https://github.com/longinYu/HSIVI .

Ground Truth

SIVI-LB

HSIVI-LB (T=5)

SIVI-SM

HSIVI-SM (T=5)

15

15

15

15

15

10

10

10

10

10

5

5

5

5

5

0

0

0

0

0

5

-5

-5

-5

-5

10

-10

-10

10

-10

15

15

-10

-5

0

5

10

15

15

15

-10

-5

0

5

10

15

15

15

-10

-5

0

5

10

15

-15

15

-10

-5

0

5

10

15

-15

15

-10

-5

0

5

10

15

0.0

0.5

1.0

1.5

2.0

2.5

3.0

s

1.4

1.2

1.0

0.8

0.6

0.4

0.2

0.0

0.2

xs

Ground Truth

95% Confidence Interval

Observations

Sample Mean

0.0

0.5

1.0

1.5

2.0

2.5

3.0

s

1.4

1.2

1.0

0.8

0.6

0.4

0.2

0.0

0.2

xs

SIVI-SM

95% Confidence Interval

Observations

Sample Mean

0.0

0.5

1.0

1.5

2.0

2.5

3.0

s

1.4

1.2

1.0

0.8

0.6

0.4

0.2

0.0

0.2

xs

HSIVI-SM (T=5)

95% Confidence Interval

Observations

Sample Mean

5.1 Target distribution approximation

Gaussian mixture model We first evaluate HSIVI and SIVI on a two-dimensional Gaussian mixture model. The target distribution p ( x ) takes the form p ( x ) = ∑ 8 i =1 1 / 8 · N ( x ; µ i , σ 2 I ) where µ i = [10 cos( iπ 4 ) , 10 sin( iπ 4 )] T , σ = 1 . For HSIVI, we construct an auxiliary bridge of T = 5 with geometric interpolation in Example 1, where p base ( x ) = N ( x ; 0 , I ) and λ t = 1 -t/ 5 . The results are presented in Figure 2. Note that the modes in this Gaussian mixture model are far apart from each other, and both SIVI-LB and SIVI-SM are trapped in local modes. In contrast, both HSIVI-LB and HSIVI-SM discover all modes and provide an accurate approximation of the target distribution with HSIVI-SM being better for recovering the right scale of variance.

High-dimensional conditioned diffusion The second example is a high-dimensional Bayesian inference problem arising from the following Langevin SDE

where x 0 = 0 and w s is a one-dimensional standard Brownian motion. This system describes the motion of a particle with negligible mass trapped in an energy potential with thermal fluctuations represented by the Brownian forcing (Cui et al., 2016). Using an Euler-Maruyama scheme with step size ∆ s = 0 . 01 on a time interval [0 , 3] , we discretize the SDE (11) into x = ( x d 1 , . . . , x d 300 ) where d i = 0 . 01 i , which gives the prior distribution p prior ( x ) of the 300-dimensional variable x . The noisy observations y is obtained by y = x + ξ , where ξ ∼ N ( 0 , σ 2 I ) with σ = 0 . 1 . Our goal is to infer the posterior distribution of the latent states p ( x | y ) ∝ p prior ( x ) p ( y | x ) . 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 after 10,000 iterations.

For HSIVI, we form the auxiliary bridge using geometric interpolation with p base ( x ) = N ( x ; y , σ 2 I ) and λ t = 1 -t/ ( T -1) for t = 0 , . . . , T -1 . Figure 3 shows the estimated posteriors obtained by different methods. We see that SIVI-SM severely underestimates the variance. With T = 5 layers, HSIVI-SM fits the variance better and hence provides more accurate posterior estimates. For both HSIVI-SM and HSIVI-LB, the estimated covariance matrix becomes more accurate as T increases (Table 4 in Appendix D.2), demonstrating the effectiveness of hierarchical models for fitting complicated distributions.

X9

X7

X5

X9

X7

X5

米

。

O

-

-2

0

-2

0

-2

X9

X7

X5

Xo

xg

X7

X5

2.0

2.0

2.0

2.0

3

0.5

1.5

1.0

1.5

1.0

0.5

0.5

1.5

1.0

1.5

1.0

0.5

0°0

0.0

0.0

0.0

S°0-

0.5

0.5

0.5

1.0

1.0

1.0

1.0

1.5

-1.5

1.5

1.5

-2

2

0

0

-2

Z-

2

0

-2

D

-2

5.2 Diffusion model acceleration

2D toy examples In this toy model example, we test four synthetic 2D datasets: Checkerboard, Circles, Moons, and Swissroll (Pedregosa et al., 2011). We first pre-train the score model S ∗ ( x , s ) for s ∈ [0 , 1] with quadratic noise schedule 1 -α ( s ) = s 2 . For constructing the T -layer diffusion bridge, we select { s t } T -1 t =0 so that 1 -α ( s t ) = [0 . 01 + ( √ 0 . 8 -0 . 01) t/T ] 2 . Figure 4 shows the sample trajectories ( x 9 , x 7 , x 5 and x 0 ) progressively generated from 10-layer HSIVI-SM. We see clearly how the semi-implicit distributions are guided towards the target distribution and all modes are discovered. We also report the Jensen-Shannon (JS) divergence between the target distributions and the estimated distributions in Table 1. We see that HSIVI-SM significantly improves upon DDIM and DDPM in both cases with 5 and 10 steps. Also, 10-layer HSIVI-SM is comparable to DDPM with 1000 full steps. See Figure 10 in Appendix D.3 for visualization of samples from different methods.

MNIST On MNIST, we use the noise model ϵ ∗ ( x , s ) instead of the score model and use ϵ -training to train HSIVI-SM. The structure of ϵ ∗ ( x , s ) follows the UNet in Ho et al. (2020) by reducing the number of input and output channels to one. With the same noise schedule employed in Song et al. (2020a), we first pre-train the noise model ϵ ∗ ( x , s ) with 1000 discretization steps and then form the T -layer diffusion bridge for HSIVI-SM by selecting T discrete time steps. Figure 5 shows the samples from DDPM, DDIM, and HSIVI-SM with T = 5 steps. We see that the samples produced by HSIVI-SM are much cleaner and more recognizable than those produced by DDPM and DDIM.

CIFAR-10, CelebA & ImageNet On both CIFAR-10 and CelebA, the structure of our pre-trained noise model ϵ ∗ ( x , s ) follows the UNet structure(Ronneberger et al., 2015) employed by Ho et al. (2020), instead of the huge VP deep continuous-time model (Song et al., 2020b) that has more channels and layers. We also provide additional results on ImageNet (64 × 64) with more powerful pre-trained score nets in (Nichol & Dhariwal, 2021)(bigger models with more parameters). Since this generative modeling has been formulated as a score-based VI problem, we do not have to use any training data for training HSIVI-SM.

Following the noise schedule employed in Song et al. (2020a), we first pre-train the noise model ϵ ∗ ( x , s ) with 1000 discretization steps and then form the T -layer diffusion bridge for HSIVI-SM

DDPM (T=5)

DDIM (T=5)

HSIVI-SM (T=5)

927835

S

199

3

85

D

406

3586635

9

3

2S51By

7462673

7432

7

968588

8

7

2

5

4

D

G

5

L

6879538

1

4

C

G

3

876438

1373（

33

73

7

6

0

7

2

h

士

3

9

2

3

9

by selecting T discrete time steps as before. For HSIVI-SM with ϵ -training, the conditional layer q t ( ·| x t +1 ; ϕ ) is modeled as a Gaussian distribution with mean µ t ( x t +1 ; ϕ µ ) and diagonal variance matrix Σ t ( ϕ σ ) where { ϕ µ , ϕ σ } = ϕ are the variational parameters. In our implementations, both µ t ( x t +1 ; ϕ µ ) and f t ( x t ; ψ ) use the same architecture as ϵ ∗ ( x , s ) . The number of layers, which is also the number of function evaluations (NFE), is set to be T = 5 , 10 , 15 in our experiments. We train HSIVI-SM with the same setting for T = 10 , 15 . The 5-layer HSIVI-SM is trained by further fine-tuning the well-trained 15-layer HSIVI-SM and we find this strategy leads to better results. During each nested training loop of f t ( x t ; ψ ) , we update ψ 20 times before each update of ϕ , since we find f t ( x t ; ψ ) needs more training empirically to provide reliable guidance.

For each method, we draw 50,000 samples and use the Fréchet inception distance (FID) score (Karras et al., 2022) to evaluate the sample quality (Table 2). We find that HSIVI-SM performs on par or better than the other baselines on both CIFAR-10 and CelebA, and the advantage is evident when the NFE is small. The sampling trajectories of 10-layer HSIVI-SM on CelebA with the same starting point but different random seeds are shown in Figure 6. We see that HSIVI-SM is capable of producing more diverse samples due to its stochastic nature, which is different from existing ODE based fast diffusion model samplers.

5.3 Additional Study

x

9

seed 1

seed 2

seed 3

seed 4

seed 5

x

6

x

3

x

0

Ablation of layers number In Figure 7, We provide a failure case on fitting the checkerboard target with diffusion bridge, demonstrating that the HSIVI-SM algorithm fails when the layer number T is small (the distances of auxiliary distributions at successive time steps are large) on a checkerboard

4

2

0

2

4

4

3

2

1

0

1

2

3

4

T=1

4

2

0

2

4

4

3

2

1

0

1

2

3

4

T=2

4

2

0

2

4

4

3

2

1

0

1

2

3

4

T=3

4

2

0

2

4

4

3

2

1

0

1

2

3

4

T=4

4

2

0

2

4

4

3

2

1

0

1

2

3

4

T=5

distribution. In fact, the score function on the checkerboard target is sharp on the boundaries but vanishes elsewhere. Therefore, fitting this target distribution is somewhat challenging.

Ablation of the variational family To validate the improvement of HSIVI-SM on diffusion models, we train HSIVI-SM with isotropic conditional layers in consistency with denoising-diffusion sampling, like DDPM and DDIM. We report the results of FID on the CIFAR-10 dataset in Table 3. These results provide further evidence for the statement outlined in Section 3.2. HSIVI-SM matches the marginal distributions q t ( x t ) and p t ( x t ) directly via score matching and would ensure a better fit for p 0 ( x 0 ) . The enhancement of HSIVI-SM over DDPM stems not only from its more expressive variational distribution but also from the direct alignment of the marginal distributions.

6 Conclusions

We introduced HSIVI, a hierarchical semi-implicit variational inference method that enables more expressive multi-layer construction of semi-implicit distributions. Given appropriate auxiliary distributions that interpolate between a simple base distribution and the target distribution, the conditional layers in hierarchical semi-implicit distributions can be progressively trained one layer after another. In experiments, we showed that HSIVI outperforms previous single-layer SIVI methods on several Bayesian inference tasks with complicated posteriors. HSIVI can also be used to accelerate the sampling process of diffusion models, where pre-trained score networks serve as a natural sequence of bridging distributions, which allows for direct acceleration of the stochastic diffusion model and does not require expensive sampling from the diffusion models during training. We showed that HSIVI can produce high quality samples comparable to or better than existing fast diffusion model samplers with few function evaluations on various datasets. Limitations are discussed in Appendix F.

References

Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems , 34:17981-17993, 2021.

Fan Bao, Chongxuan Li, Jun Zhu, and Bo Zhang. Analytic-dpm: an analytic estimate of the optimal reverse variance in diffusion probabilistic models. arXiv preprint arXiv:2201.06503 , 2022.

Espen Bernton, Jeremy Heng, Arnaud Doucet, and Pierre E. Jacob. Schrödinger bridge samplers. arXiv preprint arXiv:1912.13170 , 2019.

Christopher M. Bishop and Michael E Tipping. Variational relevance vector machines. In Proceedings of the Sixteenth conference on Uncertainty in artificial intelligence , pp. 46-53, 2000.

David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe. Variational inference: A review for statisticians. Journal of the American Statistical Association , 112:859 - 877, 2016.

Tiangang Cui, Kody JH Law, and Youssef M Marzouk. Dimension-independent likelihood-informed mcmc. Journal of Computational Physics , 304:109-137, 2016.

Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. In International Conference on Learning Representations , 2017.

Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Score-based generative modeling with criticallydamped langevin diffusion. In International Conference on Learning Representations (ICLR) , 2022.

J. Giordano, T. Broderick, and M. I. Jordan. Linear response methods for accurate covariance estimates from mean field variational bayes. In Advances in Neural Information Processing Systems , 2015.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Tianyang Hu, Zixiang Chen, Hanxi Sun, Jincheng Bai, Mao Ye, and Guang Cheng. Stein neural sampler. arXiv preprint arXiv:1810.03545 , 2018.

Ferenc Huszár. Variational inference using implicit distributions. arXiv preprint arXiv: 1702.08235 , 2017.

Aapo Hyvärinen. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research , 6(4):695-709, 2005. URL http://jmlr.org/papers/v6/ hyvarinen05a.html .

S. Jaakkola and M. I. Jordan. Improving the mean field approximation via the use of mixture distributions. In Learning in Graphical Models , pp. 173-173, 1998.

Michael I. Jordan, Zoubin Ghahramani, T. Jaakkola, and Lawrence K. Saul. An introduction to variational methods for graphical models. Machine Learning , 37:183-233, 1999.

Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusionbased generative models. In Proc. NeurIPS , 2022.

P. Kingma and J. Ba. Adam: A method for stochastic optimization. In ICLR , 2015.

P. Kingma and M. Welling. Auto-encoding variational bayes. In ICLR , 2014.

Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Advances in neural information processing systems , 34:21696-21707, 2021.

Diederik P. Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. In Advances in Neural Information Processing Systems , pp. 4743-4751, 2016.

Zhifeng Kong and Wei Ping. On fast sampling of diffusion probabilistic models. arXiv preprint arXiv:2106.00132 , 2021.

Qiang Liu, Jason Lee, and Michael Jordan. A kernelized stein discrepancy for goodness-of-fit tests. In International conference on machine learning , pp. 276-284. PMLR, 2016.

Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. arXiv preprint arXiv:2206.00927 , 2022.

Eric Luhman and Troy Luhman. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprint arXiv:2101.02388 , 2021.

M. Mescheder, S. Nowozin, and A. Geiger. Adversarial variational bayes: Unifying variational autoencoders and generative adversarial networks. In Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017 , 2017.

Moens, H. Ren, A. Maraval, R. Tutunov, J. Wang, and H. Ammar. Efficient semi-implicit variational inference. arXiv preprint arXiv:2101.06070 , 2021.

Radford Neal. MCMC using hamiltonian dynamics. In S Brooks, A Gelman, G Jones, and XL Meng (eds.), Handbook of Markov Chain Monte Carlo , Chapman & Hall/CRC Handbooks of Modern Statistical Methods. Taylor & Francis, 2011. ISBN 9781420079425. URL http://books. google.com/books?id=qfRsAIKZ4rIC .

Radford M Neal. Annealed importance sampling. Statistics and computing , 11:125-139, 2001.

Alex Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models, 2021.

Chenhao Niu, Yang Song, Jiaming Song, Shengjia Zhao, Aditya Grover, and Stefano Ermon. Permutation invariant graph generation via score-based generative modeling. In International Conference on Artificial Intelligence and Statistics , 2020.

J. Nott, S. L. Tan, M. Villani, and R. Kohn. Regression density estimation with variational methods and stochastic approximation. Journal of Computational and Graphical Statistics , 21(3):797-820, 2012.

W. Paisley, D. M. Blei, and M. I. Jordan. Variational bayesian inference with stochastic search. In Proceedings of the 29th International Conference on Machine Learning ICML , 2012.

Papamakarios, E. Nalisnick, D. Rezende, S. Mohamed, and B. Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. ArXiv Preprint arXiv:1912.02762 , 2019.

Fabian Pedregosa, Gaël Varoquaux, Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, et al. Scikit-learn: Machine learning in python. the Journal of machine Learning research , 12:2825-2830, 2011.

Ranganath, S. Gerrish, and D. M. Blei. Black box variational inference. In AISTATS , pp. 814-822, 2014.

Rajesh Ranganath, Dustin Tran, and David Blei. Hierarchical variational models. In International conference on machine learning , pp. 324-333. PMLR, 2016.

Rezende and S. Mohamed. Variational inference with normalizing flows. In Proceedings of The 32nd International Conference on Machine Learning , pp. 1530-1538, 2015.

J. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International Conference on Machine Learning , 2014.

Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computerassisted intervention . Springer, 2015.

Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. ArXiv , abs/2202.00512, 2022.

K. Saul and M. I. Jordan. Exploiting tractable substructures in intractable networks. In Advances in Neural Information Processing Systems , 1996.

Shi, S. Sun, and J. Zhu. Kernel implicit variational inference. In International Conference on Learning Representations , 2018a.

Shi, S. Sun, and J. Zhu. A spectral approach to gradient estimation for implicit distributions. In International Conference on Machine Learning , 2018b.

Artem Sobolev and Dmitry P. Vetrov. Importance weighted hierarchical variational inference. In Neural Information Processing Systems , 2019.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502 , 2020a.

Song, S. Garg, J. Shi, and S. Ermon. Sliced score matching: A scalable approach to density and score estimation. In Proceedings of the Thirty-Fifth Conference on Uncertainty in Artificial Intelligence , 2019.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems , 32, 2019.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020b.

Masashi Sugiyama, Taiji Suzuki, and Takafumi Kanamori. Density ratio estimation in machine learning . Cambridge University Press, 2012.

Zoltán Szabó. Information theoretical estimators toolbox. Journal of Machine Learning Research , 15:283-287, 2014.

Michalis K. Titsias and Francisco J. R. Ruiz. Unbiased implicit variational inference. In The 22nd International Conference on Artificial Intelligence and Statistics , pp. 167-176. PMLR, 2019.

Tran, D. M. Blei, and E. M. Airoldi. Copula variational inference. In Advances in Neural Information Processing Systems , 2015.

Tran, R. Ranganath, and D. M. Blei. Hierarchical implicit models and likelihood-free variational inference. In Advances in Neural Information Processing Systems , 2017.

Arash Vahdat and Jan Kautz. NVAE: A deep hierarchical variational autoencoder. In Neural Information Processing Systems (NeurIPS) , 2020.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Computation , 23(7):1661-1674, 2011. doi: 10.1162/NECO_a_00142.

J. Wainwright and M. I. Jordan. Graphical models, exponential families, and variational inference. Foundations and Trends in Maching Learning , 1(1-2):1-305, 2008.

Mingzhang Yin and Mingyuan Zhou. Semi-implicit variational inference. In International Conference on Machine Learning , pp. 5646-5655, 2018.

Longlin Yu and Cheng Zhang. Semi-implicit variational inference via score matching. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/ forum?id=sd90a2ytrt .

Zhang, B. Shahbaba, and H. Zhao. Variational hamiltonian monte carlo via score matching. Bayesian Analysis , 13(2):485-506, 2018.

Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. In NeurIPS 2022 Workshop on Score-Based Methods , 2022.

Hongkai Zheng, Weili Nie, Arash Vahdat, Kamyar Azizzadenesheli, and Anima Anandkumar. Fast sampling of diffusion models via operator learning. In NeurIPS 2022 Workshop on Score-Based Methods , 2022.

A Details of diffusion models

Diffusion models work by adding noise to the training data in the forward process and then removing the noise to recover the data in the backward process, which can be integrated into a general stochastic differential equation (SDE) framework (Song et al., 2020b). The forward process { u s } s ∈ [0 ,L ] is usually described by the SDE

where p 0 ( · ) is the data distribution, w s is a standard Brownian motion, f ( u s , s ) and g ( s ) are the drift and diffusion coefficient respectively. To generate samples from the data distribution, one can run the following reversed SDE

where p s ( · ) is the probability density function (pdf) of u s and ¯ w s is a standard Brownian motion when time flows from L to 0 . There exists deterministic process shares the same marginal probability densities { p s ( · ) } s ∈ [0 ,L ] described by the following ordinary differential equation (ODE)

called probability flow (PF) ODE.

In practice, Song et al. (2020b) and Kingma et al. (2021) designed several examples of the forward process such that it diffuses the data distribution p 0 ( · ) to a fixed unstructured distribution p L ( · ) . Here we mainly consider the Variance Preserving SDE (VP-SDE) used in DDPM (Ho et al., 2020; Song et al., 2020b). Let the drift coefficient f ( u s , s ) = dlog α ( s ) 2d s u s and the diffusion coefficient g 2 ( s ) = -dlog α ( s ) d s , where α ( s ) ∈ R + is a decreasing smooth function with α (0) = 1 , α ( L ) ≈ 0 . Then the distribution of u s conditioned on u 0 is explicit as

where ϵ is a standard Gaussian noise. In practice, diffusion models use a neural network S θ ( u s , s ) to approximate the score function S θ ( u s , s ) by optimizing the denoising score matching objective (Vincent, 2011)

where ω ( s ) is a positive weighting function. Instead of modeling the score function, Ho et al. (2020) proposed to predict the conditional noise ϵ based on u t . This leads to the following DDPM loss

where ¯ ω ( s ) is a positive weighting function. In fact, we have the relationship

We call L dsm 'score-prediction' training and L ddpm ' ϵ -prediction' training.

With the pre-trained score model S θ ( u s , s ) or noise model ϵ θ ( u s , s ) , Song et al. (2020b) shows that the samples of p 0 ( · ) can be generated by simulating the backward SDE, e.g. the sampling scheme of DDPM (Ho et al., 2020). Moreover, Bao et al. (2022) proposed Analytic-DPM, the optimal discretization form responding to the KL divergence of the joint distribution on the discrete time steps. Also, several high-order ODE solvers (Song et al., 2020a; Zhang & Chen, 2022; Lu et al., 2022) were proposed to achieve faster sampling.

B More details of UIVI

Unlike optimizing the surrogate ELBO, Titsias & Ruiz (2019) proposed unbiased implicit variational inference (UIVI) which relies on an unbiased gradient estimator for the exact ELBO. To elaborate further, reparametrize the conditional q ϕ ( x | z ) such as x = T ϕ ( z , ϵ ) , ϵ ∼ q ϵ ( ϵ ) , then

where

The second gradient term g ent ϕ involves an expectation w.r.t. the reverse conditional q ϕ ( z ′ | x ) which is estimated by an MCMC sampler in UIVI. However, the inner-loop MCMC runs may require long iterations for convergence.

C More details of HSIVI

C.1 Score-based training of HSIVI-LB

In the sequential training of HSIVI-LB, although the objective L SIVI-LB ( p t ( x t ) ∥ q t ( x t ; ϕ ≥ t )) is calculated based on p t ( x ) , the gradient of it w.r.t. ϕ t has a closed form containing only the score function S t ( x ) without knowing the corresponding pdfs. This derivation is important in the tasks where score functions of the auxiliary distributions are tractable while pdfs (up to a constant) of them are unavailable (for example, the diffusion bridge in Example 2). Concretely, assume the t -th conditional layer q t ( x t | x t +1 ; ϕ t ) is induced by a parametrized transform x t = h t ( x t +1 , ϵ ; ϕ t ) where ϵ ∼ p ϵ ( ϵ ) is a random noise, since q t ( x t | x t +1 ; ϕ t ) is reparametrizable according to Definition 1. The only term in L SIVI-LB ( p t ( x t ) ∥ q t ( x t ; ϕ ≥ t )) containing p t ( x t ) is E q t ( x t ; ϕ ≥ t ) log p t ( x t ) (see equation (2)) whose gradient takes the form

by the chain rule, where ∇ ϕ t h t ( x t +1 , ϵ ; ϕ t ) is the jacobian matrix of h t ( x t +1 , ϵ ; ϕ t ) .

In our implementation of HSIVI (in both sequential training and joint training), we generally assume the conditional layer q t ( ·| x t +1 ; ϕ t ) is induced by

where Σ t ( x t +1 ; ϕ t ) is a positive definite covariance matrix and ϵ ∼ N ( 0 , I ) is a standard multivariate gaussian variable. In equation (17), ϕ t should be replaced by ϕ in the joint training case.

C.2 Proof of Proposition 1

Proposition 1. Let q t ( x t , x t +1 ; ϕ ≥ t ) = q t ( x t | x t +1 ; ϕ t ) q t +1 ( x t +1 ; ϕ ≥ t +1 ) . The minimax optimization of L SIVI-SM ( p t ( x t ) ∥ q t ( x t ; ϕ ≥ t )) is equivalent to

Proof of Propsition 1 The minimax optimization problem of L SIVI-SM ( p t ( x t ) ∥ q t ( x t ; ϕ ≥ t )) is

according to equation (4). For minimization w.r.t. ϕ t , this target is equivalent to

For maximization w.r.t. ψ t , this target is equivalent to

where C is a term that does not contain ψ t . Therefore, the minimax optimization problem is equivalent to

C.3 Joint training of HSIVI

As mentioned in Section 4.2, when parameter sharing scheme is used in the conditional layers for application to diffusion model acceleration, sequential training from t = T -1 to t = 0 is not feasible. Therefore, we consider the following training objective

An intuitive method is to randomly sample a batch of time steps { t k } K k =1 and for each t k train L SIVIf ( p t k ( x t k ) ∥ q t k ( x t k ; ϕ )) directly. However, sequentially sampling x t k through q ( x i | x i +1 ; ϕ ) from i = T -1 to i = t k is still necessary in this case, making it memory-consuming to preserve the computation graphs of the entire sampling process.

In order to reduce the cost of accumulating computation graphs, for each t , we treat q t +1 ( x t +1 ; ϕ ) as a fixed mixing layer denoted by ˜ q t +1 ( x t +1 ) and only fit the conditional layer q t ( x t | x t +1 ; ϕ ) . More specifically, for HSIVI-SM, we consider the following optimization problem

where ˜ q t ( x t , x t +1 ; ϕ ) = q t ( x t | x t +1 ; ϕ )˜ q t +1 ( x t +1 ) . In what follows, we demonstrate that the above problem also ensures an accurate approximation of the target score function.

For equation (19), by the denoising score matching trick (Hyvärinen, 2005), the optimal point of ψ , denoted by ψ ∗ ( ϕ ) , satisfies

where ˜ q t ( x t ; ϕ ) = ∫ q ( x t | x t +1 ; ϕ )˜ q ( x t +1 )d x t +1 . By plugging in the optimal point ψ ∗ ( ϕ ) , each term in equation (18) is equivalent to

Therefore, the global optimal point ϕ ∗ also ensures that the score of the variational distribution fits the target score function.

Based on the training objectives (18) (19) mentioned above, we propose Algorithm 2 for joint training, which does not need to store the computation graphs of the sample sequences. Moreover, by assuming an increasing weighting function β ( t ) , we assign larger weights β ( t ) for those t close to T -1 , which tends to train the conditional layers that are close to T -1 first during the training, resembling the sequential training.

T -1 (0)

Input: Auxiliary bridge { p t ( x ) } t =0 ; a weighting function β ( t ) ; initial value of parameters ϕ . Output: The optimal parameters ϕ ∗ .

Initialization: ϕ ← ϕ (0) .

while not converge do

Uniformly sample K time steps { t k } K k =0 with replacement from { 0 , . . . , T -1 } .

Sample a minibatch { x ( k ) T } K k =1 from the base distribution q T ( x ) .

for

k

= 1

, . . . , K

and

t

k

< T

-

1

do

Sequentially sample x ( k ) t k +1 through q ( x i | x i +1 ; ϕ i ) from i = T -1 to i = t k +1 .

Detach the computation graphs from { x ( k ) t k +1 } K k =1 .

Update ϕ by optimizing the objective ∑ K k =1 β ( t k ) L SIVIf ( p t k ( x t k ) ∥ q t k ( x t k ; ϕ )) , where the k -th term is computed based on a single sample x ( k ) t k +1 .

end for end while

ϕ

∗

ϕ

.

C.4 ϵ -training of HSIVI-SM

Another popular formulation of diffusion models is modeling the conditional noise ϵ θ ( u s , s ) by optimizing the DDPM loss in equation (14) where u s = √ α ( s ) u 0 + √ 1 -α ( s ) ϵ , introduced as ' ϵ -prediction' in Appendix A. Now, let us assume the diffusion bridge is constructed with VP-SDE and we have a pre-trained model of conditional noise ϵ ∗ ( u , s ) . Similarly, we construct a sequence of noise models { ϵ ∗ t ( x t ) } T -1 t =0 by letting x t = u s t and ϵ ∗ t ( x ) = ϵ ∗ t ( x , s t ) which forms a (generalized) T -layer diffusion bridge. We only discuss how ϵ -training can be applied to joint training and the derivation for sequential training is similar. In what follows, we consider the transformation of the joint training objective L HSIVI-SM for diffusion model acceleration.

By letting the weighting function β ( t ) = 1 -α ( s t ) and considering the reparametrization form (17) where ϕ t is replaced by ϕ , the objective of HSIVI-SM takes the form

where S ∗ t ( x t ) is a pre-trained score model. Note that we have √ β ( t ) S ∗ t ( x t ) = -ϵ ∗ t ( x t ) by equation (15). Define

The HSIVI-SM objective (20) then takes the form

and we call it the objective for ϵ -training. In our implementation of ϵ -training, we directly parametrize ˜ f t ( x t ; ψ ) and ˜ Σ t ( x t +1 ; ϕ ) instead of f t ( x t ; ψ ) and Σ t ( x t +1 ; ϕ ) . The objective (21) is more numerically stable since the magnitude of ˜ Σ t ( x t +1 ; ϕ ) is generally larger than Σ t ( x t +1 ; ϕ ) .

C.5 Complexity comparison of HSIVI

For methods that we discussed in SIVI variants, which use a single conditional layer (i.e., T=1) and hence would be much cheaper to sample from than HSIVI-SM that uses multiple layers T > 1 . For methods that we discussed for diffusion models, the computational complexity would be similar if they had the same T . That is because we used the same neural network architecture for the conditional layers in HSIVI and the score nets in diffusion models. We have a comparison of the sampling time of different methods in Figure 13.

HSIVI-LB (t=4)

HSIVI-LB (t=3)

HSIVI-LB (t=2)

HSIVI-LB (t=1)

HSIVI-LB (t=0)

15

15

15

15

15

10

10

10

10

10

5

5

5

5

0

0

0

0

0

5

-5

-5

-5

5

10

10

10

10

10

15

10

0

10

15

10

0

10

15

-10

0

10

15

10

0

10

15

10

0

10

HSIVI-SM (t=4)

HSIVI-SM (t=3)

HSIVI-SM (t=2)

HSIVI-SM (t=1)

HSIVI-SM (t=0)

15

15

15

15

15

10

10

10

10

10

5

5

5

5

5

0

0

0

0

0

-5

-5

-5

-5

-5

10

10

10

10

10

15

-10

0

10

15

-10

0

10

15

-10

0

10

15

-10

0

10

15

-10

0

10

0.0

0.5

1.0

1.5

2.0

2.5

3.0

s

1.4

1.2

1.0

0.8

0.6

0.4

0.2

0.0

0.2

xs

Ground Truth

95% Confidence Interval

Observations

Sample Mean

0.0

0.5

1.0

1.5

2.0

2.5

3.0

s

1.4

1.2

1.0

0.8

0.6

0.4

0.2

0.0

0.2

xs

SIVI-LB

95% Confidence Interval

Observations

Sample Mean

0.0

0.5

1.0

1.5

2.0

2.5

3.0

s

1.4

1.2

1.0

0.8

0.6

0.4

0.2

0.0

0.2

xs

HSIVI-LB (T=5)

95% Confidence Interval

Observations

Sample Mean

D Additional results of experiments

D.1 Gaussian mixture model

For HSIVI on the Gaussian mixture model, the auxiliary distributions can also be constructed with diffusion bridge in Example 2. Concretely, the diffusion bridge is constructed by

where α t = α ( s t ) with α ( s ) defined in equation (12). In this example, the score function S t ( x t ) = ∇ x t log p t ( x t ) has an analytical form

where S 0 ( x ; µ , σ 2 I ) is the score function of the Gaussian mixture model p ( x ; µ , σ 2 I ) = ∑ 8 i =1 1 / 8 · N ( x ; µ i , σ 2 I ) . We set the number of layers T = 5 and α t = 1 -t/ 5 for t = 0 , . . . , 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,

(a) 5 steps

(b) 10 steps

Ground Truth

DDPM (T=5)

DDIM (T=5)

HSIVI-SM (T=5)

Ground Truth

DDPM (T=5)

DDIM (T=5)

HSIVI-SM (T=5)

4

3

O

3

3

3

Z

2

2

1

1

1

1

1

0

0

0

0

0

0

0

-1

-1

1

-1

-1

-3

E=

-2

2

E

-2

4

0

Z=

0

Z

Z=

0

2

Z

0

2

=

0

2

-2

0

2

-2

0

2

Ground Truth

DDPM (T=5)

DDIM (T=5)

HSIVI-SM (T=5)

Ground Truth

DDPM (T=5)

DDIM (T=5)

HSIVI-SM (T=5)

2.0

2.0

2.0

2.0

3

3

3

3

0.5

1.5

1.0

0.0

1.5

1.0

0.5

0.0

0.5

0.0

1.5

1.0

1.5

1.0

0.5

0.0

0

2

1

0

0

1

0.5

0.5

0.5

0.5

-1

-1

-1

1.0

1.0

1.0

1.0

-2

-Z

2

-2

1.5

1.5

1.5

1.5

-2

-2

0

-2

-2

-2

0

Z-

0

Ground Truth

DDPM (T=10)

DDIM (T=10)

HSIVI-SM (T=10)

Ground Truth

DDPM (T=10)

DDIM (T=10)

HSIVI-SM (T=10)

3

2

4

2

2

O

2

E

2

O

1

1

1

0

0

0

0

0

0

-1

-1

-1

-1

-1

Z-

Z=

Z

2

2

E-

E-

E:

3

2

0

2

-2

0

2

2

0

2

2

0

2

2

0

2

0

2

2

0

2

2

2

Ground Truth

DDPM (T=10)

DDIM (T=10)

HSIVI-SM (T=10)

Ground Truth

DDPM (T=10)

DDIM (T=10)

HSIVI-SM (T=10)

2.0 -

2.0

2.0

2.0

3

3 -

3

3

1.5

1.0

1.0

1.5

1.5

1.0

1.5

1.0

1

Z

2

0.0

0.5

0.5

0.0

0.5

0.0

0.5

0.0

0

0

0.5

0.5

0.5

0.5

1

1.0

1.0

1.0

1.0

1.5

1.5

1.5

1.5

E-

Z=

0

2

Z=

0

2

Z=

0

-2

0

2

2

2

-2

0

-2

0

2

-2

0

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.

i.e.

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)

HSIVI-SM (T=10)

98388

794835

5

7093

3

694

2

4662

08

42829

3586℃635

G

9

3

？

2

77

女

A

8

8

248 d

8

B

73

79 G

8

3

力

8

7

13

24

3263

3

4

9

8

9723

631655

03

4

1

63

の1

6

/

O

/S

一

9

a

2

4

44952

9923

3

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 × 28 ), CelebA ( 64 × 64 ) and ImageNet ( 64 × 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 HSIVISM, the number of parameters in the noise model equals that in the conditional layer q t ( x t | x t +1 ; ϕ ) . We find that our model with the same parameters reaches better results in Table 2.

E Experimental details

E.1 Target distribution approximation

In this part, we set the conditional layer to be q ϕ ( x | z ) = N ( x ; µ ( z ; ϕ µ ) , diag { exp( ϕ σ ) } ) and the mixing layer to be N ( 0 , I ) for SIVI. Here, { ϕ µ , ϕ σ } = ϕ are the variational parameters. For T -layer hierarchical semi-implicit variational distribution with T ≥ 2 , the variational prior q T ( x T ) is set to be N ( 0 , I ) . Each conditional layer q t ( x t | x t +1 ; ϕ t ) for t = 0 , . . . , T -1 is a conditional Gaussian distribution

Note that the ϕ σ and { ϕ σ t } T -1 t =0 above are all vectors with the same dimension as x . We use sequential training for HSIVI in the two experiments in this part. The parameters { ϕ t } T -1 t =0 are independent across different t . If not otherwise specified, we use the Adam optimizer (Kingma & Ba, 2015) with β = (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 µ ( z ; ϕ µ ) in SIVI or µ ( x t +1 ; ϕ µ t ) in HSIVI has a residual form, i.e. µ ( z ; ϕ µ ) = z + ¯ µ ( z ; ϕ µ ) and µ ( x t +1 ; ϕ µ t ) = x t +1 + ¯ µ ( x t +1 ; ϕ µ t ) , for t = 0 , . . . , T -1 . ¯ µ ( z ; ϕ µ ) in SIVI and { ¯ µ ( x t +1 ; ϕ µ t ) } 4 t =0 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 , f t ( x t ; ψ t ) in HSIVI-SM and f ( x ; ψ ) in SIVI-SM are parameterized by MLPs with layer widths [2 , 128 , 128 , 2] and ReLU activation functions.

T = 5

T = 10

T = 15

(a) CIFAR-10 (32

×

32)

T = 5

T = 10

T = 15

(b) CelebA (64

×

64)

T = 5

T = 10

T = 15

(c) ImageNet (64

×

64)

The noise levels in the diffusion bridge are 1 -α ( s t ) = 1 -t/ 5 for t ∈ { 0 , 1 , · · · , 4 } . We set the learning rate of variational parameters ϕ t (or ϕ ) to 0.001 and the learning rate of ψ t (or ψ ) 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 × 80000 variational parameter updates. For HSIVI-SM and SIVI-SM, in each nested training loop of f t ( x t ; ψ t ) (or f ( x ; ψ ) ), we update ψ t (or ψ ) one time after each update of ϕ t (or ϕ ). 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

0.00

0.25

0.50

0.75

1.00

1.25

1.50

1.75

Time (seconds)

DDIM

HSIVI-SM (ours)

DPM-Solver-fast

FastDPM

Analytic-DPM

CIFAR-10

0

1

2

3

4

Time (seconds)

HSIVI-SM (ours)

DDIM

Analytic-DPM

DPM-Solver-fast

FastDPM

CelebA

after 10,000 iterations. For t = 0 , . . . , T -2 , the mean of each conditional layer µ ( x t +1 ; ϕ µ t ) in HSIVI has a residual form, i.e. µ ( x t +1 ; ϕ µ t ) = x t +1 + ¯ µ ( x t +1 ; ϕ µ t ) . For SIVI and t = T -1 in HSIVI, we assume µ ( z ; ϕ µ ) = ¯ µ ( z ; ϕ µ ) and µ ( x t +1 ; ϕ µ t ) = ¯ µ ( x t +1 ; ϕ µ t ) . For each t , ¯ µ t ( x ; ϕ µ t ) in HSIVI and ¯ µ ( z ; ϕ µ ) in SIVI are MLPs with layer widths [300 , 512 , 512 , 300] and ReLU activation functions. For each t , f t ( x t ; ψ t ) in HSIVI-SM and f ( x ; ψ ) 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 f t ( x t ; ψ t ) (or f ( x ; ψ ) ), we update ψ t (or ψ ) one time after each update of ϕ t (or ϕ ). We set the learning rate to be 0.0001 for ϕ t (or ϕ ) and 0.0005 for ψ t (or ψ ).

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 ( x T ) is set to be N ( 0 , I ) . To avoid the large memory consumption, we use the joint training method where the parameters of the conditional layers q t ( x t | x t +1 ; ϕ ) and f t ( x t ; ψ ) are the same across different t . The t -th conditional layer is a conditional Gaussian distribution

where { ϕ µ , ϕ σ } = ϕ are the variational parameters, ϕ σ is a vector with the same dimension as x , and σ t is a fixed scalar value. We use the generalized inference process in DDIM (Song et al., 2020a) with the noise level η > 0 to initialize µ t ( x t +1 ; ϕ µ ) and determine the value of σ t for each t . If not otherwise specified, we use the Adam optimizer (Kingma & Ba, 2015) with β = (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 ∗ ( x , s ) , we consider quadratic noise levels 1 -α ( s ) = s 2 for s ∈ [0 , 1] . We then train S ∗ ( x , s ) on 1000 fixed noise levels { 1 -α ( i/ 1000) } 1000 i =1 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 -1 t =0 so that 1 -α ( s t ) = [0 . 01 + ( √ 0 . 8 -0 . 1) t/T ] 2 for t = 0 , 1 , . . . , T -1 .

Model architecture The model architecture of S ∗ ( x , s ) is

where MLP dec is a decoder implemented as MLPs with layer widths [128 , 128 , 128 , 2] , MLP embx is a data embedding block implemented as MLPs with layer widths [2 , 128] , and MLP 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 -α ( s ) as the input of MLP embt . All these three MLPs use GELU as activation functions. We use the generalized inference process with noise level η = 1 . 0 to initialize the conditional layers. The architecture of f t ( x t ; ψ ) is the same as that of S ∗ ( x , s ) . We initialize f t ( x t ; ψ ) with S ∗ t ( x t ) := S ∗ ( x t , s t ) .

Training setting The learning rate is set to be 0.0002 for q t ( x t | x t +1 ; ϕ ) and 0.0005 for f t ( x t ; ψ ) 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 ( ·| x t +1 ; ϕ ) and 0.00005 (0.0001) for f t ( x t ; ψ ) 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 f t ( x t ; ψ ) , we update ψ 3 times after each update of ϕ .

E.2.2 MNIST

For the experiment on MNIST, we use the pre-trained noise model ϵ ∗ ( x , s ) and train HSIVI-SM with ϵ -training introduced in Section C.4. The following construction of noise schedule comes from Song et al. (2020a). Let β j = β min + β max -β min 999 j for j = 0 , . . . , 999 , where β min = 0 . 0001 , β max = 0 . 02 . We pre-train the noise model on the 1000 fixed noise levels 1 -α ( s ) := ∏ s j =0 β j for s = 0 , . . . , 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 = ⌊ 800 · t 2 T 2 ⌋ for t = 0 , . . . , T -1 to construct the T -layer diffusion bridge.

Model architecture The pre-trained noise model ϵ ∗ ( 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 × 32 to fit ϵ ∗ ( x , s ) . We use the generalized inference process with noise level η = 0 . 2 to initialize the conditional layers. The architecture of f t ( x t ; ψ ) is the same as that of ϵ ∗ ( x , s ) . We initialize f t ( x t ; ψ ) with -ϵ ∗ ( x t , s t ) / √ 1 -α ( s t ) .

Training setting For both T = 5 , 10 , the learning rate is set to be 1 . 6 × 10 -5 for ϕ and 6 . 4 × 10 -5 for ψ . 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 ; ψ ) , we update ψ 20 times after each update of ϕ .

E.2.3 CIFAR-10, CelebA & ImageNet

For experiments on CIFAR-10 and CelebA, we use the pre-trained noise model ϵ ∗ ( x , s ) and train HSIVI-SM with ϵ -training introduced in Section C.4. We use the same noise schedule as in the experiment on MNIST. Let β j = β min + β max -β min 999 j for j = 0 , . . . , 999 , where β min = 0 . 0001 , β max = 0 . 02 . We take the pretrained noise model for CIFAR10 and ImageNet seperately 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 -α ( s ) := ∏ s j =0 β j for s = 0 , . . . , 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 = ⌊ 800 · t 2 T 2 ⌋ for t = 0 , . . . , T -1 to construct the T -layer diffusion bridge.

Model architecture On CIFAR-10 and CelebA, the structure of ϵ ∗ ( x , s ) is exactly the UNet 3 employed in Ho et al. (2020) without modification; on ImageNet, the structure of ϵ ∗ ( x , s ) is exactly the UNet in Nichol & Dhariwal (2021). We use the generalized inference process with noise level η = 0 . 2 to initialize the conditional layers. The architecture of f t ( x t ; ψ ) is the same as that of ϵ ∗ ( x , s ) . We initialize f t ( x t ; ψ ) with -ϵ ∗ ( x t , s t ) / √ 1 -α ( 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 × 10 -5 for q t ( ·| x t +1 ; ϕ ) and 8 × 10 -5 for f t ( x t ; ψ ) ; on CelebA, the learning rate is set to be 1 . 2 × 10 -6 for q t ( ·| x t +1 ; ϕ ) and 6 × 10 -6 for f t ( x t ; ψ ) ; on ImageNet, the learning rate is set to be 1 × 10 -5

for q t ( ·| x t +1 ; ϕ ) and 5 × 10 -5 for f t ( x t ; ψ ) . We trained HSIVI-SM for 10,000 iterations with a batch size of 128. During each nested training loop of f t ( x t ; ψ ) , we update ψ 20 times after each update of ϕ , since we find f t ( x t ; ψ ) 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 × 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.

F Limitations

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 ( x t | x t +1 ; ϕ ) , and f t ( x t ; ψ ) . 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces hierarchical semi-implicit variational inference (HISIVI) as an extension of semi-implicit variational inference (SIVI). HISIVI incorporates interpolating distributions between prior and target data to train conditional layers progressively, resulting in accelerated sampling in diffusion models.\n\nSTRENGTHS:\nThe paper is well-structured and easy to follow.\n\nWEAKNESSES:\n- The novelty of HISIVI compared to related work is not clearly demonstrated. It is recommended to include more related works and comparison: cascaded diffusion models[1], diffusion schrodinger bridge[2], flow matching[3], or rectified flow[4], etc.\n- The experiment section is relatively weak. Comparisons could strengthen the evaluation (aforementioned references and more [5],[6], etc.)\n\n[1] Ho, J., Saharia, C., Chan, W., Fleet, D.J., Norouzi, M. and Salimans, T., 2022. Cascaded diffusion models for high fidelity image generation. The Journal of Machine Learning Research, 23(1), pp.2249-2281.\n[2] De Bortoli, V., Thornton, J., Heng, J. and Doucet, A., 2021. Diffusion Schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 34, pp.17695-17709.\n[3] Lipman, Y., Chen, R.T., Ben-Hamu, H., Nickel, M. and Le, M., 2022. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747.\n[4] Liu, X., Gong, C. and Liu, Q., 2022. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003.\n[5] Xiao, Z., Kreis, K. and Vahdat, A., 2021. Tackling the generative learning trilemma with denoising diffusion gans. arXiv preprint arXiv:2112.07804.\n[6] Zheng, H., Nie, W., Vahdat, A., Azizzadenesheli, K. and Anandkumar, A., 2022. Fast sampling of diffusion models via operator learning. arXiv preprint arXiv:2211.13449.\n\nQUESTIONS:\nHow does the proposed optimization for score function in HISIVI outperform or differ from the conventional implicit score matching or denoising score matching[7]?\n\n[7] Vincent, P., 2011. A connection between score matching and denoising autoencoders. Neural computation, 23(7), pp.1661-1674.\n\nLIMITATIONS:\nThis work only deals with the geometric interpolation of distributions without providing justification. Exploring the impact of choice of interpolation on generation would enhance the quality of this work.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nSemi-implicit variational inference increases the expressiveness of variational posteriors but introduces intractabilities to their inference. This work adds to the semi-implicit variational inference (SIVI) work of Yu and Zhang (https://openreview.net/forum?id=sd90a2ytrt) which uses an alternative objective (the Fisher divergence) that, combined with the hierarchical nature of semi-implicit variational distributions, can be approximated as the solution of a mini-max optimization problem. This work takes inspiration from simulated annealing and denoising-diffusion models by proposing to expand the single-layer semi-implicit variational family into multiple-layer variants by specifying hierarchical auxiliary distributions to guide the semi-implicit variational distribution towards the target distribution. The authors denote this method Hierarchical Semi-Implicit Variational Inference (HSIVI). The auxiliary distributions are specified either by assuming the marginal densities of a geometric interpolation between a target and base distribution are available analytically or by using pre-trained scores from a diffusion bridge. The main contribution to this paper is stated in the author's claim that \"when used for diffusion model acceleration, we show that HSIVI can produce high quality samples comparable to or better than the existing fast diffusion model based samplers with small number of function evaluations on various datasets.\" This claim is tested empirically on the CIFAR-10 (32x32) and CelebA (64x64) datasets with results showing comparable or better FID scores to existing fast diffusion model samplers with few (between 5 and 15) function evaluations. As well as the main contribution, an algorithm for training is provided by noting that the conditional layers can be trained sequentially until convergence by exploiting the hierarchical structure of the semi-implicit distribution. Also, an efficient parameterization of the neural networks is provided to make the algorithm computationally feasible.\n\nSTRENGTHS:\nQuality\n\nThe work is a novel combination of well-known, peer-reviewed techniques including simulated annealing and SIVI. It is clear how this work differs from the work of Yu and Zhang (https://openreview.net/forum?id=sd90a2ytrt) using the hierarchical semi-implicit variational posteriors, and the value of this contribution is demonstrated on a multi-modal problem in Figure 2. The methods used seem to work well judging by the experimental results. \n\nClarity\n\nThe submission is well-organized and clearly written.\n\nWEAKNESSES:\nOriginality\n\nThere is limited technical novelty in this paper as the technical content on the work is almost entirely covered by the work of Yu and Zhang (https://openreview.net/forum?id=sd90a2ytrt) that rewrites the Fisher divergence leading to a form that does not require computing the score of the hiearchical variational posterior. The paper extends this by sequentially score matching to noised versions of the target distribution, which have been obtained analytically or approximated from pretrained denoising diffusion models. The main novelty in this work is instead an experimental one, showing that HSIVI can produce high quality samples comparable to or better than the existing fast diffusion model based samplers with a small number of function evaluations on two datasets by comparing FID scores (CIFAR-10 (32x32) and CelebA (64x64)) and one dataset by comparing samples visually (MNIST). The experiments were made computationally feasible by a trick to parameterize the neural networks that is detailed in Proposition 1.\n\nQuality of results\n\nRelatively simple experiments are provided that on one hand demonstrate the contribution well, but on the other hand do not cover a very diverse range of cases. No empirical failure case is provided for a complicated distribution that requires many function evaluations to sample from, which may provide an interesting avenue for future work.\n\nOn the reproducibility of the results, whilst the experimental setups are well detailed, no code is provided with the paper, nor indication it will become available upon publication. The reviewer would like to see that the code is made available since this is crucial for reproducibility. Furthermore, there may not be enough information to replicate the experiments. For example, as far as I can see, there is no discussion on the parameterization/values of the positive weighting function, beta(t) for joint HSIVI-SM training on CIFAR-10 and MNIST, which makes it impossible to reproduce without any code provided.\n\n\nSignificance\n\nThe work's significance would be improved if it addressed why we should expect improvements in sampling efficiency over denoising-diffusion sampling and where these improvements can be derived from (the particular form of variational family, or something else?).\n\nQUESTIONS:\n1. It would be interesting if the authors could demonstrate a failure case for their hierarchical variational posterior approximating the denoising-diffusion model, what properties of the score in the backwards process (equation 9) would cause the proposed variational posterior to fail with a small number of function evaluations?\n\n2. No interpretation of the neural network f_t(x_t) is given, apart from that it provides \"guidance\". Could the authors please clarify what is being guided, and to what?\n\n3. Comparing the proposed sampling algorithm to denoising-diffusion: the sampling algorithm for denoising-diffusion uses a reverse-time Markov chain that has specified isotropic variance that appears due to the assumed forward-time transition densities. Whereas the proposed HIVI sampling algorithm in this work is a Markov chain with transition densities that are the conditional variational posterior distributions (called the conditional layer in this work).  The difference between this Markov chain and the denoising-diffusion one is the trainable non-isotropic variance of the conditional layer q_t(\\cdot|x_{t+1}; \\phi). If the conditional layer is instead modelled as a Gaussian distribution with mean and an isotropic scale-Identity matrix, then do you expect that there is no improvement over denoising-diffusion with the same number of steps? It would be good to see an experiment that showed the trade off between sample quality and number of function evaluations using this simpler variational family, which is on par with the denoising-diffusion sampling.\n\n\nTypos:\n\nFisher divergence is not symmetric in its arguments and has been incorrectly written with arguments transposed (the expectation is taken over the distribution of the first argument).\nSee for example how the Fisher divergence is written in the papers that this work has cited\nhttps://arxiv.org/pdf/1602.03253.pdf\nor\nhttps://arxiv.org/pdf/1810.03545.pdf\n\nLine 510, 511, 512 there are typos in the norm regularizing the f network (written as f_{t}(x; \\psi_t) and should be f_{t}(x_{t}; \\psi_t))\n\nLIMITATIONS:\nThe limitation discussion picks up on the important points, but it would be good to know if higher dimensional problems were attempted, or problems with more parameters (such as the \"huge VP deep continuous-time model (Song et al., 2020b) that has more channels and layers\"), and did the authors run into memory issues that were discussed in Appendix F?\n\nAn ablation of the variational family for the conditional layers (as described in question 3) comparing directly to the denoising-diffusion sampling remains unexplored, and would be useful to see.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAuthors propose a hierarchical semi-implicit variational inference framework by extending the existing SIVI. Authors showed that the proposed HSIVI, given pre-trained score networks, can be used to accelerate the sampling process of diffusion models with the score matching objective. The numerical results show more enhanced performance in faithful modeling of complex distributions and diffusion model acceleration.\n\nSTRENGTHS:\nThe paper is fairly well-written, clearly motivated and generally addresses an important problem. The extension of SIVI to a hierarchical model, although not very novel, is intuitive and makes sense. I haven't checked the proofs but the overall approach seems to hold up.\n\nWEAKNESSES:\n- Extending SIVI to a hierarchical model, by itself, is not very novel. However, the application to diffusion acceleration is interesting. \n- The experiments and presented numerical results could be improved. (check next section for more)\n\nQUESTIONS:\n- It is mentioned in the manuscript that HSIVI requires more memory. How much more memory (with respect to data dimension) does it need compared to baselines? \n- ImageNet (64x64) is a common benchmark dataset for diffusion models. I would be nice to have the results for that as well.\n- How does the proposed model performs compared to Diffusion Normalizing Flows [1] both in terms of target distribution approximation and sample generation.\n- [A suggestion] It would be interesting to extend HSIVI on graphical models (in a setup like [2]).\n\n\n\n[1] Diffusion Normalizing Flow, Qinsheng Zhang and Yongxin Chen, NeurIPS 202.\n[2] Efficient Inference Amortization in Graphical Models using Structured Continuous Conditional Normalizing Flows, Christian Weilbach, 2019.\n\nLIMITATIONS:\nThe limitations to some extent have been discussed but more discussion (specially around memory usage) is needed.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}

comments: SUMMARY:
The authors extend semi-implicit variational inference to have multiple layers of latent variables, vastly increasing the. expressiveness of the model. A nice formulation of the asymptotic lower bound plus a training scheme is introduced as well.

STRENGTHS:
This paper is an extension that on the surface seems straight forward. But the authors demonstrate that this is far from the case and introduce a novel formulation of the semi-vi ELBO.

The paper was also written well and easy to follow. The experiments section was great too.

WEAKNESSES:
My only qualm is the flow between the first half of the paper and the second half where the focus is on diffusion models. Previously, the authors argued that training using a sequence of distributions would lead to better marginal distributions. But for speeding up sampling of diffusion models, it seems like training the joint distribution leads to impressive results. This opens up the question of whether the sequence of distributions is even needed at akk.

QUESTIONS:
I don't have any questions.

LIMITATIONS:
It seems like the main text is missing a limitations section.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-11

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a hierarchical semi-implicit variational inference that stacks multiple semi-implicit layers to construct a flexible generative distribution. The paper introduces a sequence of distributions that interpolate between the target and a base distribution. Each pair of intermediate inference and target distributions are matched by optimizing a SIVI bound. The method is applied in the diffusion model for sampling acceleration. \n\n##Post-rebuttal\n\nThanks for the authors' response. Most of my questions have been addressed and I will maintain my current score.\n\nSTRENGTHS:\n- The idea of using semi-implicit distribution for the progressive approximation is interesting and intuitive. \n\n- The empirical results show HSIVI improves the distribution approximation and sample quality significantly on several tasks. \n\n- Though the modeling and training of HSIVI consists of many components, the paper explains the ideas concisely and clearly.\n\nWEAKNESSES:\n- The validity of the training procedure needs more concrete explanations or theoretical results\n\n- The empirical study needs more evidence for the improvement of acceleration.\n\n- Some closely related works need more discussions of similarity and improvement. \n\nPlease see Questions for more details.\n\nQUESTIONS:\n- To train HSIVI sequentially, the paper mentions, “Let the parameters ϕ_t in the t-th conditional layer be independent across different ts.” But since q_t(x_t; ϕ_t) depends on x_{t-1} that depends on ϕ_{t-1}, It is understandable that the graph detaching in the algorithm benefits computation efficiency. Does it have any downside for not training ϕs jointly? For example, will there be convergence issues, or will some ϕs overfit while others underfit? \n\n- In Appendix C.3., the author shows that the detaching operation keeps the optimality of \\phi unchanged in the joint training. Does a similar conclusion apply to the sequential training in the main paper?\n\n- Table 1 shows with the same layers/steps, HSIVI achieves better approximation than DDPM and DDIM. Would the computation with one layer of HSIVI be more than one step of DDPM?  If so, it might not be a fair comparison. Can the author discuss the computational complexity and provide wall clock time? \n\n- [1] and [2] design a hierarchical model using semi-implicit distributions, which seem to be closely related to HSIVI. The paper might need more detailed discussions of these papers. \n\n[1] Importance weighted hierarchical variational inference (NeurIPS 2019)\n[2] Structured Semi-Implicit Variational Inference (AABI 2019)\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a new method called Hierarchical Semi-Implicit Variational Inference (HSIVI) that enhances the expressiveness of Semi-Implicit Variational Inference (SIVI) on complex target distributions. HSIVI works by applying SIVI to a hierarchy of latent variables, and using the variational distribution of the previous step as the implicit prior for the later step. The authors apply HSIVI for sampling from diffusion models, learned on a variety of synthetic and real-world datasets.\n\nSTRENGTHS:\n1. Novelty: The paper presents a new method called Hierarchical Semi-Implicit Variational Inference (HSIVI) that enhances the expressiveness of Semi-Implicit Variational Inference (SIVI) on complicated target distributions. This is a novel approach that has not been explored before, to the best of my knowledge.\n\n2. Effectiveness: The authors demonstrate the effectiveness of HSIVI on a variety of synthetic and real-world datasets, including accelerating the sampling process of diffusion models. The results show that HSIVI outperforms existing methods in terms of accuracy and efficiency.\n\n3. Clarity: The paper provides a detailed explanation of the HSIVI method, including its mathematical formulation and implementation. The authors also provide clear and concise descriptions of the experiments and results.\n\nWEAKNESSES:\nThe main weakness of this work is that the method seem fairly complicated to implement in a practical setting. The method is presented in a general way, with an algorithm box. However, later another objective function is presented with some parameter sharing. Then, in the larger-scale experiments, HSIVI-SM is trained by fune-tuning a larger model. Either the method is not extremely stable / usable, or there are too many details but those are necessary for the experiment to be setup properly? (in the latter case, perhaps one experiment should go to the supplements?).\n\nQUESTIONS:\n(1) Why isn't the performance of HSIVI-LB reported on the other experiments?\n(2) In terms of computational efficiency, it wasn't clear to me that if two compared methods had the same \"T\", then their computational burden were comparable. Are the computational complexities exactly the same for sampling from HSIVI-SM and SIVI, for equal \"T\"?\n(3) In principle, if HSIVI was used to train the model, we could reach a better data likelihood? Why aren't the authors restricting the application of HSIVI to sampling, and not learning the score function as well?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-18"}