Command Palette
Search for a command to run...
DiffusionART: 拡散モデルの効率的・安定的・スケーラブルな強化学習
DiffusionART: 拡散モデルの効率的・安定的・スケーラブルな強化学習
Zhenglin Wan Jiayu Chen Zhengfeiyang Zhang Pengfei Zhou Wangbo Zhao Fangyuan Kong Pengfei Wan Kun Gai Yang You
概要
オンライン強化学習(RL)はブラックボックス報酬を用いて拡散モデルをファインチューニングするが、Flow-GRPOのような逆過程手法はクレジット割り当てが不十分なため効率が制限され、DiffusionNFTのような順過程手法は効率的であるものの、注意深く調整された指数移動平均(EMA)オフポリシー学習なしでは不安定である。本研究では、NFT型目的関数には順過程–逆過程間の隠れた整合性の力が含まれており、正例と負例を併用した学習ではその力が部分的にしか相殺されないこと、またEMAはオフポリシーのロールアウトではなく復元を通じて学習を安定化することを示す。我々はDiffusionART(Diffusion Reinforcement with Aggregation-First Reward Tilting)を提案する。DiffusionARTは各プロンプトの報酬を逆過程状態に対する共有の報酬改善方向へ集約することで整合性の力を構造的に除去し、さらに信頼領域でマスクされた第2の更新を加えることで、EMAの調整なしにドリフト適応的な復元を実現する。微分可能な報酬モデルを必要とせず、プロンプトあたりのサンプル数が少ない場合でも有効であり、ペアワイズ報酬モデル(DiffusionART-Pairwise)にも対応する。画像生成と動画生成のいずれにおいても、DiffusionARTはタスクごとの調整なしで効率的かつ安定的に学習する。特に、視覚品質報酬と動き品質報酬を等重み付けしたHunyuanVideo-13Bでは、DiffusionARTはDiffusionNFTおよびFlow-GRPO-Fastの後期平均報酬水準に、それぞれ12.02分の1、11.72分の1のGPU時間で到達し、平均訓練報酬はそれぞれ28.8%、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× 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.