Variational Inference

In Bayesian inference one usually specifies a model as follows: given data ,

where denotes that the samples are identically independently distributed. Our goal in Bayesian inference is then to find the posterior

In general one cannot obtain a closed form expression for , but one might still be able to sample from with guarantees of converging to the target posterior as the number of samples go to , e.g. MCMC.

As you are hopefully already aware, Turing.jl provides a lot of different methods with asymptotic exactness guarantees that we can apply to such a problem!

Unfortunately, these unbiased samplers can be prohibitively expensive to run. As the model increases in complexity, the convergence of these unbiased samplers can slow down dramatically. Still, in the infinite limit, these methods should converge to the true posterior! But infinity is fairly large, like, at least more than 12, so this might take a while.

In such a case it might be desirable to sacrifice some of these asymptotic guarantees, and instead approximate the posterior using some other model which we’ll denote .

There are multiple approaches to take in this case, one of which is variational inference (VI).

In VI, we’re looking to approximate using some approximate or variational posterior .

To approximate something you need a notion of what "close" means. In the context of probability densities a standard such "measure" of closeness is the Kullback-Leibler (KL) divergence , though this is far from the only one. The KL-divergence is defined between two densities and as

It’s worth noting that unfortunately the KL-divergence is not a metric/distance in the analysis-sense due to its lack of symmetry. On the other hand, it turns out that minimizing the KL-divergence that it’s actually equivalent to maximizing the log-likelihood! Also, under reasonable restrictions on the densities at hand,

Therefore one could (and we will) attempt to approximate using a density by minimizing the KL-divergence between these two!

One can also show that , which we’ll need later. Finally notice that the KL-divergence is only well-defined when in fact is zero everywhere is zero, i.e.

Otherwise, there might be a point such that , resulting in which doesn’t make sense!

One major problem: as we can see in the definition of the KL-divergence, we need for any if we want to compute the KL-divergence between this and . We don’t have that. The entire reason we even do Bayesian inference is that we don’t know the posterior! Cleary this isn’t going to work. Or is it?!

First off, recall that

so we can write

where in the last equality we used the fact that is independent of .

Now you’re probably thinking "Oh great! Now you’ve introduced which we also can’t compute (in general)!". Woah. Calm down human. Let’s do some more algebra. The above expression can be rearranged to

See? The left-hand side is constant and, as we mentioned before, . What happens if we try to maximize the term we just gave the completely arbitrary name ? Well, if goes up while stays constant then has to go down! That is, the which minimizes the KL-divergence is the same which maximizes :

where

and denotes the (differential) entropy of .

Assuming joint and the entropy are both tractable, we can use a Monte-Carlo for the remaining expectation. This leaves us with the following tractable expression

where

Hence, as long as we can sample from somewhat efficiently, we can indeed minimize the KL-divergence! Neat, eh?

Sidenote: in the case where is tractable but is not , we can use an Monte-Carlo estimate for this term too but this generally results in a higher-variance estimate.

Also, I fooled you real good: the ELBO isn’t an arbitrary name, hah! In fact it’s an abbreviation for the expected lower bound (ELBO) because it, uhmm, well, it’s the expected lower bound (remember ). Yup.

Finding the optimal over all possible densities of course isn’t feasible. Instead we consider a family of parameterized densities where denotes the space of possible parameters. Each density in this family is parameterized by a unique . Moreover, we’ll assume

(1) is fairly straight-forward, but (2) is a bit tricky. What does it even mean for a sampling process to be differentiable? This is quite an interesting problem in its own right and would require something like a 50-page paper to properly review the different approaches (highly recommended read).

We’re going to make use of a particular such approach which goes under a bunch of different names: reparametrization trick, path derivative, etc. This refers to making the assumption that all elements can be considered as reparameterizations of some base density, say . That is, if then

for some function differentiable wrt. . So all are using the same reparameterization-function but each correspond to different choices of for .

Under this assumption we can differentiate the sampling process by taking the derivative of wrt. , and thus we can differentiate the entire wrt. ! With the gradient available we can either try to solve for optimality either by setting the gradient equal to zero or maximize stepwise by traversing in the direction of steepest ascent. For the sake of generality, we’re going to go with the stepwise approach.

With all this nailed down, we eventually reach the section on Automatic Differentiation Variational Inference (ADVI).

So let’s revisit the assumptions we’ve made at this point:

All of these are not necessary to do VI, but they are very convenient and results in a fairly flexible approach. One distribution which has a density satisfying all of the above assumptions except (7) (we’ll get back to this in second) for any tractable and differentiable is the good ole' Gaussian/normal distribution:

where with obtained from the Cholesky-decomposition. Abusing notation a bit, we’re going to write

With this assumption we finally have a tractable expression for ! Well, assuming (7) is holds. Since a Gaussian has non-zero probability on the entirety of , we also require to have non-zero probability on all of .

Though not necessary, we’ll often make a mean-field assumption for the variational posterior , i.e. assume independence between the latent variables. In this case, we’ll write