CORAL: Disentangling Latent Representations in Long-Tailed Diffusion

2025-06-19

AI Summary

Diffusion models suffer from poor generation quality and low diversity for tail classes under long-tailed data distributions—not solely due to insufficient samples, but primarily because of severe latent representation overlap between tail and head classes in the U-Net bottleneck layer, leading to feature borrowing. This work is the first to identify class-relative imbalance itself as a fundamental cause of representation confusion. We propose CORAL, a supervised contrastive latent-space alignment framework: it introduces a class-labeled supervised contrastive loss (SupCon) at the bottleneck layer to enforce category-aware latent disentanglement—inter-class separation and intra-class compactness. Evaluated on multiple long-tailed image generation benchmarks, CORAL significantly improves tail-class performance: FID decreases by 23%, LPIPS improves by 31%, and all diversity metrics surpass state-of-the-art methods.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Natural Language Processing: Generation

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

Diffusion models have achieved impressive performance in generating high-quality and diverse synthetic data. However, their success typically assumes a class-balanced training distribution. In real-world settings, multi-class data often follow a long-tailed distribution, where standard diffusion models struggle -- producing low-diversity and lower-quality samples for tail classes. While this degradation is well-documented, its underlying cause remains poorly understood. In this work, we investigate the behavior of diffusion models trained on long-tailed datasets and identify a key issue: the latent representations (from the bottleneck layer of the U-Net) for tail class subspaces exhibit significant overlap with those of head classes, leading to feature borrowing and poor generation quality. Importantly, we show that this is not merely due to limited data per class, but that the relative class imbalance significantly contributes to this phenomenon. To address this, we propose COntrastive Regularization for Aligning Latents (CORAL), a contrastive latent alignment framework that leverages supervised contrastive losses to encourage well-separated latent class representations. Experiments demonstrate that CORAL significantly improves both the diversity and visual quality of samples generated for tail classes relative to state-of-the-art methods.

Problem

• Diffusion models struggle with long-tailed class distributions
• Latent representations overlap in tail and head classes
• Class imbalance degrades generation quality and diversity

Innovation

• Contrastive latent alignment framework
• Supervised contrastive losses
• Separated latent class representations