Сопоставление с весами преимуществ: согласование обучения с подкреплением и предварительного обучения в диффузионных моделях
Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models
September 29, 2025
Авторы: Shuchen Xue, Chongjian Ge, Shilong Zhang, Yichen Li, Zhi-Ming Ma
cs.AI
Аннотация
Обучение с подкреплением (Reinforcement Learning, RL) стало ключевой парадигмой для развития больших языковых моделей (Large Language Models, LLMs), где предварительное обучение и последующее обучение с подкреплением используют одну и ту же формулировку логарифмического правдоподобия. В отличие от этого, современные подходы RL для диффузионных моделей, в частности Оптимизация политики денойзинга диффузии (Denoising Diffusion Policy Optimization, DDPO), оптимизируют целевую функцию, отличную от целей предварительного обучения — потерю согласования оценок/потоков. В данной работе мы предлагаем новый теоретический анализ: DDPO является неявной формой согласования оценок/потоков с зашумленными целями, что увеличивает дисперсию и замедляет сходимость. На основе этого анализа мы представляем метод Согласование с взвешиванием преимуществ (Advantage Weighted Matching, AWM), который использует ту же потерю согласования оценок/потоков, что и предварительное обучение, для получения целевой функции с меньшей дисперсией и перевзвешивает каждый образец по его преимуществу. По сути, AWM повышает влияние образцов с высокой наградой и подавляет образцы с низкой наградой, сохраняя при этом целевую функцию моделирования идентичной предварительному обучению. Это объединяет предварительное обучение и RL как концептуально, так и практически, согласуется с теорией градиента политики, снижает дисперсию и обеспечивает более быструю сходимость. Этот простой, но эффективный подход приносит значительные преимущества: на бенчмарках GenEval, OCR и PickScore AWM обеспечивает ускорение до 24 раз по сравнению с Flow-GRPO (который основан на DDPO) при применении к Stable Diffusion 3.5 Medium и FLUX, без ущерба для качества генерации. Код доступен по адресу https://github.com/scxue/advantage_weighted_matching.
English
Reinforcement Learning (RL) has emerged as a central paradigm for advancing
Large Language Models (LLMs), where pre-training and RL post-training share 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 same score/flow-matching loss as pretraining to obtain a
lower-variance objective 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 unifies
pretraining and RL conceptually and practically, is consistent with
policy-gradient theory, reduces variance, and yields faster convergence. This
simple yet effective design yields substantial benefits: on GenEval, OCR, and
PickScore benchmarks, AWM delivers up to a 24times 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.