Designing Reinforcement Learning for Diffusion Models: A Unified Path-Space View

2026-08-14

AI Summary

This study addresses the fragmentation and theoretical deficiencies in reinforcement learning (RL) for diffusion models by unifying loss functions via path space principles, revealing that existing methodological discrepancies fundamentally stem from variance reduction effects. We propose a multi-sample KDE gradient estimator and a bounded weight family to construct a unified algorithmic design space. Experiments on SD3.5-M and Qwen-Image models validate this theoretical interpretation, demonstrating that our approach significantly outperforms current diffusion RL baselines and enhances training efficacy. Ultimately, this work establishes a rigorous theoretical foundation and a generalized optimization framework for RL in diffusion models, effectively bridging the gap between disparate empirical methods and principled theory.

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.

Problem

• Diffusion Models
• Reinforcement Learning
• Path-Space View
• Algorithm Unification
• Variance Reduction

Innovation

• Unified Path-Space View
• Variance Reduction
• Multi-sample KDE Estimator
• Value-Gradient Estimation
• Diffusion Reinforcement Learning