Research map / Reinforcement learning / RL from feedback

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

NeurIPS 2026 · Reinforcement learning › RL from feedback · approach: RL fine-tuning of generative models

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes DRIFT, an RL fine-tuning framework for diffusion/flow image models that counters diversity collapse via reward-concentrated sampling, noise-conditioned prompting, and potential-based diversity reward shaping within GRPO.

Problem

RL fine-tuning of image generation models (diffusion/flow) under KL-regularized reward maximization collapses the policy toward a Dirac delta, yielding repetitive outputs and losing diversity.

Approach

Combines three diversity-preserving mechanisms in GRPO: training on a reward-concentrated subset selected by DreamSim neighborhoods, annealed Gaussian noise on prompt embeddings during training, and a potential-based shaping reward equal to intra-group diversity (provably preserving optimal policy), with clipping and decoupled advantages.

Key results

• On SDv1.5 with PickScore, DRIFT reaches DreamSim 0.5698 vs 0.4567 for GARDO and reward 0.9187 vs 0.6176.
• Average gains over baselines reported as 7.19%-93.40% diversity at matched alignment and 13.23%-60.13% reward at matched diversity.
• On SD3.5-M with PickScore, DRIFT improves Generalized Recall by +93.40% over baselines (0.1719 vs 0.1385 for GARDO).

Surprising findings

• Adding the same noise only at test time yields no diversity gain; prompt noise helps only when applied during training.
• CADS, an inference-time method, attains higher raw diversity than DRIFT but with a notable reward drop.

Limitations

• Relies on an auxiliary encoder (DreamSim) for intra-group diversity, adding computational overhead.
• Diversity reward is high-variance and can cause reward hacking, requiring clipping.
• Evaluation limited to SDv1.5 and SD3.5-M with two reward models and small prompt sets (40 prompts).

Experimental setup

Datasets: 45 common animal prompts compositional prompts for SD3.5-M

Benchmarks: DreamSim diversity CLIP diversity Generalized Recall Vendi Score

Metrics: DreamSim diversity CLIP diversity generalized recall Vendi score normalized reward (PickScore/HPSv2)

Baselines: Base model CADS DDPO GRPO GRPO-KL GARDO

Builds on: GRPO DDPO Flow-GRPO CADS GARDO

Closest papers by meaning

• Training Diffusion Models Towards Diverse Image Generation with Reinforcement Learning (CVPR 2024 · similarity 0.90)
• Quality Diversity through Human Feedback: Towards Open-Ended Diversity-Driven Optimization (ICML 2024 · similarity 0.89)
• TMPO: Trajectory Matching Policy Optimization for Diverse and Efficient Diffusion Alignment (NeurIPS 2026 · similarity 0.89)
• VIGOR: Zero-Shot Visual Generalization via Latent-Space Consistency in Model-Based Reinforcement Learning (NeurIPS 2026 · similarity 0.88)
• Towards Better Alignment: Training Diffusion Models with Reinforcement Learning Against Sparse Rewards (CVPR 2025 · similarity 0.87)
• Directly Fine-Tuning Diffusion Models on Differentiable Rewards (ICLR 2024 · similarity 0.87)