AI Summary

DiffusionART fixes a hidden flaw in diffusion RL and trains video models with 12x fewer GPU hours

Researchers from NUS and Kling AI show that DiffusionNFT, a fast RL method for diffusion models, contains a hidden force that destabilizes training. Their fix, DiffusionART, needs no per-task EMA tuning and reaches baseline reward levels on HunyuanVideo-13B with 11.7x to 12x fewer GPU hours.

There are two main ways to run reinforcement learning on image and video diffusion models, and each has a serious drawback. One is slow. The other is fast but collapses unless you tune it carefully for each task. A new paper from NUS and Kling AI traces that instability to a specific term hidden in the faster method’s loss and removes it. The resulting method, DiffusionART, trained HunyuanVideo-13B to the same reward levels as the two strongest baselines with 12.02x and 11.72x fewer GPU hours, with no per-task tuning.

  • 12x cheaper on 13B text-to-video: with visual and motion quality weighted equally, DiffusionART matched the late-stage rewards of DiffusionNFT and Flow-GRPO-Fast, then finished 28.8% and 53.0% higher
  • 4.8x cheaper on 14B image-to-video: reached a motion-quality score of 6 in 45.8 GPU hours, versus 220.3 for Flow-GRPO-Fast
  • 5x to 10x fewer steps than Flow-GRPO on images: reached Flow-GRPO’s step-500 reward 10.0x faster on GenEval, 5.2x on OCR, and 9.6x on PickScore with SD3.5-M
  • One pair per prompt is enough: with pairwise feedback, 2 samples per prompt matched 8 samples while using 3.4x less compute
  • Caveat: every score in the paper comes from an automatic reward model, often the same one used for training; there is no human evaluation

Why diffusion RL has been fast or stable, but not both

Diffusion models learn to match a data distribution, not to maximize what you care about, such as text legibility, motion quality, or prompt alignment. Online RL fixes that by generating samples, scoring them with a reward model, and updating the generator. The hard part is credit assignment: you get one score for a finished video, and you need to turn it into updates across every denoising step that produced it.

Reverse-process methods such as Flow-GRPO treat each denoising step as a policy action and apply a PPO-style policy gradient. They train on the states the model actually visits, but each step receives a single trajectory-level advantage that it only partly influenced. That weak credit assignment makes them slow.

Forward-process methods such as DiffusionNFT take the finished samples, add noise again, and run supervised velocity regression: pull toward high-reward samples and push away from low-reward ones. This is much more efficient, but it only stays stable with two tricks. It needs both a positive and a negative branch, and it needs an EMA copy of the model to generate rollouts. The EMA decay schedule must be tuned per task. On OCR, the default schedule is unstable and needs a retention of 0.999. When the authors transferred NFT’s GenEval settings to FLUX.1-dev, training became unstable.

The hidden force inside DiffusionNFT

The paper starts with a puzzle. In theory, NFT’s positive and negative branches point the same way, so either should work alone. In practice, each one collapses when trained alone.

The explanation is that NFT builds its training targets from re-noised endpoints. That assumes the model’s denoising velocity matches the “ideal” velocity implied by its own samples, and in practice it doesn’t. The gap adds a consistency force that has nothing to do with reward. The positive branch pulls the model toward the ideal velocity, while the negative branch pushes it away. Training both branches together cancels this force on average across samples, but not at each individual state. The leftover force grows at low-noise steps, where a state already reveals whether its sample will score well.

The authors also show that EMA does not stabilize training because rollouts are off-policy. If you keep the EMA rollouts but remove the old-policy anchor from the loss, training still collapses. What actually helps is restoration: the loss contains a quadratic term that pulls the model back toward the lagged reference. Pure on-policy training has zero drift, so it gets zero restoration and leaves the leftover force unopposed.

The fix: aggregate first, restore second

DiffusionART changes two things.

One DiffusionART training iteration

1 ¡ ROLL OUT

Sample K outputs per prompt with a plain ODE solver and score them

→

2 ¡ AGGREGATE

Reward-vs-output covariance gives one normalized direction per prompt

→

3 ¡ UPDATE 1

Regress visited states toward old velocity + that direction

→

4 ¡ UPDATE 2

Same target on a second minibatch; drift adds restoration, trust-region mask caps it

Aggregation-first guidance. For each prompt, DiffusionART calculates the covariance between the rewards and the generated outputs. Mathematically, this is the direction in which tilting the distribution toward higher reward moves its mean. The covariance is RMS-normalized, and the same direction is applied to every state visited during that prompt’s rollout. The target never involves the re-noised forward velocity, so the consistency force is absent by design rather than cancelled after the fact. RMS normalization also removes the guidance-strength hyperparameter.

Two-update restoration. Each rollout batch is split into two minibatches with one update per minibatch. On the first update, the model has not moved, so it receives only guidance. On the second update, it has drifted, and the same MSE loss automatically adds a pull back toward the original velocity in proportion to that drift. That replaces EMA and removes its schedule. A trust-region mask, inspired by PPO clipping, drops updates that push states with too much drift even further out. The threshold is fixed at 1e-5 for every experiment.

Because the method only needs a reward value per sample, it also works with pairwise preference models. Pairs receive winner and loser labels of 1 and 0, and those labels go into the same covariance calculation.

Results across image and video

The biggest gains came on video, where each rollout is most expensive. On HunyuanVideo-13B with matched rewards (text alignment, visual quality, and motion quality weighted 7:2:1), DiffusionART reached a visual quality score of 5.5 in 63.4 GPU hours. DiffusionNFT needed 1,401.4, which is 22.1x more.

HunyuanVideo-13B text-to-video, evaluation after training

VideoReward visual quality score. Higher is better. Training weights TA:VQ:MQ = 7:2:1.

DiffusionART8.97
DiffusionNFT (EMA)6.83
Pretrained model5.20
DanceGRPO3.92
Flow-GRPO-Fast3.15

DiffusionART raised motion quality from 1.94 to 4.22 and kept text alignment level with the best baseline (2.11 vs. 2.10). The two GRPO-style methods ended below the pretrained model on visual quality, and on-policy DiffusionNFT’s text alignment fell to −4.48. Overall, DiffusionART improved evaluation visual quality by 31.4% and motion quality by 57.4% over the strongest baseline.

The image-to-video results on Wan2.1-I2V-14B, a 4-step distilled model, show a larger gap: visual quality was 11.20 and motion quality was 10.48, compared with 5.77 and 4.71 for Flow-GRPO-Fast. On FLUX.1-dev (12B), DiffusionART reached a GenEval score of 0.968 at step 500, compared with 0.918 for Flow-DPPO and 0.849 for DiffusionNFT, which had not been retuned for FLUX.

The pairwise ablation is useful in practice. With only 2 samples per prompt, DiffusionART reached a final visual quality of 12.24 in 260.8 GPU hours. With 8 samples, it reached 12.09 in 883.3 GPU hours. For video RL, where rollouts are the main cost, one pair per prompt may be enough.

Caveats

  • Reward models were used both to train and to score. VideoReward was used both to train and to evaluate the video models, and there was no human evaluation. Very large gains, such as I2V motion quality rising from about 4.7 to 10.5, can mean real improvement or a model learning to game the scorer. The paper includes before-and-after frames but no human preference study.
  • The image gains are mostly about speed. On SD3.5-M, DiffusionART matched tuned DiffusionNFT (EMA) rather than beating it. Its advantage there is that it needs no per-task tuning.
  • The comparisons used one run each. The paper reports no seeds or variance, so the training curves show one run per method.
  • Some stability claims cover limited training windows. The paper says a variant “avoids collapse over the observed training window,” which is not the same as stability over long training runs.
  • This is the first version of the paper. The code is on GitHub, but the results have not been replicated independently.

The paper’s main contribution is the diagnosis, more than another leaderboard result. If you have used DiffusionNFT and found that each new reward or model needed its own EMA schedule, this paper explains why: EMA was compensating for a force that should not be in the loss. Removing that force at its source, rather than tuning around it, made the method faster and portable across tasks.

#research #reinforcement-learning #training #video-generation

Liked this? Engineer's Codex sends one deep dive and a link roundup every week.

No spam. Unsubscribe anytime.