Variational Autoencoders = Minimizing the Prior Distribution + Maximizing Mutual Information
This post is a short one; it mainly describes a fact that's quite useful and not at all complicated, yet somehow it took me this long to notice it.
In 《Mutual Information in Deep Learning: Unsupervised Feature Extraction》], we obtained the final loss of the Deep INFOMAX model as a weighted combination of two losses: one for the prior distribution and one for maximizing mutual information. That post told the whole story, but in a sense, the resulting loss was still something cobbled together by hand. This post, on the other hand, will show that this loss can actually be derived naturally from the variational autoencoder.
Derivation
At the risk of being repetitive, let's restate that the loss the variational autoencoder (VAE) needs to optimize is
\begin{equation}\begin{aligned}&KL(\tilde{p}(x)p(z|x)\Vert q(z)q(x|z))\\ =&\iint \tilde{p}(x)p(z|x)\log \frac{\tilde{p}(x)p(z|x)}{q(x|z)q(z)} dzdx\end{aligned}\end{equation}
Related discussions have already appeared many times on this blog. A VAE contains both an encoder and a decoder, but if all we need is to encode features, then training a decoder as well seems like an unnecessary burden. So the key question is how to get rid of the decoder.
It turns out this is straightforward: just split the VAE loss into two parts, more precisely
\begin{equation}\begin{aligned}&KL(\tilde{p}(x)p(z|x)\Vert q(z)q(x|z))\\ =&\iint \tilde{p}(x)p(z|x)\log \frac{p(z|x)}{q(z)} dzdx-\iint \tilde{p}(x)p(z|x)\log \frac{q(x|z)}{\tilde{p}(x)} dzdx\end{aligned}\end{equation}
The first term is the KL divergence with respect to the prior distribution, and as for the second term, isn't $\log \frac{q(x|z)}{\tilde{p}(x)}$ actually just the pointwise mutual information of $x,z$? If $q(x|z)$ has unlimited fitting capacity, then in the end we will necessarily have $\tilde{p}(x)p(z|x) = q(x|z)p(z)$ (by Bayes' rule), so the second term becomes
\begin{equation}KL(q(x|z)p(z)\Vert \tilde{p}(x)p(z))=KL(\tilde{p}(x)p(z|x)\Vert \tilde{p}(x)p(z))\end{equation}
which is exactly the mutual information between the two random variables $x,z$. The minus sign in front means we want to maximize this mutual information.
The rest of the derivation proceeds exactly as in 《Mutual Information in Deep Learning: Unsupervised Feature Extraction》], so I'll skip it here.
Conclusion
As stated at the outset, this post is short and doesn't contain much. Its main purpose is to offer a new way of understanding the VAE loss — as minimizing the prior distribution plus maximizing mutual information — from which the Deep INFOMAX loss then follows quite naturally.
If I hadn't already written 《Mutual Information in Deep Learning: Unsupervised Feature Extraction》], I would certainly have used this as the starting point for explaining Deep INFOMAX. But since that post was written days ago already, I had no choice but to open this short new note to fill in the gap.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.