Versterking van generatoren met weinig stappen via beloningsgekantelde distributie-matching

Samenvatting

Recente vooruitgang in weinig-staps diffusiedestillatie heeft efficiënte beeldgeneratie mogelijk gemaakt, maar het afstemmen van deze modellen op menselijke voorkeuren blijft uitdagend. Wij stellen Beloning-gekantelde Distributie-Matching-Destillatie (RTDMD) voor, een tweefasenraamwerk dat distributie-matching-destillatie verenigt met beloning-gestuurd versterkend leren voor weinig-staps flowgeneratoren. We laten zien dat het minimaliseren van de KL-divergentie naar een beloning-gekantelde docentdistributie op natuurlijke wijze uiteenvalt in een distributie-matchingterm en een beloningsmaximalisatieterm. In de eerste fase introduceren we Omgevingsconsistente Distributie-Matching-Destillatie (AC-DMD), die subinterval-gewijze distributie-matching uitvoert en de nepscore-doelstelling aanvult met een consistentieregularisator om het nepscoremodel te helpen de verschuivende generatorverdeling te volgen onder beperkte updates. In de tweede fase optimaliseren we beide termen gezamenlijk: voor de beloningsmaximalisatieterm leiden we een hybride beleidsgradiënt af die een GRPO-achtige schatter voor de stochastische tussentijdse overgangen combineert met directe beloningsterugpropagatie door de deterministische laatste stap, en introduceren we verder stap-subset GRPO (SubGRPO) om de variantie te reduceren. Experimenten op SD3, SD3.5 en FLUX.2 tonen aan dat RTDMD nieuwe state-of-the-art resultaten behaalt op preferentie-, esthetische en compositiemetrieken met slechts 4 inferentiestappen, waarmee het eerdere weinig-staps tekst-naar-beeld generatiemethoden overtreft. Code en modellen zijn beschikbaar op 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.