Beyond Autoregression: Fast LLMs via Self-Distillation Through Time

Authors: Justin Deschenaux, Caglar Gulcehre
Primary area: generative models
Keywords: language modeling, LLM, diffusion models, discrete diffusion models, diffusion language models, distillation

Abstract:
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.

Decision: accept (poster)
Reviewer scores: 6, 8, 8, 6

Reviews and discussion:
<thread>
creation date:
1730508257666

thread_type:
official_review

<thread_object>
creator: reviewer_fzms

<content>
summary:
This paper explores progressive distillation, a technique commonly used in continuous image diffusion models, in the context of discrete diffusion language models. Starting with a trained discrete diffusion model that might take many denoising steps to generate a text sequence, the proposed framework iteratively distills the slower teacher into a faster student diffusion model, effectively halving the number of sampling steps in each distillation iteration. The distillation process involves collecting predicted logits for all tokens from the teacher model and then training the student diffusion model to match the teacher’s logit predictions according to a defined divergence measure. This self-distillation yields a significantly faster diffusion model, capable of generating text with up to 32 times fewer sampling steps, all while improving generation quality.

strengths:
- This work makes a valuable contribution by extending progressive distillation techniques to diffusion language models, achieving much faster generation with improved generation quality.
- This work also includes extensive ablation experiments that carefully examine each design component of the distillation framework, demonstrating the solid empirical performance of SDTT. These results offer useful insights for future advancements in discrete diffusion language models.

weaknesses:
- The paper appears to be written hastily and lacks clarity and organization. Beyond typos and formatting issues (examples below), key concepts are introduced ambiguously. For instance, only in the experiments section is it clarified that SDTT consists of multiple distillation rounds rather than a direct distillation; and the progressive nature of SDTT is not immediately evident in the main text or Algorithms 1 and 2, becoming clear only in experimental results. In addition, the number of distillation iterations can only be inferred indirectly from `dt`, which is quite unintuitive and causes more confusion. It would greatly improve understanding by emphasizing the progressive formulation early in the methods section, and annotating the number of distillation rounds explicitly instead of relying on `dt=1/k` .
- Inconsistent use of x-axis labels in figures also adds to the confusion. Some figures, such as Figures 1(a), 1(b), 8, 9(b), 15(a), and 15(b), use “sampling step size” as the x-axis (e.g., `1/512`) while others use the “number of sampling steps” (e.g., 1024). If these two metrics are equivalent to each other (e.g., 512 sampling steps correspond to a sampling step size of 1/512), unifying the x-axis to the “number of sampling steps” across all figures would greatly enhance clarity; and also prevent confusion between terms like “sampling step size = 1/512” and “dt=1/512”, which, despite similar notation, have different meanings.
- In Figure 2(a), the relationship between the figure and various distillation rounds is ambiguous. If I understand correctly, “Round 1” and “Round 2” in the figure actually denote different denoising steps within a single distillation round, rather than separate distillation rounds.
- The concept of progressive self-distillation is not new within diffusion models. The proposed framework SDTT largely resembles progressive distillation [1] with modifications such as the construction of targets and distillation objectives, both of which are straightforward following the formulation of discrete diffusion. While SDTT adapts distillation techniques to a different application and presents strong results, this paper does not sufficiently discuss its related work [1] throughout the introduction and the methods section, making it difficult to evaluate how SDTT differs from previous contributions [1].

[1] Salimans, T., & Ho, J. (2022). Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00512.

questions:
1. Typos:
    1. Line 96: `(MDLM.)` → `(MDLM).`
    2. Line 785: `the the MAUVE` → `the MAUVE`
    3. Line 785 `resuts` → `results`
2. Clarifications on contributions:
    1. In Line 79, “Unlike many distillation methods for continuous diffusion models, SDTT does not rely on deterministic mappings such as DDIM” might require further elaboration. Given that discrete diffusion models inherently lack a clear connection to deterministic processes akin to DDIM, what is the implication of this distinction? Specifically, how does the absence of deterministic mappings influence the theoretical foundations or empirical performance of SDTT?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730672685929

thread_type:
official_review

<thread_object>
creator: reviewer_ijrt

<content>
summary:
In this work, the authors propose Self-Distillation Through Time (SDTT), a training strategy for decreasing the number of sampling steps required during discrete text diffusion. Given a pretrained discrete diffusion model, the method works by using the token distribution of the base model at later denoising timesteps as the target for a student network at an early denoising timestep. The student model is initialized to be the same as the base diffusion model. The authors demonstrate that through multiple rounds of SDTT they can sample with 32x fewer steps while maintaining quality. They also compare to standard auto-regressive decoding, finding that they can achieve the same perplexity with an 8x reduction in wall clock time during generation.

strengths:
* The SDTT methodology appears to significantly improve the decoding speed of discrete diffusion models, and the authors provide evidence of this over a large set of experiments. Compressing discrete diffusion sampling to this degree seems like an important contribution that will be used by future text diffusion works.

* The paper conducts a thorough set of ablations justifying key design choices such as the particular divergence measure and the number of steps compressed per SDTT round.

* The authors demonstrate that text diffusion models can scale up to 860 million parameters and promise to release said model which is a valuable contribution for future research on text-based diffusion models.

weaknesses:
* Performing multiple rounds of SDTT seems integral to the empirical success the paper has but the detail that SDTT is performed over multiple rounds seems to be introduced in the experiments section. It would help the clarity of the paper if this technique was introduced as a core part of the algorithm.

* As mentioned by the authors, previous work has explored distilling multiple diffusion steps in the continuous diffusion setting ([1], [2]). While the technique in SDTT has slight differences and is applied to the discrete diffusion setting, its similarity to [2] in particular may limit the novelty of SDTT.

questions:
* It would be interesting to look at why more SDTT rounds work but distilling a larger number of steps in a single round doesn’t. Did you investigate progressively growing the distillation step size throughout training?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730703401896

thread_type:
official_review

<thread_object>
creator: reviewer_7pdo

<content>
summary:
This paper proposes self-distillation through time to speed up sampling from text diffusion models. The method trains student reconstruction models to predict ahead in time, given a teacher model. This method can be applied iteratively, using students to train even coarser students. Experiments demonstrate that the method results in achieving a given generative perplexity at fewer function evaluations.

strengths:
The method is simple and intuitive. The results for generative perplexity and number of function evaluations look promising. However, wall-clock time / latency is most important.

weaknesses:
The primary weakness is presentation and writing. The abstract claims that the generates tokens 8x faster than an AR baseline with KV-caching. This must be presented in a figure as early as possible (see comment 10). See comments below for more suggestions.

questions:
1. Notation nit: In the discrete case, with fully-factored reverse model, $p(z_s \mid z_t) = \sum_x p(z_s \mid x)p(x \mid z_t)$ can be computed efficiently (see D3PM) and (also in the discrete, fully-factored case) $x_\theta(z_t, t) = p(x \mid z_t)$.
2. Notation nit: Equation 3 can be written as $\log p(x \mid z_t)$, rather than writing a probability as the inner-product of the mean-parameters with a one-hot vector.
3. Writing: Please move iterated SDTT to the methods section and formally describe it with one equation and a few sentences. Something like a telescoping sum of divergences. Be sure to introduce terminology like "rounds" formally and do not use both "times" and "rounds" in order to minimize terms.
4. Writing: Are figure 2 and 3 referred to in the body of the text?
5. Writing: Notation such as `dt` should be formally defined in the methods section. Could you explain what `dt` is? Is it $\Delta$? Could you also define $\Delta$? In MDLM, $\Delta$ refers to a simplex rather than time-step or increment.
6. Writing: In general, "steps" can be ambiguous and can refer to sampling steps, decoding steps, or training steps. Please be clear which one is being used, for example in the beginning of section 4.1 it should be training steps. Also, likely sampling = decoding steps.
7. Question: In this setting, is MSE roughly chi-square divergence? It's surprising that using MSE results in better generative perplexity. Any thoughts on why?
8. Experiment: Could you also report the conditional likelihoods / perplexity alongside the MAUVE scores for conditional generation? It's not completely clear what MAUVE measures [1].
9. Question: It's very interesting that distilling more than 2 steps at a time, as well as more training steps per round hurt the student (Figure 11). 

10: Figure request and section 4.4: In section 4.4, the following claim is made: "We successfully reproduce the results of Deschenaux & Gulcehre (2024), which showed a 4x improvement when sampling with 32 steps, only this time, we retain the text quality because of SDTT." Please add a figure that clearly shows this. A figure with latency vs generative perplexity is extremely important for this paper. Additionally, Section 4.4 is likely the most important experimental section, and should be the first result discussed.

[1] Pimentel, Tiago, Clara Meister and Ryan Cotterell. “On the Usefulness of Embeddings, Clusters and Strings for Text Generator Evaluation.” ICLR 2023.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730723109037

thread_type:
official_review

<thread_object>
creator: reviewer_lqck

<content>
summary:
The paper introduces a novel distillation method for discrete diffusion models: SDTT (Self distillation through time), which significantly reduces inference steps, leading to a generation process that is up to 8 times faster than AR models with KV-caching. They ablate with different distillation alternatives. They evaluate the effectiveness of MDLMs with model sizes up to 860M parameters under unconditional generation, conditional generation, and downstream LAMBADA benchmark.

strengths:
- Improved Generation Quality and Speed: The paper achieves remarkable improvements in both generation quality and efficiency, with fewer decoding steps required. This is a significant advancement for the field.
- Clarity of Writing: The paper is well-written, with clear explanations of the methodologies and results, making it accessible to a wide audience.

weaknesses:
- Insufficient Validation on Generation Speed: While generation speed is a key advantage highlighted in the paper, the validation seems inadequate in Section 4.4 and Figure 2b. The experimental settings lack clarity. For instance, it's not clear if the reported 8x speedup is based on 32 steps, 1.3B, and a batch size of 8. Additionally, the paper does not clarify whether the quality (generation perplexity) is comparable between 1.3B DLM and AR models.
- Impact of Batch Size and Model Size: There is a need for further validation on how batch size affects memory usage and decoding speed for both diffusion and AR models. Similarly, the paper should explore how model size impacts generation speed.
- Effect of Flash-Attention: The paper does not address whether the advantage of DLMs over AR models persists when flash-attention is employed.
- More downstream tasks such as math reasoning are encouraged.

questions:
- Does the self-distillation process affect the generation diversity? There is an ongoing discussion regarding potential model collapse when AI models are trained on recursively generated data, which may be relevant here.
- How was the decision made to set the number of rounds to 7? 
- The paper focuses on discrete diffusion models and generation speed for reasoning tasks but fails to cite these related works: 
(1) A Reparameterized Discrete Diffusion Model for Text Generation;
(2) Diffusion of Thoughts: Chain-of-Thought Reasoning in Diffusion Language Models.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732191531295

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General reply

comment:
Dear reviewers,

We appreciate your insightful comments. In response to your suggestions, we have improved our manuscript and conducted additional experiments, which are now included in the updated version. For ease of review, the new text is highlighted in red. Here's a summary of the updates we've made to address the reviewers' questions:

- We specify that SDTT consists of multiple rounds in the Methods section, and updated the title of Algorithm 2 to clarify.
- We updated the label for the horizontal axis for the plots that used the `1/k` notation, as it confused the reviewers. We hope that using the number of rounds of distillation is clearer.
- We have further emphasized the relationship between continuous and discrete distillation in the introduction and methods sections. Most importantly, we further discuss the progressive nature of SDTT and describe how it differs with progressive distillation in the introduction.
- We clarify in the introduction why a deterministic sampling algorithm cannot exist for discrete diffusion models, unlike continuous diffusion models.
- We evaluate our models on six additional downstream tasks, including a multiple-choice question on math.
- We measure the diversity of generated samples using self-BLEU and observe that SDTT minimally reduces the diversity. In particular, the reduction is significantly smaller than on-policy distillation of autoregressive language models [1](https://arxiv.org/abs/2306.13649). Importantly, SDTT does not cause a collapse of the output distribution.
- We evaluate the latency on smaller and larger scales and with a different batch size. We also report the trade-off between perplexity and wall-time latency in the appendix.

These new results reinforce our claim that SDTT is an effective distillation method for diffusion language models, and makes our masked language models faster than the AR baselines.

If you have any questions about the new experiments or any others, please reach out to us.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper presents a novel distillation method, SDTT, aimed at reducing the number of sampling steps required for discrete diffusion models, thereby improving the speed of text generation. The authors have demonstrated that their method can achieve significant speedups while maintaining or improving the quality of text generation. All reviewers vote for acceptance of this paper. The main concerns lie in that authors ignore the non-autoregressive text generation literature, which are quite related to the acceleration of text generation and authors should include them for comparison and discussion.