Semi-supervised classification objective from Kingma et al

Question

In this 2014 paper, Kingma et al. develop different methods to do semi-supervised learning with VAEs. In one of their proposed solutions ("M2"), they approach this problem by incorporating the labels (y) as a causal factor in the generative model, s.t. images depend (through a generative model parameterized by a neural network; the decoder of the VAE) on both these labels (e.g. MNIST digit classes) and other latent variables (z). The labels (y) are partially observed: for some items in the data set, we have labels, while for others we don't. The other latents are never observed - discovering them is part of the learning process.

They then derive the ELBO loss for a VAE with this generative model. They observe that, for labeled images, this loss does not end up training the encoder network to produce good class labels (since the class labels are given for these cases, the ELBO does depend on the decoder's ability to use these class labels to reconstruct good images, but not on the encoder's ability to generate good class labels). This is undesirable, as it means that the classification part of the network effectively only learns from unlabeled examples. So, they augment their loss function with a classification loss (equation 9), which is basically just the cross-entropy between the class probabilities produced by the encoder, and the true labels, but scaled by a certain factor.

They then say the following:

While we have obtained this objective function by motivating the need for all model components to learn at all times, the objective 9 can also be derived directly using the variational principle by instead performing inference over the parameters $\pi$ of the categorical distribution, using a symmetric Dirichlet prior over these parameters.

And this is where they lose me. I would really like to understand how such a derivation would work, as I assume this also produces the scale factor that they apply to this term in the loss (which otherwise seems like a hyperparameter that would need to be tuned empirically). However, I just don't understand what they mean with this sentence.

Specifically, $q_\phi(y|x)$ is a categorical distribution produced by the encoder. Its parameters are just the category probabilities, which they denote by the vector $\pi_\phi(x)$. Thus, the entries in $\pi$ simply are the probabilities in $q_\phi(y|x)$. So "performing inference over the parameters $\pi$", presumably just means "estimating those numbers", or perhaps "optimizing variational beliefs over those numbers". I just don't see how you get from that idea to a loss term $E_{\tilde{p_l}(\mathbf{x},y)}\left[ -\log q_\phi(y|\mathbf{x}) \right]$.

Can anyone explain this?

Answer 1

I think they're describing a VAE with $\pi$ and $z$ as latent variables, and $x$ and $y$ as observed variables. The prior on $z$ is normal, as usual, and the prior on $\pi$ is $\text{Dirichlet}(\alpha)$.

(As a note, I don't think they mean literally that the objective in eq. 9 will exactly fall out of this model. I think they're just suggesting that something similar looking can be derived, as is described below.)

The generative process remains the same, as described by eq. 2. The evidence lower bound becomes: $$ \begin{align} \log p_\theta(x,y) &\geq \mathbb{E}_{q_\phi(z,\pi|x,y)} \left[ \log p_\theta(x,y|z,\pi) \right] - \mathcal{D}_{KL}\left[ q_\phi(z,\pi|x,y) \vert \vert p(z,\pi)\right] \\ &\geq \mathbb{E}_{q_\phi(z,\pi|x,y)} \left[ \log p_\theta(x,y|z,\pi) + \log p(z,\pi) - \log q_\phi(z,\pi|x,y) \right] \\ &\geq \mathbb{E}_{q_\phi(z,\pi|x,y)} \left[ \log p_\theta(x|y,z)+\log p(y|\pi) + \log p(z) + \log p(\pi) \\- \log q_\phi(z|x,y) - \log q_\phi(\pi|x) \right] \end{align} $$

This looks very similar to eq. 6 -- in fact, we only have three extra terms*: $\log p(y|\pi) + \log p(\pi) - \log q_\phi(\pi|x)$. The last two terms are the KL divergence between the posterior and prior distributions over $\pi$, which has a tractable form in the case that they're both dirichlet, or just an irrelevant constant if $q_\phi(\pi|x)$ is a degenerate point distribution.

The first term $\log p(y|\pi)$ is essentially the classification loss. In eq. 9, it's written as $\log q_\phi(y|x)$, but really it's the same thing, just written differently because the authors have $y$ as a latent, and infer it from $x$, whereas in the proposed model, $\pi$ is the latent which is inferred from $x$, and $\pi_k$ is the probability which the model assigns to the label $y$ being from the k-th class. Hence, we have the typical classification log-loss.

*There's one additional term in eq. 6, $p_\theta(y)$ which is missing from my derivation -- the authors don't seem to define this anywhere, so I'm not sure where it came from, but I think it's just a red herring.

As for $\alpha$ in eq. 9, this is probably a random hyperparameter they threw in, and not at all related to the $\alpha$ parameter of the dirichlet prior.