DiffusionART: Efficient, Stable and Scalable Reinforcement of Diffusion Models

Efficient, Stable and Scalable
Reinforcement of Diffusion Models

Zhenglin Wan^1,2Jiayu Chen^2,†Zhengfeiyang Zhang¹Pengfei Zhou¹Wangbo Zhao³Fangyuan Kong²Pengfei Wan²Kun Gai²Yang You¹

¹ HPC-AI Lab, NUS² Kling AI³ HKUST

Paper

Code Image model Video model

Demo Contact

Aggregate reward guidance first.
Restore as the policy moves.

12.02× fewer GPU hours and 28.8% higher training reward
vs. DiffusionNFT (EMA) · HunyuanVideo-13B · VQ:MQ = 1:1 · final-50-update mean reward

A waterfall pours into a turquoise rainforest pool.

Sunlight breaks through mist above an alpine lake.

A sea turtle swims over a colorful coral reef.

Understanding conventional forward-process RL

We study its hidden forces through DiffusionNFT, then rethink guidance and restoration.

Why do two promising branches collapse?

Ideally, the positive and negative branches follow the same reward-improving direction. Yet either branch alone collapses in our on-policy ablation, while joint training learns.

DiffusionNFT's objective contains an implicit forward–reverse consistency term. Its uncancelled force can destabilize training and slow reward improvement.

F_{NFT} = 2 g + (2 q - 1) e

Reward guidanceHidden consistency force

The branches reinforce guidance but oppose the consistency force. Centering rewards across samples does not balance the forces at every state, so joint training cancels them only partially.

Remove the consistency force state by state, rather than rely on cancellation across samples.

The forward–reverse gap and symbols

Generated endpoints induce an ideal forward velocity that can differ from the reverse sampler's velocity. Their difference is the consistency gap e. The formula is the normalized output-space force at on-policy initialization for the NFT objective. With goodness weight r in [0,1], q = E[r | sₜ] and g = Cov(r, uₜ | sₜ), where uₜ = x₀ − ε is the rectified-flow target. See Section 3.1 of the paper.

On-policy GenEval ablation: positive-only and negative-only NFT branches collapse, while their combination improves reward. — Each branch alone collapses, despite the same ideal guidance direction.

NFT's positive pull and negative push only partially cancel their forward-reverse consistency forces. — The forward–reverse gap introduces opposing forces that only partially cancel.

GenEval reward collapses when EMA rollouts are retained but the old-policy restoration anchor is removed; NFT with EMA remains stable. — EMA rollouts alone do not prevent collapse when the restoring anchor is removed.

Restoration is crucial, not just EMA data.

What does EMA actually stabilize? Keeping EMA-policy rollouts but removing the restoring reference still causes collapse in our ablation. Off-policy data alone is insufficient.

The NFT objective already contains a restoring anchor: rewriting its loss reveals a quadratic penalty on drift from the frozen rollout policy. EMA keeps this reference lagged, making restoration active as the policy moves.

Our NFT ablations identify restoration as the crucial stabilizer for this form of diffusion RL.

L_{NFT} = E_{D_{old}} [- 2 β (2 r - 1) ⟨ d, h ⟩ + β^{2} {∥ d ∥}_{2}^{2}] + const

\begin{array}{l} L_{NFT} = E_{D_{old}} [- 2 β (2 r - 1) ⟨ d, h ⟩ \\ + β^{2} {∥ d ∥}_{2}^{2}] + const \end{array}

Reward & consistency forcesRestoration to the reference

d = v_θ − v_oldh = u_t − v_old

Loss decomposition and definitions

The expectation is over old-policy training data. v_old is the frozen rollout-policy prediction, u_t = x₀ − ε the forward velocity target, r ∈ [0,1] the reward-derived goodness weight, and β the NFT tilt coefficient. The constant does not depend on trainable parameters. The linear term carries both reward guidance and consistency force; the quadratic term restores the prediction to the reference. Its gradient vanishes at d = 0. This term is part of NFT's objective, not an extra regularizer added by EMA. See Section 3.2.

Stable does not always mean fast.

Without EMA, two updates on disjoint minibatches retain one frozen reference. The first creates drift; the second adds restoration. Collapse disappears over the observed window, but reaching OCR 0.9 takes 15.2× more steps.

Left: without EMA, NFT with two disjoint updates reaches OCR 0.9 at step 320, versus step 21 for one update before it collapses. Right: EMA rollouts without the restoring anchor also collapse. — Left: no EMA, β = 1; 21 versus 320 updates to OCR 0.9. Right: removing the anchor despite EMA-generated data.

A small force can dominate the gradient.

Restoration has 9.4× smaller output-force RMS than guidance, yet a 4.9× larger aggregated parameter-gradient norm. This suggests heterogeneous guidance cancels across states while restoration survives aggregation.

Relative guidance and restoration scales: restoration is 9.4 times smaller in output-force RMS but 4.9 times larger in aggregated parameter-gradient norm. — Second-update gradients, averaged over rollouts 0–17. These are observed ratios, not a universal scaling law.

Preserve restoration without letting it overwhelm reward guidance.
This is the balance DiffusionART is designed to achieve.

Change the order. Remove the force.

Aggregate endpoints and rewards before regression, then reuse a frozen target.

DiffusionART aggregates endpoint reward covariance, normalizes once, and shares this direction across states of reverse ODE trajectories.

Aggregation-first guidance removes the hidden consistency force.

One direction for each prompt

Reward–endpoint covariance supplies shared guidance on the states actually visited by the reverse sampler. Normalize after aggregation, rather than relying on samplewise forces to cancel.

G (c) = \frac{\bar{Cov} (R, x_{0} ∣ c)}{RMS (\bar{Cov})}

For nonzero covariance; otherwise G(c) = 0.

R is reward, x₀ a clean latent endpoint, and c the prompt. The covariance is estimated over its K samples.

Two updates. One frozen target.

Fit v_old + G on two disjoint minibatches. The first update follows guidance; the second also restores the model in proportion to its output drift.

First update

G

Second update

G - d

Directions in output space, up to the MSE scale and state weight; d = v_θ − v_old. A trust-region mask suppresses excessive outward updates.

No differentiable reward model.
No EMA schedule to tune.

Less compute. Higher rewards.

Matched HunyuanVideo-13B experiments with VQ:MQ = 1:1.

Mean raw VQ/MQ reward against training steps and GPU hours. DiffusionART reaches each baseline's late-stage mean reward earlier, with 12.02 times and 11.72 times fewer GPU hours.

The same reward, reached earlier.
Then a higher reward.

DiffusionART versus each baseline
Baseline	Fewer GPU hours	Higher reward
DiffusionNFT (EMA)	12.02×	+28.8%
Flow-GRPO-Fast	11.72×	+53.0%

Mean raw reward is (VQ + MQ) / 2. Speedups compare the first crossings of each baseline's final-50-update mean reward. Reward gains compare the final 50 logged updates of each method.

Measurement protocol and data

All methods use an eight-update causal EMA for crossing measurement. Targets are unsmoothed baseline means; crossings use observed points, without interpolation. Elapsed GPU hours include evaluation and waits. The baseline levels are empirical late-stage rewards, not a guarantee of asymptotic convergence.

Curve data / Crossing audit / Figure sources

Held-out video quality also improves.

+31.4% VQ and +57.4% MQ over the strongest baseline on these metrics.

Separate TA:VQ:MQ = 7:2:1 experiment. Evaluation scores: VQ 8.970, MQ 4.224. These are not the training-reward percentages above.

Beyond the headline comparison

Setting	Result
Text to image · SD3.5-M	10.0× / 5.2× / 9.6× fewer training steps to Flow-GRPO's step-500 rewards on GenEval / OCR / PickScore.
Image to video · Wan2.1-I2V-14B	MQ 6 in 45.8 vs. 220.3 GPU hours: 4.8× fewer than Flow-GRPO-Fast.
Pairwise video feedback · K = 2	3.4× fewer GPU hours than K = 8 over 400 updates, with comparable step-400 evaluation VQ (12.235 vs. 12.092).

Separate experiments with their own settings: I2V uses VQ:MQ = 1:1; pairwise video uses VQ-only feedback. Details in the paper.

See what the model learns.

20 video pairs and 6 image pairs, with their prompts.

Explore demos

Citation

@misc{wan2026diffusionartefficientstablescalable,
  title={DiffusionART: Efficient, Stable and Scalable Reinforcement of Diffusion Models},
  author={Zhenglin Wan and Jiayu Chen and Zhengfeiyang Zhang and Pengfei Zhou and Wangbo Zhao and Fangyuan Kong and Pengfei Wan and Kun Gai and Yang You},
  year={2026},
  eprint={2610.12572},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2610.12572},
}