Learnable Sampler Distillation for Discrete Diffusion Models

Thank you for your positive evaluation of our paper and for your valuable comments and suggestions. Below are our responses to the main points raised.

(Theoretical justification is lacking. )

Thank you for the comment and question. We agree that the theoretical guarantee concerning the matching of teacher and student distributions is beneficial and providing theoretical justification for our distillation methods would strengthen the work. However, our current focus is primarily empirical: We aim to develop practical and effective techniques for accelerating discrete diffusion models. Given the time constraints of the short rebuttal period, we are unable to establish rigorous theoretical guarantees for our distillation methods at this stage, and we plan to address this with more thorough theoretical analysis in future work. That said, we would like to offer an intuitive explanation for the effectiveness of our method, which may help lay the groundwork for subsequent theoretical framing. Specifically, the final discrepancy between the outputs of the student and teacher samplers stems from the accumulation of small local errors made at each step. Each of these local errors, in turn, is tied to the specific score predicted by the model at that step. Our training objective directly enforces alignment between the score predictions of the student sampler and those of the teacher sampler at every step. By implicitly correcting these small local errors throughout the process, our approach guides the student sampler to produce final outputs that closely match the high-quality results of the teacher sampler.

(The method is not intuitively easy to understand. It is unclear why having similar scores at different points is beneficial.)

Thank you for the comment. An intuitive explanation of why having similar scores at different locations is helpful is provided as follows.

In our view, the training objective that aligns scores at corresponding time steps functions as a mechanism for continuous error correction, addressing two main sources of error. First, by aligning scores at each step, we mitigate the accumulation of discretization errors that arise from large step sizes. Second, this approach helps alleviate the training-sampling mismatch, a common issue that leads to compounding decoding errors. In our method, the teacher sampler that operates with a large number of sampling steps effectively emulates the dynamics of the training process. Conversely, the student sampler, which uses far fewer steps, is designed to mimic the inference process. By enforcing that the scores of the student sampler match the scores of the teacher sampler at each corresponding time step, we partially bridge the gap between training and sampling phases. Additionally, as noted in the Related Work Section of our main paper, aligning scores across the full trajectory is computationally intractable due to the non-differentiable nature of sampling from categorical distributions. Thus, we adopt per-step alignment as a practical alternative to circumvent this non-differentiability, and empirical results demonstrate its effectiveness in enhancing sample quality at low NFEs.

(What do the learned coefficients and timesteps look like?)

Thank you for the question. We provide a visualization of the learned parameters for a 16-step student sampler, which is distilled from a 1024-step teacher sampler on the SEDD-small backbone.

The learned timesteps are: [1.0000, 0.9219, 0.8459, 0.7719, 0.7000, 0.6302, 0.5625, 0.4969, 0.4334, 0.3719, 0.3125, 0.2552, 0.2000, 0.1469, 0.0959, 0.0469] Notably, similar to JYS, the learned schedule concentrates more steps near $t=1$. However, the learned distribution of time step is not as drastic as that of JYS, since we also adjust the student sampler through coefficients.

The learned coefficients for these timesteps are: [1.1313, 1.1305, 1.1289, 1.1270, 1.1245, 1.1210, 1.1178, 1.1141, 1.1102, 1.1068, 1.1031, 1.0999, 1.0962, 1.0925, 1.0881, 1.0839] The coefficients are all greater than 1. Additionally, the coefficients decrease with respect to timesteps, which indicates the sampler learns to apply a stronger corrective amplification to the score when taking larger steps at earlier stages, effectively compensating for the increased discretization error.

(Additional baselines.)

Thank you for this suggestion.

We conduct experiments on the sampling method of DNDM [1] you mentioned, adhering to the original experimental settings outlined in the DNDM paper, and the perplexity results for unconditional text generation using the FairSeq backbone are reported in Table E1. As for the SEDD backbone, we have not performed comparable experiments, as no public implementation currently exists for applying DNDM to SEDD.

Table E1: Comparisons on the FairSeq backbone (in terms of Perplexity$\downarrow$).

Regarding the specific work [2] you mentioned, we politely clarify that we included a direct performance comparison with the reported results of it in Table 3 of our main manuscript. We could only compare our results with their work on the RADD backbone since the public implementation of their work is currently unavailable.

Thanks for your positive assessment of this paper and the valuable comments and suggestions. Our responses to the main concerns are given as follows.

(Lack of justification for the training objective.)

Thank you for your comment and question. In our view, the training objective that aligns scores at corresponding time steps functions as a mechanism for continuous error correction, addressing two main sources of error. First, by aligning scores at each step, we mitigate the accumulation of discretization errors that arise from large step sizes. Second, this approach helps alleviate the training-sampling mismatch, a common issue that leads to compounding decoding errors. In our method, the teacher sampler that operates with a large number of sampling steps effectively emulates the dynamics of the training process. Conversely, the student sampler, which uses far fewer steps, is designed to mimic the inference process. By enforcing that the scores of the student sampler match the scores of the teacher sampler at each corresponding time step, we partially bridge the gap between training and sampling phases. Additionally, as noted in the Related Work Section of our main paper, aligning scores across the full trajectory is computationally intractable due to the non-differentiable nature of sampling from categorical distributions. Thus, we adopt per-step alignment as a practical alternative to circumvent this non-differentiability, and empirical results demonstrate its effectiveness in enhancing sample quality at low NFEs.

(Underexplored design choices.)

Thank you for raising this point about our design choice for the learnable coefficients. Our decision to follow prior works like S4S and LD3 in adopting a scalar form for the coefficient is a trade-off between expressiveness and computational feasibility. While a more expressive vector-form coefficient is theoretically feasible, directly modulating transition probabilities for each token would require the vector to match the vocabulary size (e.g., 50258 for GPT-2). This would result in a large number of learnable parameters, calculated as the product of the number of sampling steps and the vocabulary size. Such a surge would significantly increase memory usage and heighten the risk of overfitting during distillation. Thus, we opted for the scalar formulation, which offers simplicity, greater computational efficiency, and empirical effectiveness across our experiments.

(Lacking theoretical grounding.)

Thank you for this question. We agree that analyzing the distributional divergence between the optimal student sampler and the teacher sampler is beneficial for fostering a deeper understanding of the proposed method and its limitations. However, our current focus is primarily empirical: We aim to develop practical and effective techniques for accelerating discrete diffusion models. Given the time constraints of the short rebuttal period, we are unable to establish rigorous theoretical guarantees for our distillation methods at this stage, and we plan to address this with more thorough theoretical analysis in future work. That said, we would like to offer an intuitive explanation for the effectiveness of our method, which may help lay the groundwork for subsequent theoretical framing. Specifically, the final discrepancy between the outputs of the student and teacher samplers stems from the accumulation of small local errors made at each step. Each of these local errors, in turn, is tied to the specific score predicted by the model at that step. Our training objective directly enforces alignment between the score predictions of the student sampler and those of the teacher sampler at every step. By implicitly correcting these small local errors throughout the process, our approach guides the student sampler to produce final outputs that closely match the high-quality results of the teacher sampler.

(For larger models, e.g., DiffuLLaMA and DiffuGPT, how well can the proposed method perform?)

Thank you for your questions regarding the scalability of our proposed methods. To demonstrate that our framework is not limited to smaller models, we conduct experiments on applying LSD+ to DiffuLLaMA and DiffuGPT. We perform sampler distillation on pre-trained DiffuGPT-S and DiffuLLaMA checkpoints. The results in terms of Perplexity are presented in Tables D1 and D2.

Table D1: Comparisons on the DiffuGPT-S backbone (in terms of Perplexity$\downarrow$).

Table D2: Comparisons on the DiffuLLaMA backbone (in terms of Perplexity$\downarrow$).

This experiment confirms that our approach is compatible with DiffuGPT and DiffuLLaMA. We will add these results to the Experiments Section of our paper.

Thanks for your recognition of this paper and the valuable feedback and suggestions. Our responses to the main concerns are given as follows. All citations refer to the reference list at the end of the responses.

(Evaluations are limited to relatively small-scale datasets (OpenWebText, CIFAR-10) and short sequences (1024 tokens).)

We thank the reviewer for the comment. To address this, we perform experiments on larger-scale datasets. For image generation, we conduct experiments on ImageNet (256x256) using the MaskGIT architecture [1] as the backbone, incorporating the recently proposed advanced Halton sampler [2] as the sampling method. The results are presented in Table C1.

For text generation, we scale up our experiments to include the comparison to the sampler with respect to a larger backbone DiffuGPT-S [3]. The results are presented in Table C2, where we report perplexity on 1024 unconditionally generated text samples.

Table C1: Comparisons on ImageNet 256x256 (in terms of FID$\downarrow$).

Table C2: Comparisons on the DiffuGPT-S backbone (in terms of Perplexity$\downarrow$).

We also present results for larger-scale experiments on the SDTT-KLD backbone [4], which utilizes Ancestral as its sampler, and the MDLM backbone [5], which utilizes ReMDM [6] as the sampler. The experimental results are presented in Tables C3 and C4.

Table C3: Comparisons on the SDTT-KLD backbone.

Table C4: Comparisons on the MDLM backbone.

These results show that LSD+ maintains its performance and acceleration benefits when applied to more complex datasets and larger models. We will integrate these results into the Experiments Section of our revised manuscript.

(The robustness of the algorithm to parameters (e.g., Hamming distance threshold) has not been discussed.)

Thank you for the comment. To investigate the robustness of the algorithm to the Hamming distance threshold, we conduct the ablation on the SEDD-small backbone using the Euler sampler with 32 inference steps. We train our LSD+ method using several different values for the Hamming distance threshold, specifically {0%, 1%, 5%, 10%, 20%} of the sequence length, while keeping all other hyperparameters unchanged. The performance, measured by Perplexity, is reported below in Table C5.

Table C5: Ablation study on the Hamming distance threshold for the relaxed objective.

We empirically choose 5% of the sequence length as the Hamming distance threshold in our main paper.

(The authors did not discuss the limitations.)

Thanks for the comment. We politely clarify that as mentioned in our NeurIPS Paper Checklist, we discussed the limitations in the supplementary material (see Appendix A).

References:

[1] Chang et al. Masked Generative Image Transformer, CVPR’22.

[2] Besnier et al. Halton Scheduler for Masked Generative Image Transformer, ICLR’25.

[3] Gong et al. Scaling Diffusion Language Models via Adaptation from Autoregressive Models, ICLR’25.

[4] Deschenaux et al. Beyond Autoregression: Fast LLMs via Self-Distillation Through Time, ICLR’25.

[5] Sahoo et al. Simple and Effective Masked Diffusion Language Models, ICLR’25.

[6] Wang et al. Remasking Discrete Diffusion Models with Inference-Time Scaling, ICLR’25.

Thanks for your comprehensive comments and questions. Our responses to the main concerns are given as follows. All citations refer to the reference list at the end of the responses.

(This work does not compare against several important recent methods and settings which focus on quality vs latency of discrete diffusion models.)

We thank the reviewer for the comment. We have included comparisons with several recent methods that focus on the quality vs latency of discrete diffusion models, such as JYS and $\theta$-trapezoidal (see Tables 2 and 3 in the main paper). At this stage, we would like to clarify that we are uncertain about the specific works the reviewer is referring to. Additionally, we would like to politely draw attention to NeurIPS policy, which states:“Papers appearing online after March 1st, 2025 are generally considered concurrent to NeurIPS submissions. Authors are not expected to compare to those” (https://neurips.cc/Conferences/2025/PaperInformation/NeurIPS-FAQ). That said, we are fully committed to strengthening the empirical rigor of our study. If the reviewer could kindly specify the relevant works or settings they have in mind, we would be happy to conduct additional comparisons regarding the quality vs latency of discrete diffusion models and incorporate them into our revised submission.

(Authors do not consider predictor-corrector solvers like in DFM nor higher order solvers like Runge Kutta or Trapezoidal though they cite relevant work.)

Thank you for your comment. We conduct experiments on predictor-corrector (PC) solvers in Discrete Flow Matching (DFM) [1] using MDLM [2] as the backbone, and find that the improvement is limited. This might be attributed to the specific dynamics of the flow-matching process in DFM, whose error characteristics may not align as well with our error correction mechanism. We report the corresponding results in Table B1. However, we successfully achieve good results on $\tau$LDR-10 which also utilizes PC solvers [3] on the SEDD-small backbone, with the results presented in Table B2.

Table B1: Comparisons on the MDLM backbone.

Table B2: Comparisons on the SEDD-small backbone.

Regarding higher-order solvers like Runge-Kutta or Trapezoidal methods, we would like to politely clarify that we have already included a comparison in Table 3 of our main paper. We directly compare our results on the RADD backbone against the numbers reported in their work. While we are unable to directly integrate LSD+ on top of their methods due to the lack of a public implementation, this comparison demonstrates the competitive performance of our approach.

(Absence of semantic quality metrics: MAUVE scores and entropy)

Thanks for your suggestion. To provide a more thorough assessment of generation quality beyond perplexity, we conduct experiments on the SDTT-KLD backbone, which utilizes Ancestral as the sampler. We report MAUVE, Perplexity and Entropy scores in Table B3, which demonstrates that our LSD+ method usually yields improved generation performance.

Table B3: Comparisons on the SDTT-KLD backbone.

(Minimal Performance gain in case of Image Generation)

Thank you for the comment. To better demonstrate the effectiveness and scalability of our method, we conduct experiments on ImageNet (256x256) using the MaskGIT [4] architecture as the backbone, incorporating the recently proposed advanced Halton [5] sampler as the sampling method. The results are reported in Table B4, which demonstrates that our LSD+ method yields improved generation performance as measured by the FID metric.

Table B4: Comparisons on ImageNet 256x256 (in terms of FID$\downarrow$).

(Integration of LSD with Effect of Temperature/Confidence)

Thank you for this question. Due to the time limit of the rebuttal period, we study the effect of softmax temperature in MaskGIT on ImageNet256x256 for the image generation task. The evaluations are produced on 50000 generated images using the Halton sampler. The reported FIDs are given in Table B5, which indicates that our method could lead to improved performance under the tested temperature conditions.

Table B5: Comparisons on ImageNet 256x256 (in terms of FID$\downarrow$ when NFE=4).

(Integration of LSD with Multi-token unmasking heuristics like in FastDLLMs.)

Thank you for this suggestion. This is a promising direction as the two approaches are highly complementary. LSD enhances the quality of score predictions in few-step sampling, which can directly empower the confidence-based parallel decoding of Fast-dLLM. We report the accuracy and throughput metrics in Table B6, which indicate our methods improve both sampling quality and speed.

Table B6: Integration with Fast-dLLM on LLaDA. We report accuracy (throughput, token/s). Higher is better.

(Integration of LSD with distilled discrete diffusion models e.g., SDTT Checkpoints.)

Following your suggestion, to show the complementarity of model distillation (SDTT) and our sampler distillation (LSD+), we apply LSD+ to an SDTT-KLD checkpoint, with the results presented in Table B3 above, which demonstrates that our LSD+ method generally yields improved generation performance.

(Integration of LSD with Remasking)

Following your suggestion, we also integrate LSD with ReMasking (ReMDM) on the MDLM backbone in Table B7. The results show that LSD effectively learns to work in conjunction with this method, further improving performance.

Table B7: Comparisons on the MDLM backbone.

(Additional metrics using LLaMA-3)

Following your suggestion, we re-evaluate our main results on the SEDD and RADD backbone using Llama-3-8B to evaluate the perplexity. The results are reported in Tables B8 and B9.

Table B8: Comparisons on SEDD-small backbone (in terms of Perplexity$\downarrow$), judged by LLaMA-3-8B.

Table B9: Comparisons on on RADD backbone (in terms of Perplexity$\downarrow$), judged by LLaMA-3-8B.

References:

[1] Gat el al. Discrete Flow Matching, NeurIPS’24

[2] Sahoo et al. Simple and Effective Masked Diffusion Language Models, ICLR’25.

[3] Chang et al. A continuous time framework for discrete denoising models, NeurIPS’22

[4] Chang el al. Masked Generative Image Transformer, CVPR’22.

[5] Besnier et al. Halton Scheduler for Masked Generative Image Transformer, ICLR’25.

Thanks for your useful comments and questions. Our responses to the main concerns are given as follows. All citations refer to the reference list provided by the reviewer.

(Examples of work that probably should be cited and briefly discussed here include but not limited to [2,3,4,5,6,7,16]. )

Thanks for the comment. We will add a paragraph discussing the mentioned papers in the revised manuscript, and we can consider citing others if the reviewer has additional suggestions. The paragraph to be added is as follows:

The distillation of continuous diffusion models is a rapidly advancing field. A prominent direction is related to the consistency model [1], which aims to learn a function that maps any point on an ODE trajectory to its origin, enabling one-step or few-step generation. This paradigm has been extended to multi-step variants [2, 3] for improved performance. Other significant works focus on directly matching student and teacher distributions, such as distilling guided diffusion models [4], proposing simplified and faster matching objectives [5], recursively distilling a deterministic diffusion sampler into a new model [6], or concentrating on one-step distillation [7]. While these methods are highly effective for continuous models, they usually rely on continuous paths in the sense that the sampling process of each step is differentiable. Our work diverges by proposing a distillation framework specifically for the discrete diffusion model, which does not assume such a continuous path, and addressing a different set of challenges like the non-differentiability of the outputs. A recent work for Di[M]O [16] also involves distilling discrete diffusion models. It distills a multi-step masked diffusion model into a one-step generator. This is achieved by training a new student model from scratch, using a sophisticated proxy objective that involves creating "pseudo-intermediate states" and training an auxiliary model to match conditional output distributions. Our approaches are significantly different with Di[M]O in both goal and mechanism. Similarly to LD3 and S4S, we focus on a few-step sampler distillation. We tackle the challenge of non-differentiability of sampling from categorical distributions, and we enhance an existing sampler rather than replacing the model, which avoids the complexity of training a new generator and an auxiliary model.

(The authors probably also need to compare the distillation procedure proposed in this work with that of [16].)

Thank you for highlighting this work. As mentioned above, we will add a detailed comparison with Di[M]O [16] in the Related Work Section of our paper to clearly position our contribution.

While both Di[M]O and our LSD+ method employ distillation for discrete diffusion models, a direct empirical comparison of the outcomes is not straightforward due to several significant differences between the two methods. Di[M]O trains a new student model from scratch to achieve one-step generation, primarily for image generation tasks. In contrast, our method is a few-step sampler distillation process that trains scalar-form coefficients and time steps for an existing sampler, avoiding the need to train a new model, and is evaluated mainly on text generation tasks. Given these differences in objective (one-step vs. few-step), mechanism (model training vs. sampler training), and application domain (image vs. text), currently, we do not include a comparison with Di[M]O.

Additionally, we note that the work for Di[M]O is publicly available on arXiv on March 19th, 2025. In this context, we would like to politely draw attention to NeurIPS policy, which states: “Papers appearing online after March 1st, 2025 are generally considered concurrent to NeurIPS submissions. Authors are not expected to compare to those”. (https://neurips.cc/Conferences/2025/PaperInformation/NeurIPS-FAQ)

(The authors are encouraged to test the proposed distillation method on large-scale images like FFHQ [9] or ImageNet [10] to better justify its validity.)

Thank you for your suggestion. We conduct experiments on ImageNet (256x256) using the MaskGIT architecture as the backbone, incorporating the recently proposed advanced Halton sampler (https://arxiv.org/abs/2503.17076) as the sampling method. The results are reported in Table A1, which demonstrates that our LSD+ method yields improved generation performance as measured by the FID metric.

Table A1: Comparisons on ImageNet 256x256 (in terms of FID$\downarrow$).

(Whether it would be possible to provide theoretical justification for the distillation method proposed in this paper by referring to exist work on the theoretical analysis of discrete diffusion models like [8,13,14,15]?)

Thank you for this question. We agree that providing theoretical justification for our distillation methods would strengthen the work. However, our current focus is primarily empirical: We aim to develop practical and effective techniques for accelerating discrete diffusion models. Given the time constraints of the short rebuttal period, we are unable to establish rigorous theoretical guarantees for our distillation methods at this stage, and we plan to address this with more thorough theoretical analysis in future work.

That said, we would like to offer an intuitive explanation for the effectiveness of our method, which may help lay the groundwork for subsequent theoretical framing. Specifically, the final discrepancy between the outputs of the student and teacher samplers stems from the accumulation of small local errors made at each step. Each of these local errors, in turn, is tied to the specific score predicted by the model at that step. Our training objective directly enforces alignment between the score predictions of the student sampler and those of the teacher sampler at every step. By implicitly correcting these small local errors throughout the process, our approach guides the student sampler to produce final outputs that closely match the high-quality results of the teacher sampler.

(For Algorithm 3 in Appendix D.2, it seems that some ambiguity exists.)

Thank you for your careful reading and for bringing this to our attention. We will revise Algorithm 3 to correct the typos therein. The parameters $\tau_k$'s will be updated accordingly after the $\kappa_k$'s get updated. The update rules for $\kappa_k$'s and $\tau_k$'s are as follows: $\kappa\_k \gets \kappa\_k - \eta \nabla\_{\kappa\_k} \tilde{L}\_k$ $\tau\_k = t\_0 - \sum_{l=1}^{k} \kappa\_l$

(Whether it would be fair to fix the number of evaluations (NFEs) and compare the performance since the distillation step take nontrivial amount of time?)

Thank you for the comment and question. Our primary focus is on accelerating the inference stage, where a one-time and offline training cost is traded for significant, repeated latency gains. The training-stage cost is relatively low. For example, under the experimental setup detailed in the Experiment Section of our main paper, distilling a 1024-step teacher sampler to a 16-step student sampler using the SEDD-small backbone takes approximately 4 minutes on a single NVIDIA RTX4090 GPU. By contrast, inference time is substantially higher. Specifically, with a batch size of 1, evaluating perplexity via the generation of 1024 unconditional text samples requires around 300 minutes of inference time on the same GPU. Given that the one-time, offline training cost is orders of magnitude smaller than even a single moderately sized evaluation, we omit training time for our distillation methods and instead focus on reporting performance metrics with fixed NFEs.

This paradigm is not unique to our work. For example, the closely related acceleration method JYS also compares performance under fixed NFEs, despite involving a one-time training cost to learn time schedules.

Regarding high-order solvers, we have included direct comparisons against the $\theta$-RK2 and $\theta$-Trapezoidal methods from [8] in Table 3 of the main paper. However, since the code for [8] is not yet open-sourced, a direct comparison of running times is currently not feasible.

The reviewer would like to thank the authors for the detailed response, most of which the reviewer is satisfied with. The authors are strongly encouraged to include all points mentioned in the reviews and rebuttals below, such as literature review on distillation for continuous/discrete diffusion models, comparison with concurrent work like Di[M]O, as well as the additional experimental results on large-scale images like ImageNet.

Moreover, the authors might also consider including examples of generated images from CIFAR/ImageNet in the appendix of the manuscript. Given that theoretical analysis seems to be a primary limitation of the manuscript, the authors should probably consider citing the articles [13,14,15] listed in the references above along with [16] to briefly discuss how one may theoretically analyze the distillation methodology proposed here for discrete diffusion models.

The reviewer will finalize the rating soon.

Added references:

[16] Chen, Hongrui, and Lexing Ying. "Convergence analysis of discrete diffusion model: Exact implementation through uniformization." arXiv preprint arXiv:2402.08095 (2024).

Dear Reviewers, the author-reviewer discussion period has been extended to August 8th AoE. Please use this extra time to read the author's rebuttal and engage in discussion with the authors. Thank you.