# Training Class-Imbalanced Diffusion Model Via Overlap Optimization

arXiv: https://arxiv.org/abs/2402.10821 · HTML: https://mlsanity.com/paper/2402.10821/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 44 of its 58 references that are in the map.
- First posted on arXiv 16 Feb 2024.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 44)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [Adding Conditional Control to Text-to-Image Diffusion Models](https://mlsanity.com/paper/2302.05543/) (2023)
- [Rethinking the Inception Architecture for Computer Vision](https://mlsanity.com/paper/1512.00567/) (2015)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [DreamFusion: Text-to-3D using 2D Diffusion](https://mlsanity.com/paper/2209.14988/) (2022)
- [Improved Techniques for Training GANs](https://mlsanity.com/paper/1606.03498/) (2016)
- [Supervised Contrastive Learning](https://mlsanity.com/paper/2004.11362/) (2020)
- [Training Generative Adversarial Networks with Limited Data](https://mlsanity.com/paper/2006.06676/) (2020)
- [Learning Imbalanced Datasets with Label-Distribution-Aware Margin Loss](https://mlsanity.com/paper/1906.07413/) (2019)
- [Improved Precision and Recall Metric for Assessing Generative Models](https://mlsanity.com/paper/1904.06991/) (2019)
- [Decoupling Representation and Classifier for Long-Tailed Recognition](https://mlsanity.com/paper/1910.09217/) (2019)
- [Long-tail learning via logit adjustment](https://mlsanity.com/paper/2007.07314/) (2020)
- [Exploring the Limits of Weakly Supervised Pretraining](https://mlsanity.com/paper/1805.00932/) (2018)
- [Assessing Generative Models via Precision and Recall](https://mlsanity.com/paper/1806.00035/) (2018)
- [BBN: Bilateral-Branch Network with Cumulative Learning for Long-Tailed Visual Recognition](https://mlsanity.com/paper/1912.02413/) (2019)
- [Balanced Meta-Softmax for Long-Tailed Visual Recognition](https://mlsanity.com/paper/2007.10740/) (2020)
- [Differentiable Augmentation for Data-Efficient GAN Training](https://mlsanity.com/paper/2006.10738/) (2020)
- [Synthetic Data from Diffusion Models Improves ImageNet Classification](https://mlsanity.com/paper/2304.08466/) (2023)
- [Long-tailed Recognition by Routing Diverse Distribution-Aware Experts](https://mlsanity.com/paper/2010.01809/) (2020)
- [Long-Tailed Classification by Keeping the Good and Removing the Bad Momentum Causal Effect](https://mlsanity.com/paper/2009.12991/) (2020)
- [Diffusion Probabilistic Modeling for Video Generation](https://mlsanity.com/paper/2203.09481/) (2022)
- [Disentangling Label Distribution for Long-tailed Visual Recognition](https://mlsanity.com/paper/2012.00321/) (2020)
- [Balanced Contrastive Learning for Long-Tailed Visual Recognition](https://mlsanity.com/paper/2207.09052/) (2022)

## Built on by (most cited first, 6 of 6)

- [GoodDrag: Towards Good Practices for Drag Editing with Diffusion Models](https://mlsanity.com/paper/2404.07206/) (2024)
- [A Comprehensive Survey on Imbalanced Data Learning](https://mlsanity.com/paper/2502.08960/) (2025)
- [Imbalance in Balance: Online Concept Balancing in Generation Models](https://mlsanity.com/paper/2507.13345/) (2025)
- [MCLR: Improving Conditional Modeling via Inter-Class Likelihood-Ratio Maximization and Unifying Classifier-Free Guidance with Alignment Objectives](https://mlsanity.com/paper/2603.22364/) (2026)
- [Contrastive Conditional-Unconditional Alignment for Long-tailed Diffusion Model](https://mlsanity.com/paper/2507.09052/) (2025)
- [NS-Pep: De novo Peptide Design with Non-Standard Amino Acids](https://mlsanity.com/paper/2510.03326/) (2025)

## Abstract (arXiv)

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.