CORAL: Disentangling Latent Representations in Long-Tailed Diffusion
Authors: Esther Rodriguez
Monica Welfert
Samuel McDowell
Nathan Stromberg
Julian Antolin Camarena
Lalitha Sankar
Venue: NeurIPS 2025 poster ()
Keywords: Diffusion
Latent
Contrastive
Metric Learning
Long-Tailed
Imbalanced
Generative Models
Classifier-Free Guidance
Primary area: deep_learning

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.

CORAL: Disentangling Latent Representations in Long-Tailed Diffusion
Esther Rodriguez Arizona State University earodr32@asu, edu
Monica Welfert Arizona State University
Samuel McDowell Arizona State University scmcdowe@asu. edu
mwelfert@asu.en edu
Nathan Stromberg Arizona State University
Julian Antolin Camarena Arizona State University jantolin@asu edu
Lalitha Sankar
Arizona State University 1sankar@asu. edu
nstrombe@asu.com edu
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 long-tailed distribution, where standard diffusion models struggleproducing low- diversity and lower-quality samples for tail classes. While this degradation welledecu its underlying cause remains poorly understood. In this work, we investigate the behavior of diffusion models trained on long-tailed datasets and identify 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 not merely due to limited data per class, but that the relative class im- balance significantly contributes 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-aa methods. The implementation code is com wwt..////aaaaaaaaaaaaarala-----e
available https:
1 Introduction
Diffusion models (DMs) [T] have achieved impressive performance in generating high-quality and diverse samples across a range domains. However, their success typically relies on class-balanced training data. In practice, many real-world datasets exhibit long-tailed class distributions, where small number of head classes contain the majority of samples, while many tail classes are significantly underrepresented (3]. Under such imbalance, DMs often fail to generate faithful and diverse outputs for tail classes, instead exhibiting feature borrowing, where samples from rare classes display a mix of tail and head features
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 regularizer that encourages balanced sampling across classes by penalizing deviations from 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 potentially reduce robustness (e.g., lead class entanglement) during training. To address these limitations, Zhang et al. [6] propose a Bayesian
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

framework with weighted denoising score motee-ng and a gating mechanism to selectively transfer information from head tail classes. Other works incorporate contrastive learning 7 a key technique in metric learning [8) to improve class separability. For example, Yan et al. 4 propose a probabilistic contrastive approach that reduces overlap among class-conditional distributions to enhance tail-class generation. While these approaches have improved performance in imbalanced settings, they primarily operate in the ambient image space or introduce latent representations external to the denoising process. In contrast, relatively little attention has been given to structuring class representations within the latent space of the denoising network itself, highlighting a key gap in current methods.
The core generative process relies on neural architecture that processes data through lower- dimensional latent space. Current models predominantly use the U-Net architecture [2 [ which www.bboteeeeec-eeeeeee
structure consisting of convolutional neural networks to downsample, followed by multilayer perceptron, and then upsampling by further convolutional neural networks back into image space. There are skip connections from the encoder to the decoder to preserve data and feature information and mitigate vanishing gradients. It has been shown that the U-Net's bottleneck output carries semantic meaning [10]. One of our key observations is that, under long-tailed distributions, tail-class samples tend to occupy regions in this latent space that overlap heavily with head classes. This overlap - which we refer to as representation entanglement undermines model's the ability to preserve class class-specific features, leading poor generative performance for tail classes. We base this observation on extensive visualizations of long-tailes datasets for diffusion using various distance- presuecvign mappings. Figure
incorporates an
T Tillustrates this effect for t-SNE 11 and shows how tail-class representations are absorbed into dominant clusters.
To address this, we propose COntrastive Regularization for Aligning Latents (CORAL), contrastive latent alignment method that operates directly on the latent representations within the denoising network. Inspired by metric learning and its applications to learning representations [2] 13 [4], CORAL augments the encoder of the denoising U-Net with a projection head applied to the bottleneck output. The resulting projected embeddings are trained using a supervised contrastive loss, which is then combined with the standard diffusion objective. This encourages the model to pull together representations of samples from the same class while pushing apart those from different classes, thereby promoting class-wise separation in the latent space. In contrast to prior work that applics contrastive losses in the ambient auxiliary latent spaces, CORAL regularizes the internal feature space of the diffusion model itselfprecisely where representation entanglement arises.
Our contributions are summarized follows:
Empirical analysis of long-tailed diffusion behavior: We provide evidence that diffusion models trained on long-tailed data prone representation entanglement in the latent space of the denoising U-Net, particularly at the bottleneck layer, which contributes to low quallty tail-class generation.
Identification of representation entanglement as root cause: We show that the generation failure for tail classes stems from entanglement in the models latent feature representations induced by severe class imbalance, revealing previously unexplored failure mode.
We introduce COntrastive Regularization for Aligning Latents (CORAL), contrastive latent alignment method that encourages separation between class-wise latent represen- tations by augmenting the diffusion model with a supervised contrastive loss applied projected bottleneck features.
Proposal of CORAL:
Improved tail-class generation: Through extensive experiments on several long-tailed datasets (CIFARIO-LT, CIFAR (15), CclebA5 [16), (17) we demonstrate that CORAL significantly improves both the diversity and visual fidelity of tail-class samples, outperforming prior approaches. Moreover, we provide qualitative and quantitative evidence that CORAL promotes class-wise separation in the latent space of the denoising network, directly addressing the class entanglement that impairs tail-class generation.
1.1 Related Work
Diffusion Models for Imbalanced Data Standard methods for class-conditioned diffusion sampling include classifier guidance (CG) (19), which requires a separately trained classifier, and classifier-free free guidance (CFG) (18), which jointly trains conditional and unconditional denoisers. While widely

---

DDPM balanced
DDPM
CORAL
oerpate
actto
hard
deer
digg
hrea
slip
trank
Figure I: t-SNE visualizations of U-Net bottleneck features. The dataset visualized is CIFARIO-LT where the tail-to-head ratio is 0.01, i.e., the head class (airplane) is 100 times more represented than the tail class (truck), with an exponential decay in-between. Real CIFARIO-LT samples are passed through models trained under different settings. Shown are (left) DDPM [18] trained on the original balanced CIFAR- 10 dataset, (middle) DDPM trained on CIFARIO-LT with an imbalance ratio of 0.01, and (right) CORAL trained under the same imbalanced setting. In the balanced case, class representations are moderately separated, though some overlap remains. Under imbalance, DDPM exhibits substantial overlap between head and tail classes, effect refer representation entanglement, which degrades generation quality for tail classes. CORAL mitigates this effect by promoting class-wise separation in the latent space.
used, both CG and CFG struggle to generate diverse, high-quality samples for underrepresented tail classes 4
Several recent approaches have been proposed to improve tail-class performance in diffusion models, most of which operate in the ambient (image) space. For example, Class Class-Balancing Balancing Diffusion Models (CBDMs) [5] introduce a regularizer that penalizes deviations from a balanced class distribution, guiding the model to allocate more capacity to underrepresented classes during training. Time- dependent importance weighting (20] adjusts the loss based on sampling time to mitigate bias, while oriented calibration (6] uses Bayesian gating mechanisms transfer knowledge from head to tail classes (H2T) during unconditional generation and from tail to head (T2H) during conditional generation. DiftROP 4 applies a contrastive regularization based on KL divergence reduce class-conditional overlap at the output level.
In contrast to these ambient-space approaches, CORAL operates directly in the latent space of the diffusion model. Specifically, CORAL introduces a supervised contrastive loss on projected bottleneck features from the denoising U-Net, encouraging class-wise separation through metric learning. This latent-space regularization provides a more direct and structured means of disentangling class representations.
Relatedly, Han et al. [2 propose LDMLR, which generates synthetic latent features for long-tailed datasets using DDIM trained on encoder representations from a fixed model. While effective for long-tailed recognition, LDMLR operates as a post hoc feature augmentation method and does not modify the generative process. For the same objective of long-tail tailed recognition, Shao et al. [22] use a chosen classifier's s feature space to guide the diffusion model for the tail classes and filter out out-of-distriben dd-stribbution samples during generation. In contrast, CORAL directly regularizes the latent space during training, promoting class separation within the diffusion model without relying on a separate inference model.
Contrastive Learning in Variable Generative Models Contrastive learning (CL) is widely adopted technique for structuring embedding spaces in supervised, self-supervised, and metric learning settings 7 g) Recent work has extended CL to generative models: DiftROP 4 applies probabilistic contrastive loss in the ambient space to reduce overlap between class-condition output distributions in diffusion models; CONFORM 2 introduces contrastive regularization over attention maps to improve semantic alignment in text-to-image generation. TVAE [13] and Tri- VAE
[4] both incorporate a triplet loss into a variational autoencoder (VAE) framework: TVAE for general representation learning and TAi for anomaly detection. While TVAE uses standard VAE architecture, Tri-VAE employs U-Net with a projection head at the bottleneck, similar in spirit to CORAL. However, neither method involves diffusion and both modify the decoder path.
3

---

2 Preliminaries and Problem Setup
2.1 Diffusion Models
Generative DMs were first introduced in [1] which formulated data generation Markovian denoising process grounded non-equilibrium thermodynamics. The approach was later popularized by Ho et al. (2, who introduced simplified objective and fixed variance schedule, significantly improving sample quality and training stability.
DMs generate data by gradually adding noise to sample in a forward (noising) process and then learning to denoise in a reverse (denoising) process. The forward process gradually corrupts data sample Xy ~ ((xo) over T discrete time steps by adding Gaussian noise:
-0) VI- B, 3,1) and ((x1T(X0)
[((xxx11) 1SIST, Am
(1)
where {Be- 1), is a predefined variance schedule. As t increases, the distribution of X, transitions from close to (xx)) to approximately standard Gaussian. One can express the marginal
distribution at any timestep f as q(x, Xx) = NNxx;; Vā (1 - T.) where ã = -II.- - B.).
The reverse process is modeled by neural network that approximates the conditional distribution q(x,-1 (x,) using a learnable Gaussian distribution given by
po(X--1(xi) (x Eo(x,, t)) and Po(Xo:)) p(xr) II po po (
(2)
where the terminal distribution is typically set to p(x)) I)) and the variance is often fixed as (x(x,,) t) = I. Rather than predicting M4 directly, is common to reparameterize the mean in terms of the added noise cØ(x,, and train the network to predict this noise.
This leads the simplified training objective introduced by (2), in which the model c t) is trained to predict the Gaussian noise ~ NO0, I) that was used to perturb the clean input into noised version Xt, for a randomly sampled timestep ~ u( (1,...,)))
Lan No. [||e Eg (xt, 1)
In practice, the noise predictor es is implemented using an image- cosseett-immm -mage U-Net parameterized by 8, and the expectation operator is replaced by the empirical sample average for every t.
Classifier-Free Guidance When label information is available, classifier-fe freef guidance (CFG) 18 has become a widely adopted technique for improving conditional diffusion models. Instead of training separate classifier to guide generation, CFG modifies the denoising model in (3 to support both conditional and unconditional generation. During training, the model 6 (xt, 4, is optimized using class labels y; for chosen fraction Puncond of samples, the training process ignores labels to learn the unconditional model with y = Ø
Finally, at sampling time, conditional guidance is applied by combining the conditional and uncondi- tional predictions to recover Xo from X,:
CFG (x,, t,y) = (1 (x,, weg (xx, t),
where w > 0 is guidance weight controlling the strength of conditioning.
2.2 Metric Learning and Contrastive Approaches
Metric learning aims to map inputs into an embedding space where semantically similar examples are close together and dissimilar ones are apart. Instead designing distance functions manually, modern methods use neural networks to learn transformations that make standard distances (e.g., Fuclidean or cosine) meaningful for the task. This learned embedding captures complex similarity structures aligned with supervision.

---

Embedding space
Caitf
(d)
SCORAL
Ub)e
x
COCCRAL loss
linear layer
normalization
Lron
(i) Training data
(1 Samples
(c) Projection head fo
Figure 2: CORAL architecture and workflow on CelebA-5, (a) The five-class CelebA-5 training
data is input the U-Net architecture. (b) Denoising U-Net. The white inset shows an actual t-SNE visualization of the U-Net latent representations due to CORAL- (c) CORAL's addition to the standard DDPM architecture: a projection head MLP consisting of single dense layer followed by normalization. (d) The output from the U-Net and the projection head are used to compute the corresponding diffusion and contrastive losses. (e) The contrastive loss is scaled by time-dependent weighting function, X(t), and added to the standard diffusion loss to obtain the CORAL loss. (f) Samples are obtained from trained CORAL model.
2.2.1 Contrastive Loss Functions
Contrastive loss functions are a fundamental tool in metric learning, designed to shape embedding spaces so that semantically similar samples are close together, while dissimilar samples pushed apart. Early formulations, such as the triplet loss [24], enforce a margin between anchor-positive and anchor-negative pairs using triplets of labeled samples where each triplet contains anchor, positive (same class as anchor), and a negative (different class from anchor). While effective, triplet loss can suffer from slow convergence and inefficient sampling. More recent advancements such as the supervised contrastive loss (SupCon) [25] generalize this idea by leveraging all positives and negatives in mini-batch, offering greater stability and improved sample efficiency during training.
Supervised Contrastive Loss SupCon [25 generalizes triplet loss by comparing each anchor to multiple positives and negatives within batch, improving both convergence stability and overall performance. Let E gd denote the ly-normalized embedding sample. The loss defined as:
exp(z; 2p/ TsC) EESS() exp(z; 2. 2./TSC)
LSupCon = -
log
(5)
P())
REI
PEP(1)
where / is the set of all indices in the batch, S(i) = IN {i} denotes the set of all sample indices the batch excluding the anchor i, P(i) C S(i) the set of indices corresponding to positive samples that share the same class the anchor, and TsC is temperature parameter that controls the concentration (sharpness) of the similarity distribution. Lower values of Tse (e.g., TsC se 0.1) sharpen the distribution, placing greater emphasis on harder positive and negative pairs and increasing the gradient magnitude ( V Lscl x 1/Tsc).
Our Method
In this section, we present our method, COntrastive Regularization for Aligning Latents (CORAL), designed to enhance class separation in diffusion models trained on long-tailed datasets. The core insight behind CORAL is that the latent space of the denoising U-Net, specifically its bottleneck layer, plays a central role in shaping generative behavior. In long-tailed settings, we observe that latent representations of tail-class samples often overlap with those of head classes, resulting in representation entanglement and degraded generation quality (see Figure T for a visualization the CIFARIO-LT dataset). Our comparisons between models trained on balanced and imbalanced data indicate that this overlap arises from head classes dominating parameter updates, resulting in less structured latent representations for tail classes.
CORAL introduces two targeted modifications to standard diffusion training: a lightweight projection head applied to the U-Net bottleneck and supervised contrastive loss term. These additions allow
5

---

CORAL to regularize the latent space directly, promoting intraclass clustering and interclass separa- tion during training. The contrastive signal complements the diffusion objective, helping maintain semantic distinctions across classes, especially for underrepresented ones. Figure 2 illustrates the CORAL framework using CelebA-5 [6]. a 5-class L.T' subset of CelebA 26 as an example dataset.
Architectural Modification While the forward diffusion process operates in the high-dimensional ambient space, the U-Net processes information through a compressed latent space, with the bottle- neck layer playing a central role in the model's representational capacity. Prior work has shown that this bottleneck encodes semantically meaningful features [T)) making it a natural point for inter- vention. CORAL leverages this insight by adding a small projection head fø.e.g. a fully-connected linear layer followed by normalization layer, to the bottleneck output.
CORAL builds on established principles of contrastive representation learning, where projection heads have been shown to capture task-related information more effectively than direct feature space constraints [12 27]. The projection head serves two critical functions: (1) decouples the contrastive objective from the main diffusion features, allowing the model to learn class-discriminative embeddings an auxiliary space while the bottleneck continues to serve the generative objective [12], and (2) it preserves intraclass diversity by preventing the contrastive loss from directly collapsing intraclass variations the bottleneck representations (27). This complementary structure allows both the diffusion and contrastive objectives to improve simultaneouslly.
During training, we apply a supervised contrastive loss on the projected embeddings to encourage class-wise separation while the main bottleneck features continue to serve the diffusion objective. Once trained, the U-Net bottleneck has learned to produce well-separated class-specific features, and the projection head is no longer needed. This design ensures that CORAL adds zero computational overhead during sampling and remains fully compatible with standard diffusion sampling procedures. The learned structure in the bottleneck representations persists and guides generation without requiring the auxiliary projection network.
Table 4 in Appendix B demonstrates that CORAL's latent space intervention consistently outperforms ambient space approaches the style of DiffROP Ambient-space approaches enforce separation on already generated outputs, which can reduce intraclass diversity by imposing constraints post-hoc. In contrast, CORAL learns inherently separated representations at the compressed bottleneck layer where class overlap occurs during the generative process itself.
Training Objective The overall training objective for CORAL augments the standard diffusion loss with contrastive alignment term applied to the projected latent representations to obtain
LCORAL Laift + A(t) Leon:
(6)
where Lai is the standard diffusion training loss, such as the noise prediction objective defined in B Ccon is contrastive loss applied to the projected bottleneck features, and X(t) is time-dependent weighting function. While we use Loon = LSapCon in our experiments, the framework is general and supports any contrastive loss. The weighting function X(t) is defined as:
(1-1/T) E (0.1...., T) Ty
X(t) w. exp
where w is the base contrastive weight, T is the total number of diffusion steps, and T, is the temperature parameter that controls the decay rate. Although in general T, > O, our results with the SupCon loss suggest that a range between [0.5, 1.0) works best.
This dynamic weighting scheme places greater emphasis on the contrastive objective during the earlier (less noisy, R 0) denoising steps, where a meaningful semantic structure more recoverable, and gradually reduces its influence at later steps (t x T), where noise dominates the input. This encourages more discriminative latent representations during the most informative stages of training.
Training Procedure To train our proposed CORAL method, we modify the standard diffusion training procedure to incorporate both contrastive latent regularization and classifier-free guidance. Algorithm [ summarizess the full training pipeline. For each mini-batch, we first sample diffusion timesteps and generate noisy inputs via the standard DDPM forward process. In line with CFG training protocol [18), we randomly drop class labels with a fixed probability to enable joint training

---

of conditional and unconditional denoising. The noisy inputs and (possibly masked) labels are passed through the U-Net to compute the standard diffusion loss. Simultaneously, we extract bottleneck features from the U-Net encoder, project them via the projection head fo, and compute supervised contrastive loss using the original (unmasked) class labels. We then compute the total loss LCORAL. given in 66 Model parameters are updated using backpropagation on LCORAL-
Algorithm CORAL Training Procedure
Input: Dataset D, model cØ, projection head fa. total diffusion steps T, guidance dropout probability Puncond- contrastive weight schedule X(t)
Initialize: Parameters 0, Φ
for each mini-batch of size do
for each sample (xg ()) ) in mini-batch do
Sample timestep t ~ u(f1.....
Sample noise ( (0,
Compute noised inputs: x; (6) Vārxo + VI c
Drop labels with probability Puncond: 7(0) = Ø w.p. Puncond
Predict noise: ( = co (x) )())
Extract bottleneck features h of x from U-Net encoder
Compute projected embeddings: z = (b (0)
end for
Compute diffusion loss: Lait 1o Ein E 10
Compute contrastive loss Loon using {(z) (6) )) (e.g., using for SupCon)
Compute total loss: LCORAL. Lai + A(t) Leon
Update (0, Φ) using gradients of CCORAL
end for
4 Experimental Setup and Results
4.1 Experimental Setup
Datasets We evaluate CORAL on long-tailed (LT) datasets: CIFARIO-LT, CIFARIO0-LT (28), CelebA-5 (66) and ImageNet-LT [17] CIFARI0/100 datasets contain 32 x 32 color images broken into 10 and 100 classes. For CIFARIO-LT and CIFARIOO-LT, we simulate long-tailed distributions by applying an exponential decay to the class frequencies, controlled by an imbalance factor P (0.01, 0.001 This results in the most frequent (head) class appearing 1/p times more often than the rarest (tail) class, with intermediate classes following an exponentially decreasing trend. CelebA-5 consists of 64 x 64 resized face images in 5 classes corresponding to hair color. CelebA-5 is naturally imbalanced. We construct ImageNet-LT by sampling a subset of ImageNet-2012 following the Pareto distribution with power value a 6. ImageNet-LT has 1,000 clases with class sizes ranging from 5 to 1, 1,280 images, we resize images to 64 x resolution. Additional details can be found in Appendix A
Implementation Our implementation builds on the codebase from [5], with modifications to support contrastive latent regularization. We use a U-Net backbone with multi-resolution attention and dropout, consistent across all experiments. We use the SupConLoss implementation from [29]. Training was run on NVIDIA A 100 (80 GB SXM) and H100 GPUs. Key training and architectural hyperparameters are summarized in Appendix
Evaluation Metrics We compute the standard FID and IS [1 to capture both quality and diversity of the generated images. We additionally compute recall for distributions (PRD) and improved recall [32] (labeled as Recall in Table I). Standard PRD uses k-means clustering on Inception features with 2000 clusters (20 times the number of classes for CIFAR-100) to compute Fs and F1/8 scores [3). F& emphasizes recall (diversity) and F1,s emphasizes precision (quality). The improved PRD metrics employ k-nearest neighbor manifold estimation (k 3) on VGG16
7

---

Table I: Comparison of methods on long-tailed image generation benchmarks.
Dataset
Method
FID (L)
IS (f)
Fs (f)
Recall (f)
F1/8 (1)
DDPM 2
6.17
9.43
0.87
0.52
0.94
CIFARIO-LT
CBDM 5
5.62
9.28
0.96
0.57
0.95
(p= 0.01)
T2H 6
7.01
9.63
0.89
0.54
0.95
32 32
CORAL (ours)
5.32
9.69
0.97
0.59
0.97
DDPM 2
13.05
9.10
0.87
0.53
0.85
CIFARIO-LT
CBDM 5
12.74
9.05
0.87
0.56
0.89
(p= 0.001)
T2H 6
12.80
8.97
0.87
0.55
0.88
32 x 32
CORAL (ours)
11.03
9.13
0.90
0.56
0.89
DDPM 2
7.70
13.20
0.87
0.50
0.89
CIFARIO0-LT
CBDM
6.02
12.92
0.91
0,56
0.90
(p 0.01)
T2H 6
6.78
12.97
0.88
0.54
0.89
32 x 32
CORAL (ours)
5.37
13.53
0.92
0.59
0.91
DDPM 2
10.28
2.90
0.90
0.52
0.89
CelebA-5 64 x 64
CBDM 5
8.74
2.74
0.92
0.57
0.90
T2H 6
9.50
2.63
0.89
0.53
0.87
CORAL (ours)
8.12
2.97
0.94
0.59
0.92
Table 2: Comparison of methods on ImageNet-LT.
Dataset
Method
FID (L)
IS (f)
Recall (f)
DDPM 2
17.08
21.03
0.39
ImageNet-LT 64 x 64
CBDM [5]
22.66
17.13
0.42
T2H 6
18.59
19.15
0.44
CORAL (ours)
16.11
24.17
0.48
features [34]. providing more robust estimates of sample quality and coverage. These metrics collectively provide a comprehensive assessment; in particular, Fi/s measures generation fidelity while improved recall and Fs capture the diversity of the generated distribution. FID captures mixture of both quality and diversity.
For overall metric calculations, use the balanced version of the datasets for the real data to ensure fair evaluation. All metrics are computed on 50,000 generated samples to ensure statistical reliability. During sampling, class labels are drawn from uniform distribution across all classes for equal representation.
Baselines We compare CORAL's performance against that of DDPM. CBDM, and T2H for the following datasets: CIFARIO-LT with = 0.01 and = 0.001, CIFAR 100-LT, CelebA-5 and ImageNet-LT. We use the publicly available implementations for DDPM (2). CBDM [5) and T2H ( to train the models with provided parameters, where available, generate synthetic samples, and report results for each of these methods.
4.2 Experimental Results
Comparison Metrics In Table and Table we compare the performance CORAL against standard DDPM, as well as state-of-the-art baselines CBDM and T2H across multiple long-tailed datasets. CORAL consistently outperforms all baselines across all datasets and evaluation metrics, demonstrating its effectiveness in improving both the quality and diversity of generated samples. CORAL demonstrates strongest improvements on metrics that capture diversity and distribution coverage while maintaining improved performance quality focused metrics
the large scale ImageNet-LT benchmark with 1,000 classes, CORAL's advantages are most evident, outperforming all basclines by significant margins. Ambient space regularization approaches

---

exhibit degraded performance scenarios with large numbers of classes, whereas CORAL's latent space intervention maintains consistent performance.
Per-Class FID Figure presents the per-class FID scores for CIFARIO-LT with p 0.001, representing a more extreme class imbalance. CORAL consistently outperforms baseline methods across nearly all classes. The gains are particularly notable for the tail classes. Whereas both CBDM and T2H exhibit degraded performance on tail classes, CORAL maintains stable performance across both head and tail classes. For per-class FID analysis in Figure 3 generate 5K samples for each class and compare against the 5K real samples from the balanced dataset for that specific class.
70
DUFM
CHIDA
1221
56
CORAL
42
28-
14-
airplane
auto
bird
cal
deer
dog
Irog
horse
ship
truck
Figure 3: Per-class FID (1) for the CIFAR10-LT dataset with an imbalance factor = 0.001
Latent Space Visualizations Our experimental results clearly show that CORAL achieves better performance by explicitly enforcing class-wise separation in the latent space. In Figure 1 we present t-SNE visualizations of U-Net bottleneck representations for CIFAR I0-LT with an imbalance ratio of 0.01, comparing models trained with DDPM (on both the original balanced CIFAR-I0 and CIFARIO-LT) and with CORAL on CIFAR IO-LT. Figure 2 visualizes the separated representations (using t-SNE) learned by CORAL for the CelebA-5 dataset. Additional visualizations for the other datasets using both t-SNE and UMAP [35 are included in Appendix D particular, our plots for a balanced dataset with a limited number of samples per class show that the observed representation entanglement arises predominantly from class imbalance in the training distribution.
Generation Quality Figure 4 presentss generated samples from CBDM, T2H, and CORAL for the tulips class (class 92) in CIFAR100-LTT. Visually, CORAL produces samples that are both more diverse and of higher fidelity compared to the other methods. These qualitative differences align with the quantitative improvements observed Table I where CORAL achieves superior performance across all evaluated metrics. CBDM suffers from mode collapse by producing smaller flowers with excessive grass backgrounds borrowed from head animal classes. T2H shows diminished class fidelity tulips resemble other flower types due over-transfer from head tail classes. In contrast, CORAL generates tulips that reflect appropriate scale and structure, with distinctive features and backgrounds consistent with the training data. This demonstrates CORALs ability to balance the trade-off between preserving tail-class characteristics and promoting sample diversity. Additional visualizations are provided in Appendix
We have performed extensive ablation studies for various hyperparameters. including the SupCon temperature Tsc, the time-dependent weighting function temperature T,, and the CFG sampling parameter w; these plots can be found in Appendix C
Ablation Studies
Concluding Remarks
Broader Impacts As generative models have become more widely utilized in practice, their representativennsss becomes more impactful. Tail class generation has become a key method to address long-tailed recognition tasks such disease detection where real data is limited. While

---

Figure 4: Comparison of generated samples from the class tulips (class 92) in CIFARIO0-LT, 0 0.01. CBDM (left), T2H (middle), and CORAL (right). CORAL shows increased diversity and fidelity relative to existing approaches.
generated images have the potential to cause harm, e.g. deepfakes or bias amplification, CORAL helps to mitigate the bias introduced by dataset imbalance.
Future Directions While our work focuses on class conditional generation with categorical imbal- ance, the core principle of latent space disentanglement through contrastive regularization extends naturally to more complex generative settings. A particularly promising application domain the fine tuning of large scale text-to-image (T2I) models for specialized tasks. When pretrained models like Stable Diffusion [36] are adapted to domain specific applications, such as medical imaging, scientific visualization, specialized industrial use cases, the fine tuning datasets often exhibit severe imbalance. For T21 models that employ U-Net architectures, CORAL's approach of applying con- trastive regularization these bottleneck layers could prevent rare concepts from becoming entangled with common ones. Exploring CORAL's applicability to parameter efficient fine tuning methods, e.g. LoRA 37 for domain adaptation and investigating whether similar contrastive interventions can improve generation quality for underrepresented concepts constitute promising directions for ensuring equitable representation across specialized domains.
While CORAL is able to produce diverse and high quality images when trained on heavily imbalanced datasets, its power comes at the cost of additional computational complexity. This limitation is shared by all comparable methods, though it can be reduced by finetuning with the CORAL loss rather than fully training.
Limitations
Conclusions Ensuring high-quality sample generation for tail classes of long-tailed datasets remains major challenge. In addressing this challenge, we have revealed previously unknown cause for the poor performance of DMs: the (U-Net) latent representations for the tail classes completely overlap with those for the head classes, thereby severely limiting the guidance of the former. Our method, CORAL, significantly enhances both the diversity and fidelity of diffusion model outputs relative to the state-of- the-art arrt by separating and realigning the latent space representations, especially for the long-tail classes using contrastive losses. We have demonstrated that CORAL performs well for datasets with both extreme imbalance and many classes, and our results suggest that disentangling in the latent space is more effective than rebalancing and increased guidance in the ambient space.
Acknowledgments and Disclosure of Funding
This work is supported part by NSF grants CIF-181525 181561, CIF-2007688, DMS-2133 2134256, and SCH-2205080.
10

---

References
[1] Jascha Sohl-Dickstein, Fric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 2256-2265. PMLR, 2015.
[2] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Proceedings of the 34th International Conference on Neural Information Processing Systems, 2020.
[3] Yang, He Jiang, Qing Song, and Jun Guo. A survey on Journal of Computer Vision, 130(7): 1877)::887788880 1872, 2022.
visual recognition. International
[4] Divin Yan, Lu Qi. Vincent Tao Hu, Ming Hsuan Yang, and Meng Tang. Training class- imbalaned diffusion model via overlap optimization. CoRR, ahs/2402. 10822, 2024. httphtttt///do.orrr org/10 18550/arXi1 2402. 10821,
[5] Viming Qin, Huangjic Zheng, Jiangchao Yao, Mingyuan Zhou, and Ya Zhang. Class-balancing diffusion models. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition, 2023.
[6] Tianjiao Zhang, Huangjie Zheng, Jiangchao Yao, Xiangfeng Wang, Mingyuan Zhou, Ya Zhang, and Yanfeng Wang. Long-tailed diffusion models with oriented calibration. In The Twelfth International Conference on Learning Representations, 2024.
[7] Phuc H. Le- Khac, Graham Healy, and Alan F. Smeaton. Contrastive representation learning: A framework and review. IEEE Access, 1173077-133334 2020, doi: 10.1109/ACCESS 20000000
[8] Deen Dayal Mohan, Bhavin Jawade, Srirangaraj Setlur, and Venu Govindaraj. Deep metric learning for computer vision: A brief overview, 2023, URI /ttp///iww.oog/////3/2.2../ 10046
19] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and computer-assisted intervention-MICCA. 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pages 234-241.. Springer, 2015.
[10] Mingi Kwon, Jaescok Jeong, and Youngjung Uh. Diffusion models already have a semantic latent space. In The Eleventh International Conference Learning Representations, 2023.
[11] Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-sn.. Journal of Machine Learning Research, 2866)2257- 2008. URI http:h/tt//mml org/ http/aaprs//a/ / waa/vannermmateene... html
[12] Ting Chen, Simon Korbllith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Proceedings of the 37th International Conference on Machine Learning, ICML'20 20. JMLR.org, 2020.
[13] Haque Ishfaq, Assaf Hoogi. and Daniel Rubin. Tvae: Triplet-based variational autoencoder using metric learning. arXiv preprint arxiv: 000.04403, 2018.
[14] Hansen Wijanarko, Evelyne Calista, Li-Fen Chen, and Long-Sheng Chen. Tri-vae: autoencoder for unsupervised anomaly detection in brain tumor mri. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops, pages 3930-3939, 2024.
Triplet variational
[15] Kaidi Cao, Colin Wei, Adrien Gaidon, Nuno Arechiga, and Tengyu Ma. Leaming imbalanced datasets with aahel-distribution-awaaa margin loss. In NeurlPS, 2019.
[16] Jachyung Kim, Jongheon Jeong, and Jinwoo Shin. M2m: Imbalanced mmbblanccecclassiffatttio classificattin via major-to-minor translation. CoRR, abs/2004.00431, 2020. URL UtLp::tt::///astt http:///aaawvv...aaaaaa/// ab0rg///b///00/..////// 00431
[17] Ziwei Liu, Zhongqi Miao, Xiaohang Zhan, Jiayun Wang. Boqing Gong, and Stella X Yu. Large-scale long-tailed recognition in open world. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 2537-2546, 2019.
[18] Jonathan Ho and Tim Salimans. Classifier free diffusion guidance. In NeurlPS 2021 Workshop on Deep Generative Models and Downstream Applications, 2021. URL http:////oppeeeeeeeeeeeeeeeeeeee
YbL
[19] Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat GANs image synthesis. In
A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan editors, Advances in Neural Information Processing Systems, 2021. URL
met/forum' /itt//rrrrum/iaAAAAAaaaaV
[20] Yeongmin Kim, Byeonghu Na, Minsang Park, JoonHo Jang, Dongjun Kim, Wanmo Kang, and chul Moon. Training unbiased diffusion models from biased dataset. In The Twelfth International Conference on Learning Representations, 2024. URL http:s::tts///r/tee
net /et/ forum 1id1d-9999PPP Bee
[21] Pengxiao Han, Changkun Ye, Jieming Zhou, Jing Zhang. Jie Hong, and Xuesong Li. Latent-based diffussion model for tailed recognition. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages
2024.
11

---

[22] Jie Shao, Ke Zhu, Hanxiao Zhang, and Jianxin Wu. DiffuLT: Diffusion for long-tail recognition without external knowledge. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[23] Tuna Han Salih Meral, Enis Simsar, Federico Tombari, and Pinar Yanardag. Conform: Contrast is all you need for
text-to-image diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). pages 9005 9014, 2024.
[24] Florian Schroff, Dmitry Kalcnichenko, and James Philbin. Facenet: A unified embedding for face recognition and clustering. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2015.
[25] Pramnay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishnan. Supervised contrastive learning. In Advances in Neural Information Processing Systems, volume 33, pages 18661-18673, 2020.
[26] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xianou Tang. Deep learning face attributes in the wild. In 2015 IEEE International Conference Computer Vision (ICCV), pages 3730--738, 2015. doi:
10.1 109/1CCV22015.425
[27] Yihao Xue, Eric Gan, Jiayi Ni, Siddharth Joshi, and Baharan Mirzasoleiman, Investigating the benefits of projection head for representation learning, 2024.
11391
[28] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report. University of Toronto, 2009. URL (ttL//ttss:// //www. cs, toronto. edu/ htt/:/rrrieeaannng---arra feattresstttreteeeee pdf
[29] Kevin Musgrave, Serge J. Belongic, and Ser-Nam Lim. Pytorch metrie learning. ArXiv, abs/2008.091644 2020.
[30] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge to local nash equilibrium. NIPS 17, page 66296640, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860964.
[31] Tim Salimans, Ian Goodfellow, Wajeiech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. In Proceedings of the 30th International Conference on Neural Information Processing Systems, page 22342242, 2016.
[32] Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metrie for assessing generative models. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32, 2019.
[33] Mehdi S. M. Sajjadi, Olivier Bachem. Mario Lucic, Olivier Bousquet, and Sylvain Gelly. Assessing generative models via precision and recall. In S. Bengio. H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Desaa-Bianchi, and Gamett, editors, Advances Neural Information Processing Systems, volume 31, 2018.
[34] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large scale image recogni- tion. CoRR, 02//1409.15555 2014. jttpps:t swwanticccooollrr..c.mm... wwwg/oorpgsliaae 14124313
[35] Leland McInnes, John Hcaly, Nathaniel Saul, and Lukas GroSSberger. Umap: Uniform manifold approx- imation and projection. Journal of Open Source Software, 3(29):861, 2018. doi: 10.21 05/josss000861
URL https::t::///do htt.///oo.o..0////// org/10. 10. 20105/ joss. 00861
[36] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffuusio models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages
2022.
[37] Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen- Zhu, Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022.
[38] Yingfan Wang, Yiyang Sun, Haiyang Huang, and Cynthia Rudin. Dimension reduction with locally adjusted graphs. Proceedings of the AAAI Conference on Artificial Intelligence, 39(20))21357-211633 Apr. 2025, doi: 10.1609/aaii v39.20.3533330 35436. httpsttss:// //ojs. aaai. aaail.or//inddee http://ppphaaaaaaarrrraa////////
12

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We propose CORAL which improves tail class generation by separating classes in the latent space.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Yes, limitations are included in the conclusion Section 5
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairnsss.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
13

---

Justification:
No theoretical results are included
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Yes, experimental details can be found in the experimental setup Section and in the appendix material, and hyperparameters are listed in the appendix material.
Guidelines:
The answer NA means that the paper does include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
14

---

Answer: [Yes]
Justification: The code is available at the GitHub repository linked in the abstract and will be released publicly after review period.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
(olscnnnoocyy for more details.
While encourage the release of code and data, understand that this might not be possible, so No is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips.
issmiissoonnn asionnolcccciiiic for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Included in experimental setup Section 4.1 and in the appendix. Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Error bars are not included due to resource limitations.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
15

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Discussion on computational resources is in the experimental setup Section 4.1 and in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
No ethical concerns.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Yes, broader impacts are included in the concluding remarks Section 5
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
16

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: Models presented are trained on publicly available data and have limited potential for misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: License respects the licenses of cited works.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the datasets..papersstt pwpprswipttttt. thoode. com/ datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
17

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: No new assets released.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
18

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: LLMs are not used in non-standard or original ways.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy battps://neurli /ttppp////.//oeeeeeeee///////// for what should or should not be described.
19

---

A Additional Dataset and Experimental Setup Details
Long-Tail Tail Datasets Balanced datasets can be artificially transformed into long-tail (LT) datasets by assigning class sample counts according a geometric progression governed by an imbalance ratio In this formulation, the head class has the maximum number of samples, N, while the tail class has approximately P.V samples. The number of samples in class is given by:
N
(8)
where N is the number of samples in the head class, P is the imbalance ratio (0 < < 1),
E {0, 10.,,...,C C C--1-11 1] is the class index, and C' is the total number of classes.
CIFARIO-LT and CIFARI00-LT The original CIFARIO and CIFARI00 datasets each consist of training set with 50k images uniformly distributed across 10 or 100 classes, respectively. Their long-tailed variants, CIFAR 10-LT and CIFAR100-13 [28]. introduce an exponential decay in class frequency from class 0 to the final class. Common long-tail imbalance ratios include p 0.01 and 0.001. Specifically for p 0.01, CIFAR10-LT contains 12,406 images, with the first head class comprising 5,000 samples and the last tail class only 50. CIFARIOO-LT has 10,847 images, with the head class containing 500 samples and the tail just 5.
Experiments on CIFAR IO-LT and CIFAR 100-LT were conducted using NVIDIA A100 80GB SXM GPUs. Training took approximately hours, and sampling required 8 hours. For DDPM, CBDM, and CORAL, the hyperparameters used were: a learning rate of 2 x 10 batch size of 128, Adam optimizer with default momentum parameters, dropout rate of 0.1, 150k training steps, and = 1000 diffusion steps. For T2H, all settings remained the same except for the number of training steps, which was increased to 200k.
CelebA-5 CelebA-5 [16] five-class subset the CclebA dataset, composed of samples labeled with exactly one of the following hair colors: black, brown, blonde, gray, or bald. Samples with multiple or missing labels are excluded. The dataset is naturally imbalanced, with black- and brown-haired individuals significantly outnumbering those with gray hair or baldness.
on NVIDIA H100 GPUs Training took approximately 18 hours, and sampling required 22 hours. All models were trained with a learning rate of 3 x 10-4 and a batch size of 128, with all remaining hyperparameters kept consistent with the CIFAR experiments.
Experiments CelebA-5
ImageNet-LT ImageNet-LT is long-tailed variant of ImageNet-2012 constructed by 11 by sampling a subset following the Pareto distribution with power value a The dataset comprises 115.8k images from 1000 categories, with a maximum of 1,280 images per class and a minimum of 5 images per class.
Experiments on ImageNet-LT were conducted using NVIDIA H100 GPUs. All models were trained with a batch size of 128, 300k training steps, and T 1000 diffusion steps. For evaluation, generated 50k samples uniformly across all classes and compared against the balanced validation set containing 20k images the real samples.
Hyperparameters Table summarizes the regularization hyperparameters and sampling guidance scale used for each method and dataset. The sub-tables correspond to DDPM (top left), CBDM (top right), T2H (bottom left), and CORAL (bottom right), respectively. For CORAL, the base contrastive weight, w, in 7 was set to 0.01.
We follow the code implementation for CBDM 5 with the regularization weight Teb = r/T,, where is the original weight defined in [.. and T' is the total number of diffusion timesteps. For T2H 6). we use the code implementation with direct distance-based weighting that does not require regularization parameter.
20

---

Table 3: Hyperparameter settings for each method.
DDPM 2
CBDM 5
Dataset
Dataset
Teb
CIFARIO-LT (p 0.01)
0.8
CIFARIO-LT (o 0.01)
1.0
1.0
CIFARIO-LT (p = 0.001)
1.0
CIFARIO-LT (p = 0.001)
1.8
1.0
CIFAR100-LJ 0.01)
0.8
CIFARIO0-LT (p- 0.01)
1.6
1.0
CelebA-5
0.6
CelebA-5
1.0
50.0
T2H [6
CORAL (ours)
Dataset
Dataset
w
TSC
CIFARIO-LT (p- 0.01)
1.0
CIFARIO-LT (p 0.01)
0.6
0.12
0.8
CIFARIO-LT (p = 0.001)
1.7
CIFARIO-LT (p 0.001)
1.0
0.10
1.0
CIFARI00-LT (p- 0.01)
1.5
CIFAR100-LE (p 0.01)
0.8
0.09
1.0
CelebA-5
1.0
CelebA-5
0.7
0.12
0.8
Additional Results
Comparing Contrastive Regularization in Ambient and Latent Space CORAL differs funda- mentally from ambient space contrastive regularization methods (similar to DifROO 4) in both its approach and effectiveness. CORAL operates directly within the diffusion model's internal latent space, specifically at the U-Net bottleneck layer augmented with a projection head, where semantic representations are formed and disermmicrie embeddings are learned. This architectural differ- ence is crucial because CORAL addresses the representation entanglement where tail-class samples overlap heavily with head-class representations. contrast, ambient space contrastive regularization methods enforce separation constraints on the image space.
Dataset
Method
FID (L)
IS (t)
Recall (t)
DDPM 2
6.17
9.43
0.52
CIFARIO-LT
Ambient Space Contrastive
5.85
9.18
0.55
0.01)
CORAL (ours)
5.32
9.69
0.59
DDPM 2
17.08
21.03
0.39
ImageNet-LT x 64
Ambient Space Contrastive
24.73
15.12
0.34
CORAL (ours)
16.11
24.17
0.48
Table 4: Comparison of contrastive regularization strategies CIFARIO-LT and ImageNet-LT.
Table 4 provides empirical validation of CORAL's design choices. The results show that learned separation in the latent space scales more effectively than imposed separation in the ambient space. As the generation task becomes more challenging, ambient space methods not only struggle but actually degrade below baseline DDPM performance in distributional coverage as the class space grows, whereas CORAL maintains its effectiveness across all metrics.
Generated Images We generate images for CIFARIO-LT, CIFAR I00-LT, and CelebA-5 to evaluate the effectiveness of CORAL in addressing the challenges of diffusion models trained on long-tailed datasets. Randomly selected examples are shown in Figures 5 to 7 illustrating how our contrastive latent alignment framework improves both the quality and diversity of generated samples, particularly for tail classes.
Figures 5 and 6 show generated samples for CIFAR 1O-LT and CIFAR 100-T,, respectively, with = 0.01. CORAL successfully disentangles latent representations to generate high-quality, diverse samples for the underrepresented classes. CORAL preserves the distinctive characteristies of each
21

---

Figure 5: Generated samples produced CORAL the CIFARIO-LT dataset with p 0.01.
all
B M
Figure 6: Generated samples produced by CORAL the CIFARI00-LT dataset with 0.01.
22

---

Figure 7: Generated samples produced by CORAL on the CelebA-5 dataset.
class, effectively mitigating feature borrowing from head to tail classes. The visual quality of these results highlights the effectiveness of CORALs latent space regularization in promoting class separation and maintaining clean, well-structured features in the generated outputs.
Figure displays generated samples for CelobA-5, demonstrating CORALs ability to handle naturally imbalanced data. The dataset exhibits pronounced class imbalance across five hair color categories (black, brown, blond, gray, and bald) with the head class containing nearly 15 times more samples than the tail class. In such imbalanced settings, latent representations for tail classes often become entangled with those of head classes. CORAL effectively preserves class-specific features, producing diverse and realistic images in all categories.
C Ablation Studies
Effects of Hyperparameters Figure 8 illustrates the effect of three key hyperparameters
CORALS performance for CIFARIO-LT with imbalance ratio = 0.01, measured by FID. The supervised contrastive temperature, TsC. in (5) achieves optimal performance TSC 0.12, beyond which FID increases sharply.
For the decay rate temperature, in 7 the best performance is observed T, 0.8. FID remains relatively stable for 0.7 < t, < 0.9, with a steep increase outside of this range. These findings support our hypothesis that contrastive regularization is most effective when applied toward the end of the denoising process (i.e., when ~ 0).
Finally, for the CFG scale, w, in FID traces a convex curve with optimal performance at = 0.6. This indicates a trade-off between leveraging class-conditional information (w > 0) and avoiding over- conditioning that could limit sample diversity (w >> 0.6). These results underscore the importance of careful hyperparameter tuning in CORAL achieve an optimal balance between sample fidelity and diversity, particularly for tailed datasets.
23

---

6.50
6.20
11.0
6.30
6.05
10.2
9.4
6.10
5.90
5.90
5.75
5.70-
5.60
5.50
5.45
6.2-
5.30
5.30
6.4
0.08
0.1
0.4
0.8
1.2
1.6
2.0
Tsc
Tr
W
Figure 8: Effect of regularization hyprpaarameters and guidance scales FID. From left to right:
FID vs TSC. FID vs Ty, and FID vs w for the CIFARIO-LT dataset. Coral stars mark the lowest FID achieved for each hyperparameter.
CORAL Balanced Datasets We also evaluate CORAL two balanced datasets, CIFARIO and CIFAR 100, using FID as the performance metric, as shown in Table Even in the absence of class imbalance, CORAL outperforms DDPM and CBDM in terms of FID. This improvement stems from CORALs contrastive loss, which promotes class-wise separation in the latent space even in balanced settings, illustrated in Figure Importantly, this separation is achieved without compromising fidelity or diversity, as reflected in the consistently strong FID scores.
Table 5: Comparison of methods on CIFARIO and CIFARIO 100 image generation.
Dataset
Method
FID (L)
DDPM
3.84
CIFAR10
CBDM 5
3.61
CORAL (ours)
3.30
DDPM 2
3.91
CIFAR100
CBDM
3.37
CORAL (ours)
2.86
Impact of Sample Size Balanced Datasets Figure illustrates the impact of total sample size on latent representations from DDPM for the balanced CIFAR10 dataset. As the number of training samples per class decreases from 5k 100, we observe a noticeable reduction cluster formation in the latent space, with representations becoming increasingly scattered. This suggests that in such highly overparameterized regime, the model memorizes individual samples rather than learning generalizable class-level structure. We note that this type of scattering is visually distinct from the entanglement observed in DDPM under class imbalance (see, for example, Figure 12). For imbalanced datasets, models have more available information on head than tail classes, allowing them to learn better representations for the former than the latter. This does not lead to scattering, but rather to distinct overlap of tail class representation clusters within head class clusters.
Figure II illustrates the memorization behavior of DDPM on the balanced CIFAR1O dataset with 50 samples per class and guidance strength w = 0.1. For each class, the panel displays two rows of 10 images: the bottom row (outlined magenta) shows samples generated by DDPM, while the top row presents the most visually similar real training samples corresponding to each generated image. The generated samples are near-identical to the training samples, differing most by horizontal reflection, highlighting the extent of memorization in this limited-data setting.
24

---

Class
nirpaaee
frog
saip
tzuck
t-SNE
UMAP
LocalMAP
Figure 9: Visualizations of U-Net bottleneck features using I-SNE, UMAP, and LocalMAP on CIFAR-10 (balanced). Each row shows different method trained on CIFAR-10, listed from top to bottom: DDPM 18. CBDM 5 and CORAL (ours).
25

---

Cluss
nirplane
hind
deer
frog horse
ship track
cat
dog
t-SNE
UMAP
LocalMA AP
Figure 10: Visualizations of U-Net bottleneck features using t-SNE, UMAP, and LocalMAP on CIFARIO (balanced) for DDPM models trained with varying amounts of data. Each row corresponds different trained model for DDPM, with increasing samples per class from top to bottom: 100, 500, 1k, and 5k. For the final row (5k samples per class), randomly selected subset of 20k samples is visualized from the full training set of 50k.
26

---

airplane
auto
bird
cat
deer
dog
frog
horse
truck ship
Figure 11: Memorization in limited-data scenarios. For each class, the top row shows the closest
matching real training samples corresponding to the generated images in the bottom in magenta), which were produced by a DDPM model trained on CIFAR- 10 with 50 samples per class and sampled using w 0.1. The high visual similarity between generated and training samples reflects the models strong memorization behavior in limited-daa -atit settings.
(outlined
27

---

D Comparison with Other Methods
Latent Space Visualization Figure 12 visualizes the latent representations from the U-Net bot- tleneck layer using t-SNE [11] (left), UMAP [35 (middle), and LocalMAP 38 (right) for DDPM [8] CBDM [5]. T2H [6], and CORAL trained on CIFARIO-LT with an imbalance ratio p 0.01. CORAL exhibits markedly improved class-wir wise separation in latent space, mitigating the representa- tional entanglement that typically causes feature mixing between head and tail classes. Figure [3 presents analogous visualizations for CelebA-5 across all methods except T2H, for which no imple- mentation is available on this dataset.
Comparison with Baseline Methods We compare CORAL with baseline methods and highlight its strengths to address long-tailed generation in diffusion models.
CBDM: Introduces distribution adjustment regularizer during training that encourages similarity between generated images across different classes, transferring knowledge from head classes to tail classes. CBDM [3] suffers from mode collapse because its regularization loss encourages the model to produce similar outputs across different class conditions.
T2H: Employs weighted denoising score matching to transfer knowledge from head classes to tail classes by using head samples denoising targets for noisy tail samples. Its performance depends on both label distribution and sample similarity. T2H's [6] score substitution mechanism could potentially lead to mode collapse when noisy tail samples are consistently mapped to the same limited set of head references due to similarity- selection.
CORAL: As demonstrated in our experimental results, CORAL consistently outperforms both CBDM and T2H across a range of evaluation metrics, with particularly strong gains in tail classes. Our experimental results across multiple datasets highlight the strengths of CORAL:
1. Mode Stability: CORAL prevents mode collapse, and generates class consistent and visually diverse samples. As can be seen in the generated samples. This is in contrast to CBDM, which often fails preserve class identity, e.g., by generating class-conditioned samples displaying attributes of other classes, as shown in Figure 4
CORAL effectively reduces undesirable interclass feature borrowing in the class labeled datasets we consider. At the same time, CORAL allows the transfer of -on-discriminative features that facilitate generalization, as shown in Figure
2. Adaptive Regularization: CORAL incorporates time-dependent regularization into the con- trastive loss. This adaptive weighting enhances separation during the later stages of denoising, when outputs are less noisy and more semantically meaningful. Figure 8 shows that contrastive regularization is most effective when applied toward the end of the denoising process.
3. Latent Disentanglement: CORALs strength lies in leveraging the lower-dimensional latent space of the denoising U-Net, which has been shown to capture semantically meaningful structure 10). CORAL achieves effective inter-class disentanglement in the latent space by employing linear projection head (see Figure 2). resulting in high-fidelit sigh-fiilityy and class-aligned generated samples. These effects are illustrated in Figure [22 for CIFAR 10-LT with 0.01 and Figure 13 for CelebA-5, using latent space visualizations from I-SNE, UMAP, and LocalMAP.
28

---

Cluss
nirplaee
hind
deer
frog
ship
dog
t-SNE
UMAP
LAP
Figure 12: Visualizations of U-Net bottleneck features using t-SNE, UMAP, and LocaIMAP on CIFARIO-LT with an imbalance ratio of p 0.01. Each row shows different method trained
CIFAR1O-LT, listed from top to bottom: DDPM [8]. CBDM [ T2H 6) and CORAL.
29

---

Class
block
heown
blando
L-SNE
UMAP
LocalMAP
Figure 13: Visualizations of U-Net bottleneck features using t-SNE, UMAP, and LocaLMAP on ColebA-5, Each row shows a different method trained CelebA-5, listed from top to bottom: DDPM 18) CBDM [5), and CORAL.
30


Reviews:
[{"review_id": "EeflqF4oMe", "reviewer": "Reviewer_gbPZ", "summary": "This 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.", "questions": "1. 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?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "The authors’ responses have addressed some of my concerns; however, I remain concerned about the lack of experiments on large-scale models and datasets. Therefore, I have decided to maintain my original score.", "rebuttal": "We thank the reviewer for their detailed evaluation. Thank you for recognizing our strong empirical observation and meaningful contribution that deepens understanding of diffusion behavior. Your feedback has provided an excellent opportunity to explain our main contributions more clearly. We have justified our design choices with additional experiments and theoretical motivation. The added details demonstrate how our work provides novel insights and demonstrates improved performance. Below we detail our responses using W(i) and Q(i) for weakness and questions, respectively:\n\n- **W1,Q1:**  \n  While the use of contrastive learning and projection heads are established techniques in representation learning [1, 2], their application to minority class generation in diffusion models is novel because we are the first to identify latent representation entanglement at the U-Net bottleneck as a key mechanism underlying poor minority class generation. Previous work on improving long-tailed generation has focused on ambient space solutions, missing this critical failure mode. Our key contributions are three-fold:  \n  (1) identifying representation entanglement at the U-Net bottleneck as the core failure mode,  \n  (2) recognizing that effective intervention requires learning inherently separated representations rather than simply enforcing separation in the ambient space, and  \n  (3) demonstrating that the projection head is essential for creating a learned embedding space optimized for class discrimination through contrastive objectives.  \nThe advantages of projection heads and supervised contrastive learning has been well-established in other domains. It has been demonstrated that projection heads are crucial for contrastive learning [1, 3], showing that removing them significantly degrades representation quality and class separation. In [2], the authors provided theoretical analysis showing that projection heads help decouple the contrastive objective from the downstream task representation. Additionally, our choice of supervised contrastive learning is motivated by its demonstrated advantages [3]: the ability to leverage multiple positives per class within each batch, superior performance in learning discriminative representations, and better class boundary separation, which directly addresses our goal of reducing representation entanglement in long-tailed diffusion models.\nHowever, we are the first to provide empirical evidence of these principles specifically for generative diffusion models in long-tailed settings. As demonstrated in our ablation study (Figure 9, left panel), lower temperatures ($\\tau_\\mathrm{sc} \\simeq 0$) achieve optimal performance by increasing sensitivity to hard positive and negative examples within each batch. We will add an ablation study examining the interaction between batch size and performance in our final submission.\nWe verified that direct contrastive loss on bottleneck features (without a projection head) as well as ambient-space approaches fail to achieve comparable improvements, and will add those ablation studies to the final submission. Our work establishes the empirical foundation for contrastive regularization in diffusion models for long-tailed image generation.\n\n- **W2,Q2:**  \n  Our approach is grounded in both established representation learning principles from other encoder settings and our experimental results that demonstrate latent contrastive loss with projection heads significantly outperforms ambient contrastive approaches for diffusion models. Recent theoretical analysis has shown that contrastive losses with projection heads capture key task-related information more effectively than direct feature space constraints in classification and self-supervised learning [1, 2].  \nWe have conducted extensive experiments that provide empirical validation across multiple long-tailed datasets confirming that these principles extend to generative diffusion models, and showing that latent space intervention with projection heads consistently improves all generation metrics (FID, IS, Recall) compared to ambient space methods. We will include these results in the final submission. \nThe theoretical foundation for why latent space intervention outperforms ambient space approaches lies in the fundamental difference between learned versus imposed separation. Our approach addresses the root cause by learning disentangled representations at the semantically rich bottleneck layer where class overlap occurs. This is theoretically sound because:  \n  (1) the bottleneck contains compressed semantic information before decoder expansion, making it the optimal intervention point, and  \n  (2) learned metric embeddings through projection heads enable the model to develop inherently separated class representations rather than merely correcting entangled outputs.  \nOur empirical results provide strong evidence supporting this theoretical framework. Our extended experiments across multiple datasets consistently demonstrated that:  \n  1. Contrastive loss in ambient space does not improve or actually reduces diversity (lower recall),  \n  2. Direct contrastive loss on bottleneck features without projection heads is suboptimal, and  \n  3. Bottleneck contrastive learning with projection heads significantly improves all generation metrics.  \n  We will include these comprehensive ablation studies in the final submission to demonstrate the limitations of existing ambient-space approaches and validate our latent space intervention strategy.\nWe have also performed extensive ablation studies on the projection head architecture that we will include in the final submission. Our architectural choices are grounded in established contrastive learning principles from [1], which demonstrate that simple linear projection heads with normalization are optimal for most contrastive learning tasks. Our experiments confirm that this architecture transfers effectively to the diffusion setting.\n\n- **W3,Q3:**  \n  We have included comprehensive ablation studies capturing the dependence of FID on key hyperparameters in Fig. 9 in the appendix. We present three plots showcasing the dependence of FID on $\\tau_r$ from Eqn. (7) (the contrastive loss decay rate), $\\tau_\\mathrm{sc}$ from Eqn. (5) (the SupCon temperature), and the CFG weight, $\\omega$. The FID is plotted for values around the found minimum, demonstrating CORAL's robustness across reasonable parameter ranges.  \n  The functional form $\\lambda(t) = w \\cdot \\exp \\left( \\frac{1 - t/T}{\\tau_r} \\right)$ ensures that:  \n  (1) the contrastive loss weight increases as we approach the original data ($t \\rightarrow 0$), where semantic structure is most recoverable and class discrimination most meaningful,   \n  (2) the exponential decay provides smooth, differentiable weighting that doesn't introduce training instabilities.  \n  Recent work on diffusion models has explored adapting regularization strength based on noise levels [4, 5], recognizing that different timesteps require different levels of structural guidance.\nWe have also performed extensive ablation studies on the projection head architecture that we will include in the final submission. Our architectural choices are grounded in established contrastive learning principles from [1], which demonstrate that simple linear projection heads with normalization are optimal for most contrastive learning tasks. Our experiments confirm that this architecture transfers effectively to the diffusion setting.\n\n- **W4,Q4:**  \n  Our results demonstrate strong class fidelity preservation through both quantitative and qualitative evidence. Figures 3 and 8 show per-class FID improvements across all classes in CIFAR10-LT ($\\rho = 0.01$ and $\\rho = 0.001$), with CORAL consistently outperforming baselines. Figure 4 provides compelling qualitative evidence of semantic preservation for the tulip class in CIFAR100-LT. We observe that CBDM and H2T suffer from feature borrowing — producing tulips with repetitive, grass-heavy backgrounds borrowed from dominant animal classes — while CORAL generates tulips with diverse, semantically appropriate backgrounds that reflect the true tulip class distribution.  \nWe will include additional qualitative comparisons in the final submission showing semantic preservation across multiple classes, providing visual evidence of improved class fidelity that complements our quantitative results.\n\n- **Additional Results:**\nWe have expanded the scope of our experiments to include ImageNet-LT, please see a table with results in our response to the weakness stated by *Reviewer aGJo*.\n---\n\nOnce again, we thank you for the constructive feedback on our work. Working on the pointers has helped us improve the quality of our analysis. We hope we were able to clarify all your concerns, and look forward to resolving any remaining concerns during the discussion phase.\n\n\n\n[1] T. Chen, S. Kornblith, M. Norouzi, and G. Hinton. A simple framework for contrastive learning of visual representations, 2020. \n\n[2] Y. Xue, E. Gan, J. Ni, S. Joshi, and B. Mirzasoleiman. Investigating the benefits of projection head for representation learning, 2024. \n\n[3] P. Khosla, P. Teterwak, C. Wang, A. Sarna, Y. Tian, P. Isola, A. Maschinot, C. Liu, and D. Krishnan. Supervised contrastive learning. In Advances in Neural Information Processing Systems, volume 33, pages 18661–18673, 2020.\n\n[4] Y. Kim, B. Na, M. Park, J. Jang, D. Kim, W. Kang, and I.-C. Moon. Training unbiased diffusion models from biased dataset. In The Twelfth International Conference on Learning Representations, 2024. \n\n[5] R. Baptista, A. Dasgupta, N. B. Kovachki, A. Oberai, and A. M. Stuart. Memorization and regularization in generative diffusion models, 2025."}, {"review_id": "dY2XYOqGPS", "reviewer": "Reviewer_AH8g", "summary": "The 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.", "questions": "- 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?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "I appreciate the authors rebuttal and the additional results. They have addressed my concerns. I will be increasing my score.", "rebuttal": "We sincerely thank you for your thoughtful feedback and for finding our paper well written. We have addressed each concern comprehensively strengthen both our empirical validation and theoretical motivation. Below we detail our responses using W(i) and Q(i) for weakness and question, respectively:\n\n- **W1:** Our choice of bottleneck intervention is theoretically and empirically motivated over ambient approaches. Contrastive losses with projection heads capture task-related information more effectively [1,3], with [3] showing projection heads decouple contrastive objectives from downstream task representations. Our approach learns disentangled representations at the semantically rich bottleneck layer. Following the theoretical framework established in [1,2,3], the projection head facilitates learning of inherently separated class representations.\nOur empirical results provide strong evidence supporting this theoretical framework. Our extended experiments across multiple datasets consistently demonstrate that: \n  1. Contrastive loss in ambient space on negative pairs only (similar to DiffROP) can reduce intra-class diversity\n  2. Direct contrastive loss on bottleneck features without projection heads can affect the diffusion objective\n  3. Bottleneck contrastive learning with projection heads (CORAL) significantly outperforms ambient methods on generation metrics. \n    We will include these comprehensive ablation studies in the final submission to demonstrate the limitations of existing ambient-space approaches and showing that effective intervention requires learning inherently separated representations rather than simply enforcing separation in the ambient space.\n    Our architectural choices are grounded in established contrastive learning principles from [1], which demonstrated that simple linear projection heads with normalization are optimal for most contrastive learning tasks. Our experiments confirmed that this architecture transfers effectively to the diffusion setting, we will include our ablation studies on projection head architecture in the final submission.\n    Our choice of $\\lambda(t)$ in Eq. (7) ensures that the contrastive loss weight increases as we approach the original data ($t \\rightarrow 0$), where semantic structure is most recoverable and class discrimination most meaningful. The exponential decay provides smooth, differentiable weighting that doesn't introduce training instabilities. Recent work in diffusion models have implemented this [4,5], as different timesteps require different levels of structural guidance. \nAdditionally, please see ImageNet-LT results in our response to Reviewer aGJo.\n\n- **W2:** Reweighting approaches (CBDM[6], TIW[4]) face limitations:  \n    1. Estimation challenges - accurately estimating class ratios is difficult and error-prone [7], with TIW requiring additional discriminator training that introduces computational overhead.  \n    2. Training instability - heavy reweighting of tail classes can cause gradient explosion, particularly problematic for diffusion models.  \n    3. Mode collapse risk - over-emphasizing tail classes reduces sample diversity, as noted in CBDM.  \n    Additionally, CBDM's reliance on head class predictions to enhance tail generation introduces bias and reduces robustness as noted in [8].  \n    Our experimental results demonstrate that CORAL's representation-based approach consistently outperforms these reweighting methods across all datasets and metrics presented.\n\n- **W3:** We thank the reviewer for this observation. The critical distinction is not simply where overlap occurs, but where contrastive intervention is most effective for learning disentangled representations.\n\n    While overlap may be observable at multiple levels (ambient space, intermediate layers, bottleneck), the effectiveness of contrastive learning depends critically on where the intervention is applied. Our experimental results demonstrate that applying contrastive loss at the bottleneck is optimal compared to interventions at other locations. This aligns with established principles from representation learning showing that contrastive objectives work best when applied to semantically rich, compressed representations [1,2].\n\n    We will include comprehensive ablation studies in the final submission showing: (1) Direct comparison with ambient space contrastive approaches, showing superior performance of bottleneck intervention, (2) Ablations demonstrating that projection heads are essential for effective contrastive learning at the bottleneck level.\n\n- **W4,Q2:** We thank the reviewer for this important observation. Our approach enhances rather than reduces within-class diversity through our use of a projection head. Recent theoretical work by [3] demonstrates that projection heads prevent class collapse and preserve within-class diversity by enabling pre-projection representations to learn subclass-level features that are not represented post-projection, thus maintaining diversity within classes.  \n    Figure 1 demonstrates inter-class separation rather than reduced intra-class diversity - the t-SNE visualization shows CORAL separates previously entangled classes while each cluster retains internal structure. Figure 4 provides visual evidence that CORAL generates more diverse tulips with varied characteristics compared to competing methods.  \n    Quantitatively, CORAL consistently outperforms baselines on diversity metrics (Recall in Table 1), demonstrating that our approach enhances rather than reduces within-class diversity, as established theoretically in [3].\n\n- **Q1:** Thank you for this important question about Figure 3. You are absolutely correct to expect different performance patterns between head and tail classes in a more severely imbalanced setting.  \n    In Figure 3, we show results for CIFAR10-LT with an imbalance factor $\\rho$ = 0.01. The classes are ordered alphabetically in the long-tail distribution: airplane (head class), auto, bird, cat, deer, dog, frog, horse, ship, and truck (tail class).  \n    The reason DDPM shows relatively uniform performance across all classes in Figure 3 is that the imbalance factor $\\rho$ = 0.01, while significant, is not extreme enough to create the level of performance degradation typically observed in severely imbalanced scenarios. The choice of the long-tail parameter $\\rho$ directly affects the magnitude of per-class FID differences. With less severe imbalance, we observe more comparable FID scores across classes.  \n    Figure 8 shows the same CIFAR10-LT dataset but with a much more severe imbalance factor of $\\rho$ = 0.001. Figure 8 demonstrates the expected long-tail generation behavior you mentioned, we observe the expected behavior for long-tail generation problems, where tail classes exhibit substantially higher FID scores compared to head classes. We will include additional visualizations in the final submission showing this enhanced diversity across multiple classes.  \n\n    This comparison between Figures 3 and 8 illustrates how the severity of class imbalance (controlled by $\\rho$) directly influences the magnitude of performance differences across the head-to-tail spectrum, and demonstrates that CORAL's improvements are consistent across different levels of imbalance severity.\n\n- **Q3:** Thank you for this excellent question about the practical applicability of our approach versus using pretrained models.  \n    Prompt engineering and textual inversion primarily help with diversity by forcing the model to sample from minority classes more frequently, but they don't inherently improve the quality of those minority class generations. These methods essentially guide the sampling process toward underrepresented regions of the learned distribution. We acknowledge that for many practical applications, especially those involving natural images with good textual descriptions, leveraging pretrained models with prompt engineering or textual inversion can be more efficient and effective. CORAL could very well be combined with these prompt-based methods to produce even better quality images, which represents an interesting direction for future research.\n\n---\n\nOnce again, we thank you for the constructive feedback on our work. Working on the pointers has helped us improve the quality of our analysis. We hope we were able to clarify all your concerns, and look forward to resolving any remaining concerns during the discussion phase.\n\n[1] T. Chen, S. Kornblith, M. Norouzi, and G. Hinton. A simple framework for contrastive learning of visual representations, 2020.  \n[2] P. Khosla, P. Teterwak, C. Wang, A. Sarna, Y. Tian, P. Isola, A. Maschinot, C. Liu, and D. Krishnan. Supervised contrastive learning. In Advances in Neural Information Processing Systems, volume 33, pages 18661–18673, 2020.  \n[3] Y. Xue, E. Gan, J. Ni, S. Joshi, and B. Mirzasoleiman. Investigating the benefits of projection head for representation learning, 2024.  \n[4] Y. Kim, B. Na, M. Park, J. Jang, D. Kim, W. Kang, and I.-C. Moon. Training unbiased diffusion models from biased dataset. In The Twelfth International Conference on Learning Representations, 2024.  \n[5] R. Baptista, A. Dasgupta, N. B. Kovachki, A. Oberai, and A. M. Stuart. Memorization and regularization in generative diffusion models, 2025.  \n[6] Y. Qin, H. Zheng, J. Yao, M. Zhou, and Y. Zhang. Class-balancing diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023.  \n[7] B. Rhodes, K. Xu, and M. U. Gutmann. Telescoping density-ratio estimation. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 4905–4916. Curran Associates, Inc., 2020.  \n[8] T. Zhang, H. Zheng, J. Yao, X. Wang, M. Zhou, Y. Zhang, and Y. Wang. Long-tailed diffusion models with oriented calibration. In The Twelfth International Conference on Learning Representations, 2024."}, {"review_id": "QbTY1GxDEt", "reviewer": "Reviewer_rXbT", "summary": "This 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.", "questions": "This 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.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths\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.", "was_revised": true, "final_justification": "The concerns I raised were appropriately addressed through the rebuttal and additional experiment of the author. So I raising my score to 4.", "rebuttal": "We sincerely thank the reviewer for their insightful feedback and thorough evaluation. Thank you for acknowledging our simplicity and effectiveness and recognizing that our approach successfully disentangles latent representations.  The recommended additions have helped address the identified limitations and provided deeper insights into our method’s behavior. These improvements positively enhance our contribution's impact. Below, we carefully address each concern, referring to weaknesses as W(i) and addressing the questions:\n\n\n- **W1** Thank you for this insightful observation about the broader applicability of our contrastive latent regularization approach. You are correct that the core mechanism could potentially benefit representation learning in diffusion models more generally. However, we specifically focus on long-tailed settings because these scenarios exhibit the most severe and well-documented problems with both fidelity and diversity of generated samples, particularly for tail classes.\nContrastive latent regularization provides clear, measurable benefits in long-tail settings. While we agree that exploring the broader applicability to general representation disentanglement would be valuable future work, we hope that our demonstration of the effectiveness of contrastive latent regularization in this challenging scenario will inspire future research to explore its potential in more general diffusion model settings and broader representation learning contexts.\n\n- **W2** \nThank you for this important question about how our latent space regularization enhances rather than interferes with generative performance. We address this concern through both theoretical grounding and comprehensive empirical validation. \nCORAL learns separation rather than imposed constraints on the generative process. While our contrastive regularization does affect the diffusion learning process, learned metric embeddings through projection heads enable the model to develop inherently separated class representations that support better generative performance.\nRecent theoretical analysis has demonstrated that contrastive losses with projection heads capture key task-related information more effectively than direct feature space constraints. The projection head enables better representation learning by decoupling the contrastive objective from the main diffusion features, allowing the model to learn disentangled representations in a auxiliary embedding space while the main bottleneck features continue to serve the diffusion objective. This complementary learning approach results in improved rather than competing objectives.\nOur comprehensive experiments across multiple datasets provide strong empirical evidence that latent regularization enhances rather than degrades generative performance. We consistently observe improvements in all generation metrics (FID, IS, Recall), suggesting that better-separated representations actually enhance the diffusion training objective. \n\n- **W3, Details on experimental setup** Thank you for this important clarification request regarding our experimental setup. We provide the following details and we will include all details in the final submission:\n  1. During sampling for generated data, we sample the class label $y$ from a uniform distribution across all classes, ensuring equal representation regardless of the training distribution imbalance. This approach provides a fair comparison across methods by evaluating each method's ability to generate high-quality samples for all classes, including underrepresented ones.\n  2. For overall metric calculations  we use the balanced version of the datasets for the real data to ensure fair evaluation. All metrics are measured with 50K generated images, except for per-class FID studies. For per-class FID analysis (Figure 3 and Figure 8), we generate 5K samples for each class and compare against 5K real samples from the balanced dataset for that specific class. \n  3. We tune the classifier-free guidance strength $\\omega$ for both baselines and our method to ensure optimal performance for each approach. \n  4. Details on hyperparameter settings are provided in Table 2 in the appendix.\nThis experimental design is used in state of the art and allows us to fairly evaluate each method's capability to generate high-quality samples across all classes, particularly focusing on the challenging tail classes.\n\n- **W4** While the use of contrastive learning and projection heads are established techniques in representation learning [1, 2], their application to minority class generation in diffusion models is novel because we are the first to identify latent representation entanglement at the U-Net bottleneck as a key mechanism underlying poor minority class generation. Previous work on improving long-tailed generation [3,4,5] has focused on ambient space solutions, missing this critical issue occurring within the diffusion architecture itself. Our key contributions are three-fold: (1) identifying representation entanglement at the U-Net bottleneck as the core failure mode, (2) recognizing that effective intervention requires \\textit{learning} inherently separated representations rather than simply enforcing separation in the ambient space, and (3) demonstrating that the projection head is essential for creating a learned embedding space optimized for class discrimination through contrastive objectives. \nThe necessity of projection heads and supervised contrastive learning has been well-established in other domains. It has been demonstrated that projection heads are crucial for contrastive learning [1, 6], showing that removing them significantly degrades representation quality and class separation. In [2], the authors provided theoretical analysis showing that projection heads help decouple the contrastive objective from the downstream task representation. Additionally, our choice of supervised contrastive learning is motivated by its demonstrated advantages [6]: the ability to leverage multiple positives per class within each batch, superior performance in learning discriminative representations, and better class boundary separation, which directly addresses our goal of reducing representation entanglement in long-tailed diffusion models.\nHowever, we are the first to provide empirical evidence of these principles specifically for generative diffusion models in long-tailed settings. As demonstrated in our ablation study (Figure 9, left panel), lower temperatures ($\\tau_\\mathrm{sc} \\simeq 0$) achieve optimal performance by increasing sensitivity to hard positive and negative examples within each batch. We will add a ablation study examining the interaction between batch size and performance in our final submission.\nWe verified that direct contrastive loss on bottleneck features as well as ambient-space approaches fail to achieve comparable improvements, and will add those ablation studies to the final submission. Our work establishes the empirical foundation for contrastive regularization in diffusion models for long-tailed image generation \n\n\n- **Sensitivity to batch size** This is an excellent question about batch size sensitivity in contrastive learning, and we appreciate the reviewer bringing up this important consideration. \nWe would like to note that our choice of temperature $\\tau_{SC}$ (Figure 9 on appendix, left panel) in the SupCon loss is motivated by the results in the appendix of  [6] which show that lower temperatures focus learning on hard positives/negatives, improving the discriminative power of learned representations even when working with the available examples in each batch.\nWe acknowledge that a comprehensive ablation study examining the interaction between batch size and performance would strengthen our analysis and provide valuable insights into the robustness of our approach. We will include this experiment in the final submission.\n\n-**Additional Results:** We have expanded the scope of our experiments to include ImageNet-LT, please see a table with results in our response to the weakness stated by Reviewer aGJo.\n\n---\n\nWe thank you for the constructive feedback on our work. Working on the pointers has helped us improve the quality of our analysis. We hope we were able to clarify all your concerns, and look forward to resolving any remaining concerns during the discussion phase.\n\n\n[1] T. Chen, S. Kornblith, M. Norouzi, and G. Hinton. A simple framework for contrastive learning of visual representations, 2020.\n\n[2] Y. Xue, E. Gan, J. Ni, S. Joshi, and B. Mirzasoleiman. Investigating the benefits of projection head for representation learning, 2024. \n\n[3] Y. Qin, H. Zheng, J. Yao, M. Zhou, and Y. Zhang. Class-balancing diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023.\n\n[4] T. Zhang et al. Long-tailed diffusion models with oriented calibration, 2024.\n\n[5] D. Yan, L. Qi, V. T. Hu, M.-H. Yang, and M. Tang. Training class-imbalanced diffusion model via overlap optimization. CoRR, abs/2402.10821, 2024. \n\n[6] P. Khosla, P. Teterwak, C. Wang, A. Sarna, Y. Tian, P. Isola, A. Maschinot, C. Liu, and D. Krishnan. Supervised contrastive learning. In Advances in Neural Information Processing Systems, volume 33, pages 18661–18673, 2020."}, {"review_id": "tlPhSRmmhz", "reviewer": "Reviewer_aGJo", "summary": "This 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.", "questions": "1. 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?", "limitations": "1. 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", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**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).", "was_revised": true, "final_justification": "I justified the score in the discussion thread.", "rebuttal": "We sincerely thank you for your thoughtful and comprehensive review. Thank you for recognizing our novel approach and noting it makes you lean toward the positive side. We have incorporated the suggested improvements into our work and found your questions thought-provoking for future work. The additional experiments have strengthened our experimental validation. We address each point below, using W(i) for weakness Q(i) for questions and L(i) for limitations:\n\n- **W1:** \nThank you for this feedback. We would like to clarify that our evaluation extends beyond CIFAR datasets and we believe CORAL can be extended beyond U-Net-based diffusion models.\n  1. Our evaluation extends beyond CIFAR datasets and includes experiments on various datasets with different imbalance ratios. We present CelebA-5 results in Table 1 and various visualizations in the appendix, Additionaly, we have expanded the scope of our experiments to include ImageNetLT, we provide a table with results here:\n| Dataset | Method | FID ↓ | IS ↑ |  Recall ↑ |\n|---------|---------|-------|------|----------|\n| ImageNet-LT | DDPM | 17.08 |  21.03 |  0.39 |\n| ImageNet-LT | CBDM | 22.66 |  17.13 |  0.42 |\n| ImageNet-LT | T2H | 18.59 |  15.12 | 0.44 |\n| ImageNet-LT | CORAL (Ours) | **16.09** |  **22.06** | **0.47** |\n  We will also include additional visualizations of our ImageNet results in the final submission.\n  2. While our current experiments focus on U-Net architectures, the core principle of contrastive regularization at bottleneck representations is architecturally agnostic and could extend to other diffusion architectures. \n  3.  The identification of representation entanglement as a failure mode in tail-class generation and our solution through contrastive latent alignment can be extended beyond specific architectures or datasets, providing insights applicable to the broader generative models.\n\n- **Q1:** That is an excellent question. We believe our approach can be generalized to Stable Diffusion-style architectures, and this is an active area of future work for our team.\nApplying contrastive regularization at the bottleneck layer to address representation entanglement can be used in Stable Diffusion models that also employ a U-Net architecture with bottleneck layers.\nExtending CORAL to Stable Diffusion would be particularly valuable for text-to-image generation with long-tailed concept distributions, where certain visual concepts or object categories are underrepresented in training data. This could improve generation quality for rare or minority concepts while maintaining the flexibility of prompt-based control.\nWe believe it represents a promising direction that could significantly broaden the impact of our contribution.\n\n- **Q2:** Thank you for this insightful question about diffusion transformer architectures. This is indeed an important consideration given the recent prominence of transformer-based diffusion models.\nDiffusion transformer architectures vary significantly in their design [1,2]. Some maintain bottleneck-like structures where information is compressed through hierarchical attention mechanisms, while others operate directly on image patches with uniform attention across all spatial locations without explicit compression stages.\nWe believe that representation entanglement should still occur in transformer architectures because they retain the encoder-decoder structure with a bottleneck. For transformer architectures with bottleneck-like structures, CORAL could potentially be adapted by applying contrastive regularization at the compressed representation layers. However, in some transformer architectures the phenomenon might manifest differently, potentially in the attention patterns or in intermediate layer representations where semantic information is aggregated. This represents an exciting research direction. \n\n- **Q3:** Thank you for this excellent question about extending CORAL to unconditional modeling with biased datasets. \nYou raise an important point about the fundamental challenge: without explicit class labels, identifying and correcting bias in unconditional generation becomes significantly more complex. CORAL's current approach leverages supervised contrastive learning, which requires class information to define positive and negative pairs for the contrastive loss.\nFor instance, one could imagine using unsupervised clustering techniques to identify latent groups in the data, then applying contrastive regularization to encourage separation between these discovered clusters. Alternatively, methods that can detect distributional imbalance in an unsupervised manner could provide the side-information needed to guide the model toward more balanced generation. As it stands, CORAL requires class labels to define the contrastive learning objective, which limits its direct application to unconditional biased settings. \nThis represents a compelling direction for future work. Self-supervised representation learning could potentially be adapted for this setting.  This might involve combining CORAL's bottleneck regularization approach with unsupervised methods or self-supervised techniques. \n\n- **L1:** Thank you for bringing these important papers to our attention. We appreciate this insightful observation about diffusion autoencoders and their relationship to our work.\nIf labeled information is available, our CORAL approach could potentially be adapted to these architectures by applying contrastive regularization to the appropriate latent spaces - whether that's the U-Net bottleneck, the auxiliary encoder's output, or both, depending on where semantic clustering occurs.\nWhile investigating these architectures would be valuable, such experiments are beyond the scope of the current work. However, we agree that these papers represent important related work in the context of semantic representation learning in diffusion models, and we will include them in our related work section.\n\n- **L2:** Thank you for pointing out this important ablation study. We direct you to Figure 9 in the appendix, which provides a comprehensive ablation on the temperature scaling parameter $\\tau_r$ that controls the time-dependent weighting function $\\lambda(t)$. This figure shows the effect of $\\tau_r$ on FID performance for the CIFAR10-LT dataset.\nWe will also be including several additional experiments and ablations in the final submission:\n  1. Ambient space contrastive loss comparisons demonstrating that ambient-space approaches can reduce intra-class diversity while CORAL consistently improves all generation metrics\n  2. Direct contrastive loss on bottleneck features without projection heads showing this approach can interfere with the diffusion objective\n  3. Comprehensive projection head architecture ablations validating our architectural choices \n  4. ImageNet-LT complete results \n  5. Batch size sensitivity analysis examining the interaction between batch size and CORAL\n\n---\n\nOnce again, we thank you for the constructive feedback on our work. Working on the pointers has helped us improve the quality of our analysis. We hope we were able to clarify all your concerns, and look forward to resolving any remaining concerns during the discussion phase.\n\n\n[1] W. Peebles and S. Xie. Scalable diffusion models with transformers, 2023\n\n[2] A. Hatamizadeh, J. Song, G. Liu, J. Kautz, and A. Vahdat. Diffit: Diffusion vision transformers\nfor image generation, 2024"}]

Author rebuttal:

Meta-review: The authors propose CORAL (Contrastive Regularization for Aligning Latents), which introduces a projection head at the bottleneck and applies a supervised contrastive loss to separate head and tail representations to address the challenges that tail classes suffer from poor generative quality from diffusion models. 

The paper identifies latent feature overlap in diffusion bottlenecks as a key cause of tail-class underperformance, which is novel and is less explored before, thus offering a valuable contribution to understanding diffusion behavior. As a solution, the authors propose a design with only a projection head and supervised contrastive loss. The authors dive deep in the empirical perspectives and provide sufficient evidence in different settings: comparing with previous works, CORAL show significant and consistent performance gain across class-long-tailed datasets, including CIFAR-10, CIFAR-100, ImageNet. The reviewers have spotted concerns regarding limited ablations and lack of theoretical contributions to fully reveal the properties.  In the rebuttal, the authors have provided comprehensive experiments to address the concerns. Although in the end the loss is designed from heuristic perspective, given the authors have provided sufficient empirical evidence, all reviewers agree the method is technical sound and give a positive rating. Currently the paper meets the bar of publication, and the authors are encouraged to clarify the exploration focus on the empirical perspectives by carefully revising the paper to incorporate the reviewers' comments, and continue the exploration on the theoretical side.
Decision: Accept (poster)