Title: CORAL: Disentangling Latent Representations in Long-Tailed Diffusion
Venue: NeurIPS 2025
Accepted: yes
Keywords: Diffusion

Latent

Contrastive

Metric Learning

Long-Tailed

Imbalanced

Generative Models

Classifier-Free Guidance
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 underrepresented (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 **CO**ntrastive **R**egularization for **A**ligning **L**atents (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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the performance degradation of diffusion models on minority classes under class-imbalanced settings and identifies the cause as overlapping latent representations at the U-Net bottleneck. To address this, the authors introduce a contrastive loss at the bottleneck to encourage more meaningful and well-separated representations, which helps mitigate the issue. Empirically, the method shows modest performance improvements on long-tailed datasets.\n\nSTRENGTHS AND WEAKNESSES:\n**Strength**\n\n- The approach of tackling the problem through embedding-based analysis is novel, and the use of semantic embedding learning appears to address the issue effectively. Overall, it seems to be a highly appropriate and well-justified direction. (This reason alone makes me lean toward the positive side.)\n\n- The method appears to improve performance across all evaluation metrics.\n\n\n**Weaknesses**\n- Limited impacts that only focus on Unet-based diffusion models for CIFAR datasets (See questions and limitations).\n\nQUESTIONS:\n1. Is there a way this idea could be applied to T2I models?\n2. Recently, the diffusion transformer architecture has become dominant. Does a similar phenomenon occur in transformer architecture as well?\n3. Suppose we're dealing with unconditional modeling rather than a class-conditional setting, but the dataset itself is biased. In that case, would attaching a contrastive loss to the UNet bottleneck still be beneficial?\n\nLIMITATIONS:\n1. There is a line of impact papers on diffusion autoencoders [1] and diffusion bridge autoencoders [2] that are trained together with an auxiliary encoder module. In modules that utilize an auxiliary encoder like this, does the issue you mentioned still occur in the U-Net bottleneck or the encoded latents? It might be good to either conduct experiments on these papers as well or at least include them as related work in the context of semantic representation learning in diffusion models.\n\n2. Ablation on $\\lambda$ is missing.\n\n\n[1] [CVPR 2022 oral] Diffusion Autoencoders: Toward a Meaningful and Decodable Representation\n\n[2] [ICLR 2025 spotlight] Diffusion Bridge AutoEncoders for Unsupervised Representation Learning", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces CORAL, a novel framework designed to address the challenges faced by diffusion models when trained on long-tailed data distributions, where tail classes are often underrepresented. To mitigate this issue, the authors attach a projection head to the encoder of the U-Net architecture and apply a supervised contrastive loss on its output, aiming to disentangle the latent space representations across classes. By promoting class-specific separation in the latent space, the model better preserves discriminative and class-relevant features. The training objective is augmented with a supervised contrastive loss, incorporated into the original diffusion loss via a time-dependent weighting factor $\\lambda(t)$, resulting in a regularized objective.\nExperimental results on various long-tailed datasets demonstrate that CORAL significantly improves generation quality and diversity, especially for tail classes. Furthermore, latent space visualizations support that the proposed method successfully achieves well-disentangled class representations.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n- One of the key strengths of this paper is its simplicity and effectiveness: by appending only a lightweight linear projection head to the encoder bottleneck of the U-Net and applying a supervised contrastive loss, the authors successfully disentangle latent representations across classes. This disentanglement alleviates the representation entanglement issue commonly observed in diffusion models trained on long-tailed datasets, particularly benefiting tail-class generation. The impact of this regularization is well-supported both numerically (e.g., improved FID, IS, Recall, F1/8 in Table 1 and 3) and visually (e.g., Figure 1, 2, 10, and 13), where t-SNE and UMAP plots demonstrate clear class-wise separation in the latent space.\n\nWeaknesses\n- Although CORAL is primarily evaluated under long-tailed data settings, its core mechanism—contrastive latent regularization—seems broadly applicable beyond class imbalance. This raises the question of whether the problem formulation should have been positioned more generally as a method for representation disentanglement in diffusion models. Broadening the scope could improve both the conceptual clarity and the perceived applicability of the method.\n\n- The paper lacks a detailed explanation of how class-wise latent disentanglement directly leads to improved sample quality, particularly in terms of FID. Introducing regularization in the latent space could, in theory, interfere with the likelihood-based training objective of DDPM and potentially steer the model away from the true data distribution. A more thorough analysis or targeted ablations would help clarify how the proposed regularization supports or enhances generative performance.\n\n- Additionally, the paper adopts classifier-free guidance (CFG) during sampling but does not specify how the class label $y$ is determined at inference time. This omission is critical, as $y$ directly influences the conditional generation process and affects class-specific evaluations such as per-class FID and qualitative comparisons (e.g., Figures 3, 4, and 15). Without a clear description of how $y$ is chosen, the reproducibility and interpretability of the results may be compromised.\n\n- Finally, while the method is effective, the idea of applying contrastive loss for latent regularization may appear somewhat marginal or incremental, given its conceptual similarity to existing contrastive learning techniques in other generative frameworks. Additional novelty or theoretical grounding could help distinguish it further from prior work.\n\nQUESTIONS:\nThis question is related to the weaknesses mentioned above. It would be helpful if the paper provided more detailed explanations of the experimental setup. For example, when using classifier-free guidance (CFG) during sampling, how were the class labels sampled or selected? In addition, since the dataset follows a long-tailed distribution, the number of samples per class differs—was this imbalance accounted for or corrected when computing class-wise FID? Lastly, since contrastive learning is often sensitive to batch size, it would be valuable to include an ablation study examining the impact of batch size on performance.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a method for improving the generations of tail classes. The method is motivated by the observation that in the bottleneck features of the diffusion model, the latent representations of classes of the head and tail classes overlap. The method involved training the diffusion model with a contrastive loss on the bottleneck features. Supportive results were shown on several long-tailed datasets.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n- The paper was well written and easy to follow.\n\n**Weakness**\n\n- There could be more thorough experiments to justify key design choices of the method. Several components of the method, e.g., where to apply the loss (e.g., bottleneck features vs. denoised images), or the selected loss function are used without clear motivation or empirical evidence. E.g., it is not clear why applying the contrastive loss at the bottleneck layer is preferred over applying it directly to the denoised output, similar to DiffROP [4]. The choice of contrastive loss over simpler alternatives such as reweighting schemes is not well-justified.\n- The paper highlights the overlap between head and tail classes in the bottleneck features. However, DiffROP made the observation of an overlap between the tail and head classes although in the image space. It makes sense that in some intermediate layer their features also overlap.\n- It seems like the contrastive loss could reduce the diversity of the generations within each class. This can also be seen from Fig. 1.\n\nQUESTIONS:\n- In Fig 3, why is the performance of DDPM poor for all classes? I would expect that it would perform similarly to the baseline for the head class. Which is the head/tail class in this example?\n- With regards to the last point above, did the authors observe a decrease in the diversity of the generations in the non tail classes with the contrastive loss?\n- In what settings would such methods (training a custom diffusion model from scratch on long-tailed data), be preferable to controlling a pre-trained model like Stable Diffusion via e.g., prompt engineering, textual inversion?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses the challenge of generating high-quality, diverse images for underrepresented classes in long-tailed datasets using diffusion models. It identifies that poor performance on tail classes is due to representation entanglement—overlap in the latent space between head and tail classes within the U-Net bottleneck. To mitigate this, the authors propose CORAL (Contrastive Regularization for Aligning Latents), which adds a projection head to the bottleneck and applies a supervised contrastive loss to promote class separation. CORAL significantly improves tail-class generation across multiple datasets, outperforming prior methods both quantitatively and qualitatively.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\nThe paper makes a strong empirical observation that latent feature overlap in diffusion models is a core failure mode for tail-class generation under long-tailed distributions. This is a meaningful contribution that deepens our understanding of diffusion behavior. CORAL introduces a straightforward architectural addition (a projection head with contrastive loss) to existing diffusion frameworks, showing consistent improvements across multiple datasets and metrics, without requiring external models or complex rebalancing schemes.\n\nWeaknesses\n\n* Limited methodological novelty: The main contribution lies in *where* contrastive learning is applied (i.e., the bottleneck of U-Net), rather than in the contrastive mechanism itself. The approach builds on well-established tools (e.g., SupCon loss, projection head), making the core technique relatively incremental.\n\n* Shallow theoretical grounding: The paper lacks theoretical justification or formal analysis of why contrastive regularization in latent space works better than in the ambient space. The intuition is sound but remains at an empirical level.\n\n* Underdeveloped ablation and sensitivity analysis**: While some ablations are mentioned, there is limited exploration of how sensitive CORAL is to design choices like the projection head architecture, temperature hyperparameters, or the exact form of φ(t). This weakens the assessment of robustness.\n\n* No analysis of class-conditional generation quality: The evaluation is mostly focused on overall metrics and per-class FID. It does not deeply assess whether generated tail-class images are semantically correct or align well with class labels, particularly under CFG.\n\nQUESTIONS:\n1. Is the contribution of CORAL substantial enough, given that it primarily repurposes existing contrastive learning techniques without introducing novel mechanisms?\n2. Why does contrastive regularization in the latent space outperform ambient-space approaches, and can this be theoretically justified beyond empirical intuition?\n3. How sensitive is CORAL to architectural and hyperparameter choices, and can more comprehensive ablations improve confidence in its robustness?\n4. Do the generated tail-class images accurately reflect their intended semantics, and how well does CORAL preserve class fidelity under conditional generation?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}