Dual knowledge-guided data augmentation for robust clinical prediction models

Abstract

Domain shift poses a critical barrier to adopting medical artificial intelligence (AI) models, a challenge that is particularly acute in data-scarce, single-source domain generalization (SSDG) settings. Conventional data-augmentation techniques, including Mixup and input masking, fail to leverage the rich structural information and expert knowledge inherent in clinical data, causing models to overfit by learning spurious correlations. To address this, we propose a dual knowledge-guided data augmentation framework that enhances model robustness by systematically embedding clinical expertise into the training process. It comprises two novel components: similarity-guided Mixup, which generates clinically plausible virtual samples by interpolating between patients with similar clinical profiles, and group-based masking, which simulates realistic missing-data patterns by concurrently masking clinically correlated features. We validated our framework on the multicenter KNOW-pedCKD cohort for pediatric chronic kidney disease, training it exclusively on a single-source domain and evaluating it on three unseen target domains. It demonstrated significant performance gains in recall, a metric of critical importance in clinical settings. This study demonstrates that embedding domain knowledge into data augmentation is a promising strategy for developing generalizable and trustworthy medical AI models capable of operating reliably across heterogeneous clinical environments. The codes are available at https://github.com/msw6468/dual_augmentation_public.