CORAL: Disentangling Latent Representations in Long-Tailed Diffusion

Diffusion models have achieved impressive performance in generating high-quality and diverse synthetic data. However, their success typically assumes a classbalanced training distribution. In real-world settings, multi-class data often follow a long-tailed distribution, where standard diffusion models struggleproducing lowdiversity 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. The implementation code is available at https://github.com/SankarLab/coral-lt-diffusion.

Recent work has sought to improve generative models under long-tailed class distributions by addressing sampling imbalance and promoting class-aware generation. Class-Balancing Diffusion Models (CBDMs) [5] introduce a regularizer that encourages balanced sampling across classes by penalizing deviations from a target distribution. In particular, the approach enhances tail generation based on the model prediction on the head class. This increased reliance on the model prediction and conditional priors introduces bias and can potentially reduce robustness (e.g., lead to class entanglement) during training. To address these limitations, Zhang et al. [6] propose a Bayesian 39th Conference on Neural Information Processing Systems (NeurIPS 2025).