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

arXiv: https://arxiv.org/abs/2411.15247 · HTML: https://mlsanity.com/paper/2411.15247/

## 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.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)
- LLM alignment / RLHF / preference optimisation / [Diffusion and multimodal model alignment](https://mlsanity.com/alignment/diffusion-and-multimodal-model-alignment/)
- Deep reinforcement learning methods / [RL with diffusion and flow models](https://mlsanity.com/deeprl/rl-with-diffusion-and-flow-models/)

## Builds on (most cited first, 30 of 71)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://mlsanity.com/paper/2204.06125/) (2022)
- [Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding](https://mlsanity.com/paper/2205.11487/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [BERTScore: Evaluating Text Generation with BERT](https://mlsanity.com/paper/1904.09675/) (2019)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis](https://mlsanity.com/paper/2307.01952/) (2023)
- [BLIP: Bootstrapping Language-Image Pre-training for Unified Vision-Language Understanding and Generation](https://mlsanity.com/paper/2201.12086/) (2022)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [LAION-5B: An open large-scale dataset for training next generation image-text models](https://mlsanity.com/paper/2210.08402/) (2022)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [DreamFusion: Text-to-3D using 2D Diffusion](https://mlsanity.com/paper/2209.14988/) (2022)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Video Diffusion Models](https://mlsanity.com/paper/2204.03458/) (2022)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Imagen Video: High Definition Video Generation with Diffusion Models](https://mlsanity.com/paper/2210.02303/) (2022)
- [Group Normalization](https://mlsanity.com/paper/1803.08494/) (2018)

## Built on by (most cited first, 16 of 16)

- [Noise Hypernetworks: Amortizing Test-Time Compute in Diffusion Models](https://mlsanity.com/paper/2508.09968/) (2025)
- [Manifold-Aware Exploration for Reinforcement Learning in Video Generation](https://mlsanity.com/paper/2603.21872/) (2026)
- [Learnable Sparsity for Vision Generative Models](https://mlsanity.com/paper/2412.02852/) (2024)
- [MIRA: Towards Mitigating Reward Hacking in Inference-Time Alignment of T2I Diffusion Models](https://mlsanity.com/paper/2510.01549/) (2025)
- [Beyond the Dirac Delta: Mitigating Diversity Collapse in Reinforcement Fine-Tuning for Versatile Image Generation](https://mlsanity.com/paper/2601.12401/) (2026)
- [Finite Difference Flow Optimization for RL Post-Training of Text-to-Image Models](https://mlsanity.com/paper/2603.12893/) (2026)
- [Faster and Better Alignment for Flow Matching Models via Step-aware Advantages](https://mlsanity.com/paper/2602.01591/) (2026)
- [DreamLite: A Lightweight On-Device Unified Model for Image Generation and Editing](https://mlsanity.com/paper/2603.28713/) (2026)
- [VARD: Efficient and Dense Fine-Tuning for Diffusion Models with Value-based RL](https://mlsanity.com/paper/2505.15791/) (2025)
- [Alignment and Safety of Diffusion Models via Reinforcement Learning and Reward Modeling: A Survey](https://mlsanity.com/paper/2505.17352/) (2025)
- [MIRO: MultI-Reward cOnditioned pretraining improves T2I quality and efficiency](https://mlsanity.com/paper/2510.25897/) (2025)
- [Rethinking Test Time Scaling for Flow-Matching Generative Models](https://mlsanity.com/paper/2511.22242/) (2025)
- [Rdm: Re-conceptualizing Distribution Matching as a Reward for Diffusion Distillation](https://mlsanity.com/paper/2603.28460/) (2026)
- [Reward-guided Fine-Tuning of One-Step Generative Models via Wasserstein Gradient Flow](https://mlsanity.com/paper/2608.29647/) (2026)
- [ReDiF: Reinforced Distillation for Few Step Diffusion](https://mlsanity.com/paper/2512.22802/) (2025)
- [Selective Timestep Weighting and Advantage-Based Replay for Sample-Efficient Diffusion RLHF](https://mlsanity.com/paper/2607.07693/) (2026)

## 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.