PaperlayerView exact revision on arXiv

Source-linked AI summary

Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models

Shuchen Xue, Chongjian Ge, Shilong Zhang, Yichen Li, Zhi-Ming Ma

arXiv:2509.25050v2cs.LG

TL;DR

Diffusion RL post-training uses an objective different from pretraining, and DDPO's noisy-target formulation increases variance and slows convergence. The paper introduces AWM, which applies advantage weighting to the pretraining score/flow-matching loss. AWM achieves up to 34× faster convergence than Flow-GRPO on SD3.5M and FLUX without degrading generation quality.

  • Problem

    Diffusion RL methods such as DDPO use an objective different from the score/flow-matching objective used in diffusion pretraining, unlike the aligned objectives in LLM training.

  • Method

    AWM applies advantage-based sample weights to the same score/flow-matching loss used during diffusion pretraining, without relying on reverse-time discretization.

  • Results

    Up to 34× faster convergence than Flow-GRPO is achieved on SD3.5M and FLUX without degrading generation quality.

  • Takeaways & Limitations

    AWM unifies diffusion pretraining and RL post-training under a single policy-gradient-consistent score/flow-matching objective.

  • Takeaways & Limitations

    DDPO's practical reverse-process discretization introduces additional bias, whereas the theoretical equivalence is characterized up to discretization error.

Abstract

from arXiv · show

Reinforcement Learning (RL) has emerged as a central paradigm for advancing Large Language Models (LLMs), where both pre-training and RL post-training stages are grounded in the same log-likelihood formulation. In contrast, recent RL approaches for diffusion models, most notably Denoising Diffusion Policy Optimization (DDPO), optimize an objective different from the pretraining objectives--score/flow matching loss. In this work, we establish a novel theoretical analysis: DDPO is an implicit form of score/flow matching with noisy targets, which increases variance and slows convergence. Building on this analysis, we introduce Advantage Weighted Matching (AWM), a policy-gradient method for diffusion. It uses the score/flow-matching loss and reweights each sample by its advantage. In effect, AWM raises the influence of high-reward samples and suppresses low-reward ones while keeping the modeling objective identical to pretraining. This simple yet effective design yields substantial benefits: on the GenEval, OCR, and PickScore benchmarks, AWM delivers up to a $34\times$ speedup over Flow-GRPO (which builds on DDPO), when applied to Stable Diffusion 3.5 Medium and FLUX, without compromising generation quality. Code is available at https://github.com/scxue/advantage_weighted_matching

1. Introduction

Diffusion RL post-training diverges from pretraining's score/flow-matching objective, and DDPO implicitly uses noisy-target DSM with higher variance. AWM restores objective alignment through advantage-weighted matching and achieves faster training without quality degradation.

  • 1. Introduction: Diffusion pretraining uses score/flow matching on forward-process pairs, whereas DDPO uses per-step reverse-transition likelihoods on noisy backward-process pairs.This objective mismatch contrasts with LLMs, whose pretraining and RL post-training share a log-likelihood objective.
  • 1. Introduction: AWM reweights the pretraining score/flow-matching loss by sample advantages, amplifying high-reward samples while preserving the original modeling objective.Unlike DDPO, AWM does not rely on reverse-time discretization and supports any sampler.
  • 1. Introduction: DDPO is theoretically equivalent, up to discretization error, to denoising score matching with noisy data targets, regardless of score or velocity parameterization.This connects DDPO's per-step Gaussian likelihood to noisy-target DSM.
  • 1. Introduction: Noisy-target DSM has higher estimator variance than clean-target DSM and converges more slowly in controlled CIFAR-10 and ImageNet-64 experiments.The experiments validate the theoretical variance analysis under identical settings.
  • 1. Introduction: Up to 34× faster training than Flow-GRPO is achieved by AWM on GenEval, OCR, ImageReward, and PickScore across SD3.5M and FLUX without generation quality degradation.AWM also surpasses Flow-GRPO on GenEval, scoring 0.98 versus 0.95 with a 16.54× speed-up on SD3.5M.

2. DDPO Is Secretly Doing Score Matching

DDPO’s reverse-transition likelihood is equivalent at the population level to denoising score matching with noisy targets, which share clean matching’s minimizers but have higher variance and slower optimization.

  • 2. DDPO Is Secretly Doing Score Matching: DDPO’s per-step reverse-transition likelihood induces denoising score matching on forward-process pairs with noisy targets, despite using a reverse-process formulation.The equivalence holds for the common joint distribution of adjacent diffusion states, up to discretization considerations.
  • 2. DDPO Is Secretly Doing Score Matching: In practice, DDPO’s discretized reverse-time transitions introduce additional bias, whereas the forward diffusion process has an exact transition.The discrepancy vanishes as the discretization step approaches zero, but the practical formulation still uses discretized reverse transitions.
  • 2. DDPO Is Secretly Doing Score Matching: The noisy-target and clean-target DSM objectives share the same population minimizer, so noisy targets do not change the population optimum.The equivalence is established at the population-objective level.
  • 2. DDPO Is Secretly Doing Score Matching: Locally at a frozen rollout policy, the advantage-weighted DDPO gradient equals a positively rescaled signed advantage-weighted clean-target DSM gradient.This local-gradient relation motivates using the clean score-matching objective while retaining advantage weighting.
  • 2. DDPO Is Secretly Doing Score Matching: Noisy DSM has strictly larger conditional covariance by d κ(s, t), increasing objective variance and slowing optimization.The added covariance term is d κ(s, t), with κ increasing in the earlier noise level and κ(0, t)=0.
  • 2. DDPO Is Secretly Doing Score Matching: Under identical architectures, noise schedules, and optimization settings, noisy DSM reaches a given quality level substantially later than the clean-data objective.The comparison is reported for the noisy-DSM and clean-data objectives under matched training conditions.

3. Advantage Weighted Matching

AWM aligns diffusion-model RL post-training with pretraining by applying advantage weights directly to the score/flow-matching objective. Its forward-process formulation emphasizes high-reward samples, supports flexible sampling, and reduces variance relative to DDPO.

  • 3. Advantage Weighted Matching: AWM replaces DDPO’s reverse-time likelihood objective with advantage-weighted score/flow matching, preserving the pretraining loss while incorporating rewards.High-advantage samples receive greater influence, while low-advantage samples are downweighted.
  • 3. Advantage Weighted Matching: AWM decouples training from sampling, allowing inference with any ODE or SDE sampler rather than sampler-specific trajectories.This flexibility follows from operating on final samples and the forward process.
  • 3. Advantage Weighted Matching: For each prompt, AWM samples a group, computes rewards and advantages, evaluates forward-process matching losses, and optimizes an advantage-weighted objective with velocity-space KL regularization.The implementation uses group-relative advantages and a frozen or periodically updated reference policy for KL control.
  • 3. Advantage Weighted Matching: The method motivates variance reduction by showing that DDPO implicitly performs noisy-target DSM, whereas AWM retains the pretraining score/flow-matching objective.The paper presents this alignment as restoring conceptual symmetry between diffusion pretraining and RL post-training.

4. Experiments

Experiments evaluate AWM on SD3.5M and FLUX across GenEval, OCR, and PickScore, finding faster training with preserved or improved generation quality. Ablations further examine Monte Carlo estimation, rollout policies, KL regularization, and sampler choice.

  • 4.1. Main Results: AWM reaches a final GenEval score of 0.985 on SD3.5M, exceeding Flow-GRPO’s 0.95 while matching that baseline 16.54× faster.It also reports strong subtask scores including Color Attribution 0.94 and Counting 0.99.
  • 4.1. Main Results: AWM achieves OCR and PickScore gains with fewer GPU hours across SD3.5M and FLUX, including 34.67× faster SD3.5M PickScore and 6.82× faster FLUX PickScore.On SD3.5M OCR, AWM reaches 0.96 using 58 GPU hours, while Flow-GRPO reaches 0.90 after over 400 GPU hours.
  • 4.1. Main Results: After 100 AWM gradient steps, FLUX better follows composition constraints and improves text rendering on the visual prompts.The examples test numerosity, color, position, and text-rendering instructions.
  • 4.2. Ablations: Using T = 6 Monte Carlo steps gives a stronger final GenEval score than T = 4, so T = 6 is adopted by default.Both settings learn stably.
  • 4.2. Ablations: β = 3 remains stable but slows learning and underperforms β = 1, while ODE and SDE sampling produce highly similar reward performance.The sampler comparison covers OCR and ImageReward objectives across different numbers of steps.

5. Conclusion

The paper concludes that AWM unifies diffusion pretraining and RL post-training through advantage-weighted score/flow matching. This produces faster convergence than Flow-GRPO without degrading generation quality.

  • 5. Conclusion: AWM keeps the pretraining score/flow-matching objective while applying advantage-based weights, unifying diffusion pretraining and RL post-training.The conclusion attributes the approach to the analysis of DDPO’s noisy-target matching behavior.
  • 5. Conclusion: AWM achieves up to 34× faster convergence than Flow-GRPO on SD3.5M and FLUX without degrading generation quality.

Impact Statement

The paper positions its analysis and AWM framework as a basis for more stable and efficient diffusion-model training under varied rewards and sampling procedures. It also links improved efficiency to reduced computational cost.

  • Impact Statement: The DDPO–DSM analysis and AWM framework may enable more stable and efficient diffusion-model training across diverse reward functions and sampling procedures.
  • Impact Statement: Improved training efficiency can help reduce computational cost, according to the paper’s stated impact rationale.

C. Proofs

The proof shows that the gradient of Denoising Score Matching equals the gradient of standard Score Matching with respect to model parameters.

  • DSM and standard Score Matching have identical parameter gradients.The proof establishes this equivalence by differentiating the Denoising Score Matching loss with respect to θ.

C.2. Proof of Theorem 2.1

This proof connects DDPO’s reverse-process likelihood to noisy-target DSM through diffusion time reversal and forward–reverse joint-distribution identities.

  • Forward and reverse diffusion processes share the same marginal and joint distributions, enabling the objective connection.The joint-distribution identity permits forward-process score matching to represent the reverse-process DDPO objective.
  • DDPO’s per-step likelihood is equivalent to a noisy-target DSM loss.The equivalence follows after expressing the reverse-step likelihood under discretization and matching it to the noisy DSM formulation.
  • The advantage-weighted DDPO gradient maps locally to an advantage-weighted clean-target DSM or flow-matching gradient.This relation holds at the frozen rollout policy and does not require the advantage to be independent of x_t.

C.4. Proof of Theorem 2.3

The theorem compares noisy and clean conditional score targets, showing equal conditional means but strictly larger noisy-target variance for positive intermediate noise.

  • Noisy and clean conditional scores are both unbiased estimators of the true score ∇x_t log p(x_t).Their conditional expectations coincide with the marginal score for every s in [0,t).
  • Conditional covariance increases by κ(s,t)I for the noisy score relative to the clean score.The trace likewise increases by dκ(s,t), so noisy targets have greater conditional variance whenever κ(s,t)>0.
  • κ(s,t) increases strictly with s, equals zero at s=0, and diverges as s approaches t.Thus the variance penalty grows with the separation between the noisy target and the clean target.

F.1. Experiment Settings

The experiments use separate LoRA configurations for SD3.5M and FLUX, with later experiments adopting a different weighting, sampler, timestep, guidance, and reference-model setup.

  • LoRA uses α=64 and r=32 for SD3.5M, versus α=128 and r=64 for FLUX.The learning rate is fixed at 3e−4, and group size is G=24.
  • Later experiments use w(t)=t, SA-Solver with η=0.4, 14 timesteps, and guidance scale cfg=1.0.They estimate the ELBO with 6 Monte Carlo samples and use an EMA reference model.
  • The main-paper experiments use the later configuration unless otherwise specified.

F.2. Additional Experiments

Additional experiments report benchmark metrics, speed comparisons, training curves, and variance measurements across diffusion backbones, datasets, and noise levels.

  • Additional results under an earlier configuration provide full training curves and final metrics across all benchmarks.They complement the ablation studies by examining an alternative training setting.
  • Table 5 compares OCR and PickScore for SD3.5M and FLUX alongside total GPU hours and relative speed-ups.These comparisons evaluate AWM across different backbones and metrics.
  • Table 4 reports per-category and overall GenEval performance for SD3.5M.
  • Figure 9 plots OCR, PickScore, and ImageReward training curves against GPU hours on SD3.5M.
  • The noisy-target formulation has substantially larger conditional target variance than clean-target DSM across measured noise levels.The comparison uses the EDM noise schedule and spans CIFAR-10 and ImageNet-64.

G. Additional Samples

Additional samples show that the method preserves broad mode coverage after 600 RL steps, maintaining visual diversity across seeds for identical prompts.

  • Four samples from identical prompts remain visually diverse across different seeds after 600 RL optimization steps.The samples cover prompts ranging from twilight scenes and animals to an electric drill and a lion.
Loading 2509.25050v2…