Beyond the Dirac Delta: Mitigating Diversity Collapse in Reinforcement Fine-Tuning for Versatile Image Generation

Jinmei Liu; Haoru Li; Zhenhong Sun; Chaofeng Chen; Yatao Bian; Bo Wang; Daoyi Dong; Chunlin Chen

cs.LGJan 18, 2026

Beyond the Dirac Delta: Mitigating Diversity Collapse in Reinforcement Fine-Tuning for Versatile Image Generation

Authors: Jinmei Liu, Haoru Li, Zhenhong Sun, Chaofeng Chen, Yatao Bian, Bo Wang, Daoyi Dong, Chunlin Chen,

Organizations: Nanjing University · Australian National University · Wuhan University · National University of Singapore ·

Abstract

Reinforcement learning (RL) has emerged as a paradigm for fine-tuning large-scale generative models, such as diffusion and flow models, to align with complex human preferences and user-specified tasks. A fundamental limitation remains the curse of diversity collapse, where the objective formulation and optimization landscape inherently collapse the policy to a Dirac delta distribution. To address this challenge, we propose DRIFT (DiveRsity-Incentivized Reinforcement Fine-Tuning for Versatile Image Generation), an innovative framework that systematically incentivizes output diversity throughout the on-policy fine-tuning process, reconciling strong task alignment with high generation diversity to enhance versatility essential for applications that demand diverse candidate generations. We approach the problem across three representative perspectives: i) sampling a reward-concentrated subset that filters out reward outliers to prevent premature collapse; ii) prompting with stochastic variations to expand the conditioning space, and iii) optimization of the intra-group diversity with a potential-based reward shaping mechanism. Experimental results show that DRIFT exhibits clear Pareto dominance in task alignment and generation diversity, achieving 3.88% $\sim$ 81.42% higher diversity at matched alignment and 14.38% $\sim$ 56.45% higher alignment at matched diversity. Our code is available at \textcolor{magenta}{https://github.com/NJU-RL/DRIFT}.

Figures & tables

Figure 1: Image generation models fine-tuned via RL often suffer from diversity collapse , resulting in repetitive outputs with similar subject appearances and scene compositions. Instead, DRIFT maintains both high fidelity and diversity. The three methods used identical prompts and random seeds for generation, with SDv1.5 and SD3.5-M fine-tuned using PickScore as the reward function.

Figure 2: Left : Diversity collapse of the policy distribution

\pi_{\theta}(\bm{x}_{0}|\bm{c})

during GRPO fine-tuning, visualized by kernel density estimation contours of DreamSim embeddings, where the support area shrinks monotonically over training. Right : Comparison of reward-diversity tradeoffs, where DRIFT outperforms GRPO-KL by preserving significantly higher diversity at equivalent reward levels.

Figure 3: Overview of DRIFT, a unified framework for mitigating diversity collapse across sampling, prompting, and optimization.

Figure 4: Reward-diversity comparison between DRIFT and baselines, with SDv1.5 and SD3.5-M fine-tuned using PickScore reward.

Backbone	Algorithm	PickScore					HPSv2
Backbone	Algorithm	DreamSim ( $\uparrow$ )	CLIP ( $\uparrow$ )	Recall ( $\uparrow$ )	Vendi ( $\uparrow$ )	Reward ( $\uparrow$ )	DreamSim ( $\uparrow$ )	CLIP ( $\uparrow$ )	Recall ( $\uparrow$ )	Vendi ( $\uparrow$ )	Reward ( $\uparrow$ )
SDv1.5	Base model	$\displaystyle 0.7476$ [1.5pt] $\pm 0.0121$	$\displaystyle 0.2678$ [1.5pt] $\pm 0.0041$	$\displaystyle 1.0000$ [1.5pt] $\pm 0.0000$	$\displaystyle 78.96$ [1.5pt] $\pm 0.9665$	$\displaystyle 0.2653$ [1.5pt] $\pm 0.01064$	$\displaystyle 0.7476$ [1.5pt] $\pm 0.0121$	$\displaystyle 0.2678$ [1.5pt] $\pm 0.0041$	$\displaystyle 1.0000$ [1.5pt] $\pm 0.0000$	$\displaystyle 78.96$ [1.5pt] $\pm 0.9665$	$\displaystyle 0.3016$ [1.5pt] $\pm 0.01810$
	CADS	$\displaystyle 0.7805$ [1.5pt] $\pm 0.0054$	$\displaystyle 0.2899$ [1.5pt] $\pm 0.0077$	$\displaystyle 0.6614$ [1.5pt] $\pm 0.0080$	$\displaystyle 83.15$ [1.5pt] $\pm 0.3352$	$\displaystyle 0.2323$ [1.5pt] $\pm 0.00648$	$\displaystyle 0.7805$ [1.5pt] $\pm 0.0054$	$\displaystyle 0.2899$ [1.5pt] $\pm 0.0077$	$\displaystyle 0.6614$ [1.5pt] $\pm 0.0080$	$\displaystyle 83.15$ [1.5pt] $\pm 0.3352$	$\displaystyle 0.2629$ [1.5pt] $\pm 0.01251$
	DDPO	$\displaystyle 0.3974$ [1.5pt] $\pm 0.0030$	$\displaystyle 0.1058$ [1.5pt] $\pm 0.0034$	$\displaystyle 0.2233$ [1.5pt] $\pm 0.0105$	$\displaystyle 39.23$ [1.5pt] $\pm 0.1550$	$\displaystyle 0.5851$ [1.5pt] $\pm 0.00886$	$\displaystyle 0.4267$ [1.5pt] $\pm 0.0041$	$\displaystyle 0.1076$ [1.5pt] $\pm 0.0021$	$\displaystyle 0.3303$ [1.5pt] $\pm 0.0072$	$\displaystyle 41.53$ [1.5pt] $\pm 0.2605$	$\displaystyle 0.5400$ [1.5pt] $\pm 0.00408$
	GRPO	$\displaystyle 0.3804$ [1.5pt] $\pm 0.0077$	$\displaystyle 0.1105$ [1.5pt] $\pm 0.0026$	$\displaystyle 0.1989$ [1.5pt] $\pm 0.0170$	$\displaystyle 36.90$ [1.5pt] $\pm 0.4305$	$\displaystyle 0.5927$ [1.5pt] $\pm 0.01331$	$\displaystyle 0.4245$ [1.5pt] $\pm 0.0038$	$\displaystyle 0.1102$ [1.5pt] $\pm 0.0024$	$\displaystyle 0.2883$ [1.5pt] $\pm 0.0090$	$\displaystyle 40.79$ [1.5pt] $\pm 0.3253$	$\displaystyle 0.5409$ [1.5pt] $\pm 0.01487$
	GRPO-KL	$\displaystyle 0.4153$ [1.5pt] $\pm 0.0086$	$\displaystyle 0.1100$ [1.5pt] $\pm 0.0047$	$\displaystyle 0.2945$ [1.5pt] $\pm 0.0095$	$\displaystyle 43.73$ [1.5pt] $\pm 0.4253$	$\displaystyle 0.5503$ [1.5pt] $\pm 0.00739$	$\displaystyle 0.4342$ [1.5pt] $\pm 0.0072$	$\displaystyle 0.1120$ [1.5pt] $\pm 0.0043$	$\displaystyle 0.3496$ [1.5pt] $\pm 0.0154$	$\displaystyle 44.14$ [1.5pt] $\pm 0.3083$	$\displaystyle 0.5961$ [1.5pt] $\pm 0.01160$

Table 1: Comparison of diversity and quality metrics under different reward functions and backbones. Diversity is evaluated using DreamSim, ClipScore Diversity, Generalized Recall, and Vendi Score, while quality is measured by the normalized reward. Percentage improvements represent the average relative gains of DRIFT over the respective baselines in each group under identical settings.

Figure 5: Comparison of reward-diversity tradeoff between reward-concentrated and reward-contrasted sampling, with SDv1.5 fine-tuned using PickScore reward. The record points on the Pareto frontier are collected from checkpoints at intervals during fine-tuning. DG denotes D iversity G ain at equivalent reward levels and RG denotes R eward G ain at equivalent diversity levels (Details in Appendix H ). Both metrics are reported in the histograms (e.g., Figure 6 ) and tables (e.g., Table 1 ).

Figure 6: Left : Qualitative results show that reward-concentrated sampling preserves fidelity while improving diversity when fine-tuning SDv1.5 with PickScore. Right : Quantitative results show higher diversity and reward gain than reward-contrasted sampling.

Figure 7: Comparison of reward-diversity tradeoff for prompting with and without noise, with SDv1.5 fine-tuned using PickScore reward.

Figure 8: Left : Qualitative results show that training with prompting noise preserves fidelity while improving diversity in PickScore fine-tuning of SDv1.5. Right : Quantitative results show higher diversity and reward gain than training without noise.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Method	DreamSim	CLIP	Recall	Vendi	Reward
Full sampling	0.4680	0.1285	0.3639	49.71	0.5363
RCS	0.5822	0.1682	0.3969	52.75	0.7297
Gain (%)	+24.40	+30.89	+9.07	+6.12	+36.06

Appendix

Table 2: Comparison between RCS and full sampling under the same candidate-generation budget on SDv1.5 with PickScore. Higher values indicate better performance.

Hyperparameters	PickScore(SDv1.5)	HPSv2(SDv1.5)	PickScore(SD3.5-M)	HPSv2(SD3.5-M)
Random seed	42	42	42	42
Denoising timesteps ( $T$ )	50	50	16	16
group size (G)	8	8	24	24
Guidance scale	5.0	5.0	4.5	4.5
Policy learning rate	$1\times 10^{-4}$	$1\times 10^{-4}$	$3\times 10^{-4}$	$3\times 10^{-4}$
Policy clipping range	$1\times 10^{-4}$	$1\times 10^{-4}$	$1\times 10^{-4}$	$1\times 10^{-4}$

Appendix

Table 3: List of hyperparameter configurations for PickScore and HPSv2.

Figure 11: Reward-diversity comparison between DRIFT and baselines, with SDv1.5 and SD3.5-M fine-tuned using HPSv2 reward.

Figure 12: Qualitative diversity comparisons show that baseline methods suffer from diversity collapse, producing repetitive samples with similar breeds, poses, and backgrounds, whereas DRIFT maintains high fidelity with substantially greater diversity. All models are fine-tuned on SDv1.5 using HPSv2 as the reward function.

Figure 13: Qualitative diversity comparisons show that baseline methods suffer from diversity collapse, producing repetitive samples with similar visual attributes and scene compositions, whereas DRIFT maintains high fidelity with substantially greater diversity. All models are fine-tuned on SD3.5-M using HPSv2 as the reward function.

SDv1.5 (100 prompts)
Algorithm	DreamSim ( $\uparrow$ )	CLIP ( $\uparrow$ )	Recall ( $\uparrow$ )	Vendi ( $\uparrow$ )	Reward ( $\uparrow$ )	FID ( $\downarrow$ )
Base model	$\displaystyle 0.7508/0.7508$ [1.5pt] $\pm 0.0081/\pm 0.0081$	$\displaystyle 0.2753/0.2753$ [1.5pt] $\pm 0.0036/\pm 0.0036$	$\displaystyle 1.0000/1.0000$ [1.5pt] $\pm 0.0000/\pm 0.0000$	$\displaystyle 94.11/94.11$ [1.5pt] $\pm 0.7093/\pm 0.7093$	$\displaystyle 0.2969/0.2755$ [1.5pt] $\pm 0.01237/\pm 0.00763$	$\displaystyle 0.00/0.00$ [1.5pt] $\pm 0.0000/\pm 0.0000$
CADS	$\displaystyle 0.7771/0.7771$ [1.5pt] $\pm 0.0032/\pm 0.0032$	$\displaystyle 0.2958/0.2958$ [1.5pt] $\pm 0.0031/\pm 0.0031$	$\displaystyle 0.6532/0.6532$ [1.5pt] $\pm 0.0046/\pm 0.0046$	$\displaystyle 97.36/97.36$ [1.5pt] $\pm 0.4540/\pm 0.4540$	$\displaystyle 0.2787/0.2400$ [1.5pt] $\pm 0.01275/\pm 0.00585$	$\displaystyle 8.26/8.26$ [1.5pt] $\pm 0.3935/\pm 0.3935$
DDPO	$\displaystyle 0.4266/0.4255$ [1.5pt] $\pm 0.0074/\pm 0.0078$	$\displaystyle 0.1345/0.1222$ [1.5pt] $\pm 0.0039/\pm 0.0027$	$\displaystyle 0.2201/0.2919$ [1.5pt] $\pm 0.0100/\pm 0.0149$	$\displaystyle 47.74/46.49$ [1.5pt] $\pm 0.5386/\pm 0.6237$	$\displaystyle 0.6363/0.6577$ [1.5pt] $\pm 0.01185/\pm 0.00761$	$\displaystyle 24.21/22.35$ [1.5pt] $\pm 0.9564/\pm 0.7033$
GRPO	$\displaystyle 0.4143/0.4324$ [1.5pt] $\pm 0.0047/\pm 0.0056$	$\displaystyle 0.1345/0.1255$ [1.5pt] $\pm 0.0023/\pm 0.0021$	$\displaystyle 0.2201/0.2781$ [1.5pt] $\pm 0.0063/\pm 0.0129$	$\displaystyle 43.67/47.49$ [1.5pt] $\pm 0.3499/\pm 0.3248$	$\displaystyle 0.7030/0.6711$ [1.5pt] $\pm 0.00883/\pm 0.00258$	$\displaystyle 23.71/20.03$

Appendix

Table 4: Evaluation on expanded prompt sets with 100 prompts for SDv1.5 and 200 prompts for SD3.5-M. Each entry reports PickScore/HPSv2, with mean and standard deviation on separate lines. Higher values indicate better performance except for FID. Percentage improvements are averaged over the original fine-tuning baselines, excluding the base model and CADS.

RL for Image Generation

Reinforcement learning for image generation trains visual models with reward feedback to improve prompt alignment, quality, and control.

See all 93 papers

DivRL: Disentangled Self-Similarity Rewards for Diverse Subject-Driven Generation
Jun 22, 2026Qian Wang, Zhenyu Li, Abdelrahman Eldesokey +1Subject-Driven Image Generation Identity-Preserving Image Generation
KAUST, Saudi Arabia
Policy-based Tuning of Autoregressive Image Models with Instance- and Distribution-Level Rewards
Mar 24, 2026Orhun Bugra Baran, Melih Kandemir, Ramazan Gokberk CinbisImage Generation Autoregressive Image Generation
Department of Computer Engineering, Middle East Technical University, Ankara, Türkiye · Department of Mathematics and Computer Science (IMADA), University of Southern Denmark, Odense, Denmark · ROMER Robotics and AI Center, Middle East Technical University, Ankara, Türkiye
TPSO: Training-Free Diverse Image Generation via Semantic Prompt Embedding Optimization
Nov 25, 2025Debin Meng, Chen Jin, Zheng Gao +3T2I Generation T2I Diffusion Models
Queen Mary University of London · Centre for AI, AstraZeneca, UK · University of Bedfordshire +1

Jun 22, 2026J/K move · Enter open · S save

DivRL: Disentangled Self-Similarity Rewards for Diverse Subject-Driven Generation

Qian Wang, Zhenyu Li, Abdelrahman Eldesokey, Peter Wonka

KAUST, Saudi Arabia

Subject-Driven Image Generation Identity-Preserving Image Generation RL for Image Generation

Open paper AskPDF arXiv

Subject-driven image generation faces an "Identity-Diversity Paradox", where strong identity preservation often leads to rigid and low-diversity outputs. We propose a post-training framework called DivRL that jointly optimizes identity consistency and structural diversity simultaneously by leveraging disentangled visual features from a robust similarity model. Specifically, we introduce a Negative Self-Similarity Measure (nSSM) to quantify structural diversity, and Visual Semantic Matching (VSM) to evaluate identity consistency. We propose an "Explore-and-Suppress" strategy that treats VSM as a gated constraint: the model freely explores structurally diverse configurations, and only samples that violate the identity threshold are penalized via a quadratic hinge loss. This converts identity preservation from a competing objective into a feasibility constraint, allowing nSSM and VSM to improve jointly. Experiments demonstrate that our method effectively pushes the model to generate both consistent and diverse images and improves structural diversity while maintaining comparable identity consistency through a gated optimization formulation.

Abstract

RL for Image Generation

Explore similar work

DivRL: Disentangled Self-Similarity Rewards for Diverse Subject-Driven Generation

Policy-based Tuning of Autoregressive Image Models with Instance- and Distribution-Level Rewards

TPSO: Training-Free Diverse Image Generation via Semantic Prompt Embedding Optimization