HyperAIHyperAI

Command Palette

Search for a command to run...

DIFFUSIONART : renforcement efficace, stable et à grande échelle des modèles de diffusion

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

Résumé

L’apprentissage par renforcement (RL) en ligne ajuste finement les modèles de diffusion à l’aide de récompenses en boîte noire, mais les méthodes fondées sur le processus inverse, telles que Flow-GRPO, voient leur efficacité limitée par une mauvaise attribution de crédit, tandis que les méthodes fondées sur le processus direct, telles que DiffusionNFT, sont efficaces mais instables sans un entraînement off-policy à moyenne mobile exponentielle (EMA) soigneusement réglé. Nous montrons que les objectifs de type NFT contiennent une force de cohérence direct–inverse cachée que l’entraînement conjoint positif–négatif n’annule que partiellement, et que l’EMA stabilise l’entraînement par restauration plutôt que par des déploiements hors politique. Nous introduisons DiffusionART (Diffusion Reinforcement with Aggregation-First Reward Tilting), qui agrège les récompenses de chaque prompt en une direction commune d’amélioration des récompenses pour les états du processus inverse, éliminant par construction la force de cohérence, et ajoute une seconde mise à jour masquée par région de confiance pour une restauration adaptative à la dérive, sans réglage de l’EMA. Elle ne nécessite aucun modèle de récompense différentiable, reste efficace avec moins d’échantillons par prompt et prend en charge les modèles de récompense par paires (DiffusionART-Pairwise). En génération d’images et de vidéos, DiffusionART s’entraîne de manière efficace et stable sans réglage par tâche. En particulier, sur HunyuanVideo-13B avec des récompenses de qualité visuelle et de mouvement équipondérées, DiffusionART atteint les niveaux de récompense moyenne de fin d’entraînement de DiffusionNFT et de Flow-GRPO-Fast avec respectivement 12,02× et 11,72× moins d’heures GPU, tout en obtenant des récompenses moyennes d’entraînement supérieures de 28,8 % et 53,0 %.

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×12.02\times12.02× and 11.72×11.72\times11.72× fewer GPU hours, respectively, while achieving 28.8%28.8\%28.8% and 53.0%53.0\%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)G(c)G(c). The authors compute the reward-endpoint covariance from KKK rollout endpoints and normalize it to obtain G(c)G(c)G(c):

Cov⁡‾(R,x0∣c)=1K∑i=1K(Ri−Rˉc)(x0(i)−xˉ0,c),G(c)=Cov⁡‾RMS⁡(Cov⁡‾)\overline{\operatorname{Cov}}(R, x_0 \mid c) = \frac{1}{K} \sum_{i=1}^{K} (R_i - \bar{R}_c)(x_0^{(i)} - \bar{x}_{0,c}), \quad G(c) = \frac{\overline{\operatorname{Cov}}}{\operatorname{RMS}(\overline{\operatorname{Cov}})}Cov(R,x0​∣c)=K1​i=1∑K​(Ri​−Rˉc​)(x0(i)​−xˉ0,c​),G(c)=RMS(Cov)Cov​

This normalized guidance is applied to every reverse-rollout state of the prompt. By freezing the target vtgtART(st)=vθ(0)(st)+G(c)v_{\mathrm{tgt}}^{\mathrm{ART}}(s_t) = v_\theta^{(0)}(s_t) + G(c)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 vtgtARTv_{\mathrm{tgt}}^{\mathrm{ART}}vtgtART​. The first update applies the guidance G(c)G(c)G(c). The second update combines the original guidance with a restoration term −d(st)-d(s_t)−d(st​) toward the initial velocity field vθ(0)v_\theta^{(0)}vθ(0)​, where d(st)d(s_t)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)=∥μ(st,vθ(st))−μ(st,vθ(0)(st))∥222Dσref2(t)m(s_t) = \frac{\|\mu(s_t, v_\theta(s_t)) - \mu(s_t, v_\theta^{(0)}(s_t))\|_2^2}{2D\sigma_{\mathrm{ref}}^2(t)}m(st​)=2Dσref2​(t)∥μ(st​,vθ​(st​))−μ(st​,vθ(0)​(st​))∥22​​

The mask identifies excessive outward drift by checking if the displacement d(st)d(s_t)d(st​) and the second update's descent direction G(c)−d(st)G(c) - d(s_t)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 KKK 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.


Créer de l'IA avec l'IA

De l'idée au lancement — accélérez votre développement IA avec le co-codage IA gratuit, un environnement prêt à l'emploi et le meilleur prix pour les GPU.

Codage assisté par IA
GPU prêts à l’emploi
Tarifs les plus avantageux

HyperAI Newsletters

Abonnez-vous à nos dernières mises à jour
Nous vous enverrons les dernières mises à jour de la semaine dans votre boîte de réception à neuf heures chaque lundi matin
Propulsé par MailChimp