Destilliertes Decodieren 1: Ein-Schritt-Sampling von Bild auto-regressiven Modellen mit Flussanpassung

Distilled Decoding 1: One-step Sampling of Image Auto-regressive Models with Flow Matching

December 22, 2024

Autoren: Enshu Liu, Xuefei Ning, Yu Wang, Zinan Lin

cs.AI

Zusammenfassung

Autoregressive (AR) Modelle haben Spitzenleistungen bei der Text- und Bildgenerierung erzielt, leiden jedoch unter einer langsamen Generierung aufgrund des tokenbasierten Prozesses. Wir stellen eine ehrgeizige Frage: Kann ein vortrainiertes AR-Modell angepasst werden, um Ausgaben in nur einem oder zwei Schritten zu generieren? Bei Erfolg würde dies die Entwicklung und Bereitstellung von AR-Modellen erheblich vorantreiben. Wir stellen fest, dass bestehende Arbeiten, die versuchen, die AR-Generierung zu beschleunigen, indem sie mehrere Tokens gleichzeitig generieren, grundlegend nicht die Ausgabeverteilung erfassen können, aufgrund der bedingten Abhängigkeiten zwischen den Tokens, was ihre Wirksamkeit für die Generierung in wenigen Schritten einschränkt. Um dies anzugehen, schlagen wir Distilled Decoding (DD) vor, das Flussabgleich verwendet, um eine deterministische Zuordnung von einer Gaußverteilung zur Ausgabeverteilung des vortrainierten AR-Modells zu erstellen. Anschließend trainieren wir ein Netzwerk, um diese Zuordnung zu destillieren, was die Generierung in wenigen Schritten ermöglicht. DD benötigt nicht die Trainingsdaten des ursprünglichen AR-Modells, was es praktischer macht. Wir evaluieren DD an Spitzenbild-AR-Modellen und präsentieren vielversprechende Ergebnisse auf ImageNet-256. Für VAR, das eine Generierung in 10 Schritten erfordert, ermöglicht DD eine Generierung in einem Schritt (6,3-fache Beschleunigung) mit einer akzeptablen Zunahme des FID von 4,19 auf 9,96. Für LlamaGen reduziert DD die Generierung von 256 Schritten auf 1, was eine 217,8-fache Beschleunigung mit einer vergleichbaren FID-Zunahme von 4,11 auf 11,35 erreicht. In beiden Fällen versagen Basismethoden vollständig mit FID>100. DD zeichnet sich auch bei der Text-zu-Bild-Generierung aus, indem es die Generierung von 256 Schritten auf 2 für LlamaGen reduziert, mit einer minimalen FID-Zunahme von 25,70 auf 28,95. Als erste Arbeit, die die Möglichkeit der Ein-Schritt-Generierung für Bild-AR-Modelle demonstriert, fordert DD die vorherrschende Vorstellung heraus, dass AR-Modelle grundsätzlich langsam sind, und eröffnet neue Möglichkeiten für eine effiziente AR-Generierung. Die Projektwebsite befindet sich unter https://imagination-research.github.io/distilled-decoding.

English

Autoregressive (AR) models have achieved state-of-the-art performance in text and image generation but suffer from slow generation due to the token-by-token process. We ask an ambitious question: can a pre-trained AR model be adapted to generate outputs in just one or two steps? If successful, this would significantly advance the development and deployment of AR models. We notice that existing works that try to speed up AR generation by generating multiple tokens at once fundamentally cannot capture the output distribution due to the conditional dependencies between tokens, limiting their effectiveness for few-step generation. To address this, we propose Distilled Decoding (DD), which uses flow matching to create a deterministic mapping from Gaussian distribution to the output distribution of the pre-trained AR model. We then train a network to distill this mapping, enabling few-step generation. DD doesn't need the training data of the original AR model, making it more practical.We evaluate DD on state-of-the-art image AR models and present promising results on ImageNet-256. For VAR, which requires 10-step generation, DD enables one-step generation (6.3times speed-up), with an acceptable increase in FID from 4.19 to 9.96. For LlamaGen, DD reduces generation from 256 steps to 1, achieving an 217.8times speed-up with a comparable FID increase from 4.11 to 11.35. In both cases, baseline methods completely fail with FID>100. DD also excels on text-to-image generation, reducing the generation from 256 steps to 2 for LlamaGen with minimal FID increase from 25.70 to 28.95. As the first work to demonstrate the possibility of one-step generation for image AR models, DD challenges the prevailing notion that AR models are inherently slow, and opens up new opportunities for efficient AR generation. The project website is at https://imagination-research.github.io/distilled-decoding.