mlsanityData to 14 Oct 2026

Reward Fine-Tuning Two-Step Diffusion Models via Learning Differentiable Latent-Space Surrogate Reward

arXiv 2411.15247 ↗First posted 22 Nov 2024cs.LGcited by 16

On the map: Diffusion models → Diffusion and flow policies for robotics

also in 2 more
Abstract (arXiv)

Recent research has shown that fine-tuning diffusion models (DMs) with arbitrary rewards, including non-differentiable ones, is feasible with reinforcement learning (RL) techniques, enabling flexible model alignment. However, applying existing RL methods to step-distilled DMs is challenging for ultra-fast (≤2-step) image generation. Our analysis suggests several limitations of policy-based RL methods such as PPO or DPO toward this goal. Based on the insights, we propose fine-tuning DMs with learned differentiable surrogate rewards. Our method, named LaSRO, learns surrogate reward models in the latent space of SDXL to convert arbitrary rewards into differentiable ones for effective reward gradient guidance. LaSRO leverages pre-trained latent DMs for reward modeling and tailors reward optimization for ≤2-step image generation with efficient off-policy exploration. LaSRO is effective and stable for improving ultra-fast image generation with different reward objectives, outperforming popular RL methods including DDPO and Diffusion-DPO. We further show LaSRO's connection to value-based RL, providing theoretical insights. See our webpage hrefhttps://sites.google.com/view/lasrohere.

Key facts. Cited by 16 papers in the map (data through 14 Oct 2026). Placed in Diffusion and flow policies for robotics, one of 14 directions in Diffusion models. Builds on 71 of its 86 references that are in the map. First posted on arXiv 22 Nov 2024.

Builds on

71 of 86 references.

Closest ancestors · 12

Same direction, newest first

  1. 2024Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review18 Jul
  2. 2024Deep Reward Supervisions for Tuning Text-to-Image Diffusion Models1 May
  3. 2024RL for Consistency Models: Faster Reward Guided Text-to-Image Generation25 Mar
  4. 2024PRDP: Proximal Reward Difference Prediction for Large-Scale Reward Finetuning of Diffusion Models13 Feb
  5. 2023Using Human Feedback to Fine-tune Diffusion Models without Any Reward Model22 Nov
  6. 2023Diffusion Model Alignment Using Direct Preference Optimization21 Nov
  7. 2023Aligning Text-to-Image Diffusion Models with Reward Backpropagation5 Oct
  8. 2023Directly Fine-Tuning Diffusion Models on Differentiable Rewards29 Sep
  9. 2023DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models25 May
  10. 2023Training Diffusion Models with Reinforcement Learning22 May
  11. 2023Human Preference Score: Better Aligning Text-to-Image Models with Human Preference25 Mar
  12. 2023Diffusion Policy: Visuomotor Policy Learning via Action Diffusion7 Mar

Other references · 59

  1. 2021Learning Transferable Visual Models From Natural Language Supervisioncited by 33,744
  2. 2020Denoising Diffusion Probabilistic Modelscited by 21,076
  3. 2021High-Resolution Image Synthesis with Latent Diffusion Modelscited by 18,863
  4. 2023GPT-4 Technical Reportcited by 18,396
  5. 2017Proximal Policy Optimization Algorithmscited by 15,215
  6. 2021LoRA: Low-Rank Adaptation of Large Language Modelscited by 14,380
  7. 2020Denoising Diffusion Implicit Modelscited by 9,781
  8. 2020Score-Based Generative Modeling through Stochastic Differential Equationscited by 9,122
  9. 2023Direct Preference Optimization: Your Language Model is Secretly a Reward Modelcited by 8,065
  10. 2015Deep Unsupervised Learning using Nonequilibrium Thermodynamicscited by 7,058
  11. 2022Hierarchical Text-Conditional Image Generation with CLIP Latentscited by 6,585
  12. 2022Photorealistic Text-to-Image Diffusion Models with Deep Language Understandingcited by 6,355
38 more
  1. 2022Flow Matching for Generative Modelingcited by 5,702
  2. 2019BERTScore: Evaluating Text Generation with BERTcited by 5,512
  3. 2018Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actorcited by 4,885
  4. 2023SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesiscited by 4,356
  5. 2022BLIP: Bootstrapping Language-Image Pre-training for Unified Vision-Language Understanding and Generationcited by 4,322
  6. 2021Zero-Shot Text-to-Image Generationcited by 4,288
  7. 2022LAION-5B: An open large-scale dataset for training next generation image-text modelscited by 4,152
  8. 2019Generative Modeling by Estimating Gradients of the Data Distributioncited by 4,101
  9. 2022Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flowcited by 3,700
  10. 2013Playing Atari with Deep Reinforcement Learningcited by 3,274
  11. 2022DreamFusion: Text-to-3D using 2D Diffusioncited by 3,025
  12. 2022Video Diffusion Modelscited by 2,305
  13. 2015High-Dimensional Continuous Control Using Generalized Advantage Estimationcited by 2,026
  14. 2022DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Stepscited by 1,979
  15. 2023Consistency Modelscited by 1,822
  16. 2022Imagen Video: High Definition Video Generation with Diffusion Modelscited by 1,778
  17. 2018Group Normalizationcited by 1,524
  18. 2020Sharpness-Aware Minimization for Efficiently Improving Generalizationcited by 1,417
  19. 2020DiffWave: A Versatile Diffusion Model for Audio Synthesiscited by 1,354
  20. 2022Scaling Laws for Reward Model Overoptimizationcited by 1,095
  21. 2022Planning with Diffusion for Flexible Behavior Synthesiscited by 1,061
  22. 2023Pick-a-Pic: An Open Dataset of User Preferences for Text-to-Image Generationcited by 977
  23. 2017Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learningcited by 950
  24. 2023Latent Consistency Models: Synthesizing High-Resolution Images with Few-Step Inferencecited by 846
  25. 2023Adversarial Diffusion Distillationcited by 759
  26. 2023RAFT: Reward rAnked FineTuning for Generative Foundation Model Alignmentcited by 686
  27. 2022LION: Latent Point Diffusion Models for 3D Shape Generationcited by 553
  28. 2023Aligning Text-to-Image Models using Human Feedbackcited by 422
  29. 2023Florence-2: Advancing a Unified Representation for a Variety of Vision Taskscited by 410
  30. 2023InstaFlow: One Step is Enough for High-Quality Diffusion-Based Text-to-Image Generationcited by 395
  31. 2022DiffusionDB: A Large-scale Prompt Gallery Dataset for Text-to-Image Generative Modelscited by 379
  32. 2024SDXL-Lightning: Progressive Adversarial Diffusion Distillationcited by 266
  33. 2019DiffTaichi: Differentiable Programming for Physical Simulationcited by 258
  34. 2024Playground v2.5: Three Insights towards Enhancing Aesthetic Quality in Text-to-Image Generationcited by 256
  35. 2021PlasticineLab: A Soft-Body Manipulation Benchmark with Differentiable Physicscited by 103
  36. 2023LLMScore: Unveiling the Power of Large Language Models in Text-to-Image Synthesis Evaluationcited by 87
  37. 2023Optimizing DDPM Sampling with Shortcut Fine-Tuningcited by 81
  38. 2024Progressive Knowledge Distillation Of Stable Diffusion XL Using Layer Level Losscited by 52

50 most cited of 59 All 59 (JSON)

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2024: 1; 2025: 7; 2026: 8.

16 papers.

Sort:

By area:

  1. 2025Noise Hypernetworks: Amortizing Test-Time Compute in Diffusion Modelscited by 35
  2. 2026Manifold-Aware Exploration for Reinforcement Learning in Video Generationcited by 15
  3. 2024Learnable Sparsity for Vision Generative Modelscited by 9
  4. 2025MIRA: Towards Mitigating Reward Hacking in Inference-Time Alignment of T2I Diffusion Modelscited by 9
  5. 2026Beyond the Dirac Delta: Mitigating Diversity Collapse in Reinforcement Fine-Tuning for Versatile Image Generationcited by 5
  6. 2026Finite Difference Flow Optimization for RL Post-Training of Text-to-Image Modelscited by 5
  7. 2026Faster and Better Alignment for Flow Matching Models via Step-aware Advantagescited by 4
  8. 2026DreamLite: A Lightweight On-Device Unified Model for Image Generation and Editingcited by 4
  9. 2025VARD: Efficient and Dense Fine-Tuning for Diffusion Models with Value-based RLcited by 3
  10. 2025Alignment and Safety of Diffusion Models via Reinforcement Learning and Reward Modeling: A Surveycited by 3
  11. 2025MIRO: MultI-Reward cOnditioned pretraining improves T2I quality and efficiencycited by 1
  12. 2025Rethinking Test Time Scaling for Flow-Matching Generative Modelscited by 1
4 more
  1. 2026Rdm: Re-conceptualizing Distribution Matching as a Reward for Diffusion Distillationcited by 1
  2. 2026Reward-guided Fine-Tuning of One-Step Generative Models via Wasserstein Gradient Flowcited by 1
  3. 2025ReDiF: Reinforced Distillation for Few Step Diffusioncited by 0
  4. 2026Selective Timestep Weighting and Advantage-Based Replay for Sample-Efficient Diffusion RLHFcited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2024Aligning Few-Step Diffusion Models with Dense Reward Difference Learning
  2. 2026TDM-R1: Reinforcing Few-Step Diffusion Models with Non-Differentiable Reward
  3. 2025CoCA: Step-level Reward for Free in RL-based T2I Diffusion Model Fine-tuning
  4. 2026Diffusion Reinforcement Learning via Centered Reward Distillation
  5. 2026CRAFT: Aligning Diffusion Models with Fine-Tuning Is Easier Than You Think
  6. 2026Reward Sharpness-Aware Fine-Tuning for Diffusion Models
  7. 2025Reward-Instruct: A Reward-Centric Approach to Fast Photo-Realistic Image Generation
  8. 2025A Simple and Effective Reinforcement Learning Method for Text-to-Image Diffusion Fine-tuning