InfoVAE addresses the aggregated posterior mismatch problem in standard VAEs by modifying the objective function to explicitly match the aggregated posterior with the prior distribution.

Key Ideas

  1. If the aggregated posterior $q_\phi(z)$ becomes a Dirac delta, all inputs $x$ map to the same latent code $z$
  2. In such a scenario, the latent code $z$ would contain no information about the input $x$
  3. The standard VAE objective does not directly encourage the matching between $q_\phi(z)$ and the prior $p(z)$

InfoVAE Objective Function

The standard VAE objective is given by:

$$ F_{\theta}(q) = \mathbb{E}_{q_{\phi}(z|x)} \left[ \log p_{\theta}(x|z) \right] - D_{KL} \left( q_{\phi}(z|x) \| p(z) \right) $$

This can also be written as(derivation skipped):

$$ L_{ELBO} = -D_{KL}(q_\phi(z) \| p(z)) - \mathbb{E}_{q_\phi(z)}[D_{KL}(q_\phi(x|z) \| p_\theta(x|z))] $$

Here $p_\theta(x|z)$ is inverted encoder distribution. The first term means that $q_\phi(z)$ should match $p(z)$. The second term means that $q_\phi(x|z)$ should match $p_\theta(x|z)$.

Incorporating Mutual Information

To further enhance information retention, the InfoVAE objective incorporates a mutual information term:

$$ L_{ELBO} = -\lambda D_{KL}(q_\phi(z) \| p(z)) - \mathbb{E}_{q_\phi(z)}[D_{KL}(q_\phi(x|z) \| p_\theta(x|z))] + \alpha I_q(x;z) $$

Terms Breakdown:

  • $I_q(x;z)$: Mutual information between input $x$ and latent code $z$, encouraging $z$ to retain meaningful information about $x$. The mutual information $I_q(x;z)$ is defined as:

$$ I_q(x;z) = \sum_x \sum_z p_\theta(x,z) \log \frac{p_\theta(x,z)}{p_\theta(x)p_\theta(z)} $$

  • $\alpha$: Coefficient controlling the importance of the mutual information term

InfoVAE addresses the aggregated posterior mismatch problem by:

  1. Explicitly minimizing the KL divergence between the aggregated posterior $q_\phi(z)$ and the prior $p_\theta(z)$
  2. Preserving mutual information between the input $x$ and the latent code $z$

Training InfoVAE

Because optimizing the mutual information term is difficult, we optimize the following objective instead(derivation skipped however it is in paper):

$$ L_{ELBO} = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - (1-\alpha)D_{KL}(q_\phi(z|x)||p(z)) - (\alpha + \lambda - 1)D_{KL}(q_\phi(z)||p(z)) $$
The first two terms are optimized similar to VAE while the third term, which is minimization of the divergence between the aggregated posterior and the latent prior, is carried out using techniques such as adversarial minimization.