Verstärkung von Generatoren mit wenigen Schritten durch belohnungsgeneigte Verteilungsanpassung

Zusammenfassung

Jüngste Fortschritte bei der Diffusionsdestillation mit wenigen Schritten ermöglichen eine effiziente Bildgenerierung, doch die Ausrichtung dieser Modelle auf menschliche Präferenzen bleibt eine Herausforderung. Wir schlagen die Reward-Tilted Distribution Matching Distillation (RTDMD) vor, ein zweistufiges Framework, das die Verteilungsanpassungsdestillation mit belohnungsgesteuertem bestärkendem Lernen für Flussgeneratoren mit wenigen Schritten vereint. Wir zeigen, dass die Minimierung der KL-Divergenz zu einer belohnungsgeneigten Lehrer-Verteilung auf natürliche Weise in einen Verteilungsanpassungsterm und einen Belohnungsmaximierungsterm zerfällt. In der ersten Stufe führen wir die Ambient-Consistent Distribution Matching Distillation (AC-DMD) ein, die eine unterintervallweise Verteilungsanpassung durchführt und die Fake-Score-Zielfunktion um einen Konsistenzregularisierer erweitert, damit das Fake-Score-Modell die sich unter begrenzten Aktualisierungen verschiebende Generatorverteilung verfolgen kann. In der zweiten Stufe optimieren wir beide Terme gemeinsam: Für den Belohnungsmaximierungsterm leiten wir einen hybriden Policy-Gradienten her, der einen GRPO-artigen Schätzer für die stochastischen Zwischenübergänge mit direkter Belohnungsrückpropagation durch den deterministischen letzten Schritt kombiniert, und führen zudem die Schritt-Teilmenge-GRPO (SubGRPO) ein, um die Varianz zu verringern. Experimente mit SD3, SD3.5 und FLUX.2 zeigen, dass RTDMD mit nur 4 Inferenzschritten neue Spitzenergebnisse in Bezug auf Präferenz-, Ästhetik- und Kompositionsmetriken erzielt und damit frühere Methoden zur Text-zu-Bild-Generierung mit wenigen Schritten übertrifft. Code und Modelle sind unter https://github.com/Harahan/RTDMD verfügbar.

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.