Contrastive Conditional-Unconditional Alignment for Long-tailed Diffusion Model

TL;DR

This work tackles long-tailed data in class-conditional diffusion models by introducing Contrastive Conditional-Unconditional Alignment (CCUA), which pairs an Unsupervised Contrastive Loss (UCL) encouraging diversity with a Conditional-Unconditional Alignment Loss (AL) promoting knowledge sharing between head and tail classes. The method is applicable to both UNet-based diffusion and Diffusion Transformer backbones, and the overall objective combines $\mathcal{L}_{ccua} = \alpha \mathcal{L}_{ucl} + \gamma \mathcal{L}_{al}$ with the standard denoising diffusion loss. Empirical results across ImageNet-LT and other long-tailed datasets show consistent improvements in tail-class fidelity and diversity (e.g., FID_tail and Recall) without sacrificing head-class quality, outperforming vanilla DDPM, score-based diffusion, CBDM, OCLT, and Dispersive Loss baselines. The approach is practical to implement (batch resampling as optional) and demonstrates strong gains for both UNet and Diffusion Transformer architectures, suggesting broad applicability to long-tailed diffusion-based generation.

Abstract

Training data for class-conditional image synthesis often exhibit a long-tailed distribution with limited images for tail classes. Such an imbalance causes mode collapse and reduces the diversity of synthesized images for tail classes. For class-conditional diffusion models trained on imbalanced data, we aim to improve the diversity and fidelity of tail class images without compromising the quality of head class images. We achieve this by introducing two simple but highly effective loss functions. Firstly, we employ an Unsupervised Contrastive Loss (UCL) utilizing negative samples to increase the distance/dissimilarity among synthetic images. Such regularization is coupled with a standard trick of batch resampling to further diversify tail-class images. Our second loss is an Alignment Loss (AL) that aligns class-conditional generation with unconditional generation at large timesteps. This second loss makes the denoising process insensitive to class conditions for the initial steps, which enriches tail classes through knowledge sharing from head classes. We successfully leverage contrastive learning and conditional-unconditional alignment for class-imbalanced diffusion models. Our framework is easy to implement as demonstrated on both U-Net based architecture and Diffusion Transformer. Our method outperforms vanilla denoising diffusion probabilistic models, score-based diffusion model, and alternative methods for class-imbalanced image generation across various datasets, in particular ImageNet-LT with 256x256 resolution.

Figures (12)

• Figure 1: Synthetic images for a tail class (‘red wine’) by (a) standard SiT and (b) proposed CCUA framework. Both are trained on long-tailed ImageNet dataset for 160 epochs with 256x256 resolution.
• Figure 2: Overview of the proposed CCUA framework of contrastive learning for diffusion model. The noise estimation network is divided into two parts, a latent encode network $e(*)$ and a decode network $d(*)$. $e(*)$ encodes an image with noise $x_t$ to a low-dimensional latent $h$, which is decoded to noise $\epsilon$ by $d(*)$. Specifically, the latent encode network $e(*)$ and decode network $d(*)$ for unet-based model and diffusion transformer are shown in Fig. \ref{['fig:method_details']} in the appendix.\ref{['sec:details']}. We increase the distance of negative pairs of unconditional encoded latents for different samples by an unsupervised contrastive loss $\mathcal{L}_{ucl}$ with negative samples only, and align unconditional and conditional generation for the same sample $x_t^i$ and utilize an alignment loss $\mathcal{L}_{al}$ to minimize their distances at initial time steps.
• Figure 3: The leftmost column shows synthetic images with different classes but with the same initial noise or random seed. We visualize the low-frequency component and the high-frequency component from the reverse processing. It shows that low-frequency components are similar during initial time steps for different classes with the same initial noise si2024freeu. More examples are in appendix \ref{['sec:reverseprocessing']}.
• Figure 4: Synthetic images of SiT and it with CCUA for ImageNet-LT tail classes 'espresso' and 'window shade'. All methods start the denoising process from the same Gaussian noise at corresponding grid cells. CCUA shows consistently higher diversity and fidelity compared to SiT.
• Figure 5: Model details of CCUA framework with UNet-based Model and Diffusion Transformer. As shown in Fig. \ref{['fig:method']}, the noise estimation network is divided into two parts, a latent encoded network $e(*)$ and a decoded network $d(*)$. (a) For UNet-based architecture, $e(*)$ is defined as the UNet encoder, while $d(*)$ is defined as the UNet decoder. (b) For Diffusion Transformer, $e(*)$ is defined as all the SiT/DiT blocks, while $d(*)$ is defined as the final linear and adaLN projection layer.