Contrastive Conditional–Unconditional Alignment for Long-tailed Diffusion Model

Training data for class-conditional image synthesis often ex-hibit a long-tailed distribution with limited amount of images for tailclasses. Such an imbalance causes mode collapse and reduces the diver-sity of synthesized images for tail classes. For class-conditional di!usionmodels trained with imbalanced data, we aim to improve the diversityand fidelity of tail class images without compromising the quality of headclass images. We propose contrastive conditional-unconditional align-ment (CCUA), which comprises two synergistic loss functions. Our firstloss is an Alignment Loss (AL) that aligns class-conditional generationwith unconditional generation at large timesteps. Alignment loss makesthe denoising process insensitive to class conditions for the initial steps,which enriches tail classes through knowledge sharing from head classes.Secondly, we diversify unconditional generation via an UnsupervisedContrastive Loss (UCL) to increase the distance/dissimilarity amongsynthetic images. We combine the two losses to implicitly diversify condi-tional generation. Our framework is easy to implement as demonstratedon both U-Net based architecture and Di!usion Transformer. Our methodoutperforms vanilla denoising di!usion probabilistic models, score-baseddi!usion model, and alternative contrastive methods for class-imbalancedimage generation across various datasets, in particular ImageNet-LT with256→256 resolution.