Is VAE based on mean-field assumption?

Question

According to the slides (section 6), mean-field variational inference (aka. MFVI) assumes the latent variables ($z=\{z_1,..,z_m\}$ are independent from each other, and on top of this assumption, we can derive the update rule for each latent variable distribution $p(z_i)$ using Lagrange multiplier, and then iteratively update all $p(z)$ using coordinate ascent algorithm.

Given this definition of MFVI, I wonder if VAE is also using mean-field assumption, because in VAE each data-point $x_i$ has its own latent variable $z_i$, and different data-points' latent variables are independent, say $(x_i, z_i)$ has nothing to do with $(x_j, z_j)$.

I'm not very sure if this understanding correct, since in VAE the latent variable $z_i$ is more like the local latent variable, whereas in the MFVI they are global?

Answer 1

It is not true that in a VAE, every data point "has its own latent variable". The same latent variables describe every data point. What is different per data point is our beliefs about the values of these variables.

For instance, say we have a VAE that performs inference on a generative model of images. An image is a vector $x$ (e.g. a list of the RGB-values of its pixels) and assumed to be generated from a set of latent causes $z$ (also a vector). The generative model is given by $p(x,z)=p(x|z)p(z)$, and the true posterior would be $p(z|x)=\frac{p(x|z)p(z)}{p(x)}$, but this is difficult to compute in practice and so we approximate it using a variational distribution $q(z|x)$, which we choose to have a simple form that will allow us to work with it easily.

You can find the parameters of $q(z|x)$ (e.g., the means and (co-)variances, if it's Gaussian) in various ways. One way is to do iterative optimization, like the coordinate ascent scheme you described. Another way is to train a neural network to output the parameters of $q(z|x)$, by running the image $x$ through the network. This is called amortized inference, because we use the same neural network (with the same neural network weights) for every image. Instead of a (possibly lengthy) optimization procedure, we use a machine that takes in the image and gives us the answer. Or at least, something close to the optimal answer, as VAEs aren't guaranteed to produce the optimal parameters of $q(z|x)$ - they are just trained to do a good job which hopefully also generalizes to new data (as with any neural net).

The important point is: you get a different $q(z|x)$ for any given image, but these are beliefs about the same set of latent variables, since every image is assumed to be caused by the same causal variables (e.g. the positions, colors, shapes etc. of objects in a scene). For instance, every image might have a handwritten digit in it, but the number might be different, it might be in a different position in the image, it might have a different stroke width, a different slant, etc..

How does the mean-field assumption enter into this? It is essentially orthogonal to this whole story. As with any variational inference problem, you can choose $q(z|x)$ to have whatever form you like. However, the typical situation with VAEs is that the variational beliefs over the latent variables are chosen to be independent. That is, $q(z|x)=\prod_i{q(z_i|x)}$. And that is precisely the mean-field assumption.

Answer 2

Mean-Field Variational Inference

In general mean-field variational inference, there is no distinction between local and global variables. The main elements are:

1. Generative Model.

$$ p(z,x) \tag{1a} $$

where the latent variables $z$ can be partitioned into $M$ components, $z = \{z_1, \dots, z_m \}$. Each $z_i$ may be multi-dimensional, but we typically assume it is a one-dimensional scalar.

2. Variational Family. Given observed data $x$, we approximate $p(z\mid x)$ using mean-field variational inference by optimizing over all distributions $q$ that take the form

$$ q(z;x) = \prod_{i=1}^{M}q_{i}(z_i;x) \tag{2a} $$

where I have included $x$ to highlight the fact that the variational distribution depends on the observed data. Let $Q$ be the collection of all probability distributions $q$ over $z$ that factorize via mean-field as in Eq (2).

3. Optimization Objective.

$$ \min_{ q \in Q} \mathrm{KL}(q(z;x) \mid\mid p(z|x)). \tag{3a} $$

Using the calculus of variations, it can be shown that the optimal mean-field distribution $q^*(z)$ satisfies

$$ \begin{align} q^*_i(z_i;x) \propto \exp\left\{ \mathbb{E}_{\prod_{j=1,j\ne l}^M q^*(z_j;x)}\left[ \log p(z, x) \right] \right\} && (i=1,\dots,m). \end{align} $$

The above equation tells us how to compute the optimal marginals $q^*(z_i)$ in terms of all the other marginals. If we are lucky enough to be able to compute the above expectations analytically, then we can find $q^*$ numerically by using coordinate ascent. See Bishop (2006) for convergence conditions of coordinate ascent.

Variational Auto-Encoder / Amortized Inference

Variational auto-encoders make more assumptions than the mean-field case stated above. There are two main specializations.

1. Generative Model.

$$ p_\theta(z,x) = \prod_{i=1}^m p_\theta(z_i)p_\theta(x_i | z_i). \tag{1b} $$

In contrast to the general form in Eq (1a), in Eq (1b) we are treating $\theta$ as "global" parameters (i.e., not latent variables) and setting $z_i$ to be the "local" latent variables. Here, each $i = 1,\dots,m$ indexes a data item, and each $z_i$ is typically multi-dimensional. Note that we are treating $\theta$ in a "non-Bayesian" way.

2. Variational Family.

$$ q(z;x,\phi) = \prod_{i=1}^mq(z_i; x_i, \phi) \tag{2b} $$

again factorizes (i.e., mean-field). However, the fundamental difference between Eq (2b) and Eq (2a) is that the factors $q(z_i; x_i, \phi)$ are all parameterized by a shared parameter $\phi$ (typically the weights of a neural network). Unlike the "parameter-free" Eq (2a), we do not specify a separate marginal $q_i$ for each $z_i$ $(i=1,\dots,m)$ but we instead learn the parameter $\phi$, which gives us the approximate distribution over all $z_1, \dots, z_m$ at once.

3. Optimization Objective.

$$ \min_{\phi}\mathrm{KL}(q(z;x,\phi) \mid\mid p_\theta(z|x)). \tag{3b} $$

Eq (3b) is again over more restricted family than Eq (3a). We have assumed above that $\theta$ is known and fixed, but we typically also optimize the parameters using the objective

$$ \min_{\theta,\phi}\mathrm{KL}(q(z;x,\phi) \mid\mid p_\theta(z|x)) $$

which is also a special case of Eq (3a) where $\theta$ is a latent variable with a uniform (improper prior) and the marginal $q(\phi)$ of the latent $\phi$ is constrained to be a dirac delta.