Training Variational Autoencoders with Buffered Stochastic Variational Inference
Rui Shu Hung H. Bui Jay Whang Stefano Ermon







Buffered Stochastic Variational Inference






Rui Shu, Hung H. Bui, Jay Whang, Stefano Ermon

[

Training Variational Autoencoders with

Buffered Stochastic Variational Inference


Rui Shu

Stanford University
Hung H. Bui

VinAI
Jay Whang

Stanford University
Stefano Ermon

Stanford University
 
]










The recognition network in deep latent variable models such as variational autoencoders (VAEs) relies on amortized inference for efficient posterior approximation that can scale up to large datasets. However, this technique has also been demonstrated to select suboptimal variational parameters, often resulting in considerable additional error called the amortization gap. 

To close the amortization gap and improve the training of the generative model, recent works have introduced an additional refinement step that applies stochastic variational inference (SVI) to improve upon the variational parameters returned by the amortized inference model. In this paper, we propose the Buffered Stochastic Variational Inference (BSVI), a new refinement procedure that makes use of SVI's sequence

of intermediate variational proposal distributions and their corresponding importance weights to construct a new generalized importance-weighted lower bound.




We demonstrate empirically that training the variational autoencoders with BSVI consistently out-performs SVI, yielding an improved training procedure for VAEs.




§ INTRODUCTION
 

Deep generative latent-variable models are important building blocks in current approaches to a host of challenging high-dimensional problems including density estimation <cit.>, semi-supervised learning <cit.> and 

representation learning for downstream tasks <cit.>. To train these models, the principle of maximum likelihood is often employed. However, maximum likelihood is often intractable due to the difficulty of marginalizing the latent variables. Variational Bayes addresses this by instead providing a tractable lower bound of the log-likelihood, which serves as a surrogate target for maximization. Variational Bayes, however, introduces a per sample optimization subroutine to find the variational proposal distribution that best matches the true posterior distribution (of the latent variable given an input observation). To amortize the cost of this optimization subroutine, the variational autoencoder introduces an amortized inference model that learns to predict the best proposal distribution given an input observation <cit.>.




Although the computational efficiency of amortized inference has enabled latent variable models to be trained at scale on large datasets <cit.>, amortization introduces an additional source of error in the approximation of the posterior distributions if the amortized inference model fails to predict the optimal proposal distribution. This additional source of error, referred to as the amortization gap <cit.>, causes variational autoencoder training to further deviate from maximum likelihood training <cit.>.

To improve

training, numerous methods have been developed to reduce the amortization gap. In this paper, we focus on a class of methods <cit.> that takes an initial proposal distribution predicted by the amortized inference model and refines this initial distribution with the application of Stochastic Variational Inference (SVI) <cit.>. Since SVI applies gradient ascent to iteratively update the proposal distribution, a by-product of this procedure is a trajectory of proposal distributions (q_0, …, q_k) and their corresponding importance weights (w_0, … w_k). The intermediate distributions are discarded, and only the last distribution q_k is retained for updating the generative model. Our key insight is that the intermediate importance weights can be repurposed to further improve training. Our contributions are as follows

    
  * We propose a new method, Buffered Stochastic Variational Inference (BSVI), that takes advantage of the intermediate importance weights and constructs a new lower bound (the BSVI bound).
    
  * We show that the BSVI bound is a special instance of a family of generalized importance-weighted lower bounds.
    
  * We show that training variational autoencoders with BSVI consistently outperforms SVI, demonstrating the effectiveness of leveraging the intermediate weights.

Our paper shows that BSVI is an attractive replacement of SVI with minimal development and computational overhead. 



§ BACKGROUND AND NOTATION


We consider a latent-variable generative model p_θ(x, z)  where x ∈ is observed, z ∈ is latent, and θ are the model's parameters. The marginal likelihood p_θ(x) is intractable but can be lower bounded by the evidence lower bound (ELBO)

    ln p_θ(x) ≥_q(z)lnp_θ(x, z)/q(z) = _q(z)ln w(z),

which holds for any distribution q(z). Since the gap of this bound is exactly the Kullback-Leibler divergence q(z)p_θ(z  x), q(z) is thus the variational approximation of the posterior.

Furthermore, by viewing q as a proposal distribution in an importance sampler, we refer to w(z)=p_θ(x, z)/q(z) as an unnormalized importance weight. Since w(z) is a random variable, the variance can be reduced by averaging the importance weights derived from i.i.d samples from q(z). This yields the Importance-Weighted Autonenocder (IWAE) bound <cit.>,

    ln p_θ(x) ≥_z_1… z_k i.i.d.∼ qln1/k∑_i=1^k w(z_i)≥,

which admits a tighter lower bound than the ELBO <cit.>.



 §.§ Stochastic Variational Inference


The generative model can be trained by jointly optimizing q and θ to maximize the lower bound over the data distribution p̂(x). Supposing the variational family = q(z λ)_λ∈Λ is parametric and indexed by the parameter space Λ (e.g. a Gaussian variational family indexed by mean and covariance parameters),



the optimization problem becomes

    max_θ_p̂(x)max_λ_q(z λ)ln w(z λ, θ) .

where importance weight w is now

    w(z λ, θ) = p_θ(x, z)/q(z λ).

For notational simplicity, we omit the dependency on x. For a fixed choice of θ and x, <cit.> proposed to optimize λ via gradient ascent, where one initializes with λ_0 and takes successive steps of

    λ_i + 1λ_i + η∇_λ_i,


for which the ELBO gradient with respect to λ_i can be approximated via Monte Carlo sampling as

    ∇_λ_i≈1/m∑_j=1^m ∇_λ_iln w(z_λ_i(_i^(j)) λ_i, θ)

where z_i^(j) = z_λ_i(_i^(j)) ∼ q(z λ_i) is reparameterized as a function of λ_i and a base distribution p_0(ϵ). We note that k applications gradient ascent generates a trajectory of variational parameters (λ_0, …, λ_k), where we use the final parameter λ_k for the approximation. Following the convention in <cit.>, we refer to this procedure as Stochastic Variational Inference (SVI).



 §.§ Amortized Inference Suboptimality

The SVI procedure introduces an inference subroutine that optimizes the proposal distribution q(z λ) per sample, which is computationally costly. <cit.> observed that the computational cost of inference can be amortized by introducing an inference model f_ϕ: →Λ, parameterized by ϕ, that directly seeks to learn the mapping x ↦λ^* from each sample x to an optimal λ^* that solves the maximization problem

    λ^* = _λ_q(z λ)lnp_θ(x, z)/q(z λ).

This yields the amortized ELBO optimization problem

    max_θ, ϕ_p̂(x)_q(z  f_ϕ(x))lnp_θ(x, z)/q(z  f_ϕ(x)),

where q(z  f_ϕ(x)) can be concisely rewritten (with a slight abuse of notation) as q_ϕ(z  x) to yield the standard variational autoencoder objective <cit.>.

While computationally efficient, the influence of the amortized inference model on the training dynamics of the generative model has recently come under scrutiny <cit.>. A notable consequence of amortization is the amortization gap

    q_ϕ(z  x)p_θ(z  x) - q(z λ^*)p_θ(z  x)

which measures the additional error incurred when the amortized inference model is used instead of the optimal λ^* for approximating the posterior <cit.>.
A large amortization gap can present a potential source of concern since it  introduces further deviation from the maximum likelihood objective <cit.>.




 §.§ Amortization-SVI Hybrids


To close the amortization gap, <cit.> proposed to blend amortized inference with SVI. Since SVI requires one to initialize λ_0, a natural solution is to set λ_0 = f_ϕ(x). Thus, SVI is allowed to fine-tune the initial proposal distribution found by the amortized inference model and reduce the amortization gap. Rather than optimizing θ, ϕ jointly with the amortized ELBO objective <ref>, the training of the inference and generative models is now decoupled; ϕ is trained to optimize the amortized ELBO objective, but θ is trained to approximately optimize <ref>, where λ^* ≈λ_k is approximated via SVI. To enable end-to-end training of the inference and generative models, <cit.> proposed to backpropagate through the SVI steps via a finite-difference estimation of the necessary Hessian-vector products. Alternatively, <cit.> adopts a learning-to-learn framework where an inference model iteratively outputs λ_i+1 as a function of λ_i and the ELBO gradient.





§ BUFFERED STOCHASTIC VARIATIONAL INFERENCE



In this paper, we focus on the simpler, decoupled training procedure described by <cit.> and identify a new way of improving the SVI training procedure (orthogonal to the end-to-end approaches in <cit.>). Our key observation is that, as part of the gradient ascent estimation in <ref>, the SVI procedure necessarily generates a sequence of importance weights (w_0, …, w_k), where w_i = w(z_i λ_i, θ). 

Since (ln w_k) likely achieves the highest ELBO, the intermediate weights (w_0, …, w_k-1) are subsequently discarded in the SVI training procedure, and only ∇_θln w_k is retained for updating the generative model parameters. However, if the preceding proposal distributions (q_k-1, q_k-2, …) are also reasonable approximations of the posterior, then it is potentially wasteful to discard their corresponding importance weights.

A natural question to ask then is whether the full trajectory of weights (w_0, …, w_k) can be leveraged to further improve the training of the generative model. 







Taking inspiration from IWAE's weight-averaging mechanism, we propose a modification to the SVI procedure where we simply keep a buffer of the entire importance weight trajectory and use an average of the importance weights ∑_i π_i w_i as the objective in training the generative model.[For simplicity, we use the uniform-weighting π_i = 1/(k+1) in our base implementation of BSVI. In <ref>, we discuss how to optimize π during training.] The generative model is then updated with the gradient ∇_θln∑_i π_i w_i. We call this procedure Buffered Stochastic Variational Inference (BSVI) and denote ln∑_i π_i w_i as the BSVI objective. We describe the BSVI training procedure in <ref> and contrast it with SVI training. For notational simplicity, we shall always imply initialization with an amortized inference model when referring to SVI and BSVI.









§ THEORETICAL ANALYSIS


An important consideration is whether the BSVI objective serves as a valid lower bound to the log-likelihood ln p_θ(x). A critical challenge in the analysis of the BSVI objective is that the trajectory of variational parameters (λ_0, …, λ_k) is actually a sequence of statistically-dependent random variables. This statistical dependency is a consequence of SVI's stochastic gradient approximation in <ref>. We capture this dependency structure in <Ref>, which shows that each λ_i+1 is only deterministically generated after z_i is sampled. When the proposal distribution parameters λ_0:k are marginalized, the resulting graphical model is a joint distribution over q(z_0:k x). To reason about such a joint distribution, we introduce the following generalization of the IWAE bound.

theoremThmGIWAE

Let p(x,z) be a distribution where z∈𝒵. Consider a joint proposal distribution q(z_0:k) over 𝒵^k. 
Let v(i)⊂0, …, k∖i for all i, and π be a categorical distribution over 0, …, k. The following construction, which we denote the Generalized IWAE Bound, is a valid lower bound of the log-marginal-likelihood



    _q(z_0:k)ln∑_i=0^k π_i p(x, z_i)/q(z_i  z_v(i))≤ln p(x),





The proof follows directly from the linearity of expectation when using q(z_0:k) for importance-sampling to construct an unbiased estimate of p_θ(x), followed by application of Jensen's inequality. A detailed proof is provided in <ref>. 



Notably, if q(z_0:k) = ∏_i q(z_i), then <ref> reduces to the IWAE bound. <ref> thus provides a generalization of IWAE, where the samples drawn are potentially non-independently and non-identically distributed. 

<Ref> thus provides a way to construct new lower bounds on the log-likelihood whenever one has access to a set of non-independent samples. 

In this paper, we focus on a special instance where a chain of samples is constructed from the SVI trajectory. We note that the BSVI objective can be expressed as

    _q(z_0:k x)ln∑_i=0^k π_i w_i = 
        _q(z_0:k x)ln∑_i=0^k π_i p_θ(x, z_i)/q(z_i  z_<i, x).

Note that since λ_i can be deterministically computed given (x, z_<i), it is therefore admissible to interchange the distributions q(z_i  z_<i, x) = q(z_i λ_i). The BSVI objective is thus a special case of the Generalized IWAE bound, where z_v(i) = z_<i with auxiliary conditioning on x. Hence, the BSVI objective is a valid lower bound of ln p_θ(x); we now refer to it as the BSVI bound where appropriate.

In the following two subsections, we address two additional aspects of the BSVI bound. First, we propose a method for ensuring that the BSVI bound is tighter than the Evidence Lower Bound achievable via SVI. Second, we provide an initial characterization of BSVI’s implicit sampling-importance-resampling distribution.






























 §.§ Buffer Weight Optimization


Stochastic variational inference uses a series of gradient ascent steps to generate a final proposal distribution q(z λ_k). As evident from <Ref>, the parameter λ_k is in fact a random variable. The ELBO achieved via SVI, in expectation, is thus

    _q(z, λ_k  x)lnp_θ(x, z)/q_ϕ(z λ_k)= _q(z_0:k x)ln w_k,

where the RHS re-expresses it in notation consistent with <ref>. We denote <ref> as the SVI bound. In general, the BSVI bound with uniform-weighting π_i = 1 / (k + 1) is not necessarily tighter than the SVI bound. For example, if SVI's last proposal distribution exactly matches posterior q_k(z) = p_θ(z  x), then assigning equal weighting to across (w_0, … w_k) would make the BSVI bound looser. 


In practice, we observe the BSVI bound with uniform-weighting to consistently achieve a tighter lower bound than SVI's last proposal distribution. We attribute this phenomenon to the effectiveness of variance-reduction from averaging multiple importance weights—even when these importance weights are generated from dependent and non-identical proposal distributions. 

To guarantee that the BSVI is tighter than the SVI bound, we propose to optimize the buffer weight π. This guarantees a tighter bound,

    max_π_q(z_0:k x)ln∑_i=0^k π_i w_i ≥_q(z_0:k x)ln w_k,

since the SVI bound is itself a special case of the BSVI bound when π = (0, …, 0, 1). It is worth noting that <ref> is concave with respect to π, allowing for easy optimization of π.

Although π is a local variational parameter, we shall, for simplicity, optimize only a single global π that we update with gradient ascent throughout the course of training. As such, π is jointly optimized with θ and  ϕ.



 §.§ Dependence-Breaking via Double-Sampling





As observed in <cit.>, taking the gradient of the log-likelihood with respect to θ results in the expression

    ∇_θln p_θ(x) = _p_θ(z  x)∇_θln p_θ(x, z).

We note that gradient of the ELBO with respect to θ results in a similar expression

    ∇_θELBO(x) = _q_ϕ(z  x)∇_θln p_θ(x, z).

As such, the ELBO gradient differs from log-likelihood gradient only in terms of the distribution applied by the expectation operator. To approximate the log-likelihood gradient, we wish to set q_ϕ(z  x) close to p_θ(z  x) under some divergence.

We now show what results from computing the gradient of the BSVI objective.
lemmaLemBSVIGrad
The BSVI gradient with θ is

    ∇_θBSVI(x)    = _q_(z  x)∇_θln p_θ(x, z),

where q_ is a sampling-importance-resampling procedure defined by the generative process

    z_0:k   ∼ q(z_0:k x)
    
        i    ∼ r(i  z_0:k) 
    
        z     z_i,

and r(i  z_0:k) = (π_i w_i)/(∑_j π_j w_j) is a probability mass function over 0, …, k. 

A detailed proof is provided in <ref>.

A natural question to ask is whether BSVI's q_ is closer to the posterior than q_k in expectation. To assist in this analysis, we first characterize a particular instance of the Generalized IWAE bound when (z_1, …, z_k) are independent but non-identically distributed.

theoremThmIndGIWAE
When q(z_0:k) = ∏_i q_i(z_i), the implicit distribution q_(z) admits the inequality

    _q_(z)lnp_θ(x, z)/q_(z)≥_q(z_0:k)ln∑_i=0^k π_i w_i 
    
    = _q(z_0:k)ln∑_i=0^k π_i p_θ(x, z)/q_i(z_i).





<Ref> extends the analysis by <cit.> from the i.i.d. case (i.e. the standard IWAE bound) to the non-identical case (proof in <ref>). It remains an open question whether the inequality holds for the non-independent case. 


Since the BSVI objective employs dependent samples, it does not fulfill the conditions for <ref>. To address this issue, we propose a variant, BSVI with double-sampling (BSVI-DS), that breaks dependency by drawing two samples at each SVI step: ẑ_i for computing the SVI gradient update and z_i for computing the BSVI importance weight w_i. 


The BSVI-DS bound is thus

    _q(ẑ_<k x)_q(z_0:kẑ_<k, x)ln∑_i=0^k π_i p_θ(x, z)/q(z_i ẑ_<k, x),

where q(z_0:kẑ_<k, x) = ∏_i q(z_i ẑ_<k, x) is a product of independent but non-identical distributions when conditioned on (ẑ_<k, x). Double-sampling now allows us to make the following comparison.

corollaryCorBSVIDS
Let q_k = q(z_k ẑ_<i, x) denote the proposal distribution found by SVI. For any choice of (ẑ_<i, x), the distribution q_ implied by BSVI-DS (with optimal weighting π^*) is at least as close to p_θ(z  x) as q_k,

    q_p_θ(z  x)≤q_kp_θ(z  x),

as measured by the Kullback-Leibler divergence.


<ref> follows from <ref> and that the BSVI-DS bound under optimal π^* is no worse than the SVI bound. Although the double-sampling procedure seems necessary for inequality in <ref> to hold, in practice we do not observe any appreciable difference between BSVI and BSVI-DS.



§ COMPUTATIONAL CONSIDERATIONS


Another important consideration is the speed of training the generative model with BSVI versus SVI. Since BSVI reuses the trajectory of weights (w_0, …, w_k) generated by SVI, the forward pass incurs the same cost. The backwards pass for BSVI, however, is O(k) 
for k SVI steps—in contrast to SVI's O(1) cost. To make the cost of BSVI's backwards pass O(1), we highlight a similar observation from the original IWAE study <cit.> that the gradient can be approximated via Monte Carlo sampling

    ∇_θBSVI(x) ≈1/m∑_i=1^m ∇_θln p_θ(x, z^(i)),

where z^(i) is sampled from BSVI's implicit distribution q_(z  x). We denote this as training BSVI with sample-importance-resampling (BSVI-SIR). Setting m = 1 allows variational autoencoder training with BSVI-SIR to have the same wall-clock speed as training with SVI.



§ EXPERIMENTS





























































































































































 §.§ Setup











We evaluated the performance of our method by training variational autoencoders with BSVI-SIR with buffer weight optimization (BSVI-SIR-π)) on the dynamically-binarized Omniglot, grayscale SVHN datasets, and FashionMNIST (a complete evaluation of all BSVI variants is available in <ref>).
Our main comparison is against the SVI training procedure (as described in <ref>). We also show the performance of the standard VAE and IWAE training procedures. Importantly, we note that 
we have chosen to compare SVI-k' and IWAE-k' trained with k' = 10 against BSVI-k-SIR trained with k = 9 SVI steps. This is because that BSVI-k-SIR generates k + 1 importance weights.

For all our experiments, we use the same architecture as <cit.> (where the decoder is a PixelCNN) and train with the AMSGrad optimizer <cit.>.

For grayscale SVHN, we follow <cit.> and replaced <cit.>'s bernoulli observation model with a discretized logistic distribution model with a global scale parameter. Each model was trained for up to 200k steps with early-stopping based on validation set performance. For the Omniglot experiment, we followed the training procedure in <cit.> and annealed the KL term multiplier <cit.> during the first 5000 iterations.
We replicated all experiments four times and report the mean and standard deviation of all relevant metrics. For additional details, refer to <ref>



 §.§ Log-Likelihood Performance













For all models, we report the log-likelihood (as measured by BSVI-500). We additionally report the SVI-500 (ELBO*) bound along with its decomposition into rate (KL*) and distortion (Reconstruction*) components <cit.>. We highlight that KL* provides a fair comparison of the rate achieved by each model without concern of misrepresentation caused by the amortized inference suboptimality.

Omniglot. <ref> shows that BSVI-SIR outperforms SVI on the test set log-likelihood. BSVI-SIR also makes greater usage of the latent space (as measured by the lower Reconstruction*). 


Interestingly, BSVI-SIR's log-likelihoods are noticeably higher than its corresponding ELBO*, suggesting that BSVI-SIR has learned posterior distributions not easily approximated by the Gaussian variational family when trained on Omniglot.





SVHN. <ref> shows that BSVI-SIR outperforms SVI on test set log-likelihood. We observe that both BSVI-SIR and SVI significantly outperform both VAE and IWAE on log-likelihood, ELBO*, and Reconstruction*, demonstrating the efficacy of iteratively refining the proposal distributions found by amortized inference model during training.

FashionMNIST. <ref> similarly show that BSVI-SIR outperforms SVI on test set log-likelihood. Here, BSVI achieves significantly better Reconstruction* as well as achieving higher ELBO* compared to VAE, IWAE, and SVI.




In <ref> (<ref>), we also observe that the use of double sampling and buffer weight optimization does not make an appreciable difference than their appropriate counterparts, demonstrating the efficacy of BSVI even when the samples (z_0:k) are statistically dependent and the buffer weight is simply uniform.








 §.§ Stochastic Gradient as Regularizer






Interestingly, <ref> shows that BSVI-SIR can outperform BSVI on the test set despite having a higher variance gradient. We show in <Ref> that this is the result of BSVI overfitting the training set. The results demonstrate the regularizing effect of having noisier gradients and thus provide informative empirical evidence to the on-going discussion about the relationship between generalization and the gradient signal-to-noise ratio in variational autoencoders <cit.>.





 §.§ Latent Space Visualization








<ref> shows that the model learned by BSVI-SIR training has better Reconstruction* than SVI, indicating greater usage of the latent variable for encoding information about the input image. We provide a visualization of the difference in latent space usage in <Ref>. Here, we sample multiple images conditioned on a fixed z. Since BSVI encoded more information into z than SVI on the Omniglot dataset, we see that the conditional distribution p_θ(x  z) of the model learned by BSVI has lower entropy (i.e. less diverse) than SVI.









 §.§ Analysis of Training Metrics
























Recall that the BSVI-k training procedure runs SVI-k as a subroutine, and therefore generates the trajectory of importance weights (w_0, …, w_k). Note that ln w_0 and ln w_k are unbiased estimates of the ELBO achieved by the proposal distribution q_0 (SVI-0 bound) and q_k (SVI-k bound) respectively. It is thus possible to monitor the health of the BSVI training procedure by checking whether the bounds adhere to the ordering

    BSVI-k ≥SVI-k ≥SVI-0

in expectation. <Ref> show that this is indeed the case. Since Omniglot was trained with KL-annealing <cit.>, we see in <Ref> that SVI plays a negligible role once the warm-up phase (first 5000 iterations) is over. In contrast, SVI plays an increasingly large role when training on the more complex SVHN and FashionMNIST datasets, demonstrating that the amortization gap is a significantly bigger issue in the generative modeling of SVHN and FashionMNIST. <Ref> further shows that BSVI-k consistently achieves a better bound than SVI-k. When the buffer weight is also optimized, we see in <Ref> that π learns to upweight the later proposal distributions in (q_0, …, q_k), as measured by the buffer weight average _π(i) i/k. For SVHN, the significant improvement of SVI-k over SVI-0 results in π being biased significantly toward the later proposal distributions. Interestingly, although <Ref> suggests that the optimal buffer weight π^* can differ significantly from naive uniform-weighting, we see from <ref> that buffer weight optimization has a negligible effect on the overall model performance.



§ CONCLUSION


In this paper, we proposed Buffered Stochastic Variational Inference (BSVI), a novel way to leverage the intermediate importance weights generated by stochastic variational inference. We showed that BSVI is effective at alleviating inference suboptimality and that training variational autoencoders with BSVI consistently outperforms its SVI counterpart, making BSVI an attractive and simple drop-in replacement for models that employ SVI. One promising line of future work is to extend the BSVI training procedure with end-to-end learning approaches in <cit.>. Additionally, we showed that BSVI procedure is a valid lower bound and belongs to general class of importance-weighted (Generalized IWAE) bounds where the importance weights are statistically dependent.
Thus, it would be of interest to study the implications of this bound for certain MCMC procedures such as Annealed Importance Sampling <cit.> and others.



  §.§.§ Acknowledgements


We would like to thank Matthew D. Hoffman for his insightful comments and discussions during this project. This research was supported by NSF (#1651565, #1522054, #1733686), ONR (N00014-19-1-2145), AFOSR (FA9550-19-1-0024), and FLI.


unsrt



§ PROOFS


*
To show the validity of this lower bound, note that

    _q(z_0, …, z_k)∑_i π_i p_θ(x, z_i)/q(z_i  z_v(i))   = ∑_i π_i _q(z_0, …, z_k)p_θ(x, z_i)/q(z_i  z_v(i))
       = ∑_i π_i _q(z_v(i))_q(z_i  z_v(i))p_θ(x, z_i)/q(z_i  z_v(i))
       = ∑_i π_i _q(z_v(i)) p_θ(x)
       = p_θ(x).

Applying Jensen's inequality shows that the lower bound in the theorem is valid.


*


    ∇_θBSVI(x)    = _q(z_0:k x)∇_θln∑_i=0^k π_i w_i
       = _q(z_0:k x)_r(i  z_0:k)∇_θln p_θ(x, z_i),

The double-expectation can now be reinterpreted as
the sampling-importance-resampling distribution q_.


*

Recall that the q_ is defined by the following sampling procedure

    (z_0, …, z_k)    ∼ q(z_0, …, z_k)
    
    i    ∼ r(i  z_0:k)
    
    z     z_i,

where 

    r(i  z_0:k) = π_i w_i/∑_j π_j w_j = π_i p(x, z_i)/q(z_i  z_<i)/∑_j π_j p(x, z_i)/q(z_j  z_<j)






























We first note that, for any distribution r(z)

    r(z) = ∫_a r(a) δ_z(a) ạ = _r(a)δ_z(a).

This provides an intuitive way of constructing the probability density function by reframing it as a sampling process (the expectation w.r.t. r(a)) paired with a filtering procedure (the dirac-delta δ_z(a)). Thus, the density under q_ is thus

    q_(z) = _q_(z_0:k)_r(i  z_0:k)δ_z(z_i).

Additionally, we shall introduce the following terms 

    (z)    = p_θ(x, z)
    _i    = π_i w_i/∑_j π_j w_j
    _i    = w_i/∑_j π_j w_j
    _i(z)    = w(z)/π_i w(z) + ∑_-iπ_j w_j.

for notational simplicity. Note that the density function q_ can be re-expressed as

    q_(z) 
       = _q(z_0:k)_r(i   z_0:k)δ_z(z_i)
       = _q(z_0:k)∑_i π_i _i δ_z(z_i)
       = _π(i)_q_z_-i_q_i(z_i  z_-i)_i δ_z(z_i)
       = _π(i)_q_z_-i_i(z) q_i(z  z_-i).

We now begin with the ELBO under q_(z) and proceed from there via

    _q_(z) ln(z)/q_(z)   = -q_(z)(z)
       = -_π(i)_q_z_-i_i(z) q_i(z  z_-i)(z)
       ≥ -_π(i)_q_z_-i_i(z) q_i(z  z_-i)(z),

where we use Jensen's Inequality to exploit the convexity of the unnormalized Kullback-Leibler divergence ··. We now do a small change of notation when rewriting the unnormalized KL as an integral to keep the notation simple

    _q_(z) ln(z)/q_(z)   ≥_π(i)_q_z_-i∫_z_i_i q(z_i  z_-i) ln(z_i)/_i q(z_i  z_-i)
       = _π(i)_q_z_-i_q(z_i  z_-i)_i ln(z_i)/_i q(z_i  z_-i)
       = _q(z_0:k)∑_i _i ln(z_i)/_i q(z_i  z_-i)
       = _q(z_0:k)∑_i _i ln∑_j π_j w_j 
    ·q(z_i  z_<i)/q(z_i  z_-i)
       = _q(z_0:k)∑_i _i ln∑_j π_j w_j 
    + lnq(z_i  z_<i)/q(z_i  z_-i)

If z_0:k are independent, then it follows that q(z_i  z_<i) = q(z_i  z_-i) = q(z_i). Thus,

    _q_(z) ln(z)/q_(z)   ≥_q(z_0:k)∑_i _i ln∑_j π_j w_j 
        + lnq(z_i  z_<i)/q(z_i  z_-i)
       = _q(z_0:k)∑_i _i ln∑_j π_j w_j
       = _q(z_0:k)ln∑_j π_j w_j.





§ MODEL PERFORMANCE ON TEST AND TRAINING DATA


Here we report various performance metrics for each type of model trained on the training set for both Omniglot and SVHN. As stated earlier, log-likelihood is estimated using BSVI-500, and ELBO* refers to the lower bound achieved by SVI-500 (i.e. z ∼ q_500).
KL* and Reconstruction* are the rate and distortion terms for ELBO*, respectively.

    Log-likelihood   = _q(z_0:500 x)ln∑_i=0^500π_i p_θ(x, z_i)/q(z_i  z_<i, x)
    ELBO*   = _q_500ln p_θ(x  z)_Reconstruction*
        + q_500(z)p_θ(z)_KL*



















§ LOG-LIKELIHOOD ESTIMATION USING BSVI AND IWAE




A popular way to approximate the true log-likelihood is to use the IWAE-k bound with a sufficiently large k during evaluation time <cit.>. Here we compare log-likelihood estimates of BSVI and IWAE in Tables <ref> and <ref> and empirically show that BSVI bounds are as tight as IWAE bounds in all of our experiments. This justifies the use of BSVI-500 for estimating log-likelihood in our reports.










§ EXPERIMENT SETUP


Here we describe our detailed experiment setup. For both Omniglot and SVHN experiments, we used a ResNet with three hidden layers of size 64 as the encoder and a 12-layer gated PixelCNN with the constant layer size of 32 as the decoder. Network parameters (ϕ,θ) were trained with the AMSGrad optimizer <cit.>. For SVI, we followed the experimental setup of <cit.> and optimized local variational parameters λ_0:k with SGD with momentum with learning rate 1.0 and momentum 0.5. To stabilize training, we applied gradient clipping to both network parameters and local variational parameters. Each model was trained for 200k steps with early-stopping based on validation loss. The best-performing models on the validation set were then evaluated on the test set. All experiments were performed four times, and we reported the mean and standard deviation of relevant metrics.

Omniglot. We used 2000 randomly-selected training images as the validation set. Each digit was dynamically binarized at training time based on the pixel intensity. We used 32-dimensional latent variable with unit Gaussian prior. Each pixel value was modeled as a Bernoulli random variable where the output of the decoder was interpreted as log probabilities. We also followed the training procedure in <cit.> and annealed the KL term multiplier <cit.> from 0.1 to 1.0 during the first 5000 iterations of training.

SVHN. We merged “train” and “extra” data in the original SVHN dataset to create our training set. We again reserved 2000 randomly-selected images as the validation set. To reuse the network architecture for the Omniglot dataset with minimal modifications, we gray-scaled all images and rescaled the pixel intensities to be in [0,1].  The only differences from Omniglot experiments are: increased latent variable dimensions (64), larger image size (32 × 32), and the use of discretized logistic distribution by <cit.> with a global scale parameter for each pixel. Similar to <cit.>, we lower-bound the scale parameter by a small positive value.

FashionMNIST. Similar to above, we used 2000 randomly-selected training images as the validation set. The network architecture and hyperparameters were identical to those of SVHN dataset, except we used 32-dimensional latent variables and did not employ KL term annealing.










Below is the list hyperparameters used in our experiments. Since we have two stochastic optimization processes (one for the model and one for SVI), we employed separate gradient clipping norms.




