Усиление генераторов с малым числом шагов посредством согласования распределений со смещением по вознаграждению

Аннотация

Недавние достижения в области многократной дистилляции диффузии позволили эффективно генерировать изображения, однако согласование этих моделей с человеческими предпочтениями остается сложной задачей. Мы предлагаем Reward-Tilted Distribution Matching Distillation (RTDMD) — двухэтапную структуру, объединяющую дистилляцию распределения с подкрепляющим обучением, направленным на вознаграждение, для генераторов с малым числом шагов. Мы показываем, что минимизация расхождения Кульбака–Лейблера до учительского распределения, смещенного в сторону вознаграждения, естественным образом распадается на член согласования распределения и член максимизации вознаграждения. На первом этапе мы вводим Ambient-Consistent Distribution Matching Distillation (AC-DMD), которая выполняет согласование распределения по подынтервалам и дополняет целевую функцию ложных оценок регуляризатором согласованности, помогающим модели ложных оценок отслеживать изменяющееся распределение генератора при ограниченных обновлениях. На втором этапе мы совместно оптимизируем оба члена: для члена максимизации вознаграждения мы выводим гибридный градиент политики, сочетающий оценщик типа GRPO для стохастических промежуточных переходов с прямым обратным распространением вознаграждения через детерминированный финальный шаг, а также вводим GRPO по подмножествам шагов (SubGRPO) для уменьшения дисперсии. Эксперименты на SD3, SD3.5 и FLUX.2 показывают, что RTDMD устанавливает новые современные результаты по метрикам предпочтений, эстетики и композиции, используя всего 4 шага вывода, превосходя предыдущие методы генерации текста в изображение с малым числом шагов. Код и модели доступны по адресу https://github.com/Harahan/RTDMD.

English

Recent advances in few-step diffusion distillation have enabled efficient image generation, yet aligning these models with human preferences remains challenging. We propose Reward-Tilted Distribution Matching Distillation (RTDMD), a two-stage framework that unifies distribution matching distillation with reward-guided reinforcement learning for few-step flow generators. We show that minimizing the KL divergence to a reward-tilted teacher distribution naturally decomposes into a distribution matching term and a reward maximization term. In the first stage, we introduce Ambient-Consistent Distribution Matching Distillation (AC-DMD), which performs subinterval-wise distribution matching and augments the fake score objective with a consistency regularizer to help the fake score model track the shifting generator distribution under limited updates. In the second stage, we jointly optimize both terms: for the reward maximization term, we derive a hybrid policy gradient that combines a GRPO-style estimator for the stochastic intermediate transitions with direct reward backpropagation through the deterministic final step, and further introduce step-subset GRPO (SubGRPO) to reduce variance. Experiments on SD3, SD3.5, and FLUX.2 demonstrate that RTDMD establishes new state-of-the-art results across preference, aesthetic, and compositional metrics with only 4 inference steps, outperforming previous few-step text-to-image generation methods. Code and models are available at https://github.com/Harahan/RTDMD.