Advantage-Weighted Divergence in RL Diffusion Models

• Advantage-Weighted Divergence Objective is a policy-gradient formulation that aligns diffusion model pretraining and RL fine-tuning through an advantage-weighted loss surrogate.
• It leverages a clean denoising score-matching (DSM) loss to reduce variance and enhance convergence compared to noisy surrogates like DDPO.
• Empirical results show significant speedups and improved sample efficiency, unifying conceptual frameworks for high-dimensional generative modeling.

Advantage-Weighted Divergence Objective (Advantage-Weighted Matching, AWM) is a policy-gradient formulation for reinforcement learning (RL) with diffusion models that aligns the RL fine-tuning objective with the pretraining score/flow-matching objective. AWM addresses the divergence between pretraining and RL post-training in diffusion models, offering reduced variance, faster convergence, and sample efficiency improvements compared to previous approaches such as Denoising Diffusion Policy Optimization (DDPO). This methodology unifies pretraining and RL both conceptually and algorithmically by constructing an advantage-weighted, policy-gradient surrogate based on the exact evidence lower bound (ELBO) or flow-matching loss used in diffusion model pretraining (Xue et al., 29 Sep 2025).

1. Formulation of the Advantage-Weighted Matching Objective

AWM treats the final sample $x_0$ generation from a prompt $\xi$ as a one-step sequence policy $\pi_\theta(x_0 \mid \xi) = p_\theta(x_0 \mid \xi)$, using the pretraining ELBO/flow-matching loss as a surrogate for the log-likelihood $\log \pi_\theta(x_0 \mid \xi)$. Given a batch of $G$ samples $\{x^i_0\}$ from the current (or reference) policy, rewards $r_i = r(x^i_0, \xi)$, and computed advantages $A_i$, the AWM surrogate objective is:

$$\mathcal J_{\mathrm{AWM}}(\theta) = \mathbb E_{\xi, \{x^i_0\} \sim \pi_{\theta_{\mathrm{old}}}} \left[\frac{1}{G}\sum_{i=1}^G \Big\{ \log \hat \pi_\theta(x^i_0 \mid \xi) A_i - \beta\, \mathrm{KL}(\pi_\theta \| \pi_{\mathrm{ref}}) \Big\} \right]$$

where the log-likelihood surrogate is:

$$\log \hat \pi_\theta(x_0 \mid \xi) = -\mathbb E_{t \sim p(t),\, \epsilon \sim \mathcal N(0, I)} \left[ w(t) \| v_\theta(x_t, t, \xi) - (\epsilon - x_0) \|^2 \right]$$

with $x_t = (1-t) x_0 + t \epsilon$ and $w(t)$ often set to 1. The total minimized loss can be written as:

$$L_{\mathrm{AWM}}(\theta) = \mathbb E_{\xi, x_0, \epsilon, t} \left[ -A(x_0, \xi) w(t) \| v_\theta(x_t, t, \xi) - (\epsilon - x_0) \|^2 \right] + \beta\, \mathbb E_{\xi, x_0, \epsilon, t} \left[ w(t) \| v_\theta(x_t, t, \xi) - v_{\mathrm{ref}}(x_t, t, \xi) \|^2 \right]$$

Typically, $\beta \ll 1$. This construction preserves the model-consistent flow-matching objective while advantage-weighting samples during RL updates (Xue et al., 29 Sep 2025).

2. Policy Gradient Interpretation

AWM yields a policy-gradient-style update under the RL objective, but with the intractable $\log \pi$ term replaced by a tight ELBO-based surrogate:

$$\nabla_\theta \mathcal J_{\mathrm{AWM}} = - \mathbb E_{x_0, \epsilon, t} \left[ A(x_0, \xi) \nabla_\theta \| v_\theta(x_t, t, \xi) - (\epsilon - x_0) \|^2 \right] - \beta \nabla_\theta \mathrm{KL}$$

This connects AWM to the classic REINFORCE gradient, showing that AWM provides unbiased estimates of the true advantage-weighted policy gradient under the sequence policy, using the exact pretraining loss as a log-likelihood proxy. This approach supports the use of off-policy samples via importance weighting or on-policy updates with $\rho^i = 1$ (Xue et al., 29 Sep 2025).

3. Variance Reduction and Convergence Benefits

Theoretical results demonstrate that DDPO, via time-reversal of the diffusion SDE, is equivalent to denoising score-matching (DSM) on $x_t$ conditioned on noisy $x_{t-\Delta t}$, thus using a noisy surrogate for the RL gradient. Theorem 2 in the cited work establishes that gradient estimators based on noisy-DSM (as used in DDPO) have strictly larger covariance than those based on the clean-DSM (pretraining) loss:

$$\mathrm{Cov}\bigl(\nabla\log p(x_t \mid s)\mid x_t\bigr) = \mathrm{Cov}\bigl(\nabla\log p(x_t \mid 0)\mid x_t\bigr) + \kappa(s, t) I, \quad \kappa(s, t) > 0 \ \forall\ s > 0$$

This increased variance in DDPO slows convergence. Empirically, clean DSM achieves a given FID on CIFAR-10 and ImageNet-64 up to 2–4$\times$ faster than the noisy analogue. AWM inherits this lower variance by using clean-DSM as its RL surrogate, resulting in faster and more stable convergence during fine-tuning (Xue et al., 29 Sep 2025).

4. Algorithmic Implementation

The core AWM algorithm iterates as follows:

1. Sample a batch $\{x_0^i\}_{i=1}^G$ from the current diffusion model.
2. Compute rewards $r_i$ and group-relative or other advantage estimates $A_i$.
3. For each $i$, sample $\epsilon^i \sim \mathcal N(0,I)$, $t^i \sim p(t)$, and compose $x^i_t = (1-t^i)x_0^i + t^i \epsilon^i$.
4. Predict velocities $v^i = v_\theta(x^i_t, t^i, \xi)$ and obtain reference $v_{\mathrm{ref}}^i$.
5. For each $i$, compute:
   • Log-likelihood surrogate: $\ell^i = -w(t^i)\|v^i - (\epsilon^i - x_0^i)\|^2$
   • Ratio for importance sampling: $\rho^i = \exp(\ell^i - \ell^i_{\mathrm{old}})$ (on-policy: $\rho^i = 1$)
   • Policy loss: $$\mathcal L_{\mathrm{pg}} = -\frac{1}{G}\sum_i \rho^i A_i$$
   • KL loss: $$\mathcal L_{\mathrm{kl}} = \frac{1}{G}\sum_i w(t^i)\|v^i - v_{\mathrm{ref}}^i\|^2$$
   • Total loss: $$\mathcal L = \mathcal L_{\mathrm{pg}} + \beta \mathcal L_{\mathrm{kl}}$$
6. Update $\theta$ via gradient descent on $\nabla_\theta \mathcal L$

This procedure supports both on-policy and importance-weighted off-policy updates (Xue et al., 29 Sep 2025).

5. Theoretical Consistency and Guarantees

While there is no standalone "AWM optimality theorem," two key statements clarify AWM's theoretical role:

• Lemma 1 asserts that minimizing DSM with noisy or clean conditioning produces the same population minimizer; the clean-DSM surrogate used by AWM does not bias the learned score.
• The REINFORCE derivation ensures that AWM gradients are unbiased estimates of the advantage-weighted policy gradient under the sequence policy, supporting both theoretical consistency and optimality in expectation.

These results establish AWM as a principled formulation for RL with diffusion models, eliminating surrogate bias and aligning the optimization landscape of RL fine-tuning with pretraining (Xue et al., 29 Sep 2025).

6. Empirical Speedups and Sample Efficiency

AWM demonstrates substantial improvements over Flow-GRPO, the DDPO-based baseline, across several benchmarks and model backbones:

Benchmark	Backbone	Flow-GRPO Time (h)	AWM Time (h)	Speedup
GenEval (0.95 score)	SD 3.5 Medium	640	80	8.0×
OCR (0.89 score)	SD 3.5 Medium	416	17.6	23.6×
OCR (0.95 score)	FLUX	343	40	8.5×
PickScore (23.01 score)	SD 3.5 Medium	956	91	10.5×

AWM matches or exceeds the reward performance of Flow-GRPO while reducing wall-clock and sample requirements by one order of magnitude. This confirms the variance-reduction and convergence benefits conferred by the method (Xue et al., 29 Sep 2025).

7. Conceptual Unification of RL and Pretraining in Diffusion

By constructing the RL fine-tuning objective to use the same flow-matching/ELBO loss as diffusion pretraining, AWM unifies the conceptual and practical frameworks of pretraining and RL for diffusion models. This contrasts with prior RL approaches that introduced variance and slowed convergence by departing from the pretraining surrogate. The AWM approach thus bridges the algorithmic gap between supervised pretraining and RL fine-tuning, preserving modeling consistency while enabling effective, low-variance advantage-based policy optimization in high-dimensional generative modeling (Xue et al., 29 Sep 2025).