Training Class-Imbalanced Diffusion Model Via Overlap Optimization
This addresses a practical limitation in generative AI for real-world applications with imbalanced data, though it appears to be an incremental improvement on existing class-conditional diffusion models.
The paper tackles the problem of diffusion models producing inferior quality images for rare classes in long-tailed datasets by proposing a contrastive learning method to minimize overlap between synthetic image distributions of different classes, showing significant improvements in image synthesis across multiple datasets.
Diffusion models have made significant advances recently in high-quality image synthesis and related tasks. However, diffusion models trained on real-world datasets, which often follow long-tailed distributions, yield inferior fidelity for tail classes. Deep generative models, including diffusion models, are biased towards classes with abundant training images. To address the observed appearance overlap between synthesized images of rare classes and tail classes, we propose a method based on contrastive learning to minimize the overlap between distributions of synthetic images for different classes. We show variants of our probabilistic contrastive learning method can be applied to any class conditional diffusion model. We show significant improvement in image synthesis using our loss for multiple datasets with long-tailed distribution. Extensive experimental results demonstrate that the proposed method can effectively handle imbalanced data for diffusion-based generation and classification models. Our code and datasets will be publicly available at https://github.com/yanliang3612/DiffROP.