Command Palette
Search for a command to run...
DIFFUSIONART: EFFICIENT, STABLE AND SCALABLE REINFORCEMENT OF DIFFUSION MODELS
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
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× fewer GPU hours, respectively, while achieving 28.8% and 53.0% higher mean training rewards.
One-sentence Summary
HPC-AI Lab, National University of Singapore, Kling AI, and HKUST introduce DiffusionART, which aggregates each prompt's rewards into a shared reward-improving direction for reverse-process states to remove the consistency force by construction and adds a trust-region-masked second update for drift-adaptive restoration without EMA tuning, enabling efficient and stable image and video generation, and on HunyuanVideo-13B reaches the late-stage mean reward levels of DiffusionNFT and Flow-GRPO-Fast with 12.02× and 11.72× fewer GPU hours, respectively, while achieving 28.8% and 53.0% higher mean training rewards.
Key Contributions
- The paper shows that NFT-style forward-process objectives contain a hidden forward-reverse consistency force that joint positive-negative training cancels only partially, and that EMA stabilization operates through restoration rather than off-policy rollouts.
- The method introduces DiffusionART, which aggregates each prompt's rewards into a shared reward-improving direction for reverse-process states, removes the consistency force by construction, and adds a trust-region-masked second update for drift-adaptive restoration without EMA tuning or a differentiable reward model.
- The paper also presents DiffusionART-Pairwise for pairwise reward models and reports efficient, stable training on image and video generation. On HunyuanVideo-13B, DiffusionART reaches the mean reward levels of DiffusionNFT and Flow-GRPO-Fast with 12.02× and 11.72× fewer GPU hours while achieving 28.8% and 53.0% higher mean training rewards.
Introduction
Reward-guided fine-tuning is important for modern image and video diffusion and flow models because standard training optimizes noise or velocity regression rather than downstream criteria such as visual quality, motion quality, or human preference. Online diffusion reinforcement learning can use black-box rewards from complete generated samples, but it faces a credit assignment problem because output-level feedback must be translated across many denoising states. Prior online methods have notable tradeoffs: reverse-process approaches such as Flow-GRPO require ODE-to-SDE conversion and suffer from noisy trajectory-level credit assignment, while forward-process approaches such as DiffusionNFT are efficient but unstable without joint positive-negative training and a task-tuned EMA rollout policy. The authors diagnose NFT-style instability as a hidden forward-reverse consistency force that joint training only partially cancels and that EMA stabilizes through restoration rather than off-policy rollouts. Their contribution, DiffusionART, uses aggregation-first reward tilting to assign a shared reward-improving direction to reverse-process states and adds a trust-region restoration update, removing the harmful force by construction and eliminating per-task EMA tuning while maintaining efficient and stable reward gains.
Method
The authors first examine the stability mechanisms in NFT-style forward-process training, specifically joint positive-negative training and off-policy exponential moving average (EMA). They reveal that while the ideal positive and negative guidance directions are collinear, a hidden forward-reverse consistency force destabilizes individual branches. Joint training partially cancels this force across samples, but not at each state, leaving a residual that can hinder optimization.
To isolate the stabilizing mechanism of EMA, the authors analyze the loss decomposition and demonstrate that EMA provides explicit restoration toward a lagged reference policy. This restoration counteracts the residual consistency force. However, relying on EMA requires task-specific tuning of the restoration schedule, as too little restoration risks collapse while too much slows reward improvement.
To address the limitations of NFT-style training, the authors introduce DiffusionART, which combines the efficiency of forward-process regression with the stability of reverse-process rollouts.
The consistency force in NFT arises because its regression target is built from the forward velocity, which differs from the reverse sampling field. DiffusionART avoids this by aggregating each prompt's endpoints and rewards into a shared reward-improving direction G(c). The authors compute the reward-endpoint covariance from K rollout endpoints and normalize it to obtain G(c):
Cov(R,x0∣c)=K1i=1∑K(Ri−Rˉc)(x0(i)−xˉ0,c),G(c)=RMS(Cov)CovThis normalized guidance is applied to every reverse-rollout state of the prompt. By freezing the target vtgtART(st)=vθ(0)(st)+G(c), every state receives an explicit reward-improving direction without the NFT consistency term.
Instead of tuning an EMA schedule, the authors implement a drift-adaptive restoration mechanism by splitting each rollout batch into two disjoint minibatches and applying one mean squared error update to each. Both updates use the frozen target vtgtART. The first update applies the guidance G(c). The second update combines the original guidance with a restoration term −d(st) toward the initial velocity field vθ(0), where d(st) is the current displacement. This ensures restoration grows with each state's drift, counteracting large displacements while leaving guidance dominant when drift is small, eliminating the need for task-specific EMA tuning.
To prevent excessive policy changes within a single rollout batch, the authors introduce a trust-region mask. Since ordinary differential equation transitions are deterministic, they define reference Gaussian transitions to compute a KL proxy for measuring drift:
m(st)=2Dσref2(t)∥μ(st,vθ(st))−μ(st,vθ(0)(st))∥22The mask identifies excessive outward drift by checking if the displacement d(st) and the second update's descent direction G(c)−d(st) have a positive inner product, masking the update if the drift exceeds a threshold.
For visual outputs lacking objective correctness criteria, the authors extend the framework to support online pairwise preference feedback. They divide K endpoints into disjoint pairs, assign winner or loser labels, and replace scalar rewards with these labels in the covariance and training pipeline.
Experiment
The experiments evaluate DiffusionART across text-to-image, text-to-video, and image-to-video generation, comparing it with baselines such as Flow-GRPO, DanceGRPO, and DiffusionNFT under LoRA and full-model tuning. DiffusionART consistently achieves faster reward improvement and greater training efficiency while avoiding the instability or collapse observed in some baselines, and it scales effectively to larger models. Ablations show that its covariance guidance, restoration updates, and trust-region masking contribute to long-term stability, while pairwise feedback and smaller group sizes remain effective in video tasks.
Pairwise feedback achieves final VQ close to scalar feedback, and reducing the group size from eight to two preserves final VQ while lowering compute substantially. The smaller group uses about a quarter of the rollout videos and training samples but still reaches a final VQ no lower than the larger group. This suggests one pair per prompt is sufficient for stable and compute-efficient optimization in this setting. Pairwise feedback with the smallest group size matches or slightly exceeds the final VQ of the larger pairwise group under the same step budget. Reducing the group size from eight to two lowers GPU hours by roughly 3.4 times while maintaining comparable VQ. Scalar feedback with eight samples per prompt yields the highest VQ in this evaluation but requires the most GPU hours.
Pairwise feedback reaches final video quality close to scalar feedback, and cutting the pairwise group size from eight to two preserves final video quality while using about a quarter of the rollout videos and training samples. The smallest pairwise group matches or slightly exceeds the larger pairwise group under the same step budget, suggesting that one pair per prompt is sufficient for stable and compute-efficient optimization. Reducing the group size from eight to two lowers GPU hours by roughly 3.4 times without hurting video quality. Scalar feedback with eight samples per prompt achieves the highest video quality in this evaluation but requires the most compute.