让扩散模型强化学习更贴近预训练目标，提升采样效率34倍。
Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models
- 用优势加权重采样，使高奖励样本影响更大。
- 在多个基准上实现最高34倍加速，生成质量不变。
- 适合希望高效优化扩散模型策略的研究者。
强化学习（RL）已成为推动大语言模型发展的核心范式，其预训练与后训练阶段均基于相同的对数似然框架。然而，当前针对扩散模型的RL方法（如去噪扩散策略优化，DDPO）优化的目标与预训练目标不同——即分数/流匹配损失。本文提出理论分析：DDPO本质上是带有噪声目标的分数/流匹配的隐式形式，导致方差增大、收敛变慢。基于此，我们提出优势加权匹配（AWM），一种用于扩散模型的策略梯度方法。该方法采用分数/流匹配损失，并按优势值重新加权每个样本。实际效果是增强高奖励样本的影响，抑制低奖励样本，同时保持建模目标与预训练一致。这一简单而有效的设计带来显著收益：在GenEval、OCR和PickScore基准上，应用于Stable Diffusion 3.5 Medium和FLUX时，相比基于DDPO的Flow-GRPO，AWM实现了最高达34倍的加速，且不牺牲生成质量。代码已开源。
原文摘要 · Abstract (English)
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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。