Conditionally structured variational Gaussian approximation with importance weights

Abstract

We develop flexible methods of deriving variational inference for models with complex latent variable structure. By splitting the variables in these models into “global” parameters and “local” latent variables, we define a class of variational approximations that exploit this partitioning and go beyond Gaussian variational approximation. This approximation is motivated by the fact that in many hierarchical models, there are global variance parameters which determine the scale of local latent variables in their posterior conditional on the global parameters. We also consider parsimonious parametrizations by using conditional independence structure and improved estimation of the log marginal likelihood and variational density using importance weights. These methods are shown to improve significantly on Gaussian variational approximation methods for a similar computational cost. Application of the methodology is illustrated using generalized linear mixed models and state space models.

Appendices

Appendix A: Derivation of stochastic gradient

Let $\otimes$ denote the Kronecker product between any two matrices. We have

$$\begin{aligned} r_\lambda (s) = \begin{bmatrix} \theta _G \\ \theta _L \end{bmatrix} = \begin{bmatrix} \mu _1 + C_1^{-T} s_1 \\ d+ C_2^{-T} (s_2 - DC_1^{-T} s_1) \end{bmatrix}, \end{aligned}$$

where $v(C_2^*) = f + F(\mu _1 + C_1^{-T} s_1)$. Differentiating $r_\lambda (s)$ with respect to $\lambda$, $\nabla _\lambda r_\lambda (s)$ is given by

$$\begin{aligned} \begin{aligned}&\begin{bmatrix} \nabla _{\mu _1} \theta _G &{} \nabla _{\mu _1} \theta _L \\ \nabla _{v(C_1^*)} \theta _G &{} \nabla _{v(C_1^*)} \theta _L \\ \nabla _{d} \theta _G &{} \nabla _{d} \theta _L \\ \nabla _{\mathrm{vec}(D)} \theta _G &{} \nabla _{\mathrm{vec}(D)} \theta _L \\ \nabla _{f} \theta _G &{} \nabla _{f} \theta _L \\ \nabla _{\mathrm{vec}(F)} \theta _G &{} \nabla _{\mathrm{vec}(F)} \theta _L \\ \end{bmatrix}. \end{aligned} \end{aligned}$$

Since $\theta _G$ does not depend on d, D, f and F, we have

$$\begin{aligned} \begin{aligned} \nabla _{d} \theta _G&= 0_{nL \times G}, \quad \nabla _{\mathrm{vec}(D)} \theta _G = 0_{nLG \times G} \\ \nabla _{f} \theta _G&= 0_{nL(nL+1)/2 \times G}, \;\; \nabla _{\mathrm{vec}(F)} \theta _G = 0_{nLG(nL+1)/2 \times G}. \end{aligned} \end{aligned}$$

It is easy to see that $\nabla _{\mu _1} \theta _G = I_G$ and $\nabla _d \theta _L = I_{nL}$. The rest of the terms are derived as follows.

Differentiating $\theta _G$ with respect to $v(C_1^*)$,

$$\begin{aligned} \begin{aligned} \mathrm{d}\theta _G&= -C_1^{-T} \mathrm{d}(C_1^T) C_1^{-T} s_1\\&= - (s_1^T C_1^{-1} \otimes C_1^{-T}) K_G E_G^T D_1^* \mathrm{d}v(C_1^*) \\&= - (C_1^{-T} \otimes s_1^T C_1^{-1}) E_G^T D_1^* \mathrm{d}v(C_1^*). \\ \therefore \; \nabla _{v(C_1^*)} \theta _G&= - D_1^* E_G (C_1^{-1} \otimes C_1^{-T} s_1 ). \end{aligned} \end{aligned}$$

Differentiating $\theta _L$ with respect to f,

$$\begin{aligned} \begin{aligned} \mathrm{d}\theta _L&= -C_2^{-T} \mathrm{d}(C_2^T) C_2^{-T} (s_2 - DC_1^{-T} s_1) \\&= - \{ (s_2 - DC_1^{-T} s_1)^T C_2^{-1} \otimes C_2^{-T} \} \\&\quad \times K_{nL} E_{nL}^T D_2^* \mathrm{d}f \\ \therefore \; \nabla _{f} \theta _L&= - D_2^* E_{nL} \{ C_2^{-1} \otimes C_2^{-T} (s_2 - DC_1^{-T} s_1) \}. \end{aligned} \end{aligned}$$

Differentiating $\theta _L$ with respect to F,

$$\begin{aligned} \begin{aligned} \mathrm{d}\theta _L&= (\nabla _{f} \theta _L)^T \mathrm{d}F \theta _G \\&= \{ \theta _G^T \otimes (\nabla _{f} \theta _L)^T \} \mathrm{d}\mathrm{vec}(F). \\ \therefore \; \nabla _{\mathrm{vec}(F)} \theta _L&= \theta _G \otimes \nabla _{f} \theta _L. \end{aligned} \end{aligned}$$

Differentiating $\theta _L$ with respect to D,

$$\begin{aligned} \begin{aligned} \mathrm{d}\theta _L&= -C_2^{-T} \mathrm{d}D C_1^{-T} s_1 \\&= - (s_1^T C_1^{-1} \otimes C_2^{-T}) \mathrm{d}\mathrm{vec}(D). \\ \therefore \; \nabla _{\mathrm{vec}(D)} \theta _L&= - (C_1^{-T} s_1 \otimes C_2^{-1}). \end{aligned} \end{aligned}$$

Differentiating $\theta _L$ with respect to $\mu _1$,

$$\begin{aligned} \begin{aligned} \mathrm{d}\theta _L&= (\nabla _{f} \theta _L)^T F \mathrm{d}\mu _1 \\ \therefore \; \nabla _{\mu _1} \theta _L&= F^T (\nabla _{f} \theta _L). \end{aligned} \end{aligned}$$

Differentiating $\theta _L$ with respect to $v(C_1)$,

$$\begin{aligned} \begin{aligned} \mathrm{d}\theta _L&= -C_2^{-T}\mathrm{d}(C_2^T) C_2^{-T}(s_2 - DC_1^{-T} s_1) \\&\quad - C_2^{-T} D \mathrm{d}(C_1^{-T}) s_1 \\&= (\nabla _{f} \theta _L)^T F\mathrm{d}(C_1^{-T}) s_1 - C_2^{-T} D \mathrm{d}(C_1^{-T}) s_1 \\&= \{ (\nabla _{f} \theta _L)^T F - C_2^{-T} D\} (\nabla _{v(C_1^*)} \theta _G )^T \mathrm{d}v(C_1^*) \\ \therefore \; \nabla _{v(C_1^*)} \theta _L&= \nabla _{v(C_1^*)}\theta _G \{ F^T \nabla _{f} \theta _L - D^T C_2^{-1} \} \\&= \nabla _{v(C_1^*)}\theta _G \{ \nabla _{\mu _1} \theta _L - D^T C_2^{-1} \}. \end{aligned} \end{aligned}$$

Since $s_1 = C_1^T(\theta _G - \mu _1)$ and $s_2 = C_2^T (\theta _L - \mu _2)$, we have

$$\begin{aligned} \begin{aligned} \log q_\lambda (\theta )&= \log q(\theta _G) + \log q(\theta _L|\theta _G) \\&= -\frac{G}{2} \log (2\pi ) + \log |C_1| \\&\quad - \frac{1}{2}(\theta _G - \mu _1)^T C_1 C_1^T (\theta _G - \mu _1) \\&\quad -\frac{nL}{2} \log (2\pi ) + \log |C_2| \\&\quad - \frac{1}{2}(\theta _L - \mu _2)^T C_2 C_2^T (\theta _L - \mu _2) \\&= - \frac{nL+G}{2} \log (2\pi ) + \log |C_1C_2| - \frac{1}{2} s^T s. \end{aligned} \end{aligned}$$

As $\mu _2 = d + C_2^{-T} D(\mu _1 - \theta _G)$ and $v(C_2^*) = f + F \theta _G$, differentiating $\log q_\lambda (\theta )$ with respect to $\theta _G$,

$$\begin{aligned} \begin{aligned} \mathrm{d}\log q_\lambda (\theta )&= - (\theta _G - \mu _1)^T C_1 C_1^T \mathrm{d}\theta _G - (\theta _L - \mu _2)^T C_2 C_2^T(-\mathrm{d}\mu _2)\\&\quad - (\theta _L - \mu _2)^T \mathrm{d}C_2 s_2 + \mathrm{tr}(C_2^{-1} \mathrm{d}C_2) \\&= -s_1^T C_1^T \mathrm{d}\theta _G + s_2^T C_2^T\{- C_2^{-T} D \mathrm{d}\theta _G \\&\quad + \mathrm{d}(C_2^{-T}) D(\mu _1 - \theta _G)\} \\&\quad - \mathrm{vec}(C_2^{-T} s_2 s_2^T)^T \mathrm{d}\mathrm{vec}(C_2) + \mathrm{vec}(C_2^{-T})^T d\mathrm{vec}(C_2) \\&= \mathrm{vec}(C_2^{-T} - \{C_2^{-T} s_2 + (\mu _2 -d)\}s_2^T)^T d\mathrm{vec}(C_2) \\&\quad -s_1^T C_1^T \mathrm{d}\theta _G - s_2^T D \mathrm{d}\theta _G \\&= \mathrm{vec}(C_2^{-T} - (\theta _L -d) s_2^T)^T E_{nL}^T D_2^* F \mathrm{d}\theta _G \\&\quad -s_1^T C_1^T \mathrm{d}\theta _G - s_2^T D \mathrm{d}\theta _G. \end{aligned} \end{aligned}$$

Therefore

$$\begin{aligned} \begin{aligned} \nabla _{\theta _G} \log q_\lambda (\theta )&=F^T D_2^* v(C_2^{-T} - (\theta _L -d) s_2^T) \\&\quad - C_1 s_1 - D^T s_2. \end{aligned} \end{aligned}$$

Note that $D_2^* v(C_2^{-T}) = v(I_{nL})$ as $C_2^{-T}$ is upper triangular and $v(C_2^{-T})$ only retains the diagonal elements of $C_2^{-T}$.

Differentiating $\log q_\lambda (\theta )$ with respect to $\theta _L$,

$$\begin{aligned} \begin{aligned} \mathrm{d}\log q_\lambda (\theta )&= - (\theta _L - \mu _2)^T C_2 C_2^T \mathrm{d}\theta _L \\&= - s_2^T C_2^T \mathrm{d}\theta _L. \\ \therefore \; \nabla _{\theta _L} \log q_\lambda (\theta )&= - C_2 s_2. \end{aligned} \end{aligned}$$

Appendix B: Gradients for generalized linear mixed models

Since $\theta = [\beta ^T, \omega ^T, {\tilde{b}}_1^T, \dots , {\tilde{b}}_n^T]^T$, we require

$$\begin{aligned}&\nabla _\theta \log p(y, \theta ) = [\nabla _\beta \log p(y, \theta ), \nabla _\omega \log p(y, \theta ), \\&\quad \nabla _{{\tilde{b}}_1} \log p(y, \theta ), \dots , \nabla _{{\tilde{b}}_n} \log p(y, \theta )]^T. \end{aligned}$$

For the centered parametrization, the components in $\nabla _\theta \log p(y, \theta )$ are given below. Note that $\beta = [\beta _{RG_1}^T, \beta _{G_2}^T]^T$.

$$\begin{aligned} \nabla _{\beta _{G_2}} \log p(y, \theta )= & {} \sum _{i=1}^n {X_i^{G_2}}^T \{ y_i - h'(\eta _i) \} - \beta _{G_2}/\sigma _\beta ^2. \\ \nabla _{\beta _{RG_1}} \log p(y, \theta )= & {} \sum _{i=1}^n C_i^T W W^T ({\tilde{b}}_i - C_i \beta _{RG_1}) \\&- \beta _{RG_1}/\sigma _\beta ^2. \end{aligned}$$

Differentiating $\log p(y, \theta )$ with respect to $\omega$,

$$\begin{aligned} \begin{aligned} \mathrm{d}\log p(y, \theta )&= -\sum _{i=1}^n ({\tilde{b}}_i - C_i \beta _{RG_1})^T \mathrm{d}W W^T ({\tilde{b}}_i - C_i \beta _{RG_1}) \\&\quad + n \mathrm{tr}(W^{-1} \mathrm{d}W) - \omega ^T \mathrm{d}\omega /\sigma _\omega ^2 \\&= \mathrm{vec}\bigg \{ - \sum _{i=1}^n ({\tilde{b}}_i - C_i \beta _{RG_1}) ({\tilde{b}}_i - C_i \beta _{RG_1})^T W \\&\quad + nW^{-T} \bigg \}^T E_L^T D_L^* \mathrm{d}\omega - \omega ^T \mathrm{d}\omega /\sigma _\omega ^2, \end{aligned} \end{aligned}$$

where $\mathrm{d}v(W) = D^*_L \mathrm{d}\omega$ and $D^*_L = \mathrm{diag}\{ v(\mathrm{dg}(W) + \mathbf{1}_L\mathbf{1}_L^T - I_L) \}$. Hence

$$\begin{aligned} \begin{aligned} \nabla _\omega \log p(y, \theta )&=- D^*_L \sum _{i=1}^n v\{ ({\tilde{b}}_i - C_i \beta _{RG_1}) ({\tilde{b}}_i \\&\quad - C_i \beta _{RG_1}) ^TW \} \\&\quad + n v(I_L) - \omega /\sigma _\omega ^2. \end{aligned} \end{aligned}$$

Note that $D_L^* v(W^{-T}) = v(I_L)$ because $W^{-T}$ is upper triangular and $v(W^{-T})$ only retains the diagonal elements.

$$\begin{aligned} \nabla _{{\tilde{b}}_i} \log p(y, \theta ) = Z_i^T \{ y_i - h'(\eta _i)\} - W W^T ({\tilde{b}}_i - C_i \beta _{RG_1}). \end{aligned}$$

Appendix C: Gradients for state space models

Since $\theta = [\alpha , \kappa , \psi , b_1^T, \dots , b_n^T]^T$, we require

$$\begin{aligned}&\nabla _\theta \log p(y, \theta ) = [\nabla _\alpha \log p(y, \theta ), \nabla _\kappa \log p(y, \theta ), \\&\quad \nabla _\psi \log p(y, \theta ), \nabla _{b_1} \log p(y, \theta ), \dots , \nabla _{b_n} \log p(y, \theta )]^T. \end{aligned}$$

The components in $\nabla _\theta \log p(y, \theta )$ are given below.

$$\begin{aligned} \nabla _\alpha \log p(y, \theta )= & {} \frac{1}{2}\sum ^n_{i=1} (b_i y_i^2 \mathrm{e}^{-\sigma b_i - \kappa } - b_i)(1-\mathrm{e}^{-\sigma }) - \frac{\alpha }{\sigma _{\alpha }^2}. \\ \nabla _\kappa \log p(y, \theta )= & {} \frac{1}{2} \bigg (\sum ^n_{i=1} y_i^2 \mathrm{e}^{-\sigma b_i - \kappa } - n \bigg ) - \kappa /\sigma _{\kappa }^2.\\ \nabla _\psi \log p(y, \theta )= & {} \bigg \{ \sum _{i=2}^{n} (b_i - \phi b_{i-1})b_{i-1} + b_1^2 \phi - \frac{\phi }{1-\phi ^2} \bigg \} \\&\quad \times \phi (1-\phi ) - \psi /\sigma _{\psi }^2.\\ \nabla _{b_1} \log p(y, \theta )= & {} \frac{\sigma }{2} (y_1^2 \mathrm{e}^{-\sigma b_1 - \kappa } - 1) \\&+ \phi (b_2 - \phi b_1)- b_1 (1-\phi )^2. \end{aligned}$$

For $2 \le i \le n-1$,

$$\begin{aligned} \begin{aligned} \nabla _{b_i} \log p(y, \theta )&= \frac{\sigma }{2} (y_i^2 \mathrm{e}^{-\sigma b_i - \kappa } - 1) +\phi (b_{i+1} - \phi b_i)\\&\quad - (b_i - \phi b_{i-1}). \\ \nabla _{b_n} \log p(y, \theta )&= \frac{\sigma }{2} (y_n^2 \mathrm{e}^{-\sigma b_n - \kappa } - 1) - (b_n - \phi b_{n-1}). \end{aligned} \end{aligned}$$