[paper (arXiv)] [code] [website]
Authors: Enshu Liu (Tsinghua)*, Xuefei Ning (Tsinghua), Yu Wang (Tsinghua), Zinan Lin (Microsoft Research)†
-
* Work mostly done during Enshu Liu's internship at Microsoft Research
-
† Project advisor: Zinan Lin
Abstract: Autoregressive (AR) models have recently achieved state-of-the-art performance in text and image generation. However, their primary limitation is slow generation speed 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 attempt 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 overcome this, we propose Distilled Decoding (DD), which leverages 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. The entire training process of DD does not need the training data of the original AR model (as opposed to some other methods), thus making DD more practical.
We evaluate DD on state-of-the-art image AR models and present promising results. For VAR, which requires 10-step generation (680 tokens), DD enables one-step generation (6.3x speed-up), with an acceptable increase in FID from 4.19 to 9.96 on ImageNet. Similarly, for LlamaGen, DD reduces generation from 256 steps to 1, achieving an 217.8x speed-up with a comparable FID increase from 4.11 to 11.35 on ImageNet. In both cases, baseline methods completely fail with FID scores >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.