[Paper Review] Accounting for Unobserved Confounding in Domain Generalization

Alexis Bellot, Mihaela van der Schaar|arXiv (Cornell University)|Jul 21, 2020

Domain Adaptation and Few-Shot Learning49 references6 citations

TL;DR

This paper proposes Derivative Invariant Risk Minimization (DIRM), a distributionally robust optimization method that accounts for unobserved confounders in domain generalization by enforcing partial invariance in error derivatives across environments. DIRM improves generalization and feature selection stability over ERM and IRM, especially under unobserved confounding, as shown in multimodal healthcare data with image, speech, and tabular inputs.

ABSTRACT

This paper investigates the problem of learning robust, generalizable prediction models from a combination of multiple datasets and qualitative assumptions about the underlying data-generating model. Part of the challenge of learning robust models lies in the influence of unobserved confounders that void many of the invariances and principles of minimum error presently used for this problem. Our approach is to define a different invariance property of causal solutions in the presence of unobserved confounders which, through a relaxation of this invariance, can be connected with an explicit distributionally robust optimization problem over a set of affine combination of data distributions. Concretely, our objective takes the form of a standard loss, plus a regularization term that encourages partial equality of error derivatives with respect to model parameters. We demonstrate the empirical performance of our approach on healthcare data from different modalities, including image, speech and tabular data.

Motivation & Objective

• To address the challenge of out-of-distribution generalization when unobserved confounders induce shifts in data distributions.
• To develop a robust learning objective that maintains performance under distributional shifts not captured by observed variables.
• To formalize a new invariance principle that accounts for unobserved confounders, bridging the gap between causal and correlation-based learning.
• To improve the reproducibility of feature selection across diverse data distributions in healthcare applications.
• To demonstrate empirical superiority of the proposed method over ERM, IRM, and REx in multimodal and clinical data settings.

Proposed method

• Proposes a relaxation of invariance to unobserved confounders by enforcing partial equality of error derivatives with respect to model parameters across environments.
• Derives a distributionally robust optimization objective that combines standard loss with a regularization term promoting derivative invariance.
• Defines the uncertainty set as affine combinations of training data distributions, enabling robustness to shifts in observed and unobserved variables.
• Uses a hyperparameter λ to interpolate between the causal solution (λ→∞) and ordinary least squares (λ=0), enabling adaptive robustness.
• Applies the method to multimodal data (image, speech, tabular) and survival prediction using real-world healthcare datasets.
• Employs a single-layer network for feature selection stability analysis, evaluating consistency of top features across 100 random data splits.

Experimental results

Research questions

• RQ1 How can we achieve robust generalization in the presence of unobserved confounders that shift the conditional distribution P(Y|X)?
• RQ2 What invariance property can be enforced when confounders are unobserved, and how does it differ from standard invariance in invariant risk minimization?
• RQ3 Can a regularization based on error derivative invariance improve generalization and feature selection stability compared to ERM and IRM?
• RQ4 How does the proposed method perform under realistic shifts in healthcare data, particularly when key predictors like age are omitted?
• RQ5 To what extent does the method maintain consistent feature selection across diverse data distributions?

Key findings

• DIRM significantly improves generalization performance on multimodal healthcare data compared to ERM, IRM, and REx, especially when spurious correlations or unobserved confounders are present.
• In speech recognition tasks, DIRM reduces error rates by correcting for spurious differences between recording types, outperforming ERM and IRM.
• On survival prediction with heart failure data, DIRM achieves competitive performance to state-of-the-art methods despite omitting age as a predictor, indicating robustness to unobserved confounding.
• DIRM recovers the top 10 influential features in 80 out of 100 experiments, compared to only 4 for ERM, demonstrating superior feature selection reproducibility.
• The method maintains stable feature selection across environments when λ ∈ [0,1], confirming theoretical stability under affine combinations of training distributions.
• Theoretical guarantees are limited to shifts within the affine hull of training distributions, indicating a need for caution in unseen, extreme distributional shifts.