Decodificação Destilada 1: Amostragem de um Passo de Modelos Auto-regressivos de Imagens com Correspondência de Fluxo

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

Resumo

Os modelos autorregressivos (AR) alcançaram um desempenho de ponta na geração de texto e imagem, mas sofrem com uma geração lenta devido ao processo token a token. Fazemos uma pergunta ambiciosa: um modelo AR pré-treinado pode ser adaptado para gerar saídas em apenas um ou dois passos? Se bem-sucedido, isso avançaria significativamente no desenvolvimento e implementação de modelos AR. Observamos que trabalhos existentes que tentam acelerar a geração AR gerando vários tokens de uma vez não conseguem capturar a distribuição de saída devido às dependências condicionais entre tokens, limitando sua eficácia para geração de poucos passos. Para abordar isso, propomos o Decodificador Destilado (DD), que utiliza correspondência de fluxo para criar um mapeamento determinístico da distribuição gaussiana para a distribuição de saída do modelo AR pré-treinado. Em seguida, treinamos uma rede para destilar esse mapeamento, possibilitando a geração de poucos passos. O DD não necessita dos dados de treinamento do modelo AR original, tornando-o mais prático. Avaliamos o DD em modelos AR de imagem de ponta e apresentamos resultados promissores no ImageNet-256. Para VAR, que requer geração de 10 passos, o DD possibilita a geração em um passo (aumento de velocidade de 6,3 vezes), com um aumento aceitável no FID de 4,19 para 9,96. Para LlamaGen, o DD reduz a geração de 256 passos para 1, alcançando um aumento de velocidade de 217,8 vezes com um aumento comparável no FID de 4,11 para 11,35. Em ambos os casos, os métodos de referência falham completamente com FID>100. O DD também se destaca na geração de texto para imagem, reduzindo a geração de 256 passos para 2 para LlamaGen com um aumento mínimo no FID de 25,70 para 28,95. Como o primeiro trabalho a demonstrar a possibilidade de geração em um passo para modelos AR de imagem, o DD desafia a noção predominante de que os modelos AR são inerentemente lentos e abre novas oportunidades para uma geração AR eficiente. O site do projeto está em 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.