Removing Biases from Molecular Representations via Information Maximization

Chenyu Wang, Sharut Gupta, Caroline Uhler, Tommi Jaakkola

Keywords: Molecular Representation, Batch Effect, Contrastive Learning, Information Maximization, Drug Discovery

TL;DR: We propose InfoCORE to deal with batch effects in drug screens and obtain refined molecular representations. It is established on a variational lower bound of the conditional mutual information between latent representations given a batch identifier.

Abstract: High-throughput drug screening -- using cell imaging or gene expression measurements as readouts of drug effect -- is a critical tool in biotechnology to assess and understand the relationship between the chemical structure and biological activity of a drug. Since large-scale screens have to be divided into multiple experiments, a key difficulty is dealing with batch effects, which can introduce systematic errors and non-biological associations in the data. We propose InfoCORE, an Information maximization approach for COnfounder REmoval, to effectively deal with batch effects and obtain refined molecular representations. InfoCORE establishes a variational lower bound on the conditional mutual information of the latent representations given a batch identifier. It adaptively reweighs samples to equalize their implied batch distribution. Extensive experiments on drug screening data reveal InfoCORE's superior performance in a multitude of tasks including molecular property prediction and molecule-phenotype retrieval. Additionally, we show results for how InfoCORE offers a versatile framework and resolves general distribution shifts and issues of data fairness by minimizing correlation with spurious features or removing sensitive attributes. The code is available at https://github.com/uhlerlab/InfoCORE.

Submission Number: 46