Neural Variational Inference: Classical Theory

As a member of Bayesian methods research group I’m heavily interested in Bayesian approach to machine learning. One of the strengths of this approach is ability to work with hidden (unobserved) variables which are interpretable. This power however comes at a cost of generally intractable exact inference, which limits the scope of solvable problems.

Another topic which gained lots of momentum in Machine Learning recently is Deep Learning, of course. With Deep Learning we can now build big and complex models that outperform most hand-engineered approaches given lots of data and computational power. The fact that Deep Learning needs a considerable amount of data also requires these methods to be scalable — a really nice property for any algorithm to have, especially in a Big Data epoch.

Given how appealing both topic are it’s not a surprise there’s been some work to marry these two recently. In this series of blogsposts I’d like to summarize recent advances, particularly in variational inference. This is not meant to be an introductory discussion as prior familiarity with classical topics (Latent variable models, Variational Inference, Mean-field approximation) is required, though I’ll introduce these ideas anyway just to remind it and setup the notation.

Latent Variables Models

Suppose you have a probabilistic model that’s easy to describe using some auxiliary variables Z$Z$ that you don’t observe directly (or even would like to infer it given the data). One classical example for this setup is Gaussian Mixture Modeling: we have K$K$ components in a mixture, and zn$z_n$ is a one-hot vector of dimensionality K$K$ indicating which component an observation xn$x_n$ belongs to. Then, conditioned on zn$z_n$ the distribution of xn$x_n$ is a usual Gaussian distribution: p(xn∣znk=1)=N(xn∣μk,Σk)$p(x_n\mid z_{nk}=1)=\mathcal{N}(x_n\mid {\mu }_k,{\mathrm{\Sigma }}_k)$ (here whenever I refer to a distribution, you should read it as its density. At least generalized one). Therefore the joint distribution of the model is

p(x,z∣Θ)=∏n=1N∏k=1KN(xn∣μk,Σk)znkπznkk$$p(x,z\mid \mathrm{\Theta })=\prod _{n=1}^N\prod _{k=1}^K\mathcal{N}(x_n\mid {\mu }_k,{\mathrm{\Sigma }}_k{)}^{z_{nk}}{\pi }_k^{z_{nk}}$$

Where π$\pi$ is a probability distribution over K$K$ outcomes, and Θ$\mathrm{\Theta }$ is a set of all model’s parameters (π$\pi$, μ$\mu$s and Σ$\mathrm{\Sigma }$s).

We’d like to do 2 things with the model: first, we obviously need to learn parameters Θ$\mathrm{\Theta }$, and second, we’d like infer these latent variables zn$z_n$ to know which cluster the observation xn$x_n$ belongs to, that is, we need to calculate the distribution of zn$z_n$ conditioned on xn$x_n$: p(zn∣xn)$p(z_n\mid x_n)$.

We want to learn the parameters Θ$\mathrm{\Theta }$ as usual by maximizing the log-likelihood. Unfortunately, we don’t know true assignments zn$z_n$, and marginalizing over it as in p(xn)=∑Kk=1πkp(xn,znk=1)$p(x_n)=\sum _{k=1}^K{\pi }_{k}p(x_n,z_{nk}=1)$ is not a good idea as the resulting optimization problem would lose its convexity. Instead we decompose the log-likelihood as follows:

logp(x)=Eq(z∣x)logp(x)const in z=Eq(z∣x)logp(x,z)q(z∣x)p(z∣x)q(z∣x)=Eq(z∣x)logp(x,z)q(z∣x)+DKL(q(z∣x)∣∣p(z∣x))$$\begin{array}{rl}\log ?p(x)& ={\mathbb{E}}_{q(z\mid x)}\stackrel{\text{const in }z}{\stackrel{?}{\log ?p(x)}}={\mathbb{E}}_{q(z\mid x)}\log ?\frac{p(x,z)q(z\mid x)}{p(z\mid x)q(z\mid x)}\\ & ={\mathbb{E}}_{q(z\mid x)}\log ?\frac{p(x,z)}{q(z\mid x)}+D_{KL}(q(z\mid x)\mid \mid p(z\mid x))\end{array}$$

The second term is a Kullback-Leibler divergence, which is always non-negative, and equals zero iff distributions are equal almost everywhere q(z∣x)=p(z∣x)$q(z\mid x)=p(z\mid x)$. Therefore putting q(z∣x)=p(z∣x)$q(z\mid x)=p(z\mid x)$ eliminates the second term, leaving us with logp(x)=Ep(z∣x)logp(x,z)p(z∣x)$\log ?p(x)={\mathbb{E}}_{p(z\mid x)}\log ?\frac{p(x,z)}{p(z\mid x)}$. Therefore all we need to be able to do is to calculate the posterior p(z∣x)$p(z\mid x)$, and maximize the expectation. This is how EM algorithm is derived: at E-step we calculate the posterior p(z∣x,Θold)$p(z\mid x,{\mathrm{\Theta }}^{\text{old}})$, and at M-step we maximize the expectation Ep(z∣x,Θold)logp(x,z∣Θ)p(z∣x,Θ)${\mathbb{E}}_{p(z\mid x,{\mathrm{\Theta }}^{\text{old}})}\log ?\frac{p(x,z\mid \mathrm{\Theta })}{p(z\mid x,\mathrm{\Theta })}$ with respect to Θ$\mathrm{\Theta }$ keeping Θold${\mathrm{\Theta }}^{\text{old}}$ fixed.

Now, all we are left to do is to find the posterior p(z∣x)$p(z\mid x)$ which is given by the following deceivingly simple formula knows as a Bayes’ rule.

p(z∣x)=p(x∣z)p(z)∫p(x∣z)p(z)dz$$p(z\mid x)=\frac{p(x\mid z)p(z)}{\int p(x\mid z)p(z)dz}$$

Of course, there’s no free lunch and computing the denominator is intractable in general case. One can compute the posterior exactly when the prior p(z)$p(z)$ and the likelihood p(x∣z)$p(x\mid z)$ are conjugate (that is, after multiplying the prior by the likelihood you get the same functional form for z$z$ as in the prior, thus the posterior comes from the same family as the prior but with different parameters), but many models of practical interest don’t have this property. This is where variational inference comes in.

Variational Inference and Mean-field

In a variational inference (VI) framework we approximate the true posterior p(z∣x)$p(z\mid x)$ with some other simpler distribution q(z∣x,Λ)$q(z\mid x,\mathrm{\Lambda })$ where Λ$\mathrm{\Lambda }$ is a set of (variational) parameters of the (variational) approximation (I may omit Λ$\mathrm{\Lambda }$ and Θ$\mathrm{\Theta }$ in a “given” clause when it’s convenient, remember, it always could be placed there). One possibility is to divide latent variables z$z$ in groups and force the groups to be independent. In extreme case each variable gets its own group, assuming independence among all variables q(z∣x)=∏Dd=1q(zd∣x)$q(z\mid x)=\prod _{d=1}^{D}q(z_d\mid x)$. If we then set about to find the best approximation to the true posterior in this fully factorized class, we will no longer have the optimal q$q$ being the true posterior itself, as the true posterior is presumably too complicated to be dealt with in analytic form (which we want from the approximation q$q$ when we say “simpler distribution”). Therefore we find the optimal q(zi)$q(z_i)$ by minimizing the KL-divergence with the true posterior (const$\text{const}$ denotes terms that are constant w.r.t. q(zi)$q(z_i)$):

DKL(q(z∣x)∣∣p(z∣x))=Eq(zi∣x)[Eq(z?i∣x)logq(z1∣x)…q(zD∣x)p(z∣x)]=Eq(zi∣x)????logq(zi∣x)?Eq(z?i∣x)logp(z∣x)logf(zi∣x)????+const=Eq(zi∣x)[logq(zi∣x)1Zf(zi∣x)]?logZ+const=DKL(q(zi∣x)∣∣1Zf(zi∣x))+const$$\begin{array}{rl}{D}_{KL}(q(z\mid x)\mid \mid p(z\mid x))& ={\mathbb{E}}_{q(z_i\mid x)}[{\mathbb{E}}_{q(z_{?i}\mid x)}\log ?\frac{q(z_1\mid x)\dots q(z_D\mid x)}{p(z\mid x)}]\\ & ={\mathbb{E}}_{q(z_i\mid x)}[\log ?q(z_i\mid x)?\underset{\log ?f(z_i\mid x)}{\underset{?}{{\mathbb{E}}_{q(z_{?i}\mid x)}\log ?p(z\mid x)}}]+\text{const}\\ & ={\mathbb{E}}_{q(z_i\mid x)}[\log ?\frac{q(z_i\mid x)}{\frac{1}{Z}f(z_i\mid x)}]?\log ?Z+\text{const}\\ & =D_{KL}(q(z_i\mid x)\mid \mid \frac{1}{Z}f(z_i\mid x))+\text{const}\end{array}$$

For many models it’s possible to look into Eq(z?i∣x)logp(z∣x)${\mathbb{E}}_{q(z_{?i}\mid x)}\log ?p(z\mid x)$ and immediately recognize logarithm of unnormalized density of some known distribution.

Another cornerstone of this framework is a notion of Evidence Lower Bound (ELBO): recall the decomposition of log-likelihood we derived above. In our current setting we can not compute the right hand side as we can not evaluate the true posterior p(z∣x)$p(z\mid x)$. However, note that the left hand side (that is, the log-likelihood) does not depend on the variational distribution q(z∣x,Λ)$q(z\mid x,\mathrm{\Lambda })$. Therefore, maximizing the first term of the right hand side w.r.t. variational parameters Λ$\mathrm{\Lambda }$ results in minimizing the second term, the KL-divergence with the true posterior. This implies we can ditch the second term, and maximize the first one w.r.t. both model parameters Θ$\mathrm{\Theta }$ and variational parameters Λ$\mathrm{\Lambda }$:

ELBO:L(Θ,Λ):=Eq(z∣x,Λ)logp(x,z∣Θ)q(z∣x,Λ)$$\text{ELBO:}\mathcal{L}(\mathrm{\Theta },\mathrm{\Lambda }):={\mathbb{E}}_{q(z\mid x,\mathrm{\Lambda })}\log ?\frac{p(x,z\mid \mathrm{\Theta })}{q(z\mid x,\mathrm{\Lambda })}$$

Okay, so this covers the basics, but before we go to the neural networks-based methods we need to discuss some general approaches to VI and how to make it scalable. This is what the next blog post is all about.