Protecting Minorities in Diffusion Models via Capacity Allocation

Feng Hong, Jiangchao Yao, Dongsheng Li, Ya Zhang, Yanfeng Wang

Keywords: Diffusion Models, Imbalance

TL;DR: We propose CALL, a method to enhance the robustness of diffusion models on imbalanced data by protecting model capacity for minority classes.

Abstract: Diffusion models have advanced quickly in image generation. However, their performance declines significantly on the imbalanced data commonly encountered in real-world scenarios. Current research on imbalanced diffusion models focuses on improving the objective function to facilitate knowledge transfer between majorities and minorities, thereby enhancing the generation of minority samples. In this paper, we make the first attempt to address the imbalanced data challenges in diffusion models from the perspective of model capacity. Specifically, majorities occupy most of the model capacity because of their larger representation, consequently restricting the capacity available for minority classes. To tackle this challenge, we propose Protecting Minorities via Capacity ALLocation (CALL). We reserve capacity for minority expertise by low-rank decomposing the model parameters and allocate the corresponding knowledge to the reserved model capacity through a capacity allocation loss function. Extensive experiments demonstrate that our method, which is orthogonal to existing methods, consistently and significantly improves the robustness of diffusion models on imbalanced data.

Primary Area: generative models

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Reciprocal Reviewing: I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 5810