Research map / Image generation / Fast diffusion sampling
Training Diffusion Models Towards Diverse Image Generation with Reinforcement Learning
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 Diversity Reward (MMD- and Gaussian-process mutual-information-based image-set rewards with per-image marginal utility) and uses policy-gradient RL (DDPO-style, LoRA) to fine-tune class-conditional and text-to-image diffusion models toward more diverse outputs guided by small reference sets.
Problem
Diffusion models amplify training-data biases and suffer conditional mode collapse, producing low-diversity outputs (e.g., gender-skewed occupations, uniform class-conditional samples under high guidance), and prior debiasing methods are task-specific.
Approach
A set-based Diversity Reward (negative MMD or GP mutual information between generated and small reference image feature sets), converted to per-image rewards via marginal utility, is used as the reward in policy-gradient RL fine-tuning of any diffusion model.
Key results
- On ImageNet 128x128 with guidance scale 4.0, the MMD diversity reward raises Recall from 36.15 to 47.66 while FID improves from 24.49 to 23.42.
- On CIFAR-10 the MMD variant reaches FID 2.98 vs 3.16 for DDPM, comparable to task-specific retraining CBDM (3.03).
- On Stable Diffusion face debiasing, the MMD reward attains KL values around 1e-3 across six attributes, competitive with ITI-GEN (e.g., 1e-6 to 1e-3).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Fine-tune with MMD Diversity Reward | ImageNet-128x128 class-conditional diffusion model · ImageNet-128x128, guidance scales s = 1.0, 2.0, 3.0, 4.0 | Recall: 36.15 -> 47.66; 40.35 -> 49.31; 45.95 -> 49.30; 55.10 -> 59.63 | Helped |
| Fine-tune with GP-MI Diversity Reward | ImageNet-128x128 class-conditional diffusion model · ImageNet-128x128, guidance scales s = 1.0, 2.0, 3.0, 4.0 | Recall: 36.15 -> 45.31; 40.35 -> 46.31; 45.95 -> 47.83; 55.10 -> 56.13 | Helped |
| Fine-tune with MMD Diversity Reward | DDPM · CIFAR-100 | Recall: 0.65 -> 0.69 | Helped |
| Fine-tune with GP-MI Diversity Reward | DDPM · CIFAR-100 | Recall: 0.65 -> 0.67 | Helped |
| Fine-tune with MMD Diversity Reward | DDPM · CIFAR-10 | Recall: 0.64 -> 0.67 | Helped |
| Fine-tune with GP-MI Diversity Reward | DDPM · CIFAR-10 | Recall: 0.64 -> 0.65 | Helped |
| Select subsets using maximum MMD reward rather than random selection | ImageNet-128x128 post-sampling selection · ImageNet-128x128 post-sampling selection | Recall: 40.22 -> 47.66 | Helped |
| Select subsets using minimum MMD reward rather than random selection | ImageNet-128x128 post-sampling selection · ImageNet-128x128 post-sampling selection | Recall: 40.22 -> 27.69 | Hurt |
| Use M = 2 or 3 generated images to estimate Diversity Reward instead of M = 5 (ablation) | ImageNet-128x128 RL fine-tuning · ImageNet-128x128 | Recall | Hurt |
| Increase generated-image count beyond M = 5 (ablation) | ImageNet-128x128 RL fine-tuning · ImageNet-128x128 | Recall | No clear effect |
| Use N = 5 reference images instead of at least N = 30 (ablation) | ImageNet-128x128 RL fine-tuning · ImageNet-128x128 | Recall | Hurt |
| Use N ≥ 30 reference images (ablation) | ImageNet-128x128 RL fine-tuning · ImageNet-128x128 | Recall | Helped |
Surprising findings
- Min-reward subsets show substantially lower diversity than random selection, confirming the reward's sign direction is meaningful.
Limitations
- Reported performance depends on the number of generated (M) and reference (N) images; M=2,3 or N=5 degrade results.
- Requires a curated small set of unbiased reference images for each task, and only some reward variants are differentiable-free.
Experimental setup
Datasets: ImageNet CIFAR-10 CIFAR-100 CelebA
Benchmarks: ImageNet 128x128 CIFAR-10 CIFAR-100 CelebA-based unbiased face generation
Metrics: Recall Precision FID IS KL divergence to uniform attribute distribution
Baselines: Stable Diffusion Ethical Intervention Hard Prompt Searching Prompts Debiasing Custom Diffusion ITI-GEN CBDM DDPM
Builds on: DDPO PPO ImageReward Classifier-guided diffusion (ADM) Latent diffusion / Stable Diffusion
Closest papers by meaning
- Beyond the Dirac Delta: Mitigating Diversity Collapse in Reinforcement Fine-Tuning for Image Generation (NeurIPS 2026 · similarity 0.90)
- Directly Fine-Tuning Diffusion Models on Differentiable Rewards (ICLR 2024 · similarity 0.90)
- Towards Better Alignment: Training Diffusion Models with Reinforcement Learning Against Sparse Rewards (CVPR 2025 · similarity 0.89)
- Training Diffusion Models with Reinforcement Learning (ICLR 2024 · similarity 0.89)
- Using Human Feedback to Fine-tune Diffusion Models without Any Reward Model (CVPR 2024 · similarity 0.89)
- Feedback Efficient Online Fine-Tuning of Diffusion Models (ICML 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.