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

TL;DR

This work addresses diversity collapse in reinforcement learning fine-tuning of diffusion- and flow-based image generators, where optimization can drive the policy toward a Dirac delta and collapse output diversity. It introduces DRIFT, a three-pronged framework spanning reward-concentrated sampling, noise-conditioned prompting, and diversity-aware reward shaping with a potential-based intrinsic reward, aimed at preserving diversity without sacrificing task alignment. The authors prove optimal policy invariance under reward shaping and demonstrate Pareto-dominant improvements in both diversity and alignment across multiple reward models and backbones, including a SDv1.5 baseline with LoRA. The findings indicate that DRIFT enables versatile image generation with richer candidate outputs while maintaining high fidelity, offering a scalable path to deploy RL-fine-tuned generative models across diverse downstream tasks.

Abstract

Reinforcement learning (RL) has emerged as a powerful 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 \textit{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 \textbf{DRIFT} (\textbf{D}ive\textbf{R}sity-\textbf{I}ncentivized Reinforcement \textbf{F}ine-\textbf{T}uning 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) \textbf{sampling} a reward-concentrated subset that filters out reward outliers to prevent premature collapse; ii) \textbf{prompting} with stochastic variations to expand the conditioning space, and iii) \textbf{optimization} of the intra-group diversity with a potential-based reward shaping mechanism. Experimental results show that DRIFT achieves superior Pareto dominance regarding task alignment and generation diversity, yielding a $ 9.08\%\!\sim\! 43.46\%$ increase in diversity at equivalent alignment levels and a $ 59.65\% \!\sim\! 65.86\%$ increase in alignment at equivalent levels of diversity.

Figures (11)

• Figure 1: Image generation models fine-tuned via RL often suffer from diversity collapse, resulting in repetitive outputs with near-identical attributes in breeds, poses, and backgrounds. Instead, DRIFT maintains both high fidelity and diversity. The three sets were sampled using identical seeds and prompts, from fine-tuning SDv1.5 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 effective 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 the sampling, prompting, and optimization.
• Figure 4: Comparison of reward-diversity tradeoff between reward-concentrated and reward-contrast sampling, with SDv1.5 fine-tuned using PickScore reward. The record points in the Pareto frontier are collected from checkpoints at regular intervals during fine-tuning. DG denotes Diversity Gain at equivalent reward levels and RG denotes Reward Gain at equivalent diversity levels. Both metrics are reported in the subsequent histograms (e.g., Figure \ref{['fig:sample_bar']}) and tables (e.g., Table \ref{['tab:diversity_quality']}).
• Figure 5: Qualitative diversity comparisons show that reward-concentrated sampling maintains high fidelity and achieves significantly greater diversity after fine-tuning SDv1.5 using PickScore.

Theorems & Definitions (3)

• Theorem 1: Optimal Policy Invariance
• Theorem 1: Optimal Policy Invariance
• proof