NI Sampling: Accelerating Discrete Diffusion Sampling by Token Order Optimization
Authors: Enshu Liu
Xuefei Ning
Yu Wang
Zinan Lin
Venue: ICLR 2026 Poster (2026)
Keywords: Discrete Diffusion Sampling; Neural Indicator
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose a general framework for sampling order optimization of discrete diffusion models by using a neural indicator.

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.

Published as a conference paper at ICLR 2026
NI SAMPLING: ACCELERATING DISCRETE DIFFU- SION SAMPLING BY TOKEN ORDER OPTIMIZATION
Enshu Liu," Xuefei Ning, Yu Wang, Department of EE, Tsinghua University les230ma i 1s. ts inghua edu.co foxdoraame@@gail com yu wangemail wtw.ll.sssnnauua edu. on
Zinan Lin Microsoft Research z nanl inêmi mnmicrrosott.com
ABSTRACT
Discrete diffusion language models (dLLMs) have recently emerged as a promis- ing alternative to traditional autoregressive approaches, offering the flexibility generate tokens arbitrary orders and the potential of parallel decoding. How- ever, existing heuristic sampling strategies remain inefficient: they choose only small part of tokens to sample at each step, leaving substantial room for improve- ment. 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 sam- pling iterations by an order of magnitude without compromising accuracy. Based on this, we propose Neural Indicator Sampling (NI Sampling), general sam- pling 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 models across multiple benchmarks show that our method achieves up to 14.3 acceleration over full-step sampling with negligible performance drop, and consis- tently outperforms confidence threshold sampling in the accuracy-step trade-off. Code is available at https: / com/ imag ietaginatiin www.inattnonnnnenneeeeeee NI- St--Sampling ing.
LLaDA and Dream
INTRODUCTION
Diffusion- based large language models (dLLMs) (Austin et al., 2021; Lou at., 2023; Sahoo al., 2024) are a recently emerged paradigm of text generative modeling, which have attracted increasing attention (Nie et al., 2024; Hoogeboom et al., 2021; He et al., 2022; Reid et al., 2022; Sun et al., 2022; Nie et al., 2025; Ye et al., 2025; Wu et al., 2025). Currently, dLLMs have demonstrated pratical impact in large-scale systems like Mercury (Khanna et al., 2025), Gemini Diffusion (Deep- Mind), and Seed Diffusion (Song et al., 2025), highlighting their potential for both deployment and methodological innovation. Similar to continuous diffusion models (Sohl-Dice bocl-Dicksttii et al., 2015; Ho et al., 2020; Song al., 2020) for image generation, which recover data from noise through reverse denoising process, dLLMMs start from a prior sequential discrete distribution and progressively trans- form into the target distribution. Compared traditional auto-regressive large language models (AR LLMs), dLLMs are not restricted to left-to-right decoding scheme and can flexibly choose the generation order of all tokens in the sequence. Additionally. dLLMs' ability of sampling multiple tokens per step suggets the potential to surpass AR LLMs in efficiency.
As mentioned above, the sampling order of tokens should be considered for dLLMs, i.c., how to determine which tokens to sample at each step. Since sampling multiple tokens simultaneously may break inter-token dependencies (Liu et al., 2024a), default samplers adopt the conservative strategy of generating only one token at each step (Nie et al., 2025; Ye et al., 2025; Kim et al., 2025),
Work mostly done during Enshu Liu's internship at Microsoft Research Project advisor: Zinan Lin

---

Published as a conference paper at ICLR 2026
denoted as full-step sampling. However, this leads to large number of sampling steps and results in significant inefficiency. Fast-dLLM (Wu et al., 2025) introduces a heuristic but effective method called confidence threshold sampling, where at each step all tokens whose predicted probabilities exceed a threshold e are unmasked simultaneously. This strategy substantially reduces the number of sampling steps while enabling trade-off between efficiency and accuracy by varying threshold.
In this paper, we observe that the efficiency of existing heuristic sampling strategies for dLLMs are sub-optimal and still leave substantial room for improvement. Specifically, existing empirical methods typically unmask a token only when its predicted confidence is large enough (Wu et al., 2025: Nie et al., 2025; Kim et al., 2025; Ye et al., 2025). However, we find that at each step, the model is actually capable of correctly predicting a large number of tokens, as the token with the highest probability aligns with the final generated token. Existing methods can reveal only a small subset of these tokens, as the predicted confidence for most positions remains below the threshold, resulting underutilization of the model's predictions each step and, consequently, a substantially larger number of sampling steps. Intuitively, if we could identify and sample all correct tokens at each step, the total number of generation steps could be significantly reduced.
Based on this insight, we propose Neural Indicator Sampling (NI Sampling), new and general framework to optimize the token sampling order in dLLMs. Our approach introduces a lightweight neural indicator, which determines whether the currently predicted token should be sampled for all masked positions in the form of a binary classification task. At each sampling step, all tokens judged positive by the neural indicator revealed, allowing the sufficient utilization of model predictions. To train this indicator, we propose to label the generated data in trajectory way, which ensures that the optimal indicator can accurately maintain the original high-quality, albeit inefficient, generated trajectory with a much faster sampling speed. Note that this framework not limited on our training strategy, leaving room for further exploration.
Our main contributions can be summarized as follows:
In Sec. 3, we demonstrate that selecting appropriate token sampling order has the potential to yield substantial acceleration. Specifically, we identify the phnomenon that the inference results of the dLLM at each step are not fully utilized. We show that if all masked positions labelled as positive according to our trajectory-preseving criterion are unmasked at every step, the model can achieve up to 24x faster than default sampling and more than 3x faster than confidence threshold sampling (Wu et al., 2025), while perfectly preseving the performance of the default full-step method. This observation highlights the remarkable space for step compression and motivates us to optimize the sampling order.
In Sec. 4, we introduce NI Sampling, a general framework for optimizing the token sampling order in dLLMs. Concretely, NI Sampling trains a lightweight indicator that makes token-wise wise binary decisions on whether each masked position should be revealed at every step. All pre- dicted tokens judged as positive are sampled simultaneously. Following Sec. 3, we leverage the criterion to generate supervision signals to train this predictor. Importantly, the predictor is generic and trained once for different tasks.
trajectory-preseet
In Sec. 5, we apply NI Sampling on LLADA-SB: -B-A-rrrtttt (Nie et al., 2025), LLaDA-1.5 (Zhu et al., 2025), and Dream-7B Pr-ase model (Ye al., 2025) and evaluate mathematical and code datasets. Compared with the full-step sampling baseline, NI Sampling achieves up to 14.3x speedup with only negligible performance degradation, outperforming confidence thresh- old sampling significantly. Additionally, our accuracy-step trade-off consistently dominates that of confidence threshold sampling across all settings. By combining with KV caching technique (Wu et al., 2025), we achieve up to 25.0x acceleration compared to full-step sampling.
2 PRELIMINARY
2.1 DISCRETE DIFFUSION MODELS
Given a discrete random variable with a finite support X' = (I....,,],, discrete diffusion models gradually perturb its distribution pdata into a prior distribution pr through a predefined forward dif- ferential equation dp: Qipt. over the interval E [o. T] (Campbell et al., 2022; Lou al., 2023). To recover Pdata. they solve the corresponding reverse differential equation, dp. = Q.P.. from T di
de
back to 0, where Q, denotes the reverse diffusion matrices parameterized by neural network.
2

---

Published as a conference paper at ICLR 2026
Token to Sample
Remain Masked
Indicator Score
0.9
Indicator Score Threshold
0.7
Probability
Probability
0.9
Confidence Threshold
Indicator
dLLM
dLLM
M
M
M
Highest Confidence
Confidence Threshold Sampling
NI Sampling (Ours)
Sampling (full-step)
Comparison between NI Sampling and previous sampling methods. At each sampling step, NI Sampling uses a neural indicator to assign scores to masked positions. Tokens with suffi- ciently high indicator scores are then sampled.
Figure 1:
Masked Diffusion Models. Among the various design choices for discrete diffusion model, set- ting pr delta distribution the mask token [MASK] and using pa (x to parameterize Qi emerges the most widely adopted approach, known masked diffusion models (MDMs). Specifically, MDMs treat the forward process as randomly masking tokens:
g(ι,|zv) = II ((<)|=a)= II Cat(af: (1- 04) + (
(1)
where n is the sequence length, / E [0, 1] denotes the diffusion timestep, and ar specifies the noise schedule. MDMs then train the model predict the conditional distribution of zo given noisy sequence T,, resulting in the following training loss:
[())= Pe.mo.m are (we) Σ lst [MASK]] log (2/(
where zo is sampled from the training set and 1 refers to the indicator function.
2.2 SAMPLING PROCESS OF MDMS
well-trained MDM samples Tp by reversing the forward process Eq. (I), i.c., by iteratively un- masking from the wholely masked sequence æ1. At each step, set of masked positions is first selected to be revealed, and tokens are then sampled according to the predicted conditional proba- bilities at these positions. Given the prevalent use of greedy sampling in prior works (Ye al., 2025; Nie et al., 2025; Zhu et al., 2025; Wu et al., 2025), which takes the token with highest probability as the sampling result, we adopt this setting in our paper.
Sampling multiple tokens will cause disalignment between the reverse and forward distributions. For example, for two masked positions i1. i2 some step f, the predictions of the MDM these positions po (z)² (æ:)) and po(2 (x)), do not account for each other's sampling results, since both x/² and xt²² are masked. Consequently, sampling them simultancously would neglect correlations (Liu et al., 2024a; Wu et al., 2025). Therefore, typical methods restrict sampling to a single token each step. Various strategies are proposed select the masked token position i" to unmask at each step. Top-1 probability (Nie et al., 2025; Chang et al., 2022) selects the token with the highest probability: = arg ( of tokens in the codebook. Top-1 probability margin (Kim et al., 2025) selects the position with the largest gap between the top-1 and the top-2 probabilities: i" arg (Hle= Ipo(z) ii (a.)) - pa (af = //2 ),) where ji and ja are tokens with top-1 and top-2 probabilities at position respectively. Top-1 entropy (Ye et al., 2025) selects the position with the highest entropy: i* arg max; Po= jl j(x;) log pa j|
max.)(1,22....V) Pe (2) = jlæ,), where V is the number
The aforementioned methods suffer from slow generation due to the large number of steps. Con- fidence threshold sampling (Wu et al., 2025) is proposed to address this issue. At timestep t, samples all tokens with probability exceeds threshold {i] max; po(=) j(æ,,) e). It can be proved that with sufficiently high threshold, the correspondence loss from sampling multiple tokens is negligible under greedy decoding (Wu et al., 2025). In practice, e typically set to 0.9.
3

---

Published as a conference paper at ICLR 2026
Algorithm Counting Mergeable Steps
Algorithm 2 Merge along the trajectory
Require:
Require:
A trajectory т defined as Eq. (4); Step & of the trajectory and æu: Pre-trained dLLM 0.
A reference trajectory т defined as Eq. (4) with steps; Pre-trained dLLM 6.
1: idæ <- +++1
1: step t- L, Tnew f- () // initialize the new trajectory with an empty list
2: while True do
3:
(E()..... E Aidz; JE Po
2: while step S n do
j|æx)) = x. then
3: idz +-run Alg. with 7,step (as the k),
4: idx + idr +
and 0.
5: else
4: Add Astep U... U Aide-1 Tnew
6: break
5: step t- idæ
end if
6: end while
8:
end while
7: refurn Trew
9:
return idæ
REVEALING THE SIGNIFICANT ACCELERATION POTENTIAL OF SAMPLING ORDER OPTIMIZATION
In this section, demonstrate the substantial potential in reducing the number of sampling steps of optimizing the token sampling order. In Sec. 3.1, we formalize the optimization problem and space of selecting token sampling order. In Sec. 3.2, we show that selecting the sampling order trajectory- preseervig manner can unlock significant acceleration ratio, thereby highlighting the value and importance of this optimization dimension.
3.1 DEFINITION OF SAMPLING ORDER OPTIMIZATION
As discussed in Sec. 2, at each step we need to determine which positions to unmask. Since MDMs do not remask tokens that have already been sampled during the sampling process, the positions selected at different sampling steps do not overlap. Therefore, sampling order optimization can be formulated dLLM 0, user-specified prompt c, and a generation length N, our goal is to find a sequence of sets A (Ai...., An) (fi1,15..... i1,,,).... J...., (t..1....,,,............ }) that maximizes the reward: max (((Acc,))
ordered partitioning of all token positions. Specifically, given pre-trained
A
(3)
ACOPan(S)
where is the set (1, , N); OPart refers to the ordered partition operation, giving OPart(S) {(A)...., An) U A, = S, Vt A; * 0, Vt, + to Ar, 7 Ar, = 0}: R(- c, denotes reward function defined over ordered partitions of S, which depends on model parameters 9 and user prompt In general, there is no limitation to the choice of R, such as the performance on specific downstream task. In our scenario, we aim to minimize the number of sampling steps n while maintaining the desired performance.
3.2 TRAJECTORY PRESERVING PRINCIPLE FOR ORDER SELECTION
Given a pre-trained model we can first apply an existing method (e.g., those described in See. 2.2) to generate a trajectory and token order (Ai.... An). We define the trajectory as follows:
T T= (L(ii,,,2) x. f1.1 ),................ (i1, Ie f1.0) ......, (ix,,, ")""), (in,an, iic. {1,
Ie trron }}),
(4)
where {is.1..... a } Aş is the set of indices of the selected positions, Z. the sampled token at position i....
V} is
We consider the k-th step of the trajectory, where the dLLM originally samples a tokens at positions fg.1.... ak.ag- At this point, consider the model's predictions at the selected positions in the next step At+1. If the model has already predicted all tokens As+1 correctly, i.e.,
Vi E As+1; arg max Po(z) j(æı) x'.,
(5)
JE(I....VV
then merging the k-th and (k: + 1)-th steps does not change the following part of the trajectory. In other words, all tokens from the original k-th and (k + 1) steps can be sampled simultancously in only one step. For the step k. this merging process can be repeated iteratively until a new step Ak-s does not meet the condition in Eq. (5), which point the merging stops. The algorithm for merging a trajectory starting from step le is presented in Alg. 1.
This merging operation can reduce the number of steps, but it requires the model to possess strong "jump-step" prediction capability. So, how well do current pre-trained dLLMs perform in this re- gard? We conducted validation experiment to reveal their potential. We firstly obtain reference
4

---

Published as a conference paper at ICLR 2026
Reference Trajectory steps)
token to sample at coment trajectory braak
B
D
Input
[M]
[M]
M
[M]
[M]
[M]
Step
Q
Inpui
B
[M]
[M]
[M]
Merged Trajectory
Output
B
Q
D
(3 steps)
Input
B
IMI
MII
Stap
8
F
D
Preserving-OOddr. At step I, tokens A, B, and C are predicted correctly and follow the order of the reference trajectory, so they can be sampled. At step 2, token D, which is the next token in the reference trajectory. can be sampled. Although token F is predicted correctly, it cannot be sampled because its preceding token the reference trajectory (E) incorrect. At step 3, token E and F are predicted correctly and can be sampled.
Figure 2: An example of
Table I: Performance of Trajectory srijectory-Preservingg Preserving Order.
Model
Method
GSM8K-256
MATH-256
MBPP-256
Step
Step
Step
Full-step
77.63
256
31,89
256
36,60)
256
LLaDA 8B. Instmatt
Threshold
77.33
(3.4x)
31,68
95.98 (2.7x)
37.60
33.48 (7.6x)
Traj- Prj-Preserving
77.63
22.36 (11.4))
31.89
27.72 (9:2x)
36.60
10.52 (24.3x)
Full-step
81.05
256
33.18
256
38.40
256
LLaDA-L55
Thresbold
(3.5x)
95.41 (2.7x)
(6.1x)
Traj- Traj-Preserving
81.05
22.31 (11.5%)
33.18
28.39 (9.0x)
38.40
11.28 (22.7x)
trajectory generated with the top-1 probability method with full-step (Chang et al., 2022; Nie et al., 2025). Starting from the first step of this trajectory, we iteratively perform step merging according to Alg. 1. When a merging round terminates due to mismatch between the predicted tokens and the reference tokens, the next merging round begins from the step where the mismatch occurred. The loop continues until all steps in the original trajectory have been merged. We denote this strategy as Trajectory-Prrservinnn-Ooo
with the full algorithm presented in Alg. 2 and Fig. 2.
Then, validate this approach with LLaDA-SB-Instruut and LLaDA-115 models GSM8K, MATh, and MBPP with generation length 256. The average numbers of steps required for the new trajectory, along with the performance, are shown in Tab. 1, Compared with confidence thresh- achieves more than 3x acceleration. When compared
old sampling, Dnajectory-Preservnng--0 with the original top-1 probability trajectory with full-step, it can accelerate up 24.3x. More- over, this ordering strategy clearly preserves all sampling results from the reference trajectory, while confidence threshold sampling may incur performance degradation.
It is worth noting that aggressive strategies may yield even greater improvements (e.g.. 36.8x speedup. See App. for details). We mainly discuss Senjectory-Preservv Order because it strictly guarantees consistency between the new outputs and the reference trajectory. While this constraint limits its potential for achieving larger speedups, the acceleration it provides is already highly attractive.
is only one of the possible strategies. More
Although these methods cannot be applied directly due to the absence of a reference trajectory, the results still demonstrate promising space for sampling order optimization, which we study next.
4 NI SAMPLING: A GENERAL FRAMEWORK FOR SAMPLING ORDER OPTIMZIATION
In this section, we introduce Neural Indicator Sampling (NI Sampling), a general framework designed to address the optimization problem defined in Eq. (3). As discussed in Sec. 4.1, the key idea is to employ a neural indicator that evaluates every masked position and determines whether it should be sampled each step. We further explain the training procedure in Sec. 4.2, and describe its input features and the model architecture of the neural indicator in Sec. 4.3.
4.1 A LIGHT NEURAL NETWORK AS TOKEN- TOSEN-WISE SAMPLING INDICATOR
We view Eq. (3) as the problem of determining which positions should be selected for sampling at each step. To this end, we propose a token-wise neural indicator, denoted as o, to make decisions. Specifically, after the inference of the dLLM, suppose there are M masked positions remaining. The neural indicator treats the decision of whether to reveal each masked position binary classifi-
5

---

Published as a conference paper at ICLR 2026
Reference
sample the current step
parallel-samplebble taken
Trajectory
A
B
D
C
token cannot be sampled due to trajestory break
Input
[M]
[M]
[M]
[M]
Step
Output
A
Input
Indicator
GT labels
Input
A
[M]
[M]
[M]
Step
Output
A
D
Input
A
B
[M]
[M]
Step
Output
A
B
D
Input
A
B
[M]
Indicator
Step
Training
Output
A
B
Figure 3: Generating training data for the indicator. At each step of the reference trajectory, labels for all tokens are assigned according to Alg. 1. For example, at step 3, token C, which is generated in the reference trajectory, is labeled positive; since the next token D is also predicted correctly, is labeled positive well. At step 2, token B, which is originally generated by the reference trajectory at this step, is labeled positive. The predicted token Q is incorrect (the correct one should be C) and therefore receives negative label. Although token D is predicted correctly, it is assigned negative label because its preceding token in the reference trajectory is not predicted correctly at this step.
cation task, producing M scores 81,-.., 8M- These indicator scores can be viewed new type of confidence, serving as the criterion for decision-making. For instance, one can simply select the token with the highest score. Since an additional network employed to process the current states, it offers greater flexibility than directly using the model outputs as confidence, and the latter can in fact be regarded as special case of our method. For the goal of acceleration, we also introduce a threshold parameter e like confidence threshold sampling (Wu et al., 2025), and reveal all positions i such that &; 2 6s- This mechanism enables a natural trade-off between speed and accuracy. To guarantee that at least one token is sampled, we first apply an existing sampling strategy to select subset of tokens, and then use the neural indicator to further select among the remaining positions. The complete procedure for sampling with NI Sampling is summarized in Alg. 4. Besides, NI Sam- pling is also compatible with random sampling, with details in App. C. Since the neural indicator trained with additional supervision signals, it encodes richer information than the raw probability vector alone, and thus has the potential to yield better sampling decisions. Note that we constrain the parameter size of the indicator to make sure the additional computational cost negligible, with detailed results shown in Sec. D.2.
4.2 TRAINING THE NEURAL INDICATOR WITH TRAJECTORY PRESERVING PRINCIPLE
In this section, we describe the training procedure for the neural indicator. Following Sec. 3.2, we adopt a
criterion to construct supervisory signals for the predictor, as it ensures that the reference trajectory is fully preserved when trained to optimality, which makes the training more stable. Specifically, we first use the pretrained dLLM to generate dataset with trajectories т. During training, we sample a trajectory from the dataset, denoted
as Td = (f(i1,1. ")" "), J..... (f1,)) xx (1 )}.... [(r,,1 24 2/221)) ",... (In.an.
x. 2lniin) }}}) Then we ran- domly mask it along the trajectory. Concretely, we randomly select an integer t' O.... 1, then replace all positions ig with
by the mask token, and denote the set of these positions as M (ie.m >>> t', Vm). Starting from this step, we follow the principle of in Alg. to determine whether subsequent steps can be merged. All positions in mergeable steps are assigned label of I, while the others are assigned O. We then collect these labels along with the input features in all masked positions, and train the indicator with cross-entropy loss. The algo- rithm is summarized in Alg. 3. Note that the data generation process can be done with any sampling method, making NI Sampling compatitable with all existing samplers.
4.3 DESIGN DETAILS OF THE NEURAL INDICATOR
In this section, we provide details of the neural indicator, including its inputs and architecture.
Inputs. The model should have explicit access at least two types of information. First, it should know the token that will be sampled at the position (i.e., the token with the highest probability under the greedy setting), since knowledge of the sampling result is necessary before deciding whether

---

Published as a conference paper at ICLR 2026
to retain it. Second, the model should have sufficient contextual information, as the appropriate decision for the same token may vary depending on the surrounding context. Inspired by existing sampling methods, dLLLM's confidence in its current prediction. This at least allows the neural indicator to reproduce existing methods. Moreover, beyond the top-1 predicted token at each position, the probabilities of other tokens can also provide valuable semantic information. Accordingly, for each position, the predictor receives the following imputs from each masked positions:
find it beneficial to provide the indicator with an explicit measure of the
The embeddings of the top-K, probability tokens, since the top-1 token is the sampled result and top-2 to K1 tokens provide additional information.
The last-layer hidden states, which obtained rich global information across the sequence.
The to logits of the current output, to help the indicator assess the dLLM's confidence.
Architecture. For simplicity, we adopt position-wise MLP architecture, where each masked po- sition is processed independently without explicit interaction between positions. This is feasible because contextual information is already embedded in the inputs. Different types of input features are first processed by separate linear layers, then concatenated and fed into the backhone of the indicator. The backbone consists of several stacked blocks, each comprising two linear layers, an activation function, and residual connections. Finally, a projection head maps the backbone outputs to 2-dimensional logits, followed by softmax to produce the output score for position i.
EXPERIMENTS
5.1 EXPERIMENTAL SETUP
Base Model and Benchmark. We mainly evaluate NI Sampling with LLaDA-SB-Intruec (Nie et al,, 2025) and LLaDA-1.5 (Zhu et al., 2025). To further demonstrate the generalization ability of our method across different model families, we additionally apply it on Dream-7B-Baa (Ye al., 2025), a model fine-tuned from an AR LLM. For benchmarks, we follow Wu et al. (2025) and select GSM8K (5-shot), MATH (4-shot), HumanEval (0-shot), and MBPP (3-shot) as evaluation datasets. We also report results under varying generation lengths, including 128, 256, and 512 tokens.
Baselines and Evaluation. We report the speedup of our method compared full-step sampling and use confidence threshold sampling as the baseline. To measure sampling efficiency, we report both the average number of steps and tokens per second. Note we have already taken the overhead of the neural indicator into account when calculating tokens per second. We further present the detailed overhead of the predictor in Sec. D.2, along with additional implementation details in Sec. D.5.
Setup of NI Sampling. We use user inputs from ShareGPT dataset as prompts and generate 204k trajectories under three different generation lengths (128, 256, and 512) to train the neural indicator. Then we test this neural indicator across all evaluation settings, demonstrating its generality across tasks and generation lengths. For the LLaDA family of models, we adopt confidence-threshold sampling with the threshold fixed at 0.8 to efficiently generate training data. For the Dream model, we find that using full-step trajectories yields better performance. Additional details in Sec. D.5.
5.2 MAIN RESULTS
We present the results of NI Sampling with the LLaDA family of models in Tab. 2. As shown in the table, compared to full-step sampling, NI Sampling achieves up to around 15x speedup, while the performance degradation of NI Sampling is negligible across all datasets. NI Sampling even show slightly better performance than full-step sampling on some datasets, which may be due to the variance in evaluation. When compared with confidence threshold sampling, NI Sampling consistently delivers higher speedups under all settings, reaching up to 2.2x faster than it (e.g., 219.3 v.s. 100.8 token/s), while also attaining superior performance under most settings. These results highlight the strong effectiveness of NI Sampling.
Trade-off between Performance and Efficiency. To more comprehensively compare NI Sampling with confidence threshold sampling, we construct trade-off curves between performance and num- ber of sampling steps for both methods by varying the probability threshold and predictor threshold (Fig. 4). More results can be found in Sec. D.3. We that method -orinadd confi- dence threshold sampling across all settings, further demonstrating its practicality under different computational constraints.
7

---

Published as a conference paper at ICLR 2026
Table Comparison of NI Sampling, full-step, and confidence threshold sampling methods.
ALaDA--B-Innsruus
LLaDA 1.5
Dataset
Method
Acc
Steps
Token/s
Acc
Steps
Token/s
Full
73.92
128
211.4
76.65
128
19.7
GSM8K-128
Threshold
73.77
49.35
(2.6%)
75.82
47.44
53.5
NI Sampling
73.69
34.36
76.5 (3.8x)
76.19
29.89
(4.3%)
Full
77.63
256
18.6
81.05
256
18.3
GSMSK-256
Threshold
77.33
74.34
62.9 (3.4%)
81,58
72.92
(3.6x)
NI Sampling
77.18
50.97
90.2 (4.9%)
80.67
53.59
85.6 (4.7x)
Full 74.83
512
14.4
80.67
512
15.0
GSM8K-512
Threshold
75.28
73.29
104.4 (7.2x)
80.89
72.38
105.6 (7.0x)
NI Sampling
76.57
51.08
147.0 (10.2x)
81.20
53.56
140.6 (9.4x)
Full 29.64
128
28.5
31.03
128
27.8
MATH-128
Threshold
29.69
60.43
57.9 (2.0x)
30,93
58.71
(2.9x)
NI Sampling
30.07
39.01
x3.1
31.41
39.65
87.1 (3.1x)
Full 31.89
256
25.0
33.18
256
25.0
MATH-256
Threshold
31.68
95.98
66.7 x2.7x)
33.18
95.41
67.0 (2.7x)
NI Sampling
31.67
61.92
98.7 (4.0x)
32.98
69.67
89.5 (3.6x)
Full
37.
2.56
44,1
43.90
256
44.7
HumanEval-256
Threshold
37.20
98.66
115.8 (2.6x)
42.68
100.7
113.3 (2.5%)
NI Sampling
37.20
55.60
192.5 (4.4x)
42.68
64.40
168.4 (3.xx)
Full 35.37
512
31,5
40.85
512
32.1
HumanEval-512
Threshold
35.37
158.6
100.8 (3.2x)
39.02
158.4
(3.1)
Sampling
35.98
219.3 (7.0x)
105.3
144.2 (4.5x)
Full 36.60
256
22.8
38,40
256
23.3
MBPP-256
Threshold
37.60
33.48
176.0 (7.7xx
38.40
41,89
143.0 (6.1x)
NI Sampling
39.00
30.02
192.8 (8.5x)
38.60
28.92
201.8 (8.7x)
Full
36. 36.80
512
19.1
37.80
512
18.7
MBPP-512
Threshold
37.60
38.75
250.2 (13.1x)
38.60
44.38
215.8 (11.5x)
NI Sampling
38.00
35.25
263.9 (13.8x)
38.80
34.62
268.0 (1433x)
Table 4: Results with Dream 7aee----BBaaa model.
Table 3: Results of NI Sampling combined with dual caching.
Datasset
Method
Acc
Steps
Token's
Full
75.05
23.0
Dataset
Method
Acc
Step
Speed
GSMBK-256
Threshold
72.78
161.95
(1.6x)
NI Sampling
74.45
108.26
52,1 (2.3x)
Full-step
74,83
512
14.4
GSM8K-512
NI Sampling
75.44
44
197.7
Full
36.46
256
29.5
NI Samplingrcache
73.84
50.76
(25.0 0x)
MATH
Threshald
36.35
(2.6x)
Full-step
35.37
512
31.5
NI Sampling
35.74
72.58
100.4 (3.4%)
HumanEvval-512
NI Sampling
35.9%
69.5 54
(7.0%)
Pull 57.60
256
29.7
NI Sumplingscache
35.98
94.39 39
247.3 (7.9x)
MBPP-256
Thrsshold
(2.8x)
Sampling
56.00
68 87
106.7 (3.6%)
Combined with Caching, We combine NI Sampling with Dual Cache method (Wu et al., 2025) to show that NI Sampling is compatible with other techniques for efficient dLL.Ms, with results shown in Tab. 3. With minimal performance loss, NI Sampling achieves more speedup (up to 25.0x).
5.3 ABLATION STUDY
Base Model Type. We test NI Sampling's 's performance on a new base model Bream-7B-Base (Ye et al., 2025), with results shown in Tab. 4. Unlike LLaDA, which is a train-from coin--foo-scratee dLLM, Dream is initialized from an AR LLM. Nonetheless, the results show that even with different base model training methods, NI Sampling still outperforms confidence threshold sampling.
Training Set Distribution. Our main experiments use the ShareGPT dataset to train a generic neural indicator. To examine the effect of training distribution, we train another neural indicator on a dataset combining the training sets of GSMSK and MATH. The results are shown in Fig. 5. We observe that the neural indicator trained on this mixed dataset performs better on the GSM8K and MATH test sets but worse on code datasets, likely because ShareGPT contains a higher proportion of code data. This suggests that neural indicator performance improves when the training distribution matches the test distribution more closely. Such a strategy can be employed to enhance NI Sampling's s performance for specific data domains, such math-focused code-focused tasks.
8

---

Published as a conference paper at ICLR 2026
Results
Results MATH-128
Results HumanEval-256
LLaDA-88-Ine 1.DA-8B-Instruct
Acc
Acc
25-
Dura
Oura
10-
Step
Step
Step
Results on GSMBK-128
Results on MATH-128
Results on HumanEval-256
LLaDA-1.5
Acc
40
DLTS
Deseline
18
Our
30-
Dasaline
20
30
30
50
20
63
150
Step
Step
Step
Figure 4: Trade-off curves between accuracy and steps with LLaDA models. More can be found in Sec. D.3.
Results G5M8K-256
Results MATH-256
Results GSM8K-128
70
50
Acc.
50
hidden
logits
on COMEE-
addition tooen
20
30
50
60
20g
10
20
40
50
Step
Step
Step
Results on HumanEval-2255
Results MBPP-256
Results HumanEval-256
40
ciiautt maaddd HISEACH
30
8
defaut
wio hilden stanes
logit
additionll
20
60
20
25
9
20
50
Step
Step
Step
Figure 5: Ablation study different training data dis-
Figure 6: Ablation study on the input features.
tribution.
As discussed in Sec. 4.3, the neural indicator uses multiple types of inputs. We perform an ablation study on hidden states, predicted logits, and additional tokens to evaluate their contributions, with results shown in Fig. 6. It is evident that hidden states and predicted logits have a significant impact, and removing additional tokens also slightly degrades performance.
Input Types.
RELATED WORKS
6.1 CACHING FOR EFFICIENT DLLMS
In addition to reducing the number of sampling steps. caching has emerged as another promising approach for accelerating dLLMs. This strategy typically partitions the target sequence into blocks and caches the hidden states in block-wise way. Wu et al. (2025) proposes caching either all non- current blocks or all prefix blocks, with the cache updated once the current block generated. Liu et al. (2025b) introduces different caching strategies for prompts and responses. Arriola et al. (2025) further explored semi-auuoregressive architecture that caches all previously generated blocks while disregarding future blocks. It further conducts additional model training to support this design.
6.2 SAMPLING TECHNIQUES OF GENERATIVE MODELS
Continuous Diffusson Models. Continuous diffusion models view the generation as solving an reverse ODE. Prior works reduce the number of ODE steps by employing higher-order ODE solvers

---

Published as a conference paper at ICLR 2026
(Lu et al., 2022; 2025; Zhao et al., 2023; Zhang & Chen, 2022). Other methods optimize sampling schedule to improve trade-offs between NFE and performance (Liu et al., 2024b; Zhou et al., 2024).
Auto-regressive (AR) Models. Speculative decoding is widely adopted method for decreasing the number of AR models' forward passes (Xia et al., 2023; Cai ct al., 2024; Li et al., 2024b;c; 2025; Teng et al., 2024), where draft model generates multiple tokens that are then verified in parallel by the large AR model. Another line of works trains a token-level router that selectively assigns tokens large and small AR model, thereby reducing the inference times of the large model. Liu al. (2024a; 2025a) introduces flow matching and successfully reducing AR sampling to only one step.
Image MDMs. Predicting for accelerating image MDM sampling (Chang et al., 2022; Li et al., 2024a). Conceptually similar this work, Token-Critic (Lezama al., 2022a) and DPC (Lezama et al., 2022b) also apply an extra neural network to decide which token to remask at each step. However, they aim for better performance of image MDMs, which is different from our work.
of tokens simultaneously and designing timestep schedule used
7 LIMITATIONS AND FUTURE WORK
Gap to the Upper Bound. The speedup ratios reported in Sec. 5.2 are still far from those achieved by the analysis in Sec. 3.2 (e.g.. 30.02 steps vs. 10.52 steps). This discrepancy arises because the indicator is not well trained. Improving the indicator's capability remains promising direction.
Potential Extensions. In this work, we empirically adopt the vaajectory-presevvin principle as the optimization target for the indicator. In fact, many other objectives can be explored. For example, the final-results- preneltiip criterion described in App. B, although potentially less stable, provides higher theoretical ceiling. Morcover, one could frame the indicator forcement learning paradigm, which takes actions over tokens at each step. This perspective would allow us to flexibly design reward functions for different targets, such reasoning performance, and enable training the indicator within an RL. framework.
agent within rein-
REPRODUCIBILITY STATEMENT
We describe the experimental details in Sec. 5.1 and Sec. D.5. The code will be open- soureeour
REFERENCES
Marianne Arriola, Aaron Gokaslan, Justin Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Sub- ham Sekhar Sahoo, and Volodymyr Kuleshov. Block diffusion: Interpolating between autore- gressive and diffusion language models. arXiv preprint 2503.0033... 2025.
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in neural information processing systems, 34:1798111 011177981999999 11993,, 2021.
Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D Lee, Deming Chen, and Tri Dao. Medusa: Simple I1m inference acceleration framework with multiple decoding heads. arXiv preprint ar arXiv:2401. 107744 2024.
Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems, 35:28266-28279, 2022.
Huiwen Chang, Han Zhang. Lu Jiang, Ce Liu, and William T Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11315- 11325, 2022.
DeepMind. Gemini diffusion - google deepmind. httpt:tt////c httpss///eeemmmm.ogggg/l//////// nd. www.ggooolleeoooooooooooooo
gemini-diffeee wii--ddiffioo////oo Accessed: 2025-09-19.
Zhengfu He, Tianxiang Sun, Kuanning Wang, Xuanjing Huang, and Xipeng Qiu. Diffusion- bert: Improving generative masked language models with diffusion models. arXiv preprint mrxiv:2211.15022 2022.
10

---

Published as a conference paper at ICLR 2026
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6833 2020.
Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling. Argmax flows and multinomial diffusion: Learning categorical distributions. Advances in neural information processing systems, 34:1 1is11244-124555 2021.
Samar Khanna, Siddhant Kharbanda, Shufan Li, Harshit Varma, Eric Wang, Sawyer Birnbaum, Ziyang Luo, Yanis Miraoui, Akash Palrecha, Stefano Ermon, al. Mercury: Ultra-fast language models based on diffusion. arXiv preprint arXiv:2506.17298 2025.
Jaeyeon Kim, Kulin Shah, Vasilis Kontonis, Sham Kakade, and Sitan Chen. Train for the worst, plan for the best: Understanding token ordering in masked diffusions. arXiv preprint arXiv:2502.06678, 2025.
José Lezama, Huiwen Chang, Lu Jiang, and Irfan Essa. Improved masked image generation with token-critic. In European Conference on Computer Vision, pp. 70-86. Springer, 2022a.
Jose Lezama, Tim Salimans, Lu Jiang, Huiwen Chang, Jonathan Ho, and Irfan Essa. Discrete predictor-correcto diffusion models for image synthesis. In The Eleventh International Confer- ence on Learning Representations, 2022b.
Tianhong Li, Yonglong Tian, He Li, Mingyang Deng. and Kaiming He. Autoregressive image generation without vector quantization. arXiv preprint arXiv:240..11838, 2024a.
Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. Eagle: Speculative sampling requires rethinking feature uncertainty. arXiv preprint ar arXiv:240..15077, 2024b.
Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. Eagle-2: Faster inference of language models with dynamic draft trees. arXiv preprint arXiv:2406.16858 2024c.
Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. Eagle-3: Scaling up inference acceler- ation of large language models via training- time test. arXiv preprint arXiv:2503.01840, 2025.
Enshu Liu, Xuefei Ning, Yu Wang, and Zinan Lin. Distilled decoding 1: One-step sampling of
image auto-regressive models with flow matching. arXiv preprint arXiv:24/2.17153, 2024a.
Enshu Liu, Xuefei Ning, Huazhong Yang, and Yu Wang. unified sampling framework for solver
searching of diffusion probabilistic models. In The Twelfth International Conference on Learning Representations, 2024b.
Enshu Liu, Qian Chen, Xuefei Ning, Shengen Yan, Guohao Dai, Zinan Lin, and Yu Wang. Distilled decoding 2: One-step sampling of image auto-regressive models with conditional score distilla- tion. arXiv preprint arXiv:2510.21000 2025a.
Zhiyuan Liu, Yieun Yang, Yaojie Zhang, Junjie Chen, Chang Zou, Qingyuan Wei, Shaobo Wang, and Linfeng Zhang. allm-cache: Accelerating diffusion large language models with adaptive caching. arXiv preprini urXiv:2506.0629) 2025b.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. arXiv preprini 231: 16834, 2023.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in neural information processing systems, 01-57755-7877 2022.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models. Machine Intelligence Research, pp. 1-22, 2025.
Shen Nie, Fengqi Zhu, Chao Du. Tianyu Pang, Qian Liu, Guangtao Zeng, Min Lin, and Chongxuan Li. Scaling up masked diffusion models text. arXiv preprint arXiv:2410.188 185/4, 2024.
11

---

Published as a conference paper at ICLR 2026
Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. Large language diffusion models. arXiv preprint arXiv:2502.09992, 2025.
Krishna Pillutla, Swabha Swayamdipta, Rowan Zellers, John Thickstun, Sean Welleck, Yejin Choi, and Zaid Harchaoui. Mauve: Measuring the gap between neural text and human text using diver- genee frontiers. Advances in Neural Information Processing Systems, 34:4816-4828, 2021.
Machel Reid, Vincent Hellendoom, and Graham Neubig. Diffuser: Discrete diffusion via edit- based reconstruction. arXiv preprint arXiv:2210.16888 2022.
Subham Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. Advances in Neural Information Processing Systems, 37: 37:130136-13018 2024.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learn- ing, pp. 2256-2265. PMLR, 2015.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 2020.
Yuxuan Song, Zheng Zhang, Cheng Luo. Pengyang Gao, Fan Xia, Hao Luo, Zheng Li, Yuchang Yang, Hongli Yu, Xingwei Qu, et Seed diffusion: large-scale diffusion language model with high-speed inference. arXiv preprint arXiv:2508.02199) 193, 2025.
Haoran Sun. Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai. Score-based continuous-time discrete diffusion models. arXiv preprint arXiv:22 16750, 2022.
Yao Teng, Han Shi, Xian Liu, Xuefei Ning, Guohao Dai, Yu Wang, Zhenguo Li, and Xihui Liu. Ac- celerating regreseeeee rext-to-image generation with training- free speculative jacobi decoding. arXiv preprint arXiv:2410.01699 2024.
Chengyue Wu, Hao Zhang, Shuchen Xue, Zhijian Liu, Shizhe Diao, Ligeng Zhu, Ping Luo, Song Han, and Enze Xie. Fast-dllm: Training- free acceleaation of diffusion 11m by enabling kv cache and parallel decoding. arXiv preprint arXiv:2505.226111 2025.
Heming Xia, Tao Ge, Peiyi Wang, Si-Qing Chen, Furu Wei, and Zhifang Sui. Speculative decod- ing: Exploiting speculative execution for accelerating seq2seq generation. In Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 3909-3925, 2023.
Jiacheng Ye, Zhihui Xie, Lin Zheng, Jiahui Gao, Zirui Wu, Xin Jiang, Zhenguo Li, and Lingpeng Kong. Dream 7b: Diffusion large language models. arXiv preprint 25ii:22508.54881 2025.
Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. arXiv preprint arXiv:2204. 1..0044.3990 13902, 2022.
Wenliang Zhao, Lujia Bai, Yongming Rao, Jie Zhou, and Jiwen Lu. Unipe: A unified predictor- corrector framework for fast sampling of diffusion models. Advances in Neural Information Processing Systems, 36:49842 49869, 2023.
Zhenyu Zhou, Defang Chen, Can Wang, and Chun Chen. Fast ode-based sampling for diffusion models in around 5 steps. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 7777-7786, 2024.
Fengqi Zhu, Rongzhen Wang, Shen Nie, Xiaolu Zhang, Chunwei Wu, Jun Hu, Jun Zhou, Jianfei Chen. Yankai Lin, Ji-Rong Wen, et al. Llada 1.5: Variance-reduced preference optimization for large language diffusion models. arXiv preprint arXiv:2505 19223, 2025.
Yaoming Zhu, Sidi Lu, Lei Zheng, Jiaxian Guo, Weinan Zhang, Jun Wang, and Yong Yu. Texygen: A benchmarking platform for text generation models. In The 41st international ACM SIGIR conference on research & development in information retrieval, pp. 1097-1100, 2018.
12

---

Published as a conference paper at ICLR 2026
A ADDITIONAL ALGORITHM
We list the pseudo algorithm of sampling with NI Sampling and train the neural indicator with Trajectory- Trajectory-Preseriing
in this section.
Algorithm 3 Train neural inidicator with Algorithm 4 Sample with the neural indica-
Trajectory-Preseet wwweerring--rriinneee
tor
Require:
Require:
Neural indicator é. Dataset D with trajecto-
A well trained neural indicator Φ: Pre-trained
ries; Pre-trained dLLM 6.
dLLM 0: Indicator threshold 6p-
1: while not converged do
Initialize æ with a sequence consisting only mask tokens.
2: Sample from dataset D.
Sample & from (l....: 13. 4: Get the æ along T.
2:
while Sampling Unfinished do
3:
Calculate po (x))
5: idæ rrun Alg. with r.step (as k), and 0.
4:
Reveal subset of tokens with an existing sampling method.
6: Label all indices that are within one of 5: Ag...., Aids-1 as I; The other indices are labled as O.
Collect all inputs at the masked positions and feed them into & to get the indicator score.
7: Collect the inputs and labels at all masked 6:
Reveal all tokens with indicator score larger than cp. resulting in new x.
position and train the indicator within it.
end while
7: end while
9:
return idr
8: return aF
B
ORDER
conservative approach in which masked position is not necessarily revealed even if the model's current prediction at that position matches the final generation result. This is because doing so might affect final predictions other masked positions if there exist intermediate positions between this position and the current step on the reference trajectory that are not aligned with the final result. Here, we relax this constraint: as long the model's prediction at the current step aligns with the pre-obtained final generation result, then the token can be sampled at this step. We refer to this ordering strategy as Final-Results- Preserving Order. Using the top-1 probability full-step sampling results as the reference, we con- ducted experiments on GSM8K-25 226 with LLaDA-8P AAAA--B-IIssreee and LLaDA 1.5 models. Results are reported in Tab. 5.
In Sec. 3.2, introduced Prajectory--Pesevving---
Table 5: Performance on GSM8K-256 of Final-Results- -reeeris- Order.
LLaDA-8B-Instrus
LLaDA-1.5
Method
Acc
Steps
Acc
Steps
Full-step
77.63
256
81.05
256
Threshold
77.33
74.34 (3.4x)
81.58
72.92 (3.5x)
Eenjectory-Preeerviine
77.63
22.36 (11.4x)
81.05
22.31 (11.5x)
Semull-Reults-Peee Preseevin
77.78
6.95 (36.8x)
80.89
9.11 (28.1x)
Although Results- Preserving -rrder Order does not strictly guarantee identical generated results with the reference results, its performance remains nearly unchanged. We further examine the answers for the first 100 problems. For the ---aDA-----nstter sults produced by Menall-Rsults-Peeeeeee cases, only a few tokens differ for these examples, while all final answers remain identical. For the LLaDA-1.5 model, 29 generations differ, among which only one produced a different answer. These observations indicate that this criterion causes negligible performance degradation while further in- creasing the acceleration ratio to remarkably high level. Nevertheless, considering the training stability and task simplicity, we adopt the simpler yet efficient enough Trajectory-Preest Preserving crite- rion the training objective for the indicator.
model, only 20 of the first 100 generated re- are not exactly identical to the reference results. In most
13

---

Published as a conference paper at ICLR 2026
C RANDOM SAMPLING WITH NI SAMPLING
As discussed in 4, the default sampling process of NI Sampling is deterministic. This does not align with the practical requirements of LLM applications, since users do not expect a single fixed output for the same prompt. In this section, we discuss the potential of NI Sampling for supporting random sampling. Specifically,
show how introduce randomness with current indicator in Sec. C.I, which is trained by the process described in Sec. 4.2. Moreover, to prevent the indicator from relying too heavily on a deterministic trajectory, we propose use random trajectory for training in Sec. C.2. Finally. we report the diversity evaluation results of NI Sampling in Sec. C.3.
C.I RANDOM SAMPLING WITH CURRENT INDICATOR
Our current indicator can be incorporated with random sampling easily without additional tuning. Specifically, similar to the implementation of full-step sampling and confidence threshold sampling, we sample token at each position randomly according to the probability vector instead of using argmax. Then, we feed all actually sampled tokens together with its probability into the existing neural indicator and let it make decisions. Other parts are the same as the description in Alg. 4.
We discuss the intuition behind the effectiveness of this design below.
Our method follows the same paradigm as the random sampling implementation of confidence threshold sampling. Specifically, under the greedy setting, both methods deterministically sample all tokens and keep the tokens whose scores exceed a threshold. When combined with random sampling, the modification remains exactly the simply replaced by sampling from the predicted distribution. The only difference between the two coaeeiidddddooiiiieen--tteeeeeee sampling uses confidence
for both methods: deterministic selection is
methods lies in how cligible tokens are values, while our method uses the indicator score.
In fact, our method may be even more diverse than confidence threshold sampling. At each step, the set of tokens that can be accepted under greedy sampling is larger in our method than meefidence-trrehhlll sampling. Consequently, under random sampling, there is a higher proba- bility of selecting tokens that deviate from the original trajectory. This increases the likelihood of trajectory divergence and thus leads to greater sampling diversity.
C.2 IMPROVING SAMPLING DIVERSITY BY TRAINING INDICATOR WITH RANDOM TRAJECTORY
The diversity of our method can be further improved by modifying the training process. Since our existing indicator has never data beyond tokens sampled by greedy sampling, trajectories other than deterministic ones, we revised the training strategy and retrained the model. The adjust- ments are:
The deterministic trajectory generation process was replaced with a stochastic process us- ing temperature of 1.0. This can address the reviewer's concern regarding potential over- fitting to deterministic trajectories.
During training, tokens are sampled according to the model's predicted probabilities to serve as inputs of the indicator, rather than being selected solely with greedy strategy.
C.3 RESULTS OF SAMPLING DIVERSITY
We evaluate the diversity of our random sampling algorithm in this section. We select LLaDA- 8B-Instruct the pre-trained model and set the temperature as 1.0 for all experiments. For each problem, we generate & different answers using multiple random seeds, and report the following metries to quantify diversity:
pass@k. For this metric, problem is considered correctly solved if at least one of the k generated answers correct. A diverse generative distribution should show a stable in- crease in pass @k k l grows. Note that we use this metric for all datasets rather than being limited to code datasets. For SQuADv2, which uses average Fl score as the evaluation met- ric, we report the average of the maximum F1 scores across the & independently generated
14

---

Published as a conference paper at ICLR 2026
Table Diversity evaluation results on HumanEval-256 dataset. "Slope" stands for the slope of the fitted line of pass@k k with respect to logk.
Token/s
NFE
Pass@1
@ass@
Pass @
Pass @
Pass@ 16
Slopet
Full-step
42.91
256
36.99
44.69
51.43
57.81
63.76
6.67
Threshold
110.82
99.13
36.82
44.35
51.02
57.40
63.10
6,56
Old Indicator
164.09
63.34
36.43
44. 13
51.09
57.63
63.62
6.79
New Indicator
155.78
66.72
36.70
44.38
51.34
58.03
64.33
6.89
Table 7: Diversity evaluation results on GSM8K-256 dataset. "Slope" stands for the slope of the fitted line of pass @k with respect logk.
Tekco'ss
NFE
Pass
Prss
Pass 4
Pass
Slopet
1-gram-BLEU)
I-gram BLEUL
2-gram BLEUL
MAUVEL
Pull-step
84.83
89.99
0,916
Threshold
65.56
74.63
77.55
85.51
91.05
93.85
5.44
0.920
0.887
0.974
Old Indicator
86.74
85.97
New Indicator
82,36
57.68
77.33
84.91
90.83
94.47
5.73
0.918
0.884
0.973
Table Diversity evaluation results on Squad-Completion- 128 dataset. "Slope" stands for the slope of the fitted line of pass @k with respect to logk.
Token''s
NHE
Passil
Pass
Pass
Pass
Slopet
--gam-BLEUU
2-gram-HLHUL
MAUVEL
Pulli-tep
48.48
83.3
0.809
0.600
0.981
Threshold
174.64
34,03
77.5
84.0
87.0
89.5
3.90
0,802
0.600
0.986
Old Indicator
266.78
21.39
76.8
83.2
89.7
4.23
0,810
0.592
0.982
New Indicator
254.98
22.38
76.8
83.2
87.3
89.8
4.31
0.811
0.589
0.981
answers for all problems. We then report pass@k under different k and further compute slope of the fitted line of pass @ k with respect logk, which serve as the primary metrics for evaluating diversity. Higher values indicate higher diversity.
Self-BLEU (Zhu et al., 2018). For each answer, we compute its BLEU score with respect answers. Then take average across all answers. We use both
to the other I-gram-BLEU and 2-gram-BLEU. Lower values indicate higher diversity.
Selr-MAUVE (Pillutla al., 2021). We randomly split the à answers into two groups and compute the MAUVE score between them. Smaller values indicate higher diversity.
We demonstrate results in Tab. 6, Tab. 7, Tab. 8, Tab. 9, and Tab. 10. The row "Old Indicator" shows the diversity of our current neural indicator, while the row "New Indicator" shows the diversity of the newly trained indicator as discussed in Sec. C.2. The key takeaways are: (1) our existing indicator has already achieved comparable diversity threshold sampling and full-step sampling across most datasets. The slope of pass @k with respect to logk is slightly higher than that of confidence thresh- old sampling most cases, and comparable full-step sampling. This indicates that our method does not introduce additional loss of distributional diversity. Additionally, in most cases, our Self- BLEU and Self-MAUVE scores are comparable to or only slightly worse than those of confidence threshold sampling, further supporting the conclusion that there is no loss of distributional diversity in our method: (2) the row "New Indicator" shows larger slope of pass @ with respect to logk, which indicates that training predictor random trajectories can boost sampling diversity.
D ADDITIONAL EXPERIMENTAL RESULTS AND DETAILS
D.1 COMPARISON WITH MORE BASELINE
In this section, we compare NI Sampling with another related work, Token-Critic (Lezama et al., 2022a). As discussed in Sec. 6, Token-Critic also introduces auxiliary neural network to decide
15

---

Published as a conference paper at ICLR 2026
Table Diversity evaluation results on Squadv2-128 dataset. All numbers of pass@k @ are Fl scores. "Slope" stands for the slope of the fitted line of pass @ k with respect to logk.
Tokcals
NFE
Pass
Pass 8:2
Pass
Pass
Slopet
1-gram-BLEU)
"gram-BLEUL
MAUVEL
Full-step
47.94
128
27.64
31.76
37.46
41,92
4.85
0.792
0.710
0.989
Threshold
202.27
31.23
28.05
33.21
37.92
42.21
4.72
0.792
0.710
0.989
Old Indicator
267,74
22.56
27.85
33.13
38.01
43,20
5.09
0.797
0.723
0,989
New
34.10
0.720
Table 10: Diversity evaluation results on TruthfulQA-128 dataset.
Taken's
NFE
blen ace
rougel noc
rouge2 acc
rongel... acc
1-gram-BLEUl
2-gram-BI.FUL
MAUVEL
Full-step
46.54
128
53.86
53.61
49.07
52.88
0.898
0.862
0,983
Threshold
218.80
Old Indicator
419.67
13:41
$7.89
56.55
52.39
$7.04
0.903
0.867
0,980
New Indicator
408,90
13.90
61.57
59.00
54.59
59.73
0.901
0.864
0.985
Table
Comparison between NI Sampling and Token-Critic (Lezama et al., 2022a).
Step
15.47
20.84
27.41
NI Sumpling
Acc
60.35
69.37
71.95
Step
15.08 (Top-8)
22.01
(Top-4)
28.58 (Top-2)
Token-Critic
Acc
33.59
42.84
53.83
Table 12: Comparison between indicator sizes of pre-trained dLLM and neural indicator.
Model Size
Indicator Size
LLaDA
8,0LB 7.62B
84.2M 96.4M
Dream
which tokens should be unmasked during sampling. For fair comparison, adopt the same neural indicator architecture and training configuration for both our method and Token-Critic. Following the original paper, we apply top-k sampling Token--ritic and combine it with confidence threshold sampling, consistent with NI sampling. Results on GSM8K are presented in Tab. 11. We observe that Token-Critic does not outperform NI sampling under accelerated sampling settings. One possible reason is that predicting whether token is originally masked may not serve as an effective principle for speeding up autoregressive decoding.
D.2 COST OF THE NEURAL INDICATOR
We first list the parameter size the neural indicator and the pre-trained dLLM in Tab. 12. We further count the inference time of the neural indicator and the pre-trained dLLM, with results ported in Tab. 13 and Tab. 14. We can see that the indicator contains only about 1/80 to 1/100 of the parameters of the dLLMM, while its inference time ranges from roughly 1/18 1/40. This clearly shows that the additional time and memory overhead introduced by NI Sampling minimal. is worth noting that we do not apply any system-level optimizations, the indicator's 's time cost could be further reduced in practice, All the speed results we report in Sec. 5 also include the indicator's runtime.
D.3 TRADE-OFF CURVES UNDER MORE SETTINGS
We present more performance-stey trade-off curves, including LLaDA-8B- Instruct and LLaDA-1.5 GSM8K-512 HumanEval-512 and MBPP-512 datasets, and Dream- In-trrucct GSM8K- 256, MATH-251 and MBPP-256 datasets, Fig. 7, Fig. and Fig. 9, respectively. We also provide trade-off curves between accuracy and inference time, in Fig. 10, Fig. 11, Fig. 12 and Fig. 13.
16

---

Published as a conference paper at ICLR 2026
Table 13: Inference time (ms) LLntt-SB-Inttrucc model and the neural indicator.
Dataset
Model
Indicator
GSMSK-128
49.0
1.15
GSM8K-256
53.8
1.56
GSM8K-512
69.4
2.43
MATH- 128
35.1
1.12
MATH-250
40.0
1,45
Human6vvl-2256
22.7
1.13
HumanEval-5112
31.7
1.77
MBPP-256
43.9
MBPP-512
52.3
1.74
Table 14: Inference time (ms) of Dream-7B-Base model and the neural indicator.
Dataset
Model
Indicator
GSM8K-256
43,4
1.82
MATH-2256
33.8
1.29
MBPP-512
33.7
1.12
NI Sampling consistently pareto-dominate baseline method, demonstrating the effectiveness of our method.
D.4 MORE ABLATION RESULTS
Parameter Size of the Indicator. increase the parameter size of indicator by simply stacking more layers onto the MLP backbone, resulting in larger indicator with 184M parameters. Fig. 14 presents a comparison between the two indicator sizes. can be seen that increasing the indicator size does not provide obvious performance gains. For efficiency considerations, we therefore adopt the smaller indicator size in our main experiments.
Training Set Size. To reduce training cost, decrease the amount of data used for training and evaluate its impact. The results are shown in Fig. 15, indicating that reducing the training data by 75% does not significantly hurt performance. Therefore, users can safely reduce the number of generated trajectories if faster training is needed.
D.5
MORE EXPERIMENTAL DETAILS
D.5.1 EVALCATION
For the evaluation of performance, we follow Wu et al. (2025) and adopt Im-eval framework htt.e//ww..oooeeeeeeee/e///////e
We use the "flexible-extract" filter for GSM8K dataset, the average of "exact_match" and verif_" for MATH dataset, and pass@1 for both HumanEval and MBPP datasets. For sampling speed, we only count the GPU inference time of the pre-trained dLLM and the neural indicator for all settings ensure clean compari- son. All inference speeds of La-D------nnttuut and LLaDA-1.5 are measured on a single NVIDIA H200 GPU, while those of Dream-7B-BBan are measured single NVIDIA H800 GPU.
D.5.2 BASELINE
For full-step baseline, we follow the default setting in the official implementation of LLaDA and Dream. We use top-1 probability criterion for LLaDA-SB-Insttun and LLaDA-1.5 to choose the token sampled at each step. For Dream-7B-Bass we use the top-1 entropy for token selection. For confidence threshold sampling, we use 0.9 as the threshold for most cases reported in the Tab. 2. For several datasets, the accuracies of using 0.9 threshold are actually so we choose the best threshold for them. Specifically, for LLaDA-8B-Ine ---nsruuttt model, we set the threshold 0.7 GSM8K-51 -12, and 0.8 on MBPP 256; for LLaDA-1.5 model, set the threshold
than a lower threshold,
17

---

Published as a conference paper at ICLR 2026
Results G5M8K226
Resuits MATH-256
Results MBPP-256
70
15
5.0
15
40
Baseline
Baseline
Baseline
20
10
45
Step
Step
Step
Results on G5MBK-512
Results on HumanEval-512
Results on MBPP-512
70
35
60
12
50
20
Ours
Ours
Baseline
Baseline
Baseline
320
120
130
Step
Step
Step
Figure 7: Performance-step trade-off curves of LLaDA-8B-Instuue
Results on GSM8K256
Results on MATH-256
Results on MBPP-256
80
70
20
20
Step
Step
Results on G9M8K-512
Results on HumnanEval-522
Results on MBPP-512
80
155
70
60
15
Dums
Ours
10-
Our
Basdiie
Basdine
Bacciin
20
50
123
130
30
35
Step
Step
Step
Figure Performance-step trade-off curves of LLaDA-1.5.
as 0.8 on GSM8K-512. For trade-off curves, we choose threshold from 0.3 0.9 with an interval 0.1.
D.5.3 NI SAMPLING
Architecture of the Neural Indicator. As discussed in Sec. 4.3, neural indicator. We set the width of MLP as 768 and the depth of MLP as 5 for all main experiments with B.aDA--B-Instruu and LLaDA-1.5 115 models. For Dream-788 7B-Base model, we change the depth of MIP to 8. For the ablation study in Fig. 14, we increase the layer number to 10 to construct larger indicator.
a MLP architecture for the
Data Generation. In our main experiments, we use ShareGPT dataset to construct training set for the indicator. Following speculative decoding methods (Cai et al., 2024; et al., 2024b), we use the first prompt from the user and ignore other data. We do not conduct any additional operations and directly feed the prompt to the generation pipeline of dLLMs. For LLaDA-8B ----trut and LLaDA- 1.5 models, generate each data sample three times, using generation lengths of 128, 256, and 512, respectively. We adopt confidence threshold sampling to construct reference trajectories and set the threshold as 0.8 for efficiency. For Dream-7B-Ban model, generate each dataset sample
18

---

Published as a conference paper at ICLR 2026
Results GSMBK-256
Results MATH-256
Results MBPP-256
80
Ous
8
Acc
60
15
50
10-
Olts
Our
Bassline
Baseline
75
100
125
155
46
60
60
150
40
50
60
Step
Step
Step
Figure Performance-step trade-off curves of Dream-7B-Base.
Results on GSM8K-128
Results on MATH-128
Results on HumanEval-256
LLaDA- 8B 28-Inntruct
70
30
Ace
Acc
50
15-
Baseire
Outs
10
Our
Suse mm
10
Base Ine
Our
1300
25000
500
1300
1320
(ms)
(ms)
(ms)
Results on GSMBK: 128
Results on MATH-128
Results on HumanEval-256
LLaDA-1.5
Acc
40
ACC
10-
Deselrre
305
1010
1510
2030
300
1003
1500
2020
300
1000
15:0
2030
Time (ms)
Time (ms)
Time (ms)
Figure 10: Performance-inferee time trade-off curves of Dream-7B -ream--B-Bas..
twice with 128 and 256 generation lengths. We use the default top-1 entropy full-step sampling to generate reference trajectories, as we find that parallel generation more challenging on Dream, possibly because Dream is initialized with an AR LLM.
Training. We train the neural indicator for 50 epochs under all settings. We apply AdamW opti- mizer, with fixed learning rate of 0.0002 and betas of (0.9, 0.95), We set the batch size 256. Empirically, to make the training task easier to learn, we additionally require that tokens labeled as positive must have probability learger than 0. 15; otherwise, they would be labeled as negative.
Sampling. As discussed in Sec. 4.1, we first sample a subset of tokens with an existing sampling method to prevent the situation where token reaches the indicator score threshold. For most cases, we choose confidence threshold sampling with threshold 0.9 for results in Tab. 2 and threshold 0.8 for other data points in all trade-off' curves on LLaDA models before we use indicator scores each step. For those datasets mentioned in Sec. D.5.2 where 0.9 is not the optimal threshold, we used the corresponding optimal threshold instead. For Dream 7 model, we consistently apply top-1 entropy sampling for all settings, since confidence threshold sampling incurs a larger performance drop on Dream than LLaDA. For the indicator score threshold used Tab. 2, we adjust this hyper- parameter around 0.9 for different datasets to ensure that our results are better than or comparable to confidence threshold sampling, thereby ensuring a fair comparison in terms efficiency. For Dream model, we directly use 0.9 for the results in Tab. 4. In all trade-off curves, select indicator score thresholds from 0.2 to 0.8 witn an interval 0.1 for the other data points.
E ANALYSIS OF INDICATOR PREDICTIONS
The speedup achieved by NI Sampling is in fact far below the speedup demonstrated in Sec. 3 (c.g., 4.8x v.s. 11.4x on GSM8K-256), To better understand the reason behind this, we investigate several practical cases GSM8K dataset and identify several problems of the current neural indicator.
19

---

Published as a conference paper at ICLR 2026
Results GSMBK-256
Resuits MATH-256
Results MBPP-256
70
ACC
ACC
20
50
15
40
Dux
10
Oum Baseline
Baseline
Baseline
1030
2303
2003
4050
2050
2503
1030
750
1050
1250 1530 1750
Time (ms)
Time (ms)
Time (ms)
Results on GSMBK-512
Results on HumanEval-512
Results on MBPP-512
70
60
Acc
AcC
50
20
Durs
10
Curs
Ours
Baseiine
Baseline
2000
4000
3000
1000
2000
3000
1558
2000
(ms)
Time (ms)
Time (ms)
Figure 11: Performance-inferrnnce time trade-off curves of LLaDA-8B-Instru
Results GSM8K-256
Results MATH-256
Results MBPP-256
Acc
Our5
2000
30.0
1000
2000
3000
1000 1250 1500
1750
Time (ms)
Time (ms)
Time (ms)
Results on GSM8K-512
Results on HumanEval-512
Results on MBPP-512
70
60
50
Ours
10-
Ours
Ours Baseline
40
Baselin
2000
3000
4000
5000
1.00
2000
3000
400o
5000
1250
1730 Time (ms)
2000
2250
Time (ms)
Time (ms)
Figure 12: Performance-inferrenee time trade-off curves of LLaDA-1.5.
Our indicator tends to select tokens following an AR schedule, even when ground truth trajectory-preserving preserving- does not follow such an order. There are two main impacts:
The indicator not identify correct tokens that are far away well. Examples are shown in Tab. 15. For GT trajectory preserving order, there are two zeros at later positions. However, NI sampling can only select the first
The indicator tends to select near tokens which actually not be sampled with trajectory preserving order. Examples are shown in Tab. 16. We can see that trajectory preserving order doesn't sample the token "of" after "The cost", because it is not aligned with the final result "to" at this position. However, NI sampling selects this position to unmask, which changes the final result.
Even though the ground truth rerjectory-preservnng-ooo follows AR schedule, the indi- cator may just not select enough tokens. Examples are shown Tab. 17. NI sampling does not select correct tokens "he needs to" due the low indicator score, although they follow AR schedule.
Although our ablation studies show that moderately increasing the parameter size does not signifi- cantly affect performance, still hypothesize that these issues are because the model capacity
20

---

Published as a conference paper at ICLR 2026
Results GSM8K-256
Results MATH-256
Results MBPP-256
80
Ours
Ours Sasaline
Bsselin
50
32
2000
3000
4000
5000
6000
7000
3000
1500
2000 Time (ms)
2500
3000
1000
2000 Time (ms)
3020
Time (ms)
Figure 13: Performance-iifference time trade-off curves of Dream-7B-Ban
Results on GSM8K-256
Results on MATH-256
80
35
smalldeeauult)
small (default)
70
large
60
Acc
20
50
10
5
20
25
30
35
40
20
30
40
50
60
70
Step
Step
Results on HumanEval
Results on MBPP-256
40
(defoult)
40
*
large
-Xi
large
30
Acc
ACC
40
20
Step
Step
Figure 14: Ablation results on parameter size of the indicator.
the amount of training data is insufficient for the model to learn these difficult patterns. Since these patterns are too hard to learn, only scaling indicator size slightly unlikely to be effective. We provide several potential solutions follows:
Modifying the model architecture. The current model is MLP that makes independent predictions for each masked position, which is likely to be suboptimal and may prevent each position from sufficiently leveraging contextual information. A possible improvement is to replace the MLP with a Transformer-lit architecture that makes joint decisions across all masked positions.
Scaling both model and data. Scaling may be more effective after replacing model archi- tecture.
Addressing the issue of accumulated errors. In practical sampling, the errors introduced by the indicator accumulate over decoding steps, whereas the current training data does not account for this. Regenerating data using the indicator itself and retraining on such data may mitigate this issue.
Designing loss formulations and weighting strategies. Modifying the loss formulation may allow the training pay more attention on the non-semi-AR positions, which are more difficult for the model learn.
21

---

Published as a conference paper at ICLR 2026
Results on GSM8K-128
Results on HumanEval-256
80
40
default
70
25% data
60
30
ACC
Acc
50
40
default
30
-X-
25% data
20
30
Step
Step
Figure 15: Ablation results on training set size.
Table 15: Examples of indicator prediction. "Original" means the sequence the current step. "Tra- jectory Preserving Order" means conducting the current step with ground truth trajectory-presert preserving order. "Indicator" means coducting the current step with NI Sampling. [M] stands for mask to- ken. The differences between "Trajectory Preserving Order", "Indicator", and "Original" are marked as orange, while the key differences between "Trajectory Preserving Order" and "Indicator" are highlighted in red.
Original
Trajectory Preserving Order
Indicator
Each implant has a base price of $2000.\n2. One implant has additional cost of $500 for the crown, so the total cost for one implant $2000 $500 = $2500. \n3. George Medd{[/MMMMMMMMMM МMИММMMМММММM
1. Each implant has a base price of $2000. In. One im- plant has of S500 for the crown, so the total cost for one implant $2000 $500 = $2.500. )n3. George needs implants, the total cost for the implants (subsequent mask tokens are omit- [MIIMIMMM[[J]]]IIIIIIIIII [MMMIIIIMMMMIIЛIIIIIIIIIICIIIIII
Each implant has a base price of $2000. One im- plant has of $500 for the crown, the one implant is
additional cost
additional cost
total cost $2000 + $500
$2500. nn.
George needs 2 implants, the total cost for the implants SIMIMMIIAMMMMMM[IIIIIIIIIMMIIMMM [MMMMМОИИ[М]M
ted)
Table 16: Examples of indicator prediction. "Original" means the sequence at the current step. "Tra- jectory Preserving Order" means conducting the current step with ground truth engjectory-prservvie order. "Indicator" means coducting the current step with NI Sampling. [M] stands for mask to- ken. The differences between "Trajectory Preserving Order", "Indicator", and "Original" are marked as orange, while the key differences between "Trajectory Preserving Order" and "Indicator" are highlighted in red
Original
Trajectory Preserving Order
Indicator
The area of the bedroom is (11--21611122112221
The area of the bedroom is
The area of the bedroom is square feet !nThe of
square feet !nThe cost of
square feet \nThe cost
new carpet is
of the new carpet 166112221111122252222 2592
new carpet is
166112-<<1161122-222 2592 AnThe cost of the[M] SIMII1GMMMMM-SSIIIIIIIIIMM
TnThe cost of the padding is 22162=$$2211222222222333 The cos[M][[Mjhhe old[M[[M]
\nThe cost of the padding 25$16 The cost of[M] the MldM MMMMMMMMMMMMMM
MIMMIMMMMMMMMM MMMMM.MMMMM.........
M
F DEMONSTRATION OF GENERATED EXAMPLES
In this section, we demonstrate examples with different sampling methods.
22

---

Published as a conference paper at ICLR 2026
Table 17: Examples of indicator prediction. "Original" means the sequence at the current step. "Tra- jectory Preserving Order" means conducting the current step with ground truth trajectory-presert preserving order. "Indicator" means coducting the current step with NI Sampling. [M] stands for the mask to- ken. The differences between "Trajectory Preserving Order", "Indicator", and "Original" are marked as orange, while the key differences between "Trajectory Preserving Order" and "Indicator" are highlighted in red
Original
Trajectory Preserving Order
Indicator
Isaias plans to sell 3/5 of his 300 chickens, which is 3/5 300 180 chickens. TToo make a profit of S2000, Isaias needs to sell the chickens at $50 per chicken, so he needs to earn 180 * $50 $9000.) AnTo make a profit of $2000. Isaias needs earn $9000 $2000 $7000. (ninne Isaias needs to MMIMMIIIMMMMMMMIIIMMM MMMMMIМMMMMMMMMMМ
Isaias plans to sell 3/5 of his 300 chickens, which is 3/5 300 180 chickens. AnTo make profit of $2000, Isaias needs to sell the chickens at $50 per chicken, so he needs to earn 180 $50 = $9000. (nTo make profit of $2000. Isaias needs earn $9000 $2000 $7000. AnSinee Isaias needs to sellthe chickens at $50 per chicken, he needs [MIMIIMMMMMMMM MMMMMMMMMMMMM
Isaias plans to sell 3/5 of his 300 chickens, which is 3/5 300 180 chickens.)nTo make a profit of $2000, Isaias needs to sell the chickens at $50 per chicken, so he needs to earn 180 $50 $9000. make profit of $2000. Isaias needs earn $9000 $2000 $7000. \nince Isaias needs to sell the chickens at $50 per chicken, МОМMМMМММММММ
Table 18: Qualitative examples across different sampling ethods. The prompt is: Janet's ducks lay 16 eggs per day. She eats three for breakfast every morning and bakes muffins for her friends every day with four. She sells the remainder at the farmers' market daily for $2 per fresh duck egg. How much in dollars does she make every day at the farmers' market?
Full-step
Threshold
NI Sampling
Janet's ducks lay 16 eggs per day. She eats 3 eggs for breakfast ev- ery morning, so she has 16 3 13 eggs left. She hakes muffins for her friends with 4 eggs ev- ery day, so she has 13 4 9 eggs left. She sells the remaining eggs the farmers' market for $2 per egg, so she makes ==="2=18">> 18. Therefore, Janet makes $18 every day at the farm- ers' market. #### 18 (NFE:128)
Janet's ducks lay 16 eggs per day. She eats 3 eggs for breakfast ev- ery morning,
Janet's ducks lay 16 eggs per day. She eats 3 eggs for breakfast ev- ery morning, so she has 16 3 13 eggs left. She hakes muffins for her friends with 4 eggs ev- ery day, so she has 13 4 eggs left. She sells the remaining eggs the farmers' market for $2 per egg, so she makes 9<9"2=18 18. Therefore, Janet makes $18 every day at the farm- ers' market. #### 18 (NFE:34)
she has 16 13 eggs left. She hakes muffins for her friends with 4 eggs ev- ery day, so she has 13 4 9 eggs left. She sells the remaining eggs at the farmers' market for $2 per egg, so she makes 9 $2 = 2=18> 18. Therefore, Janet makes $18 every day at the farm- ers' market. #### 18 (NFE:49)
$2 =
$2 =
G THE USE OF LARGE LANGUAGE MODELS (LLMS)
We used LLM to support us in grammar refinement and language polishing. The paper was mainly written by the authors.
23

---

Published as a conference paper at ICLR 2026
Table Qualitative examples across different sampling ethods. The prompt The great dragon, Perg, sat high atop mount Farbo. breathing fire upon anything within a distance of 1000 feet. Polly could throw the gold javelin, the only known weapon that could sleigh the dragon, for a distance of 400 feet, well within the reach of the dragon's flames. But when Polly held the sapphire gemstone, she could throw the javelin three times farther than when not holding the gemstone. If holding the gemstone, how far outside of the reach of the dragon's flames could Polly stand and still hit the dragon with the gold javelin?
Full-step
Threshold
NI Sampling
When Polly was holding gemstone, she could throw the javelin 400 feet. When holding the gemstone, she could throw the javelin three times farther, she 1200 feet.
When Polly was not holding the gemstone, she could throw the javelin 400 feet. When holding the gemstone, she could throw the javelin three times farther, she 1200 feet. The dragon's 's flames could reach a distance of 1000 feet. Therefore, when holding the gemstone. Polly could throw the javelin from 1200 1000 feet outside the reach of the dragon's 's flames. #### 200 (NFF:55)
gem- When not holding stone, she could throw the javelin 400 feet. When holding the gem- stone, she could throw the javelin three times farther, she could throw
could throw 400 The dragon's 's flames could reach up to 1000 feet. Therefore, to still hit the dragon while holding the gemstone, Polly could stand 1200 1000
could throw it 400
1200 feet. The dragon's flames could reach 1000 feet, so she could throw the javelin 1200 feet, she could stand 1200 1000 200 outside of reach of the dragon' 's flames. #**# 200 (NFE:38)
feet outside of the reach of the dragon's 's flames. #### 200 (NFF:128)
24


Reviews:
[{"review_id": "fGNVSdrtu2", "reviewer": "Reviewer_VVhC", "summary": "This paper proposes NI Sampling, an efficient decoding framework for pre-trained diffusion LLMs that optimizes token sampling order with a tiny neural indicator predicting which masked positions can be simultaneously and safely revealed at each step. Labels for the indicator are derived from a trajectory-preserving criterion that merges steps when the current predictions already match a reference trajectory, turning correct predictions into parallel. On LLaDA-8B/1.5 and Dream-7B across GSM8K, MATH, MBPP, and HumanEval, NI Sampling delivers significant speedups over full-step sampling and threshold-based sampling with negligible accuracy loss.", "questions": "Please see Weaknesses", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- I enjoyed reading this paper. It's well-motivated. The trajectory-preserving case is great and clear, demonstrating the great potential of merging sampling steps for acceleration\n- The proposed solution is conceptually simple and practical. A lightweight MLP as the indicator brings marginal extra cost.\n- The proposed method achieves consistent and significant speedups across multiple benchmarks compared to threshold-based sampling and full-step sampling, with negligible accuracy loss.\n- Detailed description of the setup of NI, analysis, and ablation studies, in addition to the main results.", "weaknesses": "I think this paper and its contributions could be further enhanced with more investigation into the transferability of the neural indicator:\n- I got the difference between LLaDA models and Dream, but for models from the same family (e.g., LLaDA-8B and LLaDA-1.5 used in this paper), would it be possible to use a shared pretrained neural indicator?\n- Would it be possible to use the same pretrained neural indicator for different generation window lengths?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "d73B2dcWTw", "reviewer": "Reviewer_B3No", "summary": "This paper proposes NI Sampling, a learnable framework to accelerate dLLMs by optimizing the token sampling order.\n\nInstead of relying on heuristic confidence-based thresholds, the method introduces a neural indicator that predicts which masked tokens can be safely revealed at each step. To train this indicator, the authors design a trajectory-preserving objective, ensuring consistency with reference sampling trajectories while achieving significant speedups.\n\nEmpirical results show up to 14× acceleration compared to full-step sampling with negligible performance drop across benchmarks such as GSM8K, MATH, HumanEval, and MBPP.", "questions": "* **Diversity and Trade-off Analysis**\n  As mentioned above, the main concern is that NI Sampling might overfit to the reference trajectory, thus harming sample diversity. To quantify this, please consider reporting:\n\n  1. **Token entropy**, to measure the diversity of generated tokens.\n  2. **pass@K**, to evaluate the diversity–accuracy trade-off more explicitly.\n\nAdditionally, the paper would benefit from a **theoretical analysis** of when and why merging steps (as done in the trajectory-preserving principle) does not harm the underlying distribution. In other words, even if the merged distribution deviates from the base model’s, a formal understanding of this trade-off would allow users to make informed decisions about when NI Sampling is appropriate. Currently, the method lacks such theoretical grounding, making it unclear under what conditions the distribution shift is acceptable.\n\n[1]: *Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Design*, [https://arxiv.org/abs/2402.04997](https://arxiv.org/abs/2402.04997)", "limitations": "", "rating": 2, "confidence": 3, "soundness": 1, "presentation": 3, "contribution": 2, "strengths": "* **Learnable Scheduler**\n  Improving the sampling efficiency of discrete diffusion models is an important and timely research direction. The idea of approaching this problem through a *learnable* mechanism for sampling order optimization is particularly interesting and potentially impactful.", "weaknesses": "* **Overfitting to Reference Trajectory**\n  (Please correct me if I have misunderstood this point.)\n  A major limitation of NI Sampling is that it appears to produce generations that overfit to the *reference trajectory* used during training. In this sense, it differs fundamentally from efficient sampling studies in *continuous* diffusion models, where speedup is achieved by reducing *numerical errors* of ODE solvers rather than altering the learned data distribution.\n\nWhile the paper claims acceleration in sampling speed, this efficiency primarily stems from treating tokens as *conditionally independent*—i.e., ignoring inter-token correlations. In contrast, methods that preserve token correlations (e.g., [1]) achieve efficiency without such independence assumptions. Consequently, NI Sampling resembles trajectory-conditioned approaches such as [2], which restrict the model’s output distribution to a small subset of trajectories, potentially degrading sample diversity.\n\n[1]: *Jump Your Steps: Optimizing Sampling Schedule of Discrete Diffusion Models*, [https://openreview.net/forum?id=pD6TiCpyDR](https://openreview.net/forum?id=pD6TiCpyDR)\n[2]: *Beyond Autoregression: Fast LLMs via Self-Distillation Through Time*, [https://arxiv.org/abs/2410.21035](https://arxiv.org/abs/2410.21035)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PgRaI62Kcq", "reviewer": "Reviewer_wmZ5", "summary": "This paper explores the potential for reducing sampling steps in discrete Large Language Models (dLLMs) by merging backward steps wherever possible. To approximate this sampling strategy, the authors propose Neural Indicator Sampling (NI Sampling), which learns to identify which positions can be unmasked at each step.", "questions": "- Is the proposed method limited to deterministic reference (trajectory)?\n- Is the sampling procedure after which positions will be unmasked at each step is deterministic (based on top-1 sampling)?\n- Could you clarify the statement \"which may be due to the variance in evaluation\" in Line 376. The meaning is unclear.", "limitations": "", "rating": 2, "confidence": 3, "soundness": 1, "presentation": 4, "contribution": 2, "strengths": "- The idea of merging multiple sampling steps into a single step is simple yet presents an interesting direction for accelerating generation.\n- The proposed merging strategy significantly speeds up sampling while maintaining performance, at least in the scenario described in Section 3.", "weaknesses": "**Concerns about Diversity of Generated Samples**\n- The approach relies heavily on a pre-computed, deterministic trajectory. I am concerned that this may lead to reduced diversity in the generated samples, especially if the sampling algorithm is entirely deterministic (see my question below).\n- Whether generation distribution given by the merging algorithm is close to that given by the full sampling in a distribution sense is not discussed.\n- While the experiments evaluate generation performance in terms of accuracy and speed, they do not assess diversity (e.g., entropy used in LLM field), which is a crucial aspect of generation quality.\n- The generation tasks on mathematical and code datasets are insufficient for evaluating diversity.\n\n**Performance Limitations of NI Sampling**\n- It is not clear what are missing components to fill the significant performance gap between NI sampling and ideal sampling presented in Section 3. The ablation study on network size suggests that this gap is not due to the network's expressive power. \n- Section 4.3 lists the network inputs, but it remains unclear what other components might be missing to fill the performance gap. It would be appreciated if insights into it can be more elaborated.\n\n**Experiments**\n- The proposed method is only compared with full sampling and the simple threshold-based method. However, there are related works [1,2] that should be discussed at least, even if they were originally proposed for vision domains.\n- It would be beneficial to present Figure 4 and related figures as time vs. accuracy plots.\n\n[1] Jose Lezama et al., \"Improved Masked Image Generation with Token-Critic,\" ECCV 2022.  \n[2] Jose Lezama et al., \"Discrete Predictor-Corrector Diffusion Models for Image Synthesis,\" ICLR 2023.\n\n**Typos**\n- \"gsample\" in Line 70", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "57eNqlPls6", "reviewer": "Reviewer_Duu3", "summary": "This paper studies acceleration methods for masked diffusion sampling. While masked diffusion models accept efficient prallel sampling/unmasking of multiple tokens without significant performance drop, how much speedup we can achieve has been unknown and perople used heuristic methods to determine the sampling order and parallelization magnitude. This paper introduces Trajectory-Preserving-Order, which measures to how much extent we can reduce the redundunt function evaluations of masked diffusion without changing the one-by-one sampling trajectory. In math and programming benchmarks with the top-1 probability argmax sampler, the authors confirm that up to 24x acceleration is theoretically possible. By learning the Trajectory-Preserving-Order by a neural network (neural indicator learning), NI sampling then achievs up to 14x acceleration, outperforming the threshold-based heuristic.", "questions": "While I lean to acceptance, I suspect that the NI sampler just conducts semi-AR sampling (that is not a major issue though). To better understand the method, I have the following question:\n- [Q1] Trajectory-Preserving-Order may contain poisitions with a random coincidence such as comma and period. However, NI sampling in my intuition, when generalized well, follows mostly the natural order of words. That can result in a gap between the two, which might be very difficult to learn. I would like to see which word is sampled in which step, in examples like Tables 9 & 10, together with the result of Trajectory-Preserving-Order.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper is well written, and offers an evidence that the optimization of sampling positions in masked diffusion can greatly accelerate the sampling process, *even without changing the output/trajectory*, through the empirical studies with Trajectory-Preserving-Order. This offers a deeper understanding of masked diffusion sampling from a novel angle. The neural indicator learning also has a straightforward learning signal and succeeds in actually accelerating masked diffusions.", "weaknesses": "- [W1] A primary weakness is that this method works only for almost deterministic sampling strategies and tasks such as argmax sampling and math/programming tasks. Thus it does not align well with the direction of reproducing randomness, such as unconditional generation and creative tasks requiring diversity.\n- [W2] As the authors admit in Secion 7, the performance of NI sampler in this paper is far from that of Trajectory-Prserving-Order and closer to the heuristic Threshold sampler without any training. I'm not sure how difficult learning the Trajectory-Preserving-Order is in reality.\n\nTypos etc:\n- L70: gsampling -> sampling\n- L106: I find many `\\cdots` for sequences (e.g., $1, \\cdots, R$), but the standard notation aligning with amsmath (when the separator symbols are low) is `\\ldots` (e.g., $1, \\ldots, R$). It is not at all essential.\n- L182: minize -> minimize\n- L209: Preversing -> Preserving (same typo at **many** other positions as well)\n- L272: Indicato -> Indicator", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)