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

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

Yixian Xu1∗, Yuanrui Zhang1∗ , Shengjie Luo2† , Liwei Wang1 , Di He1† 1 State Key Laboratory of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University 2 ByteDance Seed xyx050@stu.pku.edu.cn, yuanruizhang25@stu.pku.edu.cn, shengjieluo@bytedance.com, {wanglw,dihe}@pku.edu.cn

arXiv:2608.14430v1 [cs.LG] 14 Aug 2026

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.

1

Introduction

Diffusion and flow models have become the backbone of modern visual generation, powering highquality image, video, and 3D synthesis [11, 16, 23, 25, 34]. The pretrained generator is optimized for distributional alignment via score matching, but not for the downstream preferences that users ultimately care about, including human aesthetic preference, compositional correctness, text rendering, or task-specific verifiability. Reinforcement learning (RL) post-training is therefore a natural next step that treats the generator as a policy and directly optimizes rewards evaluated on generated samples, including black-box or non-differentiable rewards that are hard to express as supervised losses. This promise has led to a rapidly growing family of diffusion-RL algorithms. Starting from the raw objective of maximizing reward under the sampling policy with a reference regularizer, the central obstacle is that diffusion models do not provide the convenient policy likelihood ratios that PPO-style RL [32] relies on. Different practical routes around this obstacle have been proposed: some discretize the reverse-sampling SDE so that each denoising step has a Gaussian transition ratio [5, 24, 37]; others avoid reverse-time likelihoods and instead build supervised flow-matching ∗ Equal contribution. † Correspondence to: Shengjie Luo <shengjieluo@bytedance.com>, Di He<dihe@pku.edu.cn>,.

Preprint.

losses from reward-labeled clean samples and their forward-noised versions [40, 45]. These routes produce loss functions that look structurally unrelated, and the community currently lacks a unified view for understanding existing recipes or designing principled diffusion-RL algorithms beyond combinations of heuristics. Our first step is to expose the common continuous-time structure beneath these routes. Although a single diffusion path measure has no elementary likelihood, ratios between path measures generated by stochastic differential equations (SDEs) with the same diffusion coefficient have explicit form [2, 29]. Applying this path-space importance sampling to the raw diffusion-RL objective yields a tractable estimator under a proposal path measure. Its policy gradient term contains a stochastic Itô integral that corresponds to the noise exploration signal used by reverse-sampling-based GRPO methods [24]. We prove that this term has an equivalent deterministic value-gradient representation. This variancereduced representation is the bridge from Flow-GRPO-style stochastic estimators to forward-matching methods: the large empirical gap between Flow-GRPO and DiffusionNFT/AWM is explained as a variance reduction effect rather than as a fundamentally different RL principle. The variance-reduced form also reveals the design space. Its only unspecified component is how to estimate the value gradient ∇xt Vt on sampled trajectories, together with how to weight the resulting policy gradient and quadratic regularization terms across time. A Bayes rule expression for ∇xt Vt yields a deterministic one-sample estimator that recovers AWM and DiffusionNFT under their forward noising proposal, while the Brownian increment estimator recovers Flow-GRPO-type methods under reverse SDE rollouts. Lifting these choices to free parameters gives a unified template whose concrete entries specify the proposal dynamics, the value gradient estimator, the off-policy weight ŵ1 , the on-policy weight ŵ2 , and the sampler. These entries are organized into three axes: value gradient estimation, weight functions, and sampling. Under this template, previously separate algorithms become a specification of each entry in the same design space rather than separate design principles. Finally, the framework turns this taxonomy into concrete design guidance. Along the estimator axis, we observe that existing methods all use one-sample estimators of ∇xt Vt ; we introduce a multisample KDE estimator that reuses the rollout group already collected for GRPO-style advantages and reduces conditional variance without additional sampling. Along the weight axis, we identify d t and the quadratic regularizer a scale-bounded principle: the effective per-step update ŵ2 (t)∇V 2 ŵ1 (t)∥∆vθ ∥ should remain bounded across time and discretization scale. This principle recovers the stable weightings used by AWM, DiffusionNFT, and GRPO-Guard, rules out singular Flow-GRPOstyle weights in off-policy regimes, and reduces algorithm design to a small family of interpretable weight shapes. Contributions.

In summary, this paper contributes:

• A continuous-time path-space framework (Sec. 3.1, Sec. 3.2). Starting from the raw objective Eq. (3), we use path-space importance sampling and a variance-reduction theorem (Thm. 2) to derive an explicit template Eq. (9) that unifies Flow-GRPO [24], DanceGRPO [41], TempFlow-GRPO [14],GRPO-Guard [37], AWM [40], and DiffusionNFT [45] as special cases (Table 1). • Principled design choices (Sec. 3.3). We introduce a multi-sample KDE value-gradient estimator Eq. (10) and a scale-bounded weight principle that yields compact families for ŵ1 and ŵ2 , recovering strong existing recipes while exposing the remaining tunable degrees of freedom. • Empirical validation (Sec. 4). On SD3.5-M [11] and Qwen-Image [38] models across PickScore [19], OCR [7], and GenEval [13] rewards, we verify that the variance-reduction trick closes the Flow-GRPO–DiffusionNFT efficiency gap, that the proposed weight families track empirical optima, and that our final recipe exceeds the baselines in Table 1.

2

Backgrounds

2.1

Diffusion and Flow-Matching Generative Models

Diffusion and Flow-Matching generative models perform step-by-step transformations between the data distribution and a tractable Gaussian prior. The forward process noises data to a Gaussian prior, 2

and the reverse process generates data by denoising. Following the Flow-Matching formulation [1, 23, 25], the forward noising process is constructed as xt = (1 − t)x0 + tϵ, t ∈ [0, 1], (1) with ϵ ∼ N (0, I) and x0 ∼ pdata . The Flow-Matching formulation adopts the vector-field prediction loss to train the diffusion model: 2 L(θ) = Et∼p(t),x0 ∼pdata ,ϵ∼N (0,I) ∥vθ (xt , t) − (ϵ − x0 )∥ , where vθ (xt , t) is parameterized by a neural network and p(t) is the distribution of the time steps. Once the model is trained, the sampling process can be performed by simulating the following SDE: r  
ηt 2tηt xt + (1 − t)vθ (xt , t) dt + dxt = vθ (xt , t) + dwt , x1 ∼ N (0, I), (2) 1−t 1−t where ηt ≥ 0 is a non-negative function and wt is the standard Brownian motion on Rd . When ηt = 1, Eq. (2) is the time-reversal of Eq. (1); when ηt = 0, it collapses to the deterministic ODE dxt dt = vθ (xt , t), i.e. the FM analogue of the probability flow ODE [34]. 2.2

Reinforcement Learning for Diffusion Models

Let C be a prompt space and let R : Rd × C → [0, 1] be a reward function evaluated on clean samples. We denote x0:1 as the trajectory of the stochastic process {xt }t∈[0,1] in Eq. (2) and p(x0:1 , c) as its likelihood3 . RL post-training seeks to maximize the following canonical objective: h i L(θ) = Ec Epθ (x[0:1] ,c) [R(x0 , c)] − βKL(pθ ∥pref ) ,

(3)

where pref is a frozen reference policy. Recently, several practical alternatives for Eq. (3) have been proposed. Flow-GRPO [24] applies an Euler–Maruyama discretization to the reverse-time sampling SDE, treating each denoising step as a tractable Gaussian transition and running clipped PPO over the discretized trajectory. Denote the current diffusion policy as pθ and the old policy as pold . Given a group of old-policy trajectories {{xitj }Tj=1 }G i=1 , Flow-GRPO optimizes: GRPO

L

G T

1 X1X (θ) = Ec,{xi0 }G min rji (θ)Ai , r̃ji (θ)Ai − βKL(pθ |pref ) . i=1 ∼pold (·|c) G T i=1 j=1

R(xi0 ,c)−mean({xi0 }G i=1 ) is the group-normalized reward of the i-th sample in the group std({xi0 }G i=1 ) i G i i {x0 }i=1 , r̃j = clip(rj (θ), 1 − ϵ, 1 + ϵ) is the clipped likelihood ratio of the j-th step for the i-th pθ (xit |xit ,c) sample where rji = pold (xij−1 |xij ,c) is Euler–Maruyama per-step likelihood ratio. t t

Here Ai =

j−1

j

A different line of methods forgoes reverse-trajectory likelihood ratios entirely and constructs training signals from reward-labeled clean samples via the forward noising path. AWM [40] samples x0 from the old policy, evaluates its advantage A(x0 , c), applies the forward noising path xt = (1 − t)x0 + tϵ, and reuses the flow-matching target with advantage weighting4 : h i 2 −LAWM (θ) = Ec,x0 ,t,ϵ A(x0 , c) ∥vθ (xt , t, c) − (ϵ − x0 )∥ . DiffusionNFT [45] follows the same forward-noising interface but uses the reward to define positive and negative velocity mixtures around the old model, v+ = (1 − β)vold + βvθ and v− = (1 + β)vold − βvθ : h i 2 2 −LNFT (θ) = Ec,x0 ,t,ϵ R(x0 , c) v+ (xt , t, c) − (ϵ − x0 ) + (1 − R(x0 , c)) v− (xt , t, c) − (ϵ − x0 ) .

At first glance, these methods appear to be driven by different principles and yield structurally unrelated loss functions. Without a unified view that connects them, the field still lacks a principled basis for understanding existing recipes or systematically designing new diffusion-RL algorithms. 3 This is a formal notation that treats the probability space defined on continuous paths as its discretization. Actually, this is known as the path measure corresponding to a stochastic differential equation (SDE). See Appx. B for background. 4 Note that for AWM and DiffusionNFT, we also use L to represent the maximum objective.

3

3

A Unified Path-Space View of Diffusion RL

3.1

A Path-Space View via Importance Sampling

The structural obstacle of Diffusion RL. Despite the variety of practical recipes, RL for diffusion models ultimately rests on the canonical objective Eq. (3). Its policy-gradient estimator ∇θ Epθ [R] formally requires the score ∇θ log pθ (x0 | c), but pθ (x0 | c), the marginal of the full generation trajectory, has no closed form. The trajectory likelihood pθ (x[0:1] , c) fares no better: as the path measure of a continuous-time SDE, the path law is naturally defined as a measure on C([0, 1], Rd ), not as a tractable density with respect to a finite-dimensional reference measure [29]. The logderivative trick that underlies PPO/GRPO [32, 33, 44] therefore does not apply directly. This is the structural obstacle that any RL algorithm on a diffusion generator must confront. Viewed through this perspective, the three methods of Sec. 2.2 are different workarounds to the same obstacle. FlowGRPO discretises the reverse SDE into tractable Gaussian steps, while AWM and DiffusionNFT bypass reverse-time likelihoods entirely and match on the forward noising path. Path-Space Importance Sampling. We try to address the obstacle directly. The key observation is that although a single diffusion path measure has no density, pairs of path measures generated by SDEs that share the same diffusion coefficient admit a closed-form importance sampling ratio [2, 29]. Concretely, if we sample trajectories from a tractable proposal q drawn from the same SDE family (Eq. (2)) with drift vbase in place of vθ and the same noise schedule ηt , the ratio pθ /q is explicit and differentiable in θ under certain conditions. This is exactly importance sampling in the path space: we sample entire trajectories {xt }t∈[0,1] rather than only the endpoint x0 . Applying this technique to Eq. (3) and replacing the reward R(x0 , c) with the group-relative advantage A(x0 , c) (which leaves the policy gradient unchanged while reducing its variance), we obtain the unbiased estimator     pθ (x[0:1] , c) pθ (x[0:1] , c) pθ (x[0:1] , c) L(θ, c) = Eq(x[0:1] ,c) A(x0 , c) − βEq(x[0:1] ,c) log , (4) q(x[0:1] , c) q(x[0:1] , c) pref (x[0:1] , c) where pθ , pref , and q correspond to the sampling SDE (Eq. (2)) with velocity fields vθ , vref , and vbase respectively, and the same noise schedule ηt . In this way, the importance sampling ratio pθ (x[0:1] ,c) M1 (θ) := log q(x is given by (see Appx. B): [0:1] ,c) M1 (θ) = −

Z 1p

Z 1 2w̃1−r ∆vθ (x1−r , 1 − r) · dwr −

2

w̃1−r ∥∆vθ (x1−r , 1 − r)∥ dr, 0

0

2

t) · 1−t where r = 1 − t, ∆vθ := vθ − vbase and the weight w̃t = (1+η 4ηt t couples the training weight to the sampler noise schedule. Substituting M1 (θ) into Eq. (4), the objective becomes a single expectation under q: h i pref (x[0:1] ,c)
L(θ, c) = Eq [exp(M1 (θ)) A(x0 , c)] − β Eq exp(M1 (θ)) M1 (θ) − log q(x . (5) [0:1] ,c)

An explicit variance-reduced template. We focus here on the policy-gradient term of Eq. (5), Lpolicy (θ, c) := Eq [exp(M1 (θ))A(x0 , c)]; the KL term is handled identically. Under mild assumptions that follow standard PPO practice [32, 33], the explicit form of Lpolicy (θ, c) can be exactly obtained (See Appx. C.2 for proof): Proposition 1 (Stochastic-integral path-space estimator). Assume clipping or a small step size keeps exp(M1 (θ− )) ≈ 1, where θ− denotes the stop-gradient version of the current parameters. Then   Z 1p Z 1     2 −Lpolicy (θ, c) =θ Eq  A(x0 , c) 2w̃1−r ∆vθ · dwr + A(x0 , c) w̃1−r ∥∆vθ ∥ dr , (6)   0 0 | {z } | {z } Term (A): noise exploration

where =θ denotes equality up to a θ-independent constant. 4

Term (B): regularisation

Intuitively, Term (B) is an advantage-weighted quadratic that anchors vθ to vbase , while Term (A) follows the Brownian increments weighted by advantage — pushing vθ toward the noise realisations that yielded higher rewards. The latter is an Itô integral, however, whose variance dominates the overall variance of Eq. (6). A natural question arises: can Term (A) be rewritten as a deterministic integral with the same expectation and with smaller variance? The next theorem provides an explicit variance-reduced template (See Appx. C.3 for the proof): Theorem 2 (Variance-reduced path-space estimator). Let q be the path measure of the sampling SDE Eq. (2) with base drift vbase and noise schedule ηt > 0, and let Vt (x, c) := Eq [A(x0 , c) | xt = x] be the associated value function. Under the regularity the trust-region condition of Prop. 1, the policy gradient of Eq. (5) admits the (locally) variance-reduced representation policy

−L

(θ, c) = θ Eq

Z 1 

2

(1 + ηt )∆vθ (xt , t) · ∇xt Vt (xt , c) + A(x0 , c)w̃t ∥∆vθ (xt , t)∥



 dt . (7)

0

From the intractable canonical objective Eq. (3), we arrive at the explicit template Eq. (7), which is both (i) tractable — every term is a deterministic integral under a proposal q in the same SDE family as pθ , with no reliance on path-likelihood densities of pθ that do not exist — and (ii) variance-reduced — the stochastic integral in Eq. (6) is replaced by a deterministic integral depending on a ∇xt Vt estimator. These two properties enable Eq. (7) to serve as a unified tool to understand different RL methods for the diffusion model. 3.2

A Unified Framework and Its Design Space

The missing ingredient. Eq. (7) leaves one quantity unspecified: the value gradient ∇xt Vt (xt , c). How this quantity is evaluated on each rollout together with the weights on the two terms of Eq. (7) turns out to distinguish each diffusion RL method of Sec. 2.2 from others. We begin by deriving an analytical Bayes-rule expression for ∇xt Vt , from which a one-sample estimator follows. Proposition 3. The gradient of the value function admits the closed-form expression ∇xt Vt (xt , c) = Eq(x0 |xt ) [∇xt log q(xt | x0 ) A(x0 , c)] − Vt (xt , c) Eq(x0 |xt ) [∇xt log q(xt | x0 )] . Let v(xt , x0 ) := (xt − x0 )/t denote the FM conditional velocity. Under q with vbase = E[v | xt ] ≈ vold and ηt ≡ 1 (i.e., q is the reverse of the forward noising process), this specializes in 
 ∇xt Vt (xt , c) = − 1−t t Eq(x0 |xt ) A(x0 , c) v(xt , x0 ) − vbase (xt , t) . The corresponding one-sample estimator on a rollout trajectory x0:1 is
det d t (xt , c) := − 1−t A(x0 , c) v(xt , x0 ) − vold (xt , t) . ∇V t det

dt See Appx. C.4 for the proof. Substituting ∇V

into Thm. 2 yields the concrete training objective

    Z 1
   2 t) ∆vθ (xt , t) · A(x0 , c) v − vold (xt , t) + A(x0 , c) w̃t ∥∆vθ (xt , t)∥  dt , (8) Eq  − (1−t)(1+η t | {z } | {z } 0 weight term 2

weight term 1

where v = (xt − x0 )/t is the single-sample conditional velocity. Eq. (8) is a fully explicit, variancecontrolled diffusion-RL loss obtained by instantiating the variance-reduced template Eq. (7) with the deterministic one-sample estimator of Prop. 3 under the noise schedule ηt = 1. A unified framework. Moreover, the stochastic integrand of Term (A) in Eq. (6) is itself an estimator of ∇xt Vt — a stochastic one-sample estimator using the Brownian increment of the rollout p sto d t (xt , c) := A(x0 , c) (1 − t)/(2ηt t∆t) ϵt . Substituting it through the structural identity SDE, ∇V of Thm. 2 gives a second, equally valid instantiation. The two instantiations share a common two-term form: an off-policy quadratic in ∆vθ weighted by a function of the advantage, plus an on-policy inner product between ∆vθ and the chosen value-gradient estimator. Viewing the estimator and the two weights as free parameters yields the following unified framework:

−Lpolicy (θ, c) =

Z 1

h i
2 d t (xt , c) dt, (9) Eq ŵ1 A(x0 , c), t ∥∆vθ (xt , t)∥ + ŵ2 (t) ∆vθ (xt , t), ∇V | {z } | {z } 0 off-policy regulariser

on-policy policy-gradient

5

Table 1: Diffusion RL methods as rows in the five-knob design space of Eq. (9). We use the FM parameterisation throughout; v and ϵ denote the FM conditional velocity and the per-step Brownian increment. See Appx. D for detailed derivations of each row. Method Flow-GRPO [24, 41]

dt ∇V

vbase (ηt ) vold (ηt )

s(t) q

1−t 2ηt t∆t

s(t) · Aϵ

GRPO-Guard [37]

vold (ηt )

s(t) · Aϵ

q

TempFlow-GRPO [14]

vold (ηt )

s(t) · Aϵ

q

E[v | xt ] (1) −s(t) · A(v−vold )

DiffusionNFT [45]

E[v | xt ] (1) −s(t) · A(v−vold ) vold (ηt )

ŵ2 (t)

Sampler

2 t ) 1−t A · (1+η 4ηt t

1 + ηt

Flow-SDE

1−t 2ηt t∆t

AWM [40]

Ours (Sec. 3.3)

ŵ1 (A, t)

d tKDE (Eq. (10)) ∇V

1−t 2ηt t∆t 1−t t 1−t t 1−t t

0 q

2

t) A · (1+η 4

2(1−t)∆t ηt t

A/∆t 1/∆t (1 − t)α1 /∆t

(1+ηt ) s(t)∆t

Flow-SDE

(1 + ηt )/s(t)

Branched-SDE

2 s(t)∆t 2 βs(t)∆t

Flow-SDE / ODE

tα2 s(t)∆t

DPM-ODE [26] Flow-SDE

d t is any estimator of ∇x Vt (xt , c) and ŵ1 , ŵ2 are arbitrary weight functions. In principle, where ∇V t different diffusion RL methods of Sec. 2.2 turn out to be instances of Eq. (9) under specific choices of these parameters. Table 1 enumerates these choices for various methods. The design space. Reading Table 1 by columns, the concrete choices in Eq. (9) group naturally into three independent axes: (i) Value-gradient estimation: the proposal q (determined by the choice of vbase and ηt ) d t evaluated under q. A typical choice takes q induced by together with the estimator ∇V a lagged velocity model vold and a specific noise schedule ηt . (ii) Weight functions: the off-policy weight ŵ1 (A, t) and the on-policy weight ŵ2 (t). When vold = vθ (on-policy), the off-policy regulariser of Eq. (9) vanishes and only ŵ2 matters; when vold ̸= vθ (off-policy), ŵ1 and ŵ2 jointly determine the training signal. (iii) Sampler: the numerical scheme used to simulate q and collect rollout samples. 3.3

Elucidating Design Choices

d t and the weight The design space of Sec. 3.2 has two primary axes: the value-gradient estimator ∇V 5 pair (ŵ1 , ŵ2 ). We elucidate each axis in turn . For the estimator, we observe that methods in Table 1 rely on a one-sample Monte Carlo estimator of ∇xt Vt and we introduce a multi-sample refinement based on kernel density estimation. For the weights, we extract a scale-bounded structural principle from recent advanced methods and show that it implies the t-dependent shape of both ŵ1 and ŵ2 given the choice of value-gradient estimator. Estimator axis: from one-sample to multi-sample.

Every row of Table 1 uses a one-sample value sto

d gradient estimator: GRPO-type methods use the original stochastic one-sample estimator ∇V built t from a single Brownian increment, while AWM and DiffusionNFT use the deterministic one-sample det d t (defined in Prop. 3) from a single (x0 , xt ) pair. Both carry the full conditional estimator ∇V variance of the integrand per rollout which remains unchanged by simply enlarging the group. The natural variance reduction is to average over K samples from q(x0 |xt ), but drawing K fresh samples at each query xt would cost K extra rollouts. Instead, we exploit the collected rollout group {(xi0 , xit )}G i=1 . For a noisy state xt , we take a kernel-weighted average over the G group samples with a kernel matched to the FM forward likelihood q(xt |x0 ) = N (xt ; (1 − t)x0 , t2 I). This gives the Nadaraya–Watson KDE estimator
PG i i i KDE i=1 K(xt , x0 ) A(x0 , c) v(xt , x0 ) − vold (xt , t) 1−t d ∇V t (xt , c) := − t · , (10) PG i i=1 K(xt , x0 ) 2

0∥ where K(xt , x0 ) := exp(− ∥xt −(1−t)x ) is the FM-matched Gaussian kernel and h is a hyper2h t2 parameter to balance the relative weight for the samples in the group, typically fixed as a constant

5We treat sampler choices as specific designs within value-gradient estimation (e.g., trajectory branching in MixGRPO and

TempFlow-GRPO).

6

depends on the dimension. The KDE estimator is also inspired by Bertrand et al. [4], which shows that including one anchor posterior sample can make the self-normalized estimator unbiased. We summarize the properties of different value-gradient estimators as follows. Proposition 4 (Properties of the value-gradient estimators). Let A(x0 , c) be square-integrable. Assume vbase = vold , then the following arguments hold: (i). The discretized stochastic one-sampleh estimator is asymptotically unbiased in the continuousi sto d time Itô limit of Thm. 2: lim∆t→0 E ∇V t,∆t | xt = ∇xt Vt (xt , c); (ii). Assume ηt ≡ 1,hvold (xt , t) =i E[v(xt , x0 ) | xt ], then the deterministic one-sample estimator det d t | xt = ∇x Vt (xt , c); is unbiased: E ∇V t

(iii). Under the assumption of (ii) and h = 1, the KDE estimator is unbiased and has smaller variance than the deterministic one-sample estimator: h i     KDE j KDE det dt dt d t | xjt . E ∇V (xt , c) | xjt = ∇xt Vt (xjt , c), Var ∇V | xjt ≤ Var ∇V See Appx. C.5 for the proof. Prop. 4 shows that the KDE estimator is unbiased and performs variance reduction for the deterministic one-sample estimator when h = 1. With the same rollout construction but h ̸= 1, the anchor is drawn from the true posterior while the normalized weights correspond to a different kernel, so the estimator is generally biased. Thus, h explicitly trades exact posterior matching for smoother finite-sample estimates; we quantify this trade-off in Sec. 4.1. Weight axis: a scale-bounded principle. From Eq. (9), the on-policy gradient contribution at d t (xt , c): this product is what scales the per-step time t has magnitude proportional to ŵ2 (t) ∇V update before the network-gradient factor ∂θ ∆vθ enters. For stable training, its magnitude should be controlled across t: neither blowing up at particular timesteps (which would amplify gradient variance and skew the update) nor vanishing where gradient signal is needed. We therefore inspect d t across the rows of Table 1: ŵ2 (t) ∇V Flow-GRPO: GRPO-Guard: AWM: DiffusionNFT:

p d t ∆t = (1 + ηt ) A (1 − t)∆t/(2ηt t) ϵt ŵ2 (t) ∇V d t ∆t = (1 + ηt ) A ϵt ŵ2 (t) ∇V d ŵ2 (t) ∇V t ∆t = −2 A (v − vold ) d t ∆t = − 2 A (v − vold ) ŵ2 (t) ∇V β

(blows up as , t → 0) (bounded) (bounded) (bounded).

The three methods that empirically outperform Flow-GRPO each pair their estimator with a weight d t ∆t stays bounded uniformly in t; Flow-GRPO alone violates this property. whose product ŵ2 ∇V We conclude the pattern to an explicit design principle: d t ∆t stays bounded uniformly in t ; Scale-bounded principle. Choose ŵ2 (t) so that ŵ2 (t) ∇V 2 analogously, choose ŵ1 (A, t) so that ŵ1 (A, t)∥∆vθ ∥ ∆t stays bounded uniformly in t. The principle applies asymmetrically to the two weights, reflecting the asymmetric t-dependence of det d t and our KDE estimator the two factors it multiplies. For ŵ2 (t), the deterministic estimator ∇V KDE

d ∇V both inherit the (1 − t)/t prefactor of Prop. 3; the principle then pins down the shape of t d t (summarized in ŵ2 (t) up to a single free exponent. Let s(t) be the scale factor of the estimator ∇V Table 1), w2 (t) := ŵ2 (t)s(t)∆t. The principle requires w2 (t) to stay bounded uniformly in t. We adopt the minimal single-exponent family w2 (t) = tα , α ≥ 0 which saturates the principle and concentrates mass more sharply near t = 1 for larger α (any multiplicative constant is omitted). 2

For ŵ1 , the situation is different: because ∥∆vθ ∥ carries no t-singularity, the principle reduces to “ŵ1 ∆t is bounded in t”. Flow-GRPO’s ŵ1 ∆t = A w̃t ∆t ∝ A(1 − t)∆t/t diverges as t → 0 and is ruled out; AWM’s ŵ1 ∆t = A and DiffusionNFT’s ŵ1 ∆t = 1 are admissible, but the principle alone does not pick the monotonicity direction in t or the presence of advantage coupling. Let w1 (A, t) = ŵ1 (A, t)∆t. Four minimal single-exponent candidates remain:  w1 (A, t) ∈ tα , (1 − t)α , A(x0 , c) tα , A(x0 , c) (1 − t)α , α ≥ 0, 7

distinguishing (a) whether the trust-region tightens toward the clean-data end ((1 − t)α ) or the noise end (tα ), and (b) whether its strength couples to the advantage. Which of the four shapes is optimal is an empirical question the principle cannot resolve; we ablate across the family in Sec. 4.2.

4

Experiments

In this section, we show the experimental results on the design of each component in our framework. In Sec. 4.1, we introduce the settings of our experiments and give an example to show how the variance reduction trick in Sec. 3.1 works. In Sec. 4.2, we show the principles of weight design for different estimators, which gives a mapping between the value-gradient estimator and its optimal weight. In Sec. 4.3, we compare our final recipe derived by following the design-space map to its optimum, and show the comparison with previous baseline models across different reward settings. 4.1

Settings

Setup. We conduct our experiments on the SD-3.5 Medium model [11] and Qwen-Image [38] at a resolution of 512 × 512. Following Flow-GRPO [24] and DiffusionNFT [45], we set group size G = 24, T = 10 denoising steps for training and T = 40 for evaluation. In addition, we choose to fine-tune the pretrained model using LoRA [17] and disable classifier-free guidance (CFG) [15] in the training and evaluation stages for fair comparison. We evaluate our methods on various rewards including PickScore [19], visual-text OCR accuracy [7], and GenEval [13]. See Appx. F for details. How the variance-reduction trick works. In Sec. 3.1 and Sec. 3.2, we show the correctness of the one-sample estimator and view it as the estimator for the variance-reduced form of Term A in Eq. (6). Although the methods exhibit strong theoretical properties, it’s unclear whether it helps to improve the training convergence of RL algorithms. Previous works AWM [40] and DiffusionNFT [45] show that this type of estimator accelerates the training convergence of RL training under specific weight design (Table 1), it’s unclear either the estimator or the specific weight design leads to the empirical performance. To find out this, we directly compare the training performance of three training objectives: (1) the original objective Eq. (6) with ηt = 1 (corresponding to the reverse of the forward process) ; (2) the one-sample estimator with inherited weighting Eq. (8); (3) A reweighted version of Eq. (8) (same as AWM [40]). The results are shown in Fig. 1(a). The one-sample estimator with inherited weighting suffers instability and ultimately fails, indicating that the variance reduction alone can not accelerate the training; the reweighting tricks are also indispensable. With the reweighted form, the one-sample estimator achieves remarkably faster convergence speed. This observation inspires us to investigate the weight design for each value-gradient estimator in the next subsection. Direct bias–variance measurement. We directly measure the estimation error in the setting that matches the conditions of Prop. 4. We set ηt = 1 and evaluate SD3.5-Medium with PickScore on four prompts, using 128 trajectories per prompt and 40 denoising steps. As shown in Fig. 1(a), the stochastic estimator has 13.8× the variance of the deterministic one-sample estimator, directly supporting the variance-reduction mechanism in Thm. 2. The matched KDE estimator (h = 1) further reduces variance by 67.5% relative to the deterministic estimator. Squared bias contributes less than 2% of the MSE in every setting, so the observed error is dominated by variance. Bias–smoothness trade-off. The bandwidth controls a practical bias–smoothness trade-off. Although h = 1 yields the lowest variance (Prop. 4 and Fig. 1(b)), its kernel weights can change abruptly across adjacent denoising steps in the high-dimensional latent space. Increasing h smooths this temporal transition and improves optimization stability, at the cost of weakening posterior matching and increasing both variance and bias. We therefore use the intermediate h that retains a meaningful variance reduction while producing smoother value-gradient targets along the trajectory. 4.2

A Mapping between the Value-Gradient Estimator and Its Optimal Weight

On-policy weighting w2 . We study three types of estimators: stochastic one-sample estimator sto det KDE d t , deterministic one-sample estimator ∇V d t and our proposed KDE estimator ∇V dt ∇V . As introduced in Sec. 3.3, we investigate the weight: w2 (t) = tα . We find that different sampler noise 8

Variance 100

10−1

Squared bias 10−2

0.780

0.057

0.041

0.060

10−3 2.35e-04

0.018

3.45e-04

1.46e-04 1.39e-04

10−4

10−2 Stoch. Det. 1-sample 1-sample

(a)

6.64e-03

KDE h=1

(b)

KDE KDE h = 32.768 h = 65.536 (training config)

Stoch. Det. 1-sample 1-sample

KDE h=1

KDE KDE h = 32.768 h = 65.536 (training config)

(c)

Figure 1: (a) Comparison between different objectives on Pickscore reward. Original objective Eq. (6) with ηt = 1 (Orange line); one-sample estimator with weighting in Eq. (8) (Blue line); reweighted version of Eq. (8) (Green line). (b,c) Empirical variance and squared bias of the value-gradient estimators. We set ηt = 1 and evaluate SD3.5-Medium with PickScore reward on four prompts, using 128 trajectories per prompt and 40 denoising steps. At every intermediate state, the reference value gradient is approximated from all 128 trajectories; each estimator is computed with G = 24 then compared with this common reference. The KDE training configuration uses h = 32.768.

(a)

(b)

(c)

(d)

(e)

(f)

Figure 2: On-policy weight (row 1) and off-policy (row 2) weight ablations for the stochastic, deterministic, and our KDE estimator. The results for the same estimator are in the same column. schedules affect the behavior (Appx. F.1), and we choose the best sampler noise schedule for weight ablation. The results are shown in Fig. 2 (first row). By adjusting the weights parameter α (see Fig. 6 for visualization), we find that a different estimator exhibits a different behavior for weight preference. The stochastic estimators are not so sensitive w.r.t. the weight adjustment, w2 = t performs slightly better (Fig. 2 (a)). Weight with larger α helps stabilize the training of the deterministic estimator, but can not fully resolve it without learning rate decay (Fig. 2 (b)). Our KDE estimator exhibits better training stability under weight with larger α without any extra design (Fig. 2 (c)). Larger α improves stability also implies that more training effort should be allocated to high-noise regions, which affects the sampling process at early stages. Off-policy weighting w1 . Based on the optimal on-policy weight, we adjust the off-policy weight w1 . As introduced in Sec. 3.3, we adopt the following weight design: w1 = {0, tα , Atα , (1 − t)α , A(1 − t)α }. For stochastic estimator, we find that w1 = 0 works well (design of GRPO-Guard in Table 1); the weight in {tα , (1 − t)α } overemphasize the regularization term and leads to slow convergence, while weights {Atα , A(1 − t)α } leads to training instability (Fig. 2 (d)). Similarly, for the deterministic estimator and our KDE estimator, tα degrades the training convergence speed and 9

{Atα , A(1 − t)α } causes training instability (Fig. 2 (e,f)). We find that w1 (t) = (1 − t)α balances the weight between the off-policy regularization and the policy gradient in Eq. (9), which implies that regularization in low-noise regions helps to improve the training stability. 4.3

Comparison with Prior Works

(a)

(b)

(c)

(d)

Figure 3: (a,b,c). Training time comparison between our method and several baselines. (d). Comparison between Flow-GRPO, DiffusionNFT, and our method on Qwen-Image, HPSv3 reward. * means using the official implementation with different training configs compared to other baselines. Based on the estimator–optimal weight mapping constructed in Sec. 4.2, we perform a training-time comparison between our best setting and previous strong baselines across several rewards. We use our KDE estimator with its optimal weighting w1 = tα , w2 = (1 − t)α . The results are summarized in Fig. 3 (a,b,c). On OCR reward, our method significantly accelerates training convergence compared to previous SOTA methods: 2× faster than AWM, 3× faster than DiffusionNFT. On Pickscore and GenEval rewards, our method achieves comparable performance with AWM and DiffusionNFT without using other designs (e.g., specific EMA-KL loss in training, elaborate sampling design, training time-step selection), and outperforms all other baselines. We also conduct experiments on the large-scale model Qwen-Image [38] with HPSv3 reward [27], and the result in Fig. 3(d) shows that our method achieves 4× faster convergence speed than DiffusionNFT under the same configurations. Ablation study. To isolate the source of the gain, we start the ablation study from DiffusionNFT without its EMA-KL stabilization. The results are summarized in Fig. 4. The without EMA-KL baseline is unstable and eventually fails after 60 GPU hours. Replacing only its deterministic one-sample estimator with our KDE estimator improves its stability but hurts the performance; variance reduction alone is therefore insufficient. Using our on-policy weight helps stabilize and improve training performance. Fi- Figure 4: Ablation studies on the components of nally, combining KDE with both our on-policy our final recipe. and off-policy weights reaches 23.52 after 60 hours. For comparison, the original DiffusionNFT (on-policy, with EMA-KL) reaches 23.382 after 160 hours, so our recipe is roughly 3× faster.

5

Conclusion

In this work, we develop a unified path-space framework that places previous diffusion-RL algorithms on a common theoretical foundation. Starting from the original RL objective, we derive a general training objective based on the variance-reduced form of the original objective. Previous reversetrajectory methods and forward-matching methods can all be viewed as instances of this general objective. The derivation also yields a unified design space, including the value-gradient estimation, weight functions, and sampling choices. We propose a multi-sample KDE value gradient estimator and find a mapping between different estimators and their optimal weight design. Experiments on SD3.5-M and Qwen-Image models validate our theoretical explanation and show the empirical improvement of our result recipe. 10

References [1] Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797, 2023. [2] Brian DO Anderson. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313–326, 1982. [3] Andreas Bergmeister, Stefanie Jegelka, Nikolas Nüsken, Carles Domingo-Enrich, and Jakiw Pidstrigach. Reinforce adjoint matching: Scaling rl post-training of diffusion and flow-matching models. arXiv preprint arXiv:2605.10759, 2026. [4] Quentin Bertrand, Anne Gagneux, Mathurin Massias, and Rémi Emonet. On the closed-form of flow matching: Generalization does not arise from target stochasticity. Advances in neural information processing systems, 38:8522–8549, 2026. [5] Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305.13301, 2023. [6] George Casella and Roger Berger. Statistical inference. CRC press, 2024. [7] Jingye Chen, Yupan Huang, Tengchao Lv, Lei Cui, Qifeng Chen, and Furu Wei. Textdiffuser: Diffusion models as text painters. Advances in Neural Information Processing Systems, 36: 9353–9387, 2023. [8] Jaemoo Choi, Yuchen Zhu, Wei Guo, Petr Molodyk, Bo Yuan, Jinbin Bai, Yi Xin, Molei Tao, and Yongxin Chen. Rethinking the design space of reinforcement learning for diffusion models: On the importance of likelihood estimation beyond loss design. arXiv preprint arXiv:2602.04663, 2026. [9] Kevin Clark, Paul Vicol, Kevin Swersky, and David J Fleet. Directly fine-tuning diffusion models on differentiable rewards. arXiv preprint arXiv:2309.17400, 2023. [10] Carles Domingo i Enrich, Michal Drozdzal, Brian Karrer, and Ricky TQ Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. In International Conference on Learning Representations, volume 2025, pages 53791–53846, 2025. [11] Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first international conference on machine learning, 2024. [12] Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. Dpok: Reinforcement learning for fine-tuning text-to-image diffusion models. Advances in Neural Information Processing Systems, 36:79858–79885, 2023. [13] Dhruba Ghosh, Hannaneh Hajishirzi, and Ludwig Schmidt. Geneval: An object-focused framework for evaluating text-to-image alignment. Advances in Neural Information Processing Systems, 36:52132–52152, 2023. [14] Xiaoxuan He, Siming Fu, Yuke Zhao, Wanli Li, Jian Yang, Dacheng Yin, Fengyun Rao, and Bo Zhang. Tempflow-grpo: When timing matters for grpo in flow models. arXiv preprint arXiv:2508.04324, 2025. [15] Jonathan Ho and Tim Salimans. arXiv:2207.12598, 2022.

Classifier-free diffusion guidance.

arXiv preprint

[16] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840–6851, 2020. [17] Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Liang Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. Iclr, 1(2):3, 2022. 11

[18] Ioannis Karatzas and Steven Shreve. Brownian motion and stochastic calculus. springer, 2014. [19] Yuval Kirstain, Adam Polyak, Uriel Singer, Shahbuland Matiana, Joe Penna, and Omer Levy. Pick-a-pic: An open dataset of user preferences for text-to-image generation. Advances in neural information processing systems, 36:36652–36663, 2023. [20] Jeongjae Lee, Jinho Chang, Jeongsol Kim, and Jong Chul Ye. Reward score matching: Unifying reward-based fine-tuning for flow and diffusion models. arXiv preprint arXiv:2604.17415, 2026. [21] Kimin Lee, Hao Liu, Moonkyung Ryu, Olivia Watkins, Yuqing Du, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, and Shixiang Shane Gu. Aligning text-to-image models using human feedback. arXiv preprint arXiv:2302.12192, 2023. [22] Junzhe Li, Yutao Cui, Tao Huang, Yinping Ma, Chun Fan, Yiming Cheng, Miles Yang, Zhao Zhong, and Liefeng Bo. Mixgrpo: Unlocking flow-based grpo efficiency with mixed ode-sde. arXiv preprint arXiv:2507.21802, 2025. [23] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022. [24] Jie Liu, Gongye Liu, Jiajun Liang, Yangguang Li, Jiaheng Liu, Xintao Wang, Pengfei Wan, Di Zhang, and Wanli Ouyang. Flow-grpo: Training flow matching models via online rl. arXiv preprint arXiv:2505.05470, 2025. [25] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003, 2022. [26] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models. Machine Intelligence Research, 22(4):730–751, 2025. [27] Yuhang Ma, Xiaoshi Wu, Keqiang Sun, and Hongsheng Li. Hpsv3: Towards wide-spectrum human preference score. In 2025 IEEE/CVF International Conference on Computer Vision (ICCV), pages 15086–15095. IEEE, 2025. [28] David McAllister, Miika Aittala, Tero Karras, Janne Hellsten, Angjoo Kanazawa, Timo Aila, and Samuli Laine. Finite difference flow optimization for rl post-training of text-to-image models. arXiv preprint arXiv:2603.12893, 2026. [29] Bernt Oksendal. Stochastic differential equations: an introduction with applications. Springer Science & Business Media, 2013. [30] Peter Potaptchik, Cheuk-Kit Lee, and Michael S Albergo. Tilt matching for scalable sampling and fine-tuning. arXiv preprint arXiv:2512.21829, 2025. [31] Mihir Prabhudesai, Anirudh Goyal, Deepak Pathak, and Katerina Fragkiadaki. Aligning text-toimage diffusion models with reward backpropagation. 2023. [32] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017. [33] Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Yang Wu, et al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300, 2024. [34] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020. [35] Masatoshi Uehara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng, Tommaso Biancalani, and Sergey Levine. Finetuning of continuous-time diffusion models as entropy-regularized control. arXiv preprint arXiv:2402.15194, 2024. 12

[36] Bram Wallace, Meihua Dang, Rafael Rafailov, Linqi Zhou, Aaron Lou, Senthil Purushwalkam, Stefano Ermon, Caiming Xiong, Shafiq Joty, and Nikhil Naik. Diffusion model alignment using direct preference optimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8228–8238, 2024. [37] Jing Wang, Jiajun Liang, Jie Liu, Henglin Liu, Gongye Liu, Jun Zheng, Wanyuan Pang, Ao Ma, Zhenyu Xie, Xintao Wang, et al. Grpo-guard: Mitigating implicit over-optimization in flow matching via regulated clipping. arXiv preprint arXiv:2510.22319, 2025. [38] Chenfei Wu, Jiahao Li, Jingren Zhou, Junyang Lin, Kaiyuan Gao, Kun Yan, Sheng-ming Yin, Shuai Bai, Xiao Xu, Yilei Chen, et al. Qwen-image technical report. arXiv preprint arXiv:2508.02324, 2025. [39] Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li, Ming Ding, Jie Tang, and Yuxiao Dong. Imagereward: Learning and evaluating human preferences for text-to-image generation. Advances in Neural Information Processing Systems, 36:15903–15935, 2023. [40] Shuchen Xue, Chongjian Ge, Shilong Zhang, Yichen Li, and Zhi-Ming Ma. Advantage weighted matching: Aligning rl with pretraining in diffusion models. arXiv preprint arXiv:2509.25050, 2025. [41] Zeyue Xue, Jie Wu, Yu Gao, Fangyuan Kong, Lingting Zhu, Mengzhao Chen, Zhiheng Liu, Wei Liu, Qiushan Guo, Weilin Huang, et al. Dancegrpo: Unleashing grpo on visual generation. arXiv preprint arXiv:2505.07818, 2025. [42] Kai Yang, Jian Tao, Jiafei Lyu, Chunjiang Ge, Jiaxin Chen, Weihan Shen, Xiaolong Zhu, and Xiu Li. Using human feedback to fine-tune diffusion models without any reward model. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8941–8951, 2024. [43] Huizhuo Yuan, Zixiang Chen, Kaixuan Ji, and Quanquan Gu. Self-play fine-tuning of diffusion models for text-to-image generation. Advances in Neural Information Processing Systems, 37: 73366–73398, 2024. [44] Chujie Zheng, Shixuan Liu, Mingze Li, Xiong-Hui Chen, Bowen Yu, Chang Gao, Kai Dang, Yuqiong Liu, Rui Men, An Yang, et al. Group sequence policy optimization. arXiv preprint arXiv:2507.18071, 2025. [45] Kaiwen Zheng, Huayu Chen, Haotian Ye, Haoxiang Wang, Qinsheng Zhang, Kai Jiang, Hang Su, Stefano Ermon, Jun Zhu, and Ming-Yu Liu. Diffusionnft: Online diffusion reinforcement with forward process. arXiv preprint arXiv:2509.16117, 2025.

13

A

Organization of Appendix

The organization of the appendix is as follows. Appx. B provides background in stochastic calculus. Appx. C gives all the proofs of the theorems in the main text. Appx. D gives the derivations in Table 1. Appx. E discusses the related works. Appx. F provides experimental details for Sec. 4 and shows some additional results. Finally, qualitative examples are shown in Appx. G.

B

Backgrounds for Stochastic Analysis

In this section, we recall some basic results from stochastic calculus. Please see Oksendal [29] or Karatzas and Shreve [18] for details. Assume the probability space is (Ω, F, P). A d-dimensional stochastic process is a parametrized collection of random vectors {xt }t∈T , where T is the time interval. For each ω ∈ Ω we can consider the function t → xt (ω), t ∈ T , which is called the path of xt . Assume x is jointly measurable in (t, ω), then the stochastic process is a probability measure P on the measurable space ((Rd )T , B), where (Rd )T contains all functions from T to Rd . In this paper, we only consider the stochastic process with continuous sample paths; then the stochastic process is a probability measure on the measurable space (C(T, Rd ), B), where C(T, Rd ) contains all continuous functions from T to Rd . In this paper, we focus on the diffusion process, which can be defined by its generator: Definition 5. The infinitesimal generator of a stochastic process (xt )t for a function ϕ(x) is E[ϕ(xt+s )|xt = x] − ϕ(x) , s where ϕ is a smooth function. For an Itô process defined as the solution to the SDE dxt = f (xt , t)dt + gt dwt , the generator is Lt ϕ(x) = lim

s→0+

Lt =

d X i=1

d

f i (x, t)∂i +

gt2 X 2 ∂ , 2 i=1 i

where

∂i =

∂ . ∂xi

In stochastic calculus, Itô’s lemma is a useful tool for calculating derivatives as the chain rule in calculus. Lemma 6 (Itô’s lemma; [29], Thm. 4.2.1). Assume xt is the solution of SDE dxt = f (xt , t) dt + gt dwt . Then for any function f ∈ C 2 (Rd × R), we have   g2 df (xt , t) = ∂t f (xt , t) + ∇x f (xt , t) · f (xt , t) + t ∆f (xt , t) dt + gt ∇x f (xt , t) · dwt . 2 For the calculation of the value function, the following Kolmogorov’s backward equation is useful. Proposition 7 (Kolmogorov’s Backward Equation; [29], Thm. 8.1.1). Let Vt (x) = E[R(xT , c)|xt = x] be the value function. Then V is the solution of the following PDE ∂V + At V = 0, VT (x) = R(x, c), ∂t where xt satisfies the SDE dxt = f (xt , t) dt + gt dwt and At is its generator: Af (x, t) = f (x, t) · ∇f (x, t) +

gt2 ∆f (x, t), 2

∀f ∈ C 2 (Rd × R).

For the calculation of the importance sampling ratio between two path measures of diffusion processes, Girsanov’s theorem is useful: Theorem 8 (Girsanov; [29], Thm. 8.6.6). Let xt , yt solve the coupled SDEs on Rd dxt = u(xt , t) dt + gt dwt ,

dyt = v(yt , t) dt + gt dwt ,

with the same initial condition and diffusion coefficient, and let P, Q be their path measures. Assume RT u, v are sufficiently smooth and satisfy Novikov’s condition E[exp( 12 0 ∥(∥u − v)/gt ∥2 dt)] < ∞. Then ! Z T Z T dP 2 1 1 = exp ∥(u − v)(yt , t)∥ dt . gt (u − v)(yt , t) · dwt − 2gt2 dQ 0 0 14

C

Proofs

C.1

Setup: path-space importance sampling

Let vbase , vθ , vref : Rd × [0, 1] → Rd be three smooth velocity fields, and write Q, Pθ , Pref for the corresponding path measures on C([0, 1], Rd ): r  
ηt 2tηt Pθ : dxt = vθ (xt , t) + xt + (1 − t)vθ (xt , t) dt + dwt , 1−t 1−t r  
2tηt ηt base base base base Q : dxt = vbase (xt , t) + xt + (1 − t)vbase (xt , t) dt + dwt , 1−t 1−t r  
ηt 2tηt ref ref ref x + (1 − t)v (x , t) dt + dwt , Pref : dxref = v (x , t) + ref t ref t t 1−t t 1−t ref where x1 , xbase 1 , x1 ∼ N (0, I).

Then the original objective Eq. (3) becomes max Ec L(θ, c), where L(θ, c) = EPθ [R(x0 , c)] − βKL(Pθ ∥Pref ). If we want to use the importance sampling technique, the objective becomes      dPθ dPθ dPθ L(θ, c) = EQ R(x0 , c) − βEQ log , dQ dQ dPref θ where dP dQ is the Radon-Nikodym derivative of Pθ with respect to Q. By Thm. 8, Z 1p Z 1 2 M1 (θ) = − 2w̃1−r ∆vθ (x1−r , 1 − r) · dwr − w̃1−r ∥∆vθ (x1−r , 1 − r)∥ dr,

0

(11)

0 2

t) · 1−t where r = 1 − t, ∆vθ := vθ − vbase and w̃t = (1+η 4ηt t . The results recover Eq. (5).

C.2

Proof of Prop. 1

Proposition 9 (Stochastic-integral path-space estimator). Assume that the clipping or a small step size keeps | exp(M1 (θ− )) − 1| < ϵ and EQ [∥A∇θ M1 (θ)∥] < M bounded, where θ− denotes the stop-gradient version of the current parameters. Then −∇θ Lpolicy (θ, c) →   Z 1p Z 1 2 ∇θ EQ A(x0 , c) 2w̃1−r ∆vθ (x1−r , 1 − r) · dwr + A(x0 , c) w̃1−r ∥∆vθ (x1−r , r)∥ dr , 0

0

as ϵ → 0. Proof. Consider the policy-gradient term Lpolicy (θ, c) = EQ [exp(M1 (θ))A(x0 , c)] of Eq. (5). The gradient of the policy-gradient term is given by ∇θ Lpolicy (θ, c) = EQ [A(x0 , c) exp(M1 (θ− ))∇θ M1 (θ)]. Under the assumption | exp(M1 (θ− )) − 1| < ϵ, ∥∇θ Lpolicy (θ, c) − EQ [A(x0 , c)∇θ M1 (θ)]∥     = EQ (exp(M1 (θ− )) − 1)A∇θ M1 (θ) ≤ EQ | exp(M1 (θ− )) − 1|∥A∇θ M1 (θ)∥ ≤ ϵM. By Eq. (11), M1 (θ) = −

Z 1p Z 1 2 2w̃1−r ∆vθ (x1−r , 1 − r) · dwr − w̃1−r ∥∆vθ (x1−r , r)∥ dr. 0

0

Then we conclude that −∇θ Lpolicy (θ, c) → −EQ [A(x0 , c)∇θ M1 (θ)] =   Z 1 Z 1p 2 ∇θ EQ A(x0 , c) 2w̃1−r ∆vθ (x1−r , 1 − r) · dwr + A(x0 , c) w̃1−r ∥∆vθ (x1−r , r)∥ dr , 0

0

as ϵ → 0. 15

C.3

Proof of Thm. 2

Convention on time direction. Before starting the proof, we fix the time convention. The FM-time t ∈ [0, 1] used in the main text has t = 0 at the clean sample and t = 1 at pure noise; under this convention, the sampling process runs backward from t = 1 to t = 0. For the Itô calculus below, we use a sampling time r := 1 − t ∈ [0, 1] and denote yr := x1−r = xt that runs forward during sampling, with r = 0 at the initial noise y0 = x1 and r = 1 at the final sample y1 = x0 . Theorem 10 (Variance-reduced path-space estimator). Let q be the path measure of the sampling SDE Eq. (2) with base drift vbase and noise schedule ηt > 0, and let Vt (x, c) := Eq [A(x0 , c) | xt = x] be the associated value function. Additionally, assume that the SDE coefficients and Vt are smooth. Under the trust-region condition of Prop. 1, the policy gradient of Eq. (5) admits the (locally) variance-reduced representation −Lpolicy (θ, c) = θ Eq

Z 1 

2

(1 + ηt )∆vθ (xt , t) · ∇xt Vt (xt , c) + A(x0 , c)w̃t ∥∆vθ (xt , t)∥



 dt . (12)

0

p sto d Let ∇V t,∆t := A(x0 , c) (1 − t)/(2ηt t∆t) ϵt be the local discretization of Term (A) in Eq. (6). sto d Assume ∇V t,∆t is square-integrable for any interior time t ∈ (0, 1). When ∆t ↓ 0, we have   sto d t,∆t ≥ VarQ (∇x Vt (xt , c)) . VarQ ∇V t

(13)

Proof. We prove the identity for the policy-gradient term; the quadratic Term (B) in Eq. (6) is carried through unchanged. To avoid mixing the two time directions, we use the forward sampling time process yr := x1−r for r ∈ [0, 1]. Then the sampling SDE (Eq. (2)) becomes r h
i η1−r 2(1 − r)η1−r dwr , Q : dyr = − vbase (xr , 1 − r) + yr + rvbase (yr , 1 − r) dr + r {z } | {zr } | :=gr

:=µr (yr )

where y0 ∼ N (0, I). Under Q, we simply denote the sampling process as dyr = µr (yr ) dr + gr dwr .

(14)

All calculations can first be performed on [ε, 1 − ε] and then extended to [0, 1] by square-integrability and a limiting argument. Let p (15) Hr := 2w̃1−r ∆vθ (yr , 1 − r), Z r Ir := Hs · dws , 0

V̄r (x, c) := EQ [A(y1 , c) | yr = x],

Vr := V̄r (yr , c).

By square-integrability, Ir is an Itô martingale. The process Vr = EQ [A(y1 , c) | Fr ] is also a martingale by the tower property, with V1 = A(y1 , c) = A(x0 , c). The Kolmogorov backward equation for V̄r under Eq. (14) is ∂r V̄r + µr · ∇x V̄r +

gr2 ∆x V̄r = 0. 2

(16)

Applying Itô’s formula and using Eq. (16) eliminates the drift term: 1 dVr = ∂r V̄r (yr , c) dr + ∇x V̄r (yr , c) · dyr + ∆V̄r (yr , c) d[yr , yr ], 2 g2 = ∂r V̄r (yr , c) dr + ∇x V̄r (yr , c) · (µr (yr ) dr + gr dwr ) + r ∆V̄r (yr , c) dr,   2 gr2 = ∂r V̄r (yr , c) + ∇x V̄r (yr , c) · µr (yr ) + ∆V̄r (yr , c) dr + gr ∇x V̄r (yr , c) · dwr 2 | {z } =0, by Eq. (16)

= gr ∇x V̄r (yr , c) · dwr ,

(17) 16

Pd ∂ 2 where ∆ = i=1 ∂x 2 is the Laplace operator and [·, ·] is the quadratic variation of the diffusion i process. It remains to identify the cross-variation of the two martingales. Eq. (15) and Eq. (17) give p d[V, I]r = gr 2w̃1−r ∆vθ (yr , 1 − r) · ∇x V̄r (yr , c) dr = (1 + η1−r ) ∆vθ (yr , 1 − r) · ∇x V̄r (yr , c) dr, where the second equality follows from s r p 2tηt (1 + ηt )2 (1 − t) gr 2w̃1−r = = 1 + ηt , 1−t 2ηt t

t = 1 − r.

Itô’s product rule now yields d(Vr Ir ) = Vr dIr + Ir dVr + d[V, I]r
= Vr Hr + Ir gr ∇x V̄r (yr , c) · dwr + (1 + η1−r )∆vθ (yr , 1 − r) · ∇x V̄r (yr , c) dr. After integration, Z 1 V1 I1 = V0 I0 +

  Vr Hr + gr Ir ∇x V̄r (yr , c) · dwr

0

Z 1 (1 + η1−r )∆vθ (yr , 1 − r) · ∇x V̄r (yr , c) dr.

+ 0

Take expectation on both sides and by I0 = 0, we get   Z 1p 2w̃1−r ∆vθ (yr , 1 − r) · dwr EQ [V1 I1 ] = EQ A(x0 , c) 0

Z 1

 (1 + η1−r )∆vθ (yr , 1 − r) · ∇V̄r (yr , c) dr

= EQ 0

Z 1 + EQ |

0 =0,


Vr Hr + gr Ir ∇x V̄r (yr , c) · dwr {z } by the martingale property of Itô integral

Z 1

 (1 + η1−r )∆vθ (yr , 1 − r) · ∇V̄r (yr , c) dr .

= EQ 0

Changing variables back to FM-time t = 1 − r turns the right-hand side into the first term of Eq. (12). Restoring the unchanged quadratic Term (B) gives −Lpolicy (θ, c) = θ Eq

Z 1 

2

(1 + ηt )∆vθ (xt , t) · ∇xt Vt (xt , c) + A(x0 , c)w̃t ∥∆vθ (xt , t)∥



 dt .

0

which is exactly Eq. (12). We now prove Eq. (13). For a positive sampling-time step ∆r,
sto A(y1 , c) d ∇V wr+∆r − wr , t,∆r = gr ∆r

t = 1 − r.

By the tower property, the martingale identity Vr+∆r = EQ [A(y1 , c) | Fr+∆r ], and Eq. (17), h i sto d t,∆r | Fr = 1 EQ [(Vr+∆r − Vr )(wr+∆r − wr ) | Fr ] EQ ∇V gr ∆r "Z # r+∆r 1 = EQ gs ∇x V̄s (ys , c) ds Fr . gr ∆r r By Jensen’s inequality and Fubini’s theorem, h i sto d t,∆r | Fr − ∇x V̄r (yr , c) EQ ∇V ≤ 2

Z r+∆r 1 gs ∇x V̄s (ys , c) − gr ∇x V̄r (yr , c) L2 (Q) ds gr ∆r r L (Q) −−−→ 0. ∆r↓0

17

sto

sto

d t,∆t | Fr ] = Eq [∇V d t,∆t | xt ]. By the property of conditional By the Markov property, EQ [∇V variance,     h  i sto sto sto d t,∆t = VarQ EQ [∇V d t,∆t | xt ] + Eq Varq ∇V d t,∆t | xt VarQ ∇V   sto d ≥ VarQ EQ [∇V −−→ VarQ (∇xt Vt (xt , c)) . t,∆t | xt ] − ∆t↓0

Identifying the positive step magnitudes ∆r = ∆t proves Eq. (13). C.4

Proof of Prop. 3

Proposition 11. The gradient of the value function admits the closed-form expression ∇xt Vt (xt , c) = Eq(x0 |xt ) [∇xt log q(xt | x0 ) A(x0 , c)] − Vt (xt , c) Eq(x0 |xt ) [∇xt log q(xt | x0 )] . Let v(xt , x0 ) := (xt − x0 )/t denote the FM conditional velocity. Under q with vbase = E[v | xt ] ≈ vold and ηt ≡ 1 (i.e., q is the reverse of the forward noising process), this specialises to 
 ∇xt Vt (xt , c) = − 1−t t Eq(x0 |xt ) A(x0 , c) v(xt , x0 ) − vbase (xt , t) . The corresponding one-sample estimator on a rollout trajectory x0:1 is
det d t (xt , c) := − 1−t A(x0 , c) v(xt , x0 ) − vold (xt , t) . ∇V t Proof. Under Q with ηt ≡ 1, the sampling SDE reduces to the reverse of the forward noising process (Sec. 2.1), and the forward transition kernel is q(xt | x0 ) = N (xt ; (1 − t)x0 , t2 I). Differentiating the value function, we get Ex0 [R(x0 , c)q(xt |x0 )] Ex0 [q(xt |x0 )] Ex0 [∇xt q(xt |x0 )R(x0 , c)] Ex0 [q(xt |x0 )R(x0 , c)] Ex0 [∇xt q(xt |x0 )] = − 2 Ex0 [q(xt |x0 )] (Ex0 [q(xt |x0 )])

∇xt Vt (xt ) = ∇xt

Ex0 [q(xt |x0 )∇xt log q(xt |x0 )R(x0 , c)] Ex0 [q(xt |x0 )] Ex0 [q(xt |x0 )R(x0 , c)] Ex0 [q(xt |x0 )∇xt log q(xt |x0 )] − 2 (Ex0 [q(xt |x0 )]) = Eq(x0 |xt ) [∇xt log q(xt |x0 )R(x0 , c)] − Vt (xt )Eq(x0 |xt ) [∇xt log q(xt |x0 )]. =

(18)

The Gaussian score of q(xt |x0 ) is xt − (1 − t)x0 1−t 1 =− · v(xt , x0 ) − xt , 2 t t t where v(xt , x0 ) := (xt − x0 )/t is the FM conditional velocity. Substituting into Eq. (18) and replacing R with the group-relative advantage A (it can be viewed as a normalized reward function), ∇x log q(xt | x0 ) = −

∇x Vt (xt , c) = − 1−t t Eq(x0 |xt ) [A(x0 , c)(v(xt , x0 ) − E[v(xt , x0 )|xt ]] . Any good approximation vold (xt , t) ≈ E[v|xt ] (which holds whenever vold is a well-trained diffusion model) gives ∇x Vt (xt , c) = − 1−t t Eq(x0 |xt ) [A(x0 , c) (v(xt , x0 ) − vold (xt , t))] , d t = − 1−t A(v(xt , x0 ) − vold (xt , t)). whose one-sample Monte Carlo estimator is ∇V t C.5

Proof of Prop. 4

Proposition 12 (Properties of the value-gradient estimators). Let A(x0 , c) be square-integrable. Assume vbase = vold , then the following arguments hold: (i). The discretized stochastic one-sampleh estimator is asymptotically unbiased in the continuousi sto d time Itô limit of Thm. 2: lim∆t→0 E ∇V t,∆t | xt = ∇xt Vt (xt , c); 18

(ii). Assume ηt ≡ 1,hvold (xt , t) =i E[v(xt , x0 ) | xt ], then the deterministic one-sample estimator det d t | xt = ∇x Vt (xt , c); is unbiased: E ∇V t

(iii). Under the assumption of (ii) and h = 1, the KDE estimator is unbiased and has smaller variance than the deterministic one-sample estimator: h i     KDE j KDE det dt dt d t | xjt . E ∇V (xt , c) | xjt = ∇xt Vt (xjt , c), Var ∇V | xjt ≤ Var ∇V Proof. Again, we use the forward sampling time r = 1 − t from the proof of Thm. 2. Proof of (i).

Let Fr be the filtration generated by the proposal process. For a step ∆r > 0, define ϵr,∆r :=

wr+∆r − wr √ . ∆r

Recall from Eq. (17) that the value martingale Vr = EQ [A(y1 , c) | Fr ] satisfies r 2tηt , t = 1 − r. gr = dVr = gr ∇x V̄r (yr , c) · dwr , 1−t By the tower property and the fact that wr+∆r − wr has zero conditional mean, EQ [A(y1 , c)(wr+∆r − wr ) | Fr ] = EQ [E[A(y1 , c)(wr+∆r − wr ) | Fr+∆r ] | Fr ] = EQ [E[A(y1 , c) | Fr+∆r ](wr+∆r − wr ) | Fr ] = EQ [Vr+∆r (wr+∆r − wr ) | Fr ] = EQ [(Vr+∆r − Vr )(wr+∆r − wr ) | Fr ] = gr ∇x V̄r (yr , c) ∆r + o(∆r), where the last equality follows from the Itô isometry and the continuity of the martingale integrand. Consequently,   ϵr,∆r Fr = gr ∇x V̄r (yr , c) + o(1). EQ A(y1 , c) √ ∆r The stochastic estimator in Sec. 3.2 can be written as d t,∆r = 1 A(y1 , c) ϵ √r,∆r . ∇V gr ∆r sto

Taking ∆r → 0 and using the Markov property, so that conditioning on Fr reduces to conditioning on yr = xt , yields h i sto d t,∆r | xt = ∇x Vt (xt , c). lim EQ ∇V t ∆r→0

Identifying the positive sampling-time step ∆r with the discretization magnitude ∆t in the proposition proves (i). Thus the stochastic one-sample estimator is asymptotically unbiased. Proof of (ii). When ηt ≡ 1, the transition kernel is q(xt | x0 ) = N (xt ; (1 − t)x0 , t2 I). Under the assumption vold (xt , t) = E[v(xt , x0 ) | xt ], by Prop. 3, ∇xt Vt (xt , c) = −

 1−t  E A(x0 , c) v(xt , x0 ) − vold (xt , t) | xt . t det

d t . Therefore The random vector inside this conditional expectation is precisely ∇V h i det d t | xt = ∇x Vt (xt , c), E ∇V t which proves the unbiasedness of the deterministic one-sample estimator. 19

Proof of (iii). The results can be directly deduced from Prop. 2 of Bertrand et al. [4] (see also Exercise 7.32 of Casella and Berger [6]), we give the detailed proof here for completeness. Fix a query sample (xjt , c) from the rollout group and write q0 for the endpoint marginal under the proposal. At h = 1, the kernel in Eq. (10) is proportional to the likelihood q(xjt | x0 ); hence Bayes’ rule gives q(x0 | xjt ) =

K(xjt , x0 )q0 (x0 ) Ez∼q0 [K(xjt , z)]

.

Conditioned on xjt , the anchor endpoint xj0 is therefore distributed according to q(x0 | xjt ), while the remaining G − 1 endpoints are independent samples from q0 . To expose the resulting Rao–Blackwell structure, introduce an auxiliary index J ∼ Unif{1, . . . , G}, sample XJ from q(x0 | xjt ), and sample {Xi }i̸=J independently from q0 . Because the KDE estimator is symmetric in the group elements, this random relabeling does not change its distribution. Define F (Xi ) := −

1−t A(Xi , c) v(xjt , Xi ) − vold (xjt , t) , t

Ki := K(xjt , Xi ).

Then by Bayes’ rule, Ki P(J = i | X1 , . . . , XG ) = PG

ℓ=1 Kℓ

.

Since the KDE estimator is the conditional expectation of the deterministic anchor estimator: PG KDE j Ki F (Xi ) d (xt , c). = ∇V E[F (XJ ) | X1 , . . . , XG ] = i=1 PG t i=1 Ki Applying the tower property and the result of (ii), h i h i KDE j dt E ∇V (xt , c) | xjt = E F (XJ ) | xjt = ∇xt Vt (xjt , c), which means the KDE estimator is unbiased under the condition. Finally, applying the property of conditional variance, we get   h   i   Var F (XJ ) | xjt = E Var F (XJ ) | X1:G , xjt xjt + Var E[F (XJ ) | X1:G , xjt ] xjt   KDE j dt ≥ Var ∇V (xt , c) xjt . det

d , then Since F (XJ ) has the same conditional distribution as ∇V t     KDE det j d d Var ∇V t | xt ≤ Var ∇V t | xjt .

D

Results in Table 1

In this section, we provide detailed derivations of results in Table 1. D.1

Flow-GRPO, GRPO-Guard and TempFlow-GRPO

Flow-GRPO [24]. The stochastic-integral path-space estimator of the Flow-GRPO objective is Eq. (6). Applying an Euler–Maruyama discretization to Eq. (6) gives   Z 1p Z 1 2 GRPO −L (θ, c) =θ Eq A(x0 , c) 2w̃1−r ∆vθ · dwr + A(x0 , c) w̃1−r ∥∆vθ ∥ dr 0 0 " N −1 #  X p p ≈ Eq A(x0 , c) 2w̃1−ri ∆vθ · ∆ri ϵi + A(x0 , c)w̃1−ri ∥∆vθ ∥2 ∆ri i=0

= Eq

" N −1 X i=0

 p 1 A(x0 , c) 2w̃1−ri ∆vθ · √ ϵi + A(x0 , c)w̃1−ri ∥∆vθ ∥2 ∆ri ∆ri

20

#

With the change of variable t = 1 − r, let ∆ri = ri+1 − ri = ti − ti+1 = ∆ti > 0. The sum can then be written entirely in terms of t as " N −1 #  X p 1 2 −LGRPO (θ, c) ≈ Eq A(x0 , c) 2w̃ti ∆vθ · √ ϵi + A(x0 , c)w̃ti ∥∆vθ ∥ ∆ti . (19) ∆ti i=0 2

t ) 1−t Substituting w̃t = (1+η 4ηt t into Eq. (19) directly gives the Flow-GRPO expressions in Table 1: r sto 1−t d ∇V t = s(t)A(x0 , c)ϵ, s(t) = , 2ηt t∆t (1 + ηt )2 1 − t ŵ1 (A, t) = A , ŵ2 (t) = 1 + ηt . 4ηt t

GRPO-Guard [37]. The difference between Flow-GRPO and GRPO-Guard is that GRPO-Guard drops the regularization Term (B) and normalizes Term (A) in Eq. (6), corresponding to changes in weight ŵ1 (A, t) and ŵ2 (t): s 2ηt t ŵ1 (A, t) = 0, ŵ2 (t) = (1 + ηt ) . (1 − t)∆t TempFlow-GRPO [14]. TempFlow-GRPO first rolls out an ODE trajectory. At each point xt along this trajectory, it takes a single SDE step to generate a group of branch states xt−∆t , each of which is then rolled out to the endpoint using an ODE solver. The resulting trajectories form a branching structure, which is pit branched-SDE sampler. TempFlow-GRPO also reweights the √ the reason we call policy loss by gt ∆t, where gt = 2ηt t/(1 − t), thus its weights are s r 2ηt t∆t (1 + ηt )2 2(1 − t)∆t , ŵ2 (t) = (1 + ηt ) . ŵ1 (A, t) = A 4 ηt t 1−t D.2

AWM and DiffusionNFT

Advantage Weighted Matching (AWM) [40]. Recall from Sec. 2.2 that AWM first draws a rollout sample x0 ∼ pold (· | c) (with Flow-SDE or Flow-ODE, based on different rewards) and then constructs an intermediate state by forward noising: xt = (1 − t)x0 + tϵ,

ϵ ∼ N (0, I).

The training objective is 6 h i 2 −LAWM (θ) = Ec,x0 ∼pold , t∼p(t), ϵ∼N (0,I) A(x0 , c)∥vθ (xt , t, c) − v(xt , x0 )∥ .

(20)

In our design space, we consider the reverse sampling process Eq. (2) to obtain x[0:1] . To close the gap between the forward and reverse process, we construct the reverse sampling process with xt − x 0 . t Under such choices, the reverse sampling process Eq. (2) is exactly the time reversal of the forward noising path above, which means that (x0 , xt ) constructed by the reverse sampling process and forward noising process share the same joint distributions. Thus, the AWM is equivalent to using vbase (xt , t) = E[v(xt , x0 ) | xt ] and ηt ≡ 1 in our design space. Under this condition, Prop. 3 gives vbase (xt , t) = E[v(xt , x0 ) | xt ] with ηt = 1,

where v(xt , x0 ) :=

det d (xt , c) = −s(t)A(x0 , c) v(xt , x0 ) − vold (xt , t) , ∇V t

s(t) :=

1−t . t

(21)

6 The original AWM objective additionally includes an importance-sampling ratio between the current and rollout policies. This ratio is exactly one in the on-policy setting used by the practical algorithm and remains close to one under the mildly off-policy updates considered in practice. We therefore omit it from the present derivation.

21

We specify the weights ŵ1 (A, t) =

A , ∆t

ŵ2 (t) =

2 . s(t)∆t

Substituting these weights and the estimator Eq. (21) into Eq. (9) gives Z 1 h i 1 2 −Lpolicy (θ) = Eq A∥∆vθ ∥ − 2A ⟨∆vθ , v − vold ⟩ dt 0 ∆t Z 1 h i 1 2 Eq A(x0 , c)∥vθ (xt , t, c) − v(xt , x0 )∥ dt, =θ 0 ∆t .

(22)

For the ∆t term, partition [0, 1] into n timestep bins Ii of widths ∆ti , and write ∆t(t) = ∆ti for t ∈ Ii . We define the density by p(t) :=

1 , n∆ti

t ∈ Ii .

Using 1/∆t(t) = np(t), Eq. (22) becomes h i 2 −Lpolicy (θ) =θ n Ec,x0 ∼pold , t∼p(t), ϵ∼N (0,I) A(x0 , c)∥vθ (xt , t, c) − v(xt , x0 )∥ . This is the same as the AWM loss Eq. (20) up to the positive constant factor n, which does not change the optimization process. The timesteps sampled by the original AWM objective from p(t) are precisely the points used to discretize the continuous-time integral in our formulation. DiffusionNFT [45]. Similar to AWM, DiffusionNFT first rolls out with the DPM-Solver [26] to obtain x0 and then samples xt = (1 − t)x0 + tϵ. It minimizes the objective h i
2 2 −LNFT (θ) = Ec, x0 ∼pold , xt ∼p(·|x0 ), t R̂(x0 , c) v+ − v + 1 − R̂(x0 , c) v− − v , where v− = (1 + β)vold (xt , t, c) − βvθ (xt , t, c),

v+ = (1 − β)vold (xt , t, c) + βvθ (xt , t, c), A(x0 , c) + 1 , R̂(x0 , c) = 2

v = (xt − x0 )/t.

Define ∆vθ := vθ − vold ,

d := v − vold .

Then v+ − v = β∆vθ − d and v− − v = −β∆vθ − d. Substituting these identities and expanding the two squares gives h i 2 2 −LNFT (θ) = E R̂∥β∆vθ − d∥ + (1 − R̂)∥−β∆vθ − d∥ h i 2 2 = E β 2 ∥∆vθ ∥ − 2β(2R̂ − 1)∆vθ · d + ∥d∥   2A 2 =θ β 2 Ec, x0 ∼pold , xt ∼p(·|x0 ), t ∥∆vθ ∥ − ∆vθ · (v − vold ) , β The scaling factor β 2 does not change the optimization process; thus, we have   2A 2 NFT −L (θ) = Ec, x0 ∼pold , xt ∼p(·|x0 ), t ∥∆vθ ∥ − ∆vθ · (v − vold ) . β Therefore, similarly to AWM, DiffusionNFT is recovered in our design space by choosing vbase = E[v | xt ], ηt = 1, and the deterministic one-sample estimator Eq. (21). The weights functions are specified by ŵ1 (t) = 1/∆t and ŵ2 (t) = 2/(βs(t)∆t). 22

E

Related Work

Reverse-sampling-based RL. DDPO [5] and DPOK [12] pioneered policy-gradient fine-tuning for DDPM-style models by treating reverse denoising as a finite-horizon MDP. Flow-GRPO [24] and Dance-GRPO [41] extend this paradigm to flow-matching models by converting the sampling ODE into a marginally equivalent SDE and applying clipped GRPO-style updates. Subsequent methods refine this recipe along different axes: GRPO-Guard [37] regulates and rescales per-step likelihood ratios, TempFlow-GRPO [14] branches trajectories at selected timesteps, and MixGRPO [22] confines stochastic exploration to a sliding time window. FDFO [28] follows a different route and estimates improvement directions through finite differences rather than analytic policy gradients. Our framework disentangles the estimator, weighting, and sampling choices underlying these methods and explains their efficiency differences through the resulting value-gradient variance and scale. Complementary analyses also seek to organize this rapidly growing design space. Choi et al. [8] separate the policy-gradient objective, likelihood estimator, and sampler, and identify likelihood estimation as the dominant empirical factor. We refine the objective axis further by decomposing it into value-gradient estimation and the associated on- and off-policy weights. The concurrent work Reward Score Matching [20] develops a score-space view that connects several reward-based finetuning objectives; in contrast, our path-space formulation explicitly recovers the training structures of AWM and DiffusionNFT and relates them to reverse-sampling policy gradients.

Forward-matching RL methods. An orthogonal line retains the analytic forward-noising process and injects reward information into a pretraining-style regression loss. Reward-Weighted Regression (RWR) [21] weights the diffusion training objective by endpoint rewards. AWM [40] and DiffusionNFT [45] similarly avoid the reverse-process MDP and construct reward- or advantage-dependent flow-matching targets from sampled clean endpoints and their analytically noised states. These methods report substantial improvements in training efficiency, but their relationship to reverse-sampling policy gradients and the role of their specialized timestep weights were previously unclear. Our pathspace derivation shows that they instantiate the same policy-gradient structure with a deterministic value-gradient estimator and different weight choices.

Stochastic-control and adjoint-based methods. A complementary line formulates KL-regularized reward fine-tuning as entropy-regularized stochastic optimal control. ELEGANT [35] develops this viewpoint for continuous-time diffusion models, while Adjoint Matching [10] introduces a memoryless noise schedule and turns the Bellman optimality condition into an on-policy regression of the value gradient; its pathwise estimator, however, requires reward gradients and a backward adjoint sweep along an SDE trajectory. Tilt Matching [30] instead relates the original and rewardtilted velocity fields through a regression objective and avoids both reward gradients and trajectory backpropagation by annealing through intermediate tilted models. Most closely related to our work, the concurrent Reinforce Adjoint Matching (RAM) method [3] takes a stochastic optimal control viewpoint, identifies the value gradient as the optimal adjoint, and combines the resulting fixed point with a REINFORCE identity to derive a forward-matching-type loss. In contrast, we derive the value-gradient form from path-space policy-gradient variance reduction, yielding a unified framework for both GRPO-style reverse-sampling algorithms and forward-based methods such as AWM and DiffusionNFT. Our framework further establishes an estimator-specific optimal weight mapping and introduces a new multi-sample KDE value-gradient estimator in the final recipe.

Reward-gradient and offline preference methods. When the reward is differentiable, it can be optimized directly by backpropagating through all or part of the sampling trajectory [9, 31, 39]. This approach provides a direct training signal but requires reward gradients and incurs substantial memory and computation costs when differentiation spans many denoising steps. A separate offline line adapts direct preference optimization to diffusion models, learning from paired human-preference data without online reward evaluation [36, 42, 43]. These methods operate under a different feedback regime from the black-box online rewards considered in our framework. 23

(a)

(b)

(c)

Figure 5: Sampling noise schedule ablations for the stochastic one-sample estimator, deterministic one-sample estimator, and our KDE estimator. For each estimator, we run experiments for sampler noise schedule ηt ∈ {0.005, 0.225} on the Pickscore reward.

(a)

(b)

Figure 6: Visualization of different weights

F

Experimental Details

F.1

Additional Results

In Fig. 5, we show the sampling noise schedule ablations for each estimator. We find that the stochastic estimator in GRPO-based methods performs better in the larger ηt setting (Fig. 5 (a)), while the deterministic estimator suffers from training instability in both settings, and the small noise scheduler performs slightly better (Fig. 5 (b)). Our estimator is not sensitive to different sampling noise scales and is stable across all these settings. In Fig. 6, we provide a visualization of different weight functions. F.2

Hyperparameters

We summarize the hyperparameters of our experiments in Table 2. In all the experiments, we disable CFG and use the batch-level advantage normalization. Table 2: Hyperparameters used throughout Sec. 4. Hyperparameter

Value

Group size G Train denoising steps T Inference denoising steps Rollout SDE noise schedule Optimizer KL coefficient β Batch size Number of batches per epoch Train CFG h in Eq. (10)

24 for SD 3.5, 16 for Qwen-Image 10 40 for SD 3.5, 50 for Qwen-Image 0.225 for stochastic estimator, 0.005 for deterministic and KDE estimator AdamW, lr = 3×10−4 , β1 = 0.9, β2 = 0.999 10−4 9 for SD 3.5, 4 for Qwen-Image 16 for SD 3.5, 4 for Qwen-Image Disable CFG in training d/2000, where d is the latent dimension

24

F.3

Robustness across random seeds

We repeat the SD3.5-Medium experiment with PickScore reward using three random seeds while keeping all other configurations fixed. Table 3 reports the mean and standard deviation at four matched training checkpoints. Table 3: PickScore across three random seeds (mean ± standard deviation). Our method achieves a higher mean at every checkpoint with small run-to-run variation. 18.58

36.44

54.28

72.12

22.85 ± 0.02886 22.95 ± 0.05240

23.06 ± 0.01618 23.28 ± 0.02239

23.16 ± 0.02153 23.48 ± 0.005547

23.22 ± 0.01274 23.56 ± 0.01461

Training time (GPU hours) DiffusionNFT Ours

At the final checkpoint, our method obtains 23.56 ± 0.01461, corresponding to a relative standard deviation of 0.062%; DiffusionNFT obtains 23.22 ± 0.01274 with a relative standard deviation of 0.055%. The consistent gap across checkpoints and seeds indicates that the improvement is not driven by a favorable random seed. F.4

Runtime comparison

The KDE estimator computes pairwise distances in O(G2 d) time. With G = 24 and latent dimension d = 65,536, this is on the order of 107 arithmetic operations per timestep and is efficiently parallelized on GPUs. It also reuses cached rollout latents, velocities, and rewards, so it requires no additional denoiser evaluation. Under matched settings, our update uses a trainable-policy forward and, when KL regularization is enabled, a frozen-reference forward. DiffusionNFT additionally evaluates the old policy at the renoised latent, requiring three rather than two denoiser forwards per optimized timestep. Since one denoiser forward costs approximately 1012 FLOPs, this saved network evaluation dominates the pairwise-distance overhead.

Table 4: Measured wall-clock training time per epoch under matched SD3.5-Medium settings. Method Training time per epoch (seconds) 115 97

DiffusionNFT KDE estimator (ours)

As shown in Table 4, our method saves 18 seconds per epoch, a 15.7% wall-clock reduction. Thus, the O(G2 d) KDE computation has little effect on end-to-end runtime, while eliminating one denoiser forward pass provides a pronounced speedup.

G

Qualitative Examples

The visualization grids (Fig. 7, Fig. 8,Fig. 9) contain prompts from the test set of Pickscore reward training. For each prompt, images are generated with the same inference configuration from SD3.5Medium after the RL training with GRPO-Guard, DiffusionNFT, and our method. The rows are organized by prompt split and the columns by method, enabling direct inspection of prompt alignment, perceptual quality, excessive saturation, and repetitive reward-hacking patterns without selecting examples solely by reward. As a complementary quantitative check, a panel of eight evaluators assesses randomly generated images drawn from a pool of 30 prompts spanning the same three splits. Evaluators are blinded to method identity and judge overall visual quality, prompt alignment, and aesthetic preference. Aggregating the evaluation decisions according to this protocol yields method-level preference shares of 51.4% for our method, 29.4% for GRPO-Guard, and 19.2% for DiffusionNFT. These results suggest that the reward gains do not merely reflect numerical over-optimization. 25

H

Limitations

Our experiments are conducted on SD3.5-M and Qwen-Image with a fixed group size G = 24 and a limited set of reward functions (PickScore, OCR, and GenEval). While these cover a representative range of reward types, the generalizability of our findings to other model architectures (e.g., video diffusion models), reward functions (e.g., human-preference models or verifiable task rewards) has not been systematically evaluated.

I

Social Impact

This work proposes a unified theoretical framework for reinforcement learning post-training of diffusion generative models, which can facilitate the alignment of generative models with human preferences and safety constraints, thereby contributing to the responsible development of generative AI systems. However, this work mainly focuses on theory and algorithms, so the work does not have a broader social impact.

26

SD3.5-M + GRPO-Guard

SD3.5-M + DiffusionNFT

SD3.5-M + Ours

a jung male cyborg with white hair sitting down on a throne in a dystopian world, digital art, epic

enigmatic black square building on top of a purple hill, smoke stacks

wide shot of detailed compact simple starship in space, battle, fast, simple volume, geometrical, star wars, sharp, dynamic, cyberpunk, revolutionary, lightweight, volvo, bertone, minecraft, funky, studio scene, high contrast, super realistic, volumetric lighting, glimmery light, product lighting, high detail, super realistic, high contrast, cinematic lighting, rtx, octane render, behance, 4k, highly detailed, professional photograph, terrazzo, advertisement, epic

Figure 7: Qualitative comparison across GRPO-Guard, DiffusionNFT and our method.

27

SD3.5-M + GRPO-Guard

SD3.5-M + DiffusionNFT

SD3.5-M + Ours

photo of Ford Focus RS, night time, city, city roads, Miami streets

Labrador in anime style, black and white photography, highly detailed,

Portrait frame Penthouse Magazine, expressive and deep symmetrical eyes, stunning face, naive and innocent facial expression, stunningly beautiful slender girl, shaggy haircut, white hair color, in a sci-fi tight-fitting jumpsuit, the girl looks up into the camera, against the background of a white sci-fi room, stunning quality, ultra-high resolution, photorealistic, 8k, hyperrealism, rembrandt lighting

Figure 8: Qualitative comparison across GRPO-Guard, DiffusionNFT and our method.

28

SD3.5-M + GRPO-Guard

SD3.5-M + DiffusionNFT

SD3.5-M + Ours

Watercolor painting of woodpecker, afternoon backlight, by greg rutkowski, by anders zorn

Raw Photo, masterpiece award winning close up of a massive timber wolf standing in the moonlight in the dark in the darkness

Still life, tullips in the vase, close up shot, high details, 32k, by Peder Monsted, Jeremy Mann, Daniel F Gerhartz, Aaron Grffin ultra detailed

Figure 9: Qualitative comparison across GRPO-Guard, DiffusionNFT and our method.

29

Related documents

• Designing Reinforcement Learning for Diffusion Models: A Unified Path-Space View#609600
• Official Review · NEURIPS 2024#787016
• TLCR: Token-Level Continuous Reward for Fine-grained Reinforcement Learning from Human Feedback#436429
• Bridging Continuous-time LQR and Reinforcement Learning via Gradient Flow of the Bellman Error#147105
• Q-Learning for Reachability in MEC-Free MDPs#1157052