Anti-causal domain generalization: Leveraging unlabeled data
Abstract
The problem of domain generalization concerns learning predictive models that are robust to distribution shifts when deployed in new, previously unseen environments. Existing methods typically require labeled data from multiple training environments, limiting their applicability when labeled data are scarce. In this work, we study domain generalization in an anti-causal setting, where the outcome causes the observed covariates. Under this structure, environment perturbations that affect the covariates do not propagate to the outcome, which motivates regularizing the model's sensitivity to these perturbations. Crucially, estimating these perturbation directions does not require labels, enabling us to leverage unlabeled data from multiple environments. We propose two methods that penalize the model's sensitivity to variations in the mean and covariance of the covariates across environments, respectively, and prove that these methods have worst-case optimality guarantees under certain classes of environments. Finally, we demonstrate the empirical performance of our approach on a controlled physical system and a physiological signal dataset.
Lay Summary
Machine learning models often perform poorly when deployed in environments different from those they were trained on. For example, a model trained to predict a patient's heart function from pressure waveforms recorded at one hospital may fail when used at another hospital, because of potential discrepancies in the measurement devices. Most existing solutions require collecting fully labeled data (the waveform measurements paired with expert ground-truth) from many different environments, which is expensive and often impractical. We show that this requirement can be relaxed when the quantity we want to predict is the underlying cause of the measurements, rather than a consequence of them. In such "anti-causal" settings, we can learn how measurements vary across environments using only unlabeled data, and use this information to build models that are less sensitive to such variation. We provide mathematical guarantees that our method performs well even on environments unseen during training. We demonstrate the approach on a controlled physical system and on a real-world medical dataset. Our work can render machine learning models more reliable in domains like healthcare, where obtaining labeled data from multiple environments may be impractical.