Technical Summary: Contrastive Conditional–Unconditional Alignment for Long-tailed Diffusion Model

Problem Statement

Training data for class-conditional image synthesis frequently exhibits a long-tailed distribution, characterized by a scarcity of images for "tail" classes compared to "head" classes. This imbalance causes diffusion models to suffer from mode collapse and reduced diversity when generating images for tail classes. Existing methods, including standard denoising diffusion probabilistic models (DDPM) and score-based diffusion models (SBDM), often fail to generate high-fidelity and diverse images for these underrepresented classes without compromising the quality of head class generation. While contrastive learning has proven effective for long-tailed recognition and GAN-based generation, its application to diffusion models for class-imbalanced generation requires novel adaptation to address the specific mechanics of the denoising process.

Methodology: Contrastive Conditional–Unconditional Alignment (CCUA)

The authors propose CCUA, a framework comprising two synergistic loss functions designed to improve the diversity and fidelity of tail class images while maintaining head class quality. The method is compatible with both U-Net-based architectures and Diffusion Transformers (DiT/SiT).

1. Unsupervised Contrastive Loss (UCL) for Unconditional Generation

To address mode collapse, the authors introduce an unsupervised contrastive loss applied to unconditional generation latents.

• Mechanism: The noise estimation network is split into a latent encoder $e(\cdot)$ and a decoder $d(\cdot)$. The encoder maps noisy images $x_t$ to low-dimensional latents $h$.
• Loss Formulation: The loss uses an InfoNCE-style formulation with negative samples only. Within a mini-batch, each image serves as an anchor, and all other images in the batch serve as negative samples. There are no positive pairs created via data augmentation; the anchor is treated as its own positive vector.
• Objective: This creates a "repulsive force" in the latent space, maximizing the distance between synthetic images regardless of their class. This forces the model to diversify unconditional latents, preventing them from collapsing to a constant vector or overfitting to limited training samples.
• Theoretical Justification: Directly applying contrastive loss to conditional latents can lead to a trivial solution where the model maximizes inter-class diversity (separating classes) while ignoring intra-class diversity. By applying the loss to unconditional latents, the model is forced to diversify all samples, which is then distilled into the conditional generation via the second loss.

2. Alignment Loss (AL) for Conditional-Unconditional Alignment

The second component aligns the conditional generation with the diversified unconditional generation.

• Mechanism: The loss minimizes the KL divergence between the conditional distribution $p_\theta(x_{t-1}|x_t, c)$ and the unconditional distribution $p_\theta(x_{t-1}|x_t)$. In practice, this is simplified to an MSE loss between the estimated noise $\epsilon_\theta(x_t, t, c)$ and $\epsilon_\theta(x_t, t, \emptyset)$.
• Time-Weighting: Crucially, this alignment is weighted more heavily at large timesteps (the initial steps of the reverse process).
• Motivation: The authors observe that during the initial denoising steps, the low-frequency components of images from different classes (and between conditional and unconditional generation) are highly similar. By aligning these initial steps, the model facilitates knowledge sharing from head classes (abundant data) to tail classes (scarce data). This allows tail classes to benefit from the rich diversity learned by head classes during the early stages of generation.

3. Synergy

The framework combines the standard DDPM loss with the proposed CCUA loss:
$$\mathcal{L}_{total} = \mathcal{L}_{ddpm} + \alpha \cdot \mathcal{L}_{ucl} + \gamma \cdot \mathcal{L}_{al}$$
The synergy lies in the fact that UCL diversifies the unconditional latent space, and AL aligns the conditional latent space to this diversified unconditional space. This indirect diversification of conditional latents is empirically shown to be more effective than directly applying contrastive loss to conditional latents.

Key Contributions

1. CCUA Framework: A novel method for class-imbalanced diffusion models that combines unsupervised contrastive learning on unconditional latents with a conditional-unconditional alignment loss.
2. Alignment Strategy: The proposal to align conditional and unconditional generation specifically at large timesteps (initial denoising steps) to leverage the similarity of low-frequency components and enable knowledge transfer from head to tail classes.
3. Unsupervised Contrastive Regularization: The use of an unsupervised contrastive loss with negative samples only on unconditional latents to maximize intra-class and inter-class diversity without requiring class labels for the contrastive pairs.
4. Architecture Agnosticism: The method is demonstrated to be effective on both U-Net-based architectures and Diffusion Transformers (SiT).

Experimental Results

The method was evaluated on ImageNet-LT (256×256), TinyImageNet-LT (64×64), Places-LT (64×64), and CIFAR-LT datasets.

• Quantitative Performance:
  • On ImageNet-LT, CCUA significantly outperforms baselines including SiT, CBDM, REPA, and Dispersive Loss.
  • At 900k training steps, CCUA achieved an FID of 15.1 (vs. 19.9 for SiT) and a FIDtail of 22.5 (vs. 33.9 for SiT), representing a ~11.4 point improvement in tail class fidelity.
  • The method also improved Recall (diversity) and Precision (fidelity) metrics across the board.
  • On TinyImageNet-LT, CCUA achieved the best FID (12.0) and FIDtail (20.8), outperforming DDPM, CBDM, and OCLT.
• Qualitative Results: Visualizations show that CCUA generates images with higher diversity and fidelity for tail classes (e.g., 'red wine', 'espresso', 'window shade') compared to standard SiT, which often produces repetitive or low-quality images for these classes.
• Ablation Studies:
  • The combination of both losses is necessary; using only one yields suboptimal results.
  • Applying the alignment loss at large timesteps (linear weighting) is superior to constant or quadratic weighting.
  • Using latent representations from the N/4-th block of the transformer generally offers the best trade-off between overall FID and spatial FID/Recall.

Significance and Claims

The paper claims that CCUA effectively addresses the long-tailed generation problem by adapting principles from long-tailed GANs (specifically unconditional training at lower resolutions/early stages) to the diffusion paradigm. By decoupling the diversification task (handled by unconditional contrastive learning) from the conditioning task (handled by alignment), the method avoids the pitfalls of direct conditional contrastive regularization.

The authors emphasize that their approach:

• Improves tail class generation without degrading head class quality.
• Is easy to implement within existing DDPM and SBDM pipelines.
• Outperforms concurrent works like Dispersive Loss and DiffROP, particularly in scenarios with extreme class imbalance (e.g., ImageNet-LT with an imbalance factor of 0.001).
• Provides a theoretical and empirical justification for why aligning conditional and unconditional generation at the initial steps of the reverse process is beneficial for knowledge sharing in long-tailed settings.