Class-Balancing Diffusion Models

Conference on Computer Vision and Pattern Recognition 2023 (CVPR 2023)

Diffusion-based models have shown the merits in generating high-quality visual data while preserving better diversity in recent studies. However, such observation is only justified on curated data distribution, where the data samples are nicely pre-processed to be uniformly distributed in terms of their labels. In practice, a long-tailed data distribution appears more often and how diffusion models perform on such label-imbalanced data remains unknown. In this work, we first investigate this problem and observe significant degradation in both diversity and fidelity when the diffusion model is trained on datasets with imbalanced class distributions. Especially on tail classes, the generations largely lose diversity and we observe severe mode-collapse issues. To tackle this problem, we set from the hypothesis that the data distribution is not class-balanced, and propose Class-Balancing Diffusion models that are trained with a distribution adjustment regularizer as a solution. Experiments show that, with this simple solution, the generated image by CBDM exhibit higher diversity and quality in both quantitative and qualitative ways. Our method benchmarked the generation results on CIFAR100/CIFAR100LT dataset, and shows outstanding performance on the downstream recognition task.