Neural Indicator Sampling (NI Sampling)

• Neural Indicator Sampling is a framework that optimizes token unmasking in discrete diffusion language models using a learned neural indicator to decide which tokens can be revealed in parallel.
• The method computes feature-based confidence scores, enabling multi-token unmasking that reduces diffusion steps up to 14.3× with minimal impact on generation quality.
• It employs a trajectory-preserving training objective with binary cross-entropy, ensuring that early token revelations do not alter the final generation outcome.

Neural Indicator Sampling (NI Sampling) is a sampling order optimization framework for discrete diffusion LLMs (dLLMs). In the formulation introduced in "NI Sampling: Accelerating Discrete Diffusion Sampling by Token Order Optimization" (Liu et al., 20 Apr 2026), the method addresses the inefficiency of heuristic reverse-process samplers by learning which masked positions can be safely unmasked in parallel at each diffusion step. It does so with a neural indicator that estimates whether the current greedy prediction at a position can be revealed without altering the final trajectory, and it is trained with a trajectory-preserving objective. The reported result is acceleration of dLLM sampling by up to $14.3\times$ over full-step sampling with negligible performance drop, with further gains when combined with KV caching (Liu et al., 20 Apr 2026).

1. Problem formulation in discrete diffusion decoding

In autoregressive LLMs, tokens are generated strictly left-to-right: at step $t$ the model conditions on $(y_1,\ldots,y_{t-1})$ to predict $y_t$. Discrete diffusion LLMs instead define a forward process that gradually masks out tokens until the entire sequence is masked, and learn a reverse denoising process to recover the original sequence. At each reverse step, the model may unmask any subset of positions, allowing flexible, even parallel, decoding orders. In principle, dLLMs are therefore not restricted to sequential dependency; in practice, naïve multi-token sampling breaks correlations between positions whose latent states are simultaneously masked, so most dLLM samplers reveal only one token per step (Liu et al., 20 Apr 2026).

The work formalizes token sampling order as an ordered partition

$$(A_1, A_2, \ldots, A_n), \qquad \bigcup_{t=1}^n A_t = \{1,\ldots,N\}, \qquad A_t \cap A_{t'} = \varnothing \ \text{for } t \ne t'.$$

At reverse step $t$, all positions in $A_t$ are unmasked. The optimization target is to minimize $n$, the total number of diffusion steps, while maintaining the same final generation quality. With a pre-trained dLLM $\theta$ and prompt $c$, the objective is written as

$t$0

with the constraint that $t$1 remains within an acceptable tolerance of the full-step baseline.

This formulation is motivated by the limitations of two existing heuristics. Full-step sampling, also called one token per step, selects the single most confident mask position at each step and therefore requires $t$2 steps to generate $t$3 tokens. Confidence-threshold sampling, described as Fast-dLLM, unmasks all positions whose top-1 probability exceeds a fixed threshold $t$4 such as $t$5, but in practice still leaves many positions masked for hundreds of steps. The central observation is that if one could identify and unmask all positions whose current argmax prediction already matches the final output, the number of reverse steps could drop by an order of magnitude.

2. Neural indicator architecture and inference procedure

At each sampling step, the dLLM first computes the token distribution $t$6 over the current masked sequence $t$7. If there are $t$8 masked positions, NI Sampling extracts three feature groups for each masked position $t$9: the Top-$(y_1,\ldots,y_{t-1})$0 predicted tokens and their embeddings, the last-layer hidden state $(y_1,\ldots,y_{t-1})$1 from the dLLM, and the Top-$(y_1,\ldots,y_{t-1})$2 logits $(y_1,\ldots,y_{t-1})$3, from which confidence scores can be computed. A small feed-forward network $(y_1,\ldots,y_{t-1})$4, called the neural indicator, maps these features to a scalar score $(y_1,\ldots,y_{t-1})$5 that estimates whether the current greedy prediction

$(y_1,\ldots,y_{t-1})$6

can safely be unmasked now without altering the final trajectory (Liu et al., 20 Apr 2026).

The reverse-process inference loop is specified in five pseudo-steps:

1. Compute $(y_1,\ldots,y_{t-1})$7 for all positions with one forward pass of the dLLM.
2. Guarantee that at least one token is revealed by applying a fallback sampler, such as highest-confidence threshold or top-1 probability, to select a seed set $(y_1,\ldots,y_{t-1})$8.
3. For each remaining masked position $(y_1,\ldots,y_{t-1})$9, gather features and compute $y_t$0.
4. Reveal all positions $y_t$1 with $y_t$2. If none qualifies, reveal the single highest-scoring $y_t$3.
5. Update $y_t$4 by unmasking those positions with their greedy predictions, and repeat until all positions are unmasked.

The computational comparison is explicit. Full-step sampling requires $y_t$5 steps and therefore $y_t$6 model evaluations. Confidence-threshold sampling also uses one forward pass per step, but multi-token unmasking reduces the number of steps to $y_t$7, with an empirical speedup of $y_t$8–$y_t$9. NI Sampling adds one forward pass per step plus one small $$(A_1, A_2, \ldots, A_n), \qquad \bigcup_{t=1}^n A_t = \{1,\ldots,N\}, \qquad A_t \cap A_{t'} = \varnothing \ \text{for } t \ne t'.$$0-evaluation per masked position, with negligible GPU cost, so the total becomes $$(A_1, A_2, \ldots, A_n), \qquad \bigcup_{t=1}^n A_t = \{1,\ldots,N\}, \qquad A_t \cap A_{t'} = \varnothing \ \text{for } t \ne t'.$$1 dLLM evaluations. For $$(A_1, A_2, \ldots, A_n), \qquad \bigcup_{t=1}^n A_t = \{1,\ldots,N\}, \qquad A_t \cap A_{t'} = \varnothing \ \text{for } t \ne t'.$$2, the reported value is $$(A_1, A_2, \ldots, A_n), \qquad \bigcup_{t=1}^n A_t = \{1,\ldots,N\}, \qquad A_t \cap A_{t'} = \varnothing \ \text{for } t \ne t'.$$3, which is about $$(A_1, A_2, \ldots, A_n), \qquad \bigcup_{t=1}^n A_t = \{1,\ldots,N\}, \qquad A_t \cap A_{t'} = \varnothing \ \text{for } t \ne t'.$$4 fewer than full-step sampling, and the $$(A_1, A_2, \ldots, A_n), \qquad \bigcup_{t=1}^n A_t = \{1,\ldots,N\}, \qquad A_t \cap A_{t'} = \varnothing \ \text{for } t \ne t'.$$5-evaluation cost is reported as less than $$(A_1, A_2, \ldots, A_n), \qquad \bigcup_{t=1}^n A_t = \{1,\ldots,N\}, \qquad A_t \cap A_{t'} = \varnothing \ \text{for } t \ne t'.$$6 of a single dLLM forward.

3. Trajectory-preserving training objective

The training signal for the neural indicator is derived from a reference trajectory generated by full-step sampling. The trajectory is defined as

$$(A_1, A_2, \ldots, A_n), \qquad \bigcup_{t=1}^n A_t = \{1,\ldots,N\}, \qquad A_t \cap A_{t'} = \varnothing \ \text{for } t \ne t'.$$7

where at step $$(A_1, A_2, \ldots, A_n), \qquad \bigcup_{t=1}^n A_t = \{1,\ldots,N\}, \qquad A_t \cap A_{t'} = \varnothing \ \text{for } t \ne t'.$$8 the sampler unmasks positions $$(A_1, A_2, \ldots, A_n), \qquad \bigcup_{t=1}^n A_t = \{1,\ldots,N\}, \qquad A_t \cap A_{t'} = \varnothing \ \text{for } t \ne t'.$$9 and obtains tokens $t$0 for $t$1. At step $t$2, let $t$3 be the partially unmasked sequence. The paper states that the next $t$4 steps can be merged into one if, for every

$t$5

the condition

$t$6

holds. This trajectory-preserving criterion guarantees that unmasking those positions earlier does not change the subsequent trajectory (Liu et al., 20 Apr 2026).

Supervision labels are constructed from a large corpus of trajectories. The procedure is: generate trajectories using an existing sampler, either full-step or threshold, over prompts from ShareGPT with generation lengths $t$7; randomly pick $t$8; compute the maximum mergeable index $t$9 by testing the criterion iteratively; label all positions in $A_t$0 as $A_t$1 and all other masked positions as $A_t$2. The reported training set size is approximately $A_t$3 trajectories.

The optimization objective is binary cross-entropy:

$A_t$4

Training uses AdamW with $A_t$5, $A_t$6, $A_t$7, batch size $A_t$8, for $A_t$9 epochs. The procedure additionally discards positive labels whose $n$0 to simplify learning.

The significance of this construction is methodological rather than purely heuristic. The indicator is not trained to imitate confidence scores; it is trained to predict whether earlier revelation preserves the full-step trajectory. This suggests that the acceleration mechanism is tied to trajectory equivalence, not merely to high local certainty.

4. Empirical performance and accuracy-step trade-off

The evaluation covers LLaDA-8B-Instruct, LLaDA-1.5, and Dream-7B-Base across GSM8K, MATH, HumanEval, and MBPP, using task accuracy or pass@1, average number of diffusion steps $n$1, and tokens/sec including $n$2 overhead (Liu et al., 20 Apr 2026).

On LLaDA-8B-Instruct with $n$3, the reported figures are:

• Full-step: $n$4 steps, $n$5 tokens/s.
• Confidence-threshold: $n$6, $n$7 tokens/s, approximately $n$8 speedup, accuracy drop $n$9.
• NI Sampling: $\theta$0, $\theta$1 tokens/s, approximately $\theta$2 speedup, with the reported accuracy changing from $\theta$3 to $\theta$4.

On MBPP-512, the reported figures are:

• Full-step: $\theta$5 steps, $\theta$6 tokens/s.
• Confidence-threshold: $\theta$7 steps, $\theta$8 tokens/s, $\theta$9.
• NI Sampling: $c$0 steps, $c$1 tokens/s, $c$2, with accuracy changing from $c$3 to $c$4.

The paper states that NI Sampling Pareto-dominates confidence-threshold sampling over all accuracy-versus-step curves. It also reports that combining NI with KV caching, as in Fast-dLLM, yields up to $c$5 end-to-end acceleration. In the reported experiments, the main empirical claim is therefore not only that NI Sampling reduces the number of reverse iterations, but that it does so more efficiently than threshold-based heuristics at comparable quality levels.

5. Ablations, implementation properties, and projected extensions

The ablation results refine the role of model family, training distribution, input features, and indicator size. On Dream-7B-Base, NI still outperforms threshold sampling, with the example on GSM8K-256 reported as $c$6 versus $c$7. An indicator trained solely on GSM8K plus MATH yields better performance on math tasks but worse on code, suggesting that one may train domain-specific indicators for further gains. Removing last-layer hidden states or predicted logits each degrades the trade-off curve noticeably, while Top-$c$8 token embeddings provide a smaller but still visible benefit. Doubling MLP depth from $c$9 layers, from $t$00 parameters, gives no significant gain, so the smaller $t$01 model is retained (Liu et al., 20 Apr 2026).

The practical deployment characteristics are also quantified. The indicator adds only $t$02–$t$03 ms per inference step versus $t$04–$t$05 ms for the dLLM itself, corresponding to $t$06–$t$07 overhead, and it occupies approximately $t$08 of the dLLM’s parameters. The method is described as a drop-in wrapper: only the sampling loop must be modified to include $t$09 and thresholding, while existing dLLM checkpoints and caching optimizations remain compatible.

Several future directions are explicitly identified. A final-results-preserving criterion, in which $t$10 is trained to match the final generation rather than the full trajectory, can achieve up to $t$11 speedup with negligible performance drop, at the cost of theoretical trajectory divergence. Other proposed directions are reinforcement learning, treating $t$12 as an RL agent that chooses action masks to directly maximize final reward; replacing the position-wise MLP with a cross-position Transformer; retraining $t$13 on trajectories produced by NI Sampling itself to mitigate drift and further reduce steps; and scaling the indicator or the trajectory dataset, for example to $t$14 samples, to close the gap to the theoretical upper bound of approximately $t$15 from trajectory analysis.

One terminological point warrants care. The summary characterizes NI Sampling as a general, training-free inference wrapper, while the method also introduces a trajectory-preserving objective to train the indicator. A plausible implication is that "training-free" refers to the underlying dLLM checkpoints rather than to the indicator $t$16, because the reported implementation modifies inference for existing checkpoints rather than retraining the base model.

6. Terminological ambiguity and distinct usage in PINNs

The abbreviation "NI-Sampling" is also used in a different context in "Failure-informed adaptive sampling for PINNs, Part II: combining with re-sampling and subset simulation" (Gao et al., 2023). There, the methodology is not a discrete diffusion decoding method. Instead, it is an adaptive sampling framework for physics-informed neural networks in which failure probability is used as the posterior error indicator, with a truncated-Gaussian posterior model or subset simulation, cosine-annealing re-sampling, a fixed total collocation size $t$17, and early stopping when $t$18 (Gao et al., 2023).

This naming overlap can create confusion because both lines of work involve a form of learned or model-based indication of where computation should be focused next, yet the objects being sampled are different. In the dLLM setting, NI Sampling selects token positions to unmask during reverse diffusion. In the PINN setting, NI-Sampling selects collocation points in the domain according to an error or failure-probability criterion. The shared abbreviation therefore does not imply a shared algorithmic framework.

Within the discrete diffusion literature, Neural Indicator Sampling refers specifically to the token order optimization framework built around the neural indicator $t$19, trajectory-preserving supervision, and accelerated reverse-process decoding.