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

AI summary

Overview

• Research area: Efficient generative modeling, specifically accelerating image autoregressive (AR) models via distillation into one-step generators.
• Technical level: Advanced. The paper builds on flow matching, score matching, and score distillation, and its central result (Proposition 1) is a distribution-matching guarantee.
• Scope: The paper proposes Distilled Decoding 2 (DD2), a conditional score distillation framework that compresses the sampling of pretrained image AR models (VAR, LlamaGen) to a single step on class-conditional ImageNet-256.

What This Paper Is About

Image autoregressive models generate an image token by token, so they need many sequential forward passes (10 steps for VAR, 256 for LlamaGen), which is slow. The prior method DD1 achieved one-step sampling but only by training a generator to imitate a pre-defined noise-to-data mapping, which hurt quality and limited flexibility. DD2 instead treats the original AR model as a teacher that reveals the correct conditional score at each token position, and trains a one-step generator whose output distribution matches the teacher's, without any pre-defined mapping.

Key Contributions

1. A reinterpretation of AR models as conditional score models. The teacher AR model's next-token probability vector is converted, via a closed-form expression for the RectFlow noise schedule (Eq. 1), into a conditional score function over the codebook embedding space conditioned on all previous clean tokens.
2. A Conditional Score Distillation (CSD) loss. The loss aligns, at every token position, the conditional score of a separately trained guidance network with the teacher's conditional score (Eq. 4). Proposition 1 states that minimizing this loss guarantees the one-step generator's joint distribution over the token sequence equals the teacher's.
3. A novel initialization strategy. Because AR models output probability vectors while the generator and guidance network output continuous values, weights cannot be directly reused. The authors replace the teacher's classification head with a lightweight MLP and fine-tune it using a Ground Truth Score (GTS) loss (Eq. 6), then initialize both the generator and guidance network from this tuned AR-diffusion model.
4. Strong empirical results and released code. DD2 achieves one-step sampling on VAR and LlamaGen with reported speedups up to 8.1x and 238x respectively, plus up to 12.3x faster training than DD1. Code is available at https://github.com/imagination-research/Distilled-Decoding-2.

Main Findings

• Minimal quality loss on VAR: Compressing VAR from 10 sampling steps to 1 increases FID by less than 2.5, for example from 4.19 to 6.21 (VAR-d16), with up to an 8.1x speedup. Absolute one-step FIDs reported are 6.21 (VAR-d16), 5.43 (VAR-d20), 5.06 (VAR-d24), and 4.91 for VAR-d24 with more training.
• Large speedup on LlamaGen: The compression goes from 256 steps to 1 with FID degradation from 4.11 to 8.59 (LlamaGen-L), a 238x speedup. With more training, LlamaGen-L reaches FID 7.58 and IS 238.7.
• Clear margin over DD1: DD2 reduces the gap between one-step sampling and the original AR model by up to 67%. DD2's one-step results also outperform DD1's two-step results for all VAR models, and on LlamaGen the one-step FID is 2.76 better than DD1. For reference, DD1's one-step FIDs are 9.94 (VAR-d16), 9.55 (VAR-d20), 8.92 (VAR-d24) and 11.35 (LlamaGen-L).
• Weak baselines fail at one step: Skipping the last steps or predicting all token distributions in one step gives FIDs of 40.09 (VAR-skip-2), 157.5 (VAR-onestep), 80.72 (LlamaGen-skip-156) and 220.2 (LlamaGen-onestep).
• Flexible multi-step refinement: Using the teacher AR model to refine the last positions progressively improves quality, e.g., DD2 VAR-d16 at 3 steps gives FID 5.24, at 4 steps 4.88, and at 6 steps 4.47.
• Much cheaper training: On 8 NVIDIA A800 GPUs, DD2 trains in 115.5 GPU-hours for VAR-d16 (2.6x faster than DD1's 296.9), 174.4 for VAR-d20 (2.8x vs 484.4), 96.1 for VAR-d24 (6.3x vs 604.2), and 52.6 for LlamaGen-L (12.3x vs 647.7).
• Initialization is critical: On LlamaGen-L, FID-5k is 14.77 when both models are initialized, 16.08 with only the guidance network initialized, and 21.76 with only the generator initialized. On VAR-d24, mismatched initialization causes collapse (FID above 200).
• Better latent smoothness: Perceptual path length is 7231.9 for DD2 versus 18437.6 for DD1.

Methodology in Plain English

The authors start from the observation that an AR model's probability output for the next token can be converted into a "score" — a direction in embedding space that points toward more likely tokens — using an explicit formula derived from flow matching (Eq. 1). This score is conditional on all previously generated, noise-free tokens.

Rather than use that score only to build a deterministic mapping as DD1 did, DD2 borrows the score distillation idea used for diffusion models. Two networks are trained alternately:

1. A one-step generator that maps a whole sequence of Gaussian noise tokens to a complete token sequence in one pass, with no teacher refinement.
2. A conditional guidance network (a causal transformer backbone plus a small MLP head) that learns the score of the generator's own distribution at each token position, given the preceding tokens. It is trained with a standard AR-diffusion loss (Eq. 5).

The generator is then updated by the CSD loss, which measures the discrepancy between the teacher's true conditional score and the guidance network's estimated one, summed over all token positions, following the SiD loss form (Eq. 3). Gradients from the guidance network into the generator are blocked, and previous tokens are detached with a stop-gradient. The authors argue intuitively that this aligns distributions progressively: first token, then second conditioned on the first, and so on.

Because such score distillation is sensitive to initialization, they first replace the teacher's classification head with a small MLP and tune it with a GTS loss that uses the teacher's own analytic score rather than a Monte Carlo estimate. Both the generator and the guidance network are then initialized from this tuned model. A full pipeline of initialization tuning followed by alternating CSD/FCS training is given in Algorithm 3.

Why This Matters

• Research impact: DD2 is presented as the first method to one-step-sample image AR models without a pre-defined mapping, and it argues that conditional score distillation for sequential, AR-structured data is fundamentally different from score distillation for diffusion models. It also cuts the training cost of such distillation substantially.
• Interactive image generation: Reducing VAR to 1 step (with inference times around 0.019 to 0.031 reported in Table 1) supports real-time or near-real-time creative tools where users iterate repeatedly.
• On-device and edge deployment: A 238x reduction for LlamaGen reduces the compute needed per image, a relevant factor for running image generation on constrained hardware.
• Large-scale batch synthesis: Faster per-image generation lowers the cost of producing large volumes of generated images for data augmentation or content pipelines.
• Practical deployment of existing models: Because DD2 distills pretrained VAR and LlamaGen models without changing them, an industry already holding such checkpoints can gain substantial latency reductions.

Future Directions

• Closing the remaining quality gap: One-step FID still rises from 4.11 to 8.59 on LlamaGen-L, so further work is needed to make one-step quality indistinguishable from the teacher.
• Extending conditional score distillation to other sequential modalities: The authors frame their problem as sequential, AR-structured data and contrast it with diffusion-only score distillation, leaving open how well the CSD formulation transfers to other token-by-token generators.
• Removing the reliance on initialization tuning: The paper reports that DD2 collapses (FID above 200) on VAR-d24 without the proposed initialization, raising the question of how to make the method robust with cheaper or no initialization.
• Generalizing across architectures and tokenizers: The evaluation covers VAR's multi-resolution tokens and LlamaGen's raster-scan tokens on ImageNet-256; whether the same behavior holds for other AR setups, datasets, and conditional settings is not reported.

Target Audience

Researchers and practitioners in generative modeling and efficient inference who are familiar with autoregressive models, flow matching, and score distillation. It is most useful for those working on accelerating image generation, distilling large generative models, or deploying AR image models in latency- or cost-sensitive systems. Readers seeking the mathematical details will find the score derivation central, while readers interested in deployment will focus on the FID, speedup and GPU-hour tables.

Authors’ 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.