I-SPEC: An End-to-End Framework for Learning Transportable, Shift-Stable Models
This addresses the challenge of building transportable models for machine learning applications in shifting environments, representing an incremental improvement over existing graph-based methods.
The paper tackles the problem of model generalization failure due to distribution shifts between development and deployment by proposing I-SPEC, an end-to-end framework that learns a partial ancestral graph from data to determine stable interventional distributions, applied to mortality prediction to achieve robustness without requiring a full causal graph.
Shifts in environment between development and deployment cause classical supervised learning to produce models that fail to generalize well to new target distributions. Recently, many solutions which find invariant predictive distributions have been developed. Among these, graph-based approaches do not require data from the target environment and can capture more stable information than alternative methods which find stable feature sets. However, these approaches assume that the data generating process is known in the form of a full causal graph, which is generally not the case. In this paper, we propose I-SPEC, an end-to-end framework that addresses this shortcoming by using data to learn a partial ancestral graph (PAG). Using the PAG we develop an algorithm that determines an interventional distribution that is stable to the declared shifts; this subsumes existing approaches which find stable feature sets that are less accurate. We apply I-SPEC to a mortality prediction problem to show it can learn a model that is robust to shifts without needing upfront knowledge of the full causal DAG.