S EMI -I MPLICIT V ARIATIONAL I NFERENCE VIA S CORE M ATCHING
A BSTRACT
Semi-implicit variational inference (SIVI) greatly enriches the expressiveness of variational families by considering implicit variational distributions deﬁned in a hierarchical manner.
However, due to the intractable densities of variational distributions, current SIVI approaches often use surrogate evidence lower bounds (ELBOs) or employ expensive inner-loop MCMC runs for direct ELBO maximization for training.
In this paper, we propose SIVI-SM, a new method for SIVI based on an alternative training objective via score matching.
Leveraging the hierarchical structure of semi-implicit variational families, the score matching objective allows a minimax formulation where the intractable variational densities can be naturally handled with denoising score matching
We show that SIVI-SM closely matches the accuracy of MCMC and outperforms ELBO-based SIVI methods in a variety of Bayesian inference tasks.
I NTRODUCTION
Variational inference(VI) is an approximate Bayesian inference approach where the inference problem is transformed into an optimization problem (Jordan et al., 1999; Wainwright & Jordan, 2008;
Blei et al., 2017).
It starts by introducing a family of variational distributions over the model parameters (or latent variables) to approximate the posterior.
The goal then is to ﬁnd the closest member from this family of distributions to the target posterior, where the closeness is usually measured by the Kullback-Leibler (KL) divergence from the posterior to the variational approximation.
In practice, this is often achieved by maximizing the evidence lower bound (ELBO), which is equivalent to minimizing the KL divergence (Jordan et al., 1999).
One of the classical VI methods is mean-ﬁeld VI (Bishop & Tipping, 2000), where the variational distributions are assumed to be factorized over the parameters (or latent variables).
When combined with conditional conjugacy, this often leads to simple optimization schemes with closed-form update rules (Blei et al., 2017).
While popular, the factorizable assumption and conjugacy condition greatly restrict the ﬂexibility and applicability of variational posteriors, especially for complicated models with high dimensional parameter space.
Recent years have witnessed much progress in the ﬁeld of VI that extends it to more complicated settings.
For example, the conjugacy condition has been removed by the black-box VI methods which allow a broad class of models via Monte carlo gradient estimators (Nott et al., 2012;
Paisley et al., 2012; Ranganath et al., 2014;
Rezende et al., 2014;
Kingma & Welling, 2014).
On the other hand, more ﬂexible variational families have been proposed that either explicitly incorporate more complicated structures among the parameters (Jaakkola &
Jordan, 1998; Saul & Jordan, 1996; Giordano et al., 2015; Tran et al., 2015) or borrow ideas from invertible transformation of probability distributions (Rezende & Mohamed, 2015; Dinh et al., 2017;
Kingma et al., 2016; Papamakarios et al., 2019).
All these methods require tractable densities for the variational distributions.
It turns out that the variational family can be further expanded by allowing implicit models that have intractable densities but are easy to sample from (Husz´ar, 2017).
One way to construct these implicit models is to transform a simple base distribution via a deterministic map, i.e., a deep neural network (Tran et al., 2017;
Mescheder et al., 2017; Shi et al., 2018a;b; Song et al., 2019).
Due to the intractable densities of implicit models, when evaluating the ELBO during training, one often resorts to density ratio estimation which is known to be challenging 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 surrogate ELBOs (asymptotically unbiased) are employed for training (Yin & Zhou, 2018; Moens et al., 2021).
Instead of surrogate ELBOs, an unbiased gradient estimator of the exact ELBO has been derived based on MCMC samples from a reverse conditional (Titsias & Ruiz, 2019).
However, the computation for the inner-loop MCMC runs can easily become expensive in high-dimensional regimes.
There are also approaches that estimate the gradients instead of the objective (Li & Turner, 2018; Shi et al., 2018b; Song et al., 2019). Besides KL divergence, score-based distance measures have also been introduced in various statistical tasks (Hyv¨arinen, 2005; Zhang et al., 2018) and have shown advantages in complicated nonlinear models (Song & Ermon, 2019;
Ding et al., 2019; Elkhalil et al., 2021).
Recently, there are also some studies that use score matching for variational inference (Yang et al., 2019;
Hu et al., 2018).
However, these methods are not designed for SIVI and hence either do not apply to or can not fully exploit the hierarchical structure of semi-implicit variational distributions.
In this paper, we propose
SIVI-SM, a new method for SIVI using an alternative training objective via score matching.
We show that the score matching objective and the semi-implicit hierarchical construction of variational posteriors can be combined in a minimax formulation where the intractability of densities is naturally handled with denoising score matching.
We demonstrate the effectiveness and efﬁciency of our method on both synthetic distributions and a variety of real data Bayesian inference tasks.
B ACKGROUND
Semi-Implicit Variational Inference Semi-implicit variational inference (SIVI) (Yin & Zhou, 2018) posits a ﬂexible variational family deﬁned hierarchically using a mixing parameter as follows
[Equation]
where ϕ = { φ, ξ } are the variational parameters.
This variational distribution is called semi-implicit as the conditional layer q φ ( x | z ) is required to be explicit but the mixing layer q ξ ( z ) can be implicit, and q ϕ ( x ) is often implicit unless q ξ ( z ) is conjugate to q φ ( x | z ) .
Compared to standard VI, the above hierarchical construction allows a much richer variational family that is able to capture complicated dependencies between parameters (Yin & Zhou, 2018).
Similar to standard VI, current SIVI methods ﬁt the model parameters by maximizing the evidence lower bound (ELBO) derived as follows
[Equation]
where D is the observed data.
As q ϕ ( x ) is no longer tractable, Yin & Zhou (2018) considered a sequence of lower bounds of ELBO
[Equation]
Note that L L is an asymptotically exact surrogate ELBO as L → ∞ .
An increasing sequence of { L t } ∞ t =1 , therefore, is often suggested, with L L t being optimized at the t - th iteration.
Instead of maximizing surrogate ELBOs, Titsias & Ruiz (2019) proposed unbiased implicit variational inference (UIVI) which is based on an unbiased gradient estimator of the exact ELBO.
More speciﬁcally, consider a ﬁxed mixing distribution q ξ ( z )
=
q ( z ) and a reparameterizable conditional q φ ( x | z ) such that x = T φ ( z , (cid:15) ) , (cid:15) ∼ q (cid:15) ( (cid:15) ) ⇔ x ∼ q φ ( x | z ) , then [Equation] where
[Equation]
The gradient term in 4 involves an expectation w.r.t.
the reverse conditional q φ ( z | x ) which can be estimated using an MCMC sampler (e.g., Hamiltonian Monte Carlo (Neal, 2011)).
However, the inner-loop MCMC runs can easily become computationally expensive in high dimensional regimes.
See a more detailed discussion on the derivation and computation issues of UIVI in Appendix A and
D.
Score Matching Score matching is ﬁrst introduced by Hyv¨arinen (2005) to learn un-normalized statistical models given i.i.d.
samples from an unknown data distribution p ( x ) .
Instead of estimating p ( x ) directly, score matching trains a score network S ( x ) to estimate the score of the data distribution, i.e. ∇ log p ( x ) , by minimizing the score matching objective E p ( x ) [ 12 (cid:107) S [Equation] x [Equation]
− ∇ x log p
[Equation] x ) (cid:107) 22 ] . Using the trick of partial integration, Hyv¨arinen (2005) shows that the score matching objective is equivalent to the following up to a constant
[Equation]
The expectation in Eq.
5 can be quickly estimated using data samples.
However, it is often challenging to scale up score matching to high dimensional data due to the computation of Tr ∇ x ( S ( x )) .
A commonly used variant of score matching that can scale up to high dimensional data is denoising score matching (DSM) (Vincent, 2011).
The ﬁrst step of DSM is to perturb the data with a known noise distribution q σ (˜ x | x ) , which leads to a perturbed data distribution q σ (˜ x ) =
(cid:82) q σ (˜ x
| x ) p ( x ) d x .
The score matching objective for q σ (˜ x ) turns out to be equivalent to the following up to a constant
[Equation]
Unlike Eq.
5, Eq.
6 does not involve the trace term and can be computed efﬁciently, as long as the score of the noise distribution ∇ ˜ x log q σ (˜ x | x ) is easy to compute.
Note that the optimal score network here estimates the score of the perturbed data distribution rather than that of the true data distribution.
A small noise, therefore, is required for accurate approximation of the true data score ∇ log p ( x ) .
Despite this, DSM is widely used in learning energy based models (Saremi et al., 2018) and score based generative models (Song & Ermon, 2019).
P ROPOSED M ETHOD
While ELBO-based training objectives prove effective for semi-implicit variational inference, current approaches either rely on surrogates of the exact ELBO or expensive inner-loop MCMC runs for unbiased gradient estimates due to the intractable variational posteriors.
In this section, we introduce an alternative training objective for SIVI based on score matching.
We show that the score matching objective can be reformulated in a minimax fashion such that the semi-implicit hierarchical construction of variational posteriors can be efﬁciently exploited using denoising score matching.
Throughout this section, we assume the conditional layer q φ ( x | z ) to be reparameterizable and its score function ∇ x log q φ ( x | z ) is easy to evaluate 1 .
3. A M INIMAX R EFORMULATION
Rather than maximizing the ELBO as in previous semi-implicit variational inference methods, we can instead minimize the following Fisher divergence that compares the score functions of the target and the variational distribution [Equation]
[Equation]
Here S ( x ) = ∇ x log p ( x | D ) =
∇ x log p ( D, x ) is the score of the target posterior distribution, and the variational distribution q ϕ ( x ) is deﬁned in Eq.
Due to the semi-implicit construction in Eq. 1, the score of variational distribution, i.e. ∇ x log q ϕ ( x ) , is intractable, making the Fisher divergence in Eq.
7 not readily computable.
Although the hierarchical structure of q ϕ ( x ) allows us to estimate its score function via denoising score matching, the estimated score function would break the dependency on the variational parameter ϕ , leading to biased gradient estimates (see an illustration in Appendix L).
Fortunately, this issue can be remedied by reformulating 7 as a minimax problem.
The key observation is that the squared norm of S ( x ) − ∇ x log q ϕ ( x ) can be viewed as the maximum value of the following nested optimization problem
[Equation]
where f ( x ) is an arbitrary function of x , and the unique optimal solution is
[Equation]
Based on this observation, we can rewrite the optimization problem in 7 as
[Equation]
Now we can take advantage of the hierarchical structure of q ϕ ( x ) to get ride of the intractable score term ∇ x log q ϕ ( x ) in Eq.
8, similarly as done in DSM.
More speciﬁcally, note that
[Equation]
We have
[Equation]
Substituting Eq.
9 into Eq.
8 completes our reformulation.
Theorem 1.
Assume the variational distribution q ϕ ( x ) is a semi-implicit distribution deﬁned by
Eq. [Equation] then the optimization problem in 7 is equivalent to the following minimax problem
[Equation]
Where ϕ = { φ, ξ } .
Moreover, assume that f can represent any function.
If ( ϕ ∗ , f ∗ ) deﬁnes a
Nash-equilibrium of Eq.
[Equation] then f ∗ , ϕ ∗ is given by
[Equation]
See a detailed proof of Theorem 1 in Appendix B. Note that the objective in equation 8 can also be derived via Stein discrepancy minimization (Liu et al., 2016; Gorham & Mackey, 2015; Ranganath et al., 2016; Grathwohl et al., 2020).
However, our reformulation in equation 10 takes a further step by utilizing the hierarchical structure of q ϕ ( x ) and hence can easily scale up to high dimensions.
See Appendix K for a more detailed discussion.
3. P RACTICAL A LGORITHMS
In practice, we parameterize f with a neural network f ψ ( x ) .
According to the above minimax reformulation, the Monte Carlo estimation of the objective function in Eq.
10 can be easily obtained by sampling from the hierarchical variational distribution
z ∼ q ξ ( z ) , x ∼ q φ ( x | z ) .
Furthermore, using the reparameterization trick (Kingma & Welling, 2013), i.e. x = T φ ( z ; (cid:15) ) , z = h ξ ( γ ) , where (cid:15) ∼ q (cid:15) ( (cid:15) ) , γ ∼ q γ ( γ ) , we can rewrite Eq.
10 as follows
[Equation]
Input: Score
Output: Variational parameter ϕ and the neural network parameter ψ .
Initialize ϕ 0 , ψfor t = 0 to N − 1 do
Sample { γ (1) , γ (2) , · · · , γ ( m ) } from prior q γ ( γ ) and let z ( i ) =
h ξ ( γ ( i ) ) , i = 1 , . . .
, m .
Sample { (cid:15) (1) , (cid:15) (2) , · · · , (cid:15) ( m ) } from N (0 , I ) .
Compute x ( i ) =
µ φ ( z ( i ) )
+ σ φ ( z ( i ) ) (cid:12) (cid:15) ( i ) .
Update ϕ by descending its stochastic gradient: ∇ ϕ 1 m m (cid:88) i
=
f ψ ( x ( i ) )
T [ S ( x ( i ) )
+ σ φ ( z ( i ) ) − 1 (cid:12) (cid:15) ( i ) ]
− 1 2 (cid:107) f ψ ( x ( i ) )
(cid:107) 22 .
for j = 1 to K do
Sample { γ (1) , γ (2) , · · · , γ ( m ) } from prior q γ ( γ ) and let z ( i ) =
h ξ ( γ ( i ) ) , i = 1 , . . .
, m .
Sample { (cid:15) (1) , (cid:15) (2) , · · · , (cid:15) ( m ) } from N (0 , I ) .
Compute x ( i ) =
µ φ ( z ( i ) )
+ σ φ ( z ( i ) ) (cid:12) (cid:15) ( i ) .
Update ψ by ascending its stochastic gradient: ∇ ψ 1 m m (cid:88) i
=
f ψ ( x ( i ) )
T [ S ( x ( i ) )
+ σ φ ( z ( i ) ) − 1 (cid:12) (cid:15) ( i ) ] − 1 2 (cid:107) f ψ ( x ( i ) ) (cid:107) 22 .
end for end for
This allows us to directly optimize the parameters ϕ, ψ with gradient based optimization methods (Goodfellow et al., 2014).
Optimizing f to completion in the inner loop of training is computational prohibitive.
Therefore, we alternate between K steps of optimizing f and one step of optimizing q ϕ .
As mentioned before, we use the multivariate Gaussian distribution with diagonal covariance matrix for the conditional q φ ( x | z ) ∼ N ( µ φ ( z ) , diag { σ 2 φ ( z ) } ) , which can be reparameterized as follows
[Equation]
where (cid:15) ∼ N (0 , I ) , where (cid:12) means the element-wise product. The corresponding score function is ∇ x log q φ ( x | z ) = − σ φ ( z ) − 1 (cid:12) (cid:15) . The complete procedure of SIVI-SM is formally presented in Algorithm 1. 3. T HEORETICAL R ESULTS R EGARDING N EURAL N ETWORK A PPROXIMATION
The inexact lower-level optimization for neural networks introduces approximation errors.
Also, neural networks themselves may introduce approximation gaps due to their approximation capacities.
These numerical errors may affect the approximation quality of variational distribution q ϕ ( x ) , which we analyze in the following proposition.
Proposition 1.
Let Ω be the feasible domain of ϕ .
∀ ϕ ∈ Ω , we say that f ˆ ψ ( ϕ ) is (cid:15)
-accurate, if
[Equation]
where
[Equation]
Let ϕ ∗ be one of the optimal variational parameters deﬁned in Eq.
11 and ˜ ϕ be the one obtained using neural network approximation deﬁned as follows with f ˆ ψ ( ϕ ) being (cid:15) -accurate
[Equation]
Then we have
[Equation]
[Equation] See a detailed proof of Proposition 1 in Appendix C. From proposition 1, we see that the approximation error of our numerical solution to the minimax problem in 10 can be controlled by two terms.
The ﬁrst term E x ∼ q ϕ ∗ ( x ) (cid:107) S ( x ) −∇ x log q ϕ ∗ ( x ) (cid:107) 22 measures the approximation ability of the variational distribution, and the second term measures the approximation/optimization error of the neural network f ψ ( x ) .
As long as the approximation error of f ψ ( x ) is small, the minimax formulation in 10 can provide variational posteriors with similar approximation accuracy to those of the original problem in 7 in terms of Fisher divergence to the target posterior.
Remark.
When the variational parameter ϕ is ﬁxed, the lower-level optimization problem on f ψ ( x ) is equivalent to the following
[Equation]
and we use this objective as a measure of approximation accuracy of f ψ ( x ) given ϕ in Eq.
E XPERIMENTS
In this section, we compare SIVI-SM to ELBO-based methods including the original SIVI and
UIVI on a range of inference tasks.
We ﬁrst show the effectiveness of our method and illustrate the role of the auxiliary network approximation f ψ on several two-dimensional toy examples.
The KL divergence from the target distributions to different variational approximations was also provided for direct comparison.
We also compare the performance of SIVI-SM with both baseline methods on several Bayesian inference tasks, including a multidimensional Bayesian logistic regression problem and a high dimensional Bayesian multinomial logistic regression problem.
Following Titsias & Ruiz (2019), we set the conditional layer to be q φ ( x | z ) =
N ( x | µ φ ( z ) , diag ( σ )) 2 and ﬁx the mixing layer as q ( z )
=
N (0 , I ) .
The variational parameters therefore are ϕ = { φ, σ } .
All experiments were implemented in Pytorch (Paszke et al., 2019).
If not otherwise speciﬁed, we use the Adam optimizer for training (Kingma & Ba, 2014).
4. T OY EXAMPLES
We ﬁrst apply SIVI-SM to approximate three synthetic distributions deﬁned on a two-dimensional space: a banana-shaped distribution, a multimodal Gaussian, and an X-shaped mixture of Gaussian.
The densities of these distributions are given in Table 2 in Appendix E.
For the convenience of comparison, we used the same conﬁguration of semi-implicit distribution family as in UIVI (Titsias & Ruiz, 2019).
The µ φ ( z ) is a multilayer perceptron (MLP) with layer widths [3 , 50 , 50 , 2] .
The network approximation f ψ ( x ) is parameterized by a 4 layers MLP with layer widths [2 , 128 , 128 , 2] .
For SIVI-SM, we set the number of inner-loop gradient steps K = 1 .
For SIVI, we set L = 50 for the surrogate ELBO deﬁned in Eq. 2.
For UIVI, we used 10 iterations for every inner-loop HMC sampling.
To facilitate exploration, for all methods, we used the annealing trick (Rezende & Mohamed, 2015) during training for the multimodal and X-shaped Gaussian distributions.
Variational approximations from all methods were obtained after 50,000 variational parameter updates.
Figure 5 in Appendix F shows the contour plots of the synthetic distributions, together withsamples from the trained variational distributions.
We see that SIVI-SM produces samples that match the target distributions well.
We also report the KL divergence from the target distributions to the variational posteriors (estimated via the ITE package (Szab´o, 2014) using 100,000 samples from each distribution) given by different methods in Table 3 in Appendix G.
We see that SIVI-SM performs better for more challenging target distributions.
To better understand the role the network approximation f ψ ( x ) played during the training process, we visualize its training dynamics on the X-shaped distribution in Figure 1.
We see that during training, f ψ ( x ) automatically detected where the current approximation is insufﬁcient and guided the variational posterior towards these areas.
Note the Nash-equilibrium of f ∗ ( x ) in Eq.
11 is the difference between the score functions of the target distribution p ( x ) and q ϕ ( x ) .
As the variational posterior gets closer to the target, the signal provided by f ψ ( x ) becomes weaker and would converge to zero in the perfect case.
More details on the convergence of f ψ ( x ) can be found in Appendix H.
B AYESIAN L OGISTIC R EGRESSION
Our second example is a Bayesian logistic regression problem where the log-likelihood function takes the following form
[Equation]
Here x i are covariates, and y i ∈ { 0 , 1 } are binary response variables.
Following Yin & Zhou (2018), we set the prior as β ∼ N ( 0 , α − 1 I ) , where α = 0 .
We consider the waveform 3 dataset where the dimension of x i is 21 which leads to a parameter space of 22 dimensions.
We used a standard 10-dimensional Gaussian prior for the q ( z ) .
For µ φ ( z ) , we used a 4 layer MLP with layer widths [10 , 100 , 100 , 22] .
The network approximation f ψ ( β ) is also a 4 layer MLP with layer width [22 , 256 , 256 , 22] .
Similarly as in section 4.1, we set the number of inner-loop gradient steps K =in SIVI-SM.
For SIVI, we set L = 100 and used the same training method as in Yin & Zhou (2018).
For UIVI, we set the length of inner-loop HMC iterations to be 10 with the ﬁrst 5 iterations discarded as burn-in, with 5 leapfrog steps in each iteration.
The results of all methods were collected after 20,000 variational parameter updates.
We collected 1000 samples of β to represent the approximated posterior distributions for all three
SIVI variants.
The ground truth was formed from a long MCMC run of 400,000 iterations using parallel stochastic gradient Langevin dynamics (SGLD) (
Welling & Teh, 2011) with 1000 independent particles, and a small stepsize of 10 − 4 .
Figure 2 shows the posterior estimates provided by different SIVI variants in contrast to the ground truth MCMC results.
We see that SIVI and UIVI tend to slightly underestimate the variance for both univariate marginal and pairwise joint posteriors (especially for β 4 , β 5 ), while SIVI-SM agreed with MCMC well.
Furthermore, we also examined the covariance estimates of β and the results were presented in Figure 3.
We see that SIVI-SM provides the best overall approximation to the posterior which achieved the smallest rooted mean square error (RMSE) to the ground truth at 0.0184.
4. B AYESIAN M ULTINOMIAL L OGISTIC R EGRESSION
Our next example is a Bayesian multinomial logistic regression problem.
For a data set of N covariate and label pairs { ( x i , y i ) :
i = 1 , . .
. , N } , where y i ∈ { 1 , . .
, R } , the categorical likelihood is
[Equation]
where β = ( β T 1 , β T 2 , · · · , β TR ) T is the model parameter and follows a standard Gaussian prior.
Following Titsias & Ruiz (2019), we used two data sets: MNIST 4 and HAPT 5 .
MNIST is a commonly used dataset in machine learning that contains 60,000 training and 10,000 test instances of 28 ×images of hand-written digits which has R = 10 classes.
HAPT (Reyes-Ortiz et al., 2016) is a human activity recognition dataset.
It contains 7,767 training and 3,162 test data points, and each one of them contains features of 561-dimensional measurements captured by inertial sensors, which correspond to R = 12 classes of static postures, dynamic activities and postural transitions.
The dimensions of the posterior distributions are 7,850 (MNIST) and 6,744 (HAPT) respectively.
We used the same variational family as before, with a 100-dimensional standard Gaussian prior for q ( z ) .
We used MLPs with two hidden layers for the mean network µ φ ( z ) of the Gaussian conditional and the network approximation f ψ ( β ) , with 200 hidden neurons for µ φ ( z ) and 256 hidden neurons for f ψ ( β ) for each of the hidden layers respectively.
We used the same initialization of variational parameters for all methods.
Following Titsias & Ruiz (2019), we used a minibatch size of 2,000 for MNIST and 863 for HAPT.
As before, we set the number of inner-loop gradient steps
K = 1 in SIVI-SM.
For SIVI, we set L = 200 as previously done by Titsias & Ruiz (2019).
For
UIVI, we set the number of inner-loop HMC iterations to be 10 and discarded the ﬁrst 5 iterations as burn-in, with 5 leapfrog steps in each iteration.
As done in UIVI, we used the RMSProp optimizer (Tieleman & Hinton, 2012) for training.
We used different batch sizes during training to investigate its effect on the quality of variational approximations for different methods.
These batch sizes were selected in such a way that the corresponding computational times are comparable between different methods.
See a detailed comparison on the computation times in Appendix I (the experiments were run on a RTX2080 GPU).
As the gradient computation for the inner-loop HMC sampling required by UIVI is not scalable 6 , the batch size for UIVI is set as m = 1 which was also used by Titsias & Ruiz (2019).
For all methods, the results were collected after 90,000 variational parameter updates for MNIST and 40,000 variational parameter updates for HAPT.
Figure 4 shows the predictive log-likelihood on the test data as a function of the number of iterations for both data sets, where the estimates were formed based on 8,000 samples from the variational
The number in parentheses speciﬁes the batch sizes used for training distributions ﬁtted using different methods, as done in UIVI (Titsias & Ruiz, 2019).
Although
SIVI-SM converges slower at the beginning due to the slow training of the network approximation f ψ ( β ) , it eventually surpasses the other ELBO-based variants and achieves better prediction on both datasets.
Compared to ELBO-based methods, SIVI-SM would beneﬁt more from large batch sizes.
4. B AYESIAN N EURAL N ETWORKS
Lastly, we compare our method with SIVI, UIVI and SGLD on sampling the posterior of Bayesian neural network on the UCI datasets.
We conduct the two-layer network with 50 hidden units and
ReLU activation function.
The datasets are all randomly partitioned into 90% for training and 10% for testing.
We use the variational family as before, with 3-dimensional standard Gaussian prior for q ( z ) , and 10 hidden neurons for µ φ ( z ) and 16 hidden neurons for f ψ ( x ) .
The results are averaged over 10 random trials.
We refer the reader to Appendix J for hyper-parameter tuning and other experiment details.
Table 1 shows the average test RMSE and NLL and their standard deviation.
We see that SIVI-SM can achieve on par or better results than SIVI and UIVI.
Although SGLD performs better for some datasets, it requires a long run to generate samples.
[Table]
C ONCLUSION
We proposed SIVI-SM, a new method for semi-implicit variational inference based on an alternative training objective via score matching.
Unlike the ELBO-based objectives, we showed that the score matching objective allows a minimax formulation where the hierarchical structure of semi-implicit variational families can be more efﬁciently exploited as the corresponding intractable variational densities can be naturally handled with denoising score matching.
In experiments, we demonstrated that SIVI-SM closely matches the accuracy of MCMC in posterior estimation and outperforms two typical ELBO-based methods (SIVI and UIVI) in a variety of Bayesian inference tasks. A CKNOWLEDGMENTS
This work was supported by National Natural Science Foundation of China (grant no. 12201014).
The research of Cheng Zhang was support in part by the Key Laboratory of Mathematics and Its Applications (LMAM) and the Key Laboratory of Mathematical Economics and Quantitative Finance (LMEQF) of Peking University.
The authors are grateful for the computational resources provided by the Megvii institute.
The authors appreciate the anonymous ICLR reviewers for their constructive feedback.
A D ERIVATION OF E Q . 2-
The gradient of ELBO is that
[Equation]
Note that the property of margin score function,
[Equation]
Then the gradient g entφ can be representd as
[Equation]
B P ROOF OF T HEOREM
Proof.
As discussed in section 3.1, by introducing the vector-valued function f
[Equation]
we can rewrite the optimization objective in Eq.
7 as
[Equation]
Compute the score of the semi-implicit distribution q ϕ ( x ) deﬁned in Eq. [Equation]
we have
[Equation]
Bring the above score of q ϕ ( x ) into Eq.
16, we have
[Equation]
Therefor, let ϕ = { ξ, φ } , we can rewrite the original score matching problem Eq.
7 as
[Equation]
If ( ϕ ∗ , f ∗ ) deﬁnes a Nash-equilibrium of the above problem, ﬁxing the parameters ϕ
= ϕ ∗ , the optimal vector-valued function
f ∗ ( x ) is invariant in the derivation of Eq.
So we can easily deduce f ∗ by Eq.
[Equation]
Bring f ∗ ( x ) into Eq.
10, we have the unbiased approximation of ϕ ∗
[Equation]
C P ROOF OF P ROPOSITION
Proof.
Consider the score matching problem with the well (cid:15) - trained f ˆ ψ ( ϕ )
[Equation]
we have
[Equation]
Therefore
[Equation]
we can estimate the upper bound of Fisher divergence between S ( x ) and ∇ x log q ˜ ϕ
( x )
[Equation]
, (19) where ϕ ∗ := arg min ϕ ∈ Ω { E x ∼ q ϕ ( x ) (cid:107) S ( x ) − ∇ x log q ϕ ( x ) (cid:107) 22 } .
And the last inequality is due to the fact that f ˆ ψ ( ϕ ) is well (cid:15) - trained and (cid:107)
ˆ R ϕ ∗ ( x ) (cid:107) 22 is non-negative.
D C OMPUTATIONAL ISSUES ON UIVI
Unlike SIVI that samples from the prior q ( z ) , UIVI samples from the posterior distribution q ( z | x ) , which can provide unbiased gradient estimate to the exact ELBO (Titsias & Ruiz, 2019).
However,
UIVI requires computing the gradient of log q ( z | x ) during the iterations of HMC sampling procedures.
If UIVI uses a minibatch of m data points x 1 , x 2 , · · · , x m in the training process, it needs to compute the Jacobian matrix [ ∇ z log q ( z | x 1 ) , ∇ z log q ( z | x 2 ) , · · · , ∇ z log q ( z | x m )] , which is not scalable for automatic differentiation using backpropagation.
Therefore, we set the batch size m = 1 for UIVI as done in Titsias & Ruiz (2019).
E D ENSITIES FOR THE TOY EXAMPLES
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
F SIVI SAMPLES ON THE TOY EXAMPLES
G KL DIVERGENCE RESULTS ON THE TOY EXAMPLES
[Table]
H C ONVERGENCE PERFORMANCE OF SIVI-SM
Here, we demonstrate the convergence behavior of SIVI-SM in our experiments.
For the topy examples in section 4. [Table] we use 500 samples from q ϕ ( x ) to form the Monte Carlo estimates of the loss ( SM loss ) in Eq. [Table]
and the L 2 -norm E q ϕ ( x ) (cid:107) f ψ ( x ) (cid:107) 2 2 of the f ψ function( fnet’s norm ) during the training process.
Figure 6 shows the estimated SM loss and fnet’s norm as a function of the number of iterations for the three synthetic toy distributions
[Table]
Similarly, Figure 7 shows the convergence traces of the SM loss and fnet’s norm in the experiments in section 4.3.
Note that although the dimensions of the posterior distributions are high, i [Table] e. 7850 for
MNIST and 6744 for HAPT, the corresponding fnet’s norms can be quite low ( 58 . 510 for MNIST and 8 . 267 for HAPT), indicating the effectiveness of our methods even in high dimensional spaces.
I S ECONDS PER ITERATION IN F IGURE
The following table shows the run times of different methods per iteration on a RTX2080 GPU.
We see that the run times for SIVI and SIVI-SM are comparable with the chosen pairs of batch sizes (i.e., 10 vs 100 and 20 vs 200).
As discussed before, the inner-loop HMC iterations make UIVI slower than other methods.
[Table]
J E XPERIMENT SETTING FOR B AYESIAN NEURAL NETWORKS
For SIVI, we set L = 100 the batch size is m = 10 in training process.
For UIVI, the setting of HMC inner loop is similar with section 4.3.
For SGLD, we choose the step size from { 10 − 4 , 10 − 5 , 10 − 6 } and iteration number in { 50000 , 100000 } by validation in training process with 100 particles.
For
SIVI-SM, we set inner-loop gradient steps K = 1 , 3 by validation and run 20,000 iterations for training.
K R ELATED M ETHODS
Consider a test functions class F , the Stein discrepancy (Gorham & Mackey, 2015) measure between p and q is deﬁned follows [Equation] [Equation]
This measure is based on the following Stein’s identity (Stein, 1972)
[Equation]
[Equation]
An early example in variational inference used Stein discrepancy is operator variational inference (OPVI)
(Ranganath et al., 2016), which constructs a variational operator (e.g. Langevin-Stein Op-
[Figure]
[Equation]
Then OPVI solves the minmax optimization problem simultaneously with q and f .
Unlike OPVI, learned Stein discrepancy (LSD) (Grathwohl et al., 2020) utilizes the L 2 regularization term to substitute the constraint of F
[Equation]
In fact, the variational objective of SIVI-SM in Eq.8 can be viewed as L LSD .
Bring Eq.21 into Eq.
and let λ = 12 , we have
[Equation]
However, OPVI and LSD both involve the Tr( ∇ x f ( x )) term which is not easy to compute for high dimensional problems.
Our method takes a further step by utilizing the hierarchical structure of q ( x ) in Eq.1.
Using a mathematical trick that is similar to denoising score matching, we arrive at a formulation that easily scales up to high dimensions.
L O N T HE B IASENESS OF THE MC G RADIENT E STIMATE OF THE F ISHER D IVERGENCE IN E Q .