Training Class-Imbalanced Diffusion Model Via Overlap Optimization

TL;DR

This paper introduces a contrastive learning-based method to improve the quality of tail class image synthesis in class-imbalanced diffusion models, addressing bias towards frequent classes and reducing overlap between class distributions.

Contribution

It proposes a novel probabilistic contrastive learning approach applicable to class-conditional diffusion models to enhance synthesis quality for long-tailed datasets.

Findings

Significant improvement in image quality for tail classes.

Effective handling of imbalanced data in diffusion-based models.

Method applicable to various datasets with long-tailed distributions.

Abstract

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.