CORAL: Disentangling Latent Representations in Long-Tailed Diffusion
This addresses a critical issue for real-world applications where data is imbalanced, though it is an incremental improvement over existing diffusion model techniques.
The paper tackles the problem of diffusion models performing poorly on tail classes in long-tailed datasets by identifying overlapping latent representations as the cause, and proposes CORAL, a contrastive regularization method that improves tail-class sample diversity and visual quality significantly over state-of-the-art methods.
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.