Supervised Contrastive Block Disentanglement

TL;DR

This paper introduces SCBD, a supervised contrastive learning algorithm that disentangles phenomena of interest from spurious correlations, improving domain generalization and batch correction in complex real-world datasets.

Contribution

The paper proposes a novel SCBD algorithm that enforces invariance to spurious correlations using supervised contrastive learning, applicable to real-world data for the first time.

Findings

SCBD improves out-of-distribution performance in domain generalization tasks.

SCBD effectively removes batch effects in large-scale biological imaging data.

Introducing hyperparameter α allows control over invariance strength.

Abstract

Real-world datasets often combine data collected under different experimental conditions. This yields larger datasets, but also introduces spurious correlations that make it difficult to model the phenomena of interest. We address this by learning two embeddings to independently represent the phenomena of interest and the spurious correlations. The embedding representing the phenomena of interest is correlated with the target variable $y$, and is invariant to the environment variable $e$. In contrast, the embedding representing the spurious correlations is correlated with $e$. The invariance to $e$ is difficult to achieve on real-world datasets. Our primary contribution is an algorithm called Supervised Contrastive Block Disentanglement (SCBD) that effectively enforces this invariance. It is based purely on Supervised Contrastive Learning, and applies to real-world data better than existing approaches. We empirically validate SCBD on two challenging problems. The first problem is domain generalization, where we achieve strong performance on a synthetic dataset, as well as on Camelyon17-WILDS. We introduce a single hyperparameter $\alpha$ to control the degree of invariance to $e$. When we increase $\alpha$ to strengthen the degree of invariance, out-of-distribution performance improves at the expense of in-distribution performance. The second problem is batch correction, in which we apply SCBD to preserve biological signal and remove inter-well batch effects when modeling single-cell perturbations from 26 million Optical Pooled Screening images.