Research map / Reinforcement learning / RL from feedback
Reward Fine-Tuning Two-Step Diffusion Models via Learning Differentiable Latent-Space Surrogate Reward
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 LaSRO, which learns a differentiable latent-space surrogate reward (on SDXL's UNet encoder) to reward fine-tune ≤2-step distilled diffusion models, using off-policy exploration and an actor-critic-style loop.
Problem
Reward fine-tuning of step-distilled diffusion models for ultra-fast (≤2-step) image generation is hard for policy-based RL (PPO, DDPO, DPO) due to limited stochasticity, degenerate objectives, and non-smooth two-step mappings.
Approach
Train a latent-space reward model (pre-trained latent diffusion encoder plus CNN head, Bradley-Terry loss on arbitrary black-box rewards) and use its differentiable gradients to fine-tune the two-step student, with off-policy sampling of initial noise and online surrogate adaptation; connected to value-based RL.
Key results
- On MJHQ-30K, LaSRO-fine-tuned SSD-1B-LCM reaches Image Reward 0.998 at 2 steps and 0.946 at 1 step, versus 0.781 for the LCM base and 0.836 for GORS-LCM.
- LaSRO on SDXL-Turbo improves Image Reward at 1 step from 0.817 to 0.904 and at 2 steps from 0.839 to 0.957 (512 resolution).
- RLCM/DDPO and PSO/Diffusion-DPO baselines reach only 0.830 and 0.821 Image Reward at 2 steps, and the paper reports them as unstable.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Fine-tune SSD-1B-LCM with LaSRO | SSD-1B-LCM (1 step) · MJHQ-30K | Image Reward: 0.609 -> 0.946 | Helped |
| Fine-tune SSD-1B-LCM with LaSRO | SSD-1B-LCM (2 steps) · MJHQ-30K | Image Reward: 0.781 -> 0.998 | Helped |
| Fine-tune SDXL-Turbo with LaSRO | SDXL-Turbo (1 step) · MJHQ-30K | Image Reward: 0.817 -> 0.904 | Helped |
| Fine-tune SDXL-Turbo with LaSRO | SDXL-Turbo (2 steps) · MJHQ-30K | Image Reward: 0.839 -> 0.957 | Helped |
| Fine-tune SSD-1B-LCM with LaSRO rather than GORS-LCM | SSD-1B-LCM (2 steps) · MJHQ-30K | Image Reward: 0.836 -> 0.998 | Helped |
| Use LaSRO instead of RLCM/DDPO | SSD-1B-LCM (2 steps) · MJHQ-30K | Image Reward: 0.830 -> 0.998 | Helped |
| Use LaSRO instead of PSO/Diffusion-DPO | SSD-1B-LCM (2 steps) · MJHQ-30K | Image Reward: 0.821 -> 0.998 | Helped |
| Use LaSRO rather than direct Image Reward gradients (ablation) | SSD-1B-LCM · MJHQ-30K | Image Reward | Helped |
| Use LaSRO for reward fine-tuning | SSD-1B-LCM (2 steps) · T2I-CompBench test split | Attribute Binding Score | Helped |
| Use LaSRO for reward fine-tuning | SSD-1B-LCM (2 steps) · MJHQ-30K test split | Text Alignment Score | Helped |
Surprising findings
- Policy-based RL methods (DDPO/RLCM, Diffusion-DPO/PSO) degrade or fail for two-step distilled models, and ablations show that directly using ground-truth Image Reward gradients underperforms the learned surrogate.
Limitations
- Results are reported mainly for LCMs distilled from SSD-1B, with other step-distilled models tested only briefly.
- Reward tuning requires careful tuning to avoid reward overoptimization, handled with qualitative checks and FID-based checkpoint selection.
Experimental setup
Datasets: T2I-prompt-873K (DiffusionDB + open-prompts) T2I-CompBench MJHQ-30K
Benchmarks: MJHQ-30K T2I-prompt-873K T2I-CompBench
Metrics: Image Reward FID Attribute Binding Score Text Alignment Score
Baselines: DDPO/RLCM Diffusion-DPO/PSO GORS-LCM LCM base SDXL-Lightning SDXL-Turbo
Builds on: Latent Consistency Models DDPO Diffusion-DPO Bradley-Terry reward modeling GORS
Closest papers by meaning
- Directly Fine-Tuning Diffusion Models on Differentiable Rewards (ICLR 2024 · similarity 0.90)
- Using Human Feedback to Fine-tune Diffusion Models without Any Reward Model (CVPR 2024 · similarity 0.89)
- Diff-Instruct with Diffused Reward: Towards Principled One-step Generator RL (NeurIPS 2026 · similarity 0.88)
- Reward-Instruct: A Reward-Centric Approach to Fast Photo-Realistic Image Generation (NeurIPS 2025 · similarity 0.88)
- Training Diffusion Models Towards Diverse Image Generation with Reinforcement Learning (CVPR 2024 · similarity 0.88)
- PRDP: Proximal Reward Difference Prediction for Large-Scale Reward Finetuning of Diffusion Models (CVPR 2024 · similarity 0.87)
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.