Conditionally structured variational Gaussian approximation with importance weights^††thanks: Linda Tan and Aishwarya Bhaskaran are supported by the start-up grant R-155-000-190-133.

In many modern statistical applications, it is necessary to model complex dependent data. In these situations, models which employ observation specific latent variables such as random effects and state space models are widely used because of their flexibility, and Bayesian approaches dealing naturally with the hierarchical structure are attractive in principle. However, incorporating observation specific latent variables leads to a parameter dimension increasing with sample size, and standard Bayesian computational methods can be challenging to implement in very high-dimensional settings. For this reason, approximate inference methods are attractive for these models, both in exploratory settings where many models need to be fitted quickly, as well as in applications involving large datasets where exact methods are infeasible. One of the most common approximate inference paradigms is variational inference (Ormerod and Wand, 2010; Blei et al., 2017), which is the approach considered here.

Our main contribution is to consider partitioning the unknowns in a local latent variable model into “global” parameters and “local” latent variables, and to suggest ways of structuring the dependence in a variational approximation that match the specification of these models. We go beyond standard Gaussian approximations by defining the variational approximation sequentially, through a marginal density for the global parameters and a conditional density for local parameters given global parameters. Each term in our approximation is Gaussian, but we allow the conditional covariance matrix for the local parameters to depend on the global parameters, which leads to an approximation that is not jointly Gaussian. We are particularly interested in improved inference on global variance and dependence parameters which determine the scale and dependence structure of local latent variables. With this objective, we suggest a parametrization of our conditional approximation to the local variables that is well-motivated and respects the exact conditional independence structure in the true posterior distribution. Our approximations are parsimonious in terms of the number of required variational parameters, which is important since a high-dimensional variational optimization is computationally burdensome. The methods suggested improve on Gaussian variational approximation methods for a similar computational cost. Besides defining a novel and useful variational family appropriate to local latent variable models, we also employ importance weighted variational inference methods (Burda et al., 2016; Domke and Sheldon, 2018) to further improve the quality of inference, and elaborate further on the connections between this approach and the use of Rényi’s divergence within the variational optimization (Li and Turner, 2016; Regli and Silva, 2018; Yang et al., 2019).

Our method is a contribution to the literature on the development of flexible variational families, and there are many interesting existing methods for this task. One fruitful approach is based on normalizing flows (Rezende and Mohamed, 2015), where a variational family is defined using an invertible transformation of a random vector with some known density function. To be useful, the transformation should have an easily computed Jacobian determinant. In the original work of Rezende and Mohamed (2015), compositions of simple flows called radial and planar flows were considered. Later authors have suggested alternatives, such as autoregressive flows (Germain et al., 2015), inverse autoregressive flows (Kingma et al., 2016), and real-valued non-volume preserving transformations (Dinh et al., 2017), among others. Spantini et al. (2018) gives a theoretical framework connecting Markov properties of a target posterior distribution to representations involving transport maps, with normalizing flows being one way to parametrize such mappings. The variational family we consider here can be thought of as a simple autoregressive flow, but carefully constructed to preserve the conditional independence structure in the true posterior and to achieve parsimony in the representation of dependence between local latent variables and global scale parameters. Our work is also related to the hierarchically structured approximations considered in Salimans and Knowles (2013, Section 7.1); these authors also consider other flexible approximations based on mixture models, and a variety of innovative numerical approaches to the variational optimization. Hoffman and Blei (2015) propose an approach called structured stochastic variational inference which is applicable in conditionally conjugate models. Their approach is similar to ours, in the sense that local variables depend on global variables in the variational posterior. However conditional conjugacy does not hold in the examples we consider.

The methods we describe can be thought of as extending the Gaussian variational approximation (GVA) of Tan and Nott (2018), where parametrization of the variational covariance matrix was considered in terms of a sparse Cholesky factor of the precision matrix. Similar approximations have been considered for state space models in Archer et al. (2016). The sparse structure reduces the number of free variational parameters, and allows matching the exact conditional independence structure in the true posterior. Tan (2018) propose an approach called reparametrized variational Bayes, where the model is reparametrized by applying an invertible affine transformation to the local variables to minimize their posterior dependency on global variables, before applying a mean field approximation. The affine transformation is obtained by considering a second order Taylor series approximation to the posterior of the local variables conditional on the global variables. One way of improving on Gaussian approximations is to consider mixtures of Gaussians (Jaakkola and Jordan, 1998; Salimans and Knowles, 2013; Miller et al., 2016; Guo et al., 2016). However, even with a parsimonious parametrization of component densities, a large number of additional variational parameters are added with each mixture component. Other flexible variational families can be formed using copulas (Tran et al., 2015; Han et al., 2016; Smith et al., 2019), hierarchical variational models (Ranganath et al., 2016) or implicit approaches (Huszár, 2017).

We specify the model and notation in Section 2 and introduce the conditionally structured Gaussian variational approximation (CSGVA) in Section 3. The algorithm for optimizing the variational parameters is described in Section 4 and Section 5 highlights the association between GVA and CSGVA. Section 6 describes how CSGVA can be improved using importance weighting. Experimental results and applications to generalized linear mixed models (GLMMs) and state space models are presented in Sections 7, 8 and 9 respectively. Section 10 gives some concluding discussion.

Let $y=(y_{1},\dots,y_{n})^{T}$ be observations from a model with global variables $\theta_{G}$ and local variables $\theta_{L}=(b_{1},\dots,b_{n})^{T}$, where $b_{i}$ contains latent variables specific to $y_{i}$ for $i=1,\dots,n$. Suppose $\theta_{G}$ is a vector of length $G$ and each $b_{i}$ is a vector of length $L$. Let $\theta=(\theta_{G}^{T},\theta_{L}^{T})^{T}$. We consider models where the joint density is of the form

The observations $\{y_{i}\}$ are conditionally independent given $\{b_{i}\}$ and $\theta_{G}$. Conditional on $\theta_{G}$, the local variables $\{b_{i}\}$ form a $\ell$th order Markov chain if $\ell\geq 1$, and they are conditionally independent if $\ell=0$. This class of models include important models such as GLMMs and state space models. Next, we define some mathematical notation before discussing CSGVA for this class of models.

We propose to approximate the posterior distribution $p(\theta|y)$ of the model defined in Section 2 by a density of the form

where $q(\theta_{G})=N(\mu_{1},\Omega_{1}^{-1})$, $q(\theta_{L}|\theta_{G})=N(\mu_{2},\Omega_{2}^{-1})$, and $\Omega_{1}$ and $\Omega_{2}$ are the precision (inverse covariance) matrices of $q(\theta_{G})$ and $q(\theta_{L}|\theta_{G})$ respectively. Here $\mu_{2}$ and $\Omega_{2}$ depend on $\theta_{G}$, but we do not denote this explicitly for notational conciseness. Let $C_{1}C_{1}^{T}$ and $C_{2}C_{2}^{T}$ be unique Cholesky factorizations of $\Omega_{1}$ and $\Omega_{2}$ respectively, where $C_{1}$ and $C_{2}$ are lower triangular matrices with positive diagonal entries. We further define $C_{1}^{*}$ and $C_{2}^{*}$ to be lower triangular matrices of order $G$ and $nL$ respectively such that $C_{r,ii}^{*}=\log(C_{r,ii})$ and $C_{r,ij}^{*}=C_{r,ij}$ if $i\neq j$ for $r=1,2$. The purpose of introducing $C_{1}^{*}$ and $C_{2}^{*}$ is to allow unconstrained optimization of the variational parameters in the stochastic gradient ascent algorithm, since diagonal entries of $C_{1}$ and $C_{2}$ are constrained to be positive. Note that $C_{2}$ and $C_{2}^{*}$ also depend on $\theta_{G}$ but again we do not show this explicitly in our notation.

As $\Omega_{2}$ depends on $\theta_{G}$, the joint distribution $q(\theta)$ is generally non-Gaussian even though $q(\theta_{G})$ and $q(\theta_{L}|\theta_{G})$ are individually Gaussian. Here we consider a first order approximation and assume that $\mu_{2}$ and $v(C_{2}^{*})$ are linear functions of $\theta_{G}$:

In (1), $d$ is a vector of length $nL$, $D$ is a $nL\times G$ matrix, $f$ is a vector of length $nL(nL+1)/2$ and $F$ is a $nL(nL+1)/2\times G$ matrix. For this specification, $q(\theta)$ is not jointly Gaussian due to dependence of the covariance matrix of $q(\theta_{L}|\theta_{G})$ on $\theta_{G}$. It is Gaussian if and only if $F\equiv 0$. The set of variational parameters to be optimized is denoted as

As motivation for the linear approximation in (1), consider the linear mixed model,

where $y_{i}$ is a vector of responses of length $n_{i}$ for the $i$th subject, $X_{i}$ and $Z_{i}$ are covariate matrices of dimensions $n_{i}\times p$ and $n_{i}\times L$ respectively, $\beta$ is a vector of coefficients of length $p$ and $\{b_{i}\}$ are random effects. Assume $\sigma_{\epsilon}^{2}$ is known. Then the global parameters $\theta_{G}$ consists of $\beta$ and $\Lambda$. The posterior of $\theta_{L}$ conditional on $\theta_{G}$ is

Thus $p(\theta_{L}|y,\theta_{G})=\prod_{i=1}^{n}p(b_{i}|y,\theta_{G})$, where $p(b_{i}|y,\theta_{G})$ is a normal density with precision matrix $Z_{i}^{T}Z_{i}/\sigma_{\epsilon}^{2}+\Lambda^{-1}$ and mean $(Z_{i}^{T}Z_{i}/\sigma_{\epsilon}^{2}+\Lambda^{-1})^{-1}Z_{i}^{T}(y_{i}-X_{i}\beta)/\sigma_{\epsilon}^{2}$. The precision matrix depends on $\Lambda^{-1}$ linearly and the mean depends on $\beta$ linearly after scaling by the covariance matrix. The linear approximation in (1) tries to mimic this dependence relationship.

The proposed variational density is conditionally structured and highly flexible. Such dependence structure is particularly valuable in constructing variational approximations for hierarchical models, where there are global scale parameters in $\theta_{G}$ which help to determine the scale of local latent variables in the conditional posterior of $\theta_{L}|\theta_{G}$. While marginal posteriors of the global variables are often well approximated by Gaussian densities, marginal posteriors of the local variables tend to exhibit more skewness and kurtosis. This deviation from normality can be captured by $q(\theta_{L})=\int q(\theta_{G})q(\theta_{L}|\theta_{G})d\theta_{G}$, which is a mixture of normal densities. The formulation in (1) also allows for a reduction in the number of variational parameters if conditional independence structure consistent with that in the true posterior is imposed on the variational approximation.

To make the dependence on $\lambda$ explicit, write $q(\theta)$ as $q_{\lambda}(\theta)$. The variational parameters $\lambda$ are optimized by minimizing the Kullback-Leibler divergence between $q_{\lambda}(\theta)$ and the true posterior $p(\theta|y)$, where

Minimizing $\text{\rm KL}\{q_{\lambda}||p(\theta|y)\}$ is therefore equivalent to maximizing an evidence lower bound ${\mathcal{L}}^{\text{VI}}$ on the log marginal likelihood $\log p(y)$, where

In (2), $E_{q_{\lambda}}$ denotes expectation with respect to $q_{\lambda}(\theta)$. We seek to maximize ${\mathcal{L}}^{\text{VI}}$ with respect to $\lambda$ using stochastic gradient ascent. Starting with some initial estimate of $\lambda$, we perform the following update at each iteration $t$,

where $\rho_{t}$ represents a small stepsize taken in the direction of the stochastic gradient $\widehat{\nabla}_{\lambda}{\mathcal{L}}^{\text{VI}}$. The sequence $\{\rho_{t}\}$ should satisfy the conditions $\sum_{t}\rho_{t}=\infty$ and $\sum_{t}\rho_{t}^{2}<\infty$ (Spall, 2003).

An unbiased estimate of the gradient $\nabla_{\lambda}{\mathcal{L}}^{\text{VI}}$ can be constructed using (2) by simulating $\theta$ from $q_{\lambda}(\theta)$. However, this approach usually results in large fluctuations in the stochastic gradients. Hence we implement the “reparametrization trick” (Kingma and Welling, 2014; Rezende et al., 2014; Titsias and Lázaro-Gredilla, 2014), which helps to reduce the variance of the stochastic gradients. This approach writes $\theta\sim q_{\lambda}(\theta)$ as a function of the variational parameters $\lambda$ and a random vector $s$ having a density not depending on $\lambda$. To explain further, let $s=[s_{1}^{T},s_{2}^{T}]^{T}$, where $s_{1}$ and $s_{2}$ are vectors of length $G$ and $nL$ corresponding to $\theta_{G}$ and $\theta_{L}$ respectively. Consider a transformation $\theta=r_{\lambda}(s)$ of the form

Since $\mu_{2}$ and $C_{2}$ are functions of $\theta_{G}$ from (1),

Hence $\mu_{2}$ and $C_{2}$ are functions of $s_{1}$, and $\theta_{L}$ is a function of both $s_{1}$ and $s_{2}$. This transformation is invertible since given $\theta$ and $\lambda$, we can first recover $s_{1}=C_{1}^{T}(\theta_{G}-\mu_{1})$, find $\mu_{2}$ and $C_{2}$, and then recover $s_{2}=C_{2}^{T}(\theta_{L}-\mu_{2})$. Applying this transformation,

where $E_{\phi}$ denotes expectation with respect to $\phi(s)$ and $\theta=r_{\lambda}(s)$.

CSGVA is an extension of Gaussian variational approximation (GVA, Tan and Nott, 2018). In both approaches, the conditional posterior independence structure of the local latent variables is used to introduce sparsity in the precision matrix of the approximation. Below we demonstrate that GVA is a special case of CSGVA when $F\equiv 0$.

Tan and Nott (2018) consider a GVA of the form

Note that $T_{LL}$ and $T_{GG}$ are lower triangular matrices. Using a vector $s=[s_{L}^{T},s_{G}^{T}]^{T}\sim N(0,I_{nL+G})$, we can write

Assuming $F\equiv 0$ for CSGVA, we have from (3) that

where $[s_{2}^{T},s_{1}^{T}]^{T}\sim N(0,I_{nL+G})$. Hence we can identify

If the standard way of initializing of Algorithm 1 (by setting $\lambda=0$) does not work well, we can use this association to initialize Algorithm 1 by using the fit from GVA. This informative initialization can reduce computation time significantly although there may be a risk of getting stuck in a local mode.

Here we discuss how CSGVA can be improved by maximizing an importance weighted lower bound (IWLB, Burda et al., 2016), which leads to a tighter lower bound on the log marginal likelihood, and a variational approximation less prone to underestimation of the true posterior variance. We also relate the IWLB with Rényi’s divergence (Rényi, 1961; van Erven and Harremos, 2014) between $q_{\lambda}(\theta)$ and $p(\theta|y)$, demonstrating that maximizing the IWLB instead of the usual evidence lower bound leads to a transition in the behavior of the variational approximation from “mode-seeking” to “mass-covering”. We first define Rényi’s divergence and the variational Rényi bound (Li and Turner, 2016), before introducing the IWLB as the expectation of a Monte Carlo approximation of the variational Rényi bound.