NI Sampling: Accelerating Discrete Diffusion Sampling by Token Order Optimization

• URL: http://arxiv.org/abs/2604.18471v2
• Date: Sun, 26 Apr 2026 18:45:11 GMT
• Title: NI Sampling: Accelerating Discrete Diffusion Sampling by Token Order Optimization
• Abstract summary: We study the problem of token sampling order optimization and demonstrate its significant potential for acceleration.<n>We propose a Neural Indicator Sampling framework that utilize a neural indicator to decide which tokens should be sampled at each step.<n>Our method consistently outperforms confidence threshold sampling in the accuracy-step trade-off.
• Score: 27.77541110403583
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Self-Rewarding Sequential Monte Carlo for Masked Diffusion Language Models [58.946955321428845]
  This work presents self-rewarding sequential Monte Carlo (SMC)<n>Our algorithm stems from the observation that most existing MDLMs rely on a confidence-based sampling strategy.<n>We introduce the trajectory-level confidence as a self-rewarding signal for assigning particle importance weights.
  arXiv  Detail & Related papers  (2026-02-02T09:21:45Z)
• One-Step Diffusion Samplers via Self-Distillation and Deterministic Flow [24.67443222055996]
  Existing sampling algorithms typically require many iterative steps to produce high-quality samples.<n>We introduce one-step diffusion samplers which learn a step-conditioned ODE.<n>We show that standard ELBO estimates degrade in the few-step regime because common discrete yield mismatched forward/backward transition kernels.
  arXiv  Detail & Related papers  (2025-12-04T20:57:53Z)
• p-less Sampling: A Robust Hyperparameter-Free Approach for LLM Decoding [10.595336643423229]
  $p$-less sampling is an information-theoretic approach to sampling.<n>It dynamically sets a truncation threshold at each decoding step based on the entire token probability distribution.<n>It consistently produces high-quality outputs as temperature increases.
  arXiv  Detail & Related papers  (2025-09-27T10:33:41Z)
• Inference-Time Scaling of Diffusion Language Models with Particle Gibbs Sampling [70.8832906871441]
  We study how to steer generation toward desired rewards without retraining the models.<n>Prior methods typically resample or filter within a single denoising trajectory, optimizing rewards step-by-step without trajectory-level refinement.<n>We introduce particle Gibbs sampling for diffusion language models (PG-DLM), a novel inference-time algorithm enabling trajectory-level refinement while preserving generation perplexity.
  arXiv  Detail & Related papers  (2025-07-11T08:00:47Z)
• Accelerating Diffusion LLMs via Adaptive Parallel Decoding [60.407727995313074]
  We introduce adaptive parallel decoding (APD), a novel method that dynamically adjusts the number of tokens sampled in parallel.<n>APD provides markedly higher throughput with minimal quality degradations on downstream benchmarks.
  arXiv  Detail & Related papers  (2025-05-31T06:10:10Z)
• Quantizing Diffusion Models from a Sampling-Aware Perspective [43.95032520555463]
  We propose a sampling-aware quantization strategy, wherein a Mixed-Order Trajectory Alignment technique is devised.<n>Experiments on sparse-step fast sampling across multiple datasets demonstrate that our approach preserves the rapid convergence characteristics of high-speed samplers.
  arXiv  Detail & Related papers  (2025-05-04T20:50:44Z)
• Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion [61.03681839276652]
  Diffusion Forcing is a new training paradigm where a diffusion model is trained to denoise a set of tokens with independent per-token noise levels.<n>We apply Diffusion Forcing to sequence generative modeling by training a causal next-token prediction model to generate one or several future tokens.
  arXiv  Detail & Related papers  (2024-07-01T15:43:25Z)
• A Block Metropolis-Hastings Sampler for Controllable Energy-based Text Generation [78.81021361497311]
  We develop a novel Metropolis-Hastings (MH) sampler that proposes re-writes of the entire sequence in each step via iterative prompting of a large language model. Our new sampler allows for more efficient and accurate sampling from a target distribution and (b) allows generation length to be determined through the sampling procedure rather than fixed in advance.
  arXiv  Detail & Related papers  (2023-12-07T18:30:15Z)
• Masked Generative Modeling with Enhanced Sampling Scheme [1.3927943269211591]
  Enhanced Sampling Scheme (ESS) ensures both sample diversity and fidelity. ESS consists of three stages: Naive Iterative Decoding, Critical Reverse Sampling, and Critical Resampling. We demonstrate significant performance gains of ESS in both unconditional sampling and class-conditional sampling.
  arXiv  Detail & Related papers  (2023-09-14T09:42:13Z)
• Entropy-based Training Methods for Scalable Neural Implicit Sampler [20.35664492719671]
  In this paper, we introduce an efficient and scalable implicit neural sampler that overcomes limitations.<n>The implicit sampler can generate large batches of samples with low computational costs.<n>By employing the two training methods, we effectively optimize the neural implicit samplers to learn and generate from the desired target distribution.
  arXiv  Detail & Related papers  (2023-06-08T05:56:05Z)
• Towards Automated Imbalanced Learning with Deep Hierarchical Reinforcement Learning [57.163525407022966]
  Imbalanced learning is a fundamental challenge in data mining, where there is a disproportionate ratio of training samples in each class. Over-sampling is an effective technique to tackle imbalanced learning through generating synthetic samples for the minority class. We propose AutoSMOTE, an automated over-sampling algorithm that can jointly optimize different levels of decisions.
  arXiv  Detail & Related papers  (2022-08-26T04:28:01Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.