Source-linked AI summary
Fast Variational Inference in the Conjugate Exponential Family
James Hensman, Magnus Rattray, Neil D. Lawrence
TL;DR
Variational coordinate ascent can be slow, and collapsed inference has been developed in model-specific forms. This paper unifies collapsed variational inference for conjugate-exponential models and uses the KL-corrected bound’s geometry to optimize faster. Experiments report gains across three models, including many-times-over improvement for large LDA.
Problem
Classical variational coordinate ascent can converge slowly, while collapsed-inference approaches have been developed for specific models rather than unified across conjugate-exponential graphs.
Source: §1 · p. 1Method
The paper derives a KL-corrected collapsed bound for conjugate-exponential models and applies natural-gradient and conjugate-gradient optimization using mean-field updates and a d-separation test.
Source: §1 · p. 1, §2.1 · p. 2, §6 · p. 11Results
Across three models, experiments found solutions unavailable to standard VBEM for Gaussian mixtures, many-times-over performance improvement for large LDA, and very large performance gains for BitSeq.
Source: §6 · p. 11Takeaways & Limitations
The method can be applied wherever mean-field update equations are available, with collapsed variables selected by d-separation, and can be incorporated into automated inference engines such as infer.net.
Source: §6 · p. 11Further second-order corrections proposed by Kurihara et al. [2007] and Teh et al. [2007] require case-by-case forms and were reported less effective than the unified methods [Asuncion et al., 2012].
Source: §3.4 · p. 5
Abstract
from arXiv · showhide
We present a general method for deriving collapsed variational inference algo- rithms for probabilistic models in the conjugate exponential family. Our method unifies many existing approaches to collapsed variational inference. Our collapsed variational inference leads to a new lower bound on the marginal likelihood. We exploit the information geometry of the bound to derive much faster optimization methods based on conjugate gradients for these models. Our approach is very general and is easily applied to any model where the mean field update equations have been derived. Empirically we show significant speed-ups for probabilistic models optimized using our bound.
1 Introduction
Variational coordinate ascent can converge slowly, and collapsed inference addresses this by analytically marginalizing part of the variational approximation. The paper unifies collapsed inference for conjugate-exponential graphs and motivates faster optimization methods for the resulting bound.
1 Introduction
- Analytically marginalizing part of the variational approximation removes it from optimization, a strategy used in several model-specific collapsed-inference approaches.The paper presents a unifying framework for applying collapsed inference to conjugate-exponential graphs.Source: §1 · p. 1
- The framework derives collapsed bounds, uses d-separation to identify applicable variables, and motivates natural-gradient and conjugate-gradient optimization.The authors report often obtaining an order of magnitude or more increase in convergence speed across models from the literature.Source: §1 · p. 1
2 The Marginalised Variational Bound
The paper relates collapsed variational bounds to the KL-corrected bound and gives a derivation through variational approximation followed by marginalization. In the mean-field setup, substituting the optimal factor removes its parameters from the optimization.
2 The Marginalised Variational Bound
- Earlier collapsed variational approaches addressed specific models, while this framework unifies them through the KL-corrected bound.For conjugate-exponential graphs, the order of marginalization and variational approximation is irrelevant, yielding a succinct derivation without approximations in bound evaluation or optimization.Source: §2 · p. 2
2.1 Variational Inference
- Mean-field variational inference factorizes q(Z, X) as q(Z)q(X) and alternately optimizes each factor.Substituting the optimal q(X) into the bound removes q(X) from the optimization; the resulting KL-corrected bound differs from the mean-field bound by a KL divergence.Source: §2.1 · p. 2
- The derivation forms a conditional variational lower bound on the log likelihood and then marginalizes the conditioned variable.The resulting object is used to compute the final bound, gradients, curvature, and the collapsed-variable distribution q⋆(X).Source: §2.1 · p. 3
3 Partial Equivalence of the Bounds
The KL-corrected and mean-field bounds agree in value and gradient after the collapsed factor is updated, but the KL-corrected bound has different curvature. The paper connects this geometry to faster optimization, clarifies its relation to collapsed variational Bayes, and gives a d-separation test for applicability.
3 Partial Equivalence of the Bounds
- After updating q(X), the KL-corrected and mean-field bounds are equal in value, while their difference before that update is KL(q∗(X)||q(X)).The KL-corrected bound reduces the number of optimization parameters and can be combined with more sophisticated optimization than fixed-step VBEM.Source: §3 · p. 3
3.1 Gradients
- At q(X) = q⋆(X), the two bounds also have equal gradients, and mean-field updates provide gradients for optimizing the KL-corrected bound.The paper characterizes VBEM as steepest ascent and proposes conjugate gradients on the KL-corrected bound for faster convergence.Source: §3.1 · p. 4
3.2 Curvature of the Bounds
- The KL-corrected bound adds positive-semidefinite curvature terms, making its curvature less negative near a maximum and permitting larger optimization steps.Figure 1(b) sketches how the bounds compare in value and gradient at q(X) = q⋆(X), and how their difference away from that point is a KL divergence.Source: §3.2 · pp. 4–5, Figure 1 · p. 6
3.3 Relationship to Collapsed VB
- The KL-corrected derivation marginalizes after variational approximation, unlike collapsed variational Bayes; under the stated approximation, the approaches are equivalent.The KLC expression is analytically tractable under an appropriately conjugate formulation, whereas the collapsed-VB expectation may not be.Source: §3.3 · p. 5
3.4 Applicability
- A d-separation test identifies variables to marginalize by checking their independence given observed data and the approximated variables.The selected set must induce a suitable factorization and leave an effective likelihood conjugate to the prior; Figure 1(a) illustrates possible variable selections.Source: §3.4 · pp. 5–6, Figure 1 · p. 6
- Further second-order corrections proposed by Kurihara et al. [2007] and Teh et al. [2007] require case-by-case forms and were reported less effective than the unified methods [Asuncion et al., 2012].Those corrections assume q(Z) is Gaussian to obtain tractability.Source: §3.4 · p. 5
4 Riemannian Gradient Based Optimisation
The KLC bound supports gradient optimization with fewer parameters, and exponential-family information geometry makes natural gradients simple to compute. The paper develops conjugate-gradient optimization from this foundation.
4 Riemannian Gradient Based Optimisation
- The KLC bound needs only the non-collapsed distribution’s parameters, and its gradients follow from mean-field updates, making gradient optimization a natural fit.The implicit distribution of the collapsed variables is determined by the non-collapsed distribution and the data.Source: §4 · pp. 6–7
4.1 Variable Transformations
- For exponential-family approximations, gradients in canonical and expectation parameters provide each other’s natural gradients, avoiding an expensive matrix inverse.The Fisher information in these two parameterizations is reciprocal, with G(θ) = ∂η/∂θ.Source: §4.1 · p. 7
4.2 Steepest Ascent is Coordinate Ascent
- VBEM is steepest ascent with unit steps in canonical parameters; its update equations also give the KLC-bound gradient.The paper confirms these relationships for the studied models in supplementary material.Source: §4.2 · p. 7
4.3 Conjugate Gradient Optimization
- Conjugate gradients combine the current natural gradient with the previous search direction; Fletcher–Reeves worked well empirically, using unit step length.The Riemannian inner product simplifies to the ordinary gradient inner product in the described calculation.Source: §4.3 · p. 8
5 Experiments
Experiments compare VBEM with conjugate-gradient optimization on Gaussian mixtures, LDA, and RNA-Seq inference. Conjugate methods were faster in the reported comparisons, while the RNA-Seq solutions had bounds within 60 nats of one another.
5 Experiments
- Experiments used shared initial conditions for each method and treated convergence as a bound change or Riemannian-gradient change below 10^-6.The Gaussian-mixture setup varied overlap from R=1 to R=5, with eight initial components and an uninformative prior over mixing proportions.Source: §5 · p. 8, §5.1 · p. 8
5.1 Mixtures of Gaussians
- For Gaussian mixtures, VBEM was at least twice as slow as the best conjugate method even at 100-nat tolerance, and sometimes missed the optimum in 500 restarts.Table 1 reports average iterations to come within 10 nats of the best-known solution across varying cluster overlap R.Source: §5.1 · p. 9, Table 1 · p. 9
5.2 Topic Models
- For LDA, Hestenes–Steifel was almost ten times as fast as standard VB, while the final bound varied little between approaches.Table 2 reports time and iterations for two conjugate methods and VBEM, and identifies Fletcher–Reeves as almost ten times as fast as VBEM.Source: §5.2 · p. 9, Table 2 · p. 10
5.3 RNA-seq alignment
- For RNA-Seq with six million latent variables, VBEM took 4600±20 iterations, versus 268 ± 4 and 265 ± 1 for Fletcher–Reeves and Hestenes–Steifel; runtime was nearly 11 hours versus about 40 minutes.The solutions for the other algorithms were fairly close, with bounds within 60 nats; all variational approaches improved on a Gibbs sampler’s approximately one-week runtime [Glaus et al., 2012].Source: §5.3 · p. 10
6 Discussion
The paper unifies collapsed and marginalized variational bounds for conjugate exponential-family models, relates the KLC bound to standard mean-field updates, and reports improved performance across three models. It attributes faster optimization to the KLC bound’s curvature and describes how the approach can be applied using existing mean-field updates.
6 Discussion
- Under conjugate exponential-family conditions, collapsed and marginalized variational bounds are equivalent, and the paper provides a unified derivation.The equivalence uses the KL-corrected perspective of King and Lawrence [2006].Source: §6 · p. 10
- Standard mean-field updates are steepest ascent on the KLC bound, making the method applicable wherever mean-field updates exist and requiring only a d-separation test to select variables to collapse.The authors report that their implementations switch between VBEM and conjugate methods with a few lines of code.Source: §6 · pp. 10–11
- The KLC bound’s curvature is always less negative than the original variational bound, allowing larger steps in variational parameters.The authors suggest this may provide a gateway to even faster second-order optimization.Source: §6 · p. 11
- Across three models, the conjugate Riemannian algorithm found MoG solutions standard VBEM did not, LDA performance improved many times over VBEM, and BitSeq showed very large improvements.The BitSeq result concerns a model with huge numbers of latent variables.Source: §6 · p. 11
Supplementary Material for:
The supplementary material explains how the paper’s general framework applies to its specific models and briefly introduces the three conjugate-gradient algorithms used for optimization.
Supplementary Material for:
- The supplement’s purpose is to detail how the general framework applies to the specific models described in the paper.Source: Supplementary Material for · p. 14
- It briefly presents the forms of three conjugate-gradient algorithms used in optimization.Source: Supplementary Material for · p. 14
- The supplement accompanies the NIPS paper on Fast Variational Inference in the Conjugate Exponential Family.Source: Supplementary Material for · p. 14
7 Conjugate gradient algorithms
The section names three methods for approximating the conjugate-gradient parameter β and specifies the Riemannian inner product used in the algorithms.
7 Conjugate gradient algorithms
- The Polak-Ribière, Fletcher-Reeves, and Hestenes-Stiefel methods are used to approximate β in the conjugate-gradient algorithm.Source: §7 · p. 14
- The three named methods are alternative approaches for approximating the conjugate-gradient parameter β.Source: §7 · p. 14
- The Riemannian inner product is given by eg⊤G(ρ)eg.Source: §7 · p. 14
8 Mixture of Gaussians
The supplement applies the KLC bound to a mixture-of-Gaussians model by parameterizing cluster assignments and marginalizing mixing proportions and cluster parameters. This leaves the assignment distribution as the variational parameterization and makes the required integrals tractable.
8 Mixture of Gaussians
- The MoG model uses N D-dimensional vectors, binary latent cluster-membership variables, and cluster parameters comprising means and precision matrices.Source: §8 · p. 14
- The model assigns multinomial cluster memberships using mixing proportions with a Dirichlet prior, and uses a conjugate Gaussian-Wishart prior for cluster parameters.Source: §8 · pp. 14–15
8.1 Applying the KLC bound
- A d-separation test permits marginalizing mixing proportions π and cluster parameters η when latent assignments L are parameterized.Source: §8.1 · p. 15, Figure 2 · p. 15
- The chosen variational approximation for cluster assignments is a factorized multinomial distribution parameterized with softmax functions.Source: §8.1 · p. 15
- After multiplying by the priors, the integrals over η and π are tractable, leaving only the parameters of q(L) to optimize.Source: §8.1 · p. 16
- Intermediate variables combine variational parameters, data, and prior parameters, and their dependence on r is retained during optimization rather than ignored as in mean-field variational inference.Source: §8.1 · p. 16
9 Latent Dirichlet Allocation
The LDA section sets out the topic-assignment model and derives a collapsed bound by parameterizing latent assignments and marginalizing document and topic proportions. It also presents example topics found in 2011 NIPS papers.
9 Latent Dirichlet Allocation
- Each word is assigned a latent topic from its document’s topic proportions, then generated from that topic’s word proportions, with symmetric Dirichlet priors on both proportions.LDA is described as a popular topic model.Source: §9 · p. 17
9.1 The collapsed bound
- The collapsed-bound derivation parameterizes latent assignments L, using a d-separation test to select variables for marginalization.The section assumes a factorizing multinomial posterior for L and proceeds through a conditional bound.Source: §9.1 · pp. 17–18
- Marginalizing θ and φ separates the integrals into normalizers for the independent Dirichlet approximations.The derivation exponentiates the conditional bound, takes its expectation under the priors, then takes the logarithm.Source: §9.1 · p. 18
9.2 Topics found by LDA
- Table 3 presents some topics found using LDA on papers from the 2011 NIPS conference.The section introduces these topics as examples.Source: §9.2 · p. 18, Table 3 · p. 19
10 BitSeq Model
The BitSeq model infers transcript abundances from RNA-seq fragments using latent read-to-transcript assignments. Its collapsed bound parameterizes those assignments and marginalizes transcript abundances under a conjugate Dirichlet prior.
10 BitSeq Model
- RNA-seq fragments are assigned to transcripts through latent membership variables, while transcript abundances θ are the primary inference target.The model treats each fragment’s transcript identity as a latent assignment.Source: §10 · p. 19
- With a known transcriptome, read-to-transcript likelihoods are precomputed from alignment quality and can correct sequence-specific or fragmentation biases.The alignment model is described in Glaus et al. [2012].Source: §10 · p. 19
- A conjugate Dirichlet prior over θ represents prior beliefs about transcript abundances and is described as relatively uninformative but proper.The prior concentrations are assumed equal with large uncertainty.Source: §10 · pp. 19–20
10.1 The collapsed bound
- Parameterizing L allows θ to be marginalized: the conditional bound is conjugate to θ’s prior, and the approximate posterior for θ is Dirichlet.The graphical model’s d-separation test identifies this collapse choice.Source: §10.1 · p. 20, Figure 4 · p. 20