[Paper Review] Contrastive Variational Autoencoder Enhances Salient Features

Abubakar Abid, James Zou|arXiv (Cornell University)|Feb 12, 2019

Generative Adversarial Networks and Image Synthesis20 references44 citations

TL;DR

The paper introduces the contrastive variational autoencoder (cVAE), which learns latent features salient to a target dataset by contrastively separating them from a background dataset, enabling denoising and generation focused on those features.

ABSTRACT

Variational autoencoders are powerful algorithms for identifying dominant latent structure in a single dataset. In many applications, however, we are interested in modeling latent structure and variation that are enriched in a target dataset compared to some background---e.g. enriched in patients compared to the general population. Contrastive learning is a principled framework to capture such enriched variation between the target and background, but state-of-the-art contrastive methods are limited to linear models. In this paper, we introduce the contrastive variational autoencoder (cVAE), which combines the benefits of contrastive learning with the power of deep generative models. The cVAE is designed to identify and enhance salient latent features. The cVAE is trained on two related but unpaired datasets, one of which has minimal contribution from the salient latent features. The cVAE explicitly models latent features that are shared between the datasets, as well as those that are enriched in one dataset relative to the other, which allows the algorithm to isolate and enhance the salient latent features. The algorithm is straightforward to implement, has a similar run-time to the standard VAE, and is robust to noise and dataset purity. We conduct experiments across diverse types of data, including gene expression and facial images, showing that the cVAE effectively uncovers latent structure that is salient in a particular analysis.

Motivation & Objective

• Identify latent factors that are enriched in a target dataset relative to a background dataset.
• Develop a nonlinear, deep generative model to isolate salient features beyond what standard VAEs capture.
• Provide a practical, end-to-end training framework with comparable runtime to VAEs.
• Demonstrate robustness to background noise and dataset purity across diverse data types.

Proposed method

• Introduce two encoders q_phi_s(s|x) and q_phi_z(z|x) to infer salient and irrelevant latents from the target and background data.
• Use a single decoder f_theta(s,z) that reconstructs target data from concatenated (s,z) and background data from (0,z).
• Derive variational lower bounds for target and background likelihoods analogous to VAE objectives (Eqs. 3 and 4).
• Train to maximize the sum of the target and background objectives plus a total correlation (TC) term to encourage independence between s and z via a density-ratio based discriminator.
• In practice, implement TC estimation with a discriminator to distinguish joint (s,z) from the product of marginals, enabling a tractable optimization.
• Optionally constrain networks to zero biases to ensure zero-inputs map to zero outputs for denoising.

Experimental results

Research questions

• RQ1 Can a contrastive VAE identify latent features that are salient in the target data but absent in the background?
• RQ2 How does cVAE compare to standard VAE in recovering or clustering salient features when background noise or confounding variation is present?
• RQ3 Can cVAE generate or denoise samples by manipulating only the salient latent variables without reintroducing irrelevant structure?
• RQ4 How robust is the method to background dataset quality and hyperparameter settings across different domains (images, genomics)?

Key findings

• On Grassy-MNIST, the cVAE’s salient s-latents cluster digits, unlike the standard VAE where dominant factors relate to grass noise (silhouette scores).
• On RNA-Seq data, the cVAE improves clustering by pre/post transplant status in the salient latent space, unlike the standard VAE.
• On CelebA, the cVAE highlights and allows manipulation of attributes like glasses or caps, while denoising can remove irrelevant background features.
• The method is robust to background noise and modest background mis-specification, with performance maintained even when the background is noisy.
• Runtime and architectural complexity are similar to standard VAEs, with the added TC term improving disentanglement.