PaperlayerView exact revision on arXiv

Source-linked AI summary

DiffusionART: Efficient, Stable and Scalable Reinforcement of Diffusion Models

Zhenglin Wan, Jiayu Chen, Zhengfeiyang Zhang, Pengfei Zhou, Wangbo Zhao, Fangyuan Kong, Pengfei Wan, Kun Gai, Yang You

arXiv:2610.12572v1cs.LG

TL;DR

Online diffusion RL must turn black-box rewards on completed outputs into effective updates across denoising states, while reverse- and forward-process methods face efficiency and stability trade-offs. DiffusionART aggregates prompt-level reward guidance for reverse-rollout states and adds drift-adaptive restoration. Across image and video generation it reports efficient, stable reward improvement, including substantially fewer GPU hours on HunyuanVideo-13B than DiffusionNFT and Flow-GRPO-Fast.

  • Problem

    Online diffusion RL faces a credit-assignment problem because black-box rewards arrive only after a complete output, while reverse-process methods are inefficient and forward-process methods require stabilization and task-tuned EMA.

    Source: §1 · p. 2

  • Method

    DiffusionART aggregates each prompt’s samples and rewards into a shared reward-improving direction for reverse-rollout states, then uses a trust-region-masked update for drift-adaptive restoration.

    Source: §1 · p. 2, §4 · p. 6, §4.1 · p. 6

  • Results

    On HunyuanVideo-13B with VQ:MQ = 1:1, DiffusionART reaches DiffusionNFT and Flow-GRPO-Fast late-stage mean rewards with 12.02× and 11.72× fewer GPU hours, respectively, while achieving 28.8% and 53.0% higher mean training rewards.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Experiments on image and video generation show efficient, stable reward improvement, including with pairwise feedback and small rollout groups.

    Source: §6 · p. 10

    The paper targets rectified-flow models with linear data–noise interpolation, which it says shares diffusion’s core principles.

    Source: §1 · p. 2

Abstract

from arXiv · show

Online reinforcement learning (RL) fine-tunes diffusion models with black-box rewards, but reverse-process methods such as Flow-GRPO are limited in efficiency by poor credit assignment, while forward-process methods such as DiffusionNFT are efficient yet unstable without carefully tuned exponential-moving-average (EMA) off-policy training. We show that NFT-style objectives contain a hidden forward--reverse consistency force that joint positive--negative training cancels only partially, and that EMA stabilizes training through restoration rather than off-policy rollouts. We introduce DiffusionART (Diffusion Reinforcement with Aggregation-First Reward Tilting), which aggregates each prompt's rewards into a shared reward-improving direction for reverse-process states, removing the consistency force by construction, and adds a trust-region-masked second update for drift-adaptive restoration without EMA tuning. It requires no differentiable reward model, remains effective with fewer samples per prompt, and supports pairwise reward models (DiffusionART-Pairwise). Across image and video generation, DiffusionART trains efficiently and stably without per-task tuning. In particular, on HunyuanVideo-13B with equally weighted visual and motion quality rewards, DiffusionART reaches the late-stage mean reward levels of DiffusionNFT and Flow-GRPO-Fast with 12.02 and 11.72 times fewer GPU hours, respectively, while achieving 28.8% and 53.0% higher mean training rewards.

1 Introduction

Online diffusion RL must convert black-box rewards on completed outputs into updates across denoising states, while existing reverse- and forward-process methods trade efficiency against stability. DiffusionART addresses forward-process instability through aggregation-first reward guidance and drift-adaptive restoration, and reports faster HunyuanVideo training than two baselines.

1 Introduction

  • Online diffusion RL uses black-box rewards on completed outputs, creating a credit-assignment problem across many denoising states; reverse-process methods have poor credit assignment, while forward-process methods need stabilizers.Reverse-process policy gradients weight each transition’s sampling noise by a trajectory-level advantage that the transition only partly determines; forward-process training is efficient but unstable without joint positive–negative training and task-tuned EMA.Source: §1 · p. 2
  • NFT-style objectives contain a hidden forward–reverse consistency force that joint positive–negative training cancels only partially, potentially destabilizing training and slowing reward optimization.The positive branch pulls each state’s velocity toward the ideal forward velocity, while the negative branch pushes it away; reward centering balances samples rather than each state.Source: §1 · p. 2
  • EMA stabilizes NFT through restoration from its lagged reference rather than off-policy rollouts, but its schedule must be tuned for each task.Joint positive–negative training only partially cancels the consistency force, leaving a residual that can destabilize training and slow reward optimization.Source: §1 · p. 2
  • DiffusionART aggregates each prompt’s samples and rewards into a shared reward-improving direction for reverse-rollout states, removing the consistency force by construction.A second trust-region-masked update adds restoration proportional to each state’s drift, replacing EMA without a schedule to tune.Source: §1 · p. 2
  • On HunyuanVideo-13B with VQ:MQ = 1:1, DiffusionART reaches DiffusionNFT and Flow-GRPO-Fast late-stage mean rewards with 12.02× and 11.72× fewer GPU hours, respectively, while achieving 28.8% and 53.0% higher mean training rewards.The respective reward improvements are over DiffusionNFT and Flow-GRPO-Fast.Source: §1 · p. 2
  • The paper targets rectified-flow models with linear data–noise interpolation, using “diffusion” as a general term.The paper says this setting shares the core principles of diffusion.Source: §1 · p. 2

2 Preliminaries

DiffusionNFT samples clean endpoints from a frozen rollout policy for each prompt, then uses reward-derived weights to define positive and negative endpoint distributions. Its forward-noising velocity regression moves toward the positive objective and away from the negative one, followed by an EMA rollout-policy update.

2 Preliminaries

  • For each prompt, DiffusionNFT samples K clean images or videos from a frozen rollout policy, calling the generated samples endpoints.Raw rewards are evaluated on these endpoints.Source: §2 · p. 3
  • DiffusionNFT converts each endpoint’s raw reward into a clipped goodness weight using the prompt’s mean reward and a positive scale parameter.Without clipping, the centered weights sum to zero across the prompt’s samples.Source: §2 · p. 3
  • Weighting the rollout distribution by each endpoint’s goodness and its complement defines positive and negative endpoint distributions whose ideal forward velocities are collinear.Moving toward the positive field and away from the negative field follows the same ideal improvement direction.Source: §2 · p. 3
  • DiffusionNFT regresses reward-weighted positive and negative velocity predictions toward forward-noising targets, then updates its rollout policy by exponential moving average.The velocity target is x0 − ϵ, and the frozen target retains an anchor to the rollout-policy velocity.Source: §2 · p. 3

3 Understanding NFT-style Diffusion Reinforcement Learning

NFT-style training is efficient, but its positive and negative objectives only partially cancel a hidden consistency force, while EMA stabilizes training through restoration that requires task-specific tuning. These findings motivate DiffusionART’s removal of that force and drift-adaptive restoration.

3.1 Joint Positive–Negative Training: Partial Cancellation of a Hidden Consistency Force

  • Each NFT branch combines reward guidance with an uncancelled consistency force, although their ideal reward-improving directions align.The force arises because the reverse sampling field can differ from the ideal forward velocity induced by the endpoints.Source: §3.1 · pp. 3–4
  • Joint training cancels the hidden consistency force only partially because reward imbalance varies across states, diluting reward-improving guidance.The residual grows as states reveal more endpoint quality and requires an additional stabilizer.Source: §3.1 · p. 4, Figure 2 · p. 4

3.2 Off-Policy EMA: Stability through Restoration, at the Cost of Task-Specific Tuning

  • EMA stabilizes NFT through restoration toward the lagged rollout reference; off-policy rollouts alone do not prevent collapse.At zero drift the restoration gradient vanishes, while EMA maintains drift that lets restoration counter the residual consistency force.Source: §3.2 · pp. 4–5
  • Two updates per rollout avoid collapse without EMA over the observed window, but slow training by 15.2×.Restoration’s output-force RMS is 9.4× smaller than guidance, yet its aggregated parameter-gradient norm is 4.9× larger.Source: §3.2 · p. 5, Figure 3 · p. 5
  • EMA controls restoration indirectly through policy lag, requiring task-specific tuning because too little risks collapse and too much slows reward improvement.The transferred GenEval-tuned EMA and β settings also destabilize FLUX training and can cause collapse.Source: §3.2 · p. 5

4 DiffusionART: Aggregation-first Reward Tilting with Adaptive Restoration

DiffusionART applies shared, reward-derived guidance to reverse-rollout states, then uses drift-adaptive restoration and trust-region masking to stabilize updates. Its pairwise variant substitutes preference labels for scalar rewards.

4 DiffusionART: Aggregation-first Reward Tilting with Adaptive Restoration

  • DiffusionART combines aggregation-first guidance on reverse-rollout states, a second MSE update for drift-adaptive restoration, and trust-region masking to limit excessive policy updates.These components address the two failure modes identified in Section 3.Source: §4 · p. 6

4.1 Aggregation-First Guidance

  • Each prompt’s endpoint rewards and samples define a shared reward-improving direction, normalized to guide every visited state without an NFT consistency term.The direction is estimated from the reward–endpoint covariance; RMS normalization fixes its magnitude across prompts and reward models.Source: §4.1 · p. 6

4.2 Two-Update Restoration: Drift-Adaptive Stability without EMA Tuning

  • A second MSE update combines the original guidance with restoration toward the rollout prediction, with restoration increasing as each state’s drift grows.This drift-adaptive restoration avoids tuning an EMA schedule while retaining guidance when drift is small.Source: §4.2 · pp. 6–7

4.3 Trust-Region Masking for ODE Rollouts

  • A trust-region mask uses a KL proxy based on reference Gaussian transitions to limit excessive outward drift while retaining inward corrections.Rollouts remain ODE trajectories, and the drift threshold is fixed at 10−5.Source: §4.3 · p. 7

4.4 Extension to Pairwise Reward Models

  • DiffusionART-Pairwise replaces scalar rewards with winner, loser, or tie labels for disjoint endpoint pairs while leaving the covariance and training pipeline unchanged.For each prompt, K endpoints are divided into K/2 disjoint pairs.Source: §4.4 · p. 7

5 Experiments

Experiments evaluate DiffusionART across text-to-image, text-to-video, and image-to-video generation. The reported results cover reward improvement, stability, restoration, and lower-sample pairwise training.

5.1 Text-to-Image Generation

  • On SD3.5-M, DiffusionART reaches Flow-GRPO’s step-500 reward with 10.0×, 5.2×, and 9.6× fewer steps on GenEval, OCR, and PickScore, respectively.It avoids the severe collapse of DiffusionNFT (on-policy) and has training efficiency comparable to DiffusionNFT (EMA), without per-task EMA tuning.Source: §5.1 · p. 8, Figure 5 · p. 8
  • On FLUX.1-dev GenEval, DiffusionART learns fastest, while DiffusionNFT fluctuates with EMA and β settings transferred from SD3.5-M rather than retuned.The comparison uses six-step rollouts and the model’s built-in guidance embedding at scale 3.5.Source: §5.1 · p. 8, Figure 6 · p. 8
  • The OCR ablation compares FLUX.1-dev GenEval training with guidance and restoration ablation, but the supplied results do not report the latter’s outcome.Figure 6 labels its panels as FLUX.1-dev on GenEval and guidance/restoration ablation on OCR.Source: Figure 6 · p. 8

5.2 Text-to-Video Generation

  • 31.4% higher evaluation VQ and 57.4% higher evaluation MQ than the strongest baseline; DiffusionART also reaches VQ 5.5 with 22.1× fewer GPU hours than DiffusionNFT (EMA).For VQ 5.5, DiffusionART uses 63.4 GPU hours versus 1,401.4 for DiffusionNFT (EMA), retaining its gains through 400 steps.Source: Figure 7 · p. 9, §5.2 · p. 9

5.3 Image-to-Video Generation

  • For I2V, DiffusionART improves both VQ and MQ while baselines improve slowly or plateau, reaching MQ 6 in 45.8 GPU hours versus 220.3 for Flow-GRPO-Fast.The comparison is reported for Figure 7(d–e); evaluation results are in Table 4.Source: Figure 7 · p. 9, §5.3 · p. 10

5.2 Text-to-Video Generation

  • With pairwise feedback, K = 2 reaches final VQ 12.235 versus 12.092 for K = 8 and uses 3.4× fewer GPU hours over 400 steps.K = 2 uses a quarter of the rollout videos and training samples, consuming 260.8 GPU hours versus 883.3 for K = 8.Source: §5.2 · p. 9, Table 1 · p. 10
  • With VQ:MQ weights of 1:1, trust-region masking yields steady improvement, whereas the run without the mask drops around step 200 before recovering.The ablation retains normalized covariance guidance, timestep weighting, and two disjoint MSE updates.Source: §5.2 · p. 9, Figure 8 · p. 10

6 Conclusion

The conclusion diagnoses an incompletely canceled forward–reverse self-consistency force in NFT-style training and presents DiffusionART as an aggregation-first alternative with drift-adaptive restoration. It reports efficient, stable reward improvement across image and video generation, including pairwise feedback and small rollout groups.

6 Conclusion

  • DiffusionART aggregates endpoint reward guidance before assigning it to reverse-rollout states, removing the unstable force by construction, then uses a trust-region second update for drift-adaptive restoration.The conclusion describes stronger restoration as improving stability at the cost of slower reward optimization.Source: §6 · p. 10
  • Experiments across image and video generation show efficient, stable reward improvement, including with pairwise feedback and small rollout groups.The conclusion frames these findings as insight into forward-process training and a path toward scalable diffusion post-training for visual generation.Source: §6 · p. 10

A.1 Scaling to FLUX.1-dev

A.2 T2V Evaluation

Under matched TA:VQ:MQ training weights of 7:2:1, DiffusionART improves evaluation VQ and MQ over the strongest baseline. Table 3 reports the T2V evaluation.

A.2 T2V Evaluation

  • The reported evaluation comparison covers both visual quality and motion quality against the strongest baseline.Source: §A.2 · p. 14
  • 31.4% higher evaluation VQ and 57.4% higher evaluation MQ than the strongest baseline.Source: Table 3 · p. 14, §A.2 · p. 14
  • Table 3 reports T2V evaluation under matched TA:VQ:MQ training weights of 7:2:1.Source: Table 3 · p. 14, §A.2 · p. 14

A.3 I2V Evaluation

The supplied passages identify an I2V evaluation table and a separate multi-reward T2I experiment. The T2I setup jointly optimizes four rewards on one SD3.5-M policy.

A.3 I2V Evaluation

  • Table 4 reports I2V evaluation with VQ:MQ weights of 1:1.Source: Table 4 · p. 14

A.4 Multi-Reward Optimization for T2I

  • The SD3.5-M multi-reward experiment jointly optimizes GenEval, PickScore, CLIPScore, and HPSv2.1 for 500 training steps.Source: §A.4 · p. 14
  • Figure 9 reports each reward separately to compare how methods balance compositional accuracy, image–text alignment, and human preference under the shared objective.Source: §A.4 · p. 14

A.5 Alternative Reward Weights for T2V

With matched VQ:MQ weights, DiffusionART reaches strong text-to-video reward levels more efficiently than DiffusionNFT and Flow-GRPO-Fast, and reward-value-only training also compares favorably with ReFL. Image-generation training remains effective without trust-region masking, although masking is more beneficial for video.

A.5 Alternative Reward Weights for T2V

  • For MQ of 3, DiffusionART needs 19 steps and 42.9 GPU hours, versus 33 steps and 165.6 GPU hours for NFT and 117 steps and 149.4 GPU hours for Fast.The reported step speedups over NFT/Fast are 1.7×/6.2×, and GPU-hour speedups are 3.9×/3.5×.Source: §A.5 · p. 14
  • With VQ:MQ = 1:1, DiffusionART reaches the late-stage mean reward levels of DiffusionNFT (EMA) and Flow-GRPO-Fast with 12.02× and 11.72× fewer GPU hours, respectively.Its final-50-step mean reward is 13.018 versus 10.105 and 8.510, improvements of 28.8% and 53.0%.Source: §A.5 · p. 14, Figure 10 · p. 15
  • DiffusionART improves both rewards with fewer optimizer updates than ReFL, using reward values alone while ReFL uses reward-model backpropagation.ReFL uses K = 1 versus DiffusionART’s K = 8 samples per update.Source: Figure 11 · p. 15, §A.5 · p. 15

A.6 Trust-Region Masking for Image Generation

  • Without the trust-region mask, image training remains effective and can improve slightly faster; FLUX.1-dev reaches a training GenEval score of 0.971 by step 180.The masking benefit is more pronounced for video generation.Source: Figure 12 · p. 16, §A.6 · p. 16

A.7 Trust-Region Masking Diagnostics

On HunyuanVideo-13B, removing the trust-region mask produces larger late-stage gradient and velocity-drift spikes. Although the mask acts on the second update, the reported diagnostics indicate that it also helps keep later first updates well behaved.

A.7 Trust-Region Masking Diagnostics

  • Across the shared first 300 steps, the maximum first-update gradient norm is 1.96 with masking versus 3.33 without it.The mask constrains only each rollout’s second update.Source: §A.7 · p. 16
  • Velocity drift before second updates also has much larger spikes without masking, while the masked run’s drift score stays at a healthy level.The authors interpret this as evidence that restricting second updates helps keep subsequent first updates well behaved.Source: §A.7 · p. 16, Figure 13 · p. 16
  • Without masking, HunyuanVideo-13B training shows larger late-stage gradient and velocity-drift spikes.The diagnostics track gradient norm before clipping and scheduler-reference-weighted velocity drift before each update.Source: Figure 13 · p. 16, §A.7 · p. 16

B Proofs

These proofs formalize the conditional reward-tilt direction underlying NFT, the on-policy descent signal and its decomposition, and the role of shared guidance and restoration. They also state the sampling and transition requirements used for reverse-process optimization and trust-region masking.

B.1 Conditional Covariance Identity for Ideal NFT

  • At a supported state, the positive and negative NFT branches point along the same conditional reward-tilt direction, the covariance of reward with sample velocity.The identity assumes the relevant moments and tilted distributions exist and 0 < q < 1.Source: §B.1 · p. 17

B.2.1 Equivalent MSE Objective and On-Policy Descent Direction

  • At on-policy initialization, NFT’s single-sample output-space descent direction is proportional to (2r − 1)(ut − vold).The proportionality factor is 2β.Source: §B.2.1 · pp. 17–18

B.2.2 Reward Guidance and Self-Consistency Force

  • Conditioned on a state, the NFT guidance term separates into twice the reward–velocity covariance and a term involving the conditional mean reward and mean velocity residual.The covariance equals g under the stated definitions.Source: §B.2.2 · p. 18

B.2.3 Guidance and Quadratic Restoration

  • The positive and negative residuals are βd − h and −βd − h; expanding their weighted squared norms yields the stated objective up to a parameter-independent constant.The proof identifies this expansion as Equation 4.Source: §B.2.3 · p. 18

B.3 Covariance as the First-Order Reward Tilt

  • A shared reward-improving direction from sampled endpoints avoids state-conditional posterior sampling, while the second MSE update supplies local restoration as states drift from rollout predictions.A prompt-level direction may not match posterior guidance at every state, so uncorrected deviations can accumulate.Source: §B.3 · p. 19

C.1 Reverse-Process Policy Optimization

  • FlowGRPO samples stochastic reverse trajectories, group-normalizes terminal rewards, and applies clipped PPO to every transition, requiring tractable transition densities and stored trajectories.The ODE-to-SDE construction is used to expose a tractable stochastic reverse transition density.Source: §C.1 · p. 19

C.2 Reference Transitions for ODE Trust-Region Masking

  • The trust-region construction uses ODE-to-SDE reference transitions; its mask is recomputed before each update and held constant during backpropagation.The reported settings are ν = 0.8 and δ = 10^-5 without further tuning.Source: §C.2 · p. 19

C.3 Pairwise Feedback: Identity and Implementation

DiffusionART-Pairwise converts pairwise outcomes into covariance guidance, with prompt-global aggregation outperforming pair-local guidance in the reported GenEval comparison. The visualization appendix shows same-prompt video and image outputs before and after DiffusionART.

C.3 Pairwise Feedback: Identity and Implementation

  • With K = 2, pairwise covariance guidance points from the worse sample to the better one, while tied pairs contribute zero guidance.For multiple pairs, the global covariance averages pair covariances.Source: §C.3 · p. 20
  • At step 400, prompt-global guidance reaches held-out GenEval 0.91 versus 0.78 for pair-local guidance.The experiment uses 24 images per prompt to form 12 disjoint pairs; prompt-global guidance pools pair covariances before normalization.Source: §C.3 · p. 20, Figure 14 · p. 20

D Visualization

  • The visual comparisons show same-prompt video frames and image outputs before and after DiffusionART, including multi-reward and PickScore image settings.The video pairs use the same prompt, seed, and sampling settings; the image pairs use the same prompt.Source: Figure 15 · p. 21, Figure 16 · p. 22, Figure 17 · p. 23

Related papers

Loading 2610.12572v1…