Designing Reinforcement Learning for Diffusion Models: A Unified Path-Space View
Abstract
Reinforcement learning (RL) post-training provides a direct way to align diffusion models with human preferences and task-specific rewards. However, current RL algorithms for diffusion models remain fragmented: reverse-trajectory methods rely on discretized likelihood ratios, whereas forward-matching methods train on reward-labeled noising versions of the rollout samples. This paper shows that these seemingly different losses arise from a single path-space principle. Starting from the regularized diffusion-RL objective, we use importance sampling between sampling SDEs to obtain an explicit policy-gradient estimator on trajectory space. The estimator contains the stochastic Itô integral underlying Flow-GRPO-type updates; we derive an equivalent variance-reduced value-gradient form that recovers the forward-matching structure of AWM and DiffusionNFT. This identifies the empirical gap between these method families as a variance-reduction effect rather than a difference in RL principle. The derivation yields a unified design space organized by value-gradient estimation, weight functions, and sampling choices. Within this space, we propose a multi-sample KDE value-gradient estimator that reuses rollout groups, together with scale-bounded weight families that retain stable existing recipes while excluding singular ones. Experiments on SD3.5-M and Qwen-Image models validate the variance-reduction explanation and show that the resulting recipe improves over prior diffusion-RL baselines.
Related
- 1D Convection-Diffusion in Porous Media: A Unified Solution from Bounded Domain to Half-Space
- Graph-Convolutional Reinforcement Learning for Finite-Time Stabilization of a Chaotic Chua Circuit via Unified Flexible Terminal Sliding-Mode Control
- A Unified Two-Stage Generative Diffusion Framework for Channel Estimation and Port Selection in Multiuser MIMO-FAS
- A Unified Space-Time Variational Framework for Wave Equations on Differential Forms
- RECOVERY OF A SOURCE OF ANOMALOUS DIFFUSION FROM SPACE-TIME SPARSE DATA
BibTeX
@misc{xu2026designing,
title = {{Designing Reinforcement Learning for Diffusion Models: A Unified Path-Space View}},
author = {Yixian Xu and Yuanrui Zhang and Shengjie Luo and Liwei Wang and Di He},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2608.14430},
eprint = {2608.14430},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2608.14430},
}
RIS
TY - UNPB
TI - Designing Reinforcement Learning for Diffusion Models: A Unified Path-Space View
AU - Yixian Xu
AU - Yuanrui Zhang
AU - Shengjie Luo
AU - Liwei Wang
AU - Di He
PY - 2026
DA - 2026/10/05
PB - arXiv
DO - 10.48550/arxiv.2608.14430
UR - https://arxiv.org/abs/2608.14430
AB - Reinforcement learning (RL) post-training provides a direct way to align diffusion models with human preferences and task-specific rewards. However, current RL algorithms for diffusion models remain fragmented: reverse-trajectory methods rely on discretized likelihood ratios, whereas forward-matching methods train on reward-labeled noising versions of the rollout samples. This paper shows that these seemingly different losses arise from a single path-space principle. Starting from the regularized diffusion-RL objective, we use importance sampling between sampling SDEs to obtain an explicit policy-gradient estimator on trajectory space. The estimator contains the stochastic Itô integral underlying Flow-GRPO-type updates; we derive an equivalent variance-reduced value-gradient form that recovers the forward-matching structure of AWM and DiffusionNFT. This identifies the empirical gap between these method families as a variance-reduction effect rather than a difference in RL principle. The derivation yields a unified design space organized by value-gradient estimation, weight functions, and sampling choices. Within this space, we propose a multi-sample KDE value-gradient estimator that reuses rollout groups, together with scale-bounded weight families that retain stable existing recipes while excluding singular ones. Experiments on SD3.5-M and Qwen-Image models validate the variance-reduction explanation and show that the resulting recipe improves over prior diffusion-RL baselines.
ER -