Semi-Supervised Learning with VAEs

Semi-supervised learning (SSL) addresses a common scenario in machine learning: an abundance of unlabeled data alongside a much smaller, often expensive-to-acquire, set of labeled data. The objective is to leverage both data types to build models that outperform those trained on labeled data alone. Variational Autoencoders, with their inherent ability to learn rich data representations from unlabeled inputs, offer a powerful framework for SSL.

Leveraging VAEs for Semi-Supervised Classification

The core idea behind using VAEs in SSL is that the latent space $z$ learned by the VAE from all available data (both labeled $X_L$ and unlabeled $X_U$) can provide a more structured and informative representation for a subsequent classification task than the raw input data $x$. This is particularly advantageous when labeled data is scarce.

A prevalent and effective strategy involves a VAE that learns a mapping from $x$ to $z$ and back to $x$, coupled with a classifier that operates on these learned latent representations.

1. VAE for Representation Learning: The VAE, consisting of an encoder $q_\phi(z|x)$ and a decoder $p_\theta(x|z)$, is trained on the entire dataset ($X_L \cup X_U$). Its objective is to maximize the Evidence Lower Bound (ELBO):

   $$\mathcal{L}_{VAE}(x) = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - D_{KL}(q_\phi(z|x) || p(z))$$

   This process forces the encoder to capture salient features of the input data within the latent space $z$. The prior $p(z)$ is typically a standard Gaussian $N(0,I)$.

2. Classifier on Latent Space: A separate classification network, $f_\psi(y|z)$, is trained to predict labels $y$ using the latent codes $z_L$ derived from the labeled portion of the data, $X_L$. The classification loss, $\mathcal{L}_{C}$, is commonly the cross-entropy between the true labels $y_L$ and the predicted labels $\hat{y}_L = f_\psi(z_L)$.

   $$\mathcal{L}_{C}(y_L, \hat{y}_L) = -\sum_{i} y_{L,i} \log \hat{y}_{L,i}$$

These two components are often trained jointly. The total objective function combines the VAE's ELBO (applied to all data) and the classification loss (applied only to labeled data), weighted by a hyperparameter $\gamma$:

$$\mathcal{J}_{SSL} = \sum_{x \in X_L \cup X_U} \left( \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - D_{KL}(q_\phi(z|x) || p(z)) \right) + \gamma \sum_{(x_l, y_l) \in D_L} \mathcal{L}_{C}(y_l, f_\psi(q_\phi(z|x_l)))$$

The hyperparameter $\gamma$ balances the contribution of the generative (VAE) task and the discriminative (classification) task. A diagram illustrating this common architecture is shown below.

A common architecture for semi-supervised learning with VAEs. The VAE processes both labeled and unlabeled data to learn representations. A classifier then uses these representations from labeled data for supervised training.

Integrated Generative Models for SSL

More deeply integrated approaches, such as the M1 and M2 models proposed by Kingma et al. (2014), treat $y$ as a random variable within the generative model itself. For instance, the M2 model posits a generative process like $p(y) \rightarrow p(z|y) \rightarrow p(x|z)$. The label $y$ influences the generation of latent variable $z$, which in turn generates $x$.

The objective function for such models becomes more involved. For labeled data $(x_l, y_l)$, the ELBO is formulated to reflect $\log p(x_l, y_l)$. For unlabeled data $x_u$, the model must marginalize over the unknown labels. This is typically achieved by introducing an auxiliary inference network $q_\psi(y|x_u)$ that predicts a distribution over labels for unlabeled samples. The ELBO for unlabeled data then involves an expectation over these predicted label distributions:

$$\mathcal{L}_{unsup}(x_u) = \sum_k q_\psi(y=k|x_u) \left( \mathbb{E}_{q_\phi(z|x_u)}[\log p_\theta(x_u|z) + \log p(z|y=k) + \log p(y=k)] - D_{KL}(q_\phi(z|x_u)||p(z|y=k)) \right) + H(q_\psi(y|x_u))$$

The term $p(z|y=k)$ is a class-conditional prior for the latent variables, and $H(q_\psi(y|x_u))$ is an entropy term that encourages confident predictions from $q_\psi(y|x_u)$ when appropriate. The classifier $q_\psi(y|x)$ is also trained discriminatively using the labeled data, often via an additional term in the total loss function, such as:

$$\mathcal{L}_{classifier\_sup} = \sum_{(x_l, y_l) \in D_L} -\log q_\psi(y_l|x_l)$$

These integrated models can, in theory, better capture the underlying joint distribution $p(x,y)$, potentially leading to improved generative and classification performance. However, they also introduce greater model complexity and more intricate training dynamics.

Architectural and Training Considerations

• Network Components: The encoder $q_\phi(z|x)$, decoder $p_\theta(x|z)$ (or $p_\theta(x|z,y)$ in some models), and classifier $f_\psi(y|z)$ (or $q_\psi(y|x)$) are typically neural networks. Their architectures (e.g., CNNs for images, RNNs for sequences) should be chosen based on the data modality.
• Parameter Sharing: The encoder $q_\phi(z|x)$ is shared for processing both labeled and unlabeled data. The classifier $f_\psi(y|z)$ operates on the latent codes produced by this shared encoder.
• Training Batches: During training, mini-batches are often constructed by sampling from both the labeled pool $D_L$ and the unlabeled pool $D_U$. This ensures that both the VAE reconstruction/regularization terms and the supervised classification term contribute to the gradient updates.
• Balancing Losses: The weighting factor $\gamma$ (and potentially other weights for different terms in more complex objectives) is important. Its value dictates the relative importance of learning good representations versus achieving high classification accuracy on the labeled set. Proper tuning of these weights is often necessary for optimal performance.

Advantages of VAE-based SSL

• Effective Use of Unlabeled Data: VAEs naturally leverage large quantities of unlabeled data to learn the underlying structure and variations in the data.
• Improved Classification with Few Labels: The learned representations can significantly boost the performance of a classifier, especially when the number of labeled examples is small.
• Better Generalization: By learning from a broader data distribution (including unlabeled samples), models may generalize better to unseen data.
• Class-Conditional Generation: Models that explicitly incorporate $y$ into the generative process (like some variants of M2) can be used to generate new samples conditioned on specific class labels, e.g., $p(x|y,z)$ or $p(x|z)$ where $z \sim p(z|y)$.

Challenges and Limitations

• Posterior Collapse: As with standard VAEs, if the decoder is too powerful or if the KL divergence term is too heavily weighted, the latent variables $z$ might be ignored by the decoder, leading to an uninformative latent space. This can hinder the performance of the classifier $f_\psi(y|z)$.
• Mismatch between Labeled and Unlabeled Data: If the distribution of unlabeled data is significantly different from that of the labeled data (domain shift), the representations learned from unlabeled data might not be beneficial, or could even be detrimental, to classifying the labeled data.
• Complexity and Tuning: Integrated SSL VAE models can be more complex to implement and train than simpler pipelines. Tuning hyperparameters like $\gamma$ and the learning rates for different model components requires careful experimentation.
• Optimization Difficulties: Balancing the multiple objectives (reconstruction, KL regularization, classification accuracy) can sometimes lead to challenging optimization landscapes.

Despite these challenges, VAEs provide a versatile and principled framework for semi-supervised learning. By effectively combining the VAE's capacity for unsupervised representation learning with supervised signals from limited labeled data, these models can achieve strong performance on tasks where labeled data is a bottleneck. The choice between a simpler feature-extraction approach and a more integrated generative model often depends on the specific problem, the amount of available data, and computational resources.