CausalShift: A Modular, Plugin-Based Framework for Dataset Shift Handling in Machine Learning

Abstract

Machine learning models deployed under distribution shift suffer performance degradation whose root cause is rarely diagnosed before adaptation is attempted. Existing frameworks either lack a closed loop from shift diagnosis to method selection, require a fully specified causal graph, or depend on multiple heterogeneous training environments—conditions rarely met in practice. We propose CausalShift, a modular, plugin-based framework for end-to-end dataset shift handling. A universal statistical core remains fully functional without any causal assumption; four optional causal plugins are activated selectively by a four-level knowledge grading system that scales from a fully specified structural causal model (Level A) down to a single observational dataset (Level D). A three-path causal attribution module estimates the relative contributions of covariate and concept shift to the performance gap, and attribution-driven routing maps the diagnosis to an adaptation strategy grounded in minimax optimality. Proof-of-concept experiments on four benchmarks—a synthetic structural causal model (SCM), Colored Modified National Institute of Standards and Technology (MNIST), Benchmarking In-the-Wild Distribution Shifts (WILDS) Camelyon17, and five American Community Survey (ACS) geographic shift tasks—demonstrate that CausalShift reduces the in-distribution to out-of-distribution (ID–OOD) accuracy gap by 19.9 percentage points on the synthetic benchmark and 72.5 percentage points on the spurious correlation benchmark, while remaining competitive on real-world image shift and achieving performance parity with ERM on mild-shift tasks. A tiered causal evaluation suite, including Λ*-robustness bounds and distribution-induced shift decomposition of error (DISDE), reveals robustness differences that scalar accuracy metrics cannot detect.