Decodifica Distillata 1: Campionamento in un Passo dei Modelli Auto-regressivi di Immagini con Corrispondenza di Flusso

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

December 22, 2024

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

cs.AI

Abstract

I modelli autoregressivi (AR) hanno raggiunto prestazioni all'avanguardia nella generazione di testo e immagini, ma soffrono di una generazione lenta a causa del processo token-per-token. Ci poniamo una domanda ambiziosa: un modello AR pre-addestrato può essere adattato per generare output in soli uno o due passaggi? In caso di successo, ciò farebbe avanzare significativamente lo sviluppo e il dispiegamento dei modelli AR. Notiamo che i lavori esistenti che cercano di accelerare la generazione AR generando più token contemporaneamente fondamentalmente non riescono a catturare la distribuzione dell'output a causa delle dipendenze condizionali tra i token, limitando la loro efficacia per la generazione a pochi passaggi. Per affrontare questo problema, proponiamo il Decodifica Distillata (DD), che utilizza il matching di flusso per creare un mapping deterministico dalla distribuzione gaussiana alla distribuzione dell'output del modello AR pre-addestrato. Successivamente addestriamo una rete a distillare questo mapping, consentendo la generazione a pochi passaggi. DD non richiede i dati di addestramento del modello AR originale, rendendolo più pratico. Valutiamo DD sui modelli AR di immagini all'avanguardia e presentiamo risultati promettenti su ImageNet-256. Per VAR, che richiede una generazione a 10 passaggi, DD consente la generazione in un passaggio (aumento della velocità di 6,3 volte), con un aumento accettabile di FID da 4,19 a 9,96. Per LlamaGen, DD riduce la generazione da 256 passaggi a 1, ottenendo un aumento della velocità di 217,8 volte con un aumento FID comparabile da 4,11 a 11,35. In entrambi i casi, i metodi di base falliscono completamente con FID>100. DD eccelle anche nella generazione testo-immagine, riducendo la generazione da 256 passaggi a 2 per LlamaGen con un aumento minimo di FID da 25,70 a 28,95. Essendo il primo lavoro a dimostrare la possibilità di generazione in un passaggio per i modelli AR di immagini, DD sfida l'idea predominante che i modelli AR siano intrinsecamente lenti e apre nuove opportunità per una generazione AR efficiente. Il sito web del progetto è disponibile su 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.