Title: Learnable Sampler Distillation for Discrete Diffusion Models
Venue: NeurIPS 2025
Accepted: yes
Keywords: Discrete diffusion models

Efficient sampling

Sampler distillation
Abstract: Discrete diffusion models (DDMs) have shown powerful generation ability for discrete data modalities like text and molecules. However, their practical application is hindered by inefficient sampling, requiring a large number of sampling steps. Accelerating DDMs by using larger step sizes typically introduces significant problems in generation quality, as it amplifies the impact of both the compounding decoding error due to factorized predictions and discretization error from numerical approximations, leading to a significant decrease in sampling quality. To address these challenges, we propose learnable sampler distillation (LSD), a novel approach to train fast and high-fidelity samplers for DDMs. LSD employs a distillation approach where a student sampler with a few steps learns to align its intermediate score trajectory with that of a high-quality teacher sampler with numerous steps. This alignment is achieved by optimizing learnable sampler coefficients that adaptively adjust sampling dynamics. Additionally, we further propose LSD+, which also learns time schedules that allocate steps non-uniformly. Experiments across text generation, image generation, and synthetic tasks demonstrate that our proposed approaches outperform existing samplers for DDMs, achieving substantially higher sampling quality with significantly fewer sampling steps.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes LSD and LSD+, methods to accelerate the sampling of discrete diffusion models (DDMs). The main contributions are:\n\n* *Learnable sampling coefficients*: making the scaling factors applied to scores learnable.\n* *Learnable timesteps*: making the sampling schedule learnable.\n\nUsing these techniques, the method achieves comparable or superior performance to baselines (Euler, JYS) with fewer NFEs.\n\nSimply generating similar outputs by increasing NFEs does not enable learning solver parameters. Therefore, the paper introduces loss functions (Eq. 6 and 8) tailored for DDMs. The core idea is to make the scores at the intermediate states $x_k$ (student solver) and $x_k^*$ (teacher solver) similar.\n\nSTRENGTHS AND WEAKNESSES:\n*Strengths*\n\n* The results significantly outperform baseline methods (Euler sampler and JYS). This is impressive.\n\n*Weaknesses*\n\n* Theoretical justification is lacking. Even if the learning of the solver is perfectly successful, is it guaranteed that the same distribution is recovered? It would be sufficient to at least show that it contributes to reducing the gap between the teacher distribution $p^*$ and the student distribution $p$.\n* The method is not intuitively easy to understand. The current objective enforces that the outputs of the student model at $x_k$ match those of the teacher model at $x_k^*$. \n* However, it is unclear why having similar scores at different points is beneficial. My intuition was that the student model should be trained to produce $x_k^*$ itself (as in SDTT \\[1]), though in that case there would be the drawback of sacrificing diversity.\n\n\\[1]: Beyond Autoregression: Fast LLMs via Self-Distillation Through Time ([https://arxiv.org/abs/2410.21035v1](https://arxiv.org/abs/2410.21035v1))\n\nQUESTIONS:\n* As noted in the weaknesses, could you provide a theoretical guarantee or an intuitive explanation of why having similar scores at different locations is helpful?\n* What do the learned coefficients and timesteps look like? As in JYS, are they more concentrated toward $T$? Are the coefficients distributed to be greater than 1 in the early timesteps?\n* Additional baselines: it would be helpful to include comparisons with recent methods such as \\[1] and \\[2]. I would be satisfied even with results applied only to SEDD.\n\n\\[1]: Fast Sampling via Discrete Non-Markov Diffusion Models with Predetermined Transition Time ([https://arxiv.org/abs/2312.09193](https://arxiv.org/abs/2312.09193))\n\\[2]: Fast Solvers for Discrete Diffusion Models: Theory and Applications of High-Order Algorithms ([https://arxiv.org/abs/2502.00234](https://arxiv.org/abs/2502.00234))\n\nLIMITATIONS:\n* The supplementary material acknowledges that the student sampler has fundamental limitations, such as an upper bound on performance.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a learnable sampler for discrete diffusion models (DDMs) that can achieve performance comparable to a multi-step teacher model through distillation. The proposed sampler is optimized using a novel, differentiable loss based on the discrepancy between the intermediate score trajectories of the student and a high-quality teacher. Extensive experiments demonstrate the effectiveness of the method in accelerating DDM sampling.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**\n1. This work addresses a critical and well-known problem in discrete diffusion models: their slow and computationally expensive sampling process. An effective solution to this problem would significantly broaden the practical applicability of DDMs.\n2. The proposed score trajectory loss is intuitive and provides a clever solution to the non-differentiability problem in discrete sampling pipelines. Additionally, the learnable sampler is lightweight with only a few parameters, making the distillation process highly efficient. The empirical results are strong and consistent across multiple tasks and model backbones.\n\n**Weaknesses:**\n1. **Lack of justification for the training objective:** It is not entirely clear why the training objective is effective in practice. The loss only aligns the scores at each step k independently, but fails to account for the dependency of a state on the full history of previous states and sampler coefficients. The student's state $x_{t_k}$ and the teacher's state $x*_{t_k}$ are generated from different trajectories, so simply aligning their scores at that instant does not guarantee a better final sample. Could the authors provide more insights into this phenomenon?\n\n2. **Underexplored design choices:** The form of the learnable sampler coefficients needs further justification. The current sampler introduces scalar-form multiplicative coefficients (`Φ(t_k)`) to modify the student's score. This assumes a single scalar is sufficient to correct for the complex, high-dimensional accumulated error. A pending ablation study comparing this design to more expressive forms, such as vector-form element-wise multiplicative coefficients or other parametric affine transformations, would be valuable.\n\n3. **Lacking theoretical grounding:** The trajectory loss appears to be more of a well-motivated heuristic rather than a principled solution derived from a solid mathematical foundation. While the method is empirically successful, a theoretical analysis (e.g., analyzing the distributional divergence between the optimal student sampler and the teacher sampler) would foster a deeper understanding of the proposed method and its limitations.\n\nQUESTIONS:\n1.  The experiments are conducted on relatively small-scale models (GPT-2 level, CIFAR-10). For larger models, e.g., DiffuLLaMA and DiffuGPT, how well can the proposed method perform? Is LSD or LSD+ compatible with those larger-scale discrete diffusion models?\n2.  The training objective aligns scores at each step independently, even though the student and teacher states are on different trajectories. Could the authors provide more intuition or analysis on why this myopic alignment is sufficient to guide the student sampler effectively over the full generation process?\n3.  Could the authors comment on the choice of a scalar coefficient for the sampler? Have more expressive alternatives (e.g., vector-based coefficients) been explored, and if so, how do they compare in performance and efficiency?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-03"}

comments: SUMMARY:
This paper proposes a framework for accelerating discrete diffusion models by distilling a fast, high-fidelity sampler, named LSD.

In general, there is a trade-off between generation speed and generaion quality. 

LSD trains a student sampler with a small number of steps to align its intermediate score trajectories with those of a high-quality teacher sampler. By introducing learnable, time-dependent coefficients, the student dynamically adjusts the influence of scores to compensate for errors.

 In addition, LSD+, additionally learns non-uniform time schedules, allocating steps adaptively over time.

Experiments demonstrate that LSD/LSD+ significantly improve the trade-off between sampling speed and fidelity across text generation (OpenWebText, GPT-2-based DDMs), image generation (CIFAR-10), and a synthetic sequence task.

STRENGTHS AND WEAKNESSES:
Strength:
- The paper designs a novel distillation framework for fast inference of DDMs.
- The paper Introduces learnable per-step coefficients and time schedules, providing adaptive control over dynamics.
- Experimental results on three diverse tasks with improvements in perplexity, FID.

Weaknesses:
- Evaluations are limited to relatively small-scale datasets (OpenWebText, CIFAR-10) and short sequences (1024 tokens)
- The robustness of the algorithm to parameters has not been discussed.

QUESTIONS:
How sensitive are the results to the choice of the relaxed objective hyperparameter (Hamming distance threshold)? Could tuning it improve results further?

LIMITATIONS:
The authors did not discuss the limitations.
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-05

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work focuses on improving sampling efficiency of discrete diffusion models by optimizing when to query diffusion model i.e., noise level and coefficiencts of solvers, an interesting direction not explored in case of discrete diffusion models.  They demonstrate good performance  in case of text generation with discrete diffusion models at lower NFEs.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n* Optimizing coefficients, solver properties in addition to where to query diffusion/flow model is interesting in context of discrete diffusion models.\n* The paper demonstrates meaningful reductions in NFE while maintaining generation quality, addressing a key limitation of discrete diffusion models.\n* The paper is well-written and accessible to readers.\n\n\n**Weakness**\nLimited Scope of Experimental Evaluation\n\nThe paper lacks comprehensive analysis of how different solver configurations perform under varying conditions, limiting understanding of when and where LSD can be effective and relative peroformance boost. \nAs paper lacks any significant analysis or theoretical formulation too and authors point out is similar to S4S but applied to discrete setting. \n*Lack of extensive empericial analysis and settings* limit potential insights and adoption guidelines for readers and broader community.\n\n- Currently this work does not compare against several important recent methods and settings which focus on quality vs latency of discrete diffusion models.\n- Currenlty, authors do not consider predictor-corrector solvers (like in DFM) nor higher order solvers like Runge Kutta or Trapezoidal though they cite relevant work.\n- How effective LSD is with other heuristics in multi-token sampling like for e.g., in FastdLLMs and also setting with Remasking during sampling which shows enhanced generation quality.\n- Absence of semantic quality metrics: No evaluation using ModernBERT-based MAUVE scores or similar metrics that assess semantic coherence and diversity (entropy)\n- Minimal Performance gain in case of Image Generation (but didn't consider this much for rating)\n\nQUESTIONS:\nFollowing Weakness\n\nIt would be informative for readers and broader community to see integration of LSD with\n- Effect of Temperature/Confidence\n- Multi-token unmasking heuristics like in FastDLLMs, etc.\n- How well LSD complement and effective speedup for distilled discrete diffusion models e.g., SDTT Checkpoints\n- ReMasking\n- Predictor-Corrector Solvers\n\n\nAlso additional metrics like use state-of-the-art language models (LLaMA-3 or alternatives) to evaluate generated text quality independently of the diffusion model.\n\nLooking forward to rebuttal.\n\n\n**References** \n\n1. G. Wang et al. Remasking Discrete Diffusion Models with Inference-Time Scaling (https://arxiv.org/abs/2503.00307)\n2. J. Deschenaux et al. Beyond Autoregression: Fast LLMs via Self-Distillation Through Time (https://arxiv.org/html/2410.21035v1)\n3. C. Wu et al. Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding (https://arxiv.org/html/2505.22618v1)\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposed learnable sampler distillation (LSD), which is a novel approach to train fast and high-fidelity samplers for discrete diffusion models. The training method is based on the teacher-student model, where the student sampler is trained via distillation by optimizing learnable sampler coefficients with respect to a fixed teacher sampler of high fidelity. An extension of LSD, which is named as LSD+, is also proposed to learn non-uniform time schedules in an adaptive way. Numerical experiments on the generation of synthetic data, text and images are provided to justify the effectiveness of the proposed method.\n\nSTRENGTHS AND WEAKNESSES:\nPros: The reviewer finds that the manuscript is presented in a clear and straightforward way. Also, this paper studies an important question, since how to accelerate the inference time is currently one of the bottlenecks of discrete diffusion models. Experiments on multiple generation tasks are provided to justify how effective the proposed methodology is. In addition, a wide range of baselines and implementation details is also provided. \n\nCons:  The reviewer's major concern is that the literature review part of the manuscript seems to be incomplete. To the best of the reviewer's knowledge, lots of paper have already studied the idea of combining knowledge distillation with continuous diffusion models or its variants like latent diffusion models and consistency models [1] (In fact, the main idea behind consistency model is also similar to that of distillation, as both methodologies essentially use training techniques to speed up the inference process). Examples of work that probably should be cited and briefly discussed here include but not limited to [2,3,4,5,6,7,16]. More essentially, the authors probably also need to compare the distillation procedure proposed in this work with that of [16] (Or maybe consider including [16] as one of the baselines).  \n\nAlso, compared to existing work on accelerating the inference of discrete diffusion models like [8], the scale of the numerical experiment on image generation included in this manuscript might be a bit small. Specifically, 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.\n\nQUESTIONS:\nFor the setting of distilling continuous diffusion models/consistency models, a few papers like [11, 12] have provided theoretical justifications. The reviewer is wondering 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]?\n\nMoreover, for Algorithm 3 in Appendix D.2, it seems that some ambiguity exists. Specifically, as the time schedules $\\tau_{k}$'s depend on the stepsizes $\\kappa_{k}$'s, would it be possible for the authors to clarify whether the stepsizes $\\kappa_{k}$'s will be updated accordingly after the $\\tau_{k}$'s get updated in line 10 of Algorithm 3? (It seems that it doesn't make sense to fix all stepsizes $\\kappa_{k}$'s)\n\nFurthermore, for the comparison between LSD/LSD+ and other baseline methods, the reviewer is wondering 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. Given that LSD and LSD+ both accelerate the inference via training, maybe the authors need to fix the running time and compare LSD/LSD+ with other methods like [8]?\n\nLIMITATIONS:\nThe reviewer is satisfied with the limitations that the authors briefly discussed in Appendix A of the manuscript. For instance, one limitation that the authors mentioned is that the performance ceiling of the student sampler is highly related to the quality of the teacher samplers.\n\nReferences:\n\n[1] Song, Y., Dhariwal, P., Chen, M., & Sutskever, I. (2023). Consistency models. International Conference on Machine Learning (ICML), PMLR, 2023.\n\n[2] Boffi, N. M., Albergo, M. S., & Vanden-Eijnden, E. (2025). How to build a consistency model: Learning flow maps via self-distillation. arXiv preprint arXiv:2505.18825.\n\n[3] Heek, J., Hoogeboom, E., & Salimans, T. (2024). Multistep consistency models. arXiv preprint arXiv:2403.06807.\n\n[4] Meng, C., Rombach, R., Gao, R., Kingma, D., Ermon, S., Ho, J., & Salimans, T. (2023). On distillation of guided diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) (pp. 14297-14306).\n\n[5] Zhou, Z., Chen, D., Wang, C., Chen, C., & Lyu, S. (2024). Simple and fast distillation of diffusion models. Advances in Neural Information Processing Systems (NeurIPS), 37, 40831-40860.\n\n[6] Salimans, T., & Ho, J. (2022). In The Tenth International Conference on Learning Representations (ICLR), 2022.\n\n[7] Xie, S., Xiao, Z., Kingma, D., Hou, T., Wu, Y. N., Murphy, K. P., ... & Gao, R. (2024). Em distillation for one-step diffusion models. Advances in Neural Information Processing Systems (NeurIPS), 37, 45073-45104.\n\n[8] Ren, Y., Chen, H., Zhu, Y., Guo, W., Chen, Y., Rotskoff, G. M., ... & Ying, L. (2025). Fast solvers for discrete diffusion models: Theory and applications of high-order algorithms. arXiv preprint arXiv:2502.00234.\n\n[9] Karras, T., Laine, S., & Aila, T. (2019). A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) (pp. 4401-4410).\n\n[10] Deng, J., Dong, W., Socher, R., Li, L. J., Li, K., & Fei-Fei, L. (2009, June). Imagenet: A large-scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition (CVPR) (pp. 248-255). \n\n[11] Jain, N., Huang, X., Ma, Y., & Zhang, T. (2025). Multi-Step Consistency Models: Fast Generation with Theoretical Guarantees. arXiv preprint arXiv:2505.01049.\n\n[12] Chen, Y., Zhang, Y., Oertell, O., & Sun, W. (2025). Convergence of consistency model with multistep sampling under general data assumptions. arXiv preprint arXiv:2505.03194.\n\n[13] Ren, Y., Chen, H., Rotskoff, G. M., and Ying, L. How discrete and continuous diffusion meet: Comprehensive analysis of discrete diffusion models via a stochastic integral framework. In The Thirteenth International Conference on Learning Representations (ICLR), 2025\n\n[14] Zhang, Z., Chen, Z., & Gu, Q. (2024). Convergence of score-based discrete diffusion models: A discrete-time analysis. In The Thirteenth International Conference on Learning Representations (ICLR), 2025\n\n[15] Liang, Y., Huang, R., Lai, L., Shroff, N., & Liang, Y. (2025). Absorb and Converge: Provable Convergence Guarantee for Absorbing Discrete Diffusion Models. arXiv preprint arXiv:2506.02318.\n\n[16] Zhu, Y., Wang, X., Lathuilière, S., & Kalogeiton, V. (2025). Di $\\mathtt {[M]} $ O: Distilling Masked Diffusion Models into One-step Generator. arXiv preprint arXiv:2503.15457.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-08"}