[Paper Review] Preventing Failures Due to Dataset Shift: Learning Predictive Models That Transport

Adarsh Subbaswamy, Peter Schulam+1 more|arXiv (Cornell University)|Dec 11, 2018

Machine Learning in HealthcareComputer Science80 citations

TL;DR

The paper proposes Graph Surgery to learn stable predictive models that transport across environments with differing data-generating mechanisms, by using interventional distributions identified from a causal DAG with selection variables.

ABSTRACT

Classical supervised learning produces unreliable models when training and target distributions differ, with most existing solutions requiring samples from the target domain. We propose a proactive approach which learns a relationship in the training domain that will generalize to the target domain by incorporating prior knowledge of aspects of the data generating process that are expected to differ as expressed in a causal selection diagram. Specifically, we remove variables generated by unstable mechanisms from the joint factorization to yield the Surgery Estimator---an interventional distribution that is invariant to the differences across environments. We prove that the surgery estimator finds stable relationships in strictly more scenarios than previous approaches which only consider conditional relationships, and demonstrate this in simulated experiments. We also evaluate on real world data for which the true causal diagram is unknown, performing competitively against entirely data-driven approaches.

Motivation & Objective

• Motivate reliability when training and deployment environments differ in data-generating processes.
• Introduce a proactive method that uses causal prior knowledge to build transportable predictors.
• Develop the Graph Surgery estimator that yields stable interventional distributions across environments.
• Prove soundness and completeness of the estimator and relate it to graph pruning and distributional robustness.

Proposed method

• Use selection diagrams to model environment-varying mechanisms and identify mutable variables M as children of selection nodes S.
• Apply the ID algorithm on ADMGs to identify interventional distributions that predict the target T.
• Enumerate possible interventional distributions by intervening on M (and possibly T) and select the best identifiable one via held-out source-environment validation.
• Show that graph pruning is a special case of graph surgery and that surgery optimizes distributional robustness over environments.
• Provide theoretical guarantees: soundness (stable estimator) and completeness (existence if stable surgery estimator exists).
• Connect the approach to minimax robustness across environment distributions Γ.

Experimental results

Research questions

• RQ1 Can a predictive model be made stable to environment-varying mechanisms without access to target-domain data?
• RQ2 Does intervening on mutable variables via interventional distributions yield stable predictions when selection mechanisms differ across environments?
• RQ3 In which scenarios does graph surgery strictly outperform graph pruning and other reactive domain adaptation methods?
• RQ4 How does graph surgery relate to distributional robustness and minimax optimality under dataset shift?

Key findings

• Graph surgery yields stable predictors invariant to selection mechanisms and can handle target-shift where the target variable is mutable.
• Graph pruning is a special case of graph surgery; surgery generalizes pruning and can succeed where pruning fails.
• In simulations, surgery generalizes across unseen environment changes better than OLS and competing methods like causal transfer learning (CT).
• On Bike Sharing data, surgery achieves competitive MSE relative to alternatives, performing best in several test splits; AR remains strong in some settings but requires anchor assumptions.
• Theoretical results show surgery is sound and complete for identifying stable interventional predictors, and is minimax-optimal among stable, directly transportable relations.