Variational Inference

Latent Variable Model

We have data points $x$, we want to model a distribution of $x$ with the help of latent variable $z$.

How do we train latent variable models?

• Model $p_\theta(\vec{x})$

• Data: $D = \{ x_1, x_2, x_3, \dots, x_N \}$

• MLE fit: $\theta \leftarrow \argmax_{\theta} \frac{1}{N} \sum_i \log p_\theta(x_i)$
  • $p(x) = \int p(x|z) p(z) dz$
  • $\theta \leftarrow \argmax_\theta \frac{1}{N} \sum_i \log (\underbrace{\int p_\theta(x_i|z)p(z)dz}_{\text{completely untractable}})$

Estimating the log-likelihood

Intuition:

• Guess most likely $z$ given $x_i$, and pretend it is the right one
  • But there are many possible values of $z$ so use the distribution $p(z|x_i)$

• But how do we calculate $p(z|x_i)$?
  • Approximate with $q_i(z) = N(\mu_i, \sigma_i)$
    • Note that with each data point $x_i$, they all have different $q_i$
    • With any $q_i(z)$, we can construct a lower bound on $\log p(x_i)$
    • Maximize this bound (assume this bound is tight), we can push up on $\log p(x_i)$

Note:

Yensen’s Inequality (For concave function)

So

Note about entropy:

Intuitions:

• How random is the random variable?

• How large is the log probability in expectation under itself?

Also note about KL-Divergence:

Intuitions:

• How different are two distributions?

• How small is the expected log probability of one distribution under another, minus entropy?

So the variational approximation

So what makes a good $q_i(z)$?

• $q_i(z)$ should approximate $p(z|x_i)$
  • Compare in terms of KL-divergence $D_{KL}(q_i(z) || p(z|x))$

Why?

Also:

In fact minimizing the KL-divergence has the effect of tightening then bound!

But how to improve $q_i$?

• We can let $q_i = N(\mu_i, \sigma_i)$ and use gradient $\nabla_{\mu_i} L_i(p,q_i)$ and $\nabla_{\sigma_i} L_i(p,q_i)$
  • But too many parameters! We have $|\theta| + (|\mu_i|+|\sigma_i|) \times N$ parameters

• Then can we use
  • $q_i(z) = q_\phi(z|x_i) \approx p(z|x_i)$?
  • Amortized Variational Inference!

Amortized Variational Inference

We can directly calculate the derivative for the entropy term but the first term $J(\phi)$ seems a bit problematic.

But this format looks a lot like policy gradient!

What’s wrong with this gradient?

• Perfectly viable approach, but not best approach

• Tends to suffer from high variance

The reparameterization trick

This makes $z$ an independent function of a random variable $\epsilon$ that is independent of $\phi$

So now:

And now to estimate $\nabla_{\phi} J(\phi)$

• sample $\epsilon_1, \dots, \epsilon_M$ from $N(0,1)$ (even a single sample works well)
  • $\epsilon$ treated as a constant when computing gradient

• $\nabla_\phi J(\phi) \approx \frac{1}{M} \sum_j \nabla_\phi r(x_i, \mu_\phi(x_i) + \epsilon_j \sigma_\phi(x_i))$

This is useful because now we are using the derivative of the $r$ function

Another Perspective of $L_i$

Reparameterization trick vs. Policy Gradient

• Policy Gradient
  • Can handle both discrete and continuous latent variables
  • High variance, requires multiple samples & small learning rates

• Reparameterization Trick
  • Only continuous latent variables
  • Simple to implement & Low variance

Example Models

VAE Variational Autoencoder

Conditional Models

Generates $y$ given $x, z$