# Reward-guided Fine-Tuning of One-Step Generative Models via Wasserstein Gradient Flow

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 35 references that are in the map.
- First posted on arXiv 30 Aug 2026.

## Builds on (most cited first, 22 of 22)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [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)
- [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)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [On Distillation of Guided Diffusion Models](https://mlsanity.com/paper/2210.03142/) (2022)
- [Flow-GRPO: Training Flow Matching Models via Online RL](https://mlsanity.com/paper/2505.05470/) (2025)
- [Mean Flows for One-step Generative Modeling](https://mlsanity.com/paper/2505.13447/) (2025)
- [Consistency Trajectory Models: Learning Probability Flow ODE Trajectory of Diffusion](https://mlsanity.com/paper/2310.02279/) (2023)
- [One Step Diffusion via Shortcut Models](https://mlsanity.com/paper/2410.12557/) (2024)
- [Aligning Text-to-Image Diffusion Models with Reward Backpropagation](https://mlsanity.com/paper/2310.03739/) (2023)
- [Optimal Entropy-Transport problems and a new Hellinger-Kantorovich distance between positive measures](https://mlsanity.com/paper/1508.07941/) (2015)
- [Generative Modeling with Optimal Transport Maps](https://mlsanity.com/paper/2110.02999/) (2021)
- [Learn Your Reference Model for Real Good Alignment](https://mlsanity.com/paper/2404.09656/) (2024)
- [Online Reward-Weighted Fine-Tuning of Flow Matching with Wasserstein Regularization](https://mlsanity.com/paper/2502.06061/) (2025)
- [Generative Modeling through the Semi-dual Formulation of Unbalanced Optimal Transport](https://mlsanity.com/paper/2305.14777/) (2023)
- [Scalable Wasserstein Gradient Flow for Generative Modeling through Unbalanced Optimal Transport](https://mlsanity.com/paper/2402.05443/) (2024)
- [Reward-Directed Score-Based Diffusion Models via q-Learning](https://mlsanity.com/paper/2409.04832/) (2024)
- [Reward Fine-Tuning Two-Step Diffusion Models via Learning Differentiable Latent-Space Surrogate Reward](https://mlsanity.com/paper/2411.15247/) (2024)
- [Efficient Generative Modeling beyond Memoryless Diffusion via Adjoint Schr"odinger Bridge Matching](https://mlsanity.com/paper/2602.15396/) (2026)

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

- [MEND: RL For Flow Models via Proximal Velocity Matching](https://mlsanity.com/paper/2610.05954/) (2026)

## Abstract (arXiv)

To mitigate the time complexity of generative models, one-step generative models have recently emerged through direct mapping from noise to data in a single forward pass. However, the reward-guided fine-tuning method of one-step generative models remains largely unexplored. To address this, we consider one-step generators from an optimal transport view, investigating Wasserstein Gradient Flow (WGF) for modeling smooth and controlled distributional evolution in probability space. We then propose a novel reward-guided fine-tuning of a one-step generative model via WGF. We derive a practical training method that requires no reward gradients, thereby handling both non-differentiable and differentiable rewards. Moreover, our method provides smooth and stable reward-guided distributional updates while mitigating reward hacking and mode collapse. Experiments on 2D synthetic data, CIFAR-10, and ImageNet 256×256 with diverse rewards, including JPEG (in)compressibility, class probability, Black-and-White and CLIP alignment, show that our method achieves better reward alignment compared to baselines.