2
$\begingroup$

I am trying to implement a variational extension of some kind of Bayesian network estimation method. The main goal is to improve speed, since the current method is pretty slow due to MCMC.

My question should be pretty simple for veterans in the field: "How to obtain the optimal variational density?" "Is it fixed a priori? Or derived analytically from the data ?".

But before that, I think it's important to validate my understanding on what is variational inference and why we use it. This post it's an attempt to summarize what I think I understand. Correct me if I did any mistakes during the process.

DISCLAIMER: I am a frequenstist-trained biostatistician. This is my first bayesian attempt in research.

What I think I understand:

Assuming we are in a general framework (linear regression for example) with one vector of parameter $\beta$ and $\sigma$, each following a given distribution.
From a bayesian perspective we are looking for the posterior distribution of some parameters $\beta, \sigma$ given data $Y$ in order to do inferences. Formally:

$$ p(\beta, \sigma|Y) = \frac{p(\beta)p(\sigma)p(Y|\beta, \sigma)}{p(Y)} \propto p(\beta)p(\sigma)p(Y|\beta, \sigma) $$

Since the denominator is basically an integral over all parameters $\int_{\beta}\int_{\sigma} p(Y, \beta, \sigma) d_{\beta} d_{\sigma}$, in presence of lot of parameters, this density is often intractable.

One approach to estimate this quantity it's by sampling (MCMC), but this is not the focus here.

Another approach is to consider variational inference (VI). VI turns the approximation problem into an optimization one. Let's move into VI more formally and check if I understand all concepts in a good way.

In VI we consider a proxy parametric distribution for parameter $q(\beta, \sigma)$ which is easy to compute while enough flexible to capture the true posterior. Traditionally, researchers try to minimize some kind of metric in order to determine the best $q(.)$. The most common is the Kullback-Leibler divergence. This is just an asymptotic Likelihood ratio test statistic between $q(\beta, \sigma)$ and $p(\beta, \sigma|Y)$. Formally:

$$ KL(q(\beta, \sigma) || p(\beta, \sigma|Y)) = \int q(\beta, \sigma) \log(\frac{q(\beta, \sigma)}{p(\beta, \sigma|Y)})d_{\beta, \sigma} $$

From

$$ p(Y) = \int_{\beta}\int_{\sigma}p(Y, \beta, \sigma) d_{\beta}d_{\sigma} $$

$$ = \int_{\beta}\int_{\sigma} q(\beta, \sigma) \frac{p(Y, \beta, \sigma)} {q(\beta, \sigma)} d_{\beta}d_{\sigma} = \int_{\beta}\int_{\sigma} q(\beta, \sigma) \frac{p(\beta, \sigma |Y) p(Y)} {q(\beta, \sigma)} d_{\beta}d_{\sigma} $$

$$ \log(P(Y)) = \log(\int_{\beta}\int_{\sigma} q(\beta, \sigma) \frac{p(\beta, \sigma |Y) p(Y)} {q(\beta, \sigma)} d_{\beta}d_{\sigma}) >= \int_{\beta}\int_{\sigma} q(\beta, \sigma) \log(\frac{p(\beta, \sigma |Y) p(Y)} {q(\beta, \sigma)}) d_{\beta}d_{\sigma}$$

The last inequality is obtained from Jensen's inequality. The right term of the inequality is the evidence lower bound (ELBO). After simple algebraic manipulations, the ELBO can be re-expressed as:

$$ \int_{\beta}\int_{\sigma} q(\beta, \sigma) \log(\frac{p(\beta, \sigma |Y) p(Y)} {q(\beta, \sigma)}) d_{\beta}d_{\sigma} = \int_{\beta}\int_{\sigma} q(\beta, \sigma) (\log(p(\beta, \sigma |Y)) - \log(q(\beta, \sigma)) + \log(p(Y))) d_{\beta}d_{\sigma} \propto -KL(q(\beta,\sigma)||p(\beta, \sigma |Y)) $$

Since $p(Y)$ is constant over $q(\beta,\sigma)$, minimizing the KL is the same as maximizing the ELBO.

From the mean-field theory $q(\beta,\sigma)$ can be rewritten as $q(\beta) q(\sigma)$. Thus,

$$ ELBO = \int_{\beta}\int_{\sigma} q(\beta, \sigma) \log(\frac{p(\beta, \sigma |Y))}{q(\beta, \sigma)}) d_{\beta}d_{\sigma} = \int_{\beta}\int_{\sigma} q(\beta)q(\sigma) \log(\frac{p(\beta, \sigma |Y)}{\log(q(\beta) q(\sigma))} ) d_{\beta}d_{\sigma} $$

Rewriting the ELBO we obtain: $$ \int_{\beta}\int_{\sigma} q(\beta)q(\sigma) \log(\frac{(p(\beta)p(\sigma)P(Y|\beta, \sigma))}{q(\beta) q(\sigma)} ) d_{\beta}d_{\sigma} $$

$$ = \int_{\beta}\int_{\sigma} q(\beta)q(\sigma) \log(p(\beta))d_{\beta}d_{\sigma} + \int_{\beta}\int_{\sigma}q(\beta)q(\sigma) \log(p(\sigma))d_{\beta}d_{\sigma} + \int_{\beta}\int_{\sigma}q(\beta)q(\sigma) \log(p(Y|\beta, \sigma))d_{\beta}d_{\sigma} - \int_{\beta}\int_{\sigma}q(\beta)q(\sigma) \log(q(\beta))d_{\beta}d_{\sigma} - \int_{\beta}\int_{\sigma}q(\beta)q(\sigma) \log (q(\sigma)) d_{\beta}d_{\sigma} $$

EDIT: We can rewrite it in terms of expectations:

$$ E_{q(\beta, \sigma)}\log(p(\beta)) + E_{q(\beta, \sigma)}\log(p(\sigma)) + E_{q(\beta, \sigma)}\log(p(Y| \beta, \sigma)) - E_{q(\beta, \sigma)}\log(q(\beta)) - E_{q(\beta, \sigma)}\log(q(\sigma)) $$

From my understanding we have two general approaches to optimize the ELBO.

  1. Analytically, where each optimal $q(z)^*$ are defined for every particular problem. Following Blei et al., 2017 for the jth parameter $q_j(z_j)^* \propto \exp{E_{-j}(log(p(Z,X))}$. This quantity is the same as integrating out over the jth parameter keeping all the other parameters constant. But practically speaking I am not sure what does that mean.
  2. Approximately, where there is no context-specific optimal $q(z)^*$ but we use approximate distribution instead, such as Gaussian or any distribution closed to the prior. I know that in PyMC3 or Stan, Gaussian Mean-field approxiamtions are used. If I translate it correctly it's the same as replacing $q(z_j) \sim N(z_j, \mu_s, \sigma_s)$, where $\mu_s, \sigma_s$ are the corresponding variational parameters. I know that other kind of approximations are available depending on the context.

My questions

Now this is the difficult part for me. I am not sure which method is better in which context.

Is the Gaussian mean-field can be extended to other kind of parametric family ? I am thinking about Laplace distribution for example.

Also, I will appreciate if a high-level explanation of the algorithm needed to obtain the variational parameters can be provided. Some authors talk about "EM-like" algorithms. But I am not sure what is done at each iteration.

$\endgroup$
8
  • $\begingroup$ Are you asking how we go about choosing functional forms for the mean-field variational distributions $q(\beta)$ and $q(\sigma)$, or are you asking how the derivation is turned into an algorithm for arriving at an optimised $q^*(\beta)$ and $q^*(\sigma)$? This is not something I have experience with professionally, it's only something I've programmed for assignments, but I'd be happy to point out some of the gaps in your exposition and point you in the direction of pedagogical papers which I found really clear, if that's what you are looking for? $\endgroup$
    – microhaus
    May 28, 2023 at 4:53
  • $\begingroup$ @microhaus Actually both, a good way to choose distribution family for $q$ and the optimal $q^*$. Mathematically, the latter is just a matter of calculus. Where I miss the point is how to implement that by myself in any context. Some kind of general recipe (if possible) should be good for example. I have already read a lot on that point. If I miss something please provide references. Thanks. $\endgroup$ May 28, 2023 at 11:11
  • $\begingroup$ @microhaus Do you have any update on this ? Thanks $\endgroup$ May 29, 2023 at 12:17
  • 1
    $\begingroup$ @microhaus Thanks for your answers and pointing out my mathematical mistakes. I corrected the typos while specifying my questions, regarding what I think I understand on the $q()$. $\endgroup$ May 31, 2023 at 14:46
  • 1
    $\begingroup$ This may be of interest arxiv.org/pdf/2304.14251.pdf $\endgroup$
    – paul
    May 31, 2023 at 15:31

1 Answer 1

1
$\begingroup$

I think in this instance it would be best to give some high-level answers to your questions and direct you to appropriate sources for details.

Your original question was:

How do we select the functional forms of the variational distribution $q$?

Recall that we wish to consider a restricted family of distributions $\mathcal{Q}$ and that we wish to find a member of this family to minimise the Kullback-Leibler diveregence. Typically, this family needs to be both tractable and sufficiently expressive.

From Chapter 10: Approximate Inference in Bishop (2006), there are two ways of thinking about this. We could either:

  • Specify $\mathcal{Q}$ to be some known parametric distribution family, consisting of distributions $q(\mathbf{Z} ; \omega)$ indexed by a parameter $\omega$.
  • Specify $\mathcal{Q}$ to be the mean-field family of distributions.

In the second case, this distribution family is defined by the assumption that $q(\mathbf{Z})$ has the following factorisation structure

$$q(\mathbf{Z}) = \prod^M_{i=1} q_j(\mathbf{Z}_j)$$

Strictly, one does not explicitly specify the functional forms/parametric distribution families of each of the $q_j$. Variational inference only requires a restriction on $\mathcal{Q}$ that it is a mean-field family.

Instead, the parametric distribution family one should use for each $q_j$ emerges naturally when deriving the update equations for each of the variational factors.

This is a position that can be found in earlier presentations of variational inference e.g. in Bishop (2006).

In practice however, the 'natural emergence' of what distribution family to use for the $q_j$ is reliant on (conditional) conjugacy, and this is essential to being able to derive closed form updates for the original version of variational inference that uses a co-ordinate ascent algorithm for updating. You will need to have a look at the worked example of variational linear regression in Bishop (2006) to see this in action.

All that being said, in more modern presentations of the variational inference and its more developed variants in the literature, such as in Hoffman et al. (2013) and Blei et al. (2016), you will find that the setting assumed for co-ordinate ascent variational inference is that all the complete conditionals $p_j(\mathbf{Z}_j \vert \mathbf{Z}_{-j}, \mathbf{X})$ in the model are in the exponential family.

From Hoffman et al. (2013):

With the assumptions that we have made about the model and variational distribution—that each conditional is in an exponential family and that the corresponding variational distribution is in the same exponential family—we can optimize each coordinate in closed form.

In this case, you just specify the distribution family of each of the $q_j(\mathbf{Z}_j; \phi_j)$ to be the same as that of the corresponding complete conditional $p_j(\mathbf{Z}_j \vert \mathbf{Z}_{-j}, \mathbf{X}; \theta_j)$, and index the $q_j$ with its own set of variational parameters $\phi_j$.

How are the optimised variational distribution $q^*$ and the iterative update equations determined?

The review by Blei et al. (2017) is slightly terse for my liking. A similar more verbose derivation of the following update equations can be found in Bishop (2006):

$$q_j(\mathbf{Z}_j) \propto \exp\{\mathbb{E}_{-j}[\log p(\mathbf{Z}_j \vert \mathbf{Z}_{-j}, \mathbf{X})]\} \quad j = 1, \dots M$$

Where the expectation is with respect to the product of all variational distributions $\prod_{i \neq j} q_i(\mathbf{Z}_i)$ except that of the $j$th factor.

Briefly, each update equation is the result of maximising the evidence lower bound functional $\mathcal{L}(q)$ with respect to the $j$th variational distribution $q_j(\mathbf{Z}_j)$ whilst holding all other variational distributions $q_{i}(\mathbf{Z}_{i})$ for $i \neq j$ fixed.

Note that in this abstract derivation, there is no functional form/parametric distribution family specified for each of the $q_j$ in advance, and that maximisation occurs over all possible functional forms that $q_j$ can take, with the only caveat being that $\mathcal{Q}$ is the mean-field family.

Rewriting the maximisation of the evidence lower bound functional $\mathcal{L}(q)$, the update equations arise as solutions to the the minimisation of a negative Kullback-Leibler divergence between $q_j$ and $\tilde{p}(\mathbf{X}, \mathbf{Z}_j)$, with the latter defined as

$$\log \tilde{p}(\mathbf{X}, \mathbf{Z}_j) = \mathbb{E}_{i \neq j}[\log p(\mathbf{X}, \mathbf{Z})] + \text{const.}$$

Where again, the expectation is with respect to $\prod_{i \neq j} q_i(\mathbf{Z}_i)$.

Another way to derive the update equations is using context from the outset, and this kind of strategy is used in Blei et al. (2003)

Under the assumption that the complete conditionals $p_j(\mathbf{Z}_j \vert \mathbf{Z}_{-j}, \mathbf{X})$ in the model belong to the exponential family, then we can just assume that each of the variational distributions $q_j(\mathbf{Z}_j)$ are in the same parametric distribution family. Meaning that we can write each of the $q_j$ as $q_j(\mathbf{Z}_j ; \phi_j)$, indexing each variational distribution with its own variational parameter $\phi_j$

In this case, we now have functional forms for all the variational distributions $q_j(\mathbf{Z}_j; \phi_j)$, and all complete conditionals $p_j$. All we do now is to substitute these into the evidence lower bound, which then becomes a function of the model parameters $\boldsymbol{\Theta}$ and variational parameters $\boldsymbol{\Phi}$.

$$\mathcal{L}(\boldsymbol{\Theta}, \boldsymbol{\Phi}) = \mathbb{E}_{q(\mathbf{Z})}[\log p(\mathbf{X}, \mathbf{Z})] + \mathbb{E}_{q(\mathbf{Z})}[\log q(\mathbf{Z}) ]$$

The update equations for each of the variational parameters then arise by maximising $\mathcal{L}(\boldsymbol{\Theta}, \boldsymbol{\Phi})$ with respect to the variational parameters $\boldsymbol{\Phi}$.

The co-ordinate ascent algorithm then proceeds with an initialisation of the variational parameters $\boldsymbol{\Phi}^{t=0}$. Using what is derived above, at time step $t = 1$, you would update $\phi_1$, given $\phi_2 = \phi_2^{t=0}, \dots , \phi_M = \phi_M^{t=0}$. After this is done for all $M$ variational parameters, you would compute the evidence lower bound using $\boldsymbol{\Phi}^{t=1}$

You would then repeat until your evidence lower bound has converged to some local maximum at $\boldsymbol{\Phi}^*$. As the evidence lower bound is in general non-convex, you will need to re-initialise the algorithm using different starting values $\boldsymbol{\Phi}^{t=0}$ and repeat.

Depending on the coupling structure of the update equations and your motivations, there is also scope for maximising this evidence lower bound with respect to model parameters $\boldsymbol{\Theta}$ also, in which case you would alternate between updating all variational parameters $\boldsymbol{\Phi}$ in a 'variational E-step', followed by all model parameters in a 'variational M-step'.

I will appreciate if you can a provide comprehensive example as answer while highlighting important concepts which are surely missing.

The application of variational inference to Bayesian linear regression as a worked example is too long to place here, and others have done more comprehensively than I ever could. Instead, see

  • Chapter 10: Approximate inference of Bishop has a simple example of Bayesian linear regression.
  • For a more general worked example of Bayesian linear regression see Drugowitsch (2013) and the supplementary note by Rapela (2017)

Some suggestions:

  • Have a look at the theoretical derivations in the sources I've outlined above, but also work through the Bayesian linear regression worked examples.
  • Personally, when looking at the derivations in abstracto, it can be difficult to see why or how certain steps are being made, only a worked example will advance your understanding.
  • In particular, for Bayesian linear regression, you would have, in your notation, $\mathbf{Z} = \{\beta, \sigma \}$ and $\boldsymbol{X} = Y$.

References

Bishop, C. (2006). Pattern Recognition and Machine Learning. Springer New York.

Blei, D. M., Kucukelbir, A., & McAuliffe, J. D. (2017). Variational inference: A review for statisticians. Journal of the American Statistical Association, 112(518), 859–877. https://doi.org/10.1080/01621459.2017.1285773

Blei, D. Variational inference: Foundations and innovations. http://www.cs.columbia.edu/~blei/talks/Blei_VI_tutorial.pdf

Blei, D. M., Ng, A. Y., & Jordan, M. I. (2003). Latent Dirichlet allocation. Journal of Machine Learning Research, 3, 993–1022. https://doi.org/10.1162/jmlr.2003.3.4-5.993

Blei, D., Ranganath, R., & Mohamed, S. (2016). Variational inference: Foundations and modern methods. NIPS 2016 Tutorial.

Drugowitsch, J. (2013). Variational Bayesian inference for linear and logistic regression. arXiv preprint arXiv:1310.5438.

Hoffman, M. D., Blei, D. M., Wang, C., & Paisley, J. (2013). Stochastic variational inference. Journal of Machine Learning Research, 14(1), 1303–1347.

Rapela, J. (2017) Derivation of a Variational-Bayes linear regression algorithm using a Normal Inverse-Gamma generative model. https://sccn.ucsd.edu/~rapela/docs/vblr.pdf

$\endgroup$
9
  • $\begingroup$ Thanks a lot ! I still need a little clarification about one thing I can read. If I summarize your post for my own understanding, can we say that by using $q_{j}(Zj) \propto \exp{E_{−j}[\log p(Z_j|Z_{−j},X)]}, j=1,…M$ we will find the optimal parametric form for the jth variational distribution ? In other words, that the optimal q is data-driven ? Moreover, can an optimal form be always derived for the joint ? I ask this question since in some papers it seems that q is fixed a priori, having the same distribution as the prior. Is it possible or I misunderstand something ? $\endgroup$ Jun 5, 2023 at 18:43
  • $\begingroup$ I understand your follow-up questions, except the last one. Have you got a link to a sample paper for the last query? $\endgroup$
    – microhaus
    Jun 5, 2023 at 19:00
  • $\begingroup$ Yes of course. You can refer to this one : 10.1080/01621459.2022.2044827 or this one 10.1016/j.ajhg.2023.03.009. Thanks ! $\endgroup$ Jun 5, 2023 at 19:05
  • $\begingroup$ Unfortunately, I can't actually see any of the links as I am not currently affiliated with an educational institution nor are they are available on certain websites :-( $\endgroup$
    – microhaus
    Jun 5, 2023 at 19:21
  • $\begingroup$ The second one can be found on biorxiv preprint server with this link biorxiv.org/content/10.1101/2022.05.10.491396v1 See Section S1. $\endgroup$ Jun 5, 2023 at 19:42

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.