NI Sampling: Accelerating Discrete Diffusion Sampling by Token Order Optimization

TL;DR

NI Sampling introduces a neural indicator framework for token order optimization in discrete diffusion models, significantly accelerating sampling by up to 14.3 times with minimal accuracy loss.

Contribution

The paper proposes a novel neural indicator sampling method that optimizes token order, enabling faster sampling in discrete diffusion language models without sacrificing accuracy.

Findings

Achieves up to 14.3× acceleration over full-step sampling.

Reduces sampling iterations by an order of magnitude.

Outperforms confidence threshold sampling in accuracy-step trade-off.

Abstract

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.