Decodificación Destilada 1: Muestreo de un Paso de Modelos Auto-regresivos de Imagen con Coincidencia de Flujos

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

December 22, 2024

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

cs.AI

Resumen

Los modelos autorregresivos (AR) han logrado un rendimiento de vanguardia en la generación de texto e imágenes, pero sufren de una generación lenta debido al proceso token por token. Planteamos una pregunta ambiciosa: ¿se puede adaptar un modelo AR pre-entrenado para generar salidas en solo uno o dos pasos? En caso de éxito, esto avanzaría significativamente en el desarrollo y despliegue de modelos AR. Observamos que los trabajos existentes que intentan acelerar la generación AR al generar múltiples tokens a la vez no pueden capturar fundamentalmente la distribución de salida debido a las dependencias condicionales entre tokens, lo que limita su efectividad para la generación de pocos pasos. Para abordar esto, proponemos Decodificación Destilada (DD), que utiliza el emparejamiento de flujo para crear un mapeo determinista de una distribución gaussiana a la distribución de salida del modelo AR pre-entrenado. Luego entrenamos una red para destilar este mapeo, permitiendo la generación de pocos pasos. DD no necesita los datos de entrenamiento del modelo AR original, lo que lo hace más práctico. Evaluamos DD en modelos AR de imagen de vanguardia y presentamos resultados prometedores en ImageNet-256. Para VAR, que requiere una generación de 10 pasos, DD permite la generación en un paso (aumento de velocidad de 6.3 veces), con un aumento aceptable en FID de 4.19 a 9.96. Para LlamaGen, DD reduce la generación de 256 pasos a 1, logrando un aumento de velocidad de 217.8 veces con un aumento FID comparable de 4.11 a 11.35. En ambos casos, los métodos de referencia fallan por completo con FID>100. DD también destaca en la generación de texto a imagen, reduciendo la generación de 256 pasos a 2 para LlamaGen con un aumento mínimo de FID de 25.70 a 28.95. Como el primer trabajo que demuestra la posibilidad de generación en un paso para modelos AR de imagen, DD desafía la noción predominante de que los modelos AR son inherentemente lentos y abre nuevas oportunidades para una generación AR eficiente. El sitio web del proyecto se encuentra en 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.