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

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 $\mathit{\text{the}}\mathit{\text{curse}}\mathit{\text{of}}\mathit{\text{diversity}}\mathit{\text{collapse}}$, where the objective formulation and optimization landscape inherently collapse the policy to a Dirac delta distribution. To address this challenge, we propose $\mathbf{\text{DRIFT}}$ ($\mathbf{\text{D}}$ive$\mathbf{\text{R}}$sity-$\mathbf{\text{I}}$ncentivized Reinforcement $\mathbf{\text{F}}$ine-$\mathbf{\text{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) $\mathbf{\text{sampling}}$ a reward-concentrated subset that filters out reward outliers to prevent premature collapse; ii) $\mathbf{\text{prompting}}$ with stochastic variations to expand the conditioning space, and iii) $\mathbf{\text{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\mathrm{\%}\text{\hspace{0.17em}⁣}\sim \text{\hspace{0.17em}⁣}43.46\mathrm{\%}$ increase in diversity at equivalent alignment levels and a $59.65\mathrm{\%}\text{\hspace{0.17em}⁣}\sim \text{\hspace{0.17em}⁣}65.86\mathrm{\%}$ increase in alignment at equivalent levels of diversity.