F LOWAWR: O NLINE A DAPTIVE F LOW R EINFORCE MENT VIA A DVANTAGE -W EIGHTED R ECTIFICATION Zheming Fu1,2,3 , Ruizhe He2,4 , Wei Shang1 , Xiaoxiao Ma2,5 , Lei Wang1∗, Chang Liu3∗, Siming Fu2∗ 1
arXiv:2606.30376v1 [cs.LG] 29 Jun 2026
5
Beihang University, 2 Joy Future Academy, 3 Zhongguancun Academy, 4 Zhejiang University, University of Science and Technology of China
A BSTRACT Aligning generative flow models on continuous spaces via online reinforcement learning is constrained by intractable trajectory likelihoods. Existing densityapproximated 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 intragroup quality. We propose Flow Advantage-Weighted Rectification (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× to 5× 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.
1
I NTRODUCTION
The empirical success of Reinforcement Learning (RL) in the post-training of Large Language Models (LLMs) (Christiano et al., 2017; Ziegler et al., 2019; Achiam et al., 2023; Comanici et al., 2025; Yang et al., 2025) has motivated its extension to continuous generative spaces. Within this domain, Flow Matching (Lipman et al., 2022; Liu et al., 2022; Song et al., 2023) has emerged as a prevalent paradigm, offering a deterministic trajectory framework for modern diffusion models (Esser et al., 2024; Labs et al., 2024; Seedream et al., 2025; Wu et al., 2025). However, adapting RL to this continuous space introduces a structural constraint: the approximation of trajectory likelihood. Unlike AR models that compute exact discrete token probabilities, continuous flow models inherently lack tractable step-wise transition densities (Song et al., 2020). Ranging from offline applications of Direct Preference Optimization (DPO) (Wallace et al., 2024; Yang et al., 2024) to recent online policy gradient approaches (Liu et al., 2025; Xue et al., 2025), this likelihood bottleneck consistently compromises the computational efficiency and theoretical guarantees of RL on continuous spaces. To circumvent this likelihood barrier, a prevalent strategy discretizes the continuous generative process into a multi-step Markov Decision Process (MDP) (Black et al., 2023; Fan et al., 2023). Building upon this, recent GRPO-based methods (Liu et al., 2025; Xue et al., 2025) inject stochasticity via SDE samplers to construct tractable Gaussian transition kernels, enabling the evaluation of policy gradients. Although subsequent improvements attempt to interleave ODE and SDE steps (Li et al., 2025), this paradigm fails to resolve the forward inconsistency between stochastic training trajectories and deterministic inference rollouts. Furthermore, by failing to implicitly distill the guidance ∗
Corresponding authors.
1
signal into the vector field, these methods retain a structural reliance on CFG (Ho & Salimans, 2022) to maintain generation quality, thereby imposing computational overhead during inference. An alternative paradigm leverages implicit optimization and velocity field guidance to bypass explicit density estimation. Early reward-weighted methods (Lee et al., 2023; Dong et al., 2023) are typically restricted to offline settings and lack explicit penalization mechanisms for low-quality samples. Recent online RL frameworks, such as DiffusionNFT (Zheng et al., 2025), have reformulated policy optimization as a supervised velocity regression task, achieving notable improvements in both training efficiency and generation performance over GRPO-style methods. However, it relies on a heuristic definition of velocity correction direction and applies modifications with a fixed magnitude. Furthermore, its utilization of intra-group samples is restricted to computing a scalar normalized reward, which merely balances the weights between preset positive and negative objectives. Consequently, this simultaneous exploration operates strictly at the individual sample level, thereby failing to leverage relative intra-group quality to directly determine the magnitude of velocity updates. Herein, we propose Flow Advantage-Weighted Rectification (FlowAWR), an online RL paradigm for continuous generative spaces that operates without stochastic trajectory formulation or heuristic field modifications. Our framework originates from the exact closed-form optimal policy derived from reward maximization under a temporal KL divergence constraint. By mapping this optimal terminal distribution onto the intermediate probability paths, we analytically derive the corresponding optimal velocity field. We demonstrate that this optimal field is equivalent to rectifying the reference velocity field by an advantage-weighted expectation. Consequently, FlowAWR recasts flow RL as a supervised regression task toward this theoretical optimum. By incorporating a group-based advantage function, our framework quantifies the continuous relative quality of intra-group samples, replacing rigid binary push-pull operations with magnitude-aware rectifications. Our method preserves the architectural benefits of likelihood-free approaches by eliminating the forward inconsistency induced by SDE sampling, while its advantage-weighted updates implicitly distill alignment trajectories into the network parameters to enable CFG-free generation. To evaluate the efficacy of FlowAWR, we conduct systematic experiments on the SD3.5-Medium (2.5B parameters) model (Esser et al., 2024). Quantitative results indicate that, in the single-reward setting, FlowAWR achieves a 2× to 5× convergence acceleration alongside improved performance compared to DiffusionNFT. For instance, when optimizing for PickScore, FlowAWR elevates the metric to 24.12 within 1.2k steps. In contrast, DiffusionNFT requires 2.0k steps to reach 23.82, and FlowGRPO necessitates over 4k steps while relying on CFG to achieve 23.50. Furthermore, under the complex distribution shifts induced by multi-reward optimization, our decoupled expert-branching strategy maintains optimization stability and mitigates out-of-domain aesthetic degradation. Our primary contributions are summarized as follows: • We introduce FlowAWR, an online flow RL framework that executes exact policy updates without SDE simulation. Our derivation validates the heuristic correction direction of DiffusionNFT while demonstrating that it constitutes a binary quantized special case. • We characterize the optimal velocity field by propagating the closed-form optimal policy to intermediate optimal distributions along the probability path, ultimately instantiating this theoretical optimum via a group-based approximated advantage function. • We conduct quantitative evaluations across both single- and multi-objective optimization settings, indicating that FlowAWR accelerates convergence rates and improves empirical alignment performance compared to baseline methodologies.
2
BACKGROUND
2.1
F LOW M ATCHING AND R ECTIFIED F LOW
Given a prior noise distribution p0 and a target data distribution p1 , Flow Matching (Lipman et al., 2022) constructs a continuous probability path pt (x) for t ∈ [0, 1] generated by a marginal vector field ut (x). Flow Matching constructs the target field via conditional probability paths pt (x|x1 ) and corresponding conditional vector fields ut (x|x1 ). Specifically, adopting the Rectified Flow (Liu et al., 2022) formulation, the forward process defines a linear interpolation between noise x0 ∼ p0 2
prompt c
FlowAWR Target Loss_AWR = Group-relative Advantage
Noise
...
...
0
Figure 1: Overview of FlowAWR Framework. The parameterized velocity vθ is trained to regress the optimal velocity field v ∗ . v ∗ is constructed by rectifying the reference field v old toward the target velocity ut using magnitude-aware, advantage-weighted residuals. and data x1 ∼ p1 : xt = tx1 + (1 − t)x0 , −xt which induces a constant conditional vector field, ut (xt |x1 ) = x1 − x0 = x11−t . The parameterized network vθ (xt , t) is then trained to regress this conditional vector field via the standard flow matching objective: (1) LFM (θ) = Et∼U [0,1],x1 ∼p1 ,x0 ∼p0 ∥vθ (xt , t) − ut (xt |x1 )∥2 . The global minimum of this objective yields the optimal marginal velocity field, which is explicitly formulated as the expectation of the conditional fields over the posterior data distribution (Lipman et al., 2022; Albergo et al., 2025): x1 − x t x̂1 (xt ) − xt v(xt , t) = Ex1 ∼p(x1 |xt ) xt = , (2) 1−t 1−t where x̂1 (xt ) denotes the posterior mean estimation of the data x1 given the intermediate state xt . 2.2
R EINFORCEMENT L EARNING AS S UPERVISED F INE - TUNING
A distinct paradigm seeks to bypass the complexities of policy gradients by treating RL as a supervised regression task. DiffusionNFT (Zheng et al., 2025) grounds this in the decomposition of the reference policy π old into positive (π + ) and negative (π − ) conditional distributions based on reward feedback. It theoretically derives a “Reinforcement Guidance” direction ∆(xt ), which represents the optimal shift from π old towards π + (and away from π − ). Instead of explicitly learning the guidance term, it employs an implicit strategy that constructs implicit positive (vθ+ ) and negative (vθ− ) policies by linearly mixing the parameterized field vθ with the reference field v old : vθ+ (xt , t) := (1 − β)v old (xt , t) + βvθ (xt , t), (3) vθ− (xt , t) := (1 + β)v old (xt , t) − βvθ (xt , t), where β is a fixed hyperparameter controlling the magnitude of policy deviation from the reference field. The optimization is then cast as a weighted regression loss that constrains these implicit policies to match the data vector field ut : LNFT (θ) = E r∥vθ+ (xt , t) − ut (xt |x1 )∥2 + (1 − r)∥vθ− (xt , t) − ut (xt |x1 )∥2 , (4) where r ∈ [0, 1] is the normalized scalar reward associated with the terminal samples x1 . This formulation embeds the reinforcement signal directly into the velocity field without requiring trajectory-level likelihood computation.
3
M ETHOD
3.1
R EVISITING O PTIMAL P OLICY FROM F IXED E XPERIENCE B UFFER
Standard policy optimization algorithms, such as PPO (Schulman et al., 2015; 2017) and GRPO (Shao et al., 2024), alternate between data sampling and policy optimization. In each iteration, samples are drawn from the current policy π old to construct an experience buffer D. While the 3
iterative update of π old enables online RL, the optimization of πθ within each individual step can be formulated as an offline reward maximization problem subject to a KL divergence constraint (Peters et al., 2010; Schulman et al., 2015; Levine, 2018; Peng et al., 2019): max Es∼D Ea∼πθ (·|s) R(s, a) − γDKL (πθ (·|s)||π old (·|s)) , (5) πθ
where s denotes the state, a represents the action, R(s, a) is the scalar reward, and γ > 0 determines the strength of the KL penalty bounding the deviation from the reference policy. Solving the Lagrangian of Eq. equation 5 with respect to the action probability normalization constraint yields the closed-form optimal policy π ∗ , which takes the form of an exponentially reweighted distribution (Peters et al., 2010; Rafailov et al., 2023): 1 old R(s, a) ∗ π (a|s) = π (a|s) exp , (6) Z(s) γ where Z(s) is the partition function. This exponentially re-weighted distribution mathematically indicates that the optimal strategy shifts probability mass toward high-reward actions while being regularized by the prior π old (Haarnoja et al., 2017; 2018). In the context of generative Flow Matching, the closed-form solution π ∗ (a|s) defines the theoretically optimal terminal distribution by mapping the MDP state s to the conditioning context c (e.g., input prompt) and the action a to the generated sample. Rather than approximating policy updates via SDE samplers, Flow Matching allows us to directly regress the optimal velocity field that induces this target distribution. 3.2
D ERIVATION OF O PTIMAL V ELOCITY F IELD
Building upon the optimal policy derived in Eq. equation 6, we first identify the optimal terminal distribution p∗1 (x1 |c), which represents the theoretical optimum for reward maximization on the current experience buffer, regardless of the exploration mechanism (e.g., SDE-based sampling with policy optimization): 1 old R(c, x1 ) ∗ p (x1 |c) = p (x1 |c) exp , Z(c) γ where c denotes an input prompt, and pold (x1 |c) represents the reference policy. To derive the corresponding velocity field, the explicit form of the intermediate marginal distribution p∗ (xt |c) is required: Proposition 1 (Intermediate Marginals). Given the optimal terminal distribution p∗1 , the corvia a value function responding marginal density at time t relates to the reference density pold t ϕ(xt , c, t): 1 old p∗t (xt |c) = p (xt |c) · ϕ(xt , c, t), (7) Z(c) t h i 1) where ϕ(xt , c, t) = Ex1 ∼pold (·|xt ,c) exp R(c,x . γ With the density relationship established, the explicit form of the optimal velocity field naturally emerges. Combining the velocity form in Eq. equation 2 with Tweedie’s formula (Robbins, 1992; Efron, 2011; Song et al., 2020), which links the posterior mean to the marginal score via x̂1 (xt , c) = (1−t)2 xt ∇xt log pt (xt |c), we observe that the velocity is strictly determined by the score function t + t of p∗t (xt |c) in Eq. equation 7: xt 1−t v ∗ (xt , c, t) = + ∇xt log p∗t (xt |c) t t 1−t xt + ∇xt log pold = t (xt |c) + log ϕ(xt , c, t) t t 1−t old = v (xt , c, t) + ∇xt log ϕ(xt , c, t). (8) t While Eq. equation 8 provides a theoretically rigorous direction for rectification, directly computing ∇xt log ϕ(xt , c, t) is computationally intractable. Herein, we expand the gradient into an expectation over samples, which reveals a connection to advantage-weighted regression (Peng et al., 2019; Zhang et al., 2025). 4
Theorem 1 (Advantage-Weighted Rectification). The optimal velocity field v ∗ (xt , c) derived in Eq. equation 8 can be equivalently expressed as the reference field rectified by an advantageweighted residual expectation: v ∗ (xt , c, t) = v old (xt , c, t) + Ex1 ∼pold (·|xt ,c) A(x1 , xt ) ut (xt |x1 ) − v old (xt , c, t) , where A(x1 , xt ) is the Centered Advantage, constructed by leveraging the zero-expectation property E[ut − v old ] = 0 to subtract a constant baseline of 1: 1) exp R(c,x γ A(x1 , xt ) = − 1. (9) ϕ(xt , c, t) This formulation reveals that the optimal field v ∗ is obtained by dynamically rectifying the reference field v old with the velocity residual (ut − v old ) based on A. This exhibits an explicit awareness of sample quality, enabling the formulation to quantitatively differentiate between positive and negative samples through the polarity of A, thereby reinforcing or repelling trajectories accordingly. 3.3
F LOW R EINFORCEMENT VIA A DVANTAGE -W EIGHTED R ECTIFICATION
Based on Theorem 1, the policy optimization process can be cast as a supervised regression task towards the optimal velocity field. Ideally, the formulation is the minimization of the distance between vθ and the theoretical optimum v ∗ , as defined in the standard objective in Eq. equation 1: h i 2 ∗ LIdeal (θ) = Et,xt ∼pold ∥v (x , c, t) − v (x , c, t)∥ . θ t t t However, directly optimizing LIdeal is computationally intractable due to the expectation over pold (x1 |xt , c) embedded within v ∗ (xt , c, t). To circumvent this, we propose the AdvantageWeighted Rectification (AWR) Loss, which learns from individual conditional samples by constructing a stochastic rectified target: LAWR (θ) = Et∼U [0,1],x1 ∼pold (·|c),x0 ∼p0 (·) (10) 2 vθ (xt , c, t) − v old (xt , c, t) + A(x1 , xt ) · ut (xt |x1 ) − v old (xt , c, t) . Theorem 2 (Optimization Equivalence). The Advantage-Weighted Rectification Loss LAWR (θ) and the Ideal Loss LIdeal (θ) share the same global optimum. Specifically, the gradients of both objectives with respect to θ are identical: ∇θ LAWR (θ) = ∇θ LIdeal (θ). which indicates that minimizing LAWR drives the parameterized velocity field vθ to converge to the theoretical optimal field v ∗ (xt , c, t). Theorem 2 establishes a novel perspective for Flow RL: directly learning the optimal terminal distribution derived from each experience buffer. This results in a new off-policy paradigm that casts policy optimization as supervised regression on advantage-weighted samples (Figure 1). Below, we introduce several practical techniques to translate this theoretical framework into an efficient online RL algorithm. Group-Based Estimation. Following GRPO-based algorithms (Shao et al., 2024; Liu et al., 2025; Xue et al., 2025; Zheng et al., 2025), we approximate the local value function ϕ(xt , c, t) using global group statistics to maintain computational efficiency. Specifically, for each prompt c, we generate (i) a group of G images {x1 }G i=1 . The advantage for the i-th sample in Eq. equation 9 is empirically computed as: (i) R(c,x1 ) exp γ (i) − 1. A(x1 ) ≈ (11) (j) PG R(c,x1 ) 1 exp j=1 G γ 5
Table 1: Our evaluation results compared with DiffusionNFT. Gray-colored : In-domain target task. ∗ OCR Performance under Standard Setting (β = 1, default EMA) Task Metric
DrawBench Metric
Method
#Iter
SD3.5-M + CFG
—
0.63
0.59
21.75
0.285
0.279
5.36
0.85
22.34
DiffusionNFT Ours
1k 0.5k
0.98 0.98
0.35 0.37
20.19 20.25
0.274 0.274
0.240 0.252
4.99 5.09
0.40 0.43
21.78 21.92
DiffusionNFT Ours
0.52k 0.26k
0.24 0.29
0.89∗ 0.97
19.38 19.43
0.240 0.246
0.183 0.193
4.90 5.00
-0.92 -0.78
20.27 20.51
DiffusionNFT Ours
2.0k 1.2k
0.55 0.61
0.67 0.72
23.82 24.12
0.272 0.275
0.316 0.315
6.19 6.28
1.31 1.36
24.09 24.29
GenEval OCR PickScore ClipScore HPSv2.1 Aesthetic ImgRwd PickScore
And the practical training objective for a prompt c can be expressed as: G 1 X (c) LAWR (θ) = Et∼U [0,1],x0 ∼p0 (·) G i=1 2 (i) (i) (i) (i) old (i) old (i) vθ (xt , t) − v (xt , t) + A(x1 ) · ut (xt |x1 ) − v (xt , t) . Adaptive γ Scaling. Instead of a fixed γ in the advantage term, we adopt the standard deviation (std) of rewards within the current experience buffer, which aligns with practices in GRPObased algorithms. Theoretically, this implies an adaptive adjustment of the KL constraint strength in Eq. equation 5. Soft Online Evolution. To ensure the robustness of FlowAWR, we employ an exponential moving average (EMA) update for the reference policy at the i-th iteration: θold ← µi θold +(1−µi )θ, which provides a steady evolving baseline for advantage estimation, mitigating the variance of immediate updates while allowing the reference policy to progressively shift towards high-reward regions (Ho et al., 2020). CFG-Free Optimization. While CFG (Ho & Salimans, 2022) is standard for inference, it becomes redundant under our AWR objective, which implicitly distills the guidance signal by rectifying the velocity field towards high-reward trajectories. We observe that the policy rapidly learns to generate high-fidelity samples without external guidance, aligning with recent findings (Zheng et al., 2025) and significantly reducing inference cost. 3.4
U NIFIED P ERSPECTIVE ON F LOW R EINFORCEMENT
Our FlowAWR framework not only provides a foundation for supervised flow policy optimization but also establishes a unifying lens through which both standard supervised fine-tuning and recent heuristic methods can be understood. As derived in Theorem 1 and 2, our framework explicitly defines the optimal velocity field, which takes the following explicit form: v ∗ (xt , c, t) = v old (xt , c, t) + A(x1 , xt ) · ut (xt |x1 ) − v old (xt , c, t) . (12) By modulating the advantage distribution A(x1 , xt ), we can formally deduce both standard supervised fine-tuning and discrete heuristic methods. Unification with Supervised Fine-Tuning. Under the condition where no reward prioritization is applied, all samples generated by the reference policy are treated as equally optimal. In our framework, this corresponds to setting a constant advantage A(x1 , xt ) ≡ 1, which arises naturally when we omit the baseline subtraction. Substituting this into Eq. equation 12, the rectification residual perfectly cancels the reference field: v ∗ (xt , c, t) = v old (xt , c, t) + 1 · ut (xt |x1 ) − v old (xt , c, t) = ut (xt |x1 ). Consequently, the objective strictly degenerates into regressing the raw conditional vector field, which demonstrates that SFT is a trivial case of FlowAWR under a uniform advantage distribution. 6
24.5
DiffusionNFT
22.5
FlowAWR (Ours) 0
500
1000
1500
Training Iterations
2000
(a)
0.370
3.7× Faster
0.360 0.350 0.212
DiffusionNFT
FlowAWR (Ours) 0
100
200
300
400
Training Iterations (b)
500
OCR Score
23.0
HPSv2.1 Score
PickScore
4.2× Faster
23.5
19.8
1.0
0.380
24.0
0.8
5.2× Faster
0.6 0.4 DiffusionNFT
0.2
FlowAWR (Ours) 0
100
200
300
400
Training Iterations
500
(c)
Figure 2: Head-to-head comparison between FlowAWR and DiffusionNFT on single-reward optimization tasks: (a) PickScore, (b) HPSv2.1, and (c) OCR. Unification with DiffusionNFT. Furthermore, we identify that DiffusionNFT (Sec. 2.2) can be interpreted as a Binary Quantized Special Case of our approach. By factoring out the coefficient β from the implicit policies in Eq. equation 3, equation 4, the optimization can be equivalently rewritten as: 2 1 LNFT ∝ E r vθ − v old + (ut − v old ) β 2 1 old old . +(1 − r) vθ − v − (ut − v ) β where the temporal and conditional variables are omitted for brevity. This factorization reveals that DiffusionNFT analytically constrains the model to approximate our theoretical v ∗ under fixed, hard-clipped advantages: ANFT = ±1/β. Although DiffusionNFT incorporates group statistics to normalize the raw reward into a scalar r ∈ [0, 1], it remains restricted to “single-sample pushing/pulling” between positive and negative directions, which does not effectively utilize intra-group relative information for rectification. This limitation becomes particularly evident under the extreme condition of r = 0.5, where both directions are redundantly updated with identical weights. Advantages of FlowAWR. By generalizing DiffusionNFT into a probabilistic framework, FlowAWR retains its structural benefits: Forward Consistency (optimizing forward ODEs), Solver Flexibility (decoupling sampling from training), and a Likelihood-free formulation (bypassing SDE approximations). Furthermore, by deriving the advantage function A in Theorem 1, FlowAWR bypasses the manual tuning of the mixing hyperparameter β. Building upon this grounded formulation, our method effectively implements Group-Awareness. While DiffusionNFT reduces group statistics to a scalar weight for isolated trajectory updates, FlowAWR leverages the Group-Based Estimation of the partition function ϕ to establish a dynamic local baseline. This relative judgment, rooted in the group P 1 normalization term G exp(R/γ), provides a more robust and statistically grounded direction for velocity rectification compared to isolated absolute feedback.
4
E XPERIMENTS
We evaluate the proposed FlowAWR framework across three dimensions: (1) head-to-head comparative analysis with DiffusionNFT on single-reward optimization, (2) multi-reward joint training performance, and (3) ablation studies isolating key algorithmic components. 4.1
E XPERIMENTAL S ETUP
To ensure a fair comparison, our experiments are based on SD3.5-Medium (Esser et al., 2024) at 512 × 512 resolution, with most experimental settings strictly aligned with FlowGRPO (Liu et al., 2025) and DiffusionNFT (Zheng et al., 2025). 7
Tasks and Evaluation Metrics. We categorize the text-to-image optimization tasks into two paradigms: (1) Rule-based constraints, utilizing GenEval (Ghosh et al., 2023) for compositional spatial reasoning and OCR accuracy for explicit text rendering. (2) Model-based human preference, employing PickScore (Kirstain et al., 2023) as the primary reward signal. To assess out-of-domain generalization, models are additionally evaluated on the DrawBench (Saharia et al., 2022) benchmark using an array of evaluation metrics: PickScore, CLIPScore (Hessel et al., 2021), HPSv2.1 (Wu et al., 2023), Aesthetic Score (Schuhmann, 2022), and ImageReward (Xu et al., 2023). Training and Evaluation. We finetune with LoRA (α = 64, r = 32). During training, each epoch processes 48 prompts, with an intra-group sample size of G = 24. For trajectory rollouts, we use T = 10 sampling steps for both head-to-head baseline comparisons, ablation studies and multireward joint training. All final evaluations are performed using a 40-step first-order ODE solver. Further implementation details and hyperparameter settings are provided in the supplementary material. 4.2
H EAD - TO -H EAD C OMPARISON ON S INGLE R EWARDS
We first compare FlowAWR with DiffusionNFT under the single-reward setting. As shown in Table 1 and Figure 2, FlowAWR achieves a 2× to 5× acceleration in training iterations. For the OCR task, the Figure 2c illustrates the baseline performance of DiffusionNFT without its extreme setting of β = 0.1 (which corresponds to an aggressive hard-clipped advantage of ANFT = ±10), further indicating that its efficacy depends heavily on heuristic tuning. As evidenced in Table 1, FlowAWR exhibits reduced reward hacking during out-of-domain evaluations. However, exclusively optimizing a single rule-based objective (e.g., GenEval or OCR) typically degrades overall image aesthetics due to the absence of holistic visual constraints. To address this, we explore multi-reward sequential training in the subsequent section. 4.3
M ULTI -R EWARD S EQUENTIAL T RAINING
To mitigate the aesthetic degradation associated with isolated rule-based objectives, we extend FlowAWR to a multi-reward setting via a sequential optimization strategy. Using the CFG-free SD3.5-Medium (2.5B parameters) as the base model, we first optimize the policy using a composite reward signal comprising PickScore, CLIPScore, and HPSv2.1 on the Pick-a-Pic dataset (Kirstain et al., 2023). Building upon this aesthetically aligned prior, we subsequently branch the training process to derive task-specific expert policies for GenEval and OCR. This decoupled approach enables improved adherence to structural rules while preserving the global visual fidelity established in the first stage as exhibited in Table 2 and Figure 3. 4.4
A BLATION S TUDY
We conduct comprehensive ablation studies to analyze our key design choices: Group Size. The group size G dictates the estimation quality of the partition function ϕ and the optimization boundaries. As derived from Eq. equation 11, the advantage can be equivalently ex(i) (i) pressed as A(x1 ) = G·Softmax(R(c, x1 )/γ)−1. A smaller G restricts the exploration magnitude and increases the estimation bias of ϕ. Due to out-of-memory constraints at G = 48, we evaluate G ∈ {8, 24, 32} and adopt G = 24, which strikes an effective balance between unbiasedness and computational cost, ensuring controlled comparisons with baselines (Figure 4). Adaptive γ Scaling. We dynamically scale the γ in Eq. 9 using the std of the current global reward. Compared to static assignments, Figure 5 demonstrates that this adaptive scaling effectively maintains stable optimization dynamics, mitigating the need for dataset-specific hyperparameter. Adaptive Advantage. To validate the theoretical continuous formulation derived in Theorem 1, we ablate our adaptive advantage A against fixed magnitude assignment (A ∈ {±1, ±10}), which correspond to β = 1 and 0.1 in Eq. equation 3. As shown in Figure 6, an excessively large advantage (e.g., A = ±10) leads to training collapse. In contrast, our rectification maintains stability and yields higher generation performance. 8
Table 2: Evaluations for the multi-reward setting. Gray-colored : In-domain reward. ‡ Evaluated under 1024 × 1024 resolution. Bold: best result. Underline: second best. Model
Task Metric
#Iter
DrawBench Metric
GenEval OCR PickScore ClipScore HPSv2.1 Aesthetic ImgRwd PickScore SD-XL‡ SD3.5-L‡ FLUX.1-Dev
— — —
0.55 0.71 0.66
0.14 0.68 0.59
21.74 22.24 22.55
0.287 0.289 0.295
0.280 0.288 0.274
5.60 5.50 5.71
0.76 0.96 0.96
22.42 22.91 22.84
SD3.5-M (w/o CFG) — + CFG — + FlowGRPO >5k 2k 4k
0.24 0.63 0.95 0.66 0.54
0.12 0.59 0.66 0.92 0.68
19.90 21.75 21.67 21.80 23.33
0.237 0.285 0.293 0.290 0.280
0.204 0.279 0.274 0.280 0.316
5.13 5.36 5.32 5.32 5.90
-0.58 0.85 1.06 0.95 1.29
20.51 22.34 22.37 22.44 23.50
+ DiffusionNFT + Ours
0.93 0.93
0.90 0.91
23.36 23.60
0.293 0.296
0.330 0.341
5.99 6.04
1.48 1.51
23.79 23.91
1.7k 1.5k
A photo of a purple scissors
A high-altitude mountain summit with a wooden signpost "Elevation 8000 Feet"
Girl on the beach, black and white, minimalist, only lines
A person planting a tree with cat.
Two cats and one dog sitting on the grass.
SD3.5-M
FlowGRPO
DiffusionNFT
FlowAWR
Figure 3: Qualitative Comparison. The text prompts are sampled from GenEval, OCR, PickScore, and DrawBench evaluations, illustrating the generation fidelity of our FlowAWR.
Soft Online Update. The EMA decay rate µi governs the evolution of the reference policy. Fully on-policy (µi = 0) accelerates early progress but destabilizes training, whereas an overly off-policy setting (large constant µ) slows convergence. As shown in Figure 7, starting with a small µ and gradually increasing it strikes an effective balance between convergence speed and training stability.
5
R ELATED W ORK
The broad application of RL in aligning LLMs Ouyang et al. (2022); Rafailov et al. (2023) has motivated efforts to align continuous generative models. Existing literature predominantly bifurcates into two paradigms. 9
1.0
0.8
23
0.8
0.6 0.4
G = 32 G = 24 G=8
0.2 0
100
22 21
G = 32 G = 24 G=8
20
200
0
Training Iterations
100
(a)
200
300
400
Training Iterations
0.6 0.4
0
100
150
200
Training Iterations
250
Figure 5: Effect of adaptive or fixed γ.
23
0.8
0.6 0.4
Adaptive = ±1 = ±10
0.2 200
Training Iterations (a)
300
22 21
Adaptive = ±1 = ±10
20 0
100
200
300
400
Training Iterations
500
OCR Score
1.0
0.8
PickScore
24
100
50
(b)
1.0
0
Adaptive = 0.001 = 0.01
0.2
500
Figure 4: Ablation on group size G.
OCR Score
OCR Score
24
PickScore
OCR Score
1.0
0.6 0.4
= min(0.001i, 0.5) = min(0.01i, 0.9) =0
0.2 0
50
100
150
200
Training Iterations
250
(b)
Figure 6: Ablation on adaptive advantage.
Figure 7: Influence of the soft online update strategy.
Implicit Optimization and Guidance-based Methods. This paradigm circumvents exact likelihoods via implicit guidance, reward-weighted fine-tuning, or score modifications. Early methods like Reward-Weighted Regression (RWR) (Lee et al., 2023; Dong et al., 2023) lack strict penalization for low-quality samples, while direct Reward Backpropagation (Xu et al., 2023; Prabhudesai et al., 2023; Clark et al., 2023) suffers from substantial memory overhead and gradient explosion through the ODE/SDE solvers. Inference-time guidance techniques (Janner et al., 2022; Jin et al., 2025) effectively steer generation but multiply inference latency. More recently, DiffusionNFT (Zheng et al., 2025) optimizes the vector field directly via a deterministic, likelihood-free framework. Nevertheless, as analyzed in Sec. 3.4, its reliance on a fixed binary scaling factor restricts its capacity to effectively exploit intra-group relative rewards. Density-Approximated Policy Gradient Methods. To leverage standard RL algorithms, this paradigm explicitly approximates the intractable continuous density by formulating the sampling process as a MDP (Black et al., 2023; Fan et al., 2023). Offline approaches like Diffusion-DPO (Wallace et al., 2024; Yang et al., 2024) require additional likelihood approximations that diverge from the AR formulation. For online RL, methods such as DDPO (Black et al., 2023) and recent Flow Matching extensions (Liu et al., 2025; Xue et al., 2025) decompose trajectories to approximate stepwise likelihoods. However, to extract tractable variance for density estimation, they necessitate the injection of noise via SDE samplers during training, which introduces computational bottlenecks and creates a structural forward inconsistency between stochastic training and deterministic ODE inference. While MixGRPO (Li et al., 2025) attempts to improve efficiency by mixing SDE and ODE steps, the underlying density approximation and issues remain unresolved.
6
C ONCLUSION
We introduce Flow Advantage-Weighted Rectification (FlowAWR), an off-policy online reinforcement learning paradigm for continuous generative models. By deriving the optimal velocity field from a KL-constrained policy formulation, FlowAWR provides a unified analytical perspective that connects flow-based RL with supervised fine-tuning. This formulation also clarifies the relationship between prior methods such as DiffusionNFT and policy optimization, revealing them as discrete approximations within a continuous framework. Empirically, FlowAWR enables stable CFG-free optimization and achieves 2× to 5× faster convergence than DiffusionNFT in single-reward tasks, while remaining robust under the distribution shifts induced by multi-reward alignment.
10
R EFERENCES Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023. Michael Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. Journal of Machine Learning Research, 26(209):1–80, 2025. Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305.13301, 2023. Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. Advances in neural information processing systems, 30, 2017. Kevin Clark, Paul Vicol, Kevin Swersky, and David J Fleet. Directly fine-tuning diffusion models on differentiable rewards. arXiv preprint arXiv:2309.17400, 2023. Gheorghe Comanici, Eric Bieber, Mike Schaekermann, Ice Pasupat, Noveen Sachdeva, Inderjit Dhillon, Marcel Blistein, Ori Ram, Dan Zhang, Evan Rosen, et al. Gemini 2.5: Pushing the frontier with advanced reasoning, multimodality, long context, and next generation agentic capabilities. arXiv preprint arXiv:2507.06261, 2025. Hanze Dong, Wei Xiong, Deepanshu Goyal, Yihan Zhang, Winnie Chow, Rui Pan, Shizhe Diao, Jipeng Zhang, Kashun Shum, and Tong Zhang. Raft: Reward ranked finetuning for generative foundation model alignment. arXiv preprint arXiv:2304.06767, 2023. Bradley Efron. Tweedie’s formula and selection bias. Journal of the American Statistical Association, 106(496):1602–1614, 2011. 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. 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. 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. Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In International conference on machine learning, pp. 1352–1361. PMLR, 2017. Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and applications. arXiv preprint arXiv:1812.05905, 2018. Jack Hessel, Ari Holtzman, Maxwell Forbes, Ronan Le Bras, and Yejin Choi. Clipscore: A reference-free evaluation metric for image captioning. In Proceedings of the 2021 conference on empirical methods in natural language processing, pp. 7514–7528, 2021. Jonathan Ho and Tim Salimans. arXiv:2207.12598, 2022.
Classifier-free diffusion guidance.
arXiv preprint
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840–6851, 2020. 11
Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.09991, 2022. Luozhijie Jin, Zijie Qiu, Jie Liu, Zijie Diao, Lifeng Qiao, Ning Ding, Alex Lamb, and Xipeng Qiu. Inference-time alignment control for diffusion models with reinforcement learning guidance. arXiv preprint arXiv:2508.21016, 2025. Yuval Kirstain, Adam Polyak, Uriel Singer, Shahbuland Matiana, Joe Penna, and Omer Levy. Picka-pic: An open dataset of user preferences for text-to-image generation. Advances in neural information processing systems, 36:36652–36663, 2023. Black Forest Labs, Stephen Batifol, Andreas Blattmann, Frederic Boesel, Saksham Consul, Cyril Diagne, Tim Dockhorn, Jack English, Zion English, Patrick Esser, et al. Flux.1 kontext: Flow matching for in-context image generation and editing in latent space, 2024. 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. Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv:1805.00909, 2018. Junzhe Li, Yutao Cui, Tao Huang, Yinping Ma, Chun Fan, Miles Yang, and Zhao Zhong. Mixgrpo: Unlocking flow-based grpo efficiency with mixed ode-sde. arXiv preprint arXiv:2507.21802, 2025. Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022. 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. 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. Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35: 27730–27744, 2022. Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-weighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv:1910.00177, 2019. Jan Peters, Katharina Mulling, and Yasemin Altun. Relative entropy policy search. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 24, pp. 1607–1612, 2010. Mihir Prabhudesai, Anirudh Goyal, Deepak Pathak, and Katerina Fragkiadaki. Aligning text-toimage diffusion models with reward backpropagation. arXiv preprint arXiv:2310.03739, 2023. Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. Advances in neural information processing systems, 36:53728–53741, 2023. Herbert E Robbins. An empirical bayes approach to statistics. In Breakthroughs in Statistics: Foundations and basic theory, pp. 388–394. Springer, 1992. Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in neural information processing systems, 35:36479–36494, 2022. Christoph Schuhmann. laion-aesthetics/, 2022.
Laion-aesthetics.
12
https://laion.ai/blog/
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, pp. 1889–1897. PMLR, 2015. John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017. Team Seedream, Yunpeng Chen, Yu Gao, Lixue Gong, Meng Guo, Qiushan Guo, Zhiyao Guo, Xiaoxia Hou, Weilin Huang, Yixuan Huang, et al. Seedream 4.0: Toward next-generation multimodal image generation. arXiv preprint arXiv:2509.20427, 2025. 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. 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. Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01469, 2023. 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, pp. 8228–8238, 2024. 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. Xiaoshi Wu, Yiming Hao, Keqiang Sun, Yixiong Chen, Feng Zhu, Rui Zhao, and Hongsheng Li. Human preference score v2: A solid benchmark for evaluating human preferences of text-toimage synthesis. arXiv preprint arXiv:2306.09341, 2023. 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. 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. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025. 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, pp. 8941–8951, 2024. Shiyuan Zhang, Weitong Zhang, and Quanquan Gu. Energy-weighted flow matching for offline reinforcement learning. arXiv preprint arXiv:2503.04975, 2025. 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. Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. arXiv preprint arXiv:1909.08593, 2019.
13
A
E XTENDED T HEORETICAL D ERIVATIONS AND P ROOFS
A.1
D ERIVATION OF THE O PTIMAL P OLICY
This section provides the analytical derivation of the closed-form optimal policy π ∗ , which is strictly optimized over the current fixed experience buffer D subject to a KL divergence constraint against the reference policy. While the iterative update of π old enables online RL, the step-wise optimization of πθ fundamentally solves a constrained reward maximization problem over the empirical states in D: J(πθ ) = Es∼D Ea∼πθ (·|s) [R(s, a)] − γDKL (πθ (·|s)||π old (·|s)) , where maximizing the global objective over D is mathematically equivalent to solving the optimization for each state s ∈ D independently. Expanding the terms for a specific state s yields: Z Z πθ (a|s) J(πθ (·|s)) = πθ (a|s)R(s, a)da − γ πθ (a|s) log old da. π (a|s) To ensure πθ (·|s) is a valid probability density, a state-dependent Lagrangian multiplier λ(s) is R introduced for the normalization constraint πθ (a|s)da = 1, yielding the Lagrangian functional: Z Z πθ (a|s) L(πθ , λ) = πθ (a|s)R(s, a)da − γ πθ (a|s) log old da π (a|s) Z + λ(s) 1 − πθ (a|s)da . Taking the functional derivative of L with respect to the density πθ (a|s) and setting it to zero yields the first-order optimality condition: δL πθ (a|s) = R(s, a) − γ log old + 1 − λ(s) = 0. δπθ (a|s) π (a|s) Rearranging the terms to isolate πθ (a|s): log
πθ (a|s) R(s, a) λ(s) = −1− , π old (a|s) γ γ R(s, a) λ(s) πθ (a|s) = π old (a|s) exp exp −1 − . γ γ
is independent of a, it acts as a normalization constant. Defining Since the term exp −1 − λ(s) γ its inverse as the partition function Z(s): Z λ(s) R(s, a) Z(s) = exp 1 + = π old (a|s) exp da. γ γ Substituting Z(s) back yields the theoretical optimal policy π ∗ corresponding to Eq. equation 6: 1 old R(s, a) ∗ . (13) π (a|s) = π (a|s) exp Z(s) γ A.2
P ROOF OF P ROPOSITION 1
Proposition 1. The optimal marginal density at time t relates to the reference density pold t via a time-dependent value function ϕ(xt , c, t): 1 old p (xt |c) · ϕ(xt , c, t), Z(c) t h i 1) where ϕ(xt , c, t) = Ex1 ∼pold (·|xt ,c) exp R(c,x . γ p∗t (xt |c) =
14
Proof. The marginal density p∗t (xt |c) is obtained by marginalizing the joint distribution over the terminal state x1 : Z p∗t (xt |c) =
p(xt |x1 )p∗1 (x1 |c)dx1 .
The forward transition kernel p(xt |x1 ) is determined by the ReFlow schedule (e.g., xt = tx1 + (1 − t)x0 ) and remains strictly independent of the policy. Substituting the optimal terminal distribution
1 p∗1 (x1 |c) = Z(c) pold 1 (x1 |c) exp
p∗t (xt |c) =
R(c,x1 ) γ
1 Z(c)
Z
:
p(xt |x1 )pold 1 (x1 |c) exp
R(c, x1 ) γ
dx1 .
Applying Bayes’ theorem to the reference policy terms, the joint probability can be factorized as old old p(xt |x1 )pold 1 (x1 |c) = p (x1 |xt , c)pt (xt |c). Substituting this factorization back into the integral formulation: Z 1 R(c, x1 ) old p∗t (xt |c) = pold dx1 . (x |c)p (x |x , c) exp t 1 t t Z(c) γ Since the intermediate reference marginal pold t (xt |c) is independent of the integration variable x1 , it can be factored out of the integral: Z 1 old R(c, x1 ) p∗t (xt |c) = pt (xt |c) pold (x1 |xt , c) exp dx1 . Z(c) γ {z } | ϕ(xt ,c,t)
The integral term mathematically corresponds to the conditional expectation over the posterior distribution pold (x1 |xt , c), which explicitly defines the value function ϕ(xt , c, t). A.3
D ERIVATION OF T WEEDIE ’ S F ORMULA IN F LOW M ATCHING
To establish the mathematical foundation, this section provides the analytical derivation of Tweedie’s formula within the context of Flow Matching, which explicitly maps the posterior mean estimation to the marginal score function. Under the ReFlow formulation, the forward probability path is constructed via a linear interpolation between the noise x0 ∼ N (0, I) and the data x1 ∼ p1 : xt = tx1 + (1 − t)x0 . Given the Gaussian conditional distribution p(xt |x1 ) = N (xt ; tx1 , (1 − t)2 I), the corresponding conditional score function evaluates to: xt − tx1 ∇xt log p(xt |x1 ) = − . (1 − t)2 The marginal score function ∇xt log pt (xt |c) can be expressed by marginalizing over the terminal data distribution p1 (x1 |c) and applying Bayes’ theorem: 1 ∇x pt (xt |c) pt (xt |c) t Z 1 = p1 (x1 |c)∇xt p(xt |x1 )dx1 pt (xt |c) Z p(xt |x1 )p1 (x1 |c) = ∇xt log p(xt |x1 )dx1 pt (xt |c) Z = p(x1 |xt , c)∇xt log p(xt |x1 )dx1 .
∇xt log pt (xt |c) =
15
Substituting the explicit form of the conditional score function into the integral gives: Z xt − tx1 ∇xt log pt (xt |c) = p(x1 |xt , c) − dx1 (1 − t)2 Z xt t =− + p(x1 |xt , c)x1 dx1 . (1 − t)2 (1 − t)2 R Substituting the posterior mean x̂1 (xt , c) = p(x1 |xt , c)x1 dx1 : ∇xt log pt (xt |c) = −
t xt + x̂1 (xt , c). 2 (1 − t) (1 − t)2
Rearranging the terms to isolate x̂1 (xt , c) yields Tweedie’s formula, connecting the posterior mean estimation directly to the marginal score: x̂1 (xt , c) = A.4
xt (1 − t)2 + ∇xt log pt (xt |c). t t
(14)
P ROOF OF T HEOREM 1
Theorem 1. The optimal velocity field v ∗ (xt , c, t) can be equivalently expressed as the reference field rectified by an advantage-weighted residual expectation: v ∗ (xt , c, t) = v old (xt , c, t) + Ex1 ∼pold (·|xt ,c) A(x1 , xt ) ut (xt |x1 ) − v old (xt , c, t) . where A(x1 , xt ) is the Centered Advantage, constructed by leveraging the zero-expectation property E[ut − v old ] = 0 to subtract a constant baseline of 1: 1) exp R(c,x γ − 1. A(x1 , xt ) = ϕ(xt , c, t) Proof. From Eq. equation 8, the optimal velocity field is strictly determined by the score function of the intermediate optimal marginal: 1−t v ∗ (xt , c, t) = v old (xt , c, t) + ∇xt log ϕ(xt , c, t). (15) t Expanding the gradient of the log value function and applying the logarithmic derivative identity ∇xt p = p∇xt log p, the gradient is derived as: ∇xt ϕ(xt , c, t) ϕ(xt , c, t) Z 1 R(c, x1 ) = ∇xt pold (x1 |xt , c) exp dx1 ϕ(xt , c, t) γ Z 1 R(c, x1 ) = pold (x1 |xt , c) exp ϕ(xt , c, t) γ
∇xt log ϕ(xt , c, t) =
· ∇xt log pold (x1 |xt , c)dx1 .
(16)
Relating the score of the posterior ∇xt log pold (x1 |xt , c) to the forward transition fields via Bayes’ theorem factorizes the gradient into a direct linear relationship with the velocity residual: ∇xt log pold (x1 |xt , c) = ∇xt log p(xt |x1 ) − ∇xt log pold t (xt |c) tx1 − xt tx̂old − xt − 1 2 (1 − t) (1 − t)2 t x1 − x̂old 1 = 1−t 1−t t = ut (xt |x1 ) − v old (xt , c, t) , 1−t
=
16
(17)
where the conditional and marginal scores are derived from Tweedie’s formula in Sec A.3, while x̂old −x −xt ut (xt |x1 ) = x11−t and v old (xt , c, t) = 11−t t denote the conditional flow target and reference velocity, respectively. Substituting Eq.equation 17 into Eq.equation 16 transforms the integral into an expectation over the posterior: " t old ∇xt log ϕ(xt , c, t) = E 1 − t x1 ∼p (·|xt ,c) # 1) exp R(c,x γ old ut (xt |x1 ) − v (xt , c, t) . ϕ(xt , c, t) Substituting this explicit gradient form back into the base velocity Eq.equation 15: v ∗ (xt , c, t) =v old (xt , c, t) 1) exp R(c,x γ old ut (xt |x1 ) − v (xt , c, t) . + Ex1 ∼pold (·|xt ,c) ϕ(xt , c, t) Since the velocity residual has zero expectation, Ex1 ∼pold (·|xt ,c) [ut (xt |x1 ) − v old (xt , c, t)] = 0, subtracting a unit baseline from the weight preserves the equality: v ∗ (xt , c, t) = v old (xt , c, t) 1) exp R(c,x γ + Ex1 ∼pold (·|xt ,c) − 1 ut (xt |x1 ) − v old (xt , c, t) . ϕ(xt , c, t) 1 )/γ) Defining the centered advantage as A(x1 , xt ) = exp(R(c,x − 1 explicitly yields the exact forϕ(xt ,c,t) mulation of the optimal velocity field, completing the proof.
A.5
P ROOF OF T HEOREM 2
Theorem 2. The gradients of the Advantage-Weighted Rectification Loss LAWR and the Ideal Loss LIdeal are identical with respect to θ: ∇θ LAWR (θ) = ∇θ LIdeal (θ). Proof. Define the stochastic target derived from Theorem 1 as Y (x1 , xt , c, t) = v old (xt , c, t) + A(x1 , xt ) ut (xt |x1 ) − v old (xt , c, t) . Theorem 1 establishes that the theoretical optimal velocity field is the conditional expectation of this stochastic target: v ∗ (xt , c, t) = Ex1 ∼pold (·|xt ,c) [Y (x1 , xt , c, t)]. The Advantage-Weighted Rectification loss is formulated as: LAWR (θ) = Et,c,xt Ex1 ∼pold (·|xt ,c) ∥vθ (xt , c, t) − Y (x1 , xt , c, t)∥2 . Expanding the squared L2 norm by introducing the optimal field v ∗ (xt , c, t) yields a standard biasvariance decomposition: ∥vθ − Y ∥2 = ∥(vθ − v ∗ ) + (v ∗ − Y )∥2 = ∥vθ − v ∗ ∥2 + ∥v ∗ − Y ∥2 + 2⟨vθ − v ∗ , v ∗ − Y ⟩. Substituting this decomposition into the expectation separates the loss function into three distinct terms: LAWR (θ) = Et,c,xt ∥vθ (xt , c, t) − v ∗ (xt , c, t)∥2 + Et,c,xt ,x1 ∥v ∗ (xt , c, t) − Y (x1 , xt , c, t)∥2 + Et,c,xt 2⟨vθ − v ∗ , Ex1 |xt ,c [v ∗ − Y ]⟩ . The three terms evaluate as follows: 17
• The first term is strictly the Ideal Loss LIdeal (θ), which measures the regression error between the parameterized velocity and the theoretical optimal field. • The second term represents the irreducible variance of the stochastic target Y . Because this term depends solely on the data and the reference policy, it is independent of the model parameters θ, yielding a gradient of zero. • In the third term, the inner expectation evaluates to Ex1 |xt ,c [v ∗ − Y ] = v ∗ − Ex1 |xt ,c [Y ]. By the definition of v ∗ , this equals v ∗ − v ∗ = 0, causing the entire cross-term to vanish. Consequently, the AWR loss simplifies to the Ideal Loss plus an independent constant term C: LAWR (θ) = LIdeal (θ) + C. Taking the gradient with respect to θ eliminates the constant C, proving the exact optimization equivalence: ∇θ LAWR (θ) = ∇θ LIdeal (θ).
B
T HEORETICAL D ISCUSSIONS WITH RWR AND DPO
As highlighted in the Introduction of the main text, both RWR and DPO represent distinct alignment paradigms that fundamentally originate from the same theoretical optimal policy. Despite their divergent practical implementations, both methods can be mathematically unified under the KLconstrained reward maximization framework, leveraging the closed-form optimal density derived in Eq.equation 13. B.1
C ONNECTION TO RWR IN F LOW M ATCHING
Under the ReFlow schedule xt = tx1 + (1 − t)x0 , aligning a parameterized velocity vθ (xt , c, t) with the ideal marginal velocity v ∗ (xt , c, t) = Ex1 ∼p∗ (·|xt ,c) [ut (xt |x1 )] yields the global matching objective: LFM (θ) = Et,xt ∼p∗t ∥vθ (xt , c, t) − v ∗ (xt , c, t)∥2 . To bypass the intractable marginal expectation in v ∗ , Conditional Flow Matching (CFM) establishes a gradient equivalence. By expanding the squared norm and marginalizing the cross-term Ext ∼p∗t [⟨vθ , v ∗ ⟩], the objective reduces to matching the conditional field: ZZ p∗1 (x1 |c)p(xt |x1 )⟨vθ , ut ⟩dx1 dxt = Ex1 ∼p∗1 Ext ∼p(·|x1 ) [⟨vθ , ut ⟩]. which recovers the conditional matching objective over the optimal target distribution: LCFM (θ) = Et,x1 ∼p∗1 ,x0 ∼π0 ∥vθ (xt , c, t) − (x1 − x0 )∥2 . Since sampling directly from p∗1 is practically infeasible, applying importance sampling from p∗ 1 (x1 |c) the accessible reference policy pold 1 re-weights the expectation by the density ratio pold (x1 |c) ∝ 1 1) exp R(c,x . Absorbing the partition function Z(c) establishes the final RWR objective: γ R(c, x1 ) 2 exp LRWR (θ) ∝ Ex1 ∼pold Et,x0 ∼π0 ∥vθ (xt , c, t) − (x1 − x0 )∥ . 1 γ which bridges the classical RWR alignment paradigm with continuous generative processes, demonstrating that exponentiated reward re-weighting naturally emerges from projecting the velocity onto the optimal policy trajectory. 18
B.2
C ONNECTION TO DPO
Unlike RWR’s regression approach, DPO bypasses explicit reward modeling. Rearranging the logarithmic form of the optimal policy analytically reparameterizes the unobserved reward R(s, a) strictly in terms of policy probabilities: R(s, a) = γ log
π ∗ (a|s) + γ log Z(s). π old (a|s)
Under the Bradley-Terry (BT) preference model, the probability that an action yw is preferred over yl is defined by the logistic function of their reward difference, P (yw ≻ yl |s) = σ (R(s, yw ) − R(s, yl )). Substituting the reparameterized rewards perfectly cancels the intractable partition function Z(s): R(s, yw ) − R(s, yl ) = γ log
π ∗ (yw |s) π ∗ (yl |s) − γ log old . old π (yw |s) π (yl |s)
Parameterizing π ∗ with a learnable policy πθ and minimizing the negative log-likelihood over a preference dataset D directly yields the exact DPO objective: πθ (yl |s) πθ (yw |s) − γ log old LDPO (θ) = −E(s,yw ,yl )∼D log σ γ log old . (18) π (yw |s) π (yl |s) which confirms that DPO structurally inherits the KL-constrained optimality criteria, simply reformulating it to accommodate pairwise preference data.
E XPERIMENT D ETAILS
C.1
E MPIRICAL R ANGE OF THE C ENTERED A DVANTAGE PickScore
GenEval
0.5
-0.60
0
200
400
600
800
Training Iterations
max min
0 1
-1.54
0
100
Advantage
Advantage
max min
0 1
-1.76
150
200
Training Iterations (d) OCR
max min
0.0 0.5
500
-0.72
0
250
max min
0.0 0.5
-0.67
400
300
Training Iterations
400
Multi-Reward-GenEval
0.5
200
200
1.0
1.06
0
100
(c) HPSv2.1
Multi-Reward-PickScore
1
100
400
1.0
1.91
50
300
0.5
(b) GenEval
OCR
0
200
Training Iterations
(a) PickScore 2
Advantage
max min
0.0
1.28
1.0
1
0.5
2
HPSv2.1
1.69
0.92
Advantage
Advantage
1.0
Advantage
C
600
Training Iterations
(e) Multi-Reward (PickScore)
800
1.05
0.5 max min
0.0 0.5
-0.74
0
100
200
300
Training Iterations
400
(f) Multi-Reward (GenEval)
Figure S1: Empirical variation of the centered advantage A(x1 , xt ) across six alignment tasks. Under both single-reward and multi-reward settings, A(x1 , xt ) remains structurally bounded. It provides discriminative learning signals while naturally restricting the upper and lower bound. To validate the numerical stability of the proposed framework, Figure S1 illustrates the empirical variation range of the centered advantage A(x1 , xt ) across six distinct alignment tasks, encompassing both single-reward (PickScore, GenEval, HPSv2.1, OCR) and multi-reward optimizations. While the advantage values exhibit the necessary fluctuations to provide discriminative learning signals, their overall dynamic range remains strictly bounded and controllable throughout the training 19
process. Specifically, the lower bound is naturally restricted near −1 for highly penalized samples, and the upper bound is effectively constrained without exhibiting the exponential explosion typical of standard reward-weighted regression. This structurally bounded behavior confirms that the statedependent value function normalization ϕ(xt , c, t) successfully regulates the scale of the learning signal, ensuring stable gradient updates across diverse objectives without requiring heuristic clipping operations. C.2
E XPERIMENTAL S ETUP
Training Configurations. The training setup follows FlowGRPO and DiffusionNFT, which adopts 48 groups per epoch, a group size of 24, and a LoRA parameterization (r = 32, α = 64) with a constant learning rate of 3 × 10−4 . For each collected clean image, forward noising and loss computation are performed on the corresponding sampling timesteps. Data collection employs a 2nd-order ODE sampler with adaptive time weighting by default. Single-Reward Optimization. For fair comparison with DiffusionNFT, we use the same 10 sampling steps and set our update rate to ηi = min(0.001i, 0.5) across the three tasks, while DiffusionNFT uses its reported configurations. By default, the advantage is simply scaled by the group size. However, for rule-based tasks (OCR, GenEval), this standard scaling struggles with highly unstable reward dynamics. To ensure stability, we replace it with an adaptive variance gating mechanism that normalizes the advantage by its batch-wise std. This dynamically damps excessive variance during early training, while a strict lower bound of 10−2 prevents numerical explosion in later stages. Multi-Reward Optimization. Following the multi-stage curriculum of DiffusionNFT, the training is structured around three dataset-specific configurations: (1) 3 rewards (PickScore, CLIPScore, and HPSv2.1) on the Pick-a-Pic dataset; (2) 4 rewards (the previous three plus GenEval) on the GenEval dataset; and (3) 4 rewards (the initial three plus OCR) on the OCR dataset. To overcome the low quality of initial CFG-free generations, the curriculum sequentially applies configuration (1) for 800 iterations, (2) for 200 iterations, (1) for 200 iterations, (2) for 200 iterations, and finally concludes with (3) for 100 iterations. Within each stage, all active rewards are equally weighted, with PickScore explicitly divided by 26 for [0, 1] normalization. For data collection, we maintain the sampling steps as 10 to ensure training efficiency. And we maintain the default update rate ηi = min(0.001i, 0.5) uniformly across all stages. Note that these iteration counts are approximate. To ensure continuous optimization during stage transitions, the complete training state is strictly resumed from the preceding stage, including the model parameters θ, the reference policy θold , the optimizer, and the gradient scaler.
D
A DDITIONAL Q UALITATIVE R ESULTS
We provide additional qualitative comparisons between the base model, FlowGRPO, and our proposed method in Figure S2, Figure S3, and Figure S4, corresponding to the GenEval, OCR, and DrawBench benchmarks, respectively.
20
SD3.5-M
FlowGRPO
DiffusionNFT
FlowAWR
A photo of a red car and an orange potted plant
A photo of a laptop below a sports ball
A photo of a scissors and a sandwich
A photo of three oranges
A photo of a white banana and a black elephant Figure S2: Qualitative comparison on the GenEval benchmark.
21
SD3.5-M
FlowGRPO
DiffusionNFT
FlowAWR
A vibrant urban wall features bold, dynamic graffiti spelling "Revolution Now" in striking colors, capturing the energy and urgency of a movement.
A vibrant hot air balloon ascends into a clear blue sky, trailing a banner that reads "Adventure Awaits" in bold, flowing letters.
A bustling train station with a vintage aesthetic, the platform speaker hanging from a metal pole, clearly announcing "Now Boarding Track 9" amidst the crowd of travelers …
A realistic movie set with a clapperboard slate clearly marked "Take 3", held by a production assistant in front of a vintage film camera, under the warm glow of studio lights.
A treehouse nestled among the branches, with a weathered sign nailed to the trunk, clearly displaying "No Adults Allowed" in bold, playful letters.
Figure S3: Qualitative comparison on OCR tasks.
22
SD3.5-M
FlowGRPO
DiffusionNFT
FlowAWR
A donut underneath a toilet.
Rainbow coloured penguin.
Two cars on the street.
A blue coloured pizza.
A cat on the right of a tennis racket. Figure S4: Qualitative comparison on the DrawBench benchmark.
23