Distilled Decoding 2: One-step Sampling of Image Auto-regressive Models with Conditional Score Distillation

TL;DR

The proposed Distilled Decoding 2 (DD2) takes a significant step toward the goal of one-step AR generation, opening up new possibilities for fast and high-quality AR modeling.

Abstract

Image Auto-regressive (AR) models have emerged as a powerful paradigm of visual generative models. Despite their promising performance, they suffer from slow generation speed due to the large number of sampling steps required. Although Distilled Decoding 1 (DD1) was recently proposed to enable few-step sampling for image AR models, it still incurs significant performance degradation in the one-step setting, and relies on a pre-defined mapping that limits its flexibility. In this work, we propose a new method, Distilled Decoding 2 (DD2), to further advances the feasibility of one-step sampling for image AR models. Unlike DD1, DD2 does not without rely on a pre-defined mapping. We view the original AR model as a teacher model which provides the ground truth conditional score in the latent embedding space at each token position. Based on this, we propose a novel \emph{conditional score distillation loss} to train a one-step generator. Specifically, we train a separate network to predict the conditional score of the generated distribution and apply score distillation at every token position conditioned on previous tokens. Experimental results show that DD2 enables one-step sampling for image AR models with an minimal FID increase from 3.40 to 5.43 on ImageNet-256. Compared to the strongest baseline DD1, DD2 reduces the gap between the one-step sampling and original AR model by 67%, with up to 12.3$\times$ training speed-up simultaneously. DD2 takes a significant step toward the goal of one-step AR generation, opening up new possibilities for fast and high-quality AR modeling. Code is available at https://github.com/imagination-research/Distilled-Decoding-2.

Figures (12)

• Figure 1: Our goal is to distill a multi-step AR model in to a one-step generator while keeping its distribution.
• Figure 2: Comparison of DD2 models, DD1 models, pre-trained models, and other acceleration methods for pre-trained models. DD2 achieves a significant speedup compared to pre-trained models while outperforming DD1 by a large margin. Other methods fail to achieve one-step sampling. For DD2, DD1, and the pre-trained model, each point corresponds to a different model size, whereas for the skip-last method, each point corresponds to a different number of skipped final steps.
• Figure 3: Training process using CSD loss. For the generator, the teacher AR model and the guidance network give the true and fake conditional score of each noisy token based on all previous clean tokens, respectively. Then the true and fake conditional score are used to calculate the score distillation loss, which produces gradient to train the generator. The guidance network learn the conditional score of each noisy token given all previous clean tokens, by optimizing with a standard AR-diffusion loss mar. The generator and the guidance model are trained alternately.
• Figure 4: Training loss for initialization.
• Figure 5: Overall pipeline. Performance alignment is an optional technique in the CSD training stage, which is introduced in \ref{['app:training_cost']}.

Theorems & Definitions (2)

• Proposition 1
• proof