Measured Coupling for Parallel Decoding in Diffusion Language Models

Tal Oved; Lior Wolf

cs.CLNov 5, 2026

Measured Coupling for Parallel Decoding in Diffusion Language Models

Authors: Tal Oved, Lior Wolf

Organizations: Tel Aviv University

Abstract

Diffusion language models can write many tokens per step, but must choose which tokens are safe to commit together. Mean-field decoding penalizes pairs of positions whose predictions look alike. We replace this proxy with the coupling itself: how much committing one token changes the model's confidence at another. Cheap probes inside the main forward pass measure it, with no training or extra model, and it plugs into the same mean-field rule. On three models and four benchmarks, our method needs far fewer steps than mean-field decoding and is faster in most settings, also with a key-value cache, at an average accuracy cost of 1.1 points. Shuffling the couplings or replacing them with a uniform bonus loses this gain.

Figures & tables

Figure 1: Overview of measured coupling. Top: mean-field decoding penalizes two positions whose predicted distributions overlap (left). We instead measure how much revealing token $j$ changes the margin of token $i$ , which shows support as well as conflict (right). Bottom: one decoding step. The main tokens and $K$ probe rows, one for each of the most confident candidates, run in a single forward pass. Each probe row reveals one candidate and recomputes only the masked positions, reading the main keys and values of the same layer everywhere else. The resulting margins give the coupling $\Delta_{ij}=c_{i}-c_{i}^{(j)}$ , which enters the mean-field update, and positions with $q_{i}\geq\tau$ are committed for the next step.

Bench.	Method	Acc. $\uparrow$	NFE $\downarrow$	Step red. $\uparrow$	TPS $\uparrow$	Time (s) $\downarrow$	$\Delta$ vs. MF [95% CI]
GSM8K	One token	77.94	256.0	1.0 $\times$	7.2	32.42
	Fast-dLLM	78.17	78.5	3.3 $\times$	23.2	10.00
	DAWN	77.79	55.8	4.6 $\times$	31.7	7.29
	MF	76.57	65.4	3.9 $\times$	28.6	8.10
	Ours	76.27	40.1	6.4 $\times$	40.3	5.86	$-0.3$ [ $-2.5$ , $+1.9$ ]
MATH	One token	33.82	256.0	1.0 $\times$	9.4	26.21

Table 1: Results on LLaDA-8B-Instruct. Step reduction is

256/\text{NFE}

, relative to one-token-per-step decoding, and is not a runtime speedup. TPS (tokens per second) is measured for all methods of a benchmark in one job on the same GPU (GSM8K on its first 200 questions, MATH on 105 problems with 15 per subject, HumanEval and MBPP on the full test sets). Time is the wall-clock time per answer in seconds, from the same job. Bold marks the best value in each efficiency column. Ours is measured coupling, MF is mean-field decoding, and one token is one-token-per-step decoding. The last column is the accuracy difference from mean-field decoding in points, with a paired 95% confidence interval.

Bench.	Method	Acc. $\uparrow$	NFE $\downarrow$	Step red. $\uparrow$	TPS $\uparrow$	Time (s) $\downarrow$	$\Delta$ vs. MF [95% CI]
GSM8K	One token	80.59	256.0	1.0 $\times$	6.9	31.82
	Fast-dLLM	81.12	75.4	3.4 $\times$	22.5	9.70
	DAWN	81.35	54.4	4.7 $\times$	30.7	7.11
	MF	81.35	64.1	4.0 $\times$	28.2	7.78
	Ours	80.36	40.2	6.4 $\times$	37.3	5.93	$-1.0$ [ $-2.9$ , $+0.9$ ]
MATH	One token	33.96	256.0	1.0 $\times$	8.8	25.62

Table 2: Results on LLaDA-1.5. Our method uses the setting chosen on LLaDA-8B-Instruct, with no new tuning. Columns and TPS measurement as in Table 1 . Bold marks the best value in each efficiency column.

Bench.	Method	Acc. $\uparrow$	NFE $\downarrow$	Step red. $\uparrow$	TPS $\uparrow$	Time (s) $\downarrow$	$\Delta$ vs. MF [95% CI]
GSM8K	One token	75.06	256.0	1.0 $\times$	4.8	27.17
	Fast-dLLM	74.98	62.2	4.1 $\times$	19.1	6.71
	DAWN	72.25	52.0	4.9 $\times$	22.3	5.79
	MF	76.19	69.2	3.7 $\times$	17.0	7.41
	Ours	75.13	41.4	6.2 $\times$	21.1	5.67	$-1.1$ [ $-3.2$ , $+1.1$ ]
MATH	One token	38.18	256.0	1.0 $\times$	11.2	22.37

Table 3: Results on Dream-v0-Instruct-7B. Our method uses the setting calibrated on Dream (

K=20

\tau=0.95

). Columns and measurement as in Table 1 . Bold marks the best value in each efficiency column.

Coupling	Acc. / NFE	$\Delta$ [95% CI]	Acc. / NFE	$\Delta$ [95% CI]
	GSM8K		MBPP
Measured (ours)	76.27 / 40.1		29.00 / 32.3
Shuffled across pairs	40.11 / 44.7	$-36.2$ [ $-39.1$ , $-33.2$ ]	8.60 / 30.9	$-20.4$ [ $-24.2$ , $-16.6$ ]
Support only	53.37 / 39.8	$-22.9$ [ $-25.7$ , $-20.1$ ]	15.40 / 29.3	$-13.6$ [ $-16.9$ , $-10.3$ ]
Conflict only	78.24 / 72.2	$+2.0$ [ $-0.2$ , $+4.1$ ]	32.20 / 51.9	$+3.2$ [ $+0.0$ , $+6.4$ ]
Partially symmetrized	74.22 / 37.4	$-2.1$ [ $-4.1$ , $+0.0$ ]	26.00 / 30.8	$-3.0$ [ $-5.5$ , $-0.5$ ]

Table 4: Controls on the measured coupling, LLaDA-8B-Instruct, full test sets (accuracy % / NFE). Each row changes only the coupling of our method on the measured pairs.

\Delta

is the accuracy difference from our method with a paired 95% confidence interval. The constants (

b=0.02

and

0.05

) were chosen on 100 GSM8K training questions. Appendix I defines each control.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Pairwise term	GSM8K	HumanEval	MBPP
None (confidence only)	78.09 / 52.1	35.98 / 49.6	26.80 / 42.7
Similarity $D$ (MF)	76.57 / 65.4	37.20 / 61.7	30.20 / 51.6
Measured $\Delta$ (ours)	76.42 / 37.4	34.76 / 37.7	28.60 / 30.7

Appendix

Table 5: Effect of the pairwise term (accuracy % / NFE). All rows use block 20 and

\tau=0.85

, so only the pairwise term changes. Our row uses the selected probes (

K=12

, inside the main pass) at this common threshold.

Figure 2: Accuracy against wall-clock time per answer on LLaDA-8B-Instruct. Curves sweep

\tau

from 0.75 to 0.95 for our method and for mean-field decoding. Gray points are Fast-dLLM and DAWN at their published settings. All points of a panel are timed on the same GPU.

Bench.	Method	Acc. $\uparrow$	NFE $\downarrow$	Time (s) $\downarrow$	Step time (s) $\downarrow$
GSM8K	Fast-dLLM	79.38	85.6	2.31	1.27
	MF	77.03	61.6	2.06	1.02
	Ours	75.44	43.6	1.96	0.92
MBPP	Fast-dLLM	26.60	68.5	1.84	0.96
	MF	24.00	44.9	1.56	0.70
	Ours	28.00	30.9	1.51	0.65

Appendix

Table 6: Methods with the dual key-value cache on LLaDA-8B-Instruct, block 32 (full test sets). Time is the wall-clock time per answer in seconds. Step time excludes the full passes at block starts, which cost about the same for every method (1.04 s per answer on GSM8K, 0.86 to 0.88 s on MBPP).

Figure 3: Accuracy against the number of steps per answer on LLaDA-8B-Instruct. Curves sweep the threshold

\tau

: 0.75 to 0.95 for our method and for mean-field decoding, and 0.7 to 0.85 for the confidence-only rule (0.8 and 0.85 on code). Gray points are Fast-dLLM and DAWN at their published settings; the dashed line is one-token-per-step decoding.

Figure 4: Accuracy against steps per answer on Dream-v0-Instruct-7B. Curves sweep

\tau

over 0.9, 0.95, and 0.97 at block 20, and up to 0.99 for our method on MATH. Our curve uses

K=12

. The star is our calibrated setting from Table 3 . Gray points are Fast-dLLM and DAWN; the dashed line is one-token-per-step decoding.

$K$	Ours
4	76.42 / 45.0
8	76.27 / 41.3
12	76.27 / 40.1
20	76.27 / 38.8

Appendix

Table 7: Ablations on GSM8K (accuracy % / NFE). Left: number of probes

K

(

\tau=0.9

, block 20). Right: block size, for our method with

K=12

, with

K

scaled with the block (

K=12

at blocks 16 and 20,

K=20

at 32,

K=40

at 64), and for mean-field decoding (

\tau=0.85

Run	Positions	$\|S\|=2$ – $3$	$\|S\|=4$ – $7$	$\|S\|\geq 8$	All
GSM8K, block 20	7,880	3.7%	7.8%	15.6%	10.0%
MBPP, block 20	7,981	4.5%	7.3%	14.2%	9.4%
GSM8K, block 32	8,199	2.4%	5.7%	9.2%	7.0%
GSM8K, block 64	13,147	1.4%	3.7%	10.8%	7.2%

Appendix

Table 8: Additive prediction against the jointly measured margin, LLaDA-8B-Instruct with our method (

K=12

\tau=0.9

;

K

scaled with the block at 32 and 64). Share of remaining (uncommitted) positions whose decision at

\tau

flips, by the size of the committed set. For the committed tokens, see the text.

	GSM8K		MBPP
Coupling	Same set	Decisions changed	Same set	Decisions changed
Block-local, all candidates	70%	4.2%	71%	4.4%
Exact, probed columns only	55%	10.6%	55%	11.2%
Ours (inside the main pass)	50%	11.8%	47%	12.5%
No coupling	33%	27.3%	36%	25.2%

Appendix

Table 9: Commit decisions on the same decoding states, against exact full-sequence probes for every candidate (50 test questions each, our method with

K=12

and

\tau=0.9

; steps with only end-of-text padding left are excluded).

Task	Decoder	$\tau$	NFE	$\|S\|=2$ – $3$	$\|S\|=4$ – $7$	$\|S\|\geq 8$	All
GSM8K	Ours	0.9	41.1	0.4%	0.8%	0.8%	0.8%
	Conf.	0.75	41.9	6.6%	4.9%	1.7%	2.9%
	Conf.	0.9	66.3	0.2%	0.3%	0.2%	0.2%
	MF	0.85	68.0	0.5%	0.3%	0.1%	0.2%
MBPP	Ours	0.9	35.0	1.2%	0.8%	0.7%	0.8%
	Conf.	0.8	41.5	5.0%	3.3%	1.0%	2.3%

Appendix

Table 10: Committed tokens overturned by the rest of their commit set, by set size (50 test questions each; LLaDA-8B-Instruct unless marked Dream). Conf. is the confidence threshold.

	Ours ( $K=20$ , $\tau=0.95$ )		Mean-field ( $\tau=0.95$ )
	GSM8K	MBPP	GSM8K	MBPP
$R$	Acc. % / NFE	Acc. % / NFE	Acc. % / NFE	Acc. % / NFE
1	75.36 / 42.9	56.00 / 36.1	75.66 / 113.3	57.80 / 118.3
2	75.13 / 41.4	53.60 / 34.4	76.19 / 69.2	54.20 / 77.2
3	74.15 / 40.2	56.00 / 34.3	75.89 / 110.3	58.40 / 113.9
4	74.22 / 39.8	54.80 / 33.8	75.82 / 70.2	54.00 / 79.2

Appendix

Table 11: Number of mean-field iterations

R

on Dream-v0-Instruct-7B (accuracy % / NFE). The main results use

R=2

(bold).

	GSM8K		HumanEval		MBPP
Method	Acc. % $\uparrow$	NFE $\downarrow$	Acc. % $\uparrow$	NFE $\downarrow$	Acc. % $\uparrow$	NFE $\downarrow$
Mean-field (MF)	76.57	65.4	37.20	61.7	30.20	51.6
Learned coupling	76.72	53.1	37.20	54.2	28.80	48.7

Appendix

Table 12: Learned coupling from the next step on LLaDA-8B-Instruct (block 20,

\rho=0.05

). Mean-field decoding is as in Table 1 .

Held-out set	Probes	$\tau$	Accuracy $\uparrow$	NFE $\downarrow$
GSM8K train (100)	12, inside main pass, masked only	0.90	91.0	39.3
	12, inside main pass, masked only	0.85	84.0	37.2
	8, inside main pass, masked only	0.90	88.0	40.1
	8, inside main pass, masked only	0.85	86.0	38.4
	8, separate call	0.85	90.0	34.3
	8, separate call	0.90	85.0	36.2

Appendix

Table 13: Every candidate setting of our method scored on each held-out set. “Masked only” probes recompute only the still-masked positions, and “separate call” probes run after the main pass instead of inside it. The selected setting, chosen on GSM8K, is in bold. MBPP validation is a check on code. Mean-field decoding is shown for reference only; its setting is the published one.

Held-out set	Probes	$\tau$	Accuracy $\uparrow$	NFE $\downarrow$
GSM8K train (100)	20, inside main pass, masked only	0.95	74.0	42.2
	20, inside main pass, masked only	0.90	73.0	37.9
	8, inside main pass, masked only	0.95	71.0	42.8
	12, inside main pass, masked only	0.90	71.0	37.6
	8, inside main pass, masked only	0.90	70.0	38.2
	12, inside main pass, masked only	0.97	69.0	44.4

Appendix

Table 14: Held-out candidates on Dream-v0-Instruct-7B, with the same held-out sets and rule as Table 13 . The selected setting, chosen on GSM8K, is in bold; MBPP validation selects the same one. Mean-field decoding is shown for reference only; its setting is the published Dream one.

Diffusion Language Model Inference

Diffusion language model inference studies how to generate, guide, and serve text efficiently through iterative denoising and parallel token updates.

See all 187 papers

Jun 14, 2026

Mean-Field Parallel Decoding for Discrete Diffusion Language Models

Discrete diffusion language models enable parallel token generation, offering a pathway to low-latency decoding. However, selecting tokens independently by marginal confidence limits effective parallelism: tokens that appear reliable in isolation can form incompatible configurations when several positions are updated at once. We introduce a training-free decoding framework that coordinates these parallel updates. At each forward pass, the method assigns a commit score to each masked position and refines these scores using pairwise interactions derived from the model's predictive distributions. A variational relaxation yields a simple fixed-point update that suppresses conflicting simultaneous commitments within a single forward pass. This mechanism allows the decoder to commit more tokens in parallel while maintaining competitive generation quality. The method is lightweight, requires no auxiliary model or retraining, and drops into existing diffusion decoding pipelines without modification. Experiments on reasoning and code-generation benchmarks show consistent improvements in the quality-latency trade-off.

Tamim Zoabi, Ameen Ali, Liran Ringel +1

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

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

Ask PDF

Jun 2, 2026

Supportive Token Revealing for Fast Diffusion Language Model Decoding

Discrete diffusion language models can generate text efficiently by updating multiple masked positions in parallel, but this parallelism introduces a quality-latency trade-off. Aggressive decoding may commit mutually dependent tokens too early, while conservative decoding requires many denoising steps. Existing methods address this tension by deciding which tokens are safe to reveal using confidence or dependency criteria. However, avoiding unsafe commits does not necessarily make the remaining masked sequence easy to decode, since uncertain tokens may depend on masked tokens, creating a bottleneck for denoising steps. We propose AXON, a training-free module that can be added on top of existing parallel decoding strategies for diffusion language models. Rather than replacing the base decoder, AXON monitors the remaining uncertain masked tokens and intervenes only when their current state suggests that additional context is needed. It then shifts the criterion from which tokens are safest to reveal to which confident reveals would best support later denoising. AXON selects anchors, confident masked tokens that uncertain positions attend to, using attention, uncertainty, and confidence signals. Experiments on reasoning and code-generation benchmarks across multiple diffusion language models show that AXON improves the quality-latency trade-off of existing parallel decoders, often reducing the number of function evaluations while maintaining or improving accuracy.

Giries Abu Ayoub, Mario Barbara, Lluís Pastor-Pérez +4

Department of Computer Science, University of Haifa · Institute for AI, University of Stuttgart · IMPRS-IS

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

Ask PDF

Oct 22, 2026

SetCommit: Coordinated Token Commitments for Few-Step Diffusion Language Model Decoding

Diffusion language models generate multiple token positions in parallel, but few-step decoding makes parallel commitment difficult: tokens that appear reliable individually may be jointly inconsistent, while premature commitments can prevent later context-dependent refinement. In this paper, we introduce SetCommit, a training-free method for coordinated parallel token commitment. Instead of selecting tokens independently by confidence, SetCommit selects a commitment set that balances token reliability against local joint commitment risk, preserving conditioning opportunities for subsequent denoising steps. It further performs budget-neutral local exchanges, replacing vulnerable earlier commitments when newly revealed context supports more reliable alternatives. SetCommit requires no retraining or architectural modification and is compatible with existing diffusion decoding strategies. Experiments with LLaDA-8B and Dream-7B on GSM8K, MATH-500, MBPP, and HumanEval show accuracy/pass@1 gains of up to 7.4 percentage points over native decoding at 32 model evaluations. At common quality targets, geometric-mean incremental speedups range from

1.44

1.74\times

. Code is available at https://github.com/Kissmetothemoon/SetCommit.

Yuannuo Feng, Yubing Ye, Wenshuai Yao +4

Beihang University · The University of Hong Kong · Peking University

Diffusion Language Model Inference Diffusion Language Models Parallel Decoding for Diffusion Language Models

Ask PDF

Abstract

Diffusion Language Model Inference

Explore similar work

Mean-Field Parallel Decoding for Discrete Diffusion Language Models

Supportive Token Revealing for Fast Diffusion Language Model Decoding

SetCommit: Coordinated Token Commitments for Few-Step Diffusion Language Model Decoding