Contrastive Mixture of Posteriors for Counterfactual Inference, Data Integration and Fairness

Adam Foster, Árpi Vezér, Craig A Glastonbury, Páidí Creed, Sam Abujudeh, Aaron Sim

Keywords: Contrastive learning, variational autoencoders, counterfactual inference, data integration, fairness, batch effect correction, scRNA-seq, RNA-seq, genomics, omics, transcriptomics, invariant representations, TCGA, CCLE, single-cell

TL;DR: A novel method for learning aligned representations with applications to data integration, counterfactual inference, and fair representation learning.

Abstract: Learning meaningful representations of data that can address challenges such as batch effect correction, data integration and counterfactual inference is a central problem in many domains including computational biology. Adopting a Conditional VAE framework, we identify the mathematical principle that unites these challenges: learning a representation that is marginally independent of a condition variable. We therefore propose the Contrastive Mixtures of Posteriors (CoMP) method that uses a novel misalignment penalty to enforce this independence. This penalty is defined in terms of mixtures of the variational posteriors themselves, unlike prior work which uses external discrepancy measures such as MMD to ensure independence in latent space. We show that CoMP has attractive theoretical properties compared to previous approaches, especially when there is complex global structure in latent space. We further demonstrate state of the art performance on a number of real-world problems, including the challenging tasks of aligning human tumour samples with cancer cell-lines and performing counterfactual inference on single-cell RNA sequencing data. Incidentally, we find parallels with the fair representation learning literature, and demonstrate CoMP has competitive performance in learning fair yet expressive latent representations.

Code Of Conduct: I certify that all co-authors of this work have read and commit to adhering to the NeurIPS Statement on Ethics, Fairness, Inclusivity, and Code of Conduct.

Community Implementations: [ 3 code implementations](https://www.catalyzex.com/paper/contrastive-mixture-of-posteriors-for/code)