journal article·2026-07-31

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

Shuang Song, Muhammad Syafiq Mohd Pozi, Nik Fatinah N. Mohd Farid

Applied Sciences· Volume 16· Issue 15
Publisher page

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.
Source: crossref · Quality: baseline

Affiliations

School of Computing, Universiti Utara Malaysia, Sintok 06010, Kedah, Malaysia

References (45)

Show source-provided references
1. WILDS: A Benchmark of In-the-Wild Distribution Shifts
2. Evaluation of Domain Generalization and Adaptation on Improving Model Robustness to Temporal Dataset Shift in Clinical Medicine10.1038/s41598-022-06484-1 ↗
3. Quiñonero-Candela, J. (2009). Dataset Shift in Machine Learning, MIT Press.10.7551/mitpress/9780262170055.001.0001 ↗
4. A Unified Framework for Dataset Shift Diagnostics10.1016/j.ins.2023.119612 ↗
5. Gardner, J., Popovic, Z., and Schmidt, L. (2023). Benchmarking Distribution Shift in Tabular Data with TableShift. Proceedings of the Advances in Neural Information Processing Systems (NeurIPS), Curran Associates, Inc.10.52202/075280-2324 ↗
6. A Unifying Causal Framework for Analyzing Dataset Shift-Stable Learning Algorithms10.1515/jci-2021-0042 ↗
7. Arjovsky, M., Bottou, L., Gulrajani, I., and Lopez-Paz, D. (2020). Invariant Risk Minimization. arXiv.
8. Rosenfeld, E., Ravikumar, P., and Risteski, A. (2021, January 3–7). The Risks of Invariant Risk Minimization. Proceedings of the International Conference on Learning Representations, Vienna, Austria.
9. Wang, L., Richardson, T., and Robins, J. (2025). Causal Inference: A Tale of Three Frameworks. arXiv.10.6339/25-JDS1211 ↗
10. Generalized Out-of-Distribution Detection: A Survey10.1007/s11263-024-02117-4 ↗
11. Making Sense of Sensitivity: Extending Omitted Variable Bias10.1111/rssb.12348 ↗
12. A Causal Framework for Distribution Generalization10.1109/TPAMI.2021.3094760 ↗
13. Causal Inference by Using Invariant Prediction: Identification and Confidence Intervals10.1111/rssb.12167 ↗
14. Stable Learning Establishes Some Common Ground Between Causal Inference and Machine Learning10.1038/s42256-022-00445-z ↗
15. Counterfactual Invariance to Spurious Correlations in Text Classification
16. Zhou, C., Ma, X., Michel, P., and Neubig, G. (2021). Examining and Combating Spurious Features under Distribution Shift. Proceedings of Machine Learning Research; ICML 2021, JMLR.
17. Kuang, K., Xiong, R., Cui, P., Athey, S., and Li, B. (2018). Stable Prediction across Unknown Environments. arXiv.10.1145/3219819.3220082 ↗
18. Shen, Z., Cui, P., Zhang, T., and Kuang, K. (2019). Stable Learning via Sample Reweighting. arXiv.
19. Limitations of the Application of Fourfold Table Analysis to Hospital Data10.2307/3002000 ↗
20. Kull, M., and Flach, P. (2014, January 21–26). Patterns of Dataset Shift. Proceedings of the ICML Workshop on Distribution Shift, Beijing, China.
21. Roschewitz, M., Mehta, R., Jones, C., and Glocker, B. (2025). Automatic Dataset Shift Identification to Support Safe Deployment of Medical Imaging AI. arXiv.10.1007/978-3-032-04981-0_7 ↗
22. Shortcut Learning in Deep Neural Networks10.1038/s42256-020-00257-z ↗
23. Reddy, A.G., and Balasubramanian, V.N. (2023). Detecting and Measuring Confounding Using Causal Mechanism Shifts. arXiv.10.52202/079017-1971 ↗
24. Multi-Source Domain Adaptation: A Causal View
25. Domain Generalization: A Survey
26. From Concept Drift to Model Degradation: An Overview on Performance-Aware Drift Detectors10.1016/j.knosys.2022.108632 ↗
27. Software Architecture Description in Original Software Publications10.1016/j.simpa.2025.100802 ↗
28. Covariate Shift Adaptation by Importance Weighted Cross Validation
29. Detecting and Correcting for Label Shift with Black Box Predictors
30. Wang, D., Shelhamer, E., Liu, S., Olshausen, B., and Darrell, T. (2021, January 3–7). Tent: Fully Test-Time Adaptation by Entropy Minimization. Proceedings of the International Conference on Learning Representations, Vienna, Austria.
31. Anchor Regression: Heterogeneous Data Meet Causality10.1111/rssb.12398 ↗
32. Thams, N., Oberst, M., and Sontag, D. (2022). Evaluating Robustness to Dataset Shift via Parametric Robustness Sets. Proceedings of the Advances in Neural Information Processing Systems (NeurIPS), Curran Associates, Inc.10.52202/068431-1228 ↗
33. Yu, H., Liu, J., Zhang, X., Wu, J., and Cui, P. (2024). A Survey on Evaluation of Out-of-Distribution Generalization. arXiv.
34. Tibshirani, R.J., Barber, R.F., Candès, E.J., and Ramdas, A. (2019). Conformal Prediction Under Covariate Shift. Proceedings of the Advances in Neural Information Processing Systems (NeurIPS), Curran Associates, Inc.
35. Park, S., Yang, S., Choo, J., and Yun, S. (2023). Label Shift Adapter for Test-Time Adaptation under Covariate and Label Shifts. Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), IEEE.10.1109/ICCV51070.2023.01505 ↗
36. Causal discovery and inference: Concepts and recent methodological advances10.1186/s40535-016-0018-x ↗
37. Implicit Data Crimes: Machine Learning Bias Arising from Misuse of Public Data10.1073/pnas.2117203119 ↗
38. A Causal Perspective on Dataset Bias in Machine Learning for Medical Imaging10.1038/s42256-024-00797-8 ↗
39. Limitations of the Application of Fourfold Table Analysis to Hospital Data10.1093/ije/dyu022 ↗
40. ELSA: Efficient Label Shift Adaptation through the Lens of Semiparametric Models
41. DoWhy-GCM: An Extension of DoWhy for Causal Inference in Graphical Causal Models
42. Estimating Individual Treatment Effect: Generalization Bounds and Algorithms
43. Johansson, F.D., Shalit, U., Kallus, N., and Sontag, D. (2023). Generalization Bounds and Representation Learning for Estimation of Potential Outcomes and Causal Effects. arXiv.
44. Retiring Adult: New Datasets for Fair Machine Learning
45. Sagawa, S., Koh, P.W., Lee, T., Gao, I., Xie, S.M., Shen, K., Kumar, A., Hu, W., Yasunaga, M., and Marklund, H. (2022). Extending the WILDS Benchmark for Unsupervised Adaptation. arXiv.