Markov Chain Score Ascent: A Unifying Framework of Variational Inference with Markovian Gradients
Abstract
Minimizing the inclusive Kullback-Leibler (KL) divergence with stochastic gradient descent (SGD) is challenging since its gradient is defined as an integral over the posterior.
Recently, multiple methods have been proposed to run SGD with biased gradient estimates obtained from a Markov chain.
This paper provides the first non-asymptotic convergence analysis of these methods by establishing their mixing rate and gradient variance.
To do this, we demonstrate that these methods–which we collectively refer to as Markov chain score ascent (MCSA) methods–can be cast as special cases of the Markov chain gradient descent framework.
Furthermore, by leveraging this new understanding, we develop a novel MCSA scheme, parallel MCSA (pMCSA), that achieves a tighter bound on the gradient variance.
We demonstrate that this improved theoretical result translates to superior empirical performance.
Introduction
Bayesian inference aims to analyze the posterior distribution of an unknown latent variable 𝐳 from which data 𝐱 is observed.
By assuming a model 𝑝 (𝐱 | 𝐳) , the posterior 𝜋 (𝐳) is given by
Bayes’ rule such that 𝜋 (𝐳) ∝ 𝑝 (𝐱 | 𝐳) 𝑝 (𝐳) where 𝑝 (𝐳) represents our prior belief on 𝐳 .
Instead of working directly with 𝜋 , variational inference (VI, Blei et al. 2017) seeks a variational approximation 𝑞 (𝐳; 𝛌) ∈ 𝒬 , where 𝒬 is a variational family and 𝛌 are the variational parameters, that is the most similar to 𝜋 according to a discrepancy measure 𝑑 (𝜋, 𝑞 (⋅; 𝛌)) .
The apparent importance of choosing the right discrepancy measure has led to a quest spanning a decade (Bamler et al. , 2017;
Dieng et al. , 2017; Geffner & Domke, 2021b; Hernandez-Lobato et al. , 2016; Li & Turner, 2016; Regli & Silva, 2018; Ruiz & Titsias, 2019; Salimans et al. , 2015;
Wan et al. , 2020; Wang et al. , 2018; Zhang et al. , 2021).
So far, the exclusive (or reverse, backward) Kullback-Leibler (KL) divergence 𝑑 KL (𝑞 (⋅; 𝛌) ∥ 𝜋) has seen “exclusive” use, partly because it is defined as an integral over 𝑞 (𝐳; 𝛌) , which can be approximated efficiently.
In contrast, the inclusive (or forward) KL is defined as an integral over 𝜋 as
[Equation]
Since our goal is to approximate 𝜋 with 𝑞 (⋅; 𝛌) but the inclusive KL involves an integral over 𝜋 , we end up facing a chicken-and-egg problem.
Despite this challenge, the inclusive KL has consistently drawn attention due to its statistical properties, such as better uncertainty estimates due to its mass covering property (MacKay, 2001; Minka, 2005; Trippe & Turner, 2017).
Recently, Naesseth et al.
(2020); Ou & Song (2020) have respectively proposed Markovian score climbing (MSC) and joint stochastic approximation (JSA).
These methods minimize the inclusive KL using stochastic gradient descent (SGD, Robbins & Monro 1951), where the gradients are estimated using a Markov chain.
The Markov chain kernel 𝐾 𝛌 𝑡 (𝐳 𝑡 , ⋅) is 𝜋 -invariant (Robert & Casella, 2004) and is chosen such that it directly takes advantage of the current variational approximation 𝑞 (⋅; 𝛌 𝑡 ) .
Thus, the quality of the gradients improves over time as the KL divergence decreases.
Still, the gradients are non-asymptotically biased and Markovian across adjacent iterations, which sharply contrasts MSC and JSA from classical black-box VI (Kucukelbir et al. , 2017; Ranganath et al. , 2014), where the gradients are unbiased and independent.
While Naesseth et al.
(2020) have shown the convergence of MSC through the work of Gu & Kong (1998), this result is only asymptotic and does not provide practical insight into the performance of
MSC.
In this paper, we address these theoretical gaps by casting MSC and JSA into a general framework we call Markov chain score ascent (MCSA), which we show is a special case of Markov chain gradient descent (MCGD, Duchi et al. 2012).
This enables the application of the nonasymptotic convergence results of MCGD (Debavelaere et al. , 2021;
Doan et al. , 2020a,b; Duchi et al. , 2012; Karimi et al. , 2019; Sun et al. , 2018; Xiong et al. , 2021).
For MCGD methods, the fundamental properties affecting the convergence rate are the ergodic convergence rate ( 𝜌 ) of the MCMC kernel and the gradient variance ( 𝐺 ).
We analyze 𝜌 and 𝐺 of MSC and JSA, enabling their practical comparison given a fixed computational budget ( 𝑁 ).
Furthermore, based on the recent insight that the mixing rate does not affect the convergence rate of MCGD (Doan et al. , 2020a,b), we propose a novel scheme, parallel MCSA (pMCSA), which achieves lower variance by trading off the mixing rate.
We verify our theoretical analysis through numerical simulations and compare MSC, JSA, and pMCSA on general Bayesian inference problems.
Our experiments show that our proposed method outperforms previous MCSA approaches.
Contribution Summary
¶ We provide the first non-asymptotic theoretical analysis of two recently proposed inclusive KL minimization methods ( Section 4 ), MSC ( Theorems 1 and 2 ) and JSA ( Theorem 3 ).
• To do this, we show that both methods can be viewed as what we call “Markov chain score ascent” (MCSA) methods ( Section 3 ), which are a special case of MCGD ( Proposition 1 ).
‚ In light of this, we develop a novel MCSA method which we call parallel MCSA (pMCSA, Section 5 ) that achieves lower gradient variance ( Theorem 4 ).
„ We demonstrate that the improved theoretical performance of pMCSA translates to superior empirical performance across a variety of Bayesian inference tasks ( Section 6 ).
Background 2.1 Inclusive Kullback-Leibler Minimization with Stochastic Gradients VI with SGD
The goal of VI is to find the optimal variational parameters 𝛌 identifying 𝑞 (⋅; 𝛌) ∈ 𝒬 that minimizes some discrepancy measure 𝐷 (𝜋, 𝑞 (⋅; 𝛌)) .
A typical way to perform
VI is to use stochastic gradient descent (SGD, Robbins & Monro 1951), provided that the optimization objective provides unbiased gradient estimates 𝐠 (𝛌) such that we can repeat the update
[Equation]
where 𝛾 1 , … , 𝛾 𝑇 is a stepsize schedule.
Inclusive KL Minimization with SGD For inclusive KL minimization, 𝐠 should be set as
[Equation]
where ℍ [𝜋, 𝑞 (⋅; 𝛌)] is the cross-entropy between 𝜋 and 𝑞 (⋅; 𝛌) , which shows the connection with cross-entropy methods (de Boer et al. , 2005), and 𝐬 (𝛌; 𝐳)
= ∇ 𝛌 log 𝑞 (𝐳; 𝛌) is known as the score gradient .
Since inclusive KL minimization with SGD is equivalent to ascending towards the direction of the score, Naesseth et al.
(2020) coined the term score climbing.
To better conform with the optimization literature, we instead call this approach score ascent as in gradient ascent.
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
2.2 Markov Chain Gradient Descent
Overview of MCGD Markov chain gradient descent (MCGD, Duchi et al.
2012; Sun et al.
2018) is a family of algorithms that minimize a function 𝑓 defined as 𝑓 (𝛌) = ∫ 𝑓 (𝛌, 𝛈) Π (𝑑𝛈) , where 𝛈 is random noise, and Π (𝑑𝛈) is its probability measure.
MCGD repeats the steps
[Equation]
(1) where 𝑃 𝛌 𝑡−1 is a Π -invariant Markov chain kernel that may depend on 𝛌 𝑡−1 .
The noise of the gradient is Markovian and non-asymptotically biased, departing from vanilla SGD.
Non-asymptotic convergence of this general algorithm has recently started to gather attention as by Debavelaere et al.
; Doan et al.
(2020a,b); Duchi et al.
Karimi et al.
; Sun et al.
Applications of MCGD MCGD encompasses an extensive range of problems, including distributed optimization (Ram et al. , 2009), reinforcement learning (Doan et al. , 2020a;
Tadić &
Doucet, 2017; Xiong et al. , 2021), and expectation-minimization (Karimi et al. , 2019), to name a few.
This paper extends this list with inclusive KL VI through the MCSA framework.
Markov Chain Score Ascent
First, we develop Markov chain score ascent (MCSA), a framework for inclusive KL minimization with MCGD.
This framework will establish the connection between MSC/JSA and MCGD.
3.1 Markov Chain Score Ascent as a
Special Case of Markov Chain Gradient Descent
As shown in Equation (1), the basic ingredients of MCGD are the target function 𝑓 (𝛌, 𝜂) , the gradient estimator 𝐠 (𝛌, 𝜂) , and the Markov chain kernel 𝑃 𝛌 (𝜂, ⋅) .
Obtaining MCSA from MCGD boils down to designing 𝐠 and 𝑃 𝛌 such that 𝑓 (𝛌) =
𝑑 KL (𝜋 ∥ 𝑞 (⋅; 𝛌)) .
The following proposition provides sufficient conditions on 𝐠 and 𝑃 𝛌 to achieve this goal.
Proposition 1.
Let 𝛈 =
(𝐳 (1) , 𝐳 (2) , … , 𝐳 (𝑁) ) and a Markov chain kernel 𝑃 𝛌 (𝛈, ⋅) be Π -invariant where Π is defined as
[Equation]
Then, by defining the objective function 𝑓 and the gradient estimator 𝐠 to be
[Equation]
where ℍ [ 𝜋 ] is the entropy of 𝜋 , MCGD results in inclusive KL minimization as
[Equation]
and 𝔼 Π [ 𝐠 (𝛌, 𝝶) ]
= ∇ 𝛌 𝑑 KL (𝜋 ∥ 𝑞 (⋅; 𝛌)).
Proof.
See the full proof in page 19.
This simple connection between MCGD and VI paves the way toward the non-asymptotic analysis of JSA and MSC.
Note that 𝑁 here can be regarded as the computational budget of each
MCGD iteration since the cost of (i) generating the Markov chain samples 𝐳 (1) , … , 𝐳 (𝑁) and (ii) computing the gradient 𝐠 will linearly increase with 𝑁 .
In addition, the MCGD framework often assumes 𝑃 to be geometrically ergodic.
An exception is the analysis of Debavelaere et al.
(2021) where they work with polynomially ergodic kernels.
Assumption 1 (Uniform Geometric Ergodicity) .
The Markov chain kernel 𝑃 is geometrically ergodic as
[Equation]
for some positive constant 𝐶 .
3.2 Non-Asymptotic Convergence of Markov Chain Score Ascent
Non-Asymptotic Convergence Through Proposition 1, Assumption 1 and some technical assumptions on the objective function, we can apply the existing convergence results of MCGD to MCSA.
Table 1 provides a list of relevant results.
Apart from properties of the objective function (such as Lipschitz smoothness), the convergence rates are stated in terms of the gradient bound 𝐺 , kernel mixing rate 𝜌 , and the number of MCGD iterations 𝑇 .
We focus on 𝐺 and 𝜌 as they are closely related to the design choices of different MCSA algorithms.
Convergence and the Mixing Rate 𝜌 Duchi et al.
(2012) was the first to provide an analysis of the general MCGD setting.
Their convergence rate is dependent on the mixing rate through the 1/ log 𝜌 −1 term.
For MCSA, this result is overly conservative since, on challenging problems, mixing can be slow such that 𝜌 ≈ 1 .
Fortunately, Doan et al.
(2020a,b) have recently shown that it is possible to obtain a rate independent of the mixing rate 𝜌 .
For example, in the result of Doan et al.
(2020b), the influence of 𝜌 decreases in a rate of 𝒪 ( 1 / 𝑇 2 ) .
This observation is critical since it implies that trading a “slower mixing rate” for “lower gradient variance” could be profitable .
We exploit this observation in our novel MCSA scheme in Section 5.
Gradient Bound 𝐺
Except for Doan et al.
(2020b), most results assume that the gradient is bounded for ∀𝛈, 𝛌 as ‖𝐠(𝛌,𝛈)‖ < 𝐺 .
Admittedly, this condition is strong, but it is similar to the bounded variance assumption 𝔼[‖𝐠‖ 2 ] <
𝐺 2 used in vanilla SGD, which is also known to be strong as it contradicts strong convexity (Nguyen et al. , 2018).
Nonetheless, assuming 𝐺 can have practical benefits beyond theoretical settings.
For example, Geffner & Domke (2020) use 𝐺 to compare the performance different VI gradient estimators.
In a similar spirit, we will obtain the gradient bound 𝐺 of different MCSA algorithms and compare their theoretical performance.
Demystifying Prior Markov Chain Score Ascent Methods
In this section, we will show that MSC and JSA both qualify as MCSA methods.
Furthermore, we establish (i) the mixing rate of their implicitly defined kernel 𝑃 and (ii) the upper bound on their gradient variance.
This will provide insight into their practical non-asymptotic performance.
4.1 Technical Assumptions
To cast previous methods into MCSA, we need some technical assumptions.
Assumption 2 (Bounded importance weight) .
The importance weight ratio 𝑤 (𝐳)
= 𝜋 (𝐳)
/𝑞
(𝐳; 𝛌) is bounded by some finite constant as 𝑤 ∗ < ∞ for all 𝛌 ∈ Λ such that 𝑟
= (1 − 1/𝑤 ∗ ) < 1 .
This assumption is necessary to ensure Assumption 1, and can be practically ensured by using a variational family with heavy tails (Domke & Sheldon, 2018) or using a defensive mixture (Hesterberg, 1995; Holden et al. , 2009) as
[Equation]
where 0 < 𝛼 < 1 and 𝜈 (⋅) is a heavy tailed distribution such that sup 𝐳∈𝒵 𝜋 (𝐳) /𝜈
(𝐳) < ∞ .
Note that 𝑞 def.
is only used in the Markov chain kernels and 𝑞 (⋅; 𝛌) is still the output of the VI procedure.
While these tricks help escape slowly mixing regions, this benefit quickly vanishes as 𝜆 converges.
Therefore, ensuring Assumption 2 seems unnecessary in practice unless we absolutely care about ergodicity.
(Think of the adaptive MCMC setting for example.
Brofos et al.
2022; Holden et al.
Model (Variational Family) Misspecification and 𝑤 ∗
Note that 𝑤 ∗ is bounded below exponentially by the inclusive KL as shown in Proposition 2.
Therefore, 𝑤 ∗ will be large (i) in the initial steps of VI and (ii) under model (variational family) misspecification.
Assumption 3.
(Bounded Score)
The score gradient is bounded for ∀𝛌 ∈ Λ and ∀𝐳 ∈
𝒵 such that
‖ 𝐬 (𝛌; 𝐳) ‖ 2 ≤ 𝐿 for some finite constant 𝐿 > 0 .
Although this assumption is strong, it enables us to compare the gradient variance of MCSA methods.
We empirically justify the bounds obtained using Assumption 3 in Section 6.2.
4.2 Markovian Score
MSC (Algorithm 4 in Appendix B) is a simple instance of MCSA where 𝛈 𝑡 = 𝐳 𝑡 and 𝑃 𝛌 𝑡 =
𝐾 𝛌 𝑡 is the conditional importance sampling (CIS) kernel (originally proposed by Andrieu et al.
where the proposals are generated from 𝑞 (⋅; 𝛌 𝑡 ) .
Although MSC uses only a single sample for the Markov chain, the CIS kernel internally operates 𝑁 − 1 proposals.
Therefore, 𝑁 in MSC has a different meaning, but it still indicates the computational budget.
Theorem 1.
MSC (Naesseth et al. , 2020) is obtained by defining
[Equation]
with 𝛈 𝑡 = 𝐳 𝑡 , where 𝐾 𝛌 (𝐳, ⋅) is the CIS kernel with 𝑞 def.
(⋅; 𝛌) as its proposal distribution.
Then, given Assumption 2 and 3, the mixing rate and the gradient bounds are given as
[Equation]
where 𝑤 ∗ = sup 𝐳 𝜋 (𝐳) /𝑞 def.
(𝐳; 𝛌) .
Proof.
See the full proof in page 21.
Discussion Theorem 1 shows that the gradient variance of MSC is insensitive to 𝑁 .
Although the mixing rate does improve with 𝑁 , when 𝑤 ∗ is large due to model misspecification and lack of convergence (see the discussion in Section 4.1), this will be marginal.
Overall, the performance of MSC cannot be improved by increasing the computational budget 𝑁 .
Rao-Blackwellization
Meanwhile, Naesseth et al.
also provide a Rao-Blackwellized version of MSC we denote as MSC-RB.
Instead of selecting a single 𝐳 𝑡 by resampling over the 𝑁 internal proposals, they suggest forming an importance-weighted estimate (Robert & Casella, 2004).
The theoretical properties of this estimator have been concurrently analyzed by Cardoso et al.
Theorem 2.
(Cardoso et al. , 2022)
The gradient variance of MSC-RB is bounded as
[Equation]
(𝜋 ∥ 𝑞 (⋅; 𝝺 𝑡−1 ))
+ 𝒪 (𝑁 −3/2 + 𝛾 𝑡−1 / 𝑁−1 )]
+ ‖𝛍‖ 22 , where 𝛍 = 𝔼 𝜋 𝐬
(𝛌; 𝘇) , 𝛾 = 2𝑤 ∗ / (2𝑤 ∗ + 𝑁 − 2) is the mixing rate of the Rao-Blackwellized CIS kernel, and 𝑑 𝜒 2 (𝜋 ∥ 𝑞) = ∫( 𝜋 / 𝑞 − 1) 2 𝑞(𝑑𝐳) is the 𝜒 2 divergence.
Proof.
See the full proof in page 22.
The variance of MSC-RB decreases as 𝒪 ( 1 / 𝑁−1 ) , which is more encouraging than vanilla MSC.
However, the first term depends on the 𝜒 2 divergence, which is bounded below exponentially by the KL (Agapiou et al. , 2017).
Therefore, the variance of MSC-RB will be large on challenging problems where the 𝜒 2 divergence is large , although linear variance reduction is possible.
4.3 Joint Stochastic Approximation
JSA (Algorithm 5 in Appendix B) was proposed for deep generative models where the likelihood factorizes into each datapoint.
Then, subsampling can be used through a random-scan version of the independent Metropolis-Hastings (IMH, Hastings 1970) kernel.
Instead, we consider the general version of JSA with a vanilla IMH kernel since it can be used for any type of likelihood.
At each MCGD step, JSA performs multiple Markov chain transitions and estimates the gradient by averaging all the intermediate states, which is closer to how traditional MCMC is used.
Independent Metropolis-Hastings Similarly to MSC, the IMH kernel in JSA generates proposals from 𝑞 (⋅; 𝛌 𝑡 ) .
To show the geometric ergodicity of the implicit kernel 𝑃 , we utilize the geometric convergence rate of IMH kernels provided by Mengersen & Tweedie (1996, Theorem 2.1) and Wang (2022).
The gradient variance, on the other hand, is difficult to analyze, especially the covariance between the samples.
However, we show that, even if we ignore the covariance terms, the variance reduction with respect to 𝑁 is severly limited in the large 𝑤 ∗ regime.
To do this, we use the exact 𝑛 -step marginal IMH kernel derived by Smith & Tierney (1996) as
[Equation]
𝐳, 𝑑𝐳 ′ )
=
𝑇 𝑛 ( 𝑤 (𝐳) ∨
𝑤
𝐳 ′ ) )
𝜋 (𝐳 ′ ) 𝑑𝐳 ′ + 𝜆 𝑛 (𝑤 (𝐳)) 𝛿 𝐳 (𝑑𝐳 ′ ) , (2) where 𝑤 (𝐳) =
𝜋 (𝐳) /𝑞
def.
(𝐳; 𝛌) , 𝑥 ∨ 𝑦 = max (𝑥, 𝑦) , and for 𝑅 (𝑣) =
{ 𝐳 ′ ∣ 𝑤 (𝐳 ′ ) ≤ 𝑣 } , 𝑇 𝑛 (𝑤) = ∫ ∞
𝑤 𝑛 𝑣 2 𝜆 𝑛−1 (𝑣) 𝑑𝑣 and 𝜆 (𝑤) =
∫ 𝑅(𝑤) (1 − 𝑤 (𝐳 ′ )
𝑤 ) 𝜋 (𝑑𝐳 ′ ) .
Theorem 3.
JSA (Ou & Song, 2020) is obtained by defining 𝑃 𝑘𝛌 (𝛈, 𝑑𝛈 ′ )
=
𝐾 𝑁 (
[Equation]
+1 𝛌
[Equation]
with 𝛈 𝑡 = (𝐳 (1)𝑡 , 𝐳 (2)𝑡 , … , 𝐳 (𝑁)𝑡 ) .
Then, given Assumption 2 and 3, the mixing rate and the gradient variance bounds are
[Equation]
where 𝛍 = 𝔼 𝜋 𝐬 (𝛌; 𝘇) , 𝐶 cov =
2 𝑁 2 ∑ 𝑁 𝑛=2 ∑ 𝑛−1 𝑚=1 Cov ( 𝐬 (𝛌; 𝘇 (𝑛)𝑡 ) , 𝐬 (𝛌; 𝘇 (𝑚)𝑡 ) ||
ℱ 𝑡−1 ) is the sum of the covariance between the samples, 𝑤 ∗ = sup 𝐳 𝜋 (𝐳)
/𝑞 def.
(𝐳; 𝛌) , and 𝐶 (𝑟, 𝑁) > 0 is a finite constant.
Proof.
See the full proof in page 24.
Discussion As shown in Theorem 3, JSA benefits from increasing 𝑁 in terms of a faster mixing rate.
However, under lack of convergence and model misspecification (large 𝑤 ∗ ), the variance improvement becomes marginal.
Specifically, in the large 𝑤 ∗ regime, the variance reduction is limited by the constant 1/2 term.
This is true even when, ideally, the covariance between the samples is ignorable such that 𝐶 cov ≈ 0 .
In practice, however, the covariance term 𝐶 cov will be positive, only increasing variance.
Therefore, in the large 𝑤 ∗ regime, JSA will perform poorly, and the variance reduction by increasing 𝑁 is fundamentally limited .
Parallel Markov Chain Score Ascent Algorithm 1: pMCSA
[Table]
)0 , … , 𝐳 (𝑁)0 , initial parameter 𝛌 0 , number of iterations 𝑇 , stepsize schedule 𝛾 𝑡 for 𝑡 = 1, 2, … , 𝑇 do for 𝑛 = 1, 2, … , 𝑁 do 𝐳 (𝑛)𝑡 ∼ 𝐾 𝛌 𝑡−1 (𝐳 (𝑛)𝑡−1 , ⋅) end 𝐠 (𝛌) =
− 1 𝑁 ∑ 𝑁 𝑛=1 𝐬 (𝛌; 𝐳 (𝑛)𝑡 )
𝛌 𝑡 = 𝛌 𝑡−1 − 𝛾 𝑡 𝐠 (𝛌 𝑡−1 ) end
Our analysis in Section 4 suggests that the statistical performance of MSC, MSC-RB, and JSA are heavily affected by model specification and the state of convergence through 𝑤 ∗ .
Furthermore, for JSA, a large 𝑤 ∗ abolishes our ability to counterbalance the inefficiency by increasing the computational budget 𝑁 .
However, 𝜌 and 𝐺 do not equally impact convergence; recent results on MCGD suggest that gradient variance is more critical than the mixing rate (see
Section 3.1).
We turn to leverage this understanding to overcome the limitations of previous methods.
5.1 Parallel Markov Chain Score Ascent
We propose a novel scheme, parallel Markov chain score ascent (pMCSA, Algorithm 1), that embraces a slower mixing rate in order to consistently achieve an 𝒪 ( 1 / 𝑁 ) variance reduction, even on challenging problems with a large 𝑤 ∗ ,
Algorithm Description Unlike JSA that uses 𝑁 sequential Markov chain states, pMCSA operates 𝑁 parallel Markov chains.
To maintain a similar per-iteration cost with JSA, it performs only a single Markov chain transition for each chain.
Since the chains are independent, the
Metropolis-Hastings rejections do not affect the variance of pMCSA.
Theorem 4.
pMCSA, our proposed scheme, is obtained by setting
[Equation]
with 𝛈 = (𝐳 (1) , 𝐳 (2) , … , 𝐳 (𝑁) ) .
Then, given Assumption 2 and 3, the mixing rate and the gradient variance bounds are
[Equation]
[Table]
where 𝛍 = 𝔼 𝜋 𝐬 (𝛌; 𝘇) , 𝑤 ∗ = sup 𝐳 𝜋 (𝐳)
/𝑞 def.
(𝐳; 𝛌) and 𝐶 (𝑁) > 0 is a finite constant.
Proof.
See the full proof in page 27.
Discussion Unlike JSA and MSC, the variance reduction rate of pMCSA is independent of 𝑤 ∗ .
Therefore, it should perform significantly better on challenging practical problems.
If we consider the rate of Duchi et al.
(2012), the combined rate is constant with respect to 𝑁 since it cancels out.
In practice, however, we observe that increasing 𝑁 accelerates convergence quite dramatically.
Therefore, the mixing rate independent convergence rates by Doan et al.
(2020a,b) appears to better reflect practical performance.
This is because (i) the mixing rate 𝜌 is a conservative global bound and (ii) the mixing rate will improve naturally as MCSA converges.
[Figure]
[Figure]
5.2 Computational Cost Comparison
[Table]
ELBO  𝑁 𝑁 𝑁
MSC 𝑁 −𝑁 𝑁 −
MSC-RB 𝑁 −𝑁 𝑁 − 𝑁
JSA 𝑁 𝑁 +𝑁 𝑁 pMCSA 𝑁 2 𝑁 𝑁 𝑁
The four schemes using the CIS and
IMH kernels have different costs depending on 𝑁 as organized in Table 2.
Cost of Sampling Proposals For the
CIS kernel used by MSC, 𝑁 controls the number of internal proposals sampled from 𝑞 (⋅; 𝛌) .
For JSA and pMCSA, the
IMH kernel only uses a single sample from 𝑞 (⋅; 𝛌) , but applies the kernel 𝑁 times.
On the other hand, pMCSA needs twice more evaluations of 𝑞 (⋅; 𝛌) .
However, this added cost is minimal since it is dominated by that of evaluating 𝑝 (𝐳, 𝐱) .
Cost of Estimating the Score When estimating the score, MSC computes ∇ 𝛌 log 𝑞 (𝐳; 𝛌) only once, while JSA and our proposed scheme compute it 𝑁 times.
However, MSC-RB also computes the score 𝑁 times.
Lastly, notice that MCSA methods do not differentiate through the likelihood 𝑝 (𝐳, 𝐱) , unlike ELBO maximization, making its per-iteration cost significantly cheaper.
Evaluations 6.1 Experimental Setup
Implementation For the realistic experiments, we implemented 4 MCSA methods on top of the Turing (Ge et al. , 2018) probabilistic programming framework.
For the variational family, we use diagonal multivariate Gaussians (mean-field family) with the support transformation of Kucukelbir et al.
We use the ADAM optimizer by Kingma & Ba (2015) with a stepsize of 0.01 in all experiments.
The budget is set to 𝑁 = 10 for all experiments unless specified.
Baselines We compare (i) pMCSA (ours, Section 5)
, (ii) JSA (Ou & Song, 2020), (iii) MSC (Naesseth et al. , 2020), (iv) MSC with with Rao-Blackwellization ( MSC-RB , Naesseth et al. 2020), and (v) evidence lower-bound maximization ( ELBO , Kucukelbir et al.
Ranganath et al. 2014) with the path derivative estimator (Roeder et al. , 2017).
[Table]
6.2 Simulations
Setup First, we verify our theoretical analysis on multivariate Gaussians with full-rank covariances sampled from Wishart distribution with 𝜈 degrees of freedom (values of 𝜈 are in the figure captions).
This problem is challenging since an IMH (used by pMCSA, JSA) or CIS (used by
MSC, MSC-RB) kernel with a diagonal Gaussian proposal will mix slowly due to a large 𝑤 ∗ .
Gradient Variance We evaluate our theoretical analysis of the gradient variance.
The variance is estimated from 512 independent Markov chains using the parameters generated by the main MCSA procedure.
The estimated variances are shown in Figure 1.
We make the following observations: (i) pMCSA has the lowest variance overall, and it consistently benefits from increasing 𝑁 .
(ii) MSC does not benefit from increasing 𝑁 whatsoever.
(iii) MSC-RB does not benefit much from increasing 𝑁 until the 𝜒 2 divergence between 𝜋 and 𝑞 (⋅; 𝛌) has become small.
(iv) JSA does not benefit from increasing 𝑁 until 𝑞 (⋅; 𝛌) has sufficiently converged (when 𝑤 ∗ has become small).
These results confirm our theoretical analyses in Sections 4 and 5.
Robustness Against Optimizers Since the convergence of most sophisticated SGD optimizers has yet to be established for MCGD, we empirically investigate their effectiveness.
The results using SGD (Bottou et al. , 2018; Robbins & Monro, 1951), Momentum (Polyak, 1964), Nesterov (Nesterov, 1983), ADAM (Kingma & Ba, 2015), and varying stepsizes are shown in Figure 2.
Clearly, pMCSA successfully converges for the broadest variety of optimizer settings.
Overall, most MCSA methods seem to be the most stable with ADAM, which points out that establishing the convergence of ADAM for MCGD will be a promising direction for future works.
6.3 Bayesian Neural Network Regression
[Figure]
Setup For realistic experiments, we train
Bayesian neural networks (BNN, Neal 1996) for regression.
We use datasets from the UCI repository (Dua & Graff, 2017) with 90% random train-test splits and run for 𝑇 = 5 ⋅ 10iterations.
We use the model and forward propagation method of Hernandez-Lobato &
Adams (2015) with a 50 -unit hidden layer (see
Appendix C.1).
Results The results are shown in Table 3.
pMCSA achieves the best performance compared to all other MCSA methods.
Also, its overall performance is comparable to exclusive KL minimization methods (ELBO) unlike other
MCSA methods.
Furthermore, on airfoil and energy , pMCSA improves over ELBO by 0.29 nat and 0.48 nat .
Even on gas where pMCSA did not beat ELBO, its performance is comparable, and it dominates all other MCSA methods by roughly 0.4 nat .
Additional experimental results, including root mean-square error (RMSE) results and plots with respect to the wall clock time, can be found in Appendix E.
[Table]
Weight Pruning When using VI, BNNs have been known to underfit data, which Hoffman (2017); MacKay (2001); Trippe & Turner (2017) associated with “weight pruning.”
That is, the variational posterior converges to the zero-mean prior.
Furthermore, Coker et al.
; Huix et al.
(2022) have shown that this is guaranteed to happen under certain conditions.
However, these results are strictly based on exclusive KL minimization.
On the other hand, Figure 3 shows that pMCSA does not suffer from weight pruning, which suggests that pruning is an artifact of using the exclusive KL.
Additional plots are shown in Figure 6 (Appendix E.1).
6.4 Robust Gaussian Process Regression
Setup We train Gaussian processes (GP) with a Student-T likelihood for robust regression.
We use datasets from the UCI repository (Dua & Graff, 2017) with 90% random train-test splits.
We use the Matérn 5/2 covariance kernel with automatic relevance determination (Neal, 1996) (see
Appendix C.2).
We run all methods with 𝑇 = 2⋅10 4 iterations.
For prediction, we use the mode of 𝑞 (⋅; 𝛌) for the hyperparameters and marginalize the latent function over 𝑞 (⋅; 𝛌) (Rasmussen &
Williams, 2006).
We consider ELBO with only 𝑁 = 1 since differentiating through the likelihood makes its per-iteration cost comparable to MCSA methods with 𝑁 = 10 .
Results The results are shown in Table 4.
Except for gas , pMCSA achieves better predictive densities than all other methods.
This suggests that, overall, the exclusive KL may be less effective in terms of uncertainty quantification for GP posteriors.
While ELBO achieves the best performance on gas , among MCSA methods, pMCSA dominates.
Our encouraging regression results suggest that incorporating methods such as inducing points (Snelson & Ghahramani, 2005) into MCSA may lead to an important new class of GP models.
Additional experimental results, including RMSE results and plots with respect to the wall clock time, can be found in Appendix E.2.
Note that in Appendix E.2, MCSA methods appear worse in terms of RMSE compared to the exclusive KL since the inclusive KL puts less probability volume around the posterior mode.
Related Works
Inclusive KL minimization Our MCSA framework generalizes MSC (Naesseth et al. , 2020) and JSA Ou & Song (2020), which are inclusive KL minimization based on SGD and Markov chains.
Similar to MCSA is the method of Li et al.
However, the convergence of this method is not guaranteed since it uses short Markov chains, disqualifying for MCSA.
Other methods based on biased gradients have been proposed by Bornschein & Bengio (2015); Le et al.
(2020), but these are specific for deep generative models.
On a different note, Jerfel et al.
(2021) use boosting instead of SGD to minimize the inclusive KL, which gradually builds a complex variational approximation from a simple variational family.
Beyond the KL Divergence Discovering alternative divergences for VI has been an active research area.
For example, the 𝜒 2 (Dieng et al. , 2017), 𝑓 (Wan et al. , 2020; Wang et al. , 2018), 𝛼 (Hernandez-Lobato et al. , 2016; Li & Turner, 2016; Regli & Silva, 2018), reguarlized importance ratio (Bamler et al. , 2017) divergences have been studied for VI.
However, for gradient estimation these methods involve the importance ratio 𝑤 (𝐳)
= 𝜋 (𝐳) /𝑞
(𝐳) , which leads to significant variance and low signal-to-noise ratio under model misspecification (Bamler et al. , 2017;
Geffner & Domke, 2021a,c).
In contrast, under stationarity, the variance of pMCSA is 𝜎 2 /𝑁
( 𝜎 is the variance of the score over the posterior) regardless of model misspecification.
Meanwhile, Geffner & Domke (2021b); Ruiz & Titsias (2019); Salimans et al.
; Zhang et al.
(2021) construct implicit divergences formed by MCMC.
With the exception of Ruiz & Titsias (2019), most of these approaches aim to maximize auxiliary representations of the classic ELBO.
Therefore, their property is likely to be similar to the exclusive KL.
Adaptive MCMC and MCSA As pointed out by Ou & Song (2020), using 𝑞 (⋅; 𝛌) within the
MCMC kernel makes MCSA structurally equivalent to adaptive MCMC.
In particular, Andrieu & Thoms (2008); Brofos et al.
; Gabrié et al.
; Garthwaite et al.
(2016) discuss the use of stochastic approximation in adaptive MCMC.
Also, Andrieu & Moulines (2006); Brofos et al.
; Giordani & Kohn (2010); Habib & Barber (2019); Holden et al.
; Keith et al.
; Neklyudov et al.
(2019) specifically discuss adapting the propsosal of IMH kernels, and some of them use KL divergence minimization.
These methods focus on showing ergodicity the samples ( 𝛈 𝑡 in our context) not the convergence of the variational approximation 𝑞 (⋅; 𝛌) .
In this work, we focused on the convergence of 𝑞 (⋅; 𝛌) , which could advance the adaptive MCMC side of the story.
Discussions
This paper presented a new theoretical framework for analyzing inclusive KL divergence minimization methods based on running SGD with Markov chains.
Furthermore, we proposed pMCSA, a new MCSA method that enjoys substantially low variance.
We have shown that this theoretical improvement translates into better empirical performance.
Limitations Our work has three main limitations.
Firstly, since our work aims to understand existing MCSA methods, it inherits their current limitations.
For example, minibatch subsampling is challenging for models with non-factorizable likelihoods (Naesseth et al. , 2020).
Secondly, our theoretical analysis in Section 4 requires Assumption 3, which is strong, but required to connect with MCGD.
An important future direction would be to relax the assumptions needed by MCGD.
Lastly, our MCSA framework does not include models with parameterized posteriors such as variational autoencoders.
Parameterized Posteriors On problems with parameterized posteriors, the target posterior moves around.
Therefore, quickly chasing the moving posterior with fast converging MCMC kernels is as important as achieving low variance.
Because of this, trading bias and variance is less straightforward compared to the “static” setting we consider.
Furthermore, the usage of expensive MCMC kernels could be beneficial as suggested by Zhang et al.
Towards Alternative Divergences
In Section 6, we have shown that minimizing the inclusive KL is competitive against minimizing the exclusive KL on general Bayesian inference problems.
Although Dhaka et al.
(2021) have shown that the inclusive KL fails on high-dimensional problems, this is only the case under the presence of strong correlations.
Before entirely ditching the inclusive KL, it is essential to ask, “how correlated are posteriors really in practice?”
Furthermore, the true performance of alternative divergences is often masked by the limitations of the inference procedure (Geffner & Domke, 2021a,c).
Given that pMCSA significantly advances the best-known performance of inclusive KL minimization, it is possible that similar improvements could be extracted from other divergences.
To conclude, our results motivate further development of better inference algorithms for alternative divergence measures.
Acknowledgments and Disclosure of Funding
We thank Hongseok Yang for pointing us to relevant related work, Guanyang Wang for insightful discussions about the independent Metropolis-Hastings algorithm, Geon Park and Kwanghee
Choi for constructive comments that enriched this paper, Christian A. Naesseth for comments about Rao-Blackwellized MSC, and Gabriel V. Cardoso for comments about BR-SNIS.
We also acknowledge the Department of Computer Science and Engineering of Sogang University for providing computational resources.
K. Kim was supported in part by the EPSRC through the Big Hypotheses grant [EP/R018537/1].
J. R. Gardner was supported by NSF award [IIS-2145644].
H. Kim was supported by the Basic
Science Research Program through the National Research Foundation of Korea (NRF) funded by the Ministry of Science and ICT under Grant [NRF-2021R1A2C1095435].
[-ELBO N=
-]