Flexible mean field variational inference using mixtures of non-overlapping exponential families

Bayesian graphical models [33] are widely used across vast swaths of engineering and science. Graphical models succinctly summarize modeling assumptions, and their posterior distributions naturally quantify uncertainty and produce optimal point estimators to minimize downstream risk. Efficient, exact posterior inference algorithms exist for special cases such as the sum-product algorithm [41] for discrete random variables in models with low tree-width [56] or particular Gaussian graphical models [27]. Outside of these special cases, however, alternative techniques must be used to approximate the posterior. Markov chain Monte Carlo (MCMC), especially Gibbs sampling, can sample from the posterior asymptotically [17, 20] but such techniques are difficult to scale to models with many observations, can have issues with mixing [39], and their convergence can be hard to assess [52].

Variational inference (VI) [26] avoids sampling and instead fits an approximate posterior via optimizing a loss function that acts as a proxy for some measure of divergence between the approximate and true posteriors. Usually this divergence is the “reverse” Kullback-Leibler (KL) divergence between the approximate posterior $Q$, and the true posterior $P$: $\text{KL}(Q||P)$. Minimizing this reverse KL is equivalent to maximizing the so-called evidence lower bound (ELBo) which depends only on the prior, the likelihood, and the approximate posterior. Importantly, the unknown true posterior does not appear in the ELBo. In many models, by choosing a particular space of approximate posteriors over which to optimize (the “variational family”), maximizing the ELBo is a tractable albeit non-convex optimization problem. In practice, such optimization problems are solved using coordinate ascent, gradient ascent, or natural gradient ascent [2]. These optimization methods and stochastic variants thereof can scale to massive models and datasets [22].

This formulation introduces a key tension in VI. On one hand, simpler variational families can be more computationally tractable to optimize. On the other hand, a large approximation error will be incurred if the true posterior does not lie within the variational family. VI is particularly tractable for fully conjugate models with all distributions belonging to exponential families, with a variational family that assumes all variables are independent. This approach, called mean field VI, has a number of desirable properties including computationally efficient, often analytic coordinate-wise updates that are equivalent to natural gradient steps [5]. Many useful models–including mixture models–are not fully conjugate, but can be made conjugate by introducing additional auxiliary variables. Unfortunately, introducing more variables and enforcing their independence in the variational family results in larger approximation gaps [50].

Many approaches have been developed to extend VI beyond the mean field approximation. Specific models, such as coupled hidden Markov models, are amenable to more expressive variational families because subsets of the variables form tractable models – an approach called structured VI [47]. For non-conjugate models or models with variational families that are otherwise not amenable to mean field VI, it may be possible to compute gradients numerically [32] or obtain an unbiased noisy estimate of the gradient via the log gradient trick (sometimes called the REINFORCE estimator) [45] or in certain special cases the reparameterization trick [31]. Because such approaches are broadly applicable they have come to be part of the toolkit of “black box VI” [44], which seeks to automatically derive gradient estimators for very general models and variational families. Stochastic gradient-based optimization has also been combined with structured VI [21]. Additionally, approximations to gradients tailored to specific distributions have been developed [24]. Much recent work has been devoted to combining deep learning with Bayesian graphical models, whether purely for speeding up inference, or by using neural networks as flexible models [4, 25, 55]. Overall, these general approaches are impressive in scope, but the gradient estimators can be noisy and hence optimizing the ELBo can require small step sizes and corresponding long compute times [46].

Here, I take an alternative approach, expanding the utility of mean field VI by showing that more flexible variational families can be constructed by combining component exponential families while maintaining desirable conjugacy properties and still forming an exponential family. This construction is particularly useful in models involving sparsity. The approach I present maintains the analytic simplicity of mean field VI but avoids the need to introduce auxiliary variables to obtain conjugacy.

I begin with a motivating example in Section 2, develop the main theoretical results in Section 3, and show their utility on two examples in Section 4. Implementation details, proofs, and additional theoretical results are presented in the Appendix.

To better understand the genetics of disease and other complex traits, it is now routine to collect genetic information and measure a trait of interest in many individuals. Initially, such studies hoped to find a small number of genomic locations associated with the trait to learn about the underlying biology or to find suitable drug targets. As more of these genome-wide association studies (GWAS) were performed, however, it became increasingly clear that enormous numbers of locations throughout the genome contribute to most traits, including many diseases [7]. While this makes it difficult to prioritize a small number of candidate genes for biological followup, one can still use GWAS data to predict a trait from an individual’s genotypes, an approach referred to as a polygenic score (PGS) [11]. For many diseases, PGSs are sufficiently accurate to be clinically relevant for stratifying patients by disease risk [29].

Typically PGSs are constructed by estimating effect sizes for each position in the genome and assuming that positions contribute to the trait additively. This assumption is justified by evolutionary arguments [48] and practical concerns about overfitting due to the small number of individuals in GWAS (tens to hundreds of thousands) relative to the number of genotyped positions (millions). More concretely, the predicted value of the trait $\hat{Y_{i}}$ in individual $i$ is

where $G_{ij}$ is the genotype of individual $i$ at position $j$ and $\beta_{j}$ is the effect size at position $j$.

Estimating the effect sizes, $\beta_{j}$, is complicated by the fact that to protect participant privacy GWAS typically only release marginal effect size estimates $\hat{\beta_{j}}$ obtained by separately regressing the value of the trait against the genotypes at each position. These marginal estimates thus ignore the correlation in individuals’ genotypes at different positions in the genome. Fortunately, these correlations are stable across sets of individuals from the same population and can be estimated from publicly available genotype data (e.g., [10]) even if that data does not contain measurements of the trait of interest.

An early approach to dealing with genotypic correlations was LDpred [53] which uses MCMC to fit the following model with a spike-and-slab prior on the effect sizes:

where $\mathbf{X}\in\mathbb{R}^{P\times P}$ is the matrix of correlations between genotypes at different positions in the genome, $\sigma^{2}_{e}$ is the variance of the estimated effect sizes from the GWAS and $\sigma^{2}_{g}$ is the prior variance of the non-zero effect sizes.

While this model was motivated by PGSs, it is also a special case of Bayesian generalized least squares with a spike-and-slab prior on the regression coefficients.

Note that if $\mathbf{X}$ is diagonal, then the likelihood is separable and the true posterior can, in fact, be determined analytically. Unfortunately, in real data genotypes are highly correlated (tightly linked) across positions and so $\mathbf{X}$ is far from diagonal and as a result integrating over the mixture components for each $\beta_{j}$ must be done simultaneously, resulting in a runtime that is exponential in the number of genotyped positions. To obtain a computationally tractable approximation to the true posterior, LDpred and a number of extensions that incorporate more flexible priors (e.g., [16, 34]) use MCMC. As noted above, however, MCMC has a number of drawbacks and so there may be advantages to deriving and implementing a VI scheme.

A natural mean field VI approach to approximating the posterior would be to split the mixture distribution using an auxiliary random variable as follows

with $\sigma^{2}_{0}=0$ and $\sigma^{2}_{1}=\sigma^{2}_{g}$ (treating a Gaussian with zero variance as a point mass for notational convenience). One can then use categorical variational distributions $q_{Z_{j}}$ for the $Z$s and Gaussian variational distributions $q_{\beta_{j}}$ for the $\beta$s. Unfortunately, this approach immediately encounters an issue: when calculating the ELBo, $\text{KL}(q_{\beta_{j}}||\mathcal{N}(0,0))$ is undefined unless $q_{\beta_{j}}$ is a point mass at zero because otherwise $q_{\beta_{j}}$ is not absolutely continuous with respect to the point mass at zero. This may seem to be merely a technical issue that vanishes if $\sigma^{2}_{0}$ is taken to be some tiny value so that $q_{\beta_{j}}$ is absolutely continuous with respect to $\mathcal{N}(0,\sigma^{2}_{0})$ and hence has well-defined KL, while for practical purposes $\mathcal{N}(0,\sigma^{2}_{0})$ acts like a point mass at zero. Yet, this superficial fix is not enough: the mean field assumption requires $Z_{j}$ and $\beta_{j}$ to be independent under the variational approximation to the posterior, which cannot capture the phenomenon that $Z_{j}=0$ forces $\beta_{j}$ to be close to zero while $Z_{j}=1$ allows $\beta_{j}$ to be far from zero. Further intuition is presented in Appendix A, where I analyze the case with only a single position (i.e., $P=1$) in detail.

Fortunately, this problem can be solved by noting that spike-and-slab distributions like $p_{0}\delta_{0}+(1-p_{0})\mathcal{N}(0,\sigma_{g}^{2})$ surprisingly still form an exponential family and are conjugate to the likelihoo. In this example, by using a spike-and-slab distribution for the variational distributions $q_{\beta_{j}}$, there is no need to use auxiliary variables to obtain analytical updates for fitting the approximate posterior. A similar approach has been considered for a specific model in a different context [8], but in the next section, I show why the analysis works for this approach and also show that it is more broadly applicable. In Section 4, I return to this motivating example to show that the naive mean field VI does indeed break down and that by using sparse conjugate priors an accurate approximation to the posterior may be obtained.

In this section I begin with a simple observation. While it is usually true that mixtures of distributions from an exponential family no longer form an exponential family, that is not the case for mixtures of distributions with distinct support. Mixtures of exponential family members with non-overlapping support always form an exponential family.

Note that Theorem 1 has an apparent asymmetry with respect to component $K$. This arises because the mixture weights $\pi_{1},\ldots,\pi_{K}$ are constrained to sum to one, and hence $\pi_{K}$ is completely determined by the other mixture weights. It would be possible to have a “symmetric” version of Theorem 1 but the constraint on the mixture weights would mean that the natural parameters would live on a strictly lower dimensional space. Such exponential families are called curved exponential families, and many of the desirable properties of exponential families do not hold for curved exponential families.

While the restriction in Theorem 1 to exponential families with non-overlapping support may seem particularly limiting, it is important to note that a new exponential family can be formed by restricting all of the distributions within any exponential family to lie within a fixed subset of their usual domain and renormalizing. In particular, one could divide the original full domain into non-overlapping subsets and form mixtures of exponential families restricted to these subsets. Another important set of exponential families are uniform distributions with fixed support. While each of these exponential families only contains a single distribution, by combining mixtures of these uniform distributions with non-overlapping support, Theorem 1 shows that piece-wise constant densities with fixed break points form an exponential family. In this paper I focus primarily on the case of mixtures of a continuous exponential family member with one or more point masses. By Theorem 1, mixtures of point masses at fixed locations form an exponential family, and diffuse distributions can be trivially restricted to falling outside of this set of measure zero, which makes clear that spike-and-slab distributions where the slab is a non-atomic distribution from an exponential family always form an exponential family.

While these substantially more flexible exponential families may be of independent interest, they are of little use in mean field VI unless they are conjugate to widely used likelihood models. The following theorem provides a recipe to take a conjugate exponential family model and create a model with a more flexible prior using Theorem 1 while maintaining conjugacy.

Intuitively, Theorem 2 says that for a given conjugate prior distribution, say $\mathcal{P}$, we can create a more flexible prior that maintains conjugacy by combining non-overlapping component distributions so long as those component distributions have the same sufficient statistics as $\mathcal{P}$ on their domain up to an affine transformation.

Note that the point mass distribution at a fixed point is a member of an exponential family with a sole member, and hence it has no sufficient statistics. Furthermore, any function is constant on the support of a point mass. This provides an immediate corollary that mixtures of any continuous distribution with a finite number of point-masses are conjugate to the same distributions as the original continuous distribution. Another example is mixtures of degenerate distributions where only a subset of parameters are fixed. For instance, consider the mean vector, $\mu$, of a multivariate Gaussian. One could put a prior on $\mu$ that is a mixture of the usual conjugate multivariate Gaussian along with degenerate distributions like a multivariate Gaussian with the condition that $\mu_{i}=0$, that $\mu_{i}=\mu_{j}$ for dimensions $i$ and $j$. In fact, any degenerate multivariate Gaussian distribution defined by affine constraints $\mathbf{M}\mu=\mathbf{v}$ for some matrix $\mathbf{M}$ and vector $\mathbf{v}$ could be included in this mixture.

Theorem 2 makes it extremely easy to add sparsity or partial sparsity to any conjugate model. A non-degenerate application is constructing asymmetric priors from symmetric ones. For example, we can construct an asymmetric prior by taking a mixture of a copy of the original prior restricted to the negative reals and a copy restricted to non-negative reals. These mixture form an exponential family because the copies have the same sufficient statistics as the original prior but have non-overlapping domains.

To show the applicability of the theoretical results presented in Section 3, I test sparse VI schemes using Theorems 1 and 2–which I refer to as the non-overlapping mixtures trick–on two models and compare these schemes to non-sparse and naive VI approximations showing the superior performance of treating sparsity exactly. The first model is the polygenic score prediction model discussed in Section 2 and the second is a spike-and-slab prior for sparse probabilistic principal component analysis (PCA) [18]. Python implementations of the fitting procedures and the results of the simulations presented here are available at https://github.com/jeffspence/non_overlapping_mixtures. All mean field results were obtained on a Late 2013 MacBook Pro and fitting the mean field VI schemes took less than five seconds per dataset for the polygenic score model and less than two minutes per dataset for sparse PCA.

VI has made it possible to scale Bayesian inference on increasingly complex models to increasingly massive datasets, but the error induced by obtaining an approximate posterior can be substantial [50]. Some of this error can be mitigated by using more flexible variational families, but doing so can require alternative methods for fitting, like numerical calculation of gradients [32], sampling-based stochastic gradient estimators [31, 45], or other approximations [24]. The results of Theorems 1 and 2 provide an alternative method, using mixtures of non-overlapping exponential families to provide a more flexible variational family while maintaining conjugacy. Even in models that are not fully conjugate, methods have been developed to exploit portions of the model that are conjugate, and the results presented here may prove useful in such schemes [28]. These schemes could be especially useful in obtaining stochastic gradient updates for bayesian neural networks with spike-and-slab priors on the weights. In particular, a ReLU applied to a Gaussian random variable is a mixture of a point mass and a Gaussian restricted to be positive, which is an exponential family by Theorem 1.

Here I focused on modeling sparse phenomena and found that the non-overlapping mixtures trick is superior to a naive approach of introducing auxiliary variables. Yet, the pitfalls I described occur whenever mean field VI is applied to mixture distributions where the mixture components are very different. This suggests that in some cases, it may be beneficial to use the sparse distributions presented here to approximate non-sparse mixture distributions and then treat the sparse approximation exactly.

Throughout, I assumed that the domains of the non-overlapping mixtures were specified a priori. This assumption could be relaxed, treating the domains as hyperparameters that could then be optimized with respect to the ELBo. Yet, it is not obvious that for arbitrary models the objective function need to be differentiable with respect to these hyperparameters, which may necessitate the use of zeroth-order optimization procedures such as Bayesian Optimization [14].

Exponential families also play an important role in other forms of variational inference including Expectation Propagation [38]. The non-overlapping mixtures trick may be useful in variational approaches beyond the usual reverse KL-minimizing mean field VI.

While the non-overlapping mixtures trick makes it easy to add sparsity to conjugate models, it is not a panacea to some of the common pitfalls of VI. For example, I also considered a sparse extension of latent Dirichlet allocation (LDA) [6, 43], where documents can have exactly one topic with positive probability. Unfortunately, the zero-avoiding nature of the reverse KL results in pathological behavior: for a document with only one topic the prior topic probabilities for each word are sparse, but in the variational posterior they must be dense. Empirically, this results in the VI posterior only putting mass on the non-sparse mixture component and hence being indistinguishable from the usual VI approach to LDA. In general, care should be taken when the sparsity added to a model results in the possibility of variables having zero likelihood conditioned on latent variables coming from the sparse component.

In spite of these drawbacks, providing a recipe to easily model sparsity in otherwise conjugate Bayesian models provides another avenue to model complex phenomena while maintaining the analytical and computational benefits of mean-field VI.

The primary contribution of this paper is theoretical and so the broader societal impact depends on how the theorems are used. The polygenic score application has the possibility to improve the overall quality of healthcare, but because the majority of GWAS are performed on individuals of European ancestries, PGSs are more accurate for individuals from those ancestry groups, potentially exacerbating health disparities between individuals of different ancestries as PGSs see clinical use [36]. The methods presented here are equally applicable to GWAS data collected from any ancestry group, however, and so efforts to diversify genetic data will ameliorate performance differences across ancestry groups. PGSs used for some traits such as sexual orientation [15], educational attainment [19], or stigmatized psychiatric disorders [11] raise thorny ethical considerations, especially when the application of such PGSs could enable genetic discrimination or fuel dangerous public misconceptions about the genetic basis of such traits [40]. On the other hand, PGSs applied to diseases have the potential to improve health outcomes and so if used responsibly could provide tremendous benefit to society.

I would like to thank Nasa Sinnott-Armstrong and Jonathan Pritchard for piquing my interest in the polygenic score application, and Jeffrey Chan, Amy Ko, Clemens Weiss, and four anonymous reviewers for helpful feedback on the manuscript. I was funded by a National Institutes of Health training grant (5T32HG000044-23).