DiffusionART: Efficient, Stable and Scalable Reinforcement of Diffusion Models

AI summary

Overview

• Research area: Online reinforcement learning (RL) for diffusion and flow-matching generative models, covering text-to-image, text-to-video, and image-to-video generation (HPC-AI Lab at the National University of Singapore, Kling AI, and HKUST).
• Technical level: Advanced. The paper builds on the Flow-GRPO and DiffusionNFT literature, uses velocity-field notation for rectified flow, and presents lemmas, propositions, and proofs.
• Scope in one sentence: The paper diagnoses why forward-process diffusion RL (NFT-style training) is unstable, then proposes DiffusionART, an aggregation-first reward-tilting method with drift-adaptive restoration that is claimed to train efficiently and stably without EMA tuning or per-task hyperparameter search.

What This Paper Is About

Online RL fine-tunes diffusion models using black-box rewards, but the two existing families each have a problem: reverse-process methods like Flow-GRPO assign credit poorly and are slow, while forward-process methods like DiffusionNFT are fast but unstable unless they use carefully tuned exponential moving average (EMA) off-policy training. DiffusionART aims to get both properties at once, by removing the hidden "forward–reverse consistency force" in NFT-style objectives and replacing EMA with a restoration mechanism that adapts to each state's drift.

Key Contributions

1. Diagnosing NFT instability. The paper shows that the NFT objective implicitly contains a harmful forward–reverse consistency force, arising because forward-noised endpoints assume the model's sampling velocity equals the ideal forward velocity, and that joint positive–negative training cancels this force only partially (Section 3.1).
2. Identifying restoration as the stabilizer. It argues that EMA stabilizes NFT through the restoring force of its lagged reference, not through off-policy rollouts, but that this requires per-task tuning (Section 3.2).
3. DiffusionART: aggregation-first reward tilting with adaptive restoration. The method aggregates each prompt's samples and rewards into a shared reward-improving direction G(c), removing the consistency force by construction, and adds a trust-region-masked second update for drift-adaptive restoration without EMA tuning (Section 4).
4. Efficiency and stability at scale. DiffusionART is reported to match the strongest baseline's efficiency on SD3.5-M without EMA or per-task tuning, and to reach late-stage reward levels far faster than baselines on HunyuanVideo-13B and a 14B image-to-video model (Section 5).

Main Findings

• Joint positive–negative training is only a partial fix. Either branch alone collapses, and joint training works because the branches reinforce each other along the reward-tilt direction g but oppose each other along the forward–reverse gap e, leaving a residual of (2q − 1)e. Prompt-level reward centering makes q average to 1/2 but does not enforce q = 1/2 at every state, so cancellation is inexact; as noise decreases, states reveal more endpoint quality and the weight |2q − 1| grows.
• Restoration, not off-policy data, is what stabilizes NFT. Retaining EMA-policy rollouts but removing the old-policy anchor from the gradient still causes collapse. Without EMA, a two-update-per-rollout variant avoids collapse in the observed window but slows training by 15.2 times.
• Restoration survives aggregation while guidance largely cancels. Averaged over rollouts 0–17, restoration has a 9.4 times smaller output-force RMS than guidance, yet a 4.9 times larger parameter-gradient norm, which the authors attribute to heterogeneous per-sample guidance cancelling across states.
• EMA is sensitive. The default NFT schedule η = min(0.001i, 0.5), where i is the iteration, is unstable on OCR, which requires a maximum retention of 0.999; transferring NFT's GenEval-tuned EMA and β settings to FLUX also leads to unstable training and can cause collapse.
• Text-to-image results. On SD3.5-M, DiffusionART reaches Flow-GRPO's step-500 reward with 10.0 times, 5.2 times, and 9.6 times fewer training steps on GenEval, OCR, and PickScore, respectively. On FLUX.1-dev (12B) for GenEval, it learns fastest among all methods while DiffusionNFT shows pronounced reward fluctuations.
• Video results. On HunyuanVideo-13B with VQ:MQ = 1:1 weights, DiffusionART reaches the late-stage mean reward levels of DiffusionNFT and Flow-GRPO-Fast with 12.02 times and 11.72 times fewer GPU hours, while achieving 28.8% and 53.0% higher mean training rewards, respectively. With VQ:MQ = 7:2:1, reaching VQ of 5.5 takes 24 steps and 63.4 GPU hours versus 273 steps and 1,401.4 GPU hours for DiffusionNFT (EMA), a 11.4 times reduction in steps and 22.1 times reduction in GPU hours.
• Evaluation scores. DiffusionART retains its gains through 400 steps and achieves evaluation VQ and MQ of 8.970 and 4.224, improving over the strongest baseline by 31.4% and 57.4%.
• Pairwise rewards work with small groups. Under a 400-step budget, pairwise K = 2 reaches a final VQ of 12.235 (260.8 GPU hours), no lower than 12.092 for K = 8 (883.3 GPU hours), a 3.4 times GPU-hour reduction. Scalar feedback with K = 8 used 891.1 GPU hours for VQ 13.637.
• Image-to-video. DiffusionART reaches MQ of 6 in 45.8 GPU hours versus 220.3 for Flow-GRPO-Fast, a 4.8 times reduction; DiffusionNFT plateaus early and stays near that plateau through 400 updates.
• Trust-region masking helps stability. Removing only the mask (with VQ:MQ = 1:1) causes the run to drop around step 200 before recovering, whereas the masked run improves steadily.
• Guidance and restoration ablation. On OCR, "ART: 1 update" learns quickly but gradually deteriorates, "ART: 2 updates" adds restoration and provides long-term stability, while "NFT: 2 updates" with tuned β = 0.1 briefly collapses to near zero before recovering.

Methodology in Plain English

The authors start from an observation about how DiffusionNFT works. NFT takes generated samples, adds noise to them (a forward step), and then trains the model to move its predicted velocity toward the reward-weighted "good" samples and away from the "bad" ones. The paper shows this target secretly asks the model to match a forward-direction velocity that the model's own sampler does not actually produce, creating a hidden consistency force that distorts learning. Because NFT's positive and negative branches push in opposite directions along this unwanted force, training both branches together cancels most of it, but only approximately, and only averaged over samples rather than at each state. The paper then shows that EMA stabilizes NFT by pulling the model back toward a lagged reference copy of itself — a restoration effect — not by providing off-policy data, and that the strength of that pull has to be hand-tuned per task.

DiffusionART therefore changes two things. First, it does not regress toward forward-noised velocities at all. Instead, for each prompt it computes a covariance between the reward deviations and the endpoint deviations across the prompt's K rollouts, normalizes it with RMS normalization to get a direction G(c) that is independent of reward scale, and applies that same direction as a target offset to every state visited during the rollout. The reward-improving direction is shared across states, so there is nothing to cancel and no per-state consistency force. Second, it splits the rollout batch into two disjoint halves and performs one MSE update on each, both against the same frozen target. The first update applies guidance; the second update's gradient contains guidance minus the current displacement, so any state that has drifted far gets pulled back proportionally to its drift, with no EMA schedule to tune. A trust-region mask turns off updates where a KL proxy between the current and old policy's reference Gaussian transitions is at or above a threshold (δ = 10⁻⁵) and the update would push outward. A pairwise variant replaces scalar rewards with winner/loser labels over K/2 pairs.

Why This Matters

Impact on research. The paper reframes why forward-process diffusion RL is unstable, replacing a vague appeal to off-policy data with a specific geometric account (a consistency force plus a restoration requirement). If the analysis holds, it means EMA schedules and joint-branch tricks are workarounds for a problem that can be removed by construction, which suggests a different design principle for future reward-tuning algorithms.

Real-world applications.

• Fine-tuning text-to-image generators for human preference, prompt following, and text rendering (OCR) without hand-tuning per reward.
• Improving visual quality and motion quality in text-to-video generation on large models such as HunyuanVideo-13B.
• Animating still images into video, where the paper reports reaching a motion-quality threshold in 45.8 GPU hours versus 220.3 for Flow-GRPO-Fast.
• Preference-based alignment pipelines where only pairwise human or model comparisons are available, since DiffusionART-Pairwise supports pairwise reward models directly.

Industry relevance. The reported GPU-hour reductions are the headline: 12.02 times and 11.72 times fewer GPU hours on a 13B video model, and a 4.8 times reduction on a 14B image-to-video model. Reducing compute for reward tuning, avoiding per-task EMA searches, and working with as few as one pair per prompt (K = 2) all lower the practical cost of adapting large generative models to new objectives. The authors also make code and a project page available.

Future Directions

• Explain why restoration dominates the aggregated parameter gradient despite having a much smaller output-force RMS; the paper explicitly leaves a detailed analysis of this to future work, calling it counterintuitive.
• Determine whether the two-update restoration mechanism can be improved so it does not slow training, given that the EMA-free two-update NFT variant cost 15.2 times slower training.
• Establish whether the covariance-based guidance and the trust-region threshold δ = 10⁻⁵ generalize beyond the tested models, rewards, and resolutions, including the truncated image-to-video evaluation and other reward models.
• Explore whether the aggregation-first view extends to other online diffusion RL settings, such as offline preference learning or differentiable reward optimization, which the paper positions as distinct alternatives.

Target Audience

Researchers and engineers working on reinforcement learning or reward fine-tuning for diffusion and flow-matching generative models, particularly those who have struggled with Flow-GRPO's sample inefficiency or DiffusionNFT's EMA sensitivity. Practitioners training large image and video models who care about GPU-hour cost will find the video results and the K = 2 pairwise finding most directly useful. Readers without background in rectified flow, velocity fields, or policy-gradient methods will find the theoretical sections demanding, though the problem framing and experimental results are accessible.

Authors’ abstract

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.