Gedistilleerd decoderen 1: Eénstaps monstername van beeld autoregressieve modellen met stroomafstemming.

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

December 22, 2024

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

cs.AI

Samenvatting

Autoregressieve (AR) modellen hebben state-of-the-art prestaties behaald op het gebied van tekst- en beeldgeneratie, maar kampen met langzame generatie vanwege het token-voor-token proces. We stellen een ambitieuze vraag: kan een vooraf getraind AR-model worden aangepast om uitvoer te genereren in slechts één of twee stappen? Indien succesvol, zou dit de ontwikkeling en implementatie van AR-modellen aanzienlijk bevorderen. We merken op dat bestaande werken die proberen AR-generatie te versnellen door meerdere tokens tegelijk te genereren, fundamenteel niet in staat zijn om de uitvoerverdeling vast te leggen vanwege de conditionele afhankelijkheden tussen tokens, wat hun effectiviteit beperkt voor generatie in enkele stappen. Om dit aan te pakken, stellen we Distilled Decoding (DD) voor, dat flow-matching gebruikt om een deterministische mapping te creëren van de Gauss-verdeling naar de uitvoerverdeling van het vooraf getrainde AR-model. Vervolgens trainen we een netwerk om deze mapping te destilleren, waardoor generatie in enkele stappen mogelijk wordt. DD heeft geen trainingsdata van het oorspronkelijke AR-model nodig, waardoor het praktischer is. We evalueren DD op state-of-the-art beeld-AR-modellen en presenteren veelbelovende resultaten op ImageNet-256. Voor VAR, dat 10-staps generatie vereist, maakt DD éénstaps generatie mogelijk (6,3 keer versnelling), met een acceptabele toename in FID van 4,19 naar 9,96. Voor LlamaGen vermindert DD de generatie van 256 stappen naar 1, wat resulteert in een 217,8 keer versnelling met een vergelijkbare FID-toename van 4,11 naar 11,35. In beide gevallen falen basismethoden volledig met FID>100. DD excelleert ook bij tekst-naar-beeldgeneratie, waarbij de generatie van 256 stappen naar 2 wordt teruggebracht voor LlamaGen met een minimale FID-toename van 25,70 naar 28,95. Als het eerste werk dat de mogelijkheid van éénstaps generatie voor beeld-AR-modellen aantoont, daagt DD de heersende opvatting uit dat AR-modellen inherent langzaam zijn, en opent het nieuwe mogelijkheden voor efficiënte AR-generatie. De projectwebsite is te vinden op 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.