Search papers, labs, and topics across Lattice.
This paper tackles the inefficiency of token sampling in discrete diffusion language models (dLLMs) by optimizing the order in which tokens are sampled during the reverse diffusion process. They find that prioritizing sampling tokens that are likely to be predicted correctly can drastically reduce the number of sampling iterations needed. To achieve this, they introduce Neural Indicator Sampling (NI Sampling), a framework that uses a neural network trained with a trajectory-preserving objective to determine the optimal token sampling order.
Discrete diffusion models can be sped up by 14x by intelligently choosing which tokens to sample at each step, without sacrificing accuracy.
Discrete diffusion language models (dLLMs) have recently emerged as a promising alternative to traditional autoregressive approaches, offering the flexibility to generate tokens in arbitrary orders and the potential of parallel decoding. However, existing heuristic sampling strategies remain inefficient: they choose only a small part of tokens to sample at each step, leaving substantial room for improvement. In this work, we study the problem of token sampling order optimization and demonstrate its significant potential for acceleration. Specifically, we find that fully leveraging correct predictions at each step can reduce the number of sampling iterations by an order of magnitude without compromising accuracy. Based on this, we propose Neural Indicator Sampling (NI Sampling), a general sampling order optimization framework that utilize a neural indicator to decide which tokens should be sampled at each step. We further propose a novel trajectory-preserving objective to train the indicator. Experiments on LLaDA and Dream models across multiple benchmarks show that our method achieves up to 14.3$\times$ acceleration over full-step sampling with negligible performance drop, and consistently outperforms confidence threshold sampling in the accuracy-step trade-off. Code is available at https://github.com/imagination-research/NI-Sampling.
Microsoft Research