Latent Variable Models: VAEs

Equivalently, we maximize the log-likelihood

\[ \max_{p\in\mathcal{P}_{\mathbf{x},\mathbf{z}}} \sum_{\mathbf{x}\in \mathcal{D}} \log p(\mathbf{x})=\sum_{\mathbf{x}\in\mathcal{D}} \log \int p(\mathbf{x},\mathbf{z}) \text{d}\mathbf{z} \]

However, since \( \mathbf{z}\sim p(\mathbf{z}) \), it’s intractable for high-dimensional \( \mathbf{z} \) as it involves an integration. One option is to estimate the objective via Monte Carlo.

\[ \log p(\mathbf{x}) \approx \log \frac{1}{k} \sum_{i=1}^{k} p(\mathbf{x}|\mathbf{z}^{(i)}), \quad \text{where } \mathbf{z}^{(i)} \sim p(\mathbf{z}) \]

In practice, however, it suffers from high variance in gradient estimates.

We notice that the marginal likelihood \( p(\mathbf{x}) \) is at least as hard as evaluating the posterior \( p(\mathbf{z}|\mathbf{x}) \) since \( p(\mathbf{z}|\mathbf{x}) = p(\mathbf{x},\mathbf{z}) / p(\mathbf{x}) \)

So, we introduce a variantional family \( \mathcal{Q} \) of distributions to approximate the true posterior \( p(\mathbf{z}|\mathbf{x}) \). Suppose the model family \( \mathcal{P}_{\mathbf{x},\mathbf{z}} \) is specified by a set of parameters \( \theta \in \Theta \), while the variantional family \( \mathcal{Q} \) is specified by parameters set \( \lambda \in \Lambda \).

Given \( \mathcal{P}_{\mathbf{x},\mathbf{z}} \) and \( \mathcal{Q} \), we can derive that

\begin{split} \log p_{\theta}(\mathbf{x}) &= \log \int p_{\theta}(\mathbf{x},\mathbf{z}) \, \text{d}\mathbf{z} \\ &= \log \int \frac{q_{\lambda}(\mathbf{z})}{q_{\lambda}(\mathbf{z})} p_{\theta}(\mathbf{x},\mathbf{z})\, \text{d}\mathbf{z}\\ & \ge \int q_{\lambda}(\mathbf{z}) \log \frac{p_{\theta}(\mathbf{x},\mathbf{z})}{q_{\lambda}(\mathbf{z})} \,\text{d}\mathbf{z}\\ &=\mathbb{E}_{\mathbf{z}\sim q_{\lambda}(\mathbf{z})}\Big[ \log \frac{p_{\theta}(\mathbf{x},\mathbf{z})}{q_{\lambda}(\mathbf{z})} \Big]\\ & \overset{\Delta}{=} \text{ELBO}(\mathbf{x}; \theta,\lambda) \\ \end{split}

Then, we can apply Monte Carlo on ELBO

\[ \frac{1}{k}\sum_{i=1}^{k} \log \frac{p_{\theta}(\mathbf{x},\mathbf{z}^{(i)})}{q_{\lambda}(\mathbf{z}^{(i)})},\quad \text{where }\mathbf{z}^{(i)}\sim q_{\lambda}(\mathbf{z}) \]

The tightness between ELBO and the original objective depends on the choice of variantional family.

The \( q_{\lambda}(\mathbf{z}) \) is called encoder, and the \( p_{\theta}(\mathbf{x}|\mathbf{z}) \) is called decoder.

We estimate gradient of ELBO w.r.t. \( \lambda^{(i)},\theta \) with Monte Carlo sampling.

For \( \theta \), we can simply pushing \( \nabla_{\theta} \) through the expectation operator. But for \( \nabla_{\lambda} \), we see that

\begin{align} \nabla_\lambda \mathbb{E}_{q_\lambda(\mathbf{z})} \left[\log \frac{p_\theta(\mathbf{x}, \mathbf{z})}{q_\lambda(\mathbf{z})} \right]= \mathbb{E}_{q_\lambda(\mathbf{z})} \left[ \left(\log \frac{p_\theta(\mathbf{x}, \mathbf{z})}{q_\lambda(\mathbf{z})} \right) \cdot \nabla_\lambda \log q_\lambda(\mathbf{z}) \right] \end{align}

Thus we can use RHS for Monte Carlo estimation. However, it suffers from high variance problem. To handle this, we use re-parameterization trick that we introduce a fixed, auxiliary distribution \( p(\epsilon) \) and a differentiable function \( T(\epsilon; \lambda) \) such that the following procedure

\begin{split} \epsilon &\sim p(\epsilon) \\ \mathbf{z} &\gets T(\epsilon; \lambda), \end{split}

is equivalent to sampling from \( q_{\lambda}(\mathbf{z}) \).

A first-order stichastic gradient method for optimizing ELBO. This kind of technique allows us to sub-sample the dataset but requires differentiability of our objective functions.

In BBVI, for each mini-batch \( \mathcal{B}=\set{\mathbf{x}^{(1)}, \dots, \mathbf{x}^{(m)} } \). We alternatively perform two steps of gradient descent

\begin{split} \lambda^{(i)} &\gets \lambda^{(i)} + \tilde{\nabla}_{\lambda} \text{ELBO}(\mathbf{x}^{(i)}; \theta,\lambda^{(i)}) \\ \theta &\gets \theta + \tilde{\nabla}_{\theta} \sum_{i} ELBO(\mathbf{x}^{(i)}; \theta, \lambda^{(i)}) \end{split}

For convenience, a popular distribution for \( p_{\theta}(\mathbf{z}) \) is unit Gaussian:

\[ p_{\theta}(\mathbf{z})=\mathcal{N}(\mathbf{z}|0,I) \]

or mixture of Gaussians with trainable means and covariance parameters.

The conditional distribution \( p_{\theta}(\mathbf{x}|\mathbf{z}) \) is where we introduce deep learning. We note that a conditional distribution can be constructed by defining a distribution family (parameterized by \(\omega \in \Omega\)) in the target space \(\mathbf{x}\), i.e., \(p_{\omega}(\mathbf{x})\) defines an unconditional distribution over \(\mathbf{x}\), and a mapping function \(g_{\theta}:\mathcal{Z}\mapsto\Omega\):

\[ p_{\theta}(\mathbf{x}|\mathbf{z})=p_{\omega}(\mathbf{x}), \quad \text{where } \omega=g_{\theta}(\mathbf{z}) \]

We specify \(g_{\theta}\) as deep neural network.