Supervised Contrastive Block Disentanglement

Taro Makino, Ji Won Park, Natasa Tagasovska, TAKAMASA KUDO, Paula Coelho, Heming Yao, Jan-Christian Huetter, Ana Carolina Leote, Burkhard Hoeckendorf, Stephen Ra, David Richmond, Kyunghyun Cho, Aviv Regev, Romain Lopez

Keywords: disentanglement, block disentanglement, out-of-distribution generalization, domain generalization, distribution shift, spurious correlations, robustness

TL;DR: Practical algorithm that block disentangles environment-invariant and environment-dependent features for domain generalization and batch correction

Abstract: Real-world datasets often combine data collected under different experimental conditions. Although this yields larger datasets, it also introduces spurious correlations that make it difficult to accurately model the phenomena of interest. We address this by learning two blocks of latent variables to independently represent the phenomena of interest and the spurious correlations. The former are correlated with the target variable $y$ and invariant to the environment variable $e$, while the latter depend on $e$. The invariance of the phenomena of interest to $e$ is highly sought-after but difficult to achieve on real-world datasets. Our primary contribution is an algorithm called Supervised Contrastive Block Disentanglement (SCBD) that is highly effective at enforcing this invariance. It is based purely on supervised contrastive learning, and scales to real-world data better than existing approaches. We empirically validate SCBD on two challenging problems. The first is domain generalization, where we achieve strong performance on a synthetic dataset, as well as on Camelyon17-WILDS. SCBD introduces a single hyperparameter $\alpha$ that controls the degree of invariance to $e$. When we increase $\alpha$ to strengthen the degree of invariance, there is a monotonic improvement in out-of-distribution performance at the expense of in-distribution performance. The second is a scientific problem of batch correction. Here, we demonstrate the utility of SCBD by learning representations of single-cell perturbations from 26 million Optical Pooled Screening images that are nearly free of technical artifacts induced by the variation across wells.

Primary Area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 3855