DOS: Dependency-Oriented Sampler for Masked Diffusion Language Models

Xueyu Zhou; Yangrong Hu; Jian Huang

cs.CLMar 16, 2026

DOS: Dependency-Oriented Sampler for Masked Diffusion Language Models

Authors: Xueyu Zhou, Yangrong Hu, Jian Huang

Organizations: Department of Data Science and AI The Hong Kong Polytechnic University, Hong Kong, China

Abstract

Masked diffusion language models (MDLMs) have recently emerged as a new paradigm in language modeling, offering flexible generation dynamics and enabling efficient parallel decoding. However, existing decoding strategies for pre-trained MDLMs predominantly rely on token-level uncertainty criteria, while largely overlooking sequence-level information and inter-token dependencies. To address this limitation, we propose Dependency-Oriented Sampler (DOS), a training-free decoding strategy that leverages inter-token dependencies to inform token updates during generation. Specifically, DOS exploits attention matrices from transformer blocks to approximate inter-token dependencies, emphasizing information from unmasked tokens when updating masked positions. Empirical results demonstrate that DOS consistently achieves superior performance on both code generation and mathematical reasoning tasks. Moreover, DOS can be seamlessly integrated with existing parallel sampling methods, leading to improved generation efficiency without sacrificing generation quality.

Figures & tables

Figure 1: Accuracy on GSM8K using LLaDA-Instruct-8B ( Nie et al., 2025 ) with a fixed generation length of 512 tokens. Block size 32 corresponds to block-wise decoding, while 512 represents the single-block setting. Existing methods (Fast-dLLM ( Wu et al., 2025 ) , KLASS ( Kim et al., 2025b ) , Confidence ( Chang et al., 2022 ) ) degrade under large block sizes, whereas DOS (ours) remains consistent and robust across both settings.

Panel of Figure 2 — Figure 2: A toy example for parallel decoding in MDLMs, where $C$ denotes the prompt or unmasked tokens and $\{X_{i}\}_{i=1}^{4}$ are masked tokens. Different factorizations of $p(X_{1},X_{2},X_{3},X_{4}|C)$ correspond to different parallel decoding orders. Only factorizations that respect the underlying dependency structure are able to recover the true joint distribution, while improper independence assumptions lead to distribution mismatch.

Models	Methods	HumanEval	MBPP	GSM8K	MATH500
LLaDA-Instruct-8B	Confidence	27.44	22.80	31.69	21.40
	Entropy	27.44	19.40	33.13	18.60
	Margin	26.83	26.40	33.76	22.20
	DOS	42.68	38.40	84.31	41.60
Dream-v0-Instruct-7B	Confidence	28.66	39.60	31.08	13.60
	Entropy	25.61	33.20	30.55	10.20

Table 1: Performance of different decoding strategies under top-1 sampling, in which all methods generate the entire sequence within a single block. For the code benchmarks (HumanEval and MBPP), the generation length is fixed to 256 tokens, while for the math benchmarks (GSM8K and MATH500), the generation length is 512 tokens.

		HumanEval		MBPP		GSM8K		MATH500
Models	Methods	Acc $\uparrow$	NFE $\downarrow$	Acc $\uparrow$	NFE $\downarrow$	Acc $\uparrow$	NFE $\downarrow$	Acc $\uparrow$	NFE $\downarrow$
LLaDA-Instruct-8B	Fast-dLLM	28.66	83.26	23.00	78.93	33.06	195.46	20.80	219.54
	KLASS	29.88	107.99	24.00	97.09	28.43	265.52	19.60	300.48
	Confidence $+$ EB	28.66	160.10	22.80	167.43	32.15	299.59	21.00	327.07
	Entropy $+$ EB	29.88	175.71	18.80	183.72	25.09	309.12	19.00	334.24
	Margin $+$ EB	26.22	159.58	26.60	166.00	38.51	293.45	21.40	324.10

Table 2: Performance of different decoding strategies within a single block. EB stands for EB-sampler ( Ben-Hamu et al., 2025 ) , a parallel decoding method that can be combined with different scores. Acc denotes task accuracy, and NFE denotes the number of model forward evaluations, reflecting decoding efficiency. For HumanEval and MBPP, the generation length is fixed to 256 tokens, while for GSM8K and MATH500, the generation length is 512 tokens.

Method	HumanEval	MBPP	GSM8K	MATH500
Confidence	40.85	38.20	82.56	38.80
Entropy	41.46	37.80	82.79	40.00
Margin	40.85	37.40	82.64	38.80
Fast-dLLM	41.46	38.80	83.40	39.60
KLASS	42.07	38.60	83.02	40.20
DOS (w/o block)	42.68 ( +0.61 )	38.40 ( -0.40 )	84.31 ( +0.91 )	41.60 ( +1.40 )

Table 3: Performance comparison on HumanEval, MBPP, GSM8K, and MATH500 using LLaDA-Instruct-8B. All baseline methods (Confidence, Entropy, Margin, Fast-dLLM, and KLASS) adopt the block diffusion with a fixed block size of 32. DOS (ours) is evaluated both without block partitioning (single-block) and with block partitioning (block size = 32).

Figure 3: Accuracy on HumanEval using LLaDA-Instruct-8B with a fixed generation length of 256 under varying block sizes. Existing decoding strategies are sensitive to block size and degrade as the block size increases, whereas DOS (ours) demonstrates strong robustness to block size variation and maintains superior consistency across all settings.

Figure 4: Accuracy of DOS and DOS+EB on the HumanEval benchmark using the LLaDA-Instruct-8B model, where the x-axis indicates the transformer layer from which the attention matrix is extracted.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Task	LLaDA		Dream
	Conf	KL	Conf	KL
MATH	0.6	0.010	0.9	0.005
GSM8K	0.6	0.015	0.9	0.001
HumanEval	0.9	0.010	0.8	0.001
MBPP	0.7	0.010	0.9	0.001

Appendix

Table 4: Threshold configurations for KLASS.

	HumanEval	MBPP	GSM8K	MATH500
LLaDA-Instruct-8B
DOS (w/o block)	13	9	7	7
DOS (w/ block)	13	9	16	30
DOS $+$ EB (w/o block)	29	5	4	3
DOS $+$ EB (w/ block)	13	1	1	7
Dream-vo-Instruct-7B

Appendix

Table 5: Transformer layers that yield the best performance for DOS under different settings. For DOS , w/o block corresponds to single-block decoding where the block length equals the generation length, while w/ block applies block partitioning with a fixed block size of 32.

	HumanEval		MBPP		GSM8K		MATH500
Methods	Acc $\uparrow$	NFE $\downarrow$	Acc $\uparrow$	NFE $\downarrow$	Acc $\uparrow$	NFE $\downarrow$	Acc $\uparrow$	NFE $\downarrow$
Top-1
Confidence	40.85	256.00	38.20	256.00	82.56	512.00	38.80	512.00
Entropy	41.46	256.00	37.80	256.00	82.79	512.00	40.00	512.00
Margin	40.85	256.00	37.40	256.00	82.64	512.00	38.80	512.00
DOS (w/o block)	42.68	256.00	38.40	256.00	84.31	512.00	41.60	512.00

Appendix

Table 6: Additional results through top-1 sampling and parallel sampling on HumanEval, MBPP, GSM8K, and MATH500 using LLaDA-Instruct-8B model ( Nie et al., 2025 ) with a fixed block size of 32. EB denotes the EB-sampler ( Ben-Hamu et al., 2025 ) , a parallel decoding method that can be combined with different scoring strategies. Acc denotes task accuracy, and NFE denotes the number of model forward evaluations. For HumanEval and MBPP, the generation length is fixed to 256 tokens, while for GSM8K and MATH500, the generation length is 512 tokens. For DOS , w/o block corresponds to single-block decoding, where the block length equals the generation length, while w/ block applies block partitioning with block size 32.

	HumanEval		MBPP		GSM8K		MATH500
Methods	Acc $\uparrow$	NFE $\downarrow$	Acc $\uparrow$	NFE $\downarrow$	Acc $\uparrow$	NFE $\downarrow$	Acc $\uparrow$	NFE $\downarrow$
Top-1
Confidence	58.54	256.00	50.02	256.00	81.50	512.00	40.08	512.00
Entropy	57.32	256.00	50.80	256.00	83.93	512.00	42.20	512.00
Margin	56.71	256.00	50.40	256.00	81.05	512.00	41.80	512.00
DOS (w/o block)	59.15	256.00	50.60	256.00	80.21	512.00	45.00	512.00

Appendix

Table 7: Additional results through top-1 sampling and parallel sampling on HumanEval, MBPP, GSM8K, and MATH500 using Dream-v0-Instruct-7B model ( Ye et al., 2025 ) with a fixed block size of 32. EB denotes the EB-sampler ( Ben-Hamu et al., 2025 ) , a parallel decoding method that can be combined with different scoring strategies. Acc denotes task accuracy, and NFE denotes the number of model forward evaluations. For HumanEval and MBPP, the generation length is fixed to 256 tokens, while for GSM8K and MATH500, the generation length is 512 tokens. For DOS , w/o block corresponds to single-block decoding, where the block length equals the generation length, while w/ block applies block partitioning with block size 32.

Figure 5: Accuracy on HumanEval using Dream-v0-Instruct-7B with a fixed generation length of 256 under varying block sizes. Existing decoding strategies are sensitive to block size and degrade as the block size increases, whereas DOS (ours) demonstrates strong robustness to block size variation and maintains superior consistency across all settings.

	HumanEval		MBPP		GSM8K		MATH500
Methods	Acc $\uparrow$	NFE $\downarrow$	Acc $\uparrow$	NFE $\downarrow$	Acc $\uparrow$	NFE $\downarrow$	Acc $\uparrow$	NFE $\downarrow$
LLaDA-Instruct-8B
DOS (w/o block)	41.46	256.00	38.20	256.00	84.23	512.00	41.00	512.00
DOS (w/ block)	42.07	256.00	37.80	256.00	84.23	512.00	40.40	512.00
DOS $+$ EB (w/o block)	43.90	104.93	38.60	76.34	83.55	164.26	39.60	258.37
DOS $+$ EB (w/ block)	40.85	107.11	39.00	80.90	84.00	170.06	40.60	260.90

Appendix

Table 8: Results of DOS and DOS

+

EB with attention matrices extracted from the first transformer layer. Experiments are conducted on LLaDA-Instruct-8B and Dream-v0-Instruct-7B under both single-block and block-based decoding settings. For HumanEval and MBPP, the generation length is fixed to 256 tokens, while for GSM8K and MATH500, the generation length is 512 tokens. DOS w/o block denotes single-block decoding where the block length equals the generation length, while DOS w/ block applies block partitioning with a fixed block size of 32. Acc denotes task accuracy, and NFE denotes the number of model forward evaluations, reflecting decoding efficiency.

	HumanEval				MBPP				GSM8K				MATH500
Methods	Acc $\uparrow$	NFE $\downarrow$	TPS $\uparrow$	Time $\downarrow$	Acc $\uparrow$	NFE $\downarrow$	TPS $\uparrow$	Time $\downarrow$	Acc $\uparrow$	NFE $\downarrow$	TPS $\uparrow$	Time $\downarrow$	Acc $\uparrow$	NFE $\downarrow$	TPS $\uparrow$	Time $\downarrow$
Top-1
Confidence	40.85	256.00	4.22	32.40	38.20	256.00	5.07	20.21	82.56	512.00	4.96	66.22	38.80	512.00	6.46	67.59
Entropy	41.46	256.00	4.28	33.07	37.80	256.00	5.04	20.62	82.79	512.00	4.78	68.37	40.00	512.00	6.28	69.57
Margin	40.85	256.00	4.07	33.10	37.40	256.00	4.84	20.56	82.64	512.00	4.78	68.01	38.80	512.00	6.34	69.35
DOS (w/o block)	42.68	256.00	4.71	32.18	38.40	256.00	5.76	19.84	84.31	512.00	5.00	65.47	41.60	512.00	6.56	66.85

Appendix

Table 9: Additional results through top-1 sampling and parallel sampling on HumanEval, MBPP, GSM8K, and MATH500 using LLaDA-Instruct-8B model ( Nie et al., 2025 ) with a fixed block size of 32. EB denotes the EB-sampler ( Ben-Hamu et al., 2025 ) , a parallel decoding method that can be combined with different scoring strategies. Acc denotes task accuracy, and NFE denotes the number of model forward evaluations. For HumanEval and MBPP, the generation length is fixed to 256 tokens, while for GSM8K and MATH500, the generation length is 512 tokens. For DOS , w/o block corresponds to single-block decoding, where the block length equals the generation length, while w/ block applies block partitioning with block size 32.

Explore similar work

Dependency-Guided Parallel Decoding in Discrete Diffusion Language Models
Apr 2, 2026Liran Ringel, Ameen Ali, Yaniv RomanoDiffusion Language Model Decoding Diffusion Language Models
Department of Computer Science, Technion – Israel Institute of Technology · Department of Electrical and Computer Engineering, Technion – Israel Institute of Technology · Blavatnik School of Computer Science and AI, Tel Aviv, Israel
Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models
Jun 9, 2026Yusuf Sahin, Ahmed Rockey Saikia, Volkan Cevher +1Masked Diffusion Models Parallel Decoding
University of Bern, Bern, Switzerland · EPFL, Lausanne, Switzerland
Improving Sampling for Masked Diffusion Models via Information Gain
Feb 20, 2026Kaisen Yang, Jayden Teoh, Kaicheng Yang +2Diffusion Sampling Diffusion Language Model Decoding
Department of Computer Science and Technology, Tsinghua University · Massachusetts Institute of Technology · Shanghai Jiao Tong University +2

Apr 2, 2026J/K move · Enter open · S save

Dependency-Guided Parallel Decoding in Discrete Diffusion Language Models

Liran Ringel, Ameen Ali, Yaniv Romano

Department of Computer Science, Technion – Israel Institute of Technology · Department of Electrical and Computer Engineering, Technion – Israel Institute of Technology · Blavatnik School of Computer Science and AI, Tel Aviv, Israel

Diffusion Language Model Decoding Diffusion Language Models Parallel Decoding Diffusion Model Inference Acceleration

Open paper AskPDF arXiv

Discrete diffusion language models (dLLMs) accelerate text generation by unmasking multiple tokens in parallel. However, parallel decoding introduces a distributional mismatch: it approximates the joint conditional using a fully factorized product of per-token marginals, which degrades output quality when selected tokens are strongly dependent. We propose DEMASK (DEpendency-guided unMASKing), a lightweight dependency predictor that attaches to the final hidden states of a dLLM. In a single forward pass, it estimates pairwise conditional influences between masked positions. Using these predictions, a greedy selection algorithm identifies positions with bounded cumulative dependency for simultaneous unmasking. Under a sub-additivity assumption, we prove this bounds the total variation distance between our parallel sampling and the model's joint. Empirically, DEMASK achieves 1.7-2.2

\times

speedup on Dream-7B while matching or improving accuracy compared to confidence-based and KL-based baselines.

Abstract

Explore similar work

Dependency-Guided Parallel Decoding in Discrete Diffusion Language Models

Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models

Improving Sampling for Masked Diffusion Models via Information Gain