Elucidating the Design Space of Regression-based Diffusion Reinforcement Learning

Toyota Li; David Zhao; Alan Zhao

cs.LGSep 27, 2026

Elucidating the Design Space of Regression-based Diffusion Reinforcement Learning

Authors: Toyota Li, David Zhao, Alan Zhao

Organizations: Tencent

Abstract

A nascent family of methods that forgoes the policy gradient and reweights a supervised regression instead has garnered momentum in reinforcement learning for diffusion and flow models. DiffusionNFT, FlowAWR, and RAM are representative regimes with contrasting motivations. It is yet opaque what, if anything, they share. We substantiate that each is the solution of one divergence-constrained reward-maximization problem, and they are differentiated only by the convex generator that defines the constraint. Under the unified modeling framework, we unravel the relaxations that prior art made during building the advantage-embedded regression target: approximating the KKT condition and posterior normalizer for the linear and exponential tilt shapes DiffusionNFT and FlowAWR respectively, while preserving the exact sparsemax projection onto the probability simplex for linear tilt leads to another superior model type in this work. Beyond the theoretical underpinnings, we further empirically investigate the design space and shed light on the training recipe for regression-style diffusion RL. Retaining the merits discovered during our exploration gives rise to DiffusionRFT, our paradigm that converges faster, trains more stably, and attains the top performance.

Figures & tables

Figure 1: Tilt $\times$ baseline grid. Three tilts (exponential, dense linear, sparsemax) crossed with two baselines (canonical constant $b$ , variance-minimizing control variate $b^{*}$ ). The sparsemax tilt reduces to the dense linear tilt when the truncation is inactive, i.e., $\min_{i}\mathcal{A}>-1$ . Yet, this criterion fails on $86.8\%$ of the dense run’s steps. The optimal control variate $b^{*}$ stabilizes and strengthens every tilt.

Figure 2: Loss geometry: the ${\bm{v}}$ -space update is misdirected, not too small. More advantage signal (b) and largest movement budget (c), but an order-of-magnitude larger residual (d), and correspondingly less reward (a). The dashed line marks the matched horizon among these runs.

Figure 3: Anchor-lock is a displacement failure, not a signal failure. (a) spans the full budget, where both the official and reimplemented RAM escape the lock with the aid of a $100\times$ advantage scale but then collapse. (b)–(d) are restricted to the matched window in which the frozen-anchor runs coexist: dashed, non-zero step updates; solid, no effective accumulation (b) and no prompt group is ever solved (c), even though the advantage is larger than under a rolling anchor (d). The three frozen-anchor curves overlap in every panel, which pinpoints the anchor as the root cause.

Figure 4: Decoupling the renoising timestep makes the batch heavy-tailed, and stratification makes it worse. (a) Training reward. (b) Within-batch spread of the policy-EMA gap $\|{\bm{f}}_{\theta}-{\bm{f}}_{\mathrm{old}}\|^{2}$ on a logarithmic axis: lowest for the trajectory-coupled schedule, $\sim 2\times$ once the timestep is decoupled, and $\sim 8\times$ once the draws are stratified. (c) Advantage magnitude: the trajectory-coupled run declines as it solves its prompt groups, while the decoupled runs remain high. (d) Fraction of prompt groups whose rewards have zero variance, which carry no advantage. A uniform distribution spends its evaluations on uninformative states the sampler hardly visits instead of the ones it resolves.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: How far the implemented target is from the exact one. By equation 45 the rescaling

r_{t}

of equation 14 has mean one at every noise level and variance bounded by

\mathbb{E}[\mathcal{A}^{2}]

, a property of the reward group alone. (a) That bound, measured on four runs and plotted as a bound on the standard deviation, so every quantity is shown as a square root. No reward model is presumed. For any centered

\mathcal{A}

(\mathbb{E}|\mathcal{A}|)^{2}\leq\mathbb{E}[\mathcal{A}^{2}]\leq\max_{i}|\mathcal{A}_{i}|\,\mathbb{E}|\mathcal{A}|

. The band is that bracket and the line its upper end. (b) The decay of

\operatorname{sd}(r_{t})

with

t

, computed in a closed form for a Gaussian-mixture endpoint law, which is exact because both exponential tilting and Gaussian smoothing preserve the family. Dots mark the first nine timesteps a

T{=}10

trajectory-coupled buffer supplies under the

\mathrm{shift}=3

schedule. When averaged over them, the dispersion is

35\%

of its worst case. It is a surrogate law and not the model, so this panel is illustrative whereas (a) and (c) are measured. (c) The same bound against the reward it was measured at. As the policy converges the group becomes homogeneous, the advantages shrink and the approximation becomes exact in the limit that matters.

Figure 6: Generalization to Pickscore reward. DiffusionRFT achieves consistent superiority over DiffusionNFT, as displayed by training reward (a) and validation reward (b).

Figure 7: The trajectory schedule covers only part of the noise axis.

t

is the noise level,

t=0

being clean data. (a) Pushing the uniform grid through the SD3.5 shift

\kappa=3

gives a density

\kappa^{2}=9\times

heavier at

t=1

than at

t=0

, and the nine loss timesteps inherit it. (b) Those nine nodes never fall below

t=0.43

, so the shaded band is off-support for the loss : drawing

t\sim\mathcal{U}[0,1]

sends

43\%

of the loss evaluations there, and stratification guarantees it. (c) The same statement as a CDF. This is the mechanism behind figure 4 : decoupling the renoising timestep does not merely re-weight the noise axis, it regresses at noise levels the policy is never asked to act on.

Explore similar work

Feb 4, 2026cs.LG

Rethinking the Design Space of Reinforcement Learning for Diffusion Models: On the Importance of Likelihood Estimation Beyond Loss Design

Reinforcement learning has been widely applied to diffusion and flow models for visual tasks such as text-to-image generation. However, these tasks remain challenging because diffusion models have intractable likelihoods, which creates a barrier for directly applying popular policy-gradient type methods. Existing approaches primarily focus on crafting new objectives built on already heavily engineered LLM objectives, using ad hoc estimators for likelihood, without a thorough investigation into how such estimation affects overall algorithmic performance. In this work, we provide a systematic analysis of the RL design space by disentangling three factors: i) policy-gradient objectives, ii) likelihood estimators, and iii) rollout sampling schemes. We show that adopting an evidence lower bound (ELBO) based model likelihood estimator, computed only from the final generated sample, is the dominant factor enabling effective, efficient, and stable RL optimization, outweighing the impact of the specific policy-gradient loss functional. We validate our findings across multiple reward benchmarks using SD 3.5 Medium, and observe consistent trends across all tasks. Our method improves the GenEval score from 0.24 to 0.95 in 90 GPU hours, which is 4.6 times more efficient than FlowGRPO and

2\times

more efficient than the SOTA method without reward hacking.

Jaemoo Choi, Yuchen Zhu, Wei Guo +6

Georgia Institute of Technology · National University of Singapore · Nanjing university.

Diffusion-Based Reinforcement Learning Methods Reward Gradients Large Language Model Reinforcement Learning Diffusion Models Policy Gradient Flow-Grpo

Ask PDF

Jun 29, 2026cs.LG

FlowAWR: Online Adaptive Flow Reinforcement via Advantage-Weighted Rectification

Aligning generative flow models on continuous spaces via online reinforcement learning is constrained by intractable trajectory likelihoods. Existing density-approximated policy gradient methods rely on stochastic SDE samplers to construct tractable transition kernels, which introduce training-inference inconsistencies and necessitates Classifier-Free Guidance (CFG). While implicit frameworks such as DiffusionNFT directly optimize forward-process velocity fields, its heuristic fixed-magnitude corrections prevent optimization strength from relative intra-group quality. We propose \textit{Flow Advantage-Weighted Rectification} (\textbf{FlowAWR}), a paradigm that recasts continuous generative policy optimization as supervised regression toward a theoretically optimal velocity field. Starting from the optimal policy of a KL-constrained reward maximization, FlowAWR derives the optimal velocity field that admits a magnitude-aware, advantage-weighted rectification form, yielding SDE-free optimization and CFG-free generation. In comparative evaluations on SD3.5-Medium, FlowAWR achieves improved alignment performance alongside a 2

\times

to 5

\times

convergence acceleration over DiffusionNFT (e.g., reaching a 24.12 PickScore in 1.2k steps, versus 23.82 in 2.0k steps for DiffusionNFT and 23.50 in

>

4k steps for FlowGRPO). Under multi-reward constraints, FlowAWR sustains generation quality, satisfying structural rules while maintaining stable out-of-domain performance.

Zheming Fu, Ruizhe He, Wei Shang +4

Beihang University · Joy Future Academy · Zhongguancun Academy +2

Generative Flow Networks Velocity Field Prompt Optimization

Ask PDF

Sep 29, 2025cs.LG

Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models

Reinforcement Learning (RL) has emerged as a central paradigm for advancing Large Language Models (LLMs), where both pre-training and RL post-training stages are grounded in the same log-likelihood formulation. In contrast, recent RL approaches for diffusion models, most notably Denoising Diffusion Policy Optimization (DDPO), optimize an objective different from the pretraining objectives--score/flow matching loss. In this work, we establish a novel theoretical analysis: DDPO is an implicit form of score/flow matching with noisy targets, which increases variance and slows convergence. Building on this analysis, we introduce Advantage Weighted Matching (AWM), a policy-gradient method for diffusion. It uses the score/flow-matching loss and reweights each sample by its advantage. In effect, AWM raises the influence of high-reward samples and suppresses low-reward ones while keeping the modeling objective identical to pretraining. This simple yet effective design yields substantial benefits: on the GenEval, OCR, and PickScore benchmarks, AWM delivers up to a

34\times

speedup over Flow-GRPO (which builds on DDPO), when applied to Stable Diffusion 3.5 Medium and FLUX, without compromising generation quality. Code is available at https://github.com/scxue/advantage_weighted_matching

Shuchen Xue, Chongjian Ge, Shilong Zhang +2

UCAS · Adobe Research · HKU +1

Diffusion Models Diffusion Policies

Ask PDF