PaperlayerView exact revision on arXiv

Source-linked AI summary

Beyond Autoregression: Fast LLMs via Self-Distillation Through Time

Justin Deschenaux, Caglar Gulcehre

arXiv:2410.21035v2cs.LGcs.CL

TL;DR

Autoregressive generation produces one token at a time, while repeated sampling can improve performance but makes inference latency a bottleneck. The paper introduces SDTT, a distillation method for discrete diffusion models. It reduces decoding steps while retaining performance, with the final student reported up to 8 times faster than KV-cached AR models.

  • Problem

    Autoregressive models generate one token at a time, and latency is a bottleneck when using inference-time computation for tasks such as reasoning.

    Source: §1 · p. 1

  • Method

    Self-Distillation Through Time (SDTT) fine-tunes pre-trained masked diffusion language models to decode in fewer steps without relying on deterministic mappings.

    Source: §1 · p. 2

  • Results

    SDTT reduces decoding steps while preserving teacher natural-language-understanding performance on LAMBADA and six multiple-choice benchmarks; at 1.3B parameters, it is up to 8 times faster than AR models using KV caching.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    Faster decoding can make it practical to explore a larger search space in applications involving search, planning, and reranking.

    Source: §1 · p. 2

    A direct optimization approach was unstable: its loss diverged after 30–50 steps regardless of sampling step size.

    Source: §3.1 · p. 6

Abstract

from arXiv · show

Autoregressive (AR) Large Language Models (LLMs) have demonstrated significant success across numerous tasks. However, the AR modeling paradigm presents certain limitations; for instance, contemporary autoregressive LLMs are trained to generate one token at a time, which can result in noticeable latency. Recent advances have indicated that search and repeated sampling can enhance performance in various applications, such as theorem proving, code generation, and alignment, by utilizing greater computational resources during inference. In this study, we demonstrate that diffusion language models are capable of generating at least 32 tokens simultaneously, while exceeding the performance of AR models in text quality and on the LAMBADA natural language understanding benchmark. This outcome is achieved through a novel distillation method for discrete diffusion models, which reduces the number of inference steps by a factor of 32-64. Practically, at the 1.3B parameters scale, diffusion models, even without caching, can generate tokens at a rate that is up to 8 times faster than AR models employing KV-caching, and we anticipate further improvements with the inclusion of caching. Moreover, we demonstrate the efficacy of our approach for diffusion language models with up to 860M parameters.

1 INTRODUCTION

Repeated sampling can improve performance, making inference latency a bottleneck. The paper introduces SDTT for discrete diffusion models, reporting faster generation, preserved understanding performance, and generation in larger parallel token batches.

1 INTRODUCTION

  • SDTT generates at least 32 tokens at a time and achieves better perplexity than GPT-2 with nucleus sampling for conditional and unconditional generation.The method does not rely on deterministic mappings such as DDIM.Source: §1 · p. 2
  • Up to 8 times faster than AR models using KV caching, SDTT generates tokens in 16 decoding steps at 1.3B parameters.The discrete diffusion model does not use activation caching; the authors note potential for greater efficiency gains.Source: §1 · p. 2
  • The authors demonstrate SDTT’s effectiveness for models with up to 860M parameters, which they describe as the largest publicly available discrete diffusion language model.The “largest publicly available” characterization is qualified by “to the best of our knowledge.”Source: §1 · p. 2
  • SDTT preserves the teacher’s natural language understanding performance on LAMBADA and six multiple-choice benchmarks.The authors evaluate distilled students on LAMBADA and six benchmarks from Gao et al. (2021).Source: §1 · p. 2

2 BACKGROUND

Masked diffusion language modeling (MDLM) corrupts sequences toward a MASK-only distribution and learns to reverse that process. Knowledge distillation trains a student to imitate a teacher, with the training loop comparing their predictions on corrupted examples.

2.1 MASKED DIFFUSION LANGUAGE MODELING

  • MDLM corrupts data toward a MASK-only stationary distribution, then learns to recover the original sequence through a reverse process.Its noise schedule decreases over continuous time, and once a token is masked it cannot return to a real token in the forward process.Source: §2.1 · p. 3
  • The MDLM denoiser approximates the unknown data and supplies reverse-process predictions, with training based on a continuous negative evidence lower bound.The objective simplifies to a weighted cross-entropy between the ground-truth sequence and model predictions; derivation imposes that denoised tokens are not re-masked and remain carried forward.Source: §2.1 · p. 4

2.2 KNOWLEDGE DISTILLATION

  • Knowledge distillation trains a student network to imitate a more complex teacher using a divergence between their predictions.The cited examples include Kullback-Leibler divergence and total variation distance.Source: §2.2 · p. 4
  • The training loop initializes the student from teacher weights, creates a corrupted training example, compares student predictions with SDTT teacher targets, and updates the student.The example is corrupted through the forward diffusion process, and the divergence loss is optimized with AdamW.Source: §2.2 · p. 5

3 METHOD

Self-Distillation Through Time (SDTT) addresses the quality loss of few-step decoding by training a student to reproduce a many-step teacher's predictions with fewer steps. The method can be applied progressively to reduce decoding steps further, while a single-model step-size-growth alternative was unstable.

3 METHOD

  • Fewer decoding steps typically lower sample quality because they approximate the reverse process less accurately.Source: §3.1 · p. 5
  • The final SDTT model has lower generative perplexity in 32 steps than its teacher with 1024 forward passes.SDTT fine-tunes a pretrained masked diffusion language model to distill inference-time computation for sampling multiple steps into the student.Source: §3.1 · p. 5
  • SDTT targets collect teacher-predicted log probabilities whenever masked tokens are denoised, training the few-step student to match the many-step teacher.The procedure can distill multiple decoding steps at once, and the full training loop uses the resulting targets.Source: §3.1 · p. 5
  • Iterated SDTT progressively uses each newly distilled student as the next teacher, increasing the reduction in decoding steps across rounds.The authors report that sequentially minimizing objectives for increasing reduction factors is easier in practice than minimizing the direct objective.Source: §3.1 · p. 6
  • Growing the teacher-target sampling step size with a single model was unstable, with loss diverging after 30–50 steps regardless of how small the step size was.Source: §3.1 · p. 6

4 EXPERIMENTS

The experiments compare distillation objectives and training choices, then test SDTT’s scaling and generation latency. They evaluate sample quality, downstream performance, and diversity across these settings.

4.1 ABLATION ON THE TRAINING DIVERGENCE

  • KLD students slightly improve LAMBADA accuracy over the teacher, while their perplexity is better than or matches AR baselines except in the final SDTT round.KLD is the only tested divergence that retains LAMBADA performance.Source: §4.1 · pp. 7–8
  • KLD seems to outperform MSE and TVD on MAUVE; its student continuations reach the lowest perplexity and match GPT-2 nucleus sampling in 32 forward passes.The MAUVE comparison uses the teacher sampled with 128 decoding steps as its baseline.Source: §4.1 · p. 7, Figure 5 · p. 8
  • SDTT minimally affects downstream results and sample diversity; KLD has higher accuracy than other students on all but one task, while self-BLEU changes by at most 2.The comparison reports that autoregressive distillation in Agarwal et al. (2024) routinely increases self-BLEU by 15.Source: §4.1 · p. 7

4.2 ADDITIONAL ABLATIONS

  • Shortening SDTT rounds to 5k or 2.5k iterations hurts conditional generation but slightly improves final generative perplexity; distilling 4 or 8 teacher steps weakens LAMBADA results.LAMBADA accuracy remains unchanged with shorter rounds; the authors attribute the difficulty of larger target steps to increased target stochasticity.Source: §4.2 · p. 8
  • Analytical target sampling makes minimal difference from ancestral sampling, while using the EMA weights as teacher may slightly improve performance.The EMA-teacher experiment uses a decay rate of 0.999.Source: §4.2 · p. 8

4.3 SCALING SDTT TO 860M PARAMETERS

  • The 169M model sampled with 64 or more steps after distillation achieves better generative perplexity than the 863M model sampled with 1024 steps.Distillation also improves MAUVE for medium and larger models and LAMBADA accuracy with KLD.Source: Figure 4 · p. 7, §4.3 · p. 9

4.4 LATENCY WITH SDTT

  • 8x faster sampling than GPT-2 with KV caching is measured for untrained models around 1.3B parameters using 16 decoding steps.The latency experiment used a batch size of 8 and a single A100 GPU; the diffusion models did not use activation caching.Source: §4.4 · p. 9

5 RELATED WORK

The related work covers diffusion language models, discrete and masked language modeling, and approaches for distilling continuous diffusion. The paper situates its method among techniques that accelerate diffusion sampling.

Discrete Diffusion & Diffusion Language Models

  • Discrete diffusion language models such as SEDD can match or surpass GPT-2 in unconditional text generation when decoded with many steps [1].Other work develops continuous-time, categorical, absorbing, and re-parameterized discrete diffusion approaches.Source: §5 · p. 10

Distillation of Continuous Diffusion models

  • Continuous-diffusion distillation includes methods that train students to reproduce multiple teacher steps, predict final samples from intermediate points, or generate samples in one to four steps [1].Salimans & Ho (2022) teach students to match multiple teacher sampling steps; consistency distillation and other approaches use different strategies.Source: §5 · p. 10
  • Diffusion distillation methods often rely on DDIM’s deterministic sampling, while this paper’s approach cannot rely on such a deterministic map [1].The passage describes DDIM-based trajectory training and precomputed noise-image pairs as prior approaches.Source: §5 · p. 10

Masked & Non Auto-Regressive Language Modeling

  • BERT uses masked language modeling for representation learning, whereas discrete diffusion language models are generative; XLNet models text over permutations for pretraining [1].Pannatier et al. (2024) adapt an XLNet-like objective for generative modeling.Source: §5 · p. 10

6 DISCUSSION

The discussion presents SDTT as a way to reduce discrete diffusion decoding steps while retaining performance. It reports faster generation than KV-cached autoregressive models and identifies evaluation on large-completion tasks as future work.

6 DISCUSSION

  • SDTT reduces decoding steps for discrete diffusion models while retaining performance, addressing their greater inference-resource needs and lack of KV caching [1].The paper reports that its final student is up to 8x faster than autoregressive models using KV caching.Source: §6 · p. 10
  • The authors plan to evaluate SDTT on tasks that involve generating many completions from a base language model [1].Source: §6 · p. 10
  • The discussion states that SDTT applies to larger models, in addition to reporting speed gains over KV-cached autoregressive models [1].Source: §6 · p. 10

8 ETHICS STATEMENT

The ethics statement describes language models as dual-use and notes potential fake-content generation and bias risks. It says the paper focuses on reducing discrete diffusion language models’ test-time computational demands.

8 ETHICS STATEMENT

  • Language models can be used to generate fake content, a dual-use concern noted by the authors.Source: §8 · p. 11
  • The authors identify bias as a risk when language models are applied to datasets that are not carefully curated.Source: §8 · p. 11
  • The paper focuses on speeding up discrete diffusion language models at test time to reduce computational demands, and reports no specific concerns about this contribution.Source: §8 · p. 11

A ADDITIONAL ABLATION RESULTS

The additional ablations examine distillation objectives, round lengths, sampling choices, model scale, precision, diversity, and latency. Results include trade-offs for shorter rounds, no clear sampler benefit, and a higher generative perplexity under 64-bit computation while retaining the reported comparison with GPT-2.

Generative perplexity and precision

  • In 64-bit sampling, the final student has lower generative perplexity than GPT-2 with nucleus sampling (p=0.95) at 64 sampling steps, although overall generative perplexity is higher.The authors report that conclusions remain similar; Figure 9 covers 64-bit distillation and sampling.Source: A ADDITIONAL ABLATION RESULTS · p. 17, Figure 9 · p. 19

Conditional perplexity with TVD

  • On conditional perplexity, TVD performs worse than KLD and MSE; KLD was the best objective for retaining LAMBADA performance and was used for most ablations.Source: A ADDITIONAL ABLATION RESULTS · p. 17

Number of steps in each SDTT round

  • Shortening SDTT rounds to 5k or 2.5k iterations hurts conditional generation and MAUVE, but slightly improves final generative perplexity and marginally improves LAMBADA perplexity.The generative-perplexity gains are reported after the latest rounds; intermediate rounds are slightly worse.Source: Figure 7 · p. 18, Figure 8 · p. 19

Measuring the diversity

  • Self-BLEU diversity decreases only minimally after distillation with KLD, MSE, or TVD; the reported decrease is marginal compared with on-policy autoregressive distillation [1].The evaluation uses 1,000 prompts and five continuations per prompt.Source: A ADDITIONAL ABLATION RESULTS · p. 18, Figure 19 · p. 24

Ablation on the analytic sampler

  • The analytical sampler and the ancestral sampler used by Lou et al. (2023) show no clear benefit over the ancestral sampler in the reported ablations.Source: Figure 11 · p. 20, Figure 12 · p. 20

Distilling more than 2 steps at once

  • Distilling more than two teacher sampling steps at once appears to hurt performance, including when rounds were extended to 15k training iterations.Source: Figure 13 · p. 21

Plots for scaled SDTT

  • Scaled-model experiments report that SDTT helps medium and large models approach teacher performance with fewer sampling steps; undertrained models can improve in perplexity after distillation.At small scale, training for 400k rather than 1M steps yields a weaker model.Source: Figure 16 · p. 22, Figure 17 · p. 23

Conditional perplexity

  • For conditional completion perplexity, the final KLD-distilled student matches GPT-2 with nucleus sampling; the caption reports ground-truth continuations at approximately 13.11.Source: Figure 18 · p. 23

B ADDITIONAL DETAILS ON THE DIVERGENCE MEASURES

The student is trained to match teacher targets generated by Algorithm 1. The targets are matched using one of three divergence measures.

B.1 KULLBACK-LEIBLER DIVERGENCE

This section describes KLD properties and the reverse KLD used to train the student, then introduces TVD and its range.

B.1 KULLBACK-LEIBLER DIVERGENCE

  • KLD has a unique minimum when the two distributions are equal, but it is not symmetric in general.Source: §B.1 · p. 26
  • The student is trained using reverse KLD, DKL(pθ||pteacher), rather than forward KLD, DKL(pteacher||pθ).Source: §B.1 · p. 26
  • Forward KLD is zero-avoiding: minimizing it encourages pθ to assign non-zero probability wherever ptarget is non-zero.Source: §B.1 · p. 26
  • Reverse KLD is zero-forcing: minimizing it encourages pθ to assign zero probability where ptarget is close to zero.Source: §B.1 · p. 26

B.2 TOTAL VARIATION DISTANCE

  • TVD compares two discrete probability distributions on the same finite sample space, and its factor of 1/2 bounds it between 0 and 1.TVD is zero if and only if the distributions are equal.Source: §B.2 · p. 26

B.3 MEAN-SQUARED ERROR

The section distinguishes log-space MSE distillation from χ2 divergence, which compares distributions in probability space with a different weighting.

B.3 MEAN-SQUARED ERROR

  • For distillation, MSE is computed in log space for numerical stability and can compare scalar quantities beyond probability distributions.Source: §B.3 · p. 26
  • χ2 divergence compares probability distributions in probability space, whereas distillation MSE penalizes error in log space.Source: §B.3 · p. 27
  • MSE weights each summed term uniformly by 1/|Ω|, while χ2 divergence weights terms by 1/q(x).Source: §B.3 · p. 27

C IMPLEMENTATION DETAILS

The diffusion models follow prior work’s training code, preprocessing, and architecture, with a zero tensor replacing the noise-level input.

D TEXT EXAMPLES

This section presents varied, non-cherry-picked text samples from the small distilled model, including fragmented excerpts across unrelated topics, to illustrate its outputs and note limitations related to model size, text-quality fine-tuning, and training on concatenated documents.

D TEXT EXAMPLES

  • The examples are non-cherry-picked outputs from the small distilled model with KLD loss after the final distillation round, sampled unconditionally with varying inference steps.The section shows the first 512 tokens of each sample.Source: D TEXT EXAMPLES · p. 27
  • The samples come from small models not fine-tuned for text quality, and may begin mid-sentence because training used concatenated documents.Source: D TEXT EXAMPLES · p. 27
  • Displayed excerpts include garbled and incoherent text, as in a sample that shifts abruptly among unrelated topics.Source: D TEXT EXAMPLES · pp. 28, 30
  • The excerpts touch on UK–EU legislative negotiations and US energy policy, including pipelines, gas storage, and low-carbon initiatives.The fragments mention proposed legislative amendments and review, alongside energy infrastructure and renewable-energy measures.Source: D TEXT EXAMPLES · pp. 33–34
  • Other excerpts discuss Coinbase and bitcoin payments, marijuana advocacy, and cannabis policy.The bitcoin passages describe merchant acceptance and service costs; the marijuana passages refer to advocacy and proposals to regulate cannabis.Source: D TEXT EXAMPLES · p. 35
  • Further fragments address American religious history and UK public-sector reform, including claims about institutions, budgets, and benefits.These excerpts present disconnected historical and political assertions rather than a unified analysis.Source: D TEXT EXAMPLES · pp. 36–37

Labour and energy policy

  • Labour welcomed MPs’ vote for further workforce policies, while criticizing the government’s attention to automation’s effects on workers and rising energy-company profits.One passage reports annual profits rose by £3.1bn over the previous decade.Source: D TEXT EXAMPLES · p. 37

Seattle music scene

  • A Seattle musician describes independent studios and community connections as helping artists develop and share music, while noting that working with local part-timers is an obstacle.The speaker portrays Seattle’s music scene as small, vibrant, and growing.Source: D TEXT EXAMPLES · p. 38

Personal statement

  • A personal post apologizes to the community, expresses gratitude for its support and feedback, and asks readers to visit the next page for more detail.Source: D TEXT EXAMPLES · p. 39
Loading 2410.21035v2…