NI Sampling: Accelerating Discrete Diffusion Sampling by Token Order Optimization

2026-04-20

AI Summary

Discrete diffusion language models suffer from inefficiency due to sampling only a few tokens per step. This work proposes a Neural Indicator Sampling framework that formulates token sampling order as a learnable task, optimizing the generation process through dynamic scheduling and a trajectory-preserving training objective. It presents the first systematic approach to end-to-end learning of sampling order, achieving up to a 14.3× speedup on LLaDA and Dream models with negligible performance degradation. Moreover, the method significantly outperforms heuristic strategies—such as confidence-threshold-based approaches—in the trade-off between accuracy and the number of inference steps.

Technology Category

• Planning, Routing, and Scheduling: Planning with Language Models
• Search and Optimization: Sampling/Simulation-based Search
• Natural Language Processing: Generation

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

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.

Problem

• discrete diffusion
• sampling efficiency
• token order
• parallel decoding
• diffusion language models

Innovation

• discrete diffusion
• sampling order optimization
• Neural Indicator Sampling
• parallel decoding
• trajectory-preserving objective