Contrastive Variational Autoencoder Enhances Salient Features

TL;DR

The paper introduces a contrastive variational autoencoder (cVAE) that enhances the detection of salient features in target datasets by modeling enriched latent structures relative to background data, applicable to diverse data types.

Contribution

It presents the first deep generative contrastive autoencoder that isolates and enhances salient latent features using unpaired datasets, extending contrastive learning beyond linear models.

Findings

cVAE effectively uncovers salient latent structures in gene expression data.

It enhances feature detection in facial image datasets.

The method is robust to noise and dataset purity.

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.