Research map / Image generation / Fast diffusion sampling
Directly Fine-Tuning Diffusion Models on Differentiable Rewards
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 DRaFT, which fine-tunes Stable Diffusion 1.4 (via LoRA) by backpropagating differentiable reward gradients through the sampling chain, with truncated (DRaFT-K) and low-variance (DRaFT-LV) variants.
Problem
Fine-tuning diffusion models to maximize differentiable reward functions (e.g., human preference scores) is sample-inefficient under RL-based methods, which ignore reward gradients.
Approach
Backpropagate the reward gradient through the full diffusion sampling chain (with LoRA and gradient checkpointing), then truncate backprop to the last K steps (DRaFT-K) and average over multiple re-noised samples at K=1 (DRaFT-LV) to improve efficiency and reduce variance.
Key results
- DRaFT-LV achieves an averaged HPSv2 benchmark score of 0.3536 vs 0.2695 for Stable Diffusion 1.4 and 0.3427 for ReFL.
- Gradient-based DRaFT is claimed to maximize the LAION aesthetic score over 200x faster than DDPO in reward queries.
- DRaFT-LV learns roughly 2x faster than DRaFT-1 and ReFL on human preference rewards.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use DRaFT-LV instead of DRaFT-1 (ablation) | Stable Diffusion 1.4 fine-tuned for HPSv2 · HPDv2 averaged | HPSv2 score: 0.3467 -> 0.3536 | Helped |
| Backpropagate through only the last 5 sampling steps instead of all 50 (ablation) | DRaFT fine-tuning of Stable Diffusion 1.4 · HPDv2 averaged | HPSv2 score: 0.3033 -> 0.3500 | Helped |
| Backpropagate through only the last 10 sampling steps instead of all 50 (ablation) | DRaFT fine-tuning of Stable Diffusion 1.4 · HPDv2 averaged | HPSv2 score: 0.3033 -> 0.3431 | Helped |
| Backpropagate through only the last 30 sampling steps instead of all 50 (ablation) | DRaFT fine-tuning of Stable Diffusion 1.4 · HPDv2 averaged | HPSv2 score: 0.3033 -> 0.3116 | Helped |
| Backpropagate through only the last sampling step instead of all 50 (ablation) | DRaFT fine-tuning of Stable Diffusion 1.4 · HPDv2 averaged | HPSv2 score: 0.3033 -> 0.3467 | Helped |
| Use DRaFT-LV rather than DRaFT-1 (ablation) | Stable Diffusion 1.4 fine-tuning · Human preference reward fine-tuning | training efficiency: roughly 2 × faster | Helped |
| Use DRaFT-LV rather than ReFL | Stable Diffusion 1.4 fine-tuning · Human preference reward fine-tuning | training speed: approximately 2 × faster | Helped |
| Use DRaFT instead of RL-based fine-tuning | Stable Diffusion 1.4 · LAION Aesthetic Classifier reward | fine-tuning speed: >\!200× faster | Helped |
| Use smaller K in DRaFT-K (ablation) | Stable Diffusion 1.4 fine-tuned for HPSv2 · HPDv2 benchmark | HPSv2 score | Helped |
| Use PickScore instead of HPSv2 as the fine-tuning reward (ablation) | DRaFT-1 · HPDv2 averaged | HPSv2 score: 0.3467 -> 0.2875 | Hurt |
| Fine-tune using Pick-a-Pic prompts instead of HPDv2 prompts (ablation) | DRaFT-1 · HPDv2 averaged | HPSv2 score: 0.3467 -> 0.3332 | Hurt |
| Use gradient clipping norm 0.001 rather than 1000 (ablation) | DRaFT-50 · HPDv2 prompt dataset with PickScore reward | PickScore reward | Helped |
Surprising findings
- Full backpropagation through all 50 sampling steps (DRaFT-50) performs worse than truncated DRaFT-1/5, with gradients exploding as K increases.
- LoRA applied only to the last K steps still changes images across the whole sampling chain, with adaptation requiring 10-20 steps to appear.
Limitations
- Requires differentiable reward functions; non-differentiable rewards need RL-based methods.
- Reward hacking and over-optimization cause diversity collapse and overgeneralized detailed images.
- Full backpropagation through long sampling chains suffers from exploding gradients.
Experimental setup
Datasets: LAION Aesthetics prompt set (45 animal prompts) HPDv2 prompts Pick-a-Pic prompts DiffusionDB COCO Captions
Benchmarks: Human Preference Score v2 benchmark (HPDv2 test categories: Animation, Concept Art, Painting, Photo) LAION aesthetic score
Metrics: HPSv2 score LAION aesthetic score reward vs reward queries PickScore CLIP score
Baselines: DDPO ReFL DOODL Re-ranking (Best of 16) Prompt engineering Stable Diffusion 1.4
Builds on: ReFL DDPO LoRA DDIM classifier-free guidance
Closest papers by meaning
- Using Human Feedback to Fine-tune Diffusion Models without Any Reward Model (CVPR 2024 · similarity 0.92)
- Reward Score Matching: Unifying Reward-based Fine-tuning for Flow and Diffusion Models (NeurIPS 2026 · similarity 0.90)
- Training Diffusion Models Towards Diverse Image Generation with Reinforcement Learning (CVPR 2024 · similarity 0.90)
- Self-Play Fine-tuning of Diffusion Models for Text-to-image Generation (NeurIPS 2024 · similarity 0.90)
- Reward Fine-Tuning Two-Step Diffusion Models via Learning Differentiable Latent-Space Surrogate Reward (CVPR 2025 · similarity 0.90)
- Ranking-based Preference Optimization for Diffusion Models from Implicit User Feedback (NeurIPS 2025 · similarity 0.89)
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.