Décodage Distillé 1 : Échantillonnage en une étape des Modèles Auto-Régressifs d'Images avec Appariement de Flux

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

Résumé

Les modèles autorégressifs (AR) ont atteint des performances de pointe dans la génération de texte et d'images, mais souffrent d'une lenteur de génération en raison du processus jeton par jeton. Nous posons une question ambitieuse : est-il possible d'adapter un modèle AR pré-entraîné pour générer des sorties en seulement une ou deux étapes ? En cas de succès, cela permettrait de faire avancer de manière significative le développement et le déploiement des modèles AR. Nous remarquons que les travaux existants qui tentent d'accélérer la génération AR en générant plusieurs jetons à la fois ne parviennent fondamentalement pas à capturer la distribution de sortie en raison des dépendances conditionnelles entre les jetons, ce qui limite leur efficacité pour une génération en quelques étapes. Pour remédier à cela, nous proposons le Décodage Distillé (DD), qui utilise la mise en correspondance des flux pour créer une cartographie déterministe de la distribution gaussienne à la distribution de sortie du modèle AR pré-entraîné. Nous entraînons ensuite un réseau à distiller cette cartographie, permettant une génération en quelques étapes. DD ne nécessite pas les données d'entraînement du modèle AR original, le rendant plus pratique. Nous évaluons DD sur des modèles AR d'images de pointe et présentons des résultats prometteurs sur ImageNet-256. Pour VAR, qui nécessite une génération en 10 étapes, DD permet une génération en une étape (accélération de 6,3 fois), avec une augmentation acceptable du FID de 4,19 à 9,96. Pour LlamaGen, DD réduit la génération de 256 étapes à 1, réalisant une accélération de 217,8 fois avec une augmentation comparable du FID de 4,11 à 11,35. Dans les deux cas, les méthodes de base échouent complètement avec un FID>100. DD excelle également dans la génération de texte vers image, réduisant la génération de 256 étapes à 2 pour LlamaGen avec une augmentation minimale du FID de 25,70 à 28,95. En tant que premier travail à démontrer la possibilité d'une génération en une étape pour les modèles AR d'images, DD remet en question l'idée prédominante selon laquelle les modèles AR sont intrinsèquement lents, et ouvre de nouvelles opportunités pour une génération AR efficace. Le site web du projet se trouve à l'adresse 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.