PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Enshu Liu, Qian Chen, Xuefei Ning, Shengen Yan, Guohao Dai, Zinan Lin, Yu Wang

arXiv:2510.21003v1cs.LG

TL;DR

Image AR models generate slowly because they sample tokens sequentially, while existing approaches have difficulty achieving one-step sampling without quality loss. DD2 distills conditional scores from a teacher AR model at each token position to train a one-step generator. On VAR, it reduces the one-step performance gap versus the original model by up to 67% compared with DD1, while achieving up to 8.1× sampling speedup.

  • Problem

    Image AR generation is slow because it samples tokens sequentially, and prior step-reduction methods generally do not achieve one-step sampling without significant performance degradation.

    Source: §1 · pp. 1–2

  • Method

    DD2 trains a one-step generator and a separate conditional guidance network to align the generator’s conditional scores with the teacher AR model’s scores at every token position.

    Source: §1 · p. 2, §3.2.3 · p. 5

  • Results

    On VAR, DD2 reduces the one-step model’s performance gap to the original AR model by up to 67% versus DD1 and achieves up to 8.1× sampling speedup.

    Source: §1 · p. 2, §5.2 · p. 8

  • Takeaways & Limitations

    At optimality, minimizing the conditional score distillation loss makes the one-step generator’s output distribution match the original AR model’s distribution.

    Source: §1 · p. 2

    DD2’s distilled models still have a performance gap relative to the original AR models.

    Source: §7 · p. 10

Abstract

from arXiv · show

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.

1 Introduction

Image AR models offer strong image synthesis but sample slowly, and existing acceleration methods rarely achieve one-step generation without substantial quality loss. DD2 replaces DD1’s predefined mapping with conditional score distillation; its reported results improve on DD1 while substantially accelerating sampling and training.

1 Introduction

  • Sequential token-by-token generation makes image AR inference slow, and prior step-reduction methods generally fail at one-step sampling without significant performance degradation.DD1 enables one-step sampling through a predefined noise-to-data mapping, but that mapping is challenging to learn and may constrain flexibility.Source: §1 · pp. 1–2
  • DD2 treats the teacher AR model as a conditional score model and trains a generator with a separate guidance network to align scores at every token position.The paper states that at optimality, the one-step generator’s output distribution exactly matches the original AR model’s distribution.Source: §1 · p. 2
  • DD2 addresses one-step sampling for AR models, a different generative setting from diffusion models and their score-distillation methods.The paper distinguishes the models’ generation processes and the goals and challenges of the two approaches.Source: §1 · p. 2
  • DD2 reduces the one-step performance gap versus the original AR model by up to 67% compared with DD1, with up to 8.1× sampling and 12.3× training speedups.Figure 2 compares DD2, DD1, pretrained models, and skip-last acceleration; DD2 outperforms DD1 by a large margin, while the other methods do not achieve one-step sampling.Source: §1 · p. 2, Figure 2 · p. 3

2 Preliminary

The preliminary section introduces image tokenization and autoregressive generation as the foundations for understanding DD2. It describes converting images into discrete codebook tokens and modeling each token conditioned on its predecessors.

2 Preliminary

  • Section 2 introduces standard image AR model formulations to establish the background for understanding DD2.The section then covers image tokenization and autoregressive modeling.Source: §2 · p. 3

2.1 Image Tokenizer

  • A VQ image tokenizer encodes an image into a lower-resolution feature map, assigns each embedding its nearest codebook vector, and decodes the resulting discrete tokens for reconstruction.Training uses a reconstruction loss to preserve fidelity between the input and reconstructed image.Source: §2.1 · p. 3

2.2 Auto-regressive Modeling

  • An AR model estimates each token’s probability conditioned on all preceding tokens, and generates a sequence by sampling tokens one by one.Generating a length-n sequence requires n autoregressive steps, which results in slow inference and limited efficiency.Source: §2.2 · p. 3

3 Distilled Decoding 2

DD2 trains a one-step generator to match an autoregressive teacher’s conditional distributions without relying on a predefined mapping, using conditional score distillation and a dedicated initialization phase.

3.1 Problem Formulation

  • The problem is to generate a sequence in one run from a latent variable, with its distribution matching that of a fixed, well-trained teacher AR model.The teacher conditions each next-token probability on all previous tokens.Source: §3.1 · p. 4

3.2.1 Teacher AR as a Conditional Score Model

  • The teacher AR model’s next-token probabilities define a closed-form conditional score for each noisy token, conditioned on previous clean tokens.The token is interpolated with Gaussian noise using the RectFlow schedule.Source: §3.2.1 · p. 4

3.2.2 Training Objective

  • DD2 minimizes conditional score distillation across token positions, matching each generator conditional distribution to the teacher’s and guaranteeing equality of their full sequence distributions at the minimum.The loss adapts score distillation to an autoregressive sequence by conditioning each noisy token on previous clean tokens.Source: §3.2.2 · p. 5

3.2.3 Learning the Conditional Score of the Generator

  • A separate conditional guidance network predicts the generator’s score; it and the generator are trained alternately, with the guidance network optimized using AR-diffusion loss.During generator updates, its fake-score term is replaced by the guidance network’s prediction, with gradients blocked from propagating into that network.Source: Figure 3 · p. 4, §3.2.3 · p. 5

3.3 Initialization of Generator and Guidance Network

  • Poor initialization can slow convergence or collapse training; experiments found that randomizing either model, even only the generator’s final layer, significantly degrades training.The paper attributes initialization’s importance to both stored network knowledge and the generator’s initial distribution.Source: §3.3 · p. 6
  • The initialization method replaces the teacher’s classification head with an MLP and fine-tunes using ground-truth scores; this improves training stability and convergence speed.The teacher’s probability-vector output and the generator and guidance network’s continuous outputs prevent direct reuse of the teacher’s output-head weights.Source: §3.3 · p. 6
  • The full pipeline first tunes an AR-diffusion model, then duplicates it as generator and guidance network before alternating CSD training; teacher-refined multi-step sampling is also described as an option.Performance alignment is optional during CSD training.Source: §3.3 · pp. 6–7, Figure 5 · p. 6

4 Related Works

Prior work reduces autoregressive sampling steps through set prediction, speculative decoding, and DD1, while score distillation has primarily been developed for diffusion models.

4.1 Reducing the Sampling Steps of AR Models

  • Set prediction reduces image AR sampling to around 10 steps, while speculative decoding achieves less than 3× compression because of weak draft-generator capacity.Set prediction struggles at very few steps because token correlations within each set are lost.Source: §4.1 · p. 7
  • DD1 maps noise tokens to target tokens through flow matching and enables one-step sampling, but its predefined mapping limits flexibility and it has performance and training-speed costs.DD2 is presented as alleviating these issues.Source: §4.1 · p. 7

4.2 Score Distillation for Diffusion Models

  • Diffusion score distillation uses a guidance network to approximate the generator score and alternately trains it and the generator to align generator and teacher distributions.Methods differ in generator loss; examples include KL divergence and score-function L2 distance.Source: §4.2 · p. 7
  • The cited video method applies DMD to temporal causal data but targets reducing diffusion-model steps, assuming a teacher diffusion model, unlike this work’s AR-model sampling focus.The passage also notes approaches that remove teacher guidance using training data and class-ratio estimation.Source: §4.2 · p. 8

5 Experiments

DD2 is evaluated on ImageNet-256 with VAR and LlamaGen, which differ in tokenization, token ordering, and sampling length. Results show strong one-step quality and speed, gains over DD1, a quality–step trade-off with teacher refinement, and faster training; initialization also matters.

5.1 Setup

  • ImageNet-256 tests VAR and LlamaGen, whose tokenization, token ordering, and original sampling lengths differ; DD1 is the main baseline.The setup also reports baselines that skip final steps or predict all token distributions in one step.Source: §5.1 · p. 8

5.2 Results of One-step Generation

  • DD2 compresses VAR teacher models across all sizes to 1 step, achieving up to 8.1× speedup with FID increase below 2.5; on LlamaGen, it achieves 238× speedup with FID increase 4.48.Source: §5.2 · p. 8
  • DD2 reduces the VAR teacher-to-one-step performance gap by up to 67% versus DD1, beats DD1’s two-step results on VAR, and improves LlamaGen FID by 2.76.The weak baselines fail to generate in one step.Source: §5.2 · p. 8, Table 1 · p. 9

5.3 Results of Multi-step Sampling

  • Using the pretrained AR model to regenerate later tokens gives a smoother quality–sampling-step trade-off, with sample quality increasing as more steps are used.DD1 uses its default multistep schedule for the comparison.Source: §5.3 · p. 8, Table 2 · p. 9

5.4 Training Efficiency

  • Compared with DD1, DD2 requires substantially fewer GPU hours to train, achieving up to a 12.3× training speedup while having better performance than DD1.Source: §5.4 · p. 8

5.5 Ablation Study: the Importance of Initialization

  • Removing proper initialization from either the generator or guidance network can cause significant performance degradation or collapse.The ablation tests LlamaGen-L and VAR-d24 with one component initialized from the tuned AR-diffusion model and the other from the teacher AR backbone with a random output head.Source: §5.5 · p. 9

6 Discussions

DD2 adapts score distillation to one-step sampling for pretrained AR models, a different generative setting from diffusion models. It avoids DD1’s predefined mapping and reports smoother latent interpolation than DD1.

6.1 Distinction with diffusion score distillation

  • DD2 distills pretrained AR models for one-step sampling, a different task and generative paradigm from diffusion-model step reduction.The discussion notes that one-step AR sampling remains underexplored despite AR models’ competitive or superior performance to diffusion models.Source: §6.1 · p. 9
  • Standard score distillation cannot be applied directly to AR models; DD2 uses a conditional generator score, an MLP output head, AR-diffusion initialization, and GTS loss.The GTS loss replaces the standard AR-diffusion loss to improve convergence.Source: §6.1 · p. 10
  • DD2’s results surpass many representative diffusion-distillation techniques, supporting AR distillation as a competitive approach to one-step generation.The comparison is reported in Table 6.Source: §6.1 · p. 10

6.2 Benefits of Eliminating Pre-defined Mapping in DD1

  • DD2 avoids DD1’s predefined mapping and supplies supervision at every token position; its latent interpolation is reported as smoother than DD1’s.The discussion also describes reduced error accumulation and smoother latent representations as potential benefits.Source: §6.2 · p. 10, Table 5 · p. 10

7 Future Works and Limitations

The discussion identifies continuous-space AR models and larger text-to-image tasks as future applications, while noting a remaining quality gap to the teacher model.

Future applications

  • Applying DD2 to continuous-space AR models is future work, although the authors describe the method as naturally compatible with models that provide conditional scores.These models generate each token through a diffusion process.Source: §7 · p. 10
  • Extending DD2 to larger-scale image AR tasks such as text-to-image is identified as a practical future direction.Source: §7 · p. 10

Limitations

  • Distilled models still have a performance gap from the original AR models; matching or surpassing pretrained-model quality remains future work.Source: §7 · p. 10

NeurIPS Paper Checklist

The checklist records the paper’s reproducibility and reporting disclosures, including that code had not yet been released and limitations were discussed. It also summarizes responses on theory, ethics, assets, and human-subjects research.

Limitations, LLM use, and societal impact

  • The checklist says limitations are discussed in Sec. 7, while the authors report that the core method does not use LLMs as an important, original, or non-standard component.The authors characterize the acceleration method as unrelated to domains that may have societal impacts.Source: NeurIPS Paper Checklist · pp. 14, 17, 19

Theory assumptions and proofs

  • The checklist marks theoretical assumptions and proofs as addressed, with the justification directing readers to Sec. A.The checklist guidance says proofs may appear in the main paper or supplement.Source: NeurIPS Paper Checklist · p. 14

Reproducibility guidance

  • The checklist’s reproducibility guidance describes possible avenues such as code, data, model access, or detailed replication instructions.It says code release is encouraged but not required by itself.Source: NeurIPS Paper Checklist · p. 15

Open access and reproducibility

  • The authors report that their code had not yet been released; they direct readers to Sec. 3, Sec. C, and Sec. E, and cite Tabs. 1 and 3 for computational cost.Source: NeurIPS Paper Checklist · pp. 15, 17

Experiment statistical significance

  • Some baselines lack error bars because their source papers do not report them; the authors say their metrics are not significantly affected by randomness.The checklist asks about error bars or other statistical-significance information.Source: NeurIPS Paper Checklist · p. 16

Ethics, assets, and human subjects

  • The authors report following the NeurIPS Code of Ethics and state that the work does not involve crowdsourcing or human-subjects research.The paper’s asset checklist response says models and datasets are properly cited and discussed in Sec. 5.Source: NeurIPS Paper Checklist · pp. 17–19

A Discussion of Prop. 1

The appendix proves Proposition 1 by induction over token positions and reports additional experiments on diffusion-distillation comparisons, training curves, and initialization. The reported initialization study favors the tuned AR-diffusion model and the proposed GTS loss, though the supplied passages do not give table values.

A Discussion of Prop. 1

  • Under sufficient network capacity, minimizing LCSD at each token position yields matching conditional distributions, and induction gives equality of the full sequence distributions.The first-token term reduces to traditional score distillation; the inductive step uses the already-matched prefix distribution.Source: A Discussion of Prop. 1 · p. 21

B More Experimental Details

  • The appendix presents additional experimental settings and results beyond the initialization and diffusion-distillation studies.It identifies these as supplemental details and results.Source: B More Experimental Details · p. 21

B.1 Comparison with Diffusion Distillation Methods

  • The appendix includes a comparison of DD2 with commonly used diffusion-distillation methods in Table 6, whose results are taken from the Shortcut model paper [7].The supplied passages do not state the table’s numerical results.Source: §B.1 · p. 21, Table 6 · p. 21

B.2 Training Curve of the Original DMD Method

  • The reproduced DMD training curves show that inappropriate score-distillation initialization can slow convergence and harm training stability, motivating AR-diffusion initialization.Figure 6 distinguishes default, random-generator, random-guidance, and partially randomized-generator initializations.Source: §B.2 · p. 21, Figure 6 · p. 22

B.3 Performance of the AR-diffusion Model for Initialization

  • The tuned AR-diffusion model provides strong sample quality for initialization, while using traditional diffusion loss degrades performance relative to the proposed GTS loss.The reported AR-diffusion model uses a 10-step Euler solver for each token’s sampling process.Source: Table 7 · p. 21, §B.3 · p. 22

B.4 Training Details and Cost

The appendix describes training costs and settings, reports that performance alignment improves convergence for VAR-d16 and VAR-d20, and outlines the derivation of the score function from flow matching. The supplied passages do not provide the training-cost values or the alignment table’s measurements.

B.4 Training Details and Cost

  • Training comprises AR-diffusion tuning and the main CSD-loss training process, with additional model-specific stages for VAR and LlamaGen.VAR training includes output-head fine-tuning followed by full-model training; LlamaGen additionally fine-tunes its teacher for continuous input.Source: §B.4 · pp. 22–23

B.5 The Effectiveness of Performance Alignment

  • The paper reports converged performance for the d16 and d20 models before and after performance alignment, presenting the comparison as demonstrating the operation’s effectiveness.Source: §B.5 · p. 23, Table 9 · p. 23

C Implementation Techniques

C.1 Computing Score Function with Probability Vector

  • The flow-matching derivation uses a linear noise schedule and relates the score function to velocity when the source distribution is Gaussian.In this problem, the target distribution is a weighted sum of Dirac functions, leaving finitely many cases for the expectation calculation.Source: §C.1 · p. 23

C.2 Multiple Noisy Samples for Fake Conditional Score Learning

The appendix details multi-sample training for the fake conditional-score network and several training procedures intended to align or stabilize generator and guidance performance. It also describes the shared transformer-and-MLP architecture.

C.2 Multiple Noisy Samples for Fake Conditional Score Learning

  • The guidance network predicts conditional scores from causal transformer features, noisy tokens, and timesteps, and multi-sample training averages its loss across several noise versions of each generated sequence.Each noisy token is formed by mixing the token with Gaussian noise according to its sampled timestep.Source: §C.2 · pp. 23–24

C.3 Performance Alignment for both Generator and Guidance network

  • Performance alignment replaces generator weights with EMA weights, then fixes the generator while retraining guidance to adapt to the updated generator distribution.The procedure is followed by resuming standard training and is reported as particularly helpful for VAR-d16 and VAR-d20.Source: §C.3 · p. 24

C.4 Larger Update Frequency for the Guidance Network

  • The guidance network is updated K times per iteration while the generator is updated once, with K greater than 1.The stated motivation is that guidance must track the generator distribution as it evolves.Source: §C.4 · p. 24

C.5 Details of Model Architectures

  • Both networks use a transformer backbone and lightweight MLP head, with the guidance head predicting conditional scores and the generator producing the sequence from shifted noise.The transformer features are causal for the guidance network.Source: §C.5 · p. 24

D Algorithms of Multi-step Sampling Method

This section presents a pseudo-algorithm for multi-step sampling and outlines the model parameterization and experimental settings used for AR-diffusion tuning.

D Algorithms of Multi-step Sampling Method

  • Figure 7 illustrates the model architectures and their corresponding inputs and outputs.Source: Figure 7 · p. 25
  • Algorithm 4 describes multi-step sampling using the distilled one-step model, the pretrained AR model, and more than one sampling step.The section reports results for the algorithm in Table 2.Source: D Algorithms of Multi-step Sampling… · pp. 24–25

E Detailed Experimental Settings

E.1 Model Parameterization

  • The models use velocity prediction, with the velocity function training the AR-diffusion model and guidance network and the generator output formed from noise and its predicted velocity.Source: §E.1 · p. 25

E.2 Stage 1: AR-diffusion Tuning

  • For VAR, tuning begins with the transformer frozen and its output MLP trainable, then all parameters are tuned; LlamaGen parameters are tuned from the start.Source: §E.2 · p. 26

E.3 Stage 2: Training with CSD Loss

This section gives DD2 training settings, describes adapting LlamaGen for continuous inputs, discusses set prediction as a baseline, and shows generated examples from several DD2 models.

E.3 Stage 2: Training with CSD Loss

  • For VAR and LlamaGen, conditional-score calculations use classifier-free guidance scale 2.0 and top-k 900, with top-p 0.95 for VAR and 1.0 for LlamaGen.Source: §E.3 · p. 26
  • Generator learning-rate warmup lasts 40K guidance-network iterations for VAR-d16, VAR-d20, and LlamaGen-L, and 20K for VAR-d24.Source: §E.3 · p. 26
  • A progressive EMA schedule uses a small early rate followed by a dynamic rate, which the authors say ensures training stability and fast convergence.The dynamic rate is min(0.9999, (iter + 1)/(iter + 10)).Source: §E.3 · p. 26

E.4 Continuous Input Adaptation for LlamaGen Models

  • LlamaGen’s embedding layer is replaced with an MLP to accept continuous latent embeddings, then the whole model is fine-tuned for 200K iterations to align performance with the original.The resulting teacher supports continuous embeddings without performance loss, according to the authors.Source: §E.4 · p. 26

E.5 Baselines

  • Set prediction cannot reduce sampling to one step; its one-step solution independently samples tokens by position frequency, ignoring token dependencies and producing the reported failure case.Source: §E.5 · pp. 26–27

F Visualizations

Related papers

Loading 2510.21003v1…