Towards IID representation learning and its application on biomedical data

Jiqing Wu, Inti Zlobec, Maxime W Lafarge, Yukun He, Viktor Koelzer

Keywords: IID, IID representation learning, OOD generalization, causality, biomedical

TL;DR: We argue that instead of being a questionable assumption, IID is a fundamental task-relevant property that needs to be learned. The code will be publicly available via the link https://github.com/CTPLab/IID_representation_learning.git.

Abstract:

Due to the heterogeneity of real-world data, the widely accepted independent and identically distributed (IID) assumption has been criticized in recent studies on causality. In this paper, we argue that instead of being a questionable assumption, IID is a fundamental task-relevant property that needs to be learned. Consider $k$ independent random vectors $\mathsf{X}^{i = 1, \ldots, k}$, we elaborate on how a variety of different causal questions can be reformulated to learning a task-relevant function $\phi$ that induces IID among $\mathsf{Z}^i := \phi \circ \mathsf{X}^i$, which we term IID representation learning.

For proof of concept, we examine the IID representation learning on Out-of-Distribution (OOD) generalization tasks. Concretely, by utilizing the representation obtained via the learned function that induces IID, we conduct prediction of molecular characteristics (molecular prediction) on two biomedical datasets with real-world distribution shifts introduced by a) preanalytical variation and b) sampling protocol. To enable reproducibility and for comparison to the state-of-the-art (SOTA) methods, this is done by following the OOD benchmarking guidelines recommended from WILDS. Compared to the SOTA baselines supported in WILDS, the results confirm the superior performance of IID representation learning on OOD tasks.

Registration: I acknowledge that publication of this at MIDL and in the proceedings requires at least one of the authors to register and present the work during the conference.

Authorship: I confirm that I am the author of this work and that it has not been submitted to another publication before.

Paper Type: both

Primary Subject Area: Unsupervised Learning and Representation Learning

Secondary Subject Area: Transfer Learning and Domain Adaptation

Confidentiality And Author Instructions: I read the call for papers and author instructions. I acknowledge that exceeding the page limit and/or altering the latex template can result in desk rejection.

Code And Data: Code: https://github.com/CTPLab/IID_representation_learning

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2203.00332/code)