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

Jun 22, 2026

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

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.

Qian Wang, Zhenyu Li, Abdelrahman Eldesokey +1

KAUST, Saudi Arabia

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

Ask PDF

Mar 24, 2026

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

Autoregressive (AR) models are highly effective for image generation, yet their standard maximum-likelihood estimation training lacks direct optimization for sample quality and diversity. While reinforcement learning (RL) has been used to align diffusion models, these methods typically suffer from output diversity collapse. Similarly, concurrent RL methods for AR models rely strictly on instance-level rewards, often trading off distributional coverage for quality. To address these limitations, we propose a lightweight RL framework that casts token-based AR synthesis as a Markov Decision Process, optimized via Group Relative Policy Optimization (GRPO). Our core contribution is the introduction of a novel distribution-level Leave-One-Out FID (LOO-FID) reward; by leveraging an exponential moving average of feature moments, it explicitly encourages sample diversity and prevents mode collapse during policy updates. We integrate this with composite instance-level rewards (CLIP and HPSv2) for strict semantic and perceptual fidelity, and stabilize the multi-objective learning with an adaptive entropy regularization term. Extensive experiments on LlamaGen and VQGAN architectures demonstrate clear improvements across standard quality and diversity metrics within only a few hundred tuning iterations. The results also show that the model can be updated to produce competitive samples even without Classifier-Free Guidance, and bypass its 2x inference cost.

Orhun Bugra Baran, Melih Kandemir, Ramazan Gokberk Cinbis

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

Image Generation Autoregressive Image Generation RL for Image Generation

Ask PDF

Nov 25, 2025

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

Image diversity remains a fundamental challenge for text-to-image diffusion models. Low-diversity generation often leads to repetitive outputs, increasing sampling redundancy and hindering both creative exploration and downstream applications. A key factor is the tendency of diffusion models to collapse toward strong modes in the learned distribution. Existing attempts to improve diversity, such as steering-based guidance, often introduce distortions that degrade image quality. To address this issue, we propose Token-Prompt Embedding Space Optimization (TPSO), a training-free and model-agnostic module. TPSO introduces learnable parameters to explore underrepresented regions of the token embedding space, reducing the tendency to repeatedly sample from strong modes of the distribution. Meanwhile, a prompt-level semantic constraint regulates distribution shifts, preventing quality degradation while preserving semantic fidelity. Extensive experiments on MS-COCO across three representative diffusion backbones demonstrate that TPSO substantially improves diversity, boosting performance from 1.10 to 4.18, while maintaining image quality with only a modest inference-time overhead of 3.6% to 8.9%. Code is available at: https://github.com/Open-Debin/TPSO.

Debin Meng, Chen Jin, Zheng Gao +3

Queen Mary University of London · Centre for AI, AstraZeneca, UK · University of Bedfordshire +1

T2I Generation T2I Diffusion Models

Ask PDF

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