Improving Diffusion Models for Class-imbalanced Training Data via Capacity Manipulation
Authors: Feng Hong
Jiangchao Yao
Yifei Shen
Dongsheng Li
Ya Zhang
Yanfeng Wang
Venue: ICLR 2026 Oral (2026)
Keywords: Imbalance
Diffusion Models
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Abstract: While diffusion models have achieved remarkable performance in image generation, they often struggle with the imbalanced datasets frequently encountered in real-world applications, resulting in significant performance degradation on minority classes. In this paper, we identify model capacity allocation as a key and previously underexplored factor contributing to this issue, providing a perspective that is orthogonal to existing research. Our empirical experiments and theoretical analysis reveal that majority classes monopolize an unnecessarily large portion of the model's capacity, thereby restricting the representation of minority classes. To address this, we propose Capacity Manipulation (CM), which explicitly reserves model capacity for minority classes. Our approach leverages a low-rank decomposition of model parameters and introduces a capacity manipulation loss to allocate appropriate capacity for capturing minority knowledge, thus enhancing minority class representation. Extensive experiments demonstrate that CM consistently and significantly improves the robustness of diffusion models on imbalanced datasets, and when combined with existing methods, further boosts overall performance.

Published as a conference paper at ICLR 2026
IMPROVING DIFFUSION MODELS FOR CLASS- IMBALANCED TRAINING DATA VIA CAPACITY MA- NIPULATION
Feng Hong Jiangchao Yao Yifei Shen
Dongsheng Lie Ya Zhang
Yanfeng Wang
CCooperative Medianet Innovation Center, Shanghai Jiao Tong University
Microsoft Research Asia
"Scoool of Artificial Intelligence, Shanghai Jiao Tong University
'Institute of Artificial Intelligence for Medicine, Shanghai Jiao Tong University School of Medicine
(feng hong, Sunarker, 2hang, wangyanfengi ) @s jtu.u. edu. on (yi fiifeishen, dongsli] @microsoft.
ABSTRACT
While diffusion models have achieved remarkable performance in image genera- tion, they often struggle with the imbalanced datasets frequently encountered in real-world applications, resulting significant performance degradation on mi- nority classes. In this paper, we identify model capacity allocation as key and previously underexplored factor contributing this issue, providing a perspective that is orthogonal to existing research. Our empirical experiments and theoretical analysis reveal that majority classes monopolize an unnecessarily large portion the model's capacity, thereby restricting the representation of minority classes. To address this, we propose Capacity Manipulation (CM), which explicitly reserves model capacity for minority classes. Our approach leverages low-rank decompo- sition of model parameters and introduces a capacity manipulation loss to allocate appropriate capacity for capturing minority knowledge, thus enhancing minority class representation. Extensive experiments demonstrate that CM consistently and significantly improves the robustness of diffusion models on imbalanced datasets. and when combined with existing methods, further boosts overall performance.
INTRODUCTION
Diffuusin models have demonstrated exceptional generative capabilities, as evidenced by their rapid adoption in both academia and industry (Ho et al., 2020; Song al., 2021b; Dhariwal and Nichol, 2021; Ramesh et al., 2022; Rombach et al., 2022; Peng et al., 2026). Nevertheless, a key challenge persists: diffusion models exhibit marked decline performance on minority classes when trained on class-imbalanced datasets (Qin et al.. 2023; Zhang et al., 2024), a concern amplified by the prevalence data imbalance in real-world scenarios (Reed, 2001; Zhang et al., 2023).
Most research on imbalanced learning has focused on enhancing the robustness of discriminative models (Buda al., 2018; He and Garcia, 2009; Wang al., 2021a; Menon et al., 202 la; Cui al., 2021). However, these techniques cannot be directly applied diffusion models due to fundamental differences model architectures, training, and inference processes. To address class imbalance in generative diffusion models, recent pioneering works (Qin et al., 2023; Zhang et al., 2024) have fo- cused on loss function design to facilitate knowledge transfer between majority and minority classes. For example, Class Balancing Diffusion Models (CBDM) (Qin et al., 2023) introduced regular- ization loss that implicitly encourages generated images to follow a balanced prior distribution at each sampling step. Oriented Calibration (OC) (Zhang al., 2024) further improved the generation quality of minority classes through knowledge transfer from majorities.
Work partially done during Feng Hong's 's internship at Microsoft Research Asia. 'Code: httggg/uhw..ooo/ooaooooo///////a

---

Published as a conference paper at ICLR 2026
In this paper, we explore an orthogonal perspective: model capacity allocation'. Our motivation stems from the empirical observation that majority classes tend to monopolize a disproportionate share of the model's 's capacity, thus limiting the capacity available for minority classes. This finding is supported by both empirical and theoretical analyses. Empirically, the visual quality of minority classes is significantly lower than majority classes. Furthermore, we find that although the training loss for minority classes is similar to majority classes, pruning model parameters with the smallest LI-Norm leads more pronounced change in relative loss for minority classes. This suggests that minorities are allocated less model capacity, making them particularly vulnerable to pruning. A gradient- based theoretical analysis further supports this, showing that the majority classes drive a significant fraction of parameter updates, thereby dominating capacity allocation.
To address this issue, we propose a dedicated model capacity allocation method for diffusion mod- els, termed Capacity Manipulation (CM). Our core idea is to reserve model capacity specifically for minority expertise in advance, preventing from being overtaken by majority classes and thus safe- guarding the learning process for minority samples. Concretely, we decompose model parameters into two components using low-rank techniques: one capturing majority and general knowledge, and the other reserved for minority expertise (Eq. (1)). By introducing capacity manipulation loss (Eq. (2)), we ensure that minority knowledge is effectively allocated its reserved capacity during training. We further provide theoretical support, proving that our method can successfully balance update proportions to ensure adequate capacity allocation for minority classes. Our key contribu- tions are summarized as follows:
A Novel Perspective: Unlike prior works, this paper takes the first step in exploring model capacity allocation to enhance diffusion models on imsalannbd training data, moti- vated by both empirical findings and theoretical analysis.
A New Method:
We propose new method, CM, which protects minority classes by re- serving model capacity for minority expertise and strategically allocating relevant knowl- edge during training. Instead of increasing model size, CM reallocates existing capacity, introducing no additional inference overhead. Furthermore, CM is orthogonal to existing methods and can be integrated complementarily.
Excellent Performance: Extensive experiments on various datasets and diffusion model architectures demonstrate that our method consistently outperforms state-of-the-art base- lines. CM substantially improves minority class performance without sacrificing generative quality on majority classes.
ON THE LIMITATIONS OF DIFFUSION MODELS ON IMBALANCED DATA
2.1 PROBLEM FORMULATION
Let X' be the image space and y = (l...., C} the label space, where C' is the class number. We consider an imbalanced training set D {(x", y") E (X, of size N. The per-class sample counts N₂ for e E y in the descending order exhibit a long-tailed distribution. The goal is to learn generative diffusion model px(x)y), parameterized by 6, from the imbalanced data D to generate realistic and diverse samples across all classes. For unconditional generation, the class condition can be set to Null, yielding px()) = po (x Null).
2.2 EMPIRICAL MOTIVATIONS
Diffusion models, despite their strengths, are notably vulnerable to data imbalance. We illustrate this in Fig. ((a) with DDPM trained on Imb. CelebA-HQ (100:1 female-to-male sample ratio). minority-class (male)
While good at generating majority-class (female) images, its performance images is severely degraded, evident from visual results and FID scores. In comparison, we illus- trate a preview of our proposed method, which significantly improves the generation quality for the minority class (male images). Fig. 2(c) demonstrates similar findings on Imb. CIFAR-100.
In this paper, we define "model capacity" as the available representational resources of the neural network capture and store data features.
2

---

Published as a conference paper at ICLR 2026
Real Images
DDPM
Ours
Sample
Relative Loss Change
Relatve Look Change
Female
(Majority)
Sample Number
430
300
FID: 6.815
100
Male (Minoriay)
110
Class Index
FID: 16.425
FID: 12.788
(a) Visualization on Imb. CelebA-HQ (IR 100)
(b) Loss change after pruning
Figure I: (a) Real images from Imb. CelebA (Imbalance Ratio IR 100) and generated images
from DDPM and our method. Images are randomly sampled. (b) Relative per-class loss change, defined as (ppraned - / /Frww/Cc C/ for DDPM model trained on Imb. CIFAR-100 (IR = 100) when 10% of its parameters with the smallest LI-Norm are pruned. Raw losses and absolute loss changes are provided in Figs. G.1 and G.2.
Fig. 1(b) reveals an intriguing finding on Imb. CIFAR-100: despite similar raw training losses across classes (Fig. G.I), minority classes are disproportionately sensitive to pruning the smallest LI-norm parameters. A model capacity perspective may offer a compelling explanation: majority classes monopolize capacity, forcing minority class information into less critical (low L1-norm) parameters. These fragile representations are thus more susceptible to pruning, underscoring standard diffusion model limitations and motivating targeted capacity allocation.
2.3 THEORETICAL INSIGHTS
We provide Theorem 2.1 serve as an illustrative example: majority classes dominate the model's capacity, leaving fewer resources for minorities, which leads to poorer performance for them. The assumption, detailed proof, and its extension the multi-class setting are presented Sec. F.
Theorem 2.1. Consider setting with two classes y = {1,2), where the training set is highly imbalanced, i.e., N₁ >> N2. Under Assumption F.1, the expected proportion of parameter matrices within the network whose update direction are dominated by the majority class = is given by
1mmjj 22n-12 A(p1-12))
N1 is the majority ratio, N1 N2 where N1+No
2a
(-)) is the standard normal cumulative distribution function. μr and μ2 expected gradients of classes y = and y = 2, > 0 and a > 0 are positive constants defined in Assumption F.I.
Remarks. We can draw the following insight from Theorem 2.1:
When a > 0.5 and approaches 1, the standard normal CDF Φ(-) (monotonic, concave, and > 0.5 for positive inputs, Eq. (F.6)) implies: (1) the majority class dominates model capacity, and (2) greater class imbalance leads to greater capacity imbalance, albeit with diminishing returns.
Mm decreases as inter-class similarity increases (smaller ((H1, suggesting greater class similarity allows minority classes to benefit more from majority-drive updates.
Umaj decreases with relative within-class diversity (g). Higher diversity helps mitigate imbalance.
Umaj increases with total samples N, suggesting capacity dominance worsens for larger datasets.
METHOD: CAPACITY MANIPULATION
Based on the above motivations, we propose allocating dedicated model capacity specifically for minority classes, reserving it in advance to prevent encroachment by majority classes
3

---

Published as a conference paper at ICLR 2026
(a) Capacity Reservation 6 @
(b) Capacity Allocation
(c) Performance
Imb. CIFAR10O,
DOPM
(W, W2 W
CRIM
@4 (W W...
Fox
Consistent for majorities (4ou)) Diverse for minorities (Cgi))
17.5
CM
15.0-
0 @ d
Pease
Split
(a.b) An overview of our method, CM. (a) An illustration of the capacity reservation part of CM. (b) An illustration of how CM allocates the corresponding knowledge to the reserved model capacity during training.
Figure 2:
(c) Many/Mediun/Few split performance Imb. CIFAR 100 with imbalance ratio IR = 100, where Many/Medium//ert represents the top, middle, and bottom thirds of classes sorted by sample number in descending order. CM significantly improves minority per- formance without sacrificing the performance of majorities.
3.1 CAPACITY RESERVATION
To allocate sufficient model capacity for minorities, we first need to explicitly partition the model ca- pacity. Here we achieve this purpose by technique similar Low-Rank Adaptations (LoRAs) (Hu et al., 2022), which has demonstrated excellent performance and versatility in the field of efficient fine-tuning. While our task and goal differ, we apply low-rank decomposition concept partition the model capacity. For a diffusion model parameterized by 0 (W1, ........ where each W E A represents a parameter matrix the network, we decompose any W RR
W W BA W + W MA
where W9 E gaxk represents the part of W to be retained for majorities and generalized knowledge, We BA E Raxk represents the part to allocated to the minority expertise, E RRdxr A E Re and the rank < min(d, k). From Eq. (1), decompose & into 6 {W]" W?....) and d" = (WT, W.....} and merge them by 8 = 8 06' where + means the clement-wise addition. An illustration of Capacity Reservation is shown in Fig. 2(a).
As a follow-up to Theorem 2.1, we present Theorem 3.1, which demonstrates how the low-rank parameter decomposition in Eq. (1) for reserving capacity affects the actual capacity among classes.
Theorem 3.1. For two classes y 21, assume the training set is highly imbalanced, i.c., N1 > N2. Suppose that class y = contributes gradients to both W9 and W while class = 1 only contributes to W9, Under Assumptions F.1 and F.2, the expected proportion of parameter matrices W E A within the network whose update direction are dominated by the majority class y =
is bounded by mmaj- ② ( (2-a)) (2 ((p1)
where Or is monotonically increasing function of rank r, with a, = 0 when r = O, and a, when W is full rank.
2x
In this bound, we omit the effect of the inter-class sample ratio by manipulating the incquality, in order clearly showcase the impact of the low-rank component. Consequently, the rank of W modulated by a,,, controls the capacity allocation between majority and minority classes. Adjusting the rank directly mitigates the capacity collapse inherent in standard training. We provide the sumptions and detailed proof of Theorem 3.1 in Sec, F. We also provide extensions of Theorem 2.1 and Theorem 3.1 the multi-class setting in Sec. F.4.
3.2 CAPACITY ALLOCATION
With the model parameters decomposed as 0 0 @ 0s, our goal during training is to allocate 0= for minority expertise and 4 for majority and generalized knowledge, enhamcing minorities through capacity manipulation. To achieve this, the diffusion model (x)l) 4e- (2|y) should perform well on all samples, both majorities and minorities. Meanwhile, pas (zly) should perform well on majorities but poorly on minorities, 0 is not intended to learn the minority expertise.
4

---

Published as a conference paper at ICLR 2026
Capacity manipulation loss. For 0 6 + 8", use base loss Lasse (D,0) to train on the entire dataset, covering both majorities and minorities. Here, we can use the commonly adopted diffusion loss as the base loss: Lrase (D, 0) x Exyed Be,, [x]e,(X,, y) el <lž] Alternatively, we can choose variants adapted for imbalanced datasets, e.g., Zhang et al. (2024); Qin et al. (2023). For the separate 0 and 0% we propose a capacity manipulation loss, which encourages 0 to learn minority expertise and d0 capture generalized knowledge:
(x y, 1,88% 8²) = Lcon (x, y. 10,00,0 8*) + L y, 0, 80), Lco (x,9. y. 09, 0²) = wCon E Ega (X,, L y) 60s (X,, 1, 4) l3
(2)
LDiv(X, 09, 0*) - wib Fe 600 (xx, 1 v) cgo (x,, 1, v) l1..
where the capacity manipulation loss gLom is composed of a consistency loss LCon and a diversity loss LDiv. We vary the consistency class weight WCon and the diversity class weight WDiv applied to different classes. For class c E y with N₂ instances, a larger Ne (majorities) results in a higher consistency class weight wonn leading to more consistent outputs between Egugee (x,,t, y) and 6g0 (Xe, y). Conversely, for the diversity class weight, a smaller N, (minorities) results in a higher "Biv leading more diverse outputs between En-@0- (x,, y) and Egs (x,,,,)) Thus, pas (2r (l) cels on majorities, as its output aligns with 0 but underperforms on minorities due to the divergence between the outputs of fa and f. Specifically,
CN Ee No
W
U Wiv
Con
N2 EE N.
Here WCon scales linearly with class sample size, while Wniv is inversely proportional class sample size, ensuring WCon = WDiv 1, LCM(x, y, 09, 0f) = 0 for a balanced dataset. WCon and wDiv adaptively assign smooth weights to classes with varying sample sizes, avoiding the need for manual separation between majorities and minorities. While formally similar to reweighting techniques in discriminative models (Cui et al., 2019b) and generative models (Xie et al., 2023; Fan et al., 2024; Kim et al., 2024; Li et al., 2025; Liu et al., 2025), Eq. (3) differs fundamentally: it allocates model capacity by modulating output distances tied to distinct network parameters, rather than weighting individual sample losses.
Joint optimization. For 0 = 09 00 we optimize the base loss Lrase and the capacity manipulation loss LcM, weighted by hyperparameter
1 $ FCM(X y, 09,000000
min LTotal (D,0) = Lasee (D,0)+1 0
Σ
(4)
(x.0)ED
where Liase optimizes 0 for both majorities and minorities, while the capacity manipulation loss LCM acts as a regularizer to allocate capacity and protect minorities. This guides & toward more balanced and effective model weights. An illustration of the training process is presented in Fig. 2(b).
Inference. For inference, we can explicitly compute and store 0 09 e 0-, and sample images from pe (xly). Thus, method does not increase model capacity, ensuring no additional inference latency compared to a standard model, which is crucial as inference speed is a key bottleneck in real-world deployment (Song al., 2021a). This advantage comes from the LoRA-like parameter decomposition in Eq. (1) and explicitly aggregating parameters during inference.
3.3 DISSCUSSION
Comparison with existing imbalanced diffusion models. Unlike current methods such as CBDM and OC, which focus on designing improved objective functions for imbalanced data, CM improves the robustness of diffusion models class imbalance from a new perspective: allocating model capacity to protect minorities. CM is orthogonal and can benefit from improved objective functions to achieve further enhancements.
Comparison with LoRA and MoE-style Adapters. While CM shares low-rank structure with LoRA, our goal is capacity reservation prior to training rather than efficient fine-tuning. In Tab. 1, we compare two variants of the CBDM based on low-rank components: (1) CBDM (0 6, (.):: This variant utilizes our decomposition approach while maintaining the original CBDM objective; (2) CBDM (LoRA): This variant incrementally fine-tunes the pre-trained CBDM model using LoRA.
5

---

Published as a conference paper at ICLR 2026
Table 1: FIDÍ comparison with baselines and low-rank variants.
Dataset
DDPM CBDM
CBDM (# = 89 e 6*)
CBDM (LoRA)
CM
CIFAR-I00. IR = 100
10.
10.051
10.231
11,424
7.519
CIFAR-10, IR. = 100
10.697
8.233
8.316
9.725
7.727
Neither variant yields improvements over the original CBDM, whereas CM significantly outper- forms all baselines. This is because low-rank parameterizations like LoRA do not increase model capacity, but merely change how parameter updates are applied to a fixed-size network. To fur- ther distinguish CM from Mixture-o--Expers (MoE) paradigms, LoRA" baseline (details in Sec. E.4), which assigns distinet LoRA adapters to Many/Mediumm/eww class groups as a deterministic MoE. As shown in Sec. E.4, CM (FID 7.519) significantly outper- forms this MoE-style baseline (FID 10.06 on Imb. CIFAR-100) confirming that our gains stem from explicit capacity manipulation rather than simply increasing parameters or dynamic routing- Detailed conceptual comparisons with reweighting and adapters are provided in Sec. E.3,
implemented "Group-Expert
Extension to LoRA-finetuning Our method can be seamlessly extended to LoRA-finetuning sce- narios by modifying Eq. (1) to the form: W = w + A + 8 A Here, of = (W/, W..... represents the frozen pre-trainets model parameters, f = (B A4, A2....... denotes the trainable parameters allocated for majorities and generalized knowledge, and 8 = {BB A5, cor- responds to the trainable parameters reserved for minority expertise. For W E Rd B E Rdxro A E R B E Rdxr", A IR" xk we have 1d و < min(d, During inference, the model parameters are merged by 0 = of e 8 eb 0 This extension preserves the structure of LoRA while enhancing the fine-tuning process by capacity manipulation.
4 EXPERIMENTS
4.1 EXPERIMENTAL SETUP
Datasets. We conduct experiments the imbalanced versions of commonly used datasets in the field of image synthesis. including CIFAR-10 10 (Krizhevsky et al., 2009), CIFAR- 100 (Krizhevsky al., 2009), CelebA-HQ (Karras al.. 2018), ImageNet-LT (Liu al., 2019), iNaturalist (Horn et al., 3018), and ArtBench-10 (Liao et al., 2022). CIFAR-10 and CIFAR-100 have a resolution of 32x32; for CelebA-HQ. we use the 64x64 version; for ImageNet-LT and iNaturalist, we use the 32x32 and 64x64 versions; and for ArtBench-10, we use the original resolution of 256 We follow Cao et al. (2019a) to construct imbalanced versions of these datasets by downsampling, resulting in an exponential decrease in the sample size of each class with its index. We refer to these imbalanced datasets as Imb. dataset, e.g... Imb. CIFAR-10. 10. We control the level of imbalance the dataset by setting different imbalance ratios (Menon al., 202 Qin al., 2023) IR E {50, 100), where IR is the ratio of the number of samples in the most populous class to that in the least populous class, defined as IR max.ey Ne For Imb. CIFAR- 10 and Imb. ArtBench- 10, we divide the dataset min.ey N. into three splits: Many (classes 0-2), Medium (classes 3-5), and Few (classes 6-9) based on class sizes in descending order. Similarly, for Imb. CIFAR- 100, the splits are Many (classes 0-32), Medium (classes 33-65), and Few (classes 66-99). This uniform three-way split follows Jiang et al. (2021).
Baselines. We consider baselines including: (1) the base denoising diffusion probabilistic model (DDPM): (2) methods specifically targeting imbalanced diffusion models: the class-balancing dif- fusion model (CBDM) (Qin et al., 2023) and Oriented Calibration (OC) (Zhang et al., 2024); (3) ap- plying imbalance learning methods from discriminative models or generative adversarial networks (GANs) to diffusion models: re-sampling (RS) (Mahajan et al., 2018), adaptive discriminator aug- mentation (ADA) (Karras et al., 2020), and focal loss (Lin et al., 2017b). Note that many imbalanced learning methods for discriminative models and GANs heavily rely on specific model architectures or training paradigms, e.g., Menon et al. (2021a); Zhou et al. (2023a); Rangwani et al. (2022), mak- ing them incompatible with imbalanced diffusion models.
Metrics. The performance is evaluated using the metrics Frechet Inception Distance (FID) (Heusel et al., 2017), Kernel Inception Distance (KID) (Binkowski et al., 2018), Recall (Kynkäänniemi et al., 2019), and Inception Score (IS) (Salimans al., 2016). All metrics are calculated based on features

---

Published as a conference paper at ICLR 2026
Table 2: FIDs (1). KIDs ()), Recalls (f), and ISs (1) of CM and various baseline methods on Imb.
CIFAR-10 10 and Imb. CIFAR-100 100. Results for imbalance ratios IR = {100, 50} are shown side-by-
side. Best and second-bes results are highlighted. Results with Mean can be found in Tab. G.I.
Imnh. CIFAR-10
IR
Method
FID
KID
Recall g
15
FID
KID
Rec
I5
DDPM (Ho et al., 2020)
10.697
0.0035
0.47
9.39
10.216
0.0035
0.47
9.37
+ADA (Karras
9.266
9.26
9.132
9.28
RS (Mahajan et al. 20188
12.132
0.0037
0.45
9.25
11.231
0.0038
0.47
9.31
+Focal (Lin al., 2017b)
10.842
0.0034
0.6
9,42
10.315
0.48
9,38
CBDM (Qin al., 2023)
8.233
0.0026
0.53
9.23
7.933
0.0026
0.54
9.42
OC (Zhang al., 2024)
8.390
0.0027
0.52
9.53
8.034
0.0027
9.65
CM
7.727
0.0023
0.53
9.52
7.372
0.0024
0.54
9.69
Inbb. CIFAR-100
IR
Method
FID
KID
Recall
15 T
FID
KID
Rec
18
DDPM (Ho et al., 2020)
10.163
0.0029
0.46
13.45
9.363
0.0032
0.47
14.27
+ADA (Karra et
9.482
12.44
8.927
12.89
+RS (Mahajan et al., 201
11.432
0.0038
0.44
12.12
10.259
0.0037
0.47
12.38
+Focal (Lin al., 2017b)
10.212
13.07
CBDM (Qin et al., 2023)
10.051
0.0036
D.51
12.35
8.946
0.0036
0.55
12.59
(Zhang al., 2024)
8.309
0.0026
13.44
7,188
13.99
CM
7.519
0,0017
0.52
13.45
6.732
0,0021
0.55
14.12
Table 3: Per-split FIDs (L) of CM and baselines on Imb. CIFAR-1O 100) and Imb. CIFAR-100 (IR 100) shown. Many, Medium, and Few are the three splits based on training imbalancedness. Best and second- best results are highlighted. Results with Mean Std can be found in Tab. G.4.
CIFAR-I0
100)
CIFAR-100
100)
Method
Many FID
Med. FID
Few FID
Overall FID
Many FID
Med. FID 1
Few FID
Overall FTD
DDPM al., 2020)
14,203
19:714
15,869
10,697
14.068
15,660
22.188
10.163
CBDM (Qin et al., 2023)
12.222
14.465
12,230
8,2334
12.585
14,042
20.667
10.051
OC (Zhang et al., 2024)
12,026
15.234
12.254
8,3896
11.731
13.133
21.083
8,309
CM
11.705
14.340
11.218
7.7273
11.713
13,043
18.729
7,519
extracted from pre-trained Inception-Vi V3 (Szegedy et al., 2016) model?. During evaluation, the metries are calculated using balanced set of real images and 50,000 generated images. The met- rics for each [many, medium, few split are computed using the corresponding split's real images and 20,000 generated images. To mitigate discrepancies from the high sensitivity of metrics like FID to implementation details (e.g., sample count, image preprocessing, input normalization, and feature dimensionality), we benchmark all methods under unified evaluation pipeline. We stress that our models trained imbalanced datasets, which are more challenging and results in higher baseline FIDs. Accordingly, our analysis prioritizes the relative performance gains within this challenging regime over absolute FID values reported on balanced benchmarks.
Implementation details. We provide comprehensive implementation details in Sec. G.I.
4.2 MAIN RESULTS
Performance on Imb. CIFAR-10 and CIFAR-100. In Tab. we summarize the FIDs, KIDs, Re- calls, ISs of CM and all baselines on Imb. CIFAR- 10 and Imb. CIFAR- 100 with different imbalance ratios IR = (50, 100). CM achieves the best results on 16 metrics across all four settings, except for two slightly lower ISs. Note that IS cannot detect mode collapse (Barratt and Sharma, 2018), e.g., the generated minority samples are overwhelmed by majority characteristics, such low-quality images would not lead to drop in IS, which explains why vanilla DDPM still performs well on some ISs. Additionally, IS lacks a reference to real images, making it generally considered a less reliable metric (Borji, 2019; Numn et al., 2021). On the most widely used metric FID, CM achieve significant improvements over DDPM with gains of 2.725, 2.844, 2.644, and 2.571. and consistent improvements over the best baseline in each setting by 0.506. 0.561, 0.790, and 0.456. Additionally, we provide a direct comparison with the concurrent work Overlap Optimization (Yan et al., 2024) in Sec. 1. where CM consistently achieves superior performance.

---

Published as a conference paper at ICLR 2026
Table 4: FIDs (4). KIDs (L), and per-class FIDs (4) of CM and baselines on Imb. CelebA-HQ with imbalance ratios IR- = {100, 50} shown sideeby-side. Female and Male are the two classes. Best and second-best results are highlighted. Results with Mean Std can be found in Tab. G.2.
Imb. Ce
IR 100
Method
Fomalc FID 1
Male FID .
Overall FID
KID
Female FID
Male FID L Overall FID L
KID
DDPM (Hoo et 2020)
7.143
16.425
8.727
0.0037
7.348
14,808
8.007
0.0034
CBDM (Qia et al., 2023)
7,043
14.273
7.823
0.0043
7.317
12.592
7.423
0.0042
OC (Zhang al., 2024)
7.092
13.962
7,871
0.0034
7.283
12.938
7,438
0.0034
CM
6.815
12.788
7,538
0,0033
7.147
11.273
7.193
0.0033
Many/Mediun/Feww analysis. In Tab. 3, we show the fine-grain fine-grained {many, mediun, few] per-split FIDs of different methods on Imb. CIFAR-10 and Imb. CIFAR-100 with imbalance ratio TR 100. Our method achieves the best results across all three splits, with the primary improvements observed in the Medium and Few classes. It is noteworthy that on Imb. 100 the generation quality for Medium classes is worse than for Few classes. Similar observations have been made on imbalanced contrastive learning (Zhou et al., 2023a). This could be attributed to the inherent difficulty differences between classes, suggesting a promising direction of addressing imbalanced diffusion models by combining inherent difficulty imbalance with quantity imbalance.
Performance on Imb. CelebA-HQ. In Tab. 4, we report the FIDs, KIDs, and per-elass FIDs of CM and baselines on Imb. CelehA-HQ with different imbalance ratios IR {100, 50}. Imb. CelebA- HQ contains two classes: Female and Male, with Female being the majority class. CM achieves the best performance across all eight metrics in both settings. Specifically, it improves the Overall FID by 1.189 and 0.814 compared to DDPM and by 0.285 and 0.230 compared to the best baselines in each setting. For the minority class (Male), CM enhances FID by 3.637 and 3.535 over DDPM and by 1.174 and 1.319 over the best baselines. In Fig. G.6, we showcase the generated results for the "Male" class (minority) with IR- 100. CM generates more realistic and diverse faces.
Performance on datasets with thousands of classes (ImageNet-LT and iNaturalist). In Tab. 5, report results on ImageNet-LT and iNaturalist, large-scale datasets featuring thousands of classes (1,000 and 8,142, respectively) and large-scale datasets featuring thousands classes and extreme imbalance, specifically ImageNet-LT (IR 256) and iNaturalist (IR 500). This verifies CM's robustness to imbalance ratios significantly beyond 100:1. Following Bartosh et al. (2024); Huang al. (2024), we adopt resolutions of x2x32 and 64 64. While the baseline CBDM shows perfor- mance degradation compared to DDPM when handling thousands of classes, exhibiting limitations at this scale, our proposed method still demonstrates significant and consistent advantages over all baselines.
Performance of Fine-tuning Stable Diffusion on Imb. ArtBench-10 On Imb. ArtBench- 10, we fine-tune the Stable Diffusion model* (Rombach et al., 2022) by LoRA (Hu et al., 2022) with rank of 128. And for 0=, the rank set to 8. We train the model in a class-conditional manner where the textual prompt is simply set as "a (class) painting" such as "a renaissance painting". The dropout rate is set to 0.1, and the model is trained for 100 epochs with a batch size of 64, using the Adam optimizer (Loshchilov and Hutter, 2019) with weight decay of 10 and an initial learning rate 3 x 10-4. During inference, we generate new images using a 50-step DDIM solver (Song et al., 2021a). In Tab. 6, we compare and OC, on Imb. ArtBench-10 with imbalance ratios IR = ={100,50 50}. Our CM achieves the best results across all eight metrics. The generated images for one of the few classes "Realism" on Imb. ArtBench-10 10 with IR = 100 are shown in Fig. 3. Our method generates more diverse images, and the generated styles are closer to the real images. This visual diversity is quantitatively supported by the superior Recall scores (Tab. 6), and we provide extensive visualizations of the intra-class diversity for tail classes in Fig. G.5. See Sec. G.8 for more visualization results.
CM against DDPM and the two strongest baselines, CBDM
4.3 FURTHER ANALYSIS
Human and Downstream Evaluations. In addition to the above automatic metrics, also present human preference studies (Sec. G.4) and downstream utility assessments (Sec. G.5).
8

---

Published as a conference paper at ICLR 2026
Table FIDs (4) and KIDs (L) on ImageNet-LT and iNaturalist 32x and 64 64 resolutions.
ImageNet-LLT 32x32
ImageNet-IIT
iNaturalist
iNaturalist
Method
FID
KID L
FID
KID
FID
KID
FID 0
KID J
DDPM (Ho ct al., 2020)
13.42
0.0082
8.84
0.0028
15.66
0.0095
7.99
0.0030
CBDM (Qin et. 2023)
14.66
0.0096
0.0068
0.0136
10.28
0.0044
OC (Zhang et al., 2024)
12.61
0,0075
8.43
0.0025
14.29
0.0092
7.67
0.0023
CM
10.94
0.0023
0.0089
Table 6: FIDs (4), KIDs (L), Recalls (1), and ISs (f) of CM and various baselines on Imb. ArtBench- 10 with IR {100, 50} shown. using LoRA to fine-tune Stable Diffusion. Best and second-best results are highlighted. Full results with Std can be found in Tab. G.3.
Imb. ArtBench-10
IR
100
IR
50
Method
FID 1
KID }.
Recall
IS
FID 1
KID 1
Recall f
IS
DDPM (Ho et al., 2020)
27.083
0.0142
0.39
8.47
25.557
0.0134
0.39
8.41
(Qin al., 2023)
25.723
7.97
24.487
OC (Zhang al., 2024)
24.315
0.0106
0.42
8.71
23.287
0.0097
0.43
8.48
CM
22.776
0.0087
0.44
8.71
21.733
0.0080
0.44
8.51
Imb. ArtBench-100 10, IR 100, Realism
DDPM
CM
Images
The visualization of generated images for the class "Realism", which is one of the few classes on Imb. ArtBench-10 with IR 100. The last row displays real images from the dataset for reference. It is evident that CM generates results that are significantly more diverse and stylistically closer to the real images compared to DDPM. The images shown are randomly selected.
Figure 3:
CM as universal framework. Fig. 4(a) summarizes the performance of our CM when integrated with various baselines (i.e., using the corresponding objective function for D in Eq. (4)) on Imb. CIFAR-100 (TR 100). CM consistently improves the performance of imbalanced generation when combined with different baselines. Due to the orthogonality of CM to existing methods, it can consistently benefit from improved objective functions, including potential future advancements.
Knowledge allocation between ()9 and 0 We evaluate knowledge allocation by comparing our full model CM (using 0 09 0c) with ablation using only 099 (denoted CM (09)). On Imb. CIFAR- 100 (IR = 100: Tab. 7), CM (69) excels on Many and Medium classes but falters on Few classes. In contrast, CM shows strong performance across all splits. This indicates that CM successfully allocates minority expertise to 6 while reserving majority and general knowledge for 69.
Ablation on the hyperparameter À in Eq. (4), We conduct ablation experiments on the hyperpa- rameter X, the weight of the CM loss in Eq. (4). Fig. 4(b) illustrates how the FID of CM changes with varying values. We observe that CM maintains a consistent advantage over OC across a wide range of À values, with its performance peaking around À 1.0.
Ablation on the rank of W E 0% We conduct an ablation study on the rank ratio of W E 6 relative to full rank, i.e., minid, Fig. 4(c) shows how the FID varies with different rank ratios. CM consistently outperforms OC across a wide range of rank ratios, with peak performance around 0.1.
Ablation on Lcon and LDVV- Tab. presents the results of CM as well as the ablation study where the consistency loss LCon and the diversity loss LDiv removed separately from CM. Since Lcon

---

Published as a conference paper at ICLR 2026
SALEV NICH
TORM ICESANDA
EDD
RD
CM
75
OC
OC
1.42)
[A443]
(1234]
o/ //mmnnfff
Configuration
(a) Universal framework.
(b) X.
(c) Rank.
(d) Network configs.
Figure 4: (a) The performance of CM when integrated with baselines. (b) Ablation study on the hyperparameter in Eq. (4). (c) Ablation study on the rank r. (d) Ablation study on various UNet configurations. All experiments are conducted on Imb. CIFAR-100 with IR = 100. In (b) and (c), we use OC as reference because it shows the best overall performance among the baselines.
Table (Left) Per-split FIDs and overall FIDs (L) of DDPM, CM (99), and CM on Imb. CIFAR-100
with imbalance ratio IR. = 100. (Right) FIDs (1), KIDs ()), Recalls (1), and ISs (1) on Imb. CIFAR- 100 with imbalance ratio IR- 100. The last two rows show the results of CM after removing LCon and Div. respectively. Full results with Std can be found in Tables G.5 and G.ó.
Method
F1D
KID L
Recall
|
IS-
Many
DDPM (1lo et 2020)
10.163
0,0029
0.46
13.45
DDPM (tlo 2020)
14.08
15,600
22.188
10.163
OC (Zhing ata 2004)
8.309
0.0026
0.52
13.44
CM (6P)
11.923
14,872
29:357
13:712
CM
7,519
0.0017
0.52
13:45
CM 8t @*))
11,713
13,043
18.729
7,519
CM wlo Cesa
8:412
0.0009
0.50
13:25
CM wio Luwe
3,073
0.0025
0.51
13.42
and Lpiv are responsible for allocating majority knowledge and minority expertise, respectively, removing either leads significant drop in performance, highlighting their necessity.
Ablation on network configurations. We conduct experiments on UNet architectures with varying widths and depths, which is achieved by setting the channel_multippheee multipliers parameter 2, 2]. [1, 2, 2, 2] (default), and [1, 2, 3, 4). Fig. 4(d) shows FIDs on Imb. CIFAR-10 10 and CIFAR- -00 with IR 100. CM consistently demonstrates clear advantages across different network configurations.
Broader settings. The main experiments in this paper are conducted on the classical DDPM net- work (trained from scratch) and Stable Diffusion (fine-tuned), where controlled comparisons demon- strate and analyze the effectiveness of CM on class- inss-imbllanced data. In addition, we validate our method in broader settings, including consistency models, integration with DPM-Solver, training a latent diffusion model from scratch, and fine-grained text-to-image datasets, detailed in Sec. G.9. Furthermore, we explore the potential of CM discriminative tasks. On Long-Tailed CIFAR-100 classification using ResNet-18. 18, applying CM to strong baselines (LDAM (Cao et al., 2019a) and Logit Adjustment (Menon et al., 2021a)) yields consistent improvements (e,g., +1.78% accuracy for LA at IR 100), suggesting that capacity reservation principle generalizes beyond generative modeling. Detailed results are provided Sec. G.11.
CONCLUSION
In this study, we seek improve the robustness of diffusion models to class-imbalanced data. Un- like previous work that focuses on improving objective functions, we aim to protect the generation performance of minorities by reserving and allocating model capacity for them. We first decompose the model parameters into parts that capture general and majority knowledge, and dedicated part for minority expertise using low-rank decomposition techniques. By introducing a capacity manip- ulation loss, we successfully allocate the corresponding knowledge to the reserved model capacity during training. Extensive experiments and empirical analyses confirm that our CM effectively pro- tects minorities in imbalanced diffusion models via capacity manipulation.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
This work is supported by National Key R&D Program of China (No. 20222D0160702), National Natural Science Foundation of China (No. 62306178), STCSM (No. 22DZ2229005), and 111 plan (No. BPO719010).
REFERENCES
Laith Alzubaidi, Jinglan Zhang. Amjad J. Humaidi, Ayad Q. Al-Dujaili, Ye Duan, Omran Al- Shamma, José Santamaria, Mohammed A. Fadhel, Muthana Al-Amidie, and Laith Farhan. Re- view of deep learning: concepts, CNN architectures, challenges, applications, future directions. J. Big Data, 8(1 8(1)533, 2021.
Fan Bao, Chongxuan Li, Jun Zhu, and Bo Zhang, Analytic-dpm: analytic estimate of the optimal reverse variance diffusion probabilistic models. ICLR. OpenReview.net, 2022.
ShiLong Bao, Qianqian Xu, Feiran Li, Boyu Han, Zhiyong Yang, Xiaochun Cao, and Qingming Huang. Towards size-invariant salient object detection: generic evaluation and optimization approach. IEEE Transactions on Puttern Analysis and Machine Intelligence, 2025.
Shane T. Barratt and Rishi Sharma. A note on the inception score. CoRR, abs/1 801.0101733 2018.
Grigory Bartosh, Dmitry Vetrov, and Christian A. Naesseth. Neural diffusion models. ICML. OpenReview.net, 2024.
Mikolaj Binkowski, Danica Sutherland, Michael Arbel, and Arthur Gretton. Demystifying MMD gans. In ICLR. OpenReview.net, 2018.
Ali Borji. Pros and cons of GAN evaluation measures. Comput. Vis. Image Underst., 179:41-65 2019.
Mateusz Buda, Atsuto Maki, and Maciej A. Mazurowski. A systematic study of the class imbalance problem convolutional neural networks. Neural Networks, 106:249-253 2018.
Kaidi Cao, Colin Wei, Adrien Gaidon, Nikos Aréchiga, and Tengyu Ma. Learning imbalanced datasets with label-distribution-aaarr margin loss. Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buc, Emily B. Fox, and Roman Garnett, editors, NeurlPS, pages 1565-1576, 2019 9a.
Kaidi Cao, Colin Wei, Adrien Gaidon, Nikos Arechiga, and Tengyu Ma. Learning imbalanced datasets with label-distributton--aaae margin loss. In NeurlPS, 2019b.
Nitesh V Chawla, Kevin W Bowyer, Lawrence O Hall, and W Philip Kegelmeyer. Smote: synthetic minority over-sampling technique. Journal of artificial intelligence research, 16:321-357, 2002.
Jiequan Cui, Zhisheng Zhong. Shu Liu, Bei Yu, and Jiaya Jia. Parametric contrastive learning. In
ICCV, pages 695-704. IEEE, 2021.
Jiequan Cui, Shu Liu, Zhuotao Tian, Zhisheng Zhong, and Jiaya Jia. Reslt: long-tailed recognition. IEEE Trans. Pattern Anal. Mach. Intell., 45(3):3695-3706, 2023.
Residual learning for
Yin Cui, Menglin Jia, Tsung- Yi Lin, Yang Song, and Serge Belongie. balsnbalanced loss based on effective number of samples. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 9268-926689 9277, 2019a.
Yin Cui, Menglin Jia, Tsung-Yi Lin, Yang Song, and Serge J. Belongie. Class-balanced loss based effective number of samples. In CVPR, pages 0268-9277, Computer Vision Foundation / IEEE, 2019b.
Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat gans image synthesis. In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, NewrIPS, pages 8780-8794, 2021.

---

Published as a conference paper at ICLR 2026
John C. Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. J. Mach. Learn. Res., 12:2121-2159, 2011.
Simin Fan, Matteo Pagliardini, and Martin Jaggi. DOGE: domain reweighting with generalization estimation. In ICML. OpenReview.net, 2024.
lan J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron C. Courville, and Yoshua Bengio. Generative adversarial networks. CoRR, ftbs/1406.2611 2014.
Shuyang Gu, Dong Chen, Jianmin Bao, Fang Wen, Bo Zhang, Dongdong Chen, Lu Yuan, and Baining Guo. Vector quantized diffusion model for text-to-image synthesis. CVPR, pages 10686-10696. IEFE, 2022.
Haibo He and Edwardo A. Garcia. Learning from imbalanced data. IEEE Trans. Knowl. Data Eng., 21(9):1263-1288) 1284, 2009.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge to a local nash equilibrium. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vish- wanathan, and Roman Garnett, editors, NeurIPS, pages 6626 6637,2015
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In NeurIPS, 2020.
Feng Hong, Jiangehao Yao, Zhihan Zhou, Ya Zhang, and Yanfeng Wang. Long-tailed partial label learning via dynamic rebalancing. In ICLR, 2023.
Feng Hong, Jiangchao Yao, Yueming Lyu, Zhihan Zhou, Ivor Tsang, Ya Zhang, and Yanfeng Wang. On harmonizing implicit subpopulations. In ICLR, 2024.
Feng Hong, Yu Huang, Zihua Zhao, Zhihan Zhou, Jiangchao Yao, Dongsheng Li, Ya Zhang, and Yanfeng Wang. Dual-granularity sinkhorn distillation for enhanced learning from long-tailed noisy data. Machine Learning. 2026.
Grant Van Horn, Oisin Mac Aodha, Yang Song, Yin Cui, Chen Sun, Alexander Shepard, Hartwig Adam, Pietro Perona, and Serge J. Belongie. The inaturalist species classification and detection dataset. In CVPR, pages 8769-8778. Computer Vision Foundation / IEEE Computer Society, 2018.
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen -hu,, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. In ICLR. OpenRe- view.net, 2022.
Nick Huang, Aaron Gokaslan, Volodymyr Kuleshov, and James Tompkin. The GAN is dead; long live the gan! A modern GAN baseline. In Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, Ulrich Paquet, Jakub M. Tomezak, and Cheng Zhang, editors, NeurIPS, 2024.
Ziyu Jiang, Tianlong Chen, Bobak J. Mortazavi. and Zhangyang Wang. Self-damaging contrastive learning. In Marina Meila and Tong Zhang, editors, ICMI., volume 139 of Proceedings of Machine Learning Research, pages 4927 4939. PMLR, 2021.
Bingyi Kang, Saining Xie, Marcus Rohrbach, Zhicheng Yan, Albert Gordo, Jiashi Feng, and Yannis Kalantidis. Decoupling representation and classifier for long-tailed recognition. In ICLR, 2020.
Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for im- proved quality, stability, and variation. In ICLR. OpenReview.net, 2018
Tero Karras, Miika Aittala, Janne Hellsten, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Training generative adversarial networks with limited data. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, NeurlPS, 2020.
Tero Karras, Miia Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, NewrIPS, 2022.
12

---

Published as a conference paper at ICLR 2026
Yeongmin Kim, Byeonghu Na, Minsang Park, JoonHo Jang, Dongjun Kim, Wanmo Kang, and Il- Chul Moon. Training unbiased diffusion models from biased dataset. In ICLR. OpenReview.net, 2024.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors, ICLR, 2015.
Diederik P Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. On density estimation with diffu- sion models. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, NeurlPS, 2021.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved pre- cision and recall metric for assessing generative models. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buc. Buch Emily B. Fox, and Roman Garnett, editors, NeurIPS, pages
2019.
Mengke Li, Yiu-ming Cheung, and Yang Lu. Long-tailed visual recognition via gaussian clouded logit adjustment. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
Mengke Li, Zhikai Hu, Yang Lu, Weichao Lan, Yiu-ming Cheung, and Hui Huang. Feature fusion from head to tail for long-tailed visual recognition. In Proceedings of the AAAI Conference Artificial Intelligence, 2024.
Yize Li. Yihua Zhang, Sijia Liu, and Xue Lin. Pruning then reweighting: Towards data-efficient training of diffusion models. In ICASSP, pages 1-5. IEEE, 2025.
Peiyuan Liao, Xiuyu Li, Xihui Liu, and Kurt Keutzer. The artbench dataset: Benchmarking genera- tive models with artworks. CoRR, abs/2206.1 11404, 2022.
Tsung-Y Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, and Piotr Dollár. Focal loss for dense object detection. In Proceedings of the IEEE international conference on computer vision, pages 2980-2988, 20
Tsung-Yi Yi Lin, Priya Goyal, Ross B. Girshick, Kaiming He, and Piotr Dollár. Focal loss for dense object detection. In ICCV, pages 2999-3007. IEEE Computer Society, 2017b.
Qian Liu, Xiaosen Zheng, Niklas Muennighoff, Guangtao Zeng. Longxu Dou, Tianyu Pang, Jing Jiang. and Min Lin. Regmix: Data mixture regression for language model pre-training. In ICLR. OpenReview.net, 2025.
Ziwei Liu, Zhongqi Miao, Xiaohang Zhan, Jiayun Wang. Boqing Gong, and Stella X. Yu. Large- scale long-tailed recognition in an open world. In CVPR, pages 2537-2546. Computer Vision Foundation IEEF, 2019.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: fast ODE solver for diffusion probabilistic model sampling in around 10 steps. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, NewrlPS, 2022.
Jiaan Luo, Feng Hong, Qiang Hu, Xiaofeng Cao, Feng Liu, and Jiangchao Yao. Long-tailed recog- nition with model rebalancing. In NeurIPS, 2025.
Dhruy Mahajan, Ross B. Girshick, Vignesh Ramanathan, Kaiming He, Manohar Paluri, Yixuan Li, Ashwin Bharambe, and Laurens der Maaten. Exploring the limits of weakly supervised pretraining. In Vittorio Ferrari, Martial Hebert, Cristian Sminchisescu, and Yair Weiss, editors, ECCV, volume 1 11206, pages 185-201. Springer, 2018.
13

---

Published as a conference paper at ICLR 2026
Aditya Krishna Menon, Sadeep Jayasumana, Ankit Singh Rawat, Himanshu Jain, Andreas Veit, and Sanjiv Kumar. Long-tail learning via logit adjustment. In ICLR. OpenReview.net, 2021a.
Aditya Krishna Menon, Sadeep Jayasumana, Ankit Singh Rawat, Himanshu Jain, Andreas Veit, and Sanjiv Kumar. Long-tail learning via logit adjustment. In ICLR, 2021b.
Eric J. Nunn, Pejman Khadivi, and Shadrokh Samavi. Compound frechet inception distance for quality assessment of GAN created images. CoRR, abs/2106.08575, 2021.
Seulki Park, Youngkyu Hong, Byeongho Heo, Sangdoo Yun, and Jin Young Choi. The majority can help the minority: Context-rich minority oversampling for long-tailed classification. In CVPR, pages 6887-6896, 2022.
Xinyu Peng, Han Li, Yuyang Huang. Ziyang Zheng, Yaoming Wang, Xin Chen, Wenrui Dai, Chenglin Li, Junni Zou, and Hongkai Xiong. Towards holistic modeling for video frame terpolation with auto-regressive diffusion transformers. arXiv preprint arXiv:2601. 00t.:2601.4499 14959, 2026.
Yiming Qin, Huangjie Zheng, Jiangchao Yao, Mingyuan Zhou, and Ya Zhang. Class-balancing diffusion models. CVPR, pages 18434-18443, IEEE, 2023.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with CLIP latents. CoRR, abs/2204.06 125, 2022.
Harsh Rangwani, Konda Reddy Mopuri, and R. Venkatesh Babu. Class balancing GAN with classifier the loop. In Cassio P. de Campos, Marloes H. Maathuis, and Erik Quaeghebeur, editors, UAL, volume 161 of Proceedings of Machine Learning Research, pages 1618-1627. Press, 2021.
Harsh Rangwani, Naman Jaswani, Tejan Karmali, Varun Jampani, and R. Venkatesh Babu. Improv- ing gans for long-taiil tailed data through group spectral regularization. In Shai Avidan, Gabriel J. Brostow, Moustapha Cissé, Giovanni Maria Farinella, and Tal Hassner, editors, ECCV, volume 13675 of Lecture Notes in Computer Science, pages 426 442. Springer, 2022.
William Reed. The pareto, zipf and other power laws. Economics Letters, 74(1):15-19, 2001. ISSN 0765-
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer: High- resolution image synthesis with latent diffusion models. In CVPR, pages 10674-10685. IEEE, 2022.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomed- ical image segmentation. In Nassir Navab, Joachim Hornegger, William M. Wells III, and Ale- jandro F. Frangi, editors, MICCAI, volume 9351 of Lecture Notes in Computer Science, pages 234-241. Springer, 2015.
Tim Salimans, Ian J. Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. In Daniel D. Lee, Masashi Sugiyama, Ulrike Luxburg. Isabelle Guyon, and Roman Garnett, editors, NeurIPS, pages 2226-2234, 2016.
Jiang-Xin Xin Shi, Tong Wei, Zhi Zhou, Jie-Jing Shao, Xin- Han, and Yu-Feng Li. Long-tail learn- ing with foundation model: Heavy fine-tuning hurts. In International Conference on Machine PMLR, 2024.
Learning, pages
Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In Francis Bach and David M. Blei, editors. ICML, volume 37 of JMLR Workshop and Conference Proceedings, pages 2256-2265. JMLR.org, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In ICLR, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 20
Yang Song. Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In ICLR. OpenReview.net, 2b.
14

---

Published as a conference paper at ICLR 2026
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In Andreas Krause, Emma Brumskill, Kyunghyun Cho, Barbara Engelhardd, Sivan Sabato, and Jonathan Scar- lett, editors, ICML, volume 202 of Proceedings of Machine Learning Research, pages 32211- 32252. PMLR, 2023.
Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jonathon Shlens, and Zbigniew Wojna. Re- thinking the inception architecture for computer vision. In CVPR, pages 2818-2826. IEEE Com- puter Society, 2016.
Tijmen Tieleman and Geoffrey Hinton. Neural networks for machine learning lecture 6: Overview of mini-batch gradient descent. http:/ //ww: cs toronto..............oo. wwwoonto.ddddeooo wttp:ttttee//ceeceedeeeeeeeeeeee slidess cidee//letuue wre Lectture_ssddeeeeeeeeeeeeeeeeeee slwdessisssssss seeiidss_ec__.____ ..f.
Jianfeng Wang, Thomas Lukasiewicz, Xiaolin Hu, Jianfei Cai, and Zhenghua Xu. RSG: A simple Computer
but effective module for learning imbalanced datasets. In CVPR, pages Vision Foundation IEEE, 2021a.
Xudong Wang, Long Lian, Zhongqi Miao, Ziwei Liu, and Stella X. Yu. Long-tailed recognition by routing diverse distribution aware experts. ICLR. OpenReview.net, 2021b.
Zitai Wang, Qianqian Xu, Zhiyong Yang, Zhikang Xu, Linchao Zhang, Xiaochun Cao, and Qing- ming Huang. A unified perspective for loss-oriented imbalanced learning via localization. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2025.
Tong Wei and Kai Gan. Towards realistic long-tailed semi-supervised learning: Consistency is all you need. In Proceedings the IEEE/CVF conference on computer vision and pattern recogni- fion, pages 3469-3478, 2023.
Tong Wei, Zhen Mao, Zi-Hao Zhou, Yuanyu Wan, and Min-Ling Zhang, Learning label shift correc- tion for test-agnostic long-tailes recognition. In Forty-first International Conference on Machine Learning. 2024.
Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy Liang, Quoc V. Le, Tengyu Ma, and Adams Wei Yu. Doremi: Optimizing data mixtures speeds up lan- guage model pretraining. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, NeurIPS, 2023.
Divin Yan, Lu Qi, Vincent Tao Hu, Ming-Hsuan Yang, and Meng Tang. Training class-imbalanced diffusion model via overlap optimization, 2024.
Zhiyong Yang, Qianqian Xu, Li Sicong, Zitai Wang. Xiaochun Cao, and Qingming Huang. Dirmixe: Harnessing test agnostic long-tail recognition with hierarchical label variations. IEEE Transac- fions Pattern Analysis and Machine Intelligence, 2026.
Yang You, Jing Li, Sashank J. Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training BERT in 76 minutes. In ICLR. OpenReview.net, 2020.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In Richard C. Wilson, Edwin R. Hancock, and William A. P. Smith, editors, BMVC. BMVA Press, 2016.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. arXiv preprint 1rXiv:1710.09412. 2017.
Ruichi Zhang, Jiacheng Yang. Ye Liu, Chikai Shang, Mengke Li, Jiang-Xin Shi, Pengkun Wang, Yanbiao Ma, Bo Han, Bo Cheung, and Yang Lu. Pop: A pipeline-oriented paradigm for long- tailed learning survey. TechRxiv preprint, 2025.
Tianjiao Zhang, Huangjic Zheng, Jiangchao Yao, Xiangfeng Wang, Mingyuan Zhou, Ya Zhang, and Yanfeng Wang. Long-tailed tailed diffusion models with oriented calibration. In ICLR, 2024.
Yifan Zhang, Bryan Hooi, Lanqing Hong, and Jiashi Feng. Self-supervised aggregation of diverse experts for test-agnostic long-tailed recognition. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, NewrIPS, 2022.
15

---

Published as a conference paper at ICLR 2026
Yifan Zhang, Bingyi Kang, Bryan Hooi, Shuicheng Yan, and Jiashi Feng. Deep long-tail tailed learning: A survey. IEEE Trans. Pattern Anal. Mach. Intell., 2023.
Boyan Zhou, Quan Cui, Xiu-Shen Wei, and Zhao-Min Chen. Bbn: Bilateral-branch network with cumulative learning for long-tailed visual recognition. In CVPR, pages 9719-9728, 2020.
Zhihan Zhou, Jiangchao Yao, Feng Hong. Ya Zhang, Bo Han, and Yanfeng Wang. Combating representation learning disparity with geometric harmonization. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, NewrIPS, 2023a.
Zhihan Zhou, Jiangchao Yao, Feng Hong, Ya Zhang, Bo Han, and Yanfeng Wang. Combating representation learning disparity with geometric harmonization. In NeurlPS, 2023b.
16

---

Published as a conference paper at ICLR 2026
APPENDIX: IMPROVING DIFFUSION MODELS FOR CLASS-IMBALANCEN TRAINING DATA VIA CAPACITY MANIPULATION
CONTENTS
Backgrounds on Diffusion Models
19
B Related Work
19
C Algorithm Pseudocode
20
D Brief Introduction to the Objectives of CBDM and OC
21
More Discussion
E.1
Definition and Discussion of Model Capacity
21
E,2 Comparison with ensemble-base based imbalanced classification methods.
22
E.3
Conceptual Distinctions...
22
E,4
Comparison with MoE-style Baseline (Group-Exxe LoRA)
22
Theoretical Supplement
22
F.1
Problem Setup for Theorem 2.1 and Theorem 3.1
22
F.2 Proof of Theorem 2.1
23
F3 Proof of Theorem 3.1
24
F.4 Multi-class Extensions of Theorem 2.1 and Theorem 3.1
25
Experimental Supplement
26
G.1 Implementation details
26
G.2 More Results about Fig. 1(b)
26
................................
G.3 Standard Deviations
27
G.4 Human Perceptual Study
................................
28
G.5 Downstream Utility Assessment
29
G.6 More Ablations
................................
29
G.7 Visual Analysis of Capacity Allocation
30
G.8 More Visualization
30
G.9 Broader Settings
38
G.10 Comparison with Overlap Optimization
38
G.11 Extension to Discriminative Tasks
39
G.12 12 On Reproducibility
39
17

---

Published as a conference paper at ICLR 2026
H Social Impact
39
Limitations, Discussions, and Future Work
40
Use Large Language Models (LLMs)
40
18

---

Published as a conference paper at ICLR 2026
A BACKGROUNDS ON DIFFUSION MODELS
We briefly review discrete-time diffusion models, specifically denoising diffusion probabilistic mod- els (DDPMs) (Ho al., 2020). Given a random variable x E X' and a forward diffusion process on x defined as X1:7 == X1.-.., Xy with T E N+, the Markov transition probability from x;-1 to X is ((x)) Nxxx vI - 1, B,I), where Xo :- x g(xo), and {B.,,- the variance schedule. The forward process allows us to sample X, at an arbitrary timestep directly from X in a closed form (xx, xo) N (xx:: vā,xo. (1 āg) d,) where a == 1 - -e and āg := II Ot.. The variance schedule is prescribed such that xT nearly an isotropic Gaussian distribution.
Training objective. The reverse process is defined as a Markov chain that aims to approxi- mate (Xu)) by gradually denoising from the standard Gaussian distribution p(xr) (xxx; I): (0(x)-1(x,) = NN (po(xx-1; Pg(x; 1), TT), where (x(x,, -az (x, Vi-ai co(x,, 1)) is pa- rameterized by time-conditioned noise prediction network eg (xex; 1) and , or are time de-
pendent constants that can be predefined or analytically computed (Bao et al., 2022). The reverse process can be learned by optimizing the variational lower bound on lig-liielihood
log (x) 2 DDI- (a(x)/x))))
+ l
DKL.
--ee,lueleee (xx, E E + C1,
(A.1)
where E 0,) x, = Vā,xa + āāE according to the forward process, W/
03 (1-aa(1-a..] and C1 is typically small and can be dropped (Ho et al., 2020; Song et al., 2021b).
The term LDItt (x, 0) E,, (ov])le((x,, <lli] is called the diffusion loss (Kingma et al., 2021). To benefit sample quality, Ho et al. (2020) apply a simplified training objective by setting wi: = 1.
Class-conditional diffusion models. When the class labels of the training set are available, the class-conditional diffusion model pe (xly) can be parameterized by E(x,, t, y). And the uncondi- tional diffusion model pe(x)) can be viewed as special case with null condition e(x,, t, Null). A similar lower bound on the conditoanii log-likelihood to Eq. (A.1) is
(A.2) where C22 is another small constant and can be dropped (Ho et al., 2020; Song et al., 2021b). The caas--ondiitinal diffusion loss can be written as LDIE y, 8) E.,llwell (x,,
log po (xly)= 2 -E,,, [x,, f, y) - < <|2 + C2.
t, (x,, <]])
B RELATED WORK
Diffusion Models. Inspired by non-equilibrium thermodynamics (Sohl-Dickstein et al., 2015), dif- fusion models have achieved remarkably effective performance at image generation (Ho et al., 2020; Dhariwal and Nichol, 2021; Rombach et al., 2022). Ho et al. (2020) conduct the training of diffusion models using a weighted variational bound. (Song et al., 2021b) propose an alternative method for constructing diffusion models by using a stochastic differential equation (SDE). Karras et al. (2022) introduce design space that clearly outlines the key design choices in previous works. Denois- ing diffusion implicit models (DDIMs) (Song al., 2021a) employs an alternative non-Markovian generation process, enabling faster sampling for diffusion models.
Imbalanced Generation. Addressing class imbalance in generative modeling was initially ex- plored with Generative Adversarial Networks (GANs) (Goodfellow et al., 2014), For instance, CB- GAN (Rangwani et al., 2021) employed a pre-trained classifier to mitigate imbalance. rangwani et al. (2022) attributed performance decline in long-tailed GAN generation to minority-class mode -aypsed by spectral explosion of conditioning parameters- and proposed group spectral regularizer to address it. More recently, with the demonstrated success of diffusion models (DMs), research has focused on their robustness to imbalanced data. CBDM (Qin et al., 2023) applies a distribution adjustment regularizer for minority class augmentation. Yan et al. (2024) introduced contrastive regularization to bolster minority representations, and OC (Zhang et al., 2024) utilizes transfer learning from majority minority classes to improve minority generation quality.
Imbalanced Classification. Imbalanced Classification has been extensively studied to address long- tailed distributions in discriminative tasks (Zhang et al., 2023; 2025; Shi al., 2024; Wei et al., 2024;
19

---

Published as a conference paper at ICLR 2026
Wei and Gan, 2023). Existing literature generally tackles this challenge through three paradigms. At the data level, strategies aim to mitigate imbalance through se-saiining or feature synthesis, ranging from classic techniques like SMOTE (Chawla et al., 2002) to advanced data augmentations such as Mixup (Zhang et al., 2017) and Context-Rich Minority Oversampling (Park et al., 2022). At the model level, research focuses on architectural innovations (Yang et al., 2026; Luo al., 2025) and feature enhancement (Li et al., 2024; Hong et al., 2024; Bao et al., 2025); a prominent milestone is the two-stage decoupling of representation learning and classifier training (Kang et al., 2020), which inspired dual-branch architectures like BBN (Zhou et al., 2020) to balance generalized feature ex- traction with unbiased classification. Finally, at the objective level, methods directly penalize ma- jority class dominance calibrate output confidence. This encompasses cost-sensitive re- weighting approaches like Focal Loss (Lin et al., 2017a) and Class-Balances Loss (Cui al., 2019a), as well as margin-based calibration (Cao et al., 2019b; Hong et al., 2026; Zhou et al., 2023b; Wang et al., 2025) and rigorous Logit Adjustment frameworks (Menon et al., 202 lb; Li et al., 2022; Hong et al., 2023). While these methods are inherently designed for discriminative tasks, their underlying principles of distribution smoothing and bias calibration deeply inspire our approach to imbalanced generation.
Reweighting and Data Mixture Optimization. Beyond loss design for specific imbalanced classes, recent research has explored reweighting strategies and data mixture optimization to enhance gen- erative models (Xie et al., 2023; Fan al., 2024; Kim et al., 2024; Li et al., 2025; Liu al., 2025). In the context of large language models, methods like DoReMi (Xie et al., 2023) and RegMix (Liu al., 2025) optimize the weights of data domains to improve pretraining efficiency and perfor- mance. DoGE (Fan et al., 2024) further introduces domain reweighting based on generalization estimation. For diffusion models, Kim et al. (2024) propose importance weighting to train unbiased models from biased datasets, while Li et al. (2025) combine structural pruning with data reweight- ing for efficient training. While these approaches focus on adjusting the importance or sampling frequency of training data, our proposed Capacity Manipulation (CM) takes orthogonal approach by explicitly reserving and allocating learnable parameters (model capacity) for minority concepts, ensuring they not overwhelmed by majority updates.
C ALGORITHM PSEUDOCODE
We summarize the procedure of our CM in Algorithm C.I, where use DDPM as the base loss, employ DDPM for sampling, and illustrate the process in a sample-wise manner as an example.
Algorithm C.1 Algorithm of CM
D Training, take DDPM as base, sample-wise
Initialize: 0 (W7. W.,...H. J. 00 = (BqA%, .
repeat
Sample data (x, y) E D
Sample a timestep t ~ Uniform (1....,T}})
Sample noise ~ IO.
Base loss: Ltase 60o V + (1 ), 1, El3
Capacity allocation loss: LCM (wy wbr )l eaxeo- (X1, 1. E (x.,xx,,,)) v)13
Take gradient descent on V g= (Lrase + ALCM until converged
>> Sampling, take DDPM for example, sasple-wise
Merge model parameters as f = 89 @ 8
Sample xT 1) for \ T. , do
N (0, if > else = 0
X voc (X+ vi co Xx, 1, v)) + o.z
end for return XD
20

---

Published as a conference paper at ICLR 2026
D BRIEF INTRODUCTION TO THE OBJECTIVES OF CBDM AND OC
Here, we briefly introduce the objective functions of CBDM (Qin et al., 2023) and OC (Zhang et al., 2024).
CBDM.
The objective loss function of CBDM is defined as:
LCBDM(D, N
Σ
(x,, - el12 + rL
Σ (11e((x,,,,)) V'EY
DI
X
-sg(eo(X),, y))1 7/lsg((e t, 8)) - Ex 3/)113
(D.1)
where sg(.) denotes the stop gradient operation: T and y are hyperparameters. | the class number. CBDM introduces regularizer to the diffusion loss to balance the generation quality across different classes.
OC The objective loss function of CBDM is defined as:
Loc(D,
Σ Ee, [llee(x,,t,)) E112 X,DED
(D.2)
where é represents knowledge transfer term within batch of samples. For any given (x, y) D and its corresponding randomly sampled and / during training, OC selects a reference sample (&, 5) from the same batch. The term é is then defined as:
C, X.
if (()) 8,
==
(D.3)
q() 2 4()).
Va.
Here, the probability of any (&, () E Batch(x, y) (where Batch (x, y) represents the batch containing the sample (x, y)) being selected as the reference sample is given by:
((xe|8. ù)
(x 9) Ex.r/ (EBac)((x)))) (xx(
(D.4)
By setting Lhase to LOBDM or Loc, our method CM can be seamlessly integrated with these im- proved objective functions.
E MORE DISCUSSION
E. DEFINITION AND DISCUSSION MODEL CAPACITY
To clarify the concept of "capacity" used throughout the paper, we define as the available rep- resentational resources of the neural network to capture and store data features. We operationalize this concept from two complementary perspectives:
Following the principles Low-Rank Adaptation (Hu et al., 2022), we treat the rank of weight matrices direct measure of capacity. A full-rank ma- trix possesses maximum capacity, whereas low-rank decomposition restricts information flow. Our method explicitly manipulates this by reserving specific rank in 0 for minor- ity expertise.
Structural Perspective (Rank):
Optimization Perspective (Gradient Dominance): Theoretically, capacity allocation is de- termined by which classes dominate the parameter updates. As proven in Theorem 2.1, majority classes contribute significantly larger aggregated gradients, thereby dominating the update direction AW and effectively monopolizing the parameters.
In Sec. 2.2, we utilize the L1-norm magnitude solely as a diagnostic proxy rather than a definition. Since parameters with larger magnitudes typically contribute more to the model's output, the sitivity minority classes to the pruning of magnitue parameters empirically indicates that they forced into "weak" connections (marginal capacity).
21

---

Published as a conference paper at ICLR 2026
E.2 COMPARISON WITH ENSEMBLE- BASED IMBALANCED CLASSIFICATION METHODS.
Several ensemble-based methods (Cui et al., 2023; Wang et al., 2021b; Zhang et al., 2022) lever- age multiple experts to capture diverse knowledge, achieving strong performance in classification tasks through prediction ensemble. However, most of these methods are tailored for classification networks terms of architecture, training paradigm, and loss functions, making them unsuitable for direct application in diffusion models. While they also involve knowledge allocation, their gain mainly comes from increased capacity and ensemble predictions. Additionally, they often require structural modifications to the network and incur higher inference latency, further limiting applica- bility. In contrast, our method introduces no changes to network structure, does not increase model capacity or inference latency, and enhances imbalanced diffusion models purely through capacity manipulation.
E.3 CONCEPTUAL DISTINCTIONS
To clarify the position of Capacity Manipulation (CM) within the broader landscape of long-tailed learning and efficient tuning, we provide a side-by-side comparison Tab. E.1. Unlike Mixture- of- of-Experts (MoE) Class- -lass-Specific Adapters, which often incur high inference overhead due to dynamic routing or loading specific weights, CM employs joint training paradigm with zero infer- ence overhead (via merged weights). Furthermore, unlike reweighting class-balanced objectives that focus solely on loss design, CM explicitly addresses the bottleneck of model capacity allocation.
Table E.1: Conceptual comparison of CM with related paradigms.
Method
Mechanism
Training Paradigm
Inference Overhead
Reweighting / Resampling
Data/Loss Weighting
Single Stage
None
Class-Balanced Objectives (e.g.. OC)
Loss Design
Single Stage
None
Adapters/ LoRA
Parameter Addition
Usually Fine-tuning
None (if merged) / Low
Mixture-of- Experts (MoE)
Dynamic Routing
Joint Training
High (Routing + Experts)
CM (Ours)
Capacity Reservation
Joint Training
Zero (Merged Weights)
E.4 COMPARISON WITH MOE-STYLE BASELINE (GROUP-EXPERT LORA)
To empirically verify the advantage CM over simple parameter addition or routing strategies, we implemented a "Group-Expert LoRA" baseline under matched compute budgets. Specifically, we assign distinct LoRA adapters to "Many", "Medium" and "Few" class groups, effectively acting as deterministic MoE where the router is determined by the class frequency group. We trained this baseline on Imb. CIFAR-100 (IR = 100) using the same rank budget as CM.
As shown in Tab. F.2, CM outperforms the Group-Expert LoRA by large margin (FID 7.52 vs. 10.06). This indicates that simply assigning experts to different groups (MoE-stylc) is insufficient for addressing the capacity dominance of majority classes in diffusion models. CM's superior per- formance attributed to its explicit capacity reservation mechanism via the proposed capacity ma- nipulation loss.
Table E.2: Comparison with MoE-style Bascline on Imb. CIFAR-100 (IR 100).
Method
Architecture
FID +
DDPM
Standard U-Net
10.16
Stouup-Exper LoRA (MoE-style)
3 Experts (Many/Med/Few)
10.06
CM (Ours)
Capacity Reservation
7.52
F THEORETICAL SUPPLEMENT
F.1 PROBLEM SETUP FOR THEOREM 2.1 AND THBOREM 3.1
We consider a simplified version of training a diffusion model on imbalanced data, involving only two classes y (1,2), where N₁ >> N2. Therefore, the total loss for all samples in the training set
22

---

Published as a conference paper at ICLR 2026
18:
N No
L Σ
Dpin (x, Vs. 0).
(F.1)
For any parameter matrix W 0, its gradient is:
Ni+N
VwL
Σ
W CuaI (x,, Uc,
i=1
(F2)
92
where and 92 denote the gradients contributed by the i-th sample of class y = 1 and the j-th sample of class y W, respectively. Here W is omitted in the notation for brevity.
F.2 PROOF OF THEOREM 2.1
Assumption F.1. For samples the same class y E V, their gradient contributions 9p to follow normal distribution with mean My and variance og i.e., "6 ~N(My o For gradients of different classes, their Gaussian distributions share the same magnitude of mean vectors but differ in direction, while having identical variance magnitudes, i.e.,
|
|| n. pi u2;
of 02 =
where p and are positive constants.
Proof of Theorem 2.1. According to Assumption F.I. Vw C follows a normal distribution with:
E w L = N1 + N2
(F33
VarlV
= (N1 + N2)0a1
The update direction is driven by the majority class y 1 if the angle between the total gradient Vw C and the expected gradient of class 1 (/1) smaller than the angle between G and the mean gradient of class 2 (H2). This condition is equivalent to requiring that Vw C.p1 > Vw LyL., which can be further rewritten as W (pi - (2)
Let Z = Vw C- (f4 - -2)). Since V C follows multivariate normal distribution, Z is a univariate normal random variable. We next compute the mean and variance Z.
B B|Z] = (N1m1 + N2H2) (M - 422) u2)
N111111112 N1/1 1/2 + Nof22 #1 - ~2llle212
NIp²² -N cos Z(m, ((1))2) + Nay² cos Z(ui, (()))))) - NN222
= (Ni N2))p cos ((H1H22))
(F:4)
Var (Z)
(M1 N2) ".TTVar( vw 122
(M1 - (2)T ((N: + N2) (mi - 122)
= (N1 + 12)001114 - 22l1
= 2(N1 + N2)0² 112(1 cos (11,
Since Z ~ (E,2), Var(Z)), the probability that the update direction is driven by the majority class (i.e., P(Z > 0)) is:
E[Z]
II mmaj = P(Z > 0) = +
(F.5)
23

---

Published as a conference paper at ICLR 2026
where (()) is the standard normal CDF, defiend
1
((x)
dt.
(F.6)
2
Substituting the expressions for B[z] and Var (), we have:
E Z
mm = Φ(
v Var Z)
(Ni N2) cos / (111, [12)
/2(N1 + N2) 11 008 / (111 12))
(F.7)
Φ (Ni N2) (1 cos L(11, p2)
NN + N2j0
+ (2a 1)py (N cos (111, 112)
20
where a Ni+No is the majority ratio, Nı + N2.
F.3 PROOF OF THEOREM 3.1
Assumption F.2 (Simplified Overall Update Direction for Low-Rank Modules), For the equation W W W in Eq. (1), where W full-rank learnable parameter and W is a low-rank module with rank r, we model the overall update direction of W using the following simplification:
Vw
W
AW
Or
EV W
Here, the normalization terms are derived from the implicit or explicit gradient normalization com- monly used modern optimizers (Kingma and Ba, 2015; Tieleman and Hinton; Duchi al., 201 You et al., 2020). a, is monotonically increasing function of r, with = when r 0 and or when W is full rank. ar provides simplified model for the low-rank effect, following general trend and strictly satisfying the boundary conditions at 0 and when We is full rank.
Proof of Theorem 3.1.
V
Vw-
Mmj = P(AW (11
11--12 112 > 0) = P((
+ Qr
(MI 122) > 0)
(F88)
EV w
|EV
Note that Vw = Σm (
Σド 06 We define auxiliary variable Vi by replacing all terms in Vw corresponding to class with those from class = 1:
No
-ド=
N= R
V
W
fi
(F99)
This operation further amplifies the dominance probability of class y 1. Therefore:
V
Vwe
Im
P((
Ws
ar
). -(Hz 42) 0)
(F.10)
EV
W
LV
Since Vws NNN,,, No2 we have:
E
V
Wo
N
14
EV
|
Wo
(F.11)
V
V
2
Var
EV
2
Wa
24

---

Published as a conference paper at ICLR 2026
Similarly, Vw NNN we have:
V We
12
E
EV Wel
1
(F.12)
V W
0²
Var
JEVw=|]
N2
Let Z =
Viy L3V we
+ Cr
Vwe [EV
(M - -22,, we have:
E|Z] = lii + ard2
(n1 - H22)
111 a,)((1 cos (111, [2)1
0² at o²
(F.13)
Var (11 -22) (
l) 122
N22 Noj2
at ) (1 cos L(p H2))
2( N
N₂
Since Z ~ (E[Z), Var(Z)), the probability that the update direction is driven by the majority class
is:
maj < P(Z
0)
B|Z]
Var Z)
141 - a1)(1 c08 ((11, H2))
(F.14)
+
2 + a²
)02(11 cos 2(1 H2))
No
Jr (2N cos 2(H2 p2))
< Φ
20
F.4 MULTI-CLASS EXTENSIONS THEOREM 2.1 AND THEOREM 3.1
Theorem F.1 (Multi-class Extension of Theorem 2.1). Consider a training set with K classes, y' (1,2,...... K}, and per-class sample counts N₃ N2 > ... >> Ng. The expected proportion of parameter matrices where the minority class K is overwhelmed by the majority class 1 is
(Ni (NI-Nx)ie Nx) cos ((1,, + [K=2 NaPK) (wn (K)) )
II ②
(F.15)
v 22 cos / (M1,HKK)
where = -EK Nr.
Here, (N1 - Nx) 0, and the cross-term
EER=1
k= N&PAk) (#r - pK) depends on inter-class
similarity. By modeling this inner product as zero-mean random variable, or by assuming mutually orthogonal classes, the bound simplifies to
(Ni NK)R cos Z(H1, px)
II
(F.16)
V2N
This highlights the impact of the low-rank component under severe imbalance, mirroring the struc- tural behavior in the two-class setting.
Theorem F2 (Multi-class Extension of Theorem 3.1), Consider the same K K-class setup. Suppose class y K contributes gradients both W, and W.. while other classes only contribute to W
25

---

Published as a conference paper at ICLR 2026
Sample number
500
0.5
Loss
Sample Number
300
0.3
L,088
100
0.1
0
20
40
60
80
100
Class Index
Figure G.1: Per-class raw losses. The setting is same as in Fig. 1(b)
Then the expected proportion of parameter matrices where class K is overwhelmed by class is bounded by
IT (1 (1-ar)14 (N(() cos ((11,
(F.17)
20
Here, - No. We deliberately omit the dependence on class-wise sample numbers order more clearly isolate the effect of the low-rank component. This demonstrates that the imbalance- mitigating effect also extends to the multi-class setting.
G EXPERIMENTAL SUPPLEMENT
G.1 IMPLEMENTATION DETAILS
Following Ho et al. (2020), we utilize a U-Net (Ronneberger et al., 2015) based on a Wide ResNet (Zagoruyko and Komodakis, 2016) the noise prediction network. We set the hyperpa- rameters for DDPM as 31 = 10 and 3r = 0.02, with maximum timestep T 1000. The Adam optimizer (Kingma and Ba, 2015) is used with betas (0.9. 0.999) and learning rate of 10-1 The dropout rate is set to 0.1. We use batch size of 64 and train the model for 300,000 steps, including warm-up period of 5,000 steps. For the rank of BA in Eq. (1), fix it at To min(d, k). We only apply the low-rank decomposition to the upsampling part of the U-Net, i.e., the latter half of the model, the shallow layers tend to capture more general knowledge (Alzubaidi al., 2021). For the byperparameter À in Eq. (4), we fix it as X = 1. For the base loss in Eq. (4), we adopt the objective function from Zhang et (2024), unless otherwise specified. During inference, new images are generated utilizing the 50-step DDIM solver (Song et al., 2021a).
G.2 MORE RESULTS ABOUT FIG. 1(B)
We provide the explicit definition of the metric used in Fig. 1(b). measure the sensitivity of different classes to capacity reduction, we calculate the Relative Loss Change. Specifically, identify the top 10% of model parameters with the smallest Ll-norms and set them to zero. Let Ciaw be the original training loss for class c, and firuned be the loss after pruning. The metric is defined as:
Relative Loss Change c Gpruned L
(G.I)
Ceaw
This metric highlights that minority classes significantly more sensitive parameter pruning, suggesting they rely on "fragile" parameters.
In addition, showcase the raw losses and absolute per-class loss changes in Figs. G.1 and G.2.
26

---

Published as a conference paper at ICLR 2026
Sample number
Loss change
0.006
Sample Number
400
Loss Change
300
0.002
100
0.000
40
100
Class Index
Figure G.2: Per-class absolute loss changes. The setting same as in Fig. 1(b)
G.3 STANDARD DEVIATIONS
Due to space limitations, we do not include the standard deviations for Tables and 7 in the main text. The complete results with standard deviations are provided in Tables G.4 to G.6.
Table G.I: FIDS (4), KIDs (4), Recalls (1), and ISs (1) of CM and various baseline methods Imb. CIFAR-10 10 and Imb. CIFAR-100 with different imbalance ratios IR = {100, 50}. All results are reported as Mean + Std. Best and second-best results are highlighted.
Dataset
Method
FID
KID
Recall
f
Imb. CIFAR-10
DDPM (Ho et al., 2020)
10.697 0.079
0.0035 = 0.0008
0.47 : 0.01
9.39 1: 0.12
100
(Kamas al., 2020)
0.0029
+RS (Mahajan et al. 2018)
12.332 0.064
0.0037 = 0.0003
0.45 = 0.02
9.25 : 0.08
+Focal (Lin et al., 2017b)
10,842 + 0.134
0.0034 = 0.0001
0.46 0.03
1.111
CBDM (Qin al., 2023)
8.233 + 0,152
0.0026 = 0.0001
0.53 0.53 0.02
9.23
OC (Zhung 2024)
8,390 0.063
0.0027 0.0002
0.52 0.03
9,53 : 0.12
CM
7.727 + 0.124
0.0023
0.0001
0.53 + 0.01
9.52 + 0.10
Imb. CIFAR-10
DDPM (Ho et al., 2020)
10,216 - 0.138
0.0035 = 0.0002
0.47 = 0.03
9.37 0.133
+ADA (Karras al., 2020)
9.132 0.215
0.0030 = 0.0002
0.51 0.04
9.28 0.21
+RS (Mahujun al., 2018)
11.231 0.177
0.0038 0.0002
0.47 1 0.02
9.31 0.14
+Focal (Lin et al., 2017b)
1: 0.263
0.0034 = 0.0003
0.48 = 0.01
9.38 0.23
CBDM (Qin al., 2023)
7.933 +: 0.082
0.0026
0.0002
0.54 0.02
9.42 0.14
(Zhang 2024)
8.034 0.225
0.0027 0.0001
0.53
9.65
CM
0,0024 1 0.0002
0.54 = 0.01
9.69 1 0.09
Imh. CIFAR-100
DDPM (Ho et al., 2020)
10.163 +: 0.077
0.0029 = 0.0005
0.46 = 0.01
13.45 = 0.15
IR 100
+ADA (Karras et al., 2020)
9.482 + 0.125
0.0032 = 0.0002
0.51 + 0.01
12.44 0.16
(Mahajun et al., 2018)
11.432
0.0007
12.12 - 0.18
+Focal (Lin al., 2017b)
0.110
0.0032 0.0004
0.47 0.02
13.07 0.26
CBDM (Qin et al., 2023)
10.051 0.399
0.0036 = 0.0003
0.51 + 0.01
12.35 - 0.12
(Zhung al., 2024)
- 0.233
0.0026 = 0.0002
0.52
13.44 - 0.20
CM
7.519 + 0.132
0.0017 0.0003
0.52 0.02
13.45 0.23
Imb. CIFAR-100
DDPM (Ho al., 2020)
9.363 +: 0.069
0.0032 0.0002
0.47 0.02
14.27
0.22
IR 50
+ADA (Karras et al,, 2020)
8.927 + 0.138
0.0033 + 0.0001
0.51 + 0.02
12.89 + 0.17
+RS (Mahajan et al., 2018)
10.259 0.217
0.0037 = 0.0003
0.47 - 0.03
12.38 - 0.23
+Focal (Lin
2017b)
9.477 1 0.114
0.0034 0.0002
0.49
13.31 0.15
CBDM (Qin al., 2023)
8.946 + 011788
0.0036 + 0.0003
0.55 0.02
12.59 + 0.19
(Zhang al., 2024)
7.188 + 0.274
0.0024 = 0.0002
0.54 + 0.01
13.99 - 0.22
CM
6.732 1: 0.052
0.0021 0.0001
0.55 0.03
0.15
27

---

Published as a conference paper at ICLR 2026
Table G.2: FIDs (U), KIDs (L), and per-class FIDs (4) of CM and baselines on Imb. CelebA-HQ with imbalance ratios IR = {100,50)
Dutuset
Method
Female FID
Male FID
Overall FID J
KID
Imb. CelebA-HQ
DDPM (Ho al., 2020)
7.143 0.147
16.425 + 0.032
8.727 0.126
0.0037
0.0001
IR 100
CBDM (Qin et al., 2023)
7,043 + 0.079
14.273 = 0.183
7,823 +: 0.115
0.0043
0.0002
OC (Zhang et al., 2024)
7.092 + 0.323
13.962 + 0.221
7.871 : 0.237
0.0034
0.0002
CM
6.815
0.316
7.538 0.201
0.0033
0.0002
Imb. CelebA-HQ
DDPM (Ho et al., 2020)
7.348 +: 0.219
14,808 +: 0.152
8.007 +: 0.265
0.0034 = 0.0002
IR 50
CBDM (Qin et al., 2023)
7.317 + 0.273
12.592 + 0.181
7,423 + 0.139
0.0042 0.0001
OC (Zhang et al., 2024)
7,283 + 0.226
12.938
+ 0.277
7.438 + 0.247
0.0034 = 0.0003
CM
7.147 0.182
11.273 0.146
7.193 +: 0.282
0.0033
0.0002
Table G.3: FIDs (L), KIDs (L), Recalls (1), and ISs (1) of CM and various baseline methods on Imb. ArtBench-10 with imbalance ratios IR {100, 50} using LoRA to fine-tune Stable Diffusion.
Dataset
Method
FID
KID
Recall
IS
Imb. ArtBench-10
DDPM (Ho et al. 2020)
27.083 + 0.438
0.0142 + 0.0003
0.39 + 0.01
8.47 = 0.19
IR 100
CBDM (Qin et al., 2023)
25.723 1 0.263
0,0122 1 0.0002
0.43 1 0.01
7.97 = 0.22
(Zhang et al., 2024)
24.315 + 0.162
0.0106 + 0.0005
0.42 + 0.01
8.71 = 0.20
CM
22.776 0.078
0.0087 0.0002
00.44 0.02
8.71 0.18
Imh. ArtBench-10
DDPM et al., 2020)
25.557 0.082
0.0134
0.0004
0.39 0.02
8.41 0.15
IR 50
CBDM (Qin et al., 2023)
24.487 = 0.153
0.0114 + 0.0002
0.43 + 0.02
8.03 = 0.23
OC (Zhang ot al., 2024)
23.287 + 0,232
0.0097 + 0.0003
0.43 + 0.02
8.48 = 0.17
CM
21.733 0.153
0.0080 0.0002
0.44 + 0.01
8.51 0.23
Table G.5: Per-split FIDs and overall FIDs (4, Mean + Std) of DDPM, CM (09), and CM on Imb. 110 with imbalance ratio IR = 100. Many, Medium, and Few are the three splits based on the training imbalance. Best results are highlighted.
Dataset
Method
Many FID 1
FID
Few FID
Overall FID
Imb. CIFAR-100
DDPM (Ho al., 2020)
14.068 +: 0.193
15.660 0.047
22.188 +: 0.241
10.163 0.077
IR 100
CM (09)
11.923 + 0.139
14,872 + 0.157
29.357 + 0.318
13:712 + 0.240
CM (0 09 @ 09)
11.713 0.247
13.043 1: 0.138
18.729 1 0.141
7.519 1: 0.132
Table G.6: FIDs (L), KIDs (L), Recalls (1), and ISs (f) Imb. CIFAR-100 with imbalance ratio IR. = 100. The last two rows show the results of CM after removing Lcon and L respectively. Best and second-best results are highlighted.
Dataset
Method
FID
KID
Recall
IS 1
Imh. CIFAR-100
DDPM (Ho et al., 2020)
10.163 + 0.077
0.0029
0.0005
0.46 + 0.01
13.45 + 0.15
IR 100
OC (Zhang al., 2024)
8,309 + 0.233
0.0026
0.0002
0.52 1: 0.02
13.44 + 0.20
CM
7,519 : 0.132
0.0017 : 0.0003
0.52 : 0.02
13.45 : 0.23
CM w/o Acon
8.412 1 0.227
0.0029 1 0.0002
0.50 1 0.01
13.23 0.22
CM w/o CDw
8.073 0.174
0.0025 + 0.0001
0.51 + 0.01
13.42 + 0.16
G.4 HUMAN PERCEPTUAL STUDY
We conducted user study focusing the minority classes of Imb. ArtBench-10 (IR=100). Eight human evaluators participated, each assessing between 20 and 36 image groups each group, provided 10 real images (explicitly identified) together with generated images: one produced by the baseline (DDPM) and the other by our method CM, presented random order. All images within group belonged to the same class. Evaluators asked to select the generated image that
28

---

Published as a conference paper at ICLR 2026
Table G.4: Per-split FIDs (J) CM and baselines on Imb. CIFAR-IO (IR 100) and Imb. CIFAR- 100 (IR = 100). Many, Medium, and Few are the three splits based on the training imbalancedness. Best and second-best results are highlighted.
Imb. CIFAR-10, IR 100
Dataset
Method
Many HID
Med. FID
Few FTD
Overall HID
Imb. CIFAR-10
DDPM (Ho et al., 2020)
14,203 +: 0.221
19.714
0.075
15.869 + 0.152
10.697 + 0.079
100
CBDM (Qin 2023)
12.222 0.087
0.107
8,2334
OC (Zhang al., 2024)
12.026 +: 0.102
15.234
0.108 14.340
12.254 = 0.178
8.3896 = 0,063
CM
11.705 + 0.105
0.093
11.218 0.101
7.7273 + 0.124
Imb. CIFAR-100
DDPM (Ho al,, 2020)
14.068 0.193
15.660 +: 0.047
22. 188 + 0.241
10.163 + 0.077
100
CBDM (Qin 2023)
12.585 0.182
14.042
20.667 0.294
10.051 0.391
OC (Zhang et al., 2024)
11.731 +: 0.221
13.133 0.073
21.053 = 0.371
8,309 + 0.233
CM
11.713 +: 0.247
13.043 = 0.138
18.729 0.141
7.519 0.132
Table G.7: User study results on minority classes of Imb. ArtBench-10 (IR=100). Eight evaluators compared generated images from the baseline (DDPM) and our method CM. Numbers indicate the count of trials where each method was preferred.
Method
Eval.
Eval.
Eval.
Eval.
Eval. 5
Eval.
Eval.
Eval.
Overall
DDPM
10
10
4
6
3
14
12
63
CM
24
26
16
19
18
17
15
16
151
demonstrated higher realism and stylistic fidelity. As shown in Tab. G.7, our method was preferred in 70.56% of the trials, indicating clear perceptual advantage over the baseline.
G.5 DOWNSTREAM UTILITY ASSESSMENT
To evaluate the utility of generated samples, augmented the Imb. CIFAR- 100 (IR=100) dataset for a downstream classification task. ResNet-18 classifier was trained under three settings: (i) augmentation, (ii) augmentation with 50k images generated by the baseline (DDPM), and (iii) augmentation with 50k images generated by our method (CM). As shown Tab. G.8, augmentation with CM-generated images yields a significantly stronger downstream classifier, demonstrating the superior quality and diversity of our minority-clas generations.
G.6 MORE ABLATIONS
In Fig. G.3, we present additional ablation studies on the hyperparameter À across more datasets.
85
5:00
FD
8.0
175
8.0
CM
7.50
CM
7.
OC
OC
7,5
OC
OC
10
13
1813
1812
(a) Imb. CIFAR-10, IR {100, 50}
(b) Imh. CIFAR-100, IR (100, 50)
Figure G.3: Ablation study on the hyperparameter À in Eq. (4), We use OC as a reference because shows the best overall performance among the baselines and serves as our base loss. Figures (a) and (b) show results on Imb. CIFAR- 10 and Imb. CIFAR-100, respectively, with imbalance ratios of IR 100 and IR 50 from left to right. We report FIDs for {0.0, 0.5, 0.8, 1.0, 1.2, 1.5).
29

---

Published as a conference paper at ICLR 2026
Table G.8: Downstream classification performance Imb. CIFAR-100 (IR=100) using ResNet- 18. Augmentation with images generated by CM leads to clear improvements over both the baseline (DDPM) and the no-augmentation setting.
Method
Precision
Recall
No augmentation w/ DDPM 50k images w/ CM 50k images
43.25
37.25
46.73
42.34
49.23
46.78
G.7 VISUAL ANALYSIS OF CAPACITY ALLOCATION
To visually validate the roles of the decomposed parameters, we conducted qualitative ablation study on Imb. ArtBench- 10 (IR = 100). We generated samples for minority classes using only the general branch (89) and compared them with the full model (A 8 8%). As shown in Fig. G.4, samples generated by 8 effectively capture the high-level semantic structure and generic features (e.g., basic object shapes and backgrounds) dominated by majority classes. However, they fail to render the specific artistic styles and fine- fine-grained details unique to the minority classes. In contrast, the full model successfully recovers these minority-specine traits. This empirically confirms that Capacity Manipulation effectively disentangles feature learning: generic knowledge is stored in 89, while minority expertise is preserved in 0
Figure G.4: Visual comparison between using only generic parameters (09, left) and the full model (09 @ 0" right) on the minority class (Realism) of Imb. ArtBench-10. 0 captures generic features but misses minority-specific details, which are restored by 0"
G.8 MORE VISUALIZATION
The generated images for one of the medium classes "surrealism" on Imb. ArtBench-10 10 with IR
100 are shown Fig. G.7. is evident that the generated styles of CM are much closer to the real images. The generated images for the minority class "Male" on Imb. CelebA-HQ with IR 100 are shown in Fig. G.6. More visualization generation results with CM are presented Fig. G.8 (Imb. -EFAR-100, IR. 100), Fig. G.9 (Imb. CelebA-HO, IR 100), Fig. G.11 (ImageNet-LT), Fig. G.12 (iNaturalist) and Fig. G.10 (Imb. ArtBench- 10, IR. = 100).
Diversity in the Tail. To further demonstrate that method avoids mode collapse in minority classes, we provide additional visualizations focusing on the intra-class diversity of the tail. As shown in Fig. G.5, our CM generates minority samples with significant variations, closely mirroring
30

---

Published as a conference paper at ICLR 2026
the diversity of the real data distribution. This aligns with our high Recall scores reported in the main experiments.
A
Figure G.5: Additional visualization intra-class diversity for the minority class (Post Impression- ism, Imb. ArtBench-10, 10, IR 100). (Left) DDPM. (Right) CM. The generated samples exhibit diverse poses, backgrounds, and styles, confirming that CM effectively mitigates mode collapse in the tail.
31

---

Published as a conference paper at ICLR 2026
Imb. CelebA-HQ. =R 100, Male
DDPM
CM
Real Images
Figure G.6: The visualization of generated images on Imb. CelebA-HQ with imbalance ratio IR 100. The figure showcases the generated outputs for the class "Male", which is the minority class, from both DDPM and CM. The last row displays real images from the dataset for reference. is evident that CM generates more realistic and diverse faces.
Imb. ArtBench-10, IR 100, Surrealism
DDPM
CM
Real Images
Figure G.7: The visualization of generated images on Imb. ArtBench- 10 with imbalance ratio IR- 100. The figure showcases the generated outputs for the class "Surrealism", which is one of the medium classes, from both DDPM and CM. The last row displays real images from the dataset for reference. is evident that CM generates results that are significantly more stylistically closer to the real images compared to DDPM. The images shown are randomly selected.
32

---

Published as a conference paper at ICLR 2026
Figure G.8: Visualization of generation results on Imb. CIFAR-100 (IR 100) with CM. The images shown randomly selected.
33

---

Published as a conference paper at ICLR 2026
Figure G.9: Visualization of generation results on Imb. CelebA-HQ (IR. 100) with CM. The images shown randomly selected.
34

---

Published as a conference paper at ICLR 2026
KUUES
Figure 1.10: Visualization of generation results on Imb. ArtBench-10 10 (IR 100) with CM. The images shown randomly selected.
35

---

Published as a conference paper at ICLR 2026
Figure G.11: Visualization generation results on ImageNet-LT with CM. The images shown are randomly selected.
36

---

Published as a conference paper at ICLR 2026
Figure G.12: Visualization generation results on iNaturalist with CM. The images shown are randomly selected.
37

---

Published as a conference paper at ICLR 2026
Table G.9: FID results Consistency Models. Incorporating our method (+CM) consistently
duces FID compared to the baseline.
Dataset (IR=100)
Baseline FID
+CM FID
Imbalanced CIFAR-10
10.523
9.257
Imbalanced CIFAR-100
9.862
8.638
Table G. FID results on DPM-Solver (15 steps). Our method (+CM) improves generation quality over the baseline.
Dataset (IR=100)
Baseline FID
+CM FID
Imbalanced CIFAR-10
10.932
7.623
Imbalanced CIFAR-100
11.524
8.033
Table G.11: FID results latent diffusion models trained from scratch fine- fineggrained text-to-
image datasets. Our method (+CM) improves generation quality over the baseline.
Dataset
Baseline FID
+CM FID
CUB-200
24.62
22.64
Oxford-102
26.47
24.31
G.9 BROADER SETTINGS
Consistency Models. Beyond DDPM, also validated method the Consistency Mod- els (Song et al., 2023). As shown in Tab. G.9, incorporating our method CM consistently improves FID scores across datasets.
Integration with DPM-Solver. We further validated our method with DPM-Solver (Lu et al., 2022) (15 steps). As shown in Tab. 10 incorporating our method (+CM) consistently reduces FII across datasets compared to the baseline.
Latent diffusion models from scratch and fine-grained text-to-image datasets. We evaluated our method on latent diffusion models trained from scratch using fine-grained text-to-image datasets. Specifically, we conducted experiments on CUB-200 and Oxford- 102, creating imbalanced versions by downsampling based on class labels. We followed the network configuration and training details of VQ-Diffusion-S (Gu et al., 2022), using textual descriptions as conditioning for image generation. As shown in Tab. 1.11, our method (+CM) consistently improves FID scores over the baseline.
G.10 COMPARISON WITH OVERLAP OPTIMIZATION
We conduct direct experimental comparison with Overlap Optimization (Yan et al., 2024) on the Imb. CIFAR- 100 (IR 100) dataset. As shown in Tab. G.12, our method (CM) outperforms Overlap Optimization across key metrics, further validating the effectiveness of our capacity manip- ulation strategy.
Table G.12: Comparison with Overlap Optimization on Imb. CIFAR-100 (IR 100).
Method
FID J
Recall
DDPM
10.16
0,46
Overlap Opt. (Yan et al., 2024)
8.98
0.48
CM (Ours)
7.52
0.52
38

---

Published as a conference paper at ICLR 2026
G.11 EXTENSION TO DISCRIMINATIVE TASKS
While this paper focuses on diffusion models, the core concept of Capacity Manipulation (CM) reserving model capacity for minority classes potentially generalizable to dis- criminative long-tailed recognition tasks. To validate this, we conducted preliminary experiments on Long-Tailed CIFAR- 100 with imbalance ratios of IR {200, 100).
Setup. We employ ResNet-188 18 as the backbone. We apply our low-rank decomposition 89 @ 8*) to the convolutional layers. We compare standard Cross-Entropy (CE), LDAM Loss, and Logit Adjustment (LA) against their CM-enhanced versions.
Results. As shown in Tab. G.13, CM yields consistent improvements over the baselines. Notably, at IR = 100, CM improves LA by 1.78% in overall accuracy, with significant gains in Medium and Few classes. This confirms that preventing majority classes from monopolizing the shared feature extractor is beneficial for discriminative tasks as well.
Table G.13: Classification accuracy (%) on Long-Tailed CIFAR-100 using ResNet-18. We compare Cross-Entropy (CH), IDAM, and Logit Adjustment (LA) with their CM-integrated versions under imbalance ratios (IR) of 200 and 100.
Imbalance Ratio
Method
Many
Med.
Few
Overall
CE
70.03
42.13
6.38
36.56
LDAM
62. 13
43.64
16.64
38.66
IR 200
LDAM+ CM
60.20
46.35
20.69
40.50
LA
60.46
45.16
20.43
40.11
LA CM
63.73
46.38
20.16
41.36
CE
70.14
40.54
6.73
40.76
LDAM
61.77
41.57
19.10
41.90
IR = 100
LDAM + CM
62.40
44.94
21.77
44.10
LA
57.91
45.80
23.53
43.36
LA CM
61.60
46.94
23.80
45.14
G.12 ON REPRODUCIBILIIT
All the experiments conducted NVIDIA A100s and H200s with Python 3.8 and Pytorch 2.0.1. We have provided experimental setups and implementation details in Sec. 4.1 and Sec. 4.2. To ensure reproducibility, the source code will be released upon acceptance.
H SOCIAL IMPACT
In this paper, we propose a method to enhance the robustness of generative diffusion models against imbalanced data distributions. This advancement holds significant social implications, both pos- itive and negative. On the positive side, our approach could democratize access to high-quality data generation, allowing marginalized communities benefit from more equitable representation in Al applications. By improving the model's performance on underrepresented classes, we can foster inclusivity in various fields, such as healtheare, finance, and education, where data-driven decisions can impact lives. Conversely, there are potential negative consequences to consider. As generative models become more powerful, they may be misused create deceptive content, leading to misinformation and erosion of trust in digital media. Additionally, our method's emphasis on underrepresented segments the training data poses a risk data poisoning supervision is lack- ing. Malicious actors could exploit this focus to introduce biased or harmful data, compromising the model's integrity. This vulnerability underscores the need for robust monitoring and validation mechanisms ensure data reliability, any compromise could lead to unintended negative conse- quences. Therefore, proactive data governance is essential to mitigate these risks while maximizing the benefits of our proposed method.
39

---

Published as a conference paper at ICLR 2026
I LIMITATIONS, DISCUSSIONS, AND FUTURE WORK
Our work introduces CM as a novel perspective for addressing class imbalance in diffusion models, demonstrating strong empirical results. However, acknowledging the scope and potential extensions of this work is important for contextualizing our findings and guiding future research. We demon- strated CM's efficacy across various image datasets and both faining-from--ee scenarios. Evaluating the applicability and potential adaptations of capacity manipulation for differ- data modalities (e.g., video, 3D data) and other types of generative models remains interesting area for future exploration.
and fine-tuning
While CM demonstrates robustness to extreme imbalance ratios (e.g., TR 500 on iNaturalist), we observe limitations in scenarios of absolute sample scarcity. When the absolute number of minority samples is extremely low (e.g., single-digit samples), the reserved capacity 0 lacks sufficient data to learn meaningful representations, limiting generation quality despite the explicit capacity alloca- tion. Addressing this "few-shot" to "zero-shot" transition remains challenging direction for future research.
J USE OF LARGE LANGUAGE MODELS (LLMs)
For transparency, note that large language models (LLMs) (e.g.. GPT-5 and Gemini 2.5 Pro) were employed during the preparation of this manuscript. The L.LM was used exclusively for language polishing, grammar checking, and formatting. It not involved in generating research ideas, designing methods, analyzing data, or drawing conclusions. All technical contributions are solely the work of the authors.
40


Reviews:
[{"review_id": "y03pIN2wNq", "reviewer": "Reviewer_NJJB", "summary": "This paper addresses the problem of class imbalance in diffusion models, which leads to poor generation performance on minority classes. The authors identify model capacity allocation as a key overlooked factor, where majority classes dominate model parameters, leaving insufficient capacity for minorities. To mitigate this, they propose Capacity Manipulation (CM), a method that reserves model capacity for minority classes via low-rank decomposition of parameters and a novel capacity manipulation loss. The method is orthogonal to existing approaches and does not increase inference cost. Extensive experiments on multiple datasets demonstrate consistent improvements in minority-class generation without sacrificing majority-class performance.", "questions": "(1) The method proposed in the paper primarily focuses on the context of known classes. A natural follow-up question is how Capacity Manipulation would perform in scenarios involving more compositional and fine-grained concepts. For example, in a dataset imbalanced towards \"photos of cats\" vs. \"paintings of dogs,\" how would the model reserve capacity for the minority concept of \"painting\" style, which is orthogonal to the object \"dog\"? Does this framework extend to reserving capacity for concepts rather than just classes? \n\n(2) The method's architecture—using a LoRA-like adapter—makes a strong, implicit assumption: that \"minority expertise\" is inherently low-rank. What is the theoretical or empirical justification for this? One could easily argue the opposite: minority classes might be more complex and have a higher intrinsic dimensionality (e.g., \"impressionist painting\" vs. \"female face\") but are simply under-sampled. If the minority knowledge is, in fact, high-rank, then the fixed low-rank of the adapter would become the primary performance bottleneck, ironically limiting the minority class's capacity more than a standard full-rank model. How does CM cope with this potential issue? \n\n(3) The current formulation appears to use a single $\\theta^e$ to capture the expertise for all minority classes collectively. On datasets with highly heterogeneous minority classes (e.g., the \"Few\" split in Imb. CIFAR-100 or ImageNet-LT, which can contain wildly different concepts), is it plausible that a single low-rank subspace can effectively represent this diverse and multimodal knowledge? Does this not create a new \"capacity collapse\" problem within the minority adapter itself? Have the authors considered a more flexible architecture, such as a Mixture-of-Experts (MoE) model for $\\theta^e$, where different \"experts\" (adapters) are dynamically allocated to different minority clusters?\n\n(4) The paper should discuss and cite relevant literature on reweighting or balancing techniques for generative models [1-5].\n\nReference:\n\n[1] Xie et al. Doremi: Optimizing data mixtures speeds up language model pretraining. NeurIPS, 2023.\n\n[2] Fan et al. DoGE: Domain Reweighting with Generalization Estimation. ICML, 2024.\n\n[3] Kim et al. Training unbiased diffusion models from biased datase. ICLR, 2024.\n\n[4] Li et al. Pruning then Reweighting: Towards Data-Efficient Training of Diffusion Models. ICASSP, 2025.\n\n[5] Liu et al. RegMix: Data Mixture as Regression for Language Model Pre-training. ICLR, 2025.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(1) The method is well-motivated, supported by both empirical observations (e.g., pruning sensitivity) and theoretical analysis (Theorems 2.1 & 3.1). The experimental setup is rigorous, covering multiple datasets, architectures, and metrics.\n\n(2) The paper offers an orthogonal viewpoint on class imbalance in diffusion models by focusing on model capacity allocation, diverging from prior works that primarily emphasize loss reweighting or knowledge transfer (e.g., CBDM and OC). The integration of low-rank decomposition with a tailored loss function represents a creative combination of ideas from parameter-efficient fine-tuning and imbalanced learning for targeted capacity reservation.", "weaknesses": "(1) The term \"capacity\" is not clearly defined. Is it the number of parameters, the magnitude (e.g., L1-norm) of the weights or something else? The pruning experiment suggests a link to weight magnitude, but this connection is not explicitly made or theoretically grounded. Therefore, \" capacity\" remains a somewhat vague concept.\n\n(2) The capacity manipulation loss is designed to force minority-specific knowledge into the low-rank adapter. A potential risk is that this adapter becomes too specialized, failing to leverage the shared, general features learned by the main model. This could limit its ability to generate diverse minority samples that still rely on common underlying features (e.g., a \"rare breed of dog\" should still benefit from general \"dog\" features). The paper does not discuss or analyze this potential limitation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "a7RgfPGjct", "reviewer": "Reviewer_7n4d", "summary": "This paper proposes Capacity Manipulation (CM), a method to improve diffusion models trained on class-imbalanced data. It identifies that majority classes dominate model capacity, limiting minority representation. CM explicitly reserves capacity for minority classes through low-rank parameter decomposition and a capacity manipulation loss that balances consistency and diversity. Experiments on multiple benchmarks show that CM consistently enhances minority-class generation quality and overall robustness.", "questions": "N/A", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is clearly written, well-structured, and easy to follow.\n2. The proposed Capacity Manipulation (CM) method is conceptually simple yet effective, relying on low-rank decomposition and a targeted regularization loss to reserve model capacity for minority expertise.\n3. Theoretical analyses provide solid intuition about how imbalance affects parameter updates and how CM mitigates this effect.\n4. Extensive experiments across small- and large-scale datasets convincingly demonstrate that CM improves minority-class quality without degrading majority-class performance.", "weaknesses": "1. The calculation of loss change in figure 1(b) is not explained.\n2. Although the authors evaluate CM across multiple datasets, there is limited discussion of failure cases or sensitivity to extreme imbalance ratios beyond 100:1.\n3. Some comparisons (e.g., with Overlap Optimization) are only mentioned in passing, a direct experimental comparison would strengthen claims of superiority.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FTiqr7l1BR", "reviewer": "Reviewer_pAbs", "summary": "This paper addresses the poor minority-class performance of diffusion models trained on long-tailed data, arguing that a key culprit is capacity misallocation—majority classes dominate parameter updates and monopolize representational space. \nTo tackle this, it proposes Capacity Manipulation (CM): each weight matrix is decomposed into a general/majority component and a reserved low-rank minority component, and training employs a capacity-manipulation loss that enforces consistency for majority classes while promoting diversity for minority classes. \nAt inference, parameters are merged, introducing no additional latency. \nAcross imbalanced CIFAR-10/100, CelebA-HQ, ImageNet-LT, iNaturalist, and ArtBench-10 (including Stable Diffusion fine-tuning), CM improves FID/KID and delivers especially strong gains on Medium/Few splits over strong baselines (e.g., CBDM, OC), while remaining orthogonal and complementary to them.", "questions": "1. How is the minority/majority split determined, and how sensitive are results to this choice under dataset drift or rebalancing?\n\n2. When merging weights at inference, how do the authors prevent cross-talk between the general and minority subspaces?\n\n3. Does reserving capacity degrade majority-class fidelity or diversity in any regimes?\n\n4. What are the exact training overheads introduced by the extra low-rank factors and CM loss? Do gains persist under tight compute budgets?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper offers a clear and original lens—capacity allocation—and introduces a simple, effective mechanism that reserves low-rank capacity for minority classes, moving beyond reweighting or oversampling.\nMethod quality is strong: the parameter split plus a consistency/diversity loss is minimally invasive, theoretically motivated by gradient/representation analyses, and incurs no inference overhead due to weight merging.\nEmpirically, results are broad and convincing across multiple datasets/backbones (including SD fine-tuning), with especially large gains on Medium/Few splits and stable ablations over ranks and loss weights.\nThe approach is practical and orthogonal to existing long-tail remedies (e.g., CBDM, OC), making it easy to adopt and combine for further improvements.", "weaknesses": "1. The paper should more clearly distinguish CM from class-balanced objectives, reweighting/oversampling, class-specific adapters/LoRA, and Mixture-of-Experts. Add a side-by-side comparison and reproduce at least one adapter/MoE-style baseline under matched compute.\n\n2. The analysis explains majority gradient dominance and motivates reserving rank, but does not specify conditions ensuring no loss of global likelihood or bounds on interference.\n\n3. Most results are class-conditional image benchmarks.\nFor instance, text-to-image (multi-attribute, compositional) and multi-label long-tails are underexplored.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AkjtQger4z", "reviewer": "Reviewer_JwQN", "summary": "This work investigates the challenge of generative modeling for imbalanced datasets. The hypothesis is that the poor generation quality for minority classes is primarily caused by an imbalance in \"model capacity,\" where the model's learning resources are disproportionately occupied by the head (majority) classes. To address this, the paper introduces a novel technique named Capacity Manipulation (CM), which explicitly reallocates and reserves model capacity for the tail (minority) classes. The proposed method employs a low-rank decomposition of the model's parameters, enabling fine-grained control over capacity allocation. A bespoke capacity manipulation loss function is introduced to ensure sufficient capacity is dedicated to learning the features of minority classes, leading to a significant enhancement in their generative representation. The claims are substantiated by comprehensive experimental results, and the overall methodology is presented with clear and coherent logic.", "questions": "Please see the weakness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. I find this approach remarkably novel in how it attributes the class imbalance problem to \"uneven model capacity allocation.\" It represents a significant departure from traditional paradigms like data resampling or loss re-weighting, introducing a fresh perspective by intervening directly at the model parameter level. By the way, I'm also curious if the author's method could be applied to long-tail recognition tasks (e.g., with ResNeXt-50 on CIFAR). No detailed explanation is needed if the implementation is complex—I'm simply wondering about its potential.\n2. The design of  loss function is exceptionally clear in its objective. By creating a \"push-pull\" dynamic between 'consistency' and 'diversity', it effectively channels distinct knowledge into separate parameter subspace.\n3. The paper doesn't just rest on solid experimental results; it also provides theoretical analysis (Theorems 2.1 and 3.1) to substantiate its core thesis: that majority classes indeed dominate parameter updates and that low-rank decomposition can effectively mitigate this dominance.\n4. The experimental validation is remarkably comprehensive. It covers a wide range of datasets (from simple to complex, low-res to high-res), various imbalance ratios, and multiple evaluation metrics, all benchmarked against strong baseline methods.", "weaknesses": "1. I'm also curious if the author's method could be applied to long-tail recognition tasks (e.g., with ResNeXt-50 on CIFAR). No detailed explanation is needed if the implementation is complex—I'm simply wondering about its potential.\n2. My main question is about the capacity 'reservation.' The structure of the parameter decomposition seems to be fixed. This makes me wonder: is this 'hard partitioning' approach truly optimal? Could there be a way for the model to dynamically and adaptively decide how much capacity to allocate to each component during training, rather than relying on a predefined split?\n3. Are there any toy experiments that can visually illustrate this? For instance, using the two-class example you mentioned, could you show how the majority class ends up occupying most of the model's parameter capacity?\n4. I think this assumption has some limitations, especially with varying balance ratios like in ImageNet-LT. For example, in an extreme case with one head class and 999 tail classes, is a single, the setting of rank still appropriate? Or does the rank itself need to be adjusted based on class frequency?\n5. To empirically validate the hypothesis, is it possible to visually demonstrate that the class-specific parameters specialize in learning features unique to minority classes, while the class-agnostic parameters focus on capturing generic features dominated by the majority (head) classes? We propose achieving this through visualization techniques.\n6. My last question is about the diversity within the tail itself. Can you visualize them?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)