Summary (Overview)

• Key contribution: DiffusionART (Diffusion Reinforcement with Aggregation-First Reward Tilting) is a novel online reinforcement learning algorithm for fine-tuning diffusion models that is both efficient and stable without per-task hyperparameter tuning.
• Theoretical insight: The paper diagnoses that NFT-style forward-process RL contains 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.
• Method: DiffusionART aggregates per-prompt rewards into a shared reward-improving direction $G(c)$ applied directly to reverse-process states, removing the consistency force by construction, and uses a second update with trust-region masking for drift-adaptive restoration.
• Efficiency gains: On HunyuanVideo-13B with equally weighted visual and motion quality rewards, DiffusionART reaches baseline reward levels with 12.02× and 11.72× fewer GPU hours than DiffusionNFT and Flow-GRPO-Fast respectively, while achieving 28.8% and 53.0% higher mean training rewards.
• Versatility: Supports pairwise reward models (DiffusionART-Pairwise), works with as few as 2 samples per prompt, and scales across text-to-image, text-to-video, and image-to-video generation.

Introduction and Theoretical Foundation

Background and Motivation

Visual diffusion and flow-matching models are typically trained to fit data distributions by regressing noise or velocity targets, rather than optimizing downstream criteria such as visual quality, motion quality, or human preference. This makes reward-guided fine-tuning central to recent image and video generation tasks.

Online diffusion RL faces a fundamental credit-assignment problem: rewards arrive only after a complete output is generated, so costly output-level feedback must be translated into effective updates across many denoising states.

Two Existing Approaches and Their Limitations

Reverse-process RL (e.g., Flow-GRPO):

• Optimizes sampled denoising transitions via group-based endpoint reward feedback
• Trains on states actually visited by the generator
• Requires ODE-to-SDE conversion for tractable likelihoods
• Poor credit assignment limits training efficiency

Forward-process RL (e.g., DiffusionNFT):

• Applies forward noising to generated clean samples and regresses velocities with reward-weighted objectives
• Highly efficient but unstable without two stabilizers:
  1. Joint positive–negative training
  2. Exponential moving average (EMA) rollout policy whose lag must be tuned per task

Key Theoretical Findings

The consistency force puzzle: NFT's positive and negative branches share the same ideal reward-improving direction, so either alone should suffice—yet each alone collapses. The paper reveals this is because NFT's targets implicitly assume the model's sampling velocity equals the ideal forward velocity induced by its own endpoints, which is false in practice.

Lemma 1 (Ideal NFT guidance is conditional reward covariance): At a fixed state $s_t$, if $q := \mathbb{E}_{Q_{old}}[r(x_0, c) \mid s_t]$ and $g := \text{Cov}_{Q_{old}}[r(x_0, c), u_t \mid s_t]$, then:

$$q(v_{fwd}^{+} - v_{fwd}^{old}) = (1-q)(v_{fwd}^{old} - v_{fwd}^{-}) = g$$

This shows positive and negative branches point along the same reward-tilt direction $g$, differing only in step size.

Proposition 1 (Practical forward-process RL contains a forward–reverse consistency force): The normalized output-space force decomposes as:

$$\mathcal{F}_{NFT}(s_t) = 2g + (2q - 1)e$$

where $e = v_{fwd}^{old}(s_t) - v_{rvs}^{old}(s_t)$ is the forward–reverse gap. The residual consistency force $(2q-1)e$ is only partially canceled across samples, not at each state.

Methodology

DiffusionART Algorithm

DiffusionART combines forward-process velocity regression with reverse-process state/action updates:

1. Aggregation-First Guidance

Instead of building regression targets from forward velocity (which introduces the consistency force), DiffusionART aggregates each prompt's endpoints and rewards into normalized guidance:

$$\overline{\text{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}), \qquad G(c) = \frac{\overline{\text{Cov}}}{\text{RMS}(\overline{\text{Cov}})}$$

The target is frozen as $v_{tgt}^{ART}(s_t) = v_\theta^{(0)}(s_t) + G(c)$, where $v_\theta^{(0)} = v_{old}$ is the cached rollout prediction. This removes the consistency force by construction.

Lemma 2 justifies this: under rectified flow, $\text{Cov}_{Q_{old}}(R, u_t \mid c) = \text{Cov}_{p_{old}}(R, x_0 \mid c)$ at every noise level, making $G(c)$ the prompt-level counterpart of ideal NFT guidance.

2. Two-Update Restoration

The rollout batch is split into two disjoint minibatches, each receiving one MSE update. The loss for each state is:

$$\mathcal{L}_\theta(s_t) = w(s_t)\text{MSE}(v_\theta(s_t), v_{tgt}^{ART}(s_t))$$

The second update combines guidance with restoration toward $v_\theta^{(0)}$:

$$-\nabla_{v_\theta}\mathcal{L}_\theta|_{v_\theta = v_\theta^{(1)}} = \frac{2w}{D}(G(c) - d(s_t))$$

where $d(s_t) = v_\theta(s_t) - v_\theta^{(0)}(s_t)$ is the drift. Restoration grows with each state's drift, counteracting large displacements while leaving guidance dominant when drift is small.

3. Trust-Region Masking

A KL proxy measures drift using reference Gaussian transitions:

$$m(s_t) = \frac{1}{D}D_{KL}(q_\theta(\cdot|s_t) \| q_{old}(\cdot|s_t)) = \frac{\|\mu(s_t, v_\theta(s_t)) - \mu(s_t, v_\theta^{(0)}(s_t))\|_2^2}{2D\sigma_{ref}^2(t)}$$

The weight masks excessive outward drift:

$$w(s_t) = t\left[1 - \mathbf{1}\{m(s_t) \geq \delta\} \cdot \mathbf{1}\{\langle d(s_t), G(c) - d(s_t)\rangle > 0\}\right]$$

with drift threshold $\delta = 10^{-5}$ fixed throughout.

4. Pairwise Extension

DiffusionART-Pairwise divides K endpoints into K/2 disjoint pairs with winner/loser labels (1,0) or (1/2,1/2) for ties, replacing scalar rewards in the covariance computation.

Empirical Validation / Results

Text-to-Image Generation (SD3.5-M)

• Single-reward tasks (GenEval, OCR, PickScore): DiffusionART avoids collapse of on-policy DiffusionNFT and matches EMA-based efficiency
• Speedup vs Flow-GRPO: reaches Flow-GRPO's step-500 reward with 10.0×, 5.2×, and 9.6× fewer training steps on GenEval, OCR, and PickScore respectively
• Multi-reward optimization (GenEval + PickScore + CLIPScore + HPSv2.1): Behaves similarly to single-reward training

Scaling to FLUX.1-dev (12B)

• DiffusionART learns fastest among all methods
• DiffusionNFT exhibits pronounced training-reward fluctuations because its tuned hyperparameters don't transfer across models

Ablation Results

• ART: 1 update learns quickly but gradually deteriorates
• ART: 2 updates (with restoration) provides long-term stability
• NFT: 2 updates briefly collapses to near zero before recovering
• Trust-region masking: The no-mask run drops around step 200 before recovering; masked run improves steadily

Text-to-Video Generation (HunyuanVideo-13B)

• DiffusionART surpasses all baselines in both VQ and MQ while maintaining comparable TA
• Reaches training VQ of 6 and MQ of 3 after 29 and 46 steps respectively; baselines never reach these levels
• For VQ of 5.5: 11.4× fewer training steps and 22.1× fewer GPU hours vs DiffusionNFT (EMA)
• Evaluation VQ and MQ improved by 31.4% and 57.4% over the strongest baseline

Image-to-Video Generation (Wan2.1-I2V-14B)

• DiffusionART reaches MQ of 6 in 45.8 GPU hours vs 220.3 for Flow-GRPO-Fast: 4.8× reduction
• Baselines improve slowly or plateau

Group Size and Pairwise Feedback

Table 1: VQ-only evaluation and GPU hours at step 400

With pairwise feedback, K = 2 achieves comparable final VQ (12.235 vs 12.092 for K = 8) while using 3.4× fewer GPU hours.

Theoretical and Practical Implications

Theoretical Contributions

1. Diagnosis of NFT instability: The paper provides the first rigorous decomposition of NFT-style training into reward-improving guidance and a harmful consistency force, showing why joint positive–negative training is insufficient.

2. Restoration as the key stabilizer: EMA stabilizes through restoration toward a lagged reference, not through off-policy data. This clarifies the mechanism and explains why EMA requires per-task tuning.

3. Aggregation-first principle: By aggregating rewards into a shared direction before applying to states, the consistency force is removed by construction rather than by cross-sample cancellation.

Practical Implications

1. No per-task tuning: DiffusionART eliminates EMA schedule tuning entirely, with the trust-region threshold $\delta = 10^{-5}$ fixed across all tasks.

2. Compute efficiency: Order-of-magnitude GPU hour reductions enable RL fine-tuning of large models (13B–14B parameters) on modest budgets.

3. Sample efficiency: Works effectively with as few as 2 samples per prompt, enabling pairwise feedback with minimal rollout cost.

4. Black-box rewards: Requires no differentiable reward model, making it applicable to arbitrary reward functions including human feedback and proprietary scorers.

Conclusion

DiffusionART addresses the instability of forward-process diffusion RL through two key mechanisms: aggregation-first reward tilting that removes the hidden consistency force by construction, and two-update restoration with trust-region masking that provides drift-adaptive stability without EMA tuning.

The work demonstrates efficient, stable reward improvement across text-to-image, text-to-video, and image-to-video generation, including with pairwise feedback and small rollout groups. Beyond the algorithm itself, the paper provides fundamental insights into forward-process training, offering a clearer understanding of diffusion reinforcement learning and a path toward scalable diffusion post-training for visual generation.

Future directions include detailed analysis of why restoration dominates gradient aggregation in NFT-style training (the paper notes restoration has 9.4× smaller output-force RMS but 4.9× larger parameter-gradient norm than guidance), and further investigation of the trust-region mechanism's theoretical properties.