Mitigating Dataset Shift via Smart Augmentation with Conditional Diffusion Models
摘要
The deployment of AI systems in real-world settings is often hindered by poor generalization to underrepresented subgroups in training data, resulting in biased or unreliable predictions. To address this issue, we propose an automated data augmentation method based on conditional diffusion models. Unlike previous approaches that rely on predefined feature conditioning, our method leverages a representation audit of the training data, based on autoencoder reconstruction errors, to identify and target underrepresented subpopulations. These reconstruction errors are used to cluster training samples, and the resulting cluster labels condition the diffusion model to generate group-specific synthetic data. We evaluate three augmentation strategies (extreme-cluster, uniform and proportional) under a dataset shift scenario, where the test set includes subgroups underrepresented in the training data. Our results show that the method improves predictive performance for these subgroups, particularly when using the proportional strategy, without degrading performance on well-represented groups. Our approach contributes to the development of more robust and trustworthy AI systems by addressing data underrepresentation through a targeted, representation-aware data augmentation. The proposed method eliminates the need for manually defined feature combinations, enabling dataset shift mitigation in an unsupervised manner. Code is available at this .