Beyond Autoregression: Fast LLMs via Self-Distillation Through Time

We genuinely appreciate Reviewer fzmS for their insightful feedback. Below, we present our responses to the questions and comments of reviewer fzmS.

For instance, only in the experiments section is it clarified that SDTT consists of multiple distillation rounds

• We understand that introducing the iterated SDTT in the experiments section may cause some confusion. To address this, we have updated the manuscript and now introduce iterated SDTT in the methodology section.

The progressive nature of SDTT is not immediately evident in the main text or Algorithms 1 and 2

• We apologize for any confusion regarding the progressive nature of SDTT. To clarify this aspect, we have revised the introduction, updated the title of Algorithm 2, and expanded the methodology section. These changes emphasize the progressive nature of SDTT, which we hope will provide a clearer understanding of our approach.

It would improve [the submission by] annotating the number of distillation rounds explicitly instead of relying on dt=1/k.

• We acknowledge that the notation could confuse some readers. Therefore, we have updated all figures and removed all use of the notation dt=1/k. We are grateful to Reviewer fzmS for helping us to improve our submission.

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).

• We have updated all figures following your comments. We hope that the updated text is clearer.

In Figure 2(a), the relationship between the figure and 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.

• You are correct that in Figure 2(a), "Round 1" and "Round 2" refer to denoising steps used to generate teacher targets for a single training step (i.e., one parameter update). We acknowledge that the use of "rounds" in this context may cause confusion, as it is also used for iterated SDTT. To clarify, we have updated Figure 2(a). We hope you will find the figure clearer (Figure 3a in the updated manuscript)

The concept of progressive distillation is not new within diffusion models. This paper does not sufficiently discuss its related work [1] throughout the introduction and the methods section.

• We acknowledge that progressive distillation is a well-known technique in continuous diffusion models. We have updated the introduction and methodology sections to further emphasize the progressive nature of our approach.

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. What is the implication of this distinction?

• Most continuous distillation methods depend on a deterministic mapping from noise to data, typically using the Ordinary Differential Equation (ODE) form of the reverse process. This means that once the model is trained, it can define a deterministic map (e.g., using DDIM) from noise to images, making it easier to learn this mapping in fewer steps. In contrast, discrete (absorbing) diffusion models do not have such a deterministic map (lines 80 and 462). The noise distribution in these models masks all tokens, meaning only a single sample can be generated deterministically.
• This distinction is important, as it means we cannot rely on matching a deterministic map in our distillation method. Instead, we must address the challenges of working with stochastic targets. As we hypothesize on lines 363-364, this stochasticity is likely why distilling more than two steps at a time yields worse performance empirically. We have updated the introduction to discuss this further (lines 80 to 83).

We sincerely thank reviewer ijRT for their thoughtful and insightful feedback. Below, we provide our responses to the questions and comments raised.

The detail that SDTT is performed over multiple rounds seems to be introduced in the experiments section.

• Thank you for the suggestion. We now introduce iterated SDTT in the methodology section, rather than the experiments section.

As mentioned by the authors, previous work has explored distilling multiple diffusion steps in the continuous diffusion setting.

• We acknowledge that distillation techniques have been explored to accelerate sampling in continuous diffusion models, and we cite several influential works on distillation in our related work section. However, we believe our contribution is novel for two primary reasons:
  • First: most continuous distillation methods rely on a deterministic mapping from noise to data, which is possible due to the Ordinary Differential Equation (ODE) formulation of the reverse diffusion process. Once the continuous model is trained, this deterministic map (e.g., with DDIM) allows training a faster model via distillation. In contrast, as noted in lines 80 and 462, no such deterministic map exists for discrete (absorbing) diffusion models. Indeed, if such a deterministic map existed, only one sample could be generated from the model, since all samples from the absorbing distribution are sequences of masked tokens.
  • Second, we are not aware of previous distillation methods for discrete diffusion models. We believe that it was not clear whether distillation techniques from continuous diffusion could be adapted to discrete diffusion. Our work shows that distillation is indeed possible, and we hope that it will motivate further research in speeding-up diffusion language models. Given the simplicity of our approach, we believe that there is significant room for further exploration in this area.

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.

• Thank you for the suggestion. We hypothesize (line 363 of the original manuscript) that since the teacher targets are stochastic, it is too difficult for the model to match targets generated over more than two steps. As evidence, we observed that the training loss is 2-4x larger when generating the teacher targets with 4 or 8 steps.

Did you investigate progressively growing the distillation step size throughout training?

• Yes, we did investigate progressively increasing the distillation step size during training. We tried several schedules, but all attempts failed, and the loss diverged after just 30-50 training steps. We have added a note about these failed experiments in the new background section, after introducing iterated SDTT.
• Final note: We noticed that the two works you referenced ([1], [2]) were not linked in your review. We assume that [2] refers to progressive distillation. We have updated the introduction and methodology to emphasize progressive distillation further and to highlight the difference regarding the existence of deterministic maps.

We sincerely appreciate reviewer 7pDo for their valuable and insightful feedback. Below, we respond to the questions and comments raised.

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

• We agree with reviewer 7pDo that the speedup is significant. Following your suggestion, we have moved the latency figure on the first page (now Figure 1). We are grateful for the suggestion

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)$.

• We are unsure if Reviewer 7pDo is suggesting a notation change or recommending we include an additional result from previous work. Could the reviewer 7pDo clarify the notation nit so that we can update the manuscript? We are currently reaching the page limit but we will do our best to include their suggestion.

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.

• We are grateful for the suggestion, Reviewer 7pDo is correct that the two notations are indeed equivalent. We agree that the suggested notation is simpler. We chose the current formulation to align with existing notation in the diffusion literature 1. This consistency helps avoid introducing additional variations, making the manuscript more accessible to readers, particularly those new to the field.
• We can modify the notation if reviewer 7pDo still believes it is necessary after reading our reply.

Please move iterated SDTT to the methods section and formally describe it with one equation.

• We appreciate Reviewer 7pDo’s suggestion. We have added a new section on iterated SDTT in the background with a new equation (equation 7 in the updated manuscript).

Writing: Are figure 2 and 3 referred to in the body of the text?

• Yes. In the original manuscript, we refer to figure 2 on lines 246-247. We refer to figure 3 on lines 202 and 416.

Notation such as dt should be formally defined in the methods section.

• We apologize for the confusion that the notation dt=1/k might have caused. We have completely removed it from the manuscript and updated all the figures to use the number of rounds of distillation instead. We hope that the reviewer 7pDo finds the updated figures clearer.

In general, "steps" can be ambiguous and can refer to sampling steps, decoding steps, or training steps.

• We agree that the term "steps" can be ambiguous. While we considered using alternatives like "iteration," the ambiguity remained. We have reviewed the paper and clarified whether we are referring to training, decoding, or sampling steps. Decoding and sampling are used interchangeably, as is common in the LLM literature. We hope this update addresses Reviewer 7pDo’s concern. If some any doubts remain, please let us know so that we can address them.

Question: In this setting, is MSE roughly $\chi^2$ divergence? It's surprising that using MSE results in better generative perplexity. Any thoughts on why?

• The MSE is indeed related to the $\chi^2$ divergence, although it is different. Most notably, the $\chi^2$ divergence assigns a non-uniform weight on the error for each $x$ in the discrete sample space (here the token indices), while the MSE assigns a uniform weight. Additionally, we apply the MSE on the log-probabilities computed by the model. Computing the loss in log space was more stable than exponentiating the predictions of the model and induced better performance. Following your question, we have included a paragraph on the $\chi^2$ divergence in appendix B4.

Could you also report the conditional likelihoods / perplexity alongside the MAUVE scores for conditional generation?

• We appreciate the suggestion. Figure 17 shows the perplexity of the continuations given the same prompts used to compute MAUVE. In summary, the student distilled with the KLD variant matches GPT2 with nucleus sampling using 32 forward passes. Since we are interested in conditional generation, we evaluate the perplexity of the continuation only and ignore the predictions for the prompt.

We are grateful to the reviewer LQck for the insightful feedback. Below, we address the questions raised by the reviewer LQck.

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.

• The reported 8x speedup is achieved with 16 sampling steps (Fig. 2b), with the model with 1.3B parameters (lines 430-431). We have updated the caption of Figure 2b (Fig. 1 in the updated manuscript) to say that the 8x speedup corresponds to the model with 1.3B parameters.

Additionally, the paper does not clarify whether the quality (generation perplexity) is comparable between 1.3B DLM and AR models.

• We could only afford to train models with up to 860M parameters with the resources available to us. We reported the sampling latency for 1.3B models because this scale is closer to autoregressive models used in production workloads. Given the increasing interest in diffusion language models, we anticipate that researchers with more resources will soon release diffusion models with billions of parameters. Therefore, we considered 1B+ models more interesting for latency evaluations. Nonetheless, following your suggestion, we included latency measurements for our small, medium and large models, and for models with 3B and 8B parameters to match the sizes of the recent Llama models. As for the 1.3B parameters model, we use untrained models for the 3B and 8B scales. In summary, the smaller models achieve greater latency gains than larger models. You can find the new results in the appendix.

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.

• We are grateful for the suggestion. We have updated the manuscript with latency measurements at all model sizes and with a varying batch size. You can find them in the appendix.

Effect of Flash-Attention: The paper does not address whether the advantage of DLMs over AR models persists when flash-attention is employed.

• We use flash attention in all latency experiments. We have updated the last sentence of section 4.4 to state it explicitly. We are grateful for your help in making our work clearer.

More downstream tasks such as math reasoning are encouraged.

• We are grateful for the suggestion. In addition to LAMBADA, we have included 6 other downstream tasks, including a multiple-choice math benchmark (mathqa). You can find the results in Table 1. We observe that distillation affects the downstream performance minimally. For example, after 7 rounds of distillation, we observe an average drop in performance of 0.31%, for the student distilled with the KLD loss, trained to sample in 8 steps.
• Let us note that all models evaluated on these tasks are relatively weak since they are only pre-trained on internet data (WebText/OpenWebText), which is not particularly designed for performance on math reasoning benchmarks.

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.

• We are grateful for the very relevant question. We measure the diversity of samples using the self-BLEU score, introduced by 1 and used in previous work on distillation 2. Less diverse samples achieve a higher self-BLEU score, and deterministic samples achieve a self-BLEU of 1. We have included a formal introduction to self-BLEU in appendix A, and reported it in figures 4a, 18a and 18b.
• We find that SDTT minimally affects diversity (i.e. there is no collapse). We observe an increase of at most 4 in self-BLEU (2 for the KLD student), while previous work 2 observed an increase of 15 or more.

How was the decision made to set the number of rounds to 7?

• The original sampling step size to generate the teacher target is 1/1024, corresponding to 1024 sampling steps to map the noise distribution to the data distribution. On each distillation round, we double the step size. Therefore, during the 7th rounds of distillation, we use a step size of 1/16 to generate the teacher targets. Because we generate the teacher targets using two sampling steps from the teacher, it means the final student is trained to generate text using 8 sampling steps. We stopped the distillation because it became too difficult for the student to match the teacher targets beyond 7 rounds.

Dear reviewers,

We sincerely appreciate your suggestions and effort in reviewing our work. As the discussion period draws to a close soon, we would like to remind you that we are more than happy to address any concern you might have. Please reach out to us if further clarifications are needed.

Best regards,

Authors