Renforcement de générateurs en quelques étapes via l'appariement de distribution incliné par la récompense

Résumé

Les récentes avancées en distillation de diffusion en peu d’étapes ont permis une génération d’images efficace, mais aligner ces modèles sur les préférences humaines reste un défi. Nous proposons Reward-Tilted Distribution Matching Distillation (RTDMD), un cadre en deux étapes qui unifie la distillation par appariement de distributions avec l’apprentissage par renforcement guidé par récompense pour les générateurs de flux en peu d’étapes. Nous montrons que la minimisation de la divergence KL par rapport à une distribution enseignante inclinée par récompense se décompose naturellement en un terme d’appariement de distributions et un terme de maximisation de récompense. Dans la première étape, nous introduisons Ambient-Consistent Distribution Matching Distillation (AC-DMD), qui effectue un appariement de distributions par sous-intervalles et augmente l’objectif de score factice avec un régularisateur de cohérence pour aider le modèle de score factice à suivre la distribution du générateur en évolution sous des mises à jour limitées. Dans la deuxième étape, nous optimisons conjointement les deux termes : pour le terme de maximisation de récompense, nous dérivons un gradient de politique hybride qui combine un estimateur de type GRPO pour les transitions intermédiaires stochastiques avec une rétropropagation directe de la récompense à travers l’étape finale déterministe, et nous introduisons en outre le GRPO par sous-ensemble d’étapes (SubGRPO) pour réduire la variance. Les expériences sur SD3, SD3.5 et FLUX.2 montrent que RTDMD établit de nouveaux résultats de pointe sur les métriques de préférence, esthétiques et compositionnelles avec seulement 4 étapes d’inférence, surpassant les précédentes méthodes de génération texte-image en peu d’étapes. Le code et les modèles sont disponibles à l’adresse 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.