Contrastive Conditional–Unconditional Alignment for Long-tailed Diffusion Model

## Abstract

Abstract Training data for class-conditional image synthesis often exhibit a long-tailed distribution with limited images for tail classes. Such an imbalance causes mode collapse and reduces the diversity of synthesized images for tail classes. For class-conditional diffusion models trained on imbalanced data, we aim to improve the diversity and fidelity of tail class images without compromising the quality of head class images. We achieve this by introducing two simple but highly effective loss functions. Firstly, we employ an Unsupervised Contrastive Loss (UCL) utilizing negative samples to increase the distance/dissimilarity among synthetic images.
Such regularization is coupled with a standard trick of batch resampling to further diversify tail-class images.
Our second loss is an Alignment Loss (AL) that aligns class-conditional generation with unconditional generation at large timesteps. This second loss makes the denoising process insensitive to class conditions for the initial steps, which enriches tail classes through knowledge sharing from head classes.
We successfully leverage contrastive learning and conditional-unconditional alignment for class-imbalanced diffusion models. Our framework is easy to implement as demonstrated on both U-Net based architecture and Diffusion Transformer. Our method outperforms vanilla denoising diffusion probabilistic models, score-based diffusion model, and alternative methods for class-imbalanced image generation across various datasets, in particular ImageNet-LT with 256 × \times 256 resolution.

## 1 Introduction

Recent advances in diffusion models (Ho et al., 2020; Song et al., 2020) have led to breakthrough in various generation tasks such as image generation (Zhang et al., 2023; Rombach et al., 2022; Ruiz et al., 2023), video generation (Ho et al., 2022a; b), image editing (Couairon et al., 2022; Kawar et al., 2023), 3D generation (Poole et al., 2022), etc.
These diffusion-based generative models rely on large-scale datasets for training, which often follow a long-tailed distribution with dominant data for head classes and limited data for tail classes.
Similar to the class-imbalanced recognition model (Liu et al., 2022) and the class-imbalanced generative adversarial network (Tan et al., 2020; Rangwani et al., 2022; 2023; Khorram et al., 2024), diffusion model is not able to generate high-quality images for tail classes due to the scarcity of training data.
As shown in Fig. [1](https://arxiv.org/html/2507.09052v2#S1.F1), the original diffusion model with transformer backbone  (Ma et al., 2024) generates inferior images for a tail class (red wine) when trained using a long-tailed version of ImageNet.
We aim to increase the fidelity and diversity of tail class images while maintaining the quality of head class images for diffusion models.
Our work has a positive societal impact on content generation of underrepresented groups.

We look into a highly related and crucial problem of long-tailed image recognition, for which contrastive learning is an effective method (Jiang et al., 2021; Li et al., 2022).
Contrastive learning is a general learning paradigm of minimizing the distance of positive pairs while maximizing the distance of negative pairs, which has attracted lots of research interest for self-supervised representation learning (Chen et al., 2020).
On the other hand, diffusion model has also emerged as an effective model for representation learning (Fuest et al., 2024; Baranchuk et al., 2021).
We explicitly incorporate contrastive learning regularization to diffusion models to further enhance representation learning towards better image generation.

Figure: (a) SiT (Ma et al., 2024)
Refer to caption: figures_iclr/sit_966_redwine.png

Our first loss is an unsupervised contrastive loss with negative samples only, serving as regularization for denoising diffusion probabilistic models (DDPM) and score-based diffusion models (SBDM).
Given that mode collapse during inference manifests as generated samples being overly similar for tail classes, we introduce unsupervised contrastive learning with negative samples to maximize the distances between generated images.
Unlike supervised contrastive learning (Khosla et al., 2020), unsupervised contrastive learning distinguishes between representations of images regardless of their class, hence increasing intra-class and inter-class image diversity.

Our second loss is an alignment loss designed to align estimated noises from conditional and unconditional generation, which effectively minimizes the KL divergence between latent distributions for conditional and unconditional generation.
While such conditional-unconditional alignment seems undesirable with less controllability, a critical aspect of our approach is that our alignment loss is weighted more for larger timesteps corresponding to the initial stage of the reverse process.
Given the known overfitting in conditional generation for tail classes for both GAN,
we are motivated by the success of conditional-unconditional alignment for long-tailed GAN (Khorram et al., 2024; Shahbazi et al., 2022), which facilitates knowledge sharing between head classes and tail classes.
Notably, unconditional GAN generation has been observed to achieve superior FID than conditional generation under limited data (Shahbazi et al., 2022).
We adapt conditional-unconditional alignment from GAN to diffusion models.
Unlike the GAN-based method (Khorram et al., 2024) which aligns conditional generation and unconditional generation for low-resolution representations of images exhibiting intra-class similarity, we leverage observed image similarity during the initial denoising steps for tail classes and head classes (Si et al., 2024).
Hence, we propose to match conditional generation with unconditional generation for large timesteps.

While contrastive learning has been proposed for diffusion models to improve adversarial robustness (Ouyang et al., 2023), find semantically meaningful directions (Dalva & Yanardag, 2024), accelerate training (Yu et al., 2024), and regularize representation (Wang & He, 2025), we effectively leveraged contrastive learning for class-imbalanced diffusion models and demonstrate the superior performance of our method on long-tailed image generation via comprehensive experiments.

The main contributions of this work are as follows:

- •
We propose Contrastive Conditional-Unconditional Alignment for Diffusion Model (CCUA) with imbalanced data. Our proposed contrastive loss with conditional-unconditional alignment are easy to implement with standard DDPM and SBDM pipeline using both UNet-based architecture and Diffusion Transformer (DiT).
- •
Our first loss, Unsupervised Contrastive Loss (UCL), employs unsupervised contrastive learning loss with negative samples only, enhancing within-class diversity.
- •
Our second Alignment Loss (AL) aligns unconditional generation and conditional generation for the initial steps in the denoising process, facilitating knowledge sharing between head and tail classes.
- •
We improved the diversity and fidelity of tail class images for conditional generation while maintaining the quality of head class images for multiple datasets and various resolutions, in particular ImageNet-LT with 256x256 resolution.

## 2 Related Work

Class-imbalanced Image Generation Generative models such as VAE (Kingma et al., 2019), GAN (Karras et al., 2019; Brock et al., 2019), and diffusion models (Ho et al., 2020; Rombach et al., 2022) generate inferior images for tail classes when trained with real-world data with a long-tailed distribution.
It has attracted a lot of research interest (Ai et al., 2023; Tan et al., 2020; Khorram et al., 2024; Qin et al., 2023; Zhang et al., 2024) to address this issue for various types of models.
Many methods address the problem of class imbalance by augmenting training data for the tail classes.
A VAE is fine-tuned on tail classes under a majority-based prior (Ai et al., 2023).
It is observed that GAN (Tan et al., 2020) can amplify biases, leading to tail classes to be barely generated during inference, highlighting that the fairness in GAN needs improvement.
Khorram et al. (2024) propose a GAN-based long-tailed generation method, named UTLO, which shares the latent representations of conditional GAN with unconditional GAN and implicitly shares knowledge between the head class and tail class.
The motivation with UTLO is that low-resolution representations of images from GAN are similar for head classes and tail classes.
We observe a similar phenomenon for denoised images in the initial steps of the denoising process, and further propose a conditional-unconditional alignment loss designed for diffusion models.
Recent work addresses class-imbalanced diffusion models (Qin et al., 2023; Yan et al., 2024; Zhang et al., 2024) by regularization losses to align or separate the distributions of synthetic images and their corresponding latent representations across different classes.
For example, CBDM (Qin et al., 2023) loss minimizes the distance of estimated noise items between these two models, the original DDPM model and a second model trained with pseudo labels which form a uniform distribution.
DiffROP (Yan et al., 2024) attempts to combine contrastive learning with diffusion model by maximizing the distance of distributions between classes.
However, DiffROP only considers pairs of images of different classes as negative pairs without regularizing images of the same class.
Conceptually, the first part of our method, the unsupervised contrastive loss designed to disperse unconditional latent space, generalizes DiffROP by including extra negative pairs of images within each class, hence promoting intra-class diversity.

Contrastive Learning for Representation Learning Both unsupervised contrastive learning (Chen et al., 2020; He et al., 2020) and supervised contrastive learning (Khosla et al., 2020; Li et al., 2022) are approaches used in representation learning to produce feature embeddings by maximizing the similarity between positive pairs and minimizing the similarity between negative pairs.
Unsupervised contrastive learning does not require labels and typically augments the same data to form positive pairs (Chen et al., 2020; He et al., 2020).
Supervised contrastive learning incorporates class labels and pulls together all samples from the same class while pushing apart samples from different classes.
Contrastive learning is an effective method for self-supervised learning (He et al., 2020) and long-tailed recognition (Jiang et al., 2021; Li et al., 2022).
We effectively utilized contrastive learning for class-imbalanced diffusion models.
Our work is different from the contrastive-guided diffusion process (Ouyang et al., 2023), which aims to improve adversarial robustness.
Another difference is that we focus on contrastive learning during training, while contrastive-guided diffusion process (Ouyang et al., 2023) is about guidance for inference.
While the diffusion model originally for generation tasks becomes an emerging technique for representation learning (Fuest et al., 2024; Baranchuk et al., 2021), we directly integrate contrastive learning regularization to diffusion models for better generation on imbalanced data.
Notably, unsupervised contrastive loss with negative samples exclusively focuses on separating dissimilar instances in the embedding space, which we leverage to address mode collapse, as detailed in Sec . [3.1](https://arxiv.org/html/2507.09052v2#S3.SS1).
In the context of generative modeling, REPA (Yu et al., 2024) aligns diffusion model features with features from a frozen vision encoder to accelerate training. Unlike REPA, we directly regularize diffusion model features without external models.
Most relevant is concurrent work of dispersive loss (Wang & He, 2025), which is similar to our unsupervised contrastive loss with negative pairs only.
Dispersive loss is directly and explicitly applied to conditional training, as designed for balanced diffusion models.
By comparison, our unsupervised contrastive loss is formulated in the unconditional latent space and is implicitly extended to conditional training through our conditional–unconditional alignment loss.
This carefully designed contrastive conditional-unconditional alignment framework for long-tailed diffusion models achieves better generation as shown in experiments.

## 3 Method

Figure: Figure 2: Overview of the proposed CCUA framework of contrastive learning for diffusion model. The noise estimation network is divided into two parts, a latent encode network $e(*)$ and a decode network $d(*)$. $e(*)$ encodes an image with noise $x_{t}$ to a low-dimensional latent $h$, which is decoded to noise $\epsilon$ by $d(*)$. Specifically, the latent encode network $e(*)$ and decode network $d(*)$ for unet-based model and diffusion transformer are shown in Fig. [5](https://arxiv.org/html/2507.09052v2#A1.F5) in the appendix.[A.1](https://arxiv.org/html/2507.09052v2#A1.SS1). We increase the distance of negative pairs of unconditional encoded latents for different samples by an unsupervised contrastive loss $\mathcal{L}_{ucl}$ with negative samples only, and align unconditional and conditional generation for the same sample $x_{t}^{i}$ and utilize an alignment loss $\mathcal{L}_{al}$ to minimize their distances at initial time steps.
Refer to caption: x1.png

We propose Contrastive Conditional-Unconditional Alignment (CCUA) framework for Diffusion Model.
Fig. [2](https://arxiv.org/html/2507.09052v2#S3.F2) gives an overview of our framework, while Fig. [5](https://arxiv.org/html/2507.09052v2#A1.F5) in the appendix [A.1](https://arxiv.org/html/2507.09052v2#A1.SS1) details how we apply CCUA framework into UNet-based model and Diffusion Transformer.
Our framework involves a unsupervised contrastive loss with negative pairs only, and a conditional-unconditional alignment loss, as outlined in Sec. [3.1](https://arxiv.org/html/2507.09052v2#S3.SS1) and Sec. [3.2](https://arxiv.org/html/2507.09052v2#S3.SS2), respectively.
Sec. [3.3](https://arxiv.org/html/2507.09052v2#S3.SS3) summarizes our overall framework.

### 3.1 Unsupervised Contrastive Loss

We observed the synthesized images of DDPM for a tail class concentrated around the limited training images in the latent space, leading to the issue of mode collapse, as shown in the Fig. [9](https://arxiv.org/html/2507.09052v2#A1.F9), Fig. [10](https://arxiv.org/html/2507.09052v2#A1.F10) in the appendix [A.3.3](https://arxiv.org/html/2507.09052v2#A1.SS3.SSS3), and Fig. [8](https://arxiv.org/html/2507.09052v2#A1.F8) in the appendix [A.3.2](https://arxiv.org/html/2507.09052v2#A1.SS3.SSS2) with visualizing the distribution of synthetic images for the tail class.
By applying the contrastive loss with negative samples only, we intend to increase the distance of each synthesized image from other images in the latent space, hence increasing the diversity of synthetic images in particular for tail classes.

We consider contrastive loss with negative samples only, where the negative samples are based on noisy images from a mini-batch.
Specifically, we divide a noise estimation network into two parts, a latent encode network $e(*)$, and decode network $d(*)$.
$e(*)$ encodes an image with noise $x_{t}$ to a low-dimensional latent $h$, which is decoded to the estimate noise item $\epsilon$ by $d(*)$.
Our contrastive loss is defined for the latents $h$ with each image in a mini-batch treated as an anchor.
All other images are treated as negative samples.
The anchor sample $x_{t}^{i}$ and negative samples $x_{t}^{j}$ are fed into the encoder $e(*)$ to get the latents $h_{anc}^{i}=e(x_{t}^{i};\emptyset)$ and $h_{neg}^{j}=e(x_{t}^{j};\emptyset)$.
Our unsupervised contrastive loss $\mathcal{L}_{ucl}$ is defined as follows:

$$ $\displaystyle\mathcal{L}_{ucl}=-\frac{1}{|B|}\sum_{i\in B}{\log\frac{\pi^{i}_{anc}}{\pi^{i}_{anc}+\sum_{j\in{B},j\neq i}{\pi^{j}_{neg}}}},\text{with }\pi^{i}_{anc}=\exp(\frac{h_{anc}^{i}\cdot h_{pos}^{i}}{\tau}),\pi^{j}_{neg}=\exp(\frac{h_{anc}^{i}\cdot h_{neg}^{j}}{\tau}),$ (1) $$

where $B$ denotes a mini-batch, and $\tau$ is a temperature for softmax which we keep $0.1$ as the default setting.
Note that it is common to augment an anchor image $x_{t}^{i}$ to be a positive sample in many unsupervised contrastive learning methods. Here, we don’t apply any data augmentation and directly treat the anchor latent $h_{anc}^{i}$ itself as the positive vector in the contrastive loss, i.e., $h_{pos}^{i}:=h_{anc}^{i}$.

Why contrastive loss encourages diversity and discourages mode collapse? In the worst case, that all embeddings collapse to a constant vector, the contrastive loss becomes $\log N$ for a batch with $N$ samples.
Such constant embeddings yield the maximum possible loss and are therefore discouraged.
Intuitively, contrastive loss introduces a repulsive force between different samples.
The numerator remains constant, as self-similarity is always maximal, while the denominator aggregates pairwise similarities across the batch. Minimizing the loss requires reducing similarities between distinct samples, effectively pushing their embeddings apart in feature space.
In this way, the model avoids collapse by maximizing inter-sample distances. Geometrically, the optimal solution (in the absence of augmentation) corresponds to embeddings being uniformly distributed on a hypersphere.

Batch Resample is a simple strategy for long-tailed recognition and generation (Zhang et al., 2024). However, it may lead to mode collapse due to repetitive images for tail classes. Since our unsupervised contrastive loss relies on pairs of images, batch resampling increases the chance of images from the same tail class appearing in the same batch, which further diversify tail class images. We choose batch resampling as an optional strategy for our framework and provide an ablation study in Sec. [4.3](https://arxiv.org/html/2507.09052v2#S4.SS3).

### 3.2 Conditional-unconditional Alignment Loss

Conditinal-unconditinal Alignment Facilitates Knowledge Sharing for Class-imbalanced GAN We take inspiration from UTLO (Khorram et al., 2024), which addresses long-tailed generation with GANs, and observes that the similarity between head class and tail class images increases at lower resolution representations.
To share knowledge between the head class and tail class, UTLO (Khorram et al., 2024) proposes unconditional GAN objectives for lower-resolution representations and conditional GAN objectives for subsequent higher-resolution images.
More specifically, the lower-resolution representations becomes less sensitive to class labels, and the second part of the generator generates the fine details of images.
Utilizing unconditional generation for lower resolution is effective in increasing the diversity and quality of tail class images.
In a different setting of training GAN with limited data, Transitional-GAN (Shahbazi et al., 2022) leverages unconditional generation in a two-staged training scheme, as unconditional generation can give a better FID than conditional generation when data is limited.

Conditional-unconditinal Alignment for Class-imbalanced Diffusion Model We aim to adapt the conditional-unconditional alignment approach used in GANs (Khorram et al., 2024; Shahbazi et al., 2022) to diffusion models.
Our key insight is that the diffusion model denoises images recursively, and the similarity between head class images and tail class images is higher during the initial timesteps of the denoising process.
This is similar but different from UTLO (Khorram et al., 2024), which proposes unconditional generation for lower-resolution representations.
We visualize the reverse processing of unconditional generation and conditional generation with different class labels of the original DDPM, starting from the same Gaussian noise, as shown in Fig. [3](https://arxiv.org/html/2507.09052v2#S3.F3) and more examples in appendix [A.3.4](https://arxiv.org/html/2507.09052v2#A1.SS3.SSS4).
Images from unconditional generation and conditional generation share similar low-frequency components, e.g., the coarse shape of synthesized objects and background region, especially at the initial time steps (Si et al., 2024).

Motivated by this observation, we propose to match unconditional generation and conditional generation at the initial time steps, to enable knowledge sharing between tail classes and head classes with abundant data.
Another motivation is that class conditions are not necessary for all timesteps. In fact, using adaptive guidance weight with a small weight in the initial steps can yield a better FID compared to using constant guidance weight (Wang et al., 2024).

From Conditional-unconditional Distribution Matching to Alignment Loss We penalize KL divergence between the conditional distribution $p_{\theta}(x_{t-1}|x_{t}^{i},c^{i})$ and unconditional distribution $p_{\theta}(x_{t-1}|x_{t}^{i})$:

$$ $\displaystyle\mathcal{L}^{i,t}_{al}=D_{KL}[p_{\theta}(x_{t-1}|x_{t}^{i},c^{i})||p_{\theta}(x_{t-1}|x_{t}^{i})].$ (2) $$

Suppose $p_{\theta}(x_{t-1}|x_{t},c)=\mathcal{N}(x_{t-1};\mu_{\theta}(x_{t},t,c),\sigma_{t}^{2}I),p_{\theta}(x_{x-t}|x_{t})=\mathcal{N}(x_{t-1};\mu_{\theta}(x_{t},t),\sigma_{t}^{2}I)$, then we have:

$$ $\displaystyle\mathcal{L}^{i,t}_{al}=\mathbb{E}[\frac{1}{2\sigma_{t}^{2}}||\mu_{\theta}(x_{t}^{i},t,c^{i})-\mu_{\theta}(x_{t}^{i},t)||^{2}]+C,$ (3) $$

where C is constant, and

$$ $\displaystyle\mu_{\theta}(x_{t}^{i},t,c^{i})$ $\displaystyle=\frac{1}{\sqrt{\alpha_{t}}}x_{t}^{i}-\frac{\beta_{t}}{\sqrt{\alpha_{t}}\sqrt{1-\bar{\alpha}_{t}}}\epsilon_{\theta}(x_{t}^{i},t,c^{i}),\quad\mu_{\theta}(x_{t}^{i},t)=\frac{1}{\sqrt{\alpha_{t}}}x_{t}^{i}-\frac{\beta_{t}}{\sqrt{\alpha_{t}}\sqrt{1-\bar{\alpha}_{t}}}\epsilon_{\theta}(x_{t}^{i},t),$ (4) $$

where $\alpha_{t}=1-\beta_{t},\bar{\alpha}_{t}\!\!=\!\!\prod_{i=1}^{t}(1-\beta_{i})$, $\{\beta_{t}\}_{1:T}$ is the variance schedule.

Then, $\mathcal{L}_{al}$ simplifies to:

$$ $\displaystyle\mathcal{L}^{i,t}_{al}\propto||\epsilon_{\theta}(x_{t}^{i},t,c^{i})-\epsilon_{\theta}(x_{t}^{i},t)||^{2}.$ (5) $$

As shown in Fig. [2](https://arxiv.org/html/2507.09052v2#S3.F2), the anchor vector $h_{anc}^{i}$ is fed into decode network $d(*)$ to get unconditional noise estimation $\epsilon_{unc}^{i}:=\epsilon_{\theta}(x_{t}^{i};t,\emptyset)$.
Meanwhile, the anchor image $x_{t}^{i}$ is also incorporated with the class label condition $\mathbf{c^{i}}$ fed into the network to get the conditional noise estimation $\epsilon_{cond}^{i}:=\epsilon_{\theta}(x_{t}^{i};t,\mathbf{c^{i}})$.
The alignment loss is defined between the unconditional and conditional noise estimation:

$$ $\displaystyle\mathcal{L}_{al}=\frac{1}{|B|}\sum_{i\in B}{\mathbb{E}_{t,x_{0}}[\frac{t}{T}\|\epsilon_{\theta}(x_{t}^{i};t,\mathbf{c^{i}})-\epsilon_{\theta}(x_{t}^{i};t,\emptyset)\|^{2}}].$ (6) $$

The loss is weighted linearly by timesteps $t$, so that the initial steps with large $t$ are weighted more.
In other words, we align conditional generation and unconditional generation for the initial steps.

Figure: Figure 3: The leftmost column shows synthetic images with different classes but with the same initial noise or random seed. We visualize the low-frequency component and the high-frequency component from the reverse processing. It shows that low-frequency components are similar during initial time steps for different classes with the same initial noise (Si et al., 2024). More examples are in appendix [A.3.4](https://arxiv.org/html/2507.09052v2#A1.SS3.SSS4).
Refer to caption: x2.png

Figure: Algorithm 1 Training algorithm of CCUA.

### 3.3 Overall Framework

Our final loss function $\mathcal{L}$ is the sum of the standard DDPM (Ho et al., 2020) loss $\mathcal{L}_{ddpm}$ and our contrastive conditional-unconditional alignment loss $\mathcal{L}_{ccua}$:

$$ $\displaystyle\mathcal{L}_{ccua}=\alpha\cdot\mathcal{L}_{ucl}+\gamma\cdot\mathcal{L}_{al},$ (7) $$

where
$\alpha$ and $\gamma$ are the hyper-parameters for our unsupervised contrastive loss and alignment loss.

The overall algorithm framework is shown as Algorithm [1](https://arxiv.org/html/2507.09052v2#alg1).
In classifier-free guidance (CFG), the class labels $\mathbf{c}$ are randomly dropped by a specific probability $p_{uncond}$.
In our setting, we keep the $p_{uncond}=10\%$, same as (Qin et al., 2023).
The parameters of the noise estimation network $\epsilon_{\theta}$ are updated by the gradient of the final loss $\nabla_{\theta}\mathcal{L}$ as in Eq. [7](https://arxiv.org/html/2507.09052v2#S3.E7).

## 4 Experiments

### 4.1 Experimental Setup

We report results on multiple datasets, including the long-tailed versions of ImageNet, TinyImageNet, and Places datasets, namely ImageNet-LT, TinyImageNet-LT, and Places-LT, which are commonly used datasets in long-tailed image generation and recognition. We also report results and conduct more analysis on CIFAR10-LT/CIFAR100-LT datasets, which are shown in the appendix [A.2](https://arxiv.org/html/2507.09052v2#A1.SS2).
We measure IS Score, FID, spatial FID (Dhariwal & Nichol, 2021), Kernel Inception Distance (KID) (Bińkowski et al., 2018) as the overall quality metrics.
We also report FID of tail classes (FIDtail), inclined to measure quality for tail classes.
Precision and Recall are reported to respectively measure fidelity and diversity, which are widely used in evaluating generative models on long-tailed scenes (Qin et al., 2023; Zhang et al., 2024; Khorram et al., 2024).
More implementation details are provided in the appendix [A.1](https://arxiv.org/html/2507.09052v2#A1.SS1).
We compare our proposed methods with multiple baselines, covering diffusion transformer and unet-based architecture.
For diffusion transformer, we compare CCUA to SiT (Ma et al., 2024) and Dispersive Loss (Wang & He, 2025).
For unet-based architecture, we compare to the original DDPM (Ho et al., 2020), CBDM (Qin et al., 2023), and OCLT (Zhang et al., 2024).

### 4.2 Quantitative Results

**Table 1: Comparison on ImageNet-LT $256\times 256$ with SiT pipeline. We use Blue ‘()’ to highlight the improvement of our method over SiT baseline on FID and Recall, which denotes the overall quality and diversity, respectively.**
| Epochs | Method | IS $\uparrow$ | FID $\downarrow$ | sFID $\downarrow$ | Prec. (%) $\uparrow$ | Recall (%) $\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- |
| 40 | SiT | 53.93 | 33.87 | 22.66 | 54.56 | 19.17 |
| +Dispersive Loss | 53.87 | 34.00 | 22.68 | 54.87 | 19.84 |  |
| +CCUA (ours) | 73.60 | 25.89 (-7.98) | 16.54 | 58.62 | 23.05 (+4.08) |  |
| 80 | SiT | 78.88 | 25.72 | 21.95 | 64.35 | 18.58 |
| +Dispersive Loss | 83.80 | 25.22 | 21.76 | 65.50 | 17.30 |  |
| +CCUA (ours) | 103.66 | 19.99 (-5.73) | 14.88 | 65.80 | 21.31 (+2.73) |  |
| 120 | SiT | 103.11 | 21.20 | 20.17 | 69.34 | 18.25 |
| +Dispersive Loss | 104.02 | 21.32 | 20.46 | 68.07 | 18.56 |  |
| +CCUA (ours) | 119.11 | 17.55 (-3.65) | 13.89 | 68.38 | 21.20 (+2.95) |  |
| 160 | SiT | 111.70 | 19.95 | 20.12 | 70.34 | 18.61 |
| +Dispersive Loss | 115.64 | 19.76 | 20.18 | 69.98 | 18.96 |  |
| +CCUA (ours) | 124.87 | 16.41 (-3.54) | 13.17 | 69.84 | 21.34 (+2.73) |  |

Class-imbalanced Generation for Diffusion Transformer We apply the proposed CCUA loss into the standard diffusion transformer training pipeline, SiT (Ma et al., 2024), on ImageNet-LT dataset, and compare it to Dispersive Loss (Wang & He, 2025), a concurrent work of applying contrastive learning for improving SiT on balanced dataset.
All models are sampled with ODE flow 50 steps for conditional generation with CFG strength 7.5.
As shown in Table [1](https://arxiv.org/html/2507.09052v2#S4.T1), our method achieves remarkable improvement compared to SiT and Dispersive Loss on various training epochs.
Our method achieves above 10% improvement on Recall, while keep the similar Precision as SiT, illustrating the higher diversity of the proposed CCUA loss without trading off fidelity.

**Table 2: Comparison on TinyImageNet-LT $64\times 64$ and Places-LT $64\times 64$ with DDPM pipeline. Blue ‘()’ shows improvement of our method over DDPM baseline. Green ‘()’ shows improvement of DDPM trained on balanced version over trained on long-tailed version.**
| Dataset | Method | FID$\downarrow$ | FID${}_{tail}\downarrow$ | KID${}_{\times 1k}\downarrow$ |
| --- | --- | --- | --- | --- |
| TinyImageNet-LT | DDPM${}^{*}_{bal}$ (Ho et al., 2020) | 15.73 (-2.94) | 25.62 (-12.26) | 3.19 (-3.12) |
| DDPM (Ho et al., 2020) | 18.67 | 40.12 | 6.31 |  |
| CBDM (Qin et al., 2023) | 20.90 | 48.07 | 6.55 |  |
| OCLT (Zhang et al., 2024) | 17.72 | 39.67 | 5.61 |  |
| CCUA (ours) | 15.24 (-3.43) | 30.39 (-9.73) | 3.83 (-2.48) |  |
| Places-LT | DDPM (Ho et al., 2020) | 13.89 | 23.74 | 5.26 |
| CBDM (Qin et al., 2023) | 15.15 | 26.07 | 5.64 |  |
| OCLT (Zhang et al., 2024) | 13.04 | 22.84 | 4.17 |  |
| CCUA (ours) | 11.99 (-1.90) | 20.84 (-2.90) | 3.63 (-1.63) |  |

Class-imbalanced Generation for DDPM For training unet-based architecture DDPM on TinyImageNet-LT and Places-LT datasets, our method also achieves the best performance, as shown in Tab. [2](https://arxiv.org/html/2507.09052v2#S4.T2).
All models are sampled with DDIM 100 steps for conditional generation with optimal CFG strength.
For TinyImageNet-LT datasets, we provide the metrics of the DDPM model trained on the balanced version, i.e., the original TinyImageNet datasets, denoted by DDPM${}^{*}_{bal}$, as the theoretical optimal reference.
The best FID and KID score of our method illustrates the high quality of images synthesized by our method(^1^11The FID score of our reproduced DDPM, CBDM (Qin et al., 2023), and OCLT (Zhang et al., 2024) are better than the numbers reported in (Qin et al., 2023; Zhang et al., 2024). We train all methods from scratch without fine-tuning and find the optimal guidance strength $\omega$ for each method. Such implementation is different from the one used in (Qin et al., 2023; Zhang et al., 2024) but is considered more fair.).

**Table 3: FID score for three ‘super-categories’: ‘Head’, ‘Body’, and ‘Tail’. All models are measured with DDIM (Song et al., 2020) 100 sampling steps for conditional generation with CFG.**
| Dataset | Method<br>$P_{category}$<br>FID$\downarrow$ | Head<br>$\sim 80\%$ | Body<br>$\sim 17\%$ | Tail<br>$\sim 3\%$ | All |
| --- | --- | --- | --- | --- | --- |
| TinyImageNet-LT | DDPM (Ho et al., 2020) | 21.27 | 33.25 | 40.12 | 18.67 |
| CBDM (Qin et al., 2023) | 24.12 | 35.02 | 48.07 | 20.90 |  |
| OCLT (Zhang et al., 2024) | 20.64 | 30.62 | 39.67 | 17.72 |  |
| CCUA (ours) | 21.32 | 27.98 | 30.39 | 15.24 |  |
| Places-LT | DDPM (Ho et al., 2020) | 19.26 | 20.31 | 23.74 | 13.89 |
| CBDM (Qin et al., 2023) | 21.26 | 21.68 | 26.07 | 15.15 |  |
| OCLT (Zhang et al., 2024) | 19.16 | 19.43 | 22.84 | 13.04 |  |
| CCUA (ours) | 18.15 | 19.27 | 20.84 | 11.99 |  |

Results Across Various Category Intervals To see the effectiveness of our method on tail classes, we divide the Places-LT and TinyImageNet-LT datasets into three super categories: ‘Head’, ‘Body’, and ‘Tail’, with classes sorted by the number of training images.
For each dataset, the top 33% classes were allocated to the ‘Head’ category, the next 34% classes to the ‘Body’ category, and the last 33% classes to the ‘Tail’ category.
The percentage $P_{category}$ of training images belonging to each category is shown in the first row of Tab. [3](https://arxiv.org/html/2507.09052v2#S4.T3).
Our method achieves the best FID score across all three categories on Places-LT and TinyImageNet-LT datasets.
For the ‘Tail’ classes, which are only $\sim 3\%$, our method improves their FID from 23.74 to 20.84 on Places-LT and 40.12 to 30.39 on TinyImageNet-LT, respectively.
Additionally, both ‘Head’ and ‘Body’ classes show improved FID scores with our method, as demonstrated in Tab. [3](https://arxiv.org/html/2507.09052v2#S4.T3).

### 4.3 Qualitative Results and Ablation Study

Visualization of Synthetic Images Fig. [4](https://arxiv.org/html/2507.09052v2#S4.F4) shows synthetic images of SiT and with CCUA for ImageNet-LT tail classes ‘espresso’ and ‘window shade’.
Compared to SiT, CCUA achieves visually higher fidelity and diversity.
More qualitative results on ImageNet-LT, TinyImageNet-LT, and CIFAR100-LT are shown in appendix [A.3.1](https://arxiv.org/html/2507.09052v2#A1.SS3.SSS1) and [A.3.3](https://arxiv.org/html/2507.09052v2#A1.SS3.SSS3).

Figure: (a) Synthetic images from SiT (Ma et al., 2024).
Refer to caption: figures_iclr/sit_coffee_windowshade.png

**Table 4: Loss Function Analysis.**
| Loss Function | FID$\downarrow$ |
| --- | --- |
| $\mathcal{L}_{ddpm}$ | 18.67 |
| $\mathcal{L}_{ddpm}$ w/ $\mathcal{L}_{ucl}$ (Eq. [1](https://arxiv.org/html/2507.09052v2#S3.E1)) | 18.05 |
| $\mathcal{L}_{ddpm}$ w/ $\mathcal{L}_{al}$ (Eq. [6](https://arxiv.org/html/2507.09052v2#S3.E6)) | 17.84 |
| $\mathcal{L}_{ddpm}$ w/ $\mathcal{L}_{ccua}$ (Eq. [7](https://arxiv.org/html/2507.09052v2#S3.E7)) | 17.16 |

Ablation Study of Loss Function Table [4](https://arxiv.org/html/2507.09052v2#S4.T4) shows an ablation study of the proposed losses including $\mathcal{L}_{ucl}$ and $\mathcal{L}_{al}$, which is conducted on TinyImageNet-LT.
For the loss function ablation study, we disable batch resample strategy which is separately ablated.
Our method achieves lower FID when applying only the $\mathcal{L}_{ucl}$ loss (Eq. [1](https://arxiv.org/html/2507.09052v2#S3.E1)) or only $\mathcal{L}_{al}$ loss (Eq. [6](https://arxiv.org/html/2507.09052v2#S3.E6)) compared to the standard $\mathcal{L}_{ddpm}$ loss only.
The best result is achieved with both losses, i.e., $\mathcal{L}_{ccua}$ (Eq. [7](https://arxiv.org/html/2507.09052v2#S3.E7)).

Ablation Study of Batch Resample Strategy As shown in Table [5](https://arxiv.org/html/2507.09052v2#S4.T5), we measure the performance of SiT and CCUA, with or without applying batch resampling strategy.
Compared to SiT, the proposed CCUA loss achieves better IS, FID, sFID and Recall with only 0.09 Precision’s decline, illustrates our method on improving diversity without trading off fidelity.
With applying batch resample strategy, the CCUA further improves all metrics.
Notes that the CCUA increases Recall from 20.66 to 23.05, while SiT suffers a decline in Recall from 19.17 to 17.26 when applying batch resampling.
As we discuss in Section. [3.1](https://arxiv.org/html/2507.09052v2#S3.SS1), the $\mathcal{L}_{ucl}$ could benefit from batch resample strategy due to the increasing instances of tail classes, leading to more diverse distribution in the latent space.

**Table 5: Batch Resample Strategy Analysis on ImageNet-LT. All models are trained from scratch to 40 epochs. Blue ‘()’ highlights the improvement of each method compared to the SiT baseline, while Red ‘()’ highlights the decline compared to SiT baseline.**
| Method | IS $\uparrow$ | FID $\downarrow$ | sFID $\downarrow$ | Prec. (%) $\uparrow$ | Recall (%) $\uparrow$ |
| --- | --- | --- | --- | --- | --- |
| SiT (baseline) | 53.93 | 33.87 | 22.66 | 54.56 | 19.17 |
| +Batch Resample | 71.39 (+17.46) | 28.05 (-5.82) | 21.57 (-1.09) | 57.61 (+3.05) | 17.26 (-1.91) |
| $\mathcal{L}_{ccua}$ (Eq. [7](https://arxiv.org/html/2507.09052v2#S3.E7)) | 57.79 (+3.86) | 32.25 (-1.62) | 20.71 (-1.95) | 54.47 (-0.09) | 20.66 (+1.49) |
| +Batch Resample | 73.60 (+19.67) | 25.89 (-7.98) | 16.54 (-6.12) | 58.62 (+4.06) | 23.05 (+3.88) |

## 5 Conclusion

Real-world data for training image generation models often exhibit long-tailed distributions.
Similar to class-imbalanced GANs (Khorram et al., 2024), class-imbalanced diffusion models generates inferior tail class images due to data scarcity.
We propose a framework with two losses.
Firstly, we employ an unsupervised contrastive loss with negative samples only to contrast the latents of different synthetic images at every denoising step, promoting intra-class diversity in particular for tail classes.
Our second loss aligns class-conditional generation with unconditional generation for large timesteps.
This encourages the initial denoising steps to be class-agnostic, thereby enriching tail classes through knowledge sharing from head classes–a principle demonstrated to enhance long-tailed GAN performance (Shahbazi et al., 2022; Khorram et al., 2024), which we successfully adapt to diffusion models.
With the two losses, our framework of contrastive conditional-unconditional alignment boosts the performane of DDPM (Ho et al., 2020) and SiT (Ma et al., 2024) for long-tailed image generation and outperforms alternative methods including CBDM (Qin et al., 2023), OCLT (Zhang et al., 2024), and Dispersive Loss (Wang & He, 2025).
Extensive experiments on multiple datasets in particular ImageNet-LT with 256x256 resolution demonstrated the effectiveness of our method on both UNet-based architecture and Diffusion Transformer.

## References

- Ai et al. (2023)
Qingzhong Ai, Pengyun Wang, Lirong He, Liangjian Wen, Lujia Pan, and Zenglin Xu.
Generative oversampling for imbalanced data via majority-guided vae.
In *International Conference on Artificial Intelligence and Statistics*, pp. 3315–3330. PMLR, 2023.
- Baranchuk et al. (2021)
Dmitry Baranchuk, Ivan Rubachev, Andrey Voynov, Valentin Khrulkov, and Artem Babenko.
Label-efficient semantic segmentation with diffusion models.
*arXiv preprint arXiv:2112.03126*, 2021.
- Bińkowski et al. (2018)
Mikołaj Bińkowski, Danica J Sutherland, Michael Arbel, and Arthur Gretton.
Demystifying mmd gans.
*arXiv preprint arXiv:1801.01401*, 2018.
- Brock et al. (2019)
Andrew Brock, Jeff Donahue, and Karen Simonyan.
Large scale gan training for high fidelity natural image synthesis.
In *ICLR*, 2019.
- Cao et al. (2019)
Kaidi Cao, Colin Wei, Adrien Gaidon, Nikos Arechiga, and Tengyu Ma.
Learning imbalanced datasets with label-distribution-aware margin loss.
*Advances in neural information processing systems*, 32, 2019.
- Chen et al. (2020)
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton.
A simple framework for contrastive learning of visual representations.
2020.
- Couairon et al. (2022)
Guillaume Couairon, Jakob Verbeek, Holger Schwenk, and Matthieu Cord.
Diffedit: Diffusion-based semantic image editing with mask guidance.
*arXiv preprint arXiv:2210.11427*, 2022.
- Dalva & Yanardag (2024)
Yusuf Dalva and Pinar Yanardag.
Noiseclr: A contrastive learning approach for unsupervised discovery of interpretable directions in diffusion models.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp. 24209–24218, 2024.
- Dhariwal & Nichol (2021)
Prafulla Dhariwal and Alexander Nichol.
Diffusion models beat gans on image synthesis.
*Advances in neural information processing systems*, 34:8780–8794, 2021.
- Fuest et al. (2024)
Michael Fuest, Pingchuan Ma, Ming Gui, Johannes S Fischer, Vincent Tao Hu, and Bjorn Ommer.
Diffusion models and representation learning: A survey.
*arXiv preprint arXiv:2407.00783*, 2024.
- He et al. (2020)
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick.
Momentum contrast for unsupervised visual representation learning.
In *CVPR*, 2020.
- Ho & Salimans (2022)
Jonathan Ho and Tim Salimans.
Classifier-free diffusion guidance.
*arXiv preprint arXiv:2207.12598*, 2022.
- Ho et al. (2020)
Jonathan Ho, Ajay Jain, and Pieter Abbeel.
Denoising diffusion probabilistic models.
*Advances in neural information processing systems*, 33:6840–6851, 2020.
- Ho et al. (2022a)
Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, et al.
Imagen video: High definition video generation with diffusion models.
*arXiv preprint arXiv:2210.02303*, 2022a.
- Ho et al. (2022b)
Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet.
Video diffusion models.
*Advances in Neural Information Processing Systems*, 35:8633–8646, 2022b.
- Jiang et al. (2021)
Ziyu Jiang, Tianlong Chen, Bobak J Mortazavi, and Zhangyang Wang.
Self-damaging contrastive learning.
2021.
- Karras et al. (2019)
Tero Karras, Samuli Laine, and Timo Aila.
A style-based generator architecture for generative adversarial networks.
In *CVPR*, 2019.
- Karras et al. (2020)
Tero Karras, Miika Aittala, Janne Hellsten, Samuli Laine, Jaakko Lehtinen, and Timo Aila.
Training generative adversarial networks with limited data.
*Advances in neural information processing systems*, 33:12104–12114, 2020.
- Kawar et al. (2023)
Bahjat Kawar, Shiran Zada, Oran Lang, Omer Tov, Huiwen Chang, Tali Dekel, Inbar Mosseri, and Michal Irani.
Imagic: Text-based real image editing with diffusion models.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp. 6007–6017, 2023.
- Khorram et al. (2024)
Saeed Khorram, Mingqi Jiang, Mohamad Shahbazi, Mohamad H Danesh, and Li Fuxin.
Taming the tail in class-conditional gans: Knowledge sharing via unconditional training at lower resolutions.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp. 7580–7590, 2024.
- Khosla et al. (2020)
Prannay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishnan.
Supervised contrastive learning.
*Advances in neural information processing systems*, 33:18661–18673, 2020.
- Kingma et al. (2019)
Diederik P Kingma, Max Welling, et al.
An introduction to variational autoencoders.
*Foundations and Trends® in Machine Learning*, 12(4):307–392, 2019.
- Li et al. (2022)
Tianhong Li, Peng Cao, Yuan Yuan, Lijie Fan, Yuzhe Yang, Rogerio S Feris, Piotr Indyk, and Dina Katabi.
Targeted supervised contrastive learning for long-tailed recognition.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp. 6918–6928, 2022.
- Liu et al. (2022)
Ziwei Liu, Zhongqi Miao, Xiaohang Zhan, Jiayun Wang, Boqing Gong, and X Yu Stella.
Open long-tailed recognition in a dynamic world.
*IEEE Transactions on Pattern Analysis and Machine Intelligence*, 2022.
- Ma et al. (2024)
Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden-Eijnden, and Saining Xie.
Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers.
In *European Conference on Computer Vision*, pp. 23–40. Springer, 2024.
- Ouyang et al. (2023)
Yidong Ouyang, Liyan Xie, and Guang Cheng.
Improving adversarial robustness through the contrastive-guided diffusion process.
In *International Conference on Machine Learning*, pp. 26699–26723. PMLR, 2023.
- Poole et al. (2022)
Ben Poole, Ajay Jain, Jonathan T. Barron, and Ben Mildenhall.
Dreamfusion: Text-to-3d using 2d diffusion.
*arXiv*, 2022.
- Qin et al. (2023)
Yiming Qin, Huangjie Zheng, Jiangchao Yao, Mingyuan Zhou, and Ya Zhang.
Class-balancing diffusion models.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp. 18434–18443, 2023.
- Rangwani et al. (2022)
Harsh Rangwani, Naman Jaswani, Tejan Karmali, Varun Jampani, and R Venkatesh Babu.
Improving gans for long-tailed data through group spectral regularization.
In *European Conference on Computer Vision*, pp. 426–442. Springer, 2022.
- Rangwani et al. (2023)
Harsh Rangwani, Lavish Bansal, Kartik Sharma, Tejan Karmali, Varun Jampani, and R Venkatesh Babu.
Noisytwins: Class-consistent and diverse image generation through stylegans.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp. 5987–5996, 2023.
- Rombach et al. (2022)
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.
High-resolution image synthesis with latent diffusion models.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pp. 10684–10695, 2022.
- Ruiz et al. (2023)
Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman.
Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation.
In *Proceedings of the IEEE/CVF conference on computer vision and pattern recognition*, pp. 22500–22510, 2023.
- Shahbazi et al. (2022)
Mohamad Shahbazi, Martin Danelljan, Danda Pani Paudel, and Luc Van Gool.
Collapse by conditioning: Training class-conditional GANs with limited data.
In *International Conference on Learning Representations*, 2022.
URL [https://openreview.net/forum?id=7TZeCsNOUB_](https://openreview.net/forum?id=7TZeCsNOUB_).
- Si et al. (2024)
Chenyang Si, Ziqi Huang, Yuming Jiang, and Ziwei Liu.
Freeu: Free lunch in diffusion u-net.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pp. 4733–4743, 2024.
- Song et al. (2020)
Jiaming Song, Chenlin Meng, and Stefano Ermon.
Denoising diffusion implicit models.
*arXiv preprint arXiv:2010.02502*, 2020.
- Szegedy et al. (2016)
Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna.
Rethinking the inception architecture for computer vision.
In *Proceedings of the IEEE conference on computer vision and pattern recognition*, pp. 2818–2826, 2016.
- Tan et al. (2020)
Shuhan Tan, Yujun Shen, and Bolei Zhou.
Improving the fairness of deep generative models without retraining.
*arXiv preprint arXiv:2012.04842*, 2020.
- Van der Maaten & Hinton (2008)
Laurens Van der Maaten and Geoffrey Hinton.
Visualizing data using t-sne.
*Journal of machine learning research*, 9(11), 2008.
- Wang & He (2025)
Runqian Wang and Kaiming He.
Diffuse and disperse: Image generation with representation regularization.
*arXiv preprint arXiv:2506.09027*, June 2025.
- Wang et al. (2024)
Xi Wang, Nicolas Dufour, Nefeli Andreou, Marie-Paule Cani, Victoria Fernández Abrevaya, David Picard, and Vicky Kalogeiton.
Analysis of classifier-free guidance weight schedulers.
*arXiv preprint arXiv:2404.13040*, 2024.
- Yan et al. (2024)
Divin Yan, Lu Qi, Vincent Tao Hu, Ming-Hsuan Yang, and Meng Tang.
Training class-imbalanced diffusion model via overlap optimization.
*arXiv preprint arXiv:2402.10821*, 2024.
- Yu et al. (2024)
Sihyun Yu, Sangkyung Kwak, Huiwon Jang, Jongheon Jeong, Jonathan Huang, Jinwoo Shin, and Saining Xie.
Representation alignment for generation: Training diffusion transformers is easier than you think.
*arXiv preprint arXiv:2410.06940*, 2024.
- Zhang et al. (2023)
Lvmin Zhang, Anyi Rao, and Maneesh Agrawala.
Adding conditional control to text-to-image diffusion models.
In *Proceedings of the IEEE/CVF International Conference on Computer Vision*, pp. 3836–3847, 2023.
- Zhang et al. (2024)
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.

## Appendix A Technical Appendices and Supplementary Material

Figure: (a) w/ UNet-based Model
Refer to caption: x3.png

**Table 6: Our method outperforms other baselines on all datasets. We also provide the results of DDPM trained on balanced datasets, which show the upper bound of performance. All models are measured with DDIM (Song et al., 2020) 100 steps for conditional generation with CFG. Blue ‘()’ shows improvement of our method over DDPM baseline (Ho et al., 2020). Green ‘()’ shows improvement of DDPM trained on balanced version over trained on long-tailed version.**
| Dataset | Method | FID$\downarrow$ | FID${}_{tail}\downarrow$ | KID${}_{\times 1k}\downarrow$ |
| --- | --- | --- | --- | --- |
| CIFAR10-LT | DDPM${}^{*}_{bal}$ (Ho et al., 2020) | 4.90 (-1.03) | 6.27 (-5.98) | 1.32 (-0.32) |
| DDPM (Ho et al., 2020) | 5.93 | 12.25 | 1.64 |  |
| CBDM (Qin et al., 2023) | 5.81 | 10.01 | 1.58 |  |
| OCLT (Zhang et al., 2024) | 6.10 | 11.13 | 1.58 |  |
| CCUA (ours) | 5.56 (-0.37) | 10.03 (-2.22) | 1.27 (-0.37) |  |
| CIFAR100-LT | DDPM${}^{*}_{bal}$ (Ho et al., 2020) | 5.15 (-1.80) | 8.97 (-8.48) | 1.05 (-0.66) |
| DDPM (Ho et al., 2020) | 6.95 | 17.45 | 1.71 |  |
| CBDM (Qin et al., 2023) | 6.50 | 17.36 | 1.41 |  |
| OCLT (Zhang et al., 2024) | 6.45 | 17.22 | 1.42 |  |
| CCUA (ours) | 6.24 (-0.71) | 16.35 (-1.10) | 1.36 (-0.35) |  |

### A.1 Implementation Details

Dataset Details We keep the original $32\times 32$ resolution for CIFAR10-LT/CIFAR100-LT, and resize images to $64\times 64$ for TinyImageNet-LT and Places-LT, while $256\times 256$ for ImageNet-LT.
Same as (Qin et al., 2023; Yan et al., 2024; Zhang et al., 2024), we adopt the same methodology presented in (Cao et al., 2019) to construct long-tailed version datasets with an imbalance factor of $0.01$.

Model Architecture Details As described in Sec. [3](https://arxiv.org/html/2507.09052v2#S3), the proposed CCUA framework can be applied into UNet-based architecture and Diffusion Transformer.
In Fig. [5](https://arxiv.org/html/2507.09052v2#A1.F5), we display our UNet-based model and Diffusion Transformer in details.
As shown in Fig. [2](https://arxiv.org/html/2507.09052v2#S3.F2), the noise estimation network is divided into two parts, a latent encoded network $e(*)$ and a decoded network $d(*)$.
In our setting, for UNet-based model, $e(*)$ is defined as the UNet encoder, while $d(*)$ is defined as the UNet decoder, as shown in Fig. [5](https://arxiv.org/html/2507.09052v2#A1.F5) (a).
For Diffusion Transformer, $e(*)$ is defined as all the SiT/DiT blocks, while $d(*)$ is defined as the final linear and adaLN projection layer, as shown in Fig. [5](https://arxiv.org/html/2507.09052v2#A1.F5) (b).

Training Details For SiT pipeline, we strictly follow the same training hyper-parameter settings as Dispersive Loss (Wang & He, 2025).
We use 4 A6000 GPUs to train all SiT based models on ImageNet-LT datasets with batch size $48$.
During the training process, all methods are trained from scratch. We report 40 epochs, 80 epochs, 120 epochs, and 160 epochs results in Table [1](https://arxiv.org/html/2507.09052v2#S4.T1).
For DDPM pipeline, we follow the same training configurations of the baseline models (Ho et al., 2020; Qin et al., 2023; Zhang et al., 2024). We use one RTX 4090 GPU to train each model on CIFAR10-LT/CIFAR100-LT datasets with batch size $64$ while using 2 A100 GPUs for TinyImageNet-LT and Places-LT, with batch size $128$.
During the training process, all methods are trained from scratch for 200k iterations on CIFAR10-LT/CIFAR100-LT datasets, 100k iterations on TinyImageNet-LT and Places-LT datasets, while they follow the classifier-free guidance (CFG) algorithm (Ho & Salimans, 2022), which randomly drops labels with a probability of 10%.
On ImageNet-LT and TinyImageNet-LT datasets, we apply batch resample strategy with the re-balanced factor $0.1$, while we don’t apply batch resample strategy on Places-LT and CIFAR10-LT/CIFAR100-LT datasets.
For TinyImageNet-LT and Places-LT datasets, we also apply timestep adaptive weight $t/T$ to unsupervised contrastive loss since we find the unsupervised contrastive loss could also benefit from such an adaptive weight, like alignment loss.

Evaluation Details In Table [2](https://arxiv.org/html/2507.09052v2#S4.T2), FIDtail denotes the FID score for synthetic images of the last 30% classes in each dataset.
Specifically, we categorize the ‘Tail’ classes as the last 66 classes for TinyImageLT (200 classes), the last 121 classes for Places-LT (365 classes), respectively.
In Table [1](https://arxiv.org/html/2507.09052v2#S4.T1), the evaluation metrics are based on 50k synthetic images generated by each method with a CFG strength 7.5.
In Table [2](https://arxiv.org/html/2507.09052v2#S4.T2), the evaluation metrics are based on 10k synthetic images generated by each method.
During inference, we perform a grid search algorithm for each method to determine the optimal guidance strength $\omega$ of CFG, ensuring each model achieves its best performance.

### A.2 More Quantitative Results

**Table 7: We compare our method with other baselines on CIFAR10-LT/CIFAR100-LT datasets, with DDPM 1000 steps for conditional generation with CFG. We also report the data augmentation method ADA (Karras et al., 2020) and $\omega$-scheduler (Wang et al., 2024), which are orthogonal to our method.**
| Method | CIFAR10-LT | CIFAR100-LT |  |  |
| --- | --- | --- | --- | --- |
| FID$\downarrow$ | IS$\uparrow$ | FID$\downarrow$ | IS$\uparrow$ |  |
| DDPM${}^{*}_{bal}$ (Ho et al., 2020) | 4.87 | 9.35 | 5.20 | 13.29 |
| DDPM (Ho et al., 2020) | 5.81 | 9.36 | 7.09 | 12.64 |
| +ADA (Karras et al., 2020) | - | - | 6.69 | 12.87 |
| +$\omega$-Scheduler (Wang et al., 2024) | 5.87 | 9.22 | 6.60 | 12.10 |
| CBDM (Qin et al., 2023) | 5.92 | 9.38 | 6.52 | 12.79 |
| OCLT (Zhang et al., 2024) | 5.69 | 9.42 | 6.23 | 13.18 |
| CCUA (ours) | 5.57 | 9.42 | 5.99 | 13.01 |

**Table 8: We compare our method with other baselines on CIFAR10-LT/CIFAR100-LT datasets with DDIM 10 sampling steps for conditional generation with CFG. Blue ‘()’ shows improvement of our method over DDPM baseline (Ho et al., 2020). Green ‘()’ shows improvement of DDPM trained on balanced version over trained on long-tailed version.**
| Dataset | Method | FID$\downarrow$ | FID${}_{tail}\downarrow$ | KID${}_{\times 1k}\downarrow$ |
| --- | --- | --- | --- | --- |
| CIFAR10-LT | DDPM${}^{*}_{bal}$ | 13.28 (-1.44) | 13.26 (-5.01) | 6.06 (-0.98) |
| DDPM | 14.72 | 18.27 | 7.04 |  |
| CBDM | 13.54 | 16.90 | 6.52 |  |
| OCLT | 15.48 | 20.73 | 7.39 |  |
| CCUA (ours) | 13.16 (-1.56) | 16.83 (-1.44) | 6.04 (-1.00) |  |
| CIFAR100-LT | DDPM${}^{*}_{bal}$ | 13.34 (-0.75) | 18.27 (-6.80) | 5.56 (-0.57) |
| DDPM | 14.09 | 25.07 | 6.13 |  |
| CBDM | 13.37 | 23.97 | 5.83 |  |
| OCLT | 13.70 | 24.48 | 5.73 |  |
| CCUA (ours) | 12.90 (-1.19) | 23.17 (-1.90) | 5.63 (-0.50) |  |

Class-imbalanced Generation on CIFAR10-LT/CIFAR100-LT datasets We conduct more experiments and analysis on CIFAR10-LT/CIFAR100-LT datasets, as shown in Table [6](https://arxiv.org/html/2507.09052v2#A1.T6) and Table [7](https://arxiv.org/html/2507.09052v2#A1.T7).
We provide the metrics of the DDPM model trained on the balanced version, i.e., the original CIFAR10/CIFAR100 datasets, denoted by DDPM${}^{*}_{bal}$, as the theoretical optimal reference.
On CIFAR10-LT/CIFAR100-LT, our method achieves the lowest FID and KID compared to baseline methods.
Note that the FID gap between DDPM and DDPM${}^{*}_{bal}$ is 1.03 on CIFAR10-LT and 1.8 on CIFAR100-LT, while our method improves FID 0.37 over 1.03 on CIFAR10-LT and 0.71 over 1.8 on CIFAR100-LT, achieving $>\textbf{35\%}$ performance improvement.
To investigate the consistency of such improvements, we compare our method and all baseline methods on CIFAR10-LT and CIFAR100-LT with DDPM 1000 steps. As shown in Tab. [7](https://arxiv.org/html/2507.09052v2#A1.T7), our method achieves consistent improvements of FID for full 1000 sampling steps.
We also compare to a widely used data augmentation technique, Adaptive Discriminator Adaption (ADA) (Karras et al., 2020) for generative models on the full DDPM 1000 sampling steps.
Besides, we apply the CFG guidance strength scheduler $\omega$-cos (Wang et al., 2024) on DDPM, which gradually increases the guidance strength during sampling time steps decreasing to force the model transfer from unconditional to conditional generation.

Consistent Improvement for Fewer Sampling Steps To investigate the model’s performance on extremely few sampling steps, We evaluate our method and all baselines for DDIM 10 steps with CFG conditional generation on CIFAR10-LT/CIFAR100-LT.
As shown in Tab. [8](https://arxiv.org/html/2507.09052v2#A1.T8), our method achieves the best FID, FIDtail and KID scores.
Our method achieves even better results than the theoretical optimal model DDPM${}^{*}_{bal}$ under such an extreme experimental setting.
Such an improvement demonstrates the effectiveness of our method for training long-tailed image generation diffusion model.

**Table 9: Unconditional generation w/ DDIM 100 steps.**
| Method | CIFAR10-LT | CIFAR100-LT |  |  |
| --- | --- | --- | --- | --- |
| FID$\downarrow$ | IS$\uparrow$ | FID$\downarrow$ | IS$\uparrow$ |  |
| DDPM | 27.52 | 6.65 | 18.53 | 8.68 |
| CBDM | 25.60 | 6.70 | 17.06 | 9.00 |
| OCLT | 31.38 | 6.34 | 18.97 | 8.73 |
| Ours | 24.16 | 6.80 | 15.97 | 9.23 |

Class-imbalanced Unconditional Generation We also evaluate all the models for class-imbalanced unconditional generation.
As shown in Tab. [9](https://arxiv.org/html/2507.09052v2#A1.T9), the proposed method reduces FID from 27.52 to 24.16 for CIFAR10-LT and from 18.53 to 15.97 for CIFAR100-LT.
Such an improvement in unconditional generation highlights the effectiveness of the proposed contrastive learning loss, particularly the unsupervised contrastive loss $\mathcal{L}_{ucl}$.

**Table 10: Statistical Analysis for random seed on CIFAR100-LT, sampling with DDIM 100 steps.**
| Seed | Method | FID$\downarrow$ | Method | FID$\downarrow$ |
| --- | --- | --- | --- | --- |
| 0 | DDPM | 7.02 | Ours | 6.26 |
| 42 | DDPM | 6.95 | Ours | 6.24 |
| 2025 | DDPM | 7.04 | Ours | 6.27 |

Statistical Significance Analysis To inspect the robustness of the proposed method, we conduct statistical significance analysis on the ddpm and our method, as shown in Tab. [10](https://arxiv.org/html/2507.09052v2#A1.T10).
With applying different random seed, the FID deviation of DDPM is about 0.04 while our method is about 0.01, illustrating the robustness of our method for different random seeds.

Limitation of our method Our MSE loss involves both conditional generation and unconditional generation, which requires two passes of the denosing network during training and increases training time. In practice, with applying our method, the training time for one epoch is 1.6x of that of DDPM, and 1.48x of that of SiT. However, our method doesn’t increase inference latency.

### A.3 More Qualitative Results

#### A.3.1 More Visualization of Synthetic Images on ImageNet-LT and TinyImageNet-LT

Fig. [6](https://arxiv.org/html/2507.09052v2#A1.F6) shows synthetic images of SiT and with the proposed CCUA for ImageNet-LT tail classes ‘bubble’, ‘redwine’, ‘comic book’ and ‘yawl’.
Fig. [7](https://arxiv.org/html/2507.09052v2#A1.F7) shows synthetic images of DDPM and with the proposed CCUA for TinyImageNet-LT tail classes ‘teapot’, ‘water tower’, ‘pretzel’, ‘mushroom’, ‘orange’, and ‘pizza’.
These methods start the denoising process from the same Gaussian noise at corresponding grid cells.
As shown in Fig. [6](https://arxiv.org/html/2507.09052v2#A1.F6) and Fig. [7](https://arxiv.org/html/2507.09052v2#A1.F7), synthetic images of CCUA show consistently higher diversity and fidelity compared to SiT.

Figure: (a) Synthetic images from SiT (Ma et al., 2024).
Refer to caption: figures_iclr/sit_vis_imagenet.png

Figure: (a) Synthetic images from DDPM (Ho et al., 2020).
Refer to caption: x5.png

Figure: (a) Images in low-dimensional embeddings.
Refer to caption: x7.png

#### A.3.2 Distribution of Synthetic Images for Tail Class

We visualize synthetic images in the feature space and their density function for our method and the original DDPM.
Specifically, we use the Inception-V3 (Szegedy et al., 2016) model to extract $2048$ dimensional features of each image, and then apply t-SNE (Van der Maaten & Hinton, 2008) to project these features into $2$ dimensions.
As shown in Fig. [8](https://arxiv.org/html/2507.09052v2#A1.F8), the original DDPM overfits and generates highly similar images, while synthetic images based on our method have more diversity.
We visualize the image distribution more specifically by using kernel density estimation in the bottom.
The grey region represents the distribution of real ‘worm’ images from the class-balanced CIFAR100 dataset, while red ‘x’ points denote all 5 ‘worm’ images from the class-imbalanced CIFAR100-LT dataset.
The blue region represents the distribution of images synthesized by our method, while the green region is for the original DDPM.
Areas with higher color saturation (dark, green, or blue) indicate regions of higher density, which correspond to modes of distribution.
Synthetic images from DDPM mostly concentrates around the training images, leading to mode collapse.
Synthetic images from our method spans the space enclosed by all training images, the distribution of which is shown in blue in (b).

#### A.3.3 Mode Collapse Issue on Tail Class

Fig. [9](https://arxiv.org/html/2507.09052v2#A1.F9) shows synthetic images of baseline methods and our method for CIFAR100-LT tail classes ‘rose’ and ‘table’.
All methods start the denoising process from the same Gaussian noise at corresponding grid cells.
Red dashed ellipses highlight the mode collapse issues observed in DDPM.
With the same initial noise, our method gives synthetic images with higher diversity and fidelity compared to DDPM.
For example, DDPM always generates rose images containing only one rose. Our method can generate an image containing two roses (see the sixth column of the last row).
To further clarify, in the ‘table’ class, DDPM repeatedly produces near-identical images as highlighted in red ellipses. In contrast, our method produces a more varied set of table images. This suggests that our method can better capture greater diversity in image generation when compared to DDPM (Ho et al., 2020).
To further illustrate how the proposed method mitigates the issue of overfitting, we visualize the top-10 nearest neighbors among 1000 synthetic images to an anchor image in the training set for the original DDPM and our method.
As shown in Fig. [10](https://arxiv.org/html/2507.09052v2#A1.F10), DDPM shows the mode collapse to the training image while our method’s generated images show much better diversity.
For example, in the first two rows, DDPM generates repeated vertical worms while our method generates worms with diverse directions.
In the 9th-10th rows, DDPM generates tables with repeated modes while our method generates tables with different styles and colors.

Figure: (a) Synthetic images from DDPM (Ho et al., 2020). Red ellipses show similar images generated.
Refer to caption: figures_iclr/ddpm.png

Figure: Figure 10: To see overfitting on tail classes, we find Top-10 nearest neighbors (3rd to 12th columns sorted by distances) among 1000 synthetic images to an anchor image (Leftmost column) in the training set. KNN is based on $L_{2}$ distances of Inception V3 embeddings. For each example, the top row is the results of DDPM, and the bottom row shows ours. The nearest neighbors from DDPM show the mode collapse to the training image, while our method generated more diverse images.
Refer to caption: x9.png

#### A.3.4 More Visualization of Reverse Process of DDPM for Different Classes

To illustrate our observation in Section [3.2](https://arxiv.org/html/2507.09052v2#S3.SS2) more clearly, we decompose $x_{t}$ into a combination of low-frequency images and high-frequency images, as shown in Fig [11](https://arxiv.org/html/2507.09052v2#A1.F11).
The low-frequency images with the same initial noise are very similar for different classes for the initial steps, which is also observed in prior work (Si et al., 2024).

Figure: Figure 11: Top: reverse process starting from the same initial Gaussian noise but with different class conditions (2nd-4th rows) or without condition (1st row). Middle: low-frequency components of each noisy image. Bottom: high-frequency components of each noisy image.

## Appendix B The Use of Large Language Models (LLMs)

We used a large language model to refine our writing and correct grammar, but not to generate the initial draft.
The LLM served only as a supportive tool to improve presentation, not as a source of content.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 4
presentation: 2
questions: 1. Regarding the training overhead, you mention a 1.6x increase in training time in the appendix. Could you please discuss this limitation in the main paper? Is it possible that the improved convergence speed (e.g., reaching a better FID in fewer epochs, as seen in Table 1) partially amortizes this per-epoch cost? What's the performance gap between baselines and CCUA given the same training budget measured in GPU$\times$Day?

2. Could you elaborate on the design choice for the latent representation h used in the UCL? Specifically for the Diffusion Transformer, why was the output of the final transformer block chosen? Did you experiment with applying the contrastive loss at different depths of the network, and how did that affect performance? Moreover, could you explain why 

3. In Table 3 (TinyImageNet-LT), CCUA shows a slightly higher FID on 'Head' classes compared to some baselines. Could you comment on this small performance degradation and whether there is a fundamental trade-off between improving tail classes and maintaining head class performance?

4. The alignment loss L_al is linearly weighted by t/T. Have you experimented with other weighting schedules (e.g., quadratic, or a step function that applies the loss only for t > T_threshold)? How sensitive is the model's performance to this specific linear schedule?  
In the appendix (line 723), you note that the UCL was also weighted by t/T for the TinyImageNet-LT and Places-LT datasets. This detail is not mentioned in the main method description. Could you clarify why this was necessary for these specific datasets and not for ImageNet-LT? This seems to suggest the method's components might require dataset-specific tuning.

5. How's the performance of CCUA w.r.t. the imbalanced factor (e.g., 0.001 as in OCLT)?
soundness: 2
strengths: - The paper tackles the practical problem of generative model performance on imbalanced, real-world data. The proposed method is simple and intuitive.
- The method is validated on a wide range of standard long-tailed datasets (ImageNet-LT, Places-LT, TinyImageNet-LT, etc.) at various resolutions, including challenging 256x256 generation. 
- The paper is well-written, clearly structured, and easy to follow.
summary: This paper addresses the challenge of training class-conditional diffusion models on long-tailed datasets, where tail classes suffer from poor sample quality and diversity (mode collapse) due to data scarcity. The authors propose a framework named Contrastive Conditional–Unconditional Alignment (CCUA), which introduces two novel loss functions to the standard diffusion model training pipeline. The authors demonstrate the effectiveness of CCUA by integrating it into both U-Net and Diffusion Transformer (DiT/SiT) architectures. Extensive experiments on long-tailed versions of CIFAR, TinyImageNet, Places, and ImageNet show that CCUA significantly improves generation quality (FID) and diversity (Recall) for tail classes, often without compromising (and sometimes improving) the quality of head classes.
weaknesses: 1. Training Cost: The alignment loss (Lal) requires a forward pass for both the conditional and unconditional predictions for each sample during training. The appendix (line 802) states this results in a ~1.6x increase in training time per epoch compared to a standard DDPM. This is a significant computational overhead that should be mentioned and discussed as a limitation in the main paper, not just in the appendix.

2. Justification for Latent Space h: The definition of the encoder e(*) that produces the latent h for the contrastive loss feels somewhat arbitrary, especially for the Diffusion Transformer. Some justification for why this specific feature level is optimal for applying the repulsive force would strengthen the paper. It is unclear if other layers were considered.

3. Unclear Rationale for Hyperparameter Choices: 
- The loss weighting parameters $\alpha$ and $\lambda$ are introduced but their values, selection process, and sensitivity are not discussed.
- The appendix (line 723) mentions that for TinyImageNet-LT and Places-LT, the UCL was also weighted by t/T, similar to the AL. This seems like a critical implementation detail. The rationale for applying this timestep weighting to UCL for some datasets but not others is not provided, which raises questions about the method's generality and the tuning required for new datasets.

---

confidence: 4
contribution: 2
initial_rating: 2
presentation: 2
questions: See weekness
soundness: 2
strengths: 1. The two proposed losses are intuitive and well-motivated and the method is also simple to implement, plus easy to integrate into existing diffusion pipelines.

2. The empirical evidence is extensive: results span multiple strong baselines, architectures (U-Net, DiT), and datasets. CCUA consistently outperforms prior work, especially in terms of FID for tail classes and diversity (Recall), as seen in Table 1, Table 2, and Table 3.
summary: The paper introduces CCUA, a framework designed to improve the performance of class-conditional diffusion models trained on long-tailed datasets. CCUA combines unsupervised contrastive loss and alignment loss together to address the problem that models trained on imbalanced data tend to suffer from mode collapse and reduced fidelity for tail classes. The method shows empirical enhancement on ImageNet-LT,Places-LT and CIFAR-LT variants.
weaknesses: The primary weakness is that the paper's contribution appears to be an assemblage of existing or concurrent ideas rather than a fundamental innovation. The first component, UCL, is a dispersive loss. The authors themselves state it is "similar to" concurrent work on dispersive losses. The second component, AL, is admittedly a direct adaptation from the long-tailed GAN literature. The paper's insight is to translate the "low-resolution alignment" from GANs to "large timestep alignment" in diffusion, which is a reasonable but straightforward adaptation. Therefore, the core contribution is limited to combining these two specific losses into an integrated loss, which feels more incremental than foundational.

The paper provides insufficient justification and no deep mathematical analysis for why this specific combination of UCL and AL is optimal. The two losses are simply added together with weighting hyperparameters. There is no deeper analysis or theoretical argument for how these two losses interact or why they form a principled, unified framework, beyond the high-level intuition.

The method introduces a major practical drawback. The paper's "Limitation" section (Appendix, page 15) explicitly states the method increases training time significantly (1.6x for DDPM, 1.48x for SiT) due to the two forward passes required for the Alignment Loss.
Furthermore, UCL (a contrastive-style loss) typically benefits from large batch sizes to see a rich set of negatives. However, the 1.5-1.6x increase in computational cost from the AL component would likely force a reduction in batch size. The paper fails to analyze or discuss this practical tension.

The paper's empirical claims are not sufficiently stress-tested, which is a significant flaw for a method that relies on a weighted-loss combination. The paper completely lacks a sensitivity analysis for its two most critical hyperparameters, $\alpha$ and $\gamma$. It is impossible to know if the strong results are due to extensive, unreported tuning or if the method is robust. The robustness analysis is very limited. It only tests robustness to random seeds (Appendix Table 10) and fails to analyze performance across different imbalance factors to see how it performs as the long-tail problem becomes more or less severe.The evaluation metrics, while good, are limited. For generative models, metrics like CLIP Score are often vital for judging image-concept alignment, especially for diverse tail classes, but this is missing from the evaluation.

The paper fails to adequately discuss or compare against highly relevant prior work on contrastive diffusion (e.g., Yan et al., 2022) and other conditional diffusion models (Wang et al., 2023).

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: - In the ablation study on the loss function, it would be helpful to include additional evaluation metrics such as Inception Score (IS) and Precision/Recall to provide a more comprehensive analysis.
- Would it be possible to include an experiment similar to Table 5 on the TinyImageNet-LT dataset to further validate the method’s generalization?
soundness: 2
strengths: - The paper effectively integrates multiple techniques to address the long-tailed generation problem. The approach is well motivated, drawing inspiration from prior research on class-imbalance GANs, and demonstrates empirical effectiveness.
- The method consistently performs well across both UNet-based architectures and Diffusion Transformers (DiT), highlighting its generalization capability.
- The paper is well-structured, with intuitive visualizations illustrating improved tail-class diversity and fidelity.
summary: This paper studies the problem of long-tailed image generation with diffusion models. The authors observe that diffusion models trained on long-tailed datasets (e.g., ImageNet-LT) tend to produce low-quality and less diverse images for tail classes. To address this, they propose a new regularization framework called CCUA (Contrastive Conditional–Unconditional Alignment), which combines (1) an Unsupervised Contrastive Loss (UCL) to promote feature dispersion and prevent latent collapse, and (2) an Alignment Loss (AL) to align conditional and unconditional noise predictions in the early timesteps of diffusion, enabling structural knowledge transfer from head to tail classes. Extensive experiments demonstrate that the proposed framework improves the generation quality of diffusion models.
weaknesses: - As acknowledged in the paper, the training process requires a longer duration.
- Loss-function ablation is evaluated only with FID, lacking complementary metrics (IS, Precision/Recall, FID-tail) to substantiate the claimed diversity–fidelity balance.
- While the results are promising, it remains somewhat unclear how much of the improvement stems from the proposed CCUA loss compared to the batch resample strategy. Since resampling already provides gains with minimal cost, a controlled comparison isolating CCUA’s effect would further clarify its independent contribution and better justify the added training overhead.

---

confidence: 4
contribution: 2
initial_rating: 4
presentation: 3
questions: 1. The paper computes the UCL in the unconditional latent space $e(\*;\emptyset)$ and transfers its effect to the conditional model via $\mathcal{L}\_{al}$​. Could the authors clarify why UCL is not directly applied in the conditional latent space $e(\*;c)$? A theoretical justification or comparison experiment would help explain whether applying UCL directly in the conditional space would indeed degrade performance.

3. In the experimental setup (Section 4.1), the compared methods for U-Net architectures using contrastive learning or explicit diversity enhancement appear to be limited to works published in or before 2024. Given the rapid progress in this area, it would be helpful to discuss or compare with more recent (2025 or latest) SOTA diffusion models that share similar motivations.

4. The proposed UCL computes the contrastive objective using only negative samples, aiming to enhance diversity by maximizing inter-sample distances. However, standard contrastive learning typically includes positive-pair attraction terms. It would be useful to provide further justification or an ablation study to clarify whether including positive pairs (same-class samples) would negatively affect long-tailed generation, and whether relying solely on repulsive forces might lead to overly dispersed latent representations (i.e., excessive intra-class variance).
soundness: 3
strengths: 1. The paper introduces an innovative approach to applying contrastive learning in diffusion models by using only negative samples, rather than both positive and negative pairs. Moreover, it creatively adapts the conditional–unconditional alignment concept from GANs to diffusion models, enabling knowledge sharing during early denoising stages — a notable conceptual transfer.

2. The proposed method achieves significant quantitative improvements across multiple long-tailed datasets, particularly in Recall (diversity) metrics, demonstrating its effectiveness in alleviating mode collapse and enhancing sample diversity.

3. The CCUA framework is clearly structured with two distinct loss components and is easy to integrate into existing DDPM/SBDM training pipelines. It is compatible with both U-Net and DiT architectures, and the paper provides a clear explanation of how UCL mitigates mode collapse.

4. The work provides a general and effective regularization strategy for addressing class imbalance in conditional diffusion models. It represents a meaningful advancement in improving diffusion models for real-world, long-tailed data distributions.
summary: This paper proposes a new framework named CCUA (Contrastive Conditional–Unconditional Alignment) to address the challenges of mode collapse and insufficient diversity or fidelity in conditional diffusion models (e.g., DDPM, SBDM) trained on long-tailed datasets. CCUA introduces two key loss components: (1) an Unsupervised Contrastive Loss (UCL) that operates in the unconditional latent space using only negative samples to enhance intra-class diversity by increasing the distance among generated samples; and (2) an Alignment Loss (AL) that aligns conditional and unconditional noise estimations during early denoising steps, facilitating knowledge transfer from head classes to tail classes. The proposed method is compatible with both U-Net and Transformer-based diffusion architectures, and demonstrates superior performance on multiple long-tailed benchmarks, including ImageNet-LT, Tiny ImageNet-LT, and Places-LT, particularly in improving diversity (Recall) and fidelity (FID) for tail classes.
weaknesses: 1. The paper claims to improve the generation quality of tail classes while maintaining the fidelity of head classes. However, key quantitative tables (e.g., Table 1) only report overall FID, IS, Precision/Recall, and $FID_{tail}$. To fully substantiate this claim, it would be important to include metrics directly reflecting the quality of head-class generations, such as $FID_{head}$.

2. In Section 2, the paper states that “our unsupervised contrastive loss is formulated in the unconditional latent space and is implicitly extended to conditional training through our conditional-unconditional alignment loss.” This statement is somewhat unclear, as Equation (7) in Section 3.3 defines $\mathcal{L}\_{ucl}$ and $\mathcal{L}\_{al}$ as two independent weighted losses. The mechanism by which $\mathcal{L}\_{al}$​ implicitly extends or interacts with $\mathcal{L}\_{ucl}$​ is not sufficiently explained.

3. As mentioned in Section 2, _Dispersive Loss_ was originally designed for balanced datasets rather than long-tailed ones. A stronger and more convincing comparison would adapt _Dispersive Loss_ to the long-tailed setting to better isolate the specific contribution of the proposed alignment loss $\mathcal{L}\_{al}$​.

4. To apply $\mathcal{L}\_{ucl}$, the paper divides the noise estimation network $\epsilon\_{\theta}$​ into an encoder and a decoder. While this separation is natural for the U-Net architecture, the definition becomes less clear for the Transformer-based SiT, where all SiT/DiT blocks are treated as the encoder and only the final linear layer as the decoder. The rationale for treating all blocks’ outputs as the contrastive latent space, rather than using a bottleneck or intermediate representation, is not well justified or empirically validated.


Decision: Reject