Training Class-Imbalanced Diffusion Model Via Overlap Optimization
Divin Yan Lu Qi Vincent Tao Hu Ming-Hsuan Yang Meng Tang
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.