Self-Repulsive Sampling for Diffusion Language Models

Michael Helcig; Martin Jaggi

cs.LGSep 30, 2026

Self-Repulsive Sampling for Diffusion Language Models

Authors: Michael Helcig, Martin Jaggi

Organizations: ETH Zürich · EPFL

Abstract

Sampling several responses and voting over their answers can improve a language model's accuracy, but repeated answers limit the benefit of additional samples. Raising temperature increases diversity at a potential cost to per-sample accuracy. We introduce Self-Repulsion (SR), a sampler for masked diffusion language models that uses peer commitments to diversify the pool. At each penalized denoising step, each path lowers a token's logit according to how many peers have committed that token at the same position. Paths share a batched forward pass and then commit in sequence, so later paths observe choices made earlier in the same step. This coupling requires no training or additional forward or backward pass and can produce distinct paths even at temperature zero. When all paths commit a position together from identical logits, the update exactly maximizes total logit minus a convex duplication cost. On LLaDA-8B-Instruct with ten paths and 128 denoising steps, deterministic SR reaches 80.38% plurality accuracy on GSM8K, compared with 70.17% for the unpenalized greedy decoder. At temperature 0.6 and matched model-evaluation budgets, the count penalty improves over self-consistency by 2.06 percentage points in blocks of 32 and 14.50 under pure diffusion. Experiments on GSM8K, MATH and TruthfulQA show that voting gains arise mainly from higher coverage of correct answers, with gains that vary by benchmark and decoding regime.

Figures & tables

Figure 1: One denoising step of Self-Repulsion, four paths, illustrative values. Left: the

K

paths decode in parallel; the column at a still-masked position

j

is the count vector

\mathbf{c}_{k}[j,\cdot]

, the peers committed to each token (here

(2,1,0,0)

). The step runs in three parts: (1) read the column, (2) subtract the penalty

g\,\mathbf{c}_{k}[j,\cdot]

from the logits, (3) commit from the tilted law. Right: the sampling law at

x_{j}

under the model, under temperature and under SR. The tilt down-weights tokens peers hold and preserves probability ratios between tokens with equal counts; temperature reshapes every token, the taken ones included. Values are schematic except the effective tilt

e^{-11.85}

per peer.

	Plurality ( $\uparrow$ )			Coverage ( $\uparrow$ )		Per-sample ( $\uparrow$ )			Disagreement
Method	GSM	MATH	TQA	GSM	MATH	GSM	MATH	TQA	GSM	MATH	TQA
Semi-autoregressive blocks of 32
Self-consistency, $\tau{=}0$ ( $=$ greedy, $K{=}1$ )	70.17	28.25	55.26	70.17	28.25	70.17	28.25	55.26	0.00	0.00	0.00
Self-consistency	77.05	40.14	59.02	89.70	56.81	70.04	30.02	54.44	27.14	62.01	34.09
TAPS	77.69	39.47	58.27	91.00	57.32	66.77	29.26	53.53	35.51	64.47	36.63
ODD ( $\alpha{=}64$ )	77.24	37.03	60.32	91.75	54.27	64.94	24.19	52.35	38.72	70.98	40.97

Table 2: Every method at 1280 network function evaluations; Table 5 repeats it at 2560. GSM8K

n{=}600

, MATH

n{=}492

, TruthfulQA

n{=}684

, two to five seeds; TruthfulQA coverage omitted (94.37 chance). Rows run at

\tau{=}0.6

except the three

\tau{=}0

rows: row 1 is the greedy decode, equal to self-consistency (Proposition 1 ), and the two SR, count rows use gate 8 (

\lambda=72

). Pure-diffusion rows carry the end-of-sequence correction (Appendix D.9 ); TAPS (

\dagger

) is uncorrected. Bold: best per column within a regime, disagreement and the

K{=}1

greedy row excluded.

Figure 2: Pass@

k

, self-consistency against SR (count statistic), Table 2 ’s campaigns; shading is the gap, the printed number its size at

k{=}10

. SR is below self-consistency only at

k\leq 2

, the per-sample end of the curve, and above it at every larger

k

, including MATH; all methods and values in Table 17 .

Conditions on	Varies	Disagr.	Per-sample	Plur.	Cov.	Temp. here	Gap
Nothing (temperature)	$\tau{=}0.6$	15.02	53.42	54.53	63.17	–	–
Nothing (temperature)	$\tau{=}1.0$	25.06	53.76	57.44	68.83	–	–
Nothing (temperature)	$\tau{=}1.2$	28.25	53.24	57.82	71.33	–	–
Nothing (temperature)	$\tau{=}1.5$	34.86	51.95	57.83	73.08	–	–
Nothing (temperature)	$\tau{=}2.0$	50.61	47.77	60.72	79.92	–	–
Own distribution	$\lambda{=}1$	24.48	53.94	57.57	69.33	53.74	$+0.20$

Table 3: Temperature against Self-Repulsion under pure diffusion, on the 1200 problem-seed pairs common to all nine configurations. The varied parameter is temperature for the first family and strength

\lambda

for the other two (gate

g=\lambda/(K-1)

). The last two columns interpolate the temperature sweep to each configuration’s own disagreement and give the per-sample accuracy it reaches there and the gap to it.

Appendix figures & tables19 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: One position on the probability simplex, three of its tokens. Each frame shows the pool after

i

paths, the fraction of the paths on each token, as a point of the triangle, against

\mathrm{sparsemax}(\ell/gK)

, marked by a cross. The cascade starts the pool at the mode’s vertex and walks it, one path at a time, onto the cross: after six paths the pool is

(3,2,1)/6

, the nearest point of the

K

-grid to

\mathrm{sparsemax}(\ell/gK)=(0.46,0.34,0.21)

(Proposition 2 ).

Figure 4: What the

K

paths commit at one position at

\tau=0

: with the penalty off (a) , at the deployed strength (b) and at a large one (c) . Each grey bar is a token’s model logit

\ell_{v}

and each blue cell is one path committed there, stacked in commit order; the outlined cell is the first path to commit, which faces no peer and pays no penalty. With the penalty off all

K

paths tower on the mode; raising it spreads them, until at large strength each of the top

K

tokens holds one path. Logits are illustrative and the allocation is computed by the update rule; the water level is the common threshold a token’s logit must clear to receive a path.

Configuration	Metric	Seed 0	Seed 1	Seed 2	Seed 3	Seed 4	Mean	SD	Range
SR, 1280	Plurality	80.93	78.89	78.79	78.06	78.89	79.11	1.08	2.87
	Coverage	94.33	92.17	94.00	93.50	93.67	93.53	0.83	2.17
	Per-sample	67.68	68.12	68.72	68.00	68.52	68.21	0.41	1.03
Self-consistency, 1280	Plurality	77.42	77.21	77.06	76.38	77.20	77.05	0.40	1.04
	Coverage	89.50	89.67	90.67	89.50	89.17	89.70	0.57	1.50
	Per-sample	70.03	69.65	70.57	70.02	69.92	70.04	0.33	0.92

Appendix

Table 4: Effect of the decode seed. Each configuration is held fixed, block 32, GSM8K,

n=600

; only the seed varies. The reference configuration and self-consistency are at 1280 evaluations; the third configuration is the same method at 2560. SD and range are taken over the five seeds.

	Plurality ( $\uparrow$ )			Coverage ( $\uparrow$ )		Per-sample ( $\uparrow$ )			Disagreement
Method	GSM	MATH	TQA	GSM	MATH	GSM	MATH	TQA	GSM	MATH	TQA
Semi-autoregressive blocks of 32
Single greedy decode ( $K{=}1$ )	73.00	–	–	73.00	–	73.00	–	–	0.00	–	–
Self-consistency	76.09	38.18	59.84	87.94	55.89	72.31	31.60	56.37	19.93	55.39	28.47
TAPS	75.37	37.54	58.46	90.17	55.39	68.88	30.26	55.42	29.36	57.89	30.80
ODD ( $\alpha{=}64$ )	77.75	37.49	59.31	91.58	56.20	69.14	26.88	54.01	31.00	66.43	37.55

Appendix

Table 5: Both decoding regimes under the protocol of Section 3 (

K=10

, 2560 NFEs, temperature 0.6, two to five seeds per configuration, five on the GSM8K SR, count row and two or three on the rest); TruthfulQA coverage is unreported, since guessing over four options already reaches 94.37. Bold: best per column within a regime, disagreement and the

K{=}1

greedy reference row excluded. The pure-diffusion rows predate the end-of-sequence correction applied in Table 2 (Appendix D.9 ). Baselines are our reproductions (Appendix B ).

$K$	$\lambda$	$g$	NFE	$n$	Seeds	Plurality	Coverage	Per-sample	Disagr.
Fixed gate $g=8.00$
2	8	8.00	512	600	3	71.47	79.72	71.47	25.39
3	16	8.00	768	600	3	74.44	84.39	72.00	26.07
5	32	8.00	1280	600	3	76.88	88.78	71.39	28.03
7	48	8.00	1792	600	3	77.57	91.50	71.48	28.55
Fixed strength $\lambda=64$ , the $K$ -invariant scaling of Corollary 3

Appendix

Table 6: Effect of ensemble size (

K

). Fixed: 256 steps, block 32, temperature 0.6, count statistic, GSM8K,

n=600

. Not budget-matched: NFE

=K\times

steps and is the price of each row. Rows below ten hold the gate at 8.00; the two conventions meet at

K=9

(

\lambda=64

). Rows above ten hold

\lambda=64

, the

K

-invariant scaling of Corollary 3 , with the fixed-gate rows as the diagnostic; self-consistency rows above ten pool seeds, which their i.i.d. draws license. The

\lambda=0

rows are self-consistency, the same loop with the penalty off; a dash is not measured.

Parameter	Setting	Seeds	Plurality	Coverage	Per-sample	Disagr.
Reference	count, $\lambda{=}64$ , $g{=}7.11$	5	79.11	93.53	68.21	35.53
Self-consistency	$\lambda{=}0$ , $g{=}0$	5	77.05	89.70	70.04	27.14
Split at 1280 NFE	5 samples $\times$ 256 steps, $\lambda{=}32$	3	76.88	88.78	71.39	28.03
Statistic	expected count	2	80.65	94.67	66.47	40.60
Statistic	collision-weighted	2	80.07	94.42	61.87	48.63

Appendix

Table 7: The small-effect parameters, as deviations from the reference configuration (GSM8K,

n=600

K=10

, 128 steps, block 32,

\tau=0.6

, count statistic, 1280 network function evaluations). The reference and self-consistency rows are shared by both blocks; the statistic rows move only the statistic, while the split row re-allocates the fixed budget to five samples of 256 steps and so also lowers the strength to

\lambda=32

(gate 8.00). Seeds as stated; higher is better in the accuracy columns; disagreement is a rate, not an objective.

	Expected count $-$ count			Collision-weighted $-$ count
Benchmark	Diffusion	Block 32	Interaction	Diffusion	Block 32	Interaction
GSM8K ( $n=600$ )	$+6.33$	$-1.53$	$+7.86$	$+6.87$	$-4.23$	$+11.10$
MATH ( $n=492$ )	$+1.73$	$-0.46$	$+2.18$	$+4.42$	$-1.97$	$+6.39$
TruthfulQA ( $n=684$ )	$+7.81$	$-3.70$	$+11.51$	$+8.91$	$-1.45$	$+10.36$

Appendix

Table 8: Penalty by decoding regime on per-sample accuracy, paired per problem and seed-averaged over seeds

\{0,1\}

, at 2560 network function evaluations. Each entry is the named penalty minus the count statistic in that regime, positive favouring the named penalty; the interaction is the pure-diffusion effect minus the block-32 effect, taken within problem. These within-problem paired contrasts on two seeds do not reproduce the seed-mean differences of Table 5 .

Regime	$\lambda$	$g$	$n$	Seeds	Plurality	Coverage	Per-sample	Disagr.
Blocks of 32	64	7.11	600	5	79.11	93.53	68.21	35.53
Blocks of 32	0	0.00	600	5	77.05	89.70	70.04	27.14
Pure diffusion	64	7.11	600	2	70.30	87.25	52.58	51.18
Pure diffusion	0	0.00	600	2	54.39	65.83	52.02	21.08

Appendix

Table 9: Effect of the decoding regime. Fixed:

K=10

, 128 steps, temperature 0.6, count statistic, GSM8K,

n=600

. Both pure-diffusion rows predate the released end-of-sequence correction; Table 2 carries the corrected rows and Appendix D.9 reports what the correction is worth. Per-sample accuracy under the parsed-answer denominator is 0.41 higher on both pure-diffusion rows.

Parameter	Recommended	Range	Note
Ensemble size $K$	10	2 to 32	Vote flat past ten at fixed $\lambda$ ; coverage still rising
Strength $\lambda$	64	–	Selected on 200 problems at one seed
Strength across $K$	$\lambda=64$	10 to 32	Hold $\lambda$ , not $g$ (Corollary 3 ); fixed $g$ at $K\geq 20$ costs per-sample accuracy
Denoising steps	128	128 to 256	Vote flat in the step count; 128 is cheaper
Temperature	0.6 or 0	0 to 0.6	Zero (gate 8, $\lambda{=}72$ ) is 1.27 above the sampled configuration ( $\lambda{=}64$ ) at 128 steps and 0.33 below it at 256; uses no randomness
Statistic	Count	–	Ordering is benchmark-dependent

Appendix

Table 10: Recommended settings from the sweeps above. The range column states where the setting is supported by a measurement; a dash means we have no range.

Method	Plur. ( $\uparrow$ )	Cover. ( $\uparrow$ )	Per-s. ( $\uparrow$ )	Distinct	Disagr.
Self-consistency	76.85	87.25	74.15	1.69	17.82
SR, count	79.31	94.00	72.33	2.26	28.22

Appendix

Table 11: Second checkpoint: LLaDA-1.5 on GSM8K,

K{=}10

, semi-autoregressive blocks of 32, 2560 evaluations,

n{=}600

over two seeds (1200 paired observations). Bare means; higher is better for plurality, coverage (pass@10) and per-sample; distinct answers per pool and disagreement are descriptive. Bold marks the better method on each accuracy column. Paired SR minus self-consistency: coverage

+6.75

, plurality

+2.46

, per-sample

-1.82

. The coverage gain differs from the LLaDA-8B-Instruct gain on the same problems by

+0.58

Figure 5: One problem decoded twice at temperature zero, the penalty off above and on below. A row is one position of the canvas, a column is one of the ten samples; in the lower half blue marks a token written under repulsion, grey one written where the block releases the penalty. The section text gives the rest.

	Plurality at own $\tau$ ( $\uparrow$ )			Difference from SR
Method	GSM	MATH	TQA	GSM	MATH	TQA
Self-consistency, $\tau=0.6$	77.05	40.14	59.02	$+2.06$	$-0.10$	$+0.76$
Self-consistency, $\tau=1.0$	78.28	39.47	58.65	$+0.83$	$+0.57$	$+1.12$
TAPS, $\tau=1.0$	77.87	39.31	57.74	$+1.24$	$+0.73$	$+2.03$
ODD ( $\alpha{=}64$ ), $\tau=1.0$	78.23	36.66	60.86	$+0.88$	$+3.38$	$-1.08$
ODD ( $\alpha{=}256$ ), $\tau=1.0$	78.61	37.52	60.12	$+0.50$	$+2.52$	$-0.35$

Appendix

Table 12: Every family at its own temperature, 1280 network function evaluations,

K=10

, blocks of 32. Left: plurality accuracy (

\uparrow

) with each family at a temperature its authors report, self-consistency at both 0.6 and 1.0 since it has no published setting for this model. Right: the mean paired difference of SR with the count statistic at

\tau=0.6

over that row, on the problems the two methods share. GSM8K

n{=}600

, MATH

n{=}492

, TruthfulQA

n{=}684

; two seeds throughout, five for self-consistency at 0.6 and for the count statistic on GSM8K; the

\tau=0

row is one deterministic decode. Baselines are our reproductions.

Configuration	Gate $g$	$\lambda$	Agreement	Distinct	Per-sample ( $\uparrow$ )	Coverage ( $\uparrow$ )	Plurality ( $\uparrow$ )
Parallel, $K=10$ , the deployed design
Penalty off	0	0	100.00	1.00	73.00	73.00	73.00
Count penalty	1	9	77.54	1.91	72.17	89.50	76.97
Count penalty	2	18	75.62	2.04	71.85	90.83	77.31
Count penalty	4	36	73.04	2.18	71.90	92.00	77.97
Count penalty	8	72	68.27	2.43	69.65	92.33	78.60

Appendix

Table 13: The penalty at temperature zero, GSM8K,

n=600

, blocks of 32, 256 steps. One decode per cell. Agreement is the mean fraction of sample pairs returning the same answer; distinct is the mean number of different answers in a pool. The last rows of each block are the sampled references at

\tau=0.6

, seed 0, and the single greedy decode.

Coupling	Gate $g$	$\lambda$	Agreement	Per-sample ( $\uparrow$ )	Coverage ( $\uparrow$ )	Plurality ( $\uparrow$ )
Triangular filtration	4	36	76.53	71.75	90.00	76.12
Triangular filtration	8	72	73.12	71.30	92.00	77.41
Triangular filtration	16	144	67.72	68.67	91.50	77.19
Mutual cascade, symmetric	8	72	68.27	69.65	92.33	78.60

Appendix

Table 14: The cost of the information structure. GSM8K,

n=600

, temperature zero, one decode per cell. Path 0 of the filtration is the released greedy decode, its answers identical across the three configurations on all 600 problems and reading the anchor’s 73.00; the symmetric control is the deployed mutual cascade at the same gate, decoded in the same campaign, reproducing the banked symmetric cell on all 600.

Benchmark	Penalty	$\Delta$ plurality	Coverage	Selectivity	Fixes	Breaks
Semi-autoregressive blocks of 32
GSM8K	count	$+3.92$	$+5.27$	$-1.35$	7.17	3.25
GSM8K	collision-wt.	$+2.92$	$+5.98$	$-3.06$	6.83	3.92
MATH	count	$+2.34$	$+3.09$	$-0.75$	6.00	3.66
MATH	collision-wt.	$+2.13$	$+1.94$	$+0.19$	5.18	3.05
TruthfulQA	count	$+0.88$	$+2.36$	$-1.48$	6.65	5.77

Appendix

Table 15: Decomposition of the gain. Each row is SR against self-consistency in the same regime at 2560 network function evaluations,

K=10

, temperature 0.6, two seeds, paired on the problems both methods share. Coverage and selectivity are the two terms of the split defined in the text; they sum to

\Delta

in every row. Fixes and breaks are the fractions of problem-seed units whose vote changes in each direction. On TruthfulQA ten uniform guesses over four options already cover the correct option with probability 94.37, which the pure-diffusion count row’s coverage matches, so its coverage column is uninformative. Note: the pure-diffusion rows predate the end-of-sequence correction (Appendix D.9 ).

Figure 6: The exchange rate against the decodes. Problems are stratified by the baseline’s per-sample accuracy; each point is one stratum of one benchmark and regime, at 2560 network function evaluations,

K=10

, temperature 0.6, two seeds, paired per problem. Predicted is

(1-a_{\mathrm{SC}})^{K}-(1-a_{SR})^{K}

, the coverage the stratum’s measured per-sample shift buys if the samples were independent. (a) coverage change against baseline accuracy, measured and predicted. (b) the same points, measured against predicted, with the identity line; open markers are strata of fewer than 15 problems. Circles and squares are GSM8K in blocks and under pure diffusion, triangles and diamonds are MATH.

Blocks of 32
	GSM8K				MATH
Sampler	$k{=}1$	2	5	10	$k{=}1$	2	5	10
Self-consistency	72.18	78.58	84.54	87.67	31.60	39.60	48.78	55.89
SR, count	70.93	81.52	90.07	93.83	30.22	40.51	52.12	60.47
SR, coll.-wt.	66.70	80.14	90.20	94.75	28.25	38.57	50.75	58.74
Pure diffusion

Appendix

Table 16: Pass@

k

at 2560 network function evaluations,

K=10

, seeds 0 and 1 pooled; GSM8K

n{=}600

, MATH

n{=}492

. Left: that estimator at four

k

; bold is the best per column within a regime. Right: the count statistic’s pass@

k

gap over self-consistency, labeled per cell. The pure-diffusion configurations predate the end-of-sequence correction (Appendix D.9 ), which moves the baseline about two points and neither the crossover nor the ordering for

k\geq 2

	GSM8K				MATH				TQA
Method	$k{=}1$	3	5	10	$k{=}1$	3	5	10	$k{=}1$
Semi-autoregressive blocks of 32
Self-consistency	70.04	82.45	86.03	89.70	30.02	44.14	49.78	56.81	54.44
TAPS	66.77	83.28	87.53	91.00	29.26	44.03	49.84	57.32	53.53
ODD ( $\alpha{=}64$ )	64.94	82.60	87.21	91.75	24.19	40.12	46.45	54.27	52.35
ODD ( $\alpha{=}256$ )	66.57	83.32	87.73	92.17	22.77	39.22	45.97	54.37	49.85

Appendix

Table 17: Pass@

k

for every method of Table 2 , same campaigns and settings; pass@1 and pass@10 are Table 2 ’s per-sample and coverage columns. TruthfulQA reports

k{=}1

only, chance alone reaching 94.37 at

k{=}10

; TAPS (

\dagger

) runs the uncorrected protocol. Bold: best per column within a regime.

Sampler	Mask	Plurality ( $\uparrow$ )	Coverage ( $\uparrow$ )	Per-sample ( $\uparrow$ )	Disagr.	$\Delta$ plur.
Self-consistency	off	57.02	67.00	54.47	20.40	—
SR, count	off	71.53	87.75	53.85	49.73	$+14.50$
SR, count	on	70.44	88.00	53.07	50.45	$+13.42$
SR, coll.-wt.	off	76.27	90.00	59.81	45.93	$+19.25$

Appendix

Table 18: The regime re-decoded with the released correction applied to every configuration. GSM8K under pure diffusion,

K=10

, 128 steps, temperature 0.6,

n=600

, two seeds. The third row additionally excludes the end-of-sequence token from the count. The rows with the mask off are the same campaign as Table 2 ’s pure-diffusion block; the first row is the baseline. The last column is the difference in plurality from that baseline.

Explore similar work

Jun 15, 2026cs.CL

LESS Is More: Mutual-Stability Sampling for Diffusion Language Models

Diffusion large language models (dLLMs) offer a promising alternative to autoregressive decoding by iteratively refining masked sequences, enabling parallel token updates and bidirectional conditioning. Their practical efficiency, however, is limited by sampling procedures that execute a fixed number of reverse denoising steps selected before decoding, spending computation on already-stable positions and sometimes committing unstable ones too early. We present \textsc{LESS}, a training-free, model-agnostic adaptive sampler that treats token commitment as an online stopping problem. \textsc{LESS} implements mutual-stability sampling through a joint stability rule that makes a masked position eligible for unmasking only when its top-1 prediction has high confidence, its top-1 token persists across recent reverse steps, and its predictive distribution is stable under top-

K

inter-step Jensen--Shannon divergence. We evaluate \textsc{LESS} on Dream-7B, LLaDA-8B, and LLaDA-1.5-8B, covering full-sequence diffusion and semi-autoregressive blockwise sampling regimes, across seven benchmarks spanning general knowledge, math, and code. \textsc{LESS} improves average accuracy over strong training-free adaptive samplers while using

72.1\%

fewer reverse steps than fixed-budget decoding. Since each reverse step requires a Transformer forward pass, these step-count reductions translate into fewer forward evaluations, lower measured wall-clock latency, and lower estimated inference compute.

Amr Mohamed, Guokan Shang, Michalis Vazirgiannis

1MBZUAI · 2Ecole Polytechnique

Diffusion Language Models Diffusion Models Accelerated Sampling Adaptive Sampling

Ask PDF

May 31, 2026cs.CL

Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models

Masked diffusion language models (MDLMs) re-predict every position at each denoising step, but standard samplers commit tokens once revealed, leaving this revision capability unused. Existing approaches either add heuristic or learned mechanisms to revise committed tokens, or remask them back to [MASK] before re-predicting; a principled sampler that directly revises visible tokens without auxiliary modules remains underexplored. We introduce D3IM, a parameter-free sampler derived as a corrector-style reverse update that permits direct visible-to-visible revision without additional modules or auxiliary passes. D3IM also reveals a model-side obstacle we term preservation bias: the model tends to reproduce its own wrong committed tokens rather than correct them. We address this with SCOPE (Self-Conditioned On Prediction Errors), a lightweight post-training procedure that simulates D3IM's sampling process. On LLaDA-8B at 64 denoising steps, SCOPE+D3IM improves over the original LLaDA-8B with standard unmasking by +13.0 on GSM8K (68.3%), +4.8 on MATH-500 (23.6%), +15.3 on HumanEval (29.3%), and +10.4 on MBPP (30.8%), with gains that increase as more denoising steps are used on math and HumanEval.

Longxuan Yu, Shaorong Zhang, Yu Fu +3

University of California, Riverside · Microsoft

Masked Diffusion Language Models Diffusion Language Models Long Visual-Token Sequences Markovian Sampling

Ask PDF

Jun 9, 2026cs.CL

Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models

Masked diffusion language models can reduce inference steps by revealing multiple tokens per denoising iteration, but this parallelism is fragile: positions that are individually confident may be unsafe to commit together when their predictions are coupled. Existing training-free samplers such as Top-

k

, Fast-dLLM, and EB-Sampler mainly control how many tokens to reveal, while often ranking candidates by token-wise scores that ignore interactions within the selected set. We propose ADAS, a training-free reranking rule that leaves the base sampler's stopping rule unchanged and greedily discounts each token-wise confidence score according to its attention to already selected positions, weighted by their prediction uncertainty. Across LLaDA-8B-Base and Dream-7B-Base on the reasoning benchmarks GSM8K and MATH500 and the code benchmarks HumanEval and MBPP, plugging ADAS into all three samplers improves low-NFE performance at matched denoiser evaluations by

9.11

and

10.46

percentage points on average, respectively, with

3.1\%

per-forward runtime overhead. Code is available at https://github.com/yusufsahin99/ADAS.

Yusuf Sahin, Ahmed Rockey Saikia, Volkan Cevher +1

University of Bern, Bern, Switzerland · EPFL, Lausanne, Switzerland

Masked Diffusion Language Models Diffusion Language Models Adaptive Sampling

Ask PDF