CORAL: Disentangling Latent Representations in Long-Tailed Diffusion

TL;DR

This paper investigates why diffusion models perform poorly on long-tailed datasets, identifies latent class overlap as a key issue, and proposes CORAL, a contrastive regularization method, to improve tail class generation quality.

Contribution

The paper introduces CORAL, a contrastive regularization framework that disentangles latent representations in diffusion models trained on long-tailed data, improving tail class sample quality.

Findings

CORAL significantly enhances diversity of tail class samples.

CORAL improves visual quality of generated data for tail classes.

Latent class overlap is a key factor in diffusion model degradation on long-tailed 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.