Learning Clinical Representations Under Systematic Distribution Shift

TL;DR

This paper introduces a practice invariant representation learning framework for clinical prediction models that enhances robustness and transferability across different healthcare institutions by explicitly addressing systematic distribution shifts.

Contribution

The authors propose a novel framework combining supervised risk minimization with adversarial regularization to learn invariant representations in multimodal clinical data, improving out-of-distribution performance.

Findings

Out-of-distribution AUROC improved by 2-3 points.

Maintains in-distribution performance and improves calibration.

Demonstrates robustness across multiple longitudinal EHR prediction tasks.

Abstract

Clinical machine learning models are increasingly trained using large scale, multimodal foundation paradigms, yet deployment environments often differ systematically from the data generating settings used during training. Such shifts arise from heterogeneous measurement policies, documentation practices, and institutional workflows, leading to representation entanglement between physiologic signal and practice specific artifacts. In this work, we propose a practice invariant representation learning framework for multimodal clinical prediction. We model clinical observations as arising from latent physiologic factors and environment dependent processes, and introduce an objective that jointly optimizes predictive performance while suppressing environment predictive information in the learned embedding. Concretely, we combine supervised risk minimization with adversarial environment regularization and invariant risk penalties across hospitals. Across multiple longitudinal EHR prediction tasks and cross institution evaluations, our method improves out of distribution AUROC by up to 2 to 3 points relative to masked pretraining and standard supervised baselines, while maintaining in distribution performance and improving calibration. These results demonstrate that explicitly accounting for systematic distribution shift during representation learning yields more robust and transferable clinical models, highlighting the importance of structural invariance alongside architectural scale in healthcare AI.