Data augmentation improves learning under changing input distributions
Rethinking Data Augmentation under Covariate Shift: Invariant-Guided Diffusion and Prototype Reweighting
Machine LearningArtificial Intelligence
Summary
When computer programs learn from data, they often struggle if the data they see later changes or if there is not enough data. The authors studied a common problem where data shifts between training and real use, and simple data copying doesn’t help. They proposed a new method that creates useful new examples by focusing on parts of the data that are steady and important for making decisions. They also developed a way to trust whole groups of data points instead of single samples, which reduces mistakes when adapting the program to new situations.
What this means in practice
- •For industrial data teams: Generate stable synthetic tabular data to improve model robustness when data distributions change after deployment.
- •For financial risk analysts: Create reliable data augmentations that help models handle evolving data patterns for credit or fraud detection decisions.
Authors
Hongyu Cao, Xinyuan Wang, Arun Vignesh Malarkkan, Kunpeng Liu, Haifeng Chen, Yanjie Fu
Abstract
In many industrial applications, 1) tabular data is scarce and imbalanced and thus requires synthetic expansion; 2) input distributions drift between training and deployment (covariate shift); 3) validation sets often diverge from unseen test environments; or 4) standard generative models simply mimic outdated source distributions. This learning setting limits the stability of standard augmentation and adaptation pipelines. We generalize the task under such setting as the Augmented and Weighted Learning under Covariate Shift problem (AWL-CS). AWL-CS imposes two critical challenges on existing methods: 1) misleading generative guidance where models optimize for source similarity rather than downstream task relevance, and 2) structural instability of distributional density where reweighting mechanisms overfit to noisy validation signals. To tackle these challenges, we propose IGDPR (Invariant-Guided Diffusion with Prototype Reweighting), a unified framework that synergizes stable synthesis and structural adaptation: i) To achieve task-relevant generation, we steer the diffusion sampling process using invariant potentials to ensure synthetic samples align with stable decision boundaries rather than outdated correlations. ii) To ensure stable adaptation, we develop a prototype-based reweighting strategy that assesses sample reliability through structural clusters instead of isolated points, effectively filtering validation noise. Extensive experiments on real data demonstrate our method improves data quality by augmenting the most beneficial data for robust learning.