Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models

# Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models
Source: [https://arxiv.org/html/2606.10829](https://arxiv.org/html/2606.10829)
Yusuf Sahin University of Bern Bern, Switzerland &Ahmed Rockey Saikia EPFL Lausanne, Switzerland &Volkan Cevher EPFL Lausanne, Switzerland &Paolo Favaro University of Bern Bern, Switzerland

###### Abstract

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\-kk, 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 for parallel masked diffusion decoding\. ADAS leaves the base sampler’s stopping rule unchanged and modifies only subset construction: it greedily discounts a candidate when it attends strongly to already selected positions whose predictions remain uncertain\. Unlike graph\-constrained methods that turn attention into hard compatibility constraints, ADAS keeps attention continuous and uses it as a soft marginal penalty\. Across LLaDA\-8B\-Base and Dream\-7B\-Base on GSM8K, MATH500, HumanEval, and MBPP, plugging ADAS into Top\-kk, Fast\-dLLM, and EB\-Sampler improves low\-NFE performance at matched denoiser evaluations by9\.119\.11and10\.4610\.46percentage points on average, respectively, with3\.1%3\.1\\%per\-forward runtime overhead\. These results show that soft attention\-discounted reranking is a simple and modular way to improve quality in highly parallel decoding for masked diffusion language models\.

## 1Introduction

Masked diffusion language models \(MDLMs\)\(Sahooet al\.,[2024](https://arxiv.org/html/2606.10829#bib.bib7)\)promise a different speed–quality tradeoff from causal autoregressive models: each denoising step can reveal a subset of positions rather than a single next token\. This parallelism is especially appealing for reasoning and code\-generation tasks, where long outputs make inference latency costly\. The central decoding question is therefore not only*what*token to place at each masked position, but also*which positions*can be safely committed in the same step\.

This second question is where aggressive parallel decoding becomes fragile\. Many practical samplers rank candidate positions by token\-wise quantities such as confidence, margin, or entropy, and then use a stopping rule or budget to decide how many positions to reveal\. Such rules are effective when only a few tokens are unmasked at once, but they implicitly treat high\-scoring candidates as compatible\. In our controlled diagnostic, forcing the model to predict two coupled literals together drops accuracy from71%71\\%to30%30\\%, showing that joint commitment can be much harder than the individual predictions suggest\.

Recent work has made this failure mode increasingly visible\. Some methods treat diffusion decoding as planning or search\(Penget al\.,[2025](https://arxiv.org/html/2606.10829#bib.bib6); Leeet al\.,[2025](https://arxiv.org/html/2606.10829#bib.bib12)\); others control the amount of parallelism through schedules\(Luxembourget al\.,[2025](https://arxiv.org/html/2606.10829#bib.bib4); Israelet al\.,[2025](https://arxiv.org/html/2606.10829#bib.bib5)\)or learned unmasking policies\(Jazbecet al\.,[2025](https://arxiv.org/html/2606.10829#bib.bib14); Honget al\.,[2025](https://arxiv.org/html/2606.10829#bib.bib13); Baoet al\.,[2025](https://arxiv.org/html/2606.10829#bib.bib11)\)\. The theoretical lens ofBen\-Hamuet al\.\([2025](https://arxiv.org/html/2606.10829#bib.bib9)\)also separates*model error*from*joint dependence error*, clarifying that good marginal predictions are not sufficient when several positions are committed simultaneously\.

We take this perspective one step further\. Widely used decoders often differ in the*constraint*imposed on the selected subset, but still rely on a token\-wise objective for constructing that subset\. Top\-kkfixes the cardinality, Fast\-dLLM uses a confidence\-threshold condition, and EB\-Sampler enforces an entropy budget\. These mechanisms decide*when to stop adding tokens*; they do not, by themselves, make the marginal value of a candidate depend on which other positions have already been selected\.

Our proposal is to make subset construction dependency\-aware while leaving the stopping rule intact\. We introduce ADAS, a training\-free Attention\-Discounted Adaptive Sampler that greedily builds the set of positions to unmask\. Starting from standard token confidence scores, ADAS discounts a remaining candidate when it attends strongly to already selected, uncertain positions\. Equivalently, ADAS performs greedy construction under a non\-additive subset utility: confidence supplies an individual marginal\-value proxy, while attention supplies a pairwise\-compatibility proxy\.

This design deliberately separates ranking from budget control\. ADAS does not train a planner, add a verifier, perform lookahead search, or redesign the sampler’s stopping criterion\. Instead, it upgrades the greedy ranking step inside existing decoders, so the same attention\-discounted update can be used with fixed\-cardinality, confidence\-threshold, or entropy\-budget rules\. This places ADAS in a complementary position to learned unmasking policies\(Jazbecet al\.,[2025](https://arxiv.org/html/2606.10829#bib.bib14); Honget al\.,[2025](https://arxiv.org/html/2606.10829#bib.bib13); Asanoet al\.,[2026](https://arxiv.org/html/2606.10829#bib.bib3)\), lookahead methods\(Leeet al\.,[2025](https://arxiv.org/html/2606.10829#bib.bib12)\), and schedule\- or verifier\-based approaches\(Luxembourget al\.,[2025](https://arxiv.org/html/2606.10829#bib.bib4); Israelet al\.,[2025](https://arxiv.org/html/2606.10829#bib.bib5)\)\.

Empirically, the benefit appears exactly where the motivation predicts: the low\-NFE, highly parallel regime in which token\-wise ranking is most brittle\. Across mathematical reasoning and code\-generation benchmarks, adding the ADAS score update consistently improves Top\-kk, EB\-Sampler, and Fast\-dLLM at matched denoiser\-evaluation budgets\. Averaged across stopping criteria, ADAS improves matched\-NFE performance by9\.119\.11points on LLaDA\-8B\-Base and10\.4610\.46points on Dream\-7B\-Base, with positive gains on8080of9090matched operating points\.

Our contributions are as follows\.

- •We diagnose joint dependence in parallel MDLM decoding: forcing coupled positions to be predicted together drops controlled\-task accuracy from71%71\\%to30%30\\%, and masked\-token self\-attention separates dependent from non\-dependent pairs\.
- •We introduce ADAS, a training\-free attention\-discounted greedy selector that decouples*ranking*from*stopping*, plugs into Top\-kk, Fast\-dLLM, and EB\-Sampler, and adds only3\.1%3\.1\\%runtime overhead per model forward\.
- •Across two MDLMs and four reasoning/code benchmarks, ADAS improves matched\-NFE performance by9\.119\.11and10\.4610\.46points on average\.

## 2Related work

Recent work on masked diffusion decoding has shown that generation quality depends strongly on the order and grouping of positions unmasked at each step\. Several approaches therefore treat decoding as an explicit planning or policy\-learning problem\.Penget al\.\([2025](https://arxiv.org/html/2606.10829#bib.bib6)\)formulate sampling as planner\-guided decoding,Asanoet al\.\([2026](https://arxiv.org/html/2606.10829#bib.bib3)\)separate where\-to\-unmask from what\-to\-unmask and learn a supervised planner,Leeet al\.\([2025](https://arxiv.org/html/2606.10829#bib.bib12)\)use lookahead search over candidate decoding trajectories, andJazbecet al\.\([2025](https://arxiv.org/html/2606.10829#bib.bib14)\); Honget al\.\([2025](https://arxiv.org/html/2606.10829#bib.bib13)\); Baoet al\.\([2025](https://arxiv.org/html/2606.10829#bib.bib11)\)learn unmasking policies or filters that improve over hand\-designed heuristics\. These methods show that token\-wise confidence order is often not enough, especially outside weakly parallel regimes\.

A complementary line of work improves the speed–quality tradeoff of diffusion decoding through better inference schedules or auxiliary signals\.Luxembourget al\.\([2025](https://arxiv.org/html/2606.10829#bib.bib4)\)propose a structural scheduler that reduces harmful interactions through dilated grouping, whileIsraelet al\.\([2025](https://arxiv.org/html/2606.10829#bib.bib5)\)use an auxiliary autoregressive model to adapt the amount of parallelism during decoding\.Ben\-Hamuet al\.\([2025](https://arxiv.org/html/2606.10829#bib.bib9)\)provide a useful theoretical framing by decomposing decoding error into model error and joint dependence error, motivating adaptive multi\-token unmasking rules such as EB\-Sampler\.Kimet al\.\([2026b](https://arxiv.org/html/2606.10829#bib.bib10)\)propose KLASS, a training\-free sampler that combines token confidence with temporal KL divergence between consecutive denoising distributions to select stable tokens for unmasking\.

The closest related method is DAPD\(Kimet al\.,[2026a](https://arxiv.org/html/2606.10829#bib.bib2)\), which also uses self\-attention to estimate dependencies among masked positions\. DAPD builds an attention\-induced dependency graph by symmetrizing pairwise attention scores, thresholding them into binary edges, and selecting an independent set of masked positions to decode in parallel\. Its decoding rule is therefore graph\-constrained: an attention edge acts as a hard conflict, and the selected parallel batch must satisfy an approximate independence condition\.

ADAS uses the same broad source of information—self\-attention—but instantiates a different decoding principle\. It does not construct a binary graph, threshold attention, or impose independent\-set constraints\. Instead, ADAS keeps attention continuous and uses it as a marginal penalty inside greedy subset construction\. Thus, ADAS does not declare two positions incompatible\. A strongly coupled candidate can still be selected when its confidence gain outweighs the estimated commitment risk\. This soft formulation is important because attention is an imperfect proxy for harmful dependence: not every high\-attention pair should be forbidden, and not every low\-attention pair is guaranteed independent\.

The two methods also differ in their integration point\. DAPD is a standalone dependency\-aware decoder that defines its own parallel batches through graph coloring\. ADAS is a sampler\-agnostic reranking module: it leaves the stopping rule, budget, and admissibility condition of Top\-kk, Fast\-dLLM, or EB\-Sampler unchanged, and only changes the order in which candidates are proposed\. In short, DAPD uses attention to construct a graph of forbidden simultaneous updates; ADAS uses attention to compute a soft marginal discount inside existing samplers\.

## 3Preliminaries

### 3\.1Notation

Let𝒱=\{1,…,K\}⊎\{m\}\\mathcal\{V\}=\\\{1,\\dots,K\\\}\\uplus\\\{m\\\}denote a finite vocabulary augmented with a special mask tokenmm\. A sequence is denoted byx∈𝒱dx\\in\\mathcal\{V\}^\{d\}, whereddis the sequence length\. A positioniiis*masked*ifxi=mx\_\{i\}=mand*unmasked*otherwise\. For a partially masked sequence, letℳ⊆\{1,…,d\}\\mathcal\{M\}\\subseteq\\\{1,\\dots,d\\\}be the set of masked positions, and letℳ¯\\bar\{\\mathcal\{M\}\}denote its complement\. We writexℳ¯x^\{\\bar\{\\mathcal\{M\}\}\}for the observed tokens\. We useq\(⋅∣xℳ¯\)q\(\\cdot\\mid x^\{\\bar\{\\mathcal\{M\}\}\}\)to denote the true conditional distribution over masked positions given the currently revealed context andpθ\(⋅∣xℳ¯\)p\_\{\\theta\}\(\\cdot\\mid x^\{\\bar\{\\mathcal\{M\}\}\}\)to denote the model’s conditional distribution\.

### 3\.2Masked Diffusion Language Models

Masked diffusion language models\(Sahooet al\.,[2024](https://arxiv.org/html/2606.10829#bib.bib7)\)generate text through iterative denoising\. Starting from a highly masked sequence, the model repeatedly predicts token distributions at masked positions and progressively unmasks a subset of them until no masks remain\. A decoding step therefore involves two coupled decisions:*what*token to place at a selected position, and*where*to unmask next\.

At a decoding step, an MDLM with parametersθ\\thetapredicts, for each masked positioni∈ℳi\\in\\mathcal\{M\}, a conditional distributionpθ​\(xi∣xℳ¯\)p\_\{\\theta\}\(x^\{i\}\\mid x^\{\\bar\{\\mathcal\{M\}\}\}\)\. Many existing decoding strategies then assign each masked position a token\-wise score derived from this distribution, and construct a subsetS⊆ℳS\\subseteq\\mathcal\{M\}of positions to unmask according to a sampler\-specific stopping rule\. A common choice is the confidence score

ci=maxxi⁡pθ​\(xi∣xℳ¯\),c\_\{i\}=\\max\_\{x^\{i\}\}p\_\{\\theta\}\(x^\{i\}\\mid x^\{\\bar\{\\mathcal\{M\}\}\}\),\(1\)
namely the probability assigned to the most likely token value at positionii\. Different samplers mainly differ in how the subsetSSis constrained, e\.g\., by fixing its size, thresholding confidence, or enforcing an entropy budget\.

### 3\.3Model error and joint dependence error

A central difficulty in parallel decoding is that the model provides token\-wise conditional marginals for masked positions, while decoding multiple positions in the same step requires reasoning about their*joint*behavior\. As emphasized byBen\-Hamuet al\.\([2025](https://arxiv.org/html/2606.10829#bib.bib9)\), the error incurred by unmasking a subsetS⊆ℳS\\subseteq\\mathcal\{M\}can be decomposed into a*model error*term and a*joint dependence error*term:

∑i∈SDKL​\(q​\(xi∣xℳ¯\),pθ​\(xi∣xℳ¯\)\)⏟model error\+DKL​\(q​\(xS∣xℳ¯\),∏i∈Sq​\(xi∣xℳ¯\)\)⏟joint dependence error\\underbrace\{\\sum\_\{i\\in S\}D\_\{\\mathrm\{KL\}\}\\\!\\bigl\(q\(x^\{i\}\\mid x^\{\\bar\{\\mathcal\{M\}\}\}\),\\,p\_\{\\theta\}\(x^\{i\}\\mid x^\{\\bar\{\\mathcal\{M\}\}\}\)\\bigr\)\}\_\{\\text\{model error\}\}\\;\+\\;\\underbrace\{\\vphantom\{\\sum\_\{i\\in S\}D\_\{\\mathrm\{KL\}\}\\\!\\bigl\(q\(x^\{i\}\\mid x^\{\\bar\{\\mathcal\{M\}\}\}\),\\,p\_\{\\theta\}\(x^\{i\}\\mid x^\{\\bar\{\\mathcal\{M\}\}\}\)\\bigr\)\}D\_\{\\mathrm\{KL\}\}\\\!\\bigl\(q\(x^\{S\}\\mid x^\{\\bar\{\\mathcal\{M\}\}\}\),\\,\\textstyle\\prod\_\{i\\in S\}q\(x^\{i\}\\mid x^\{\\bar\{\\mathcal\{M\}\}\}\)\\bigr\)\}\_\{\\text\{joint dependence error\}\}The model error captures inaccuracies in the model’s token\-wise predictions, while the joint dependence error measures the discrepancy introduced by treating the selected positions as conditionally independent\. The latter is the multi\-information of the true conditional distributionqqoverSS: it is model\-independent and characterizes the intrinsic coupling among positions selected for simultaneous commitment\.

Sinceqqis unobserved at decoding time, ADAS uses the denoiser self\-attentionAi​jA\_\{ij\}as a model\-side proxy for coupling strength\. When this proxy is informative, discounting candidates that attend strongly to already selected, uncertain positions should reduce the expected dependence cost of the next parallel commitment\. Section[4](https://arxiv.org/html/2606.10829#S4)tests this proxy assumption\.

## 4Controlled dependency diagnostics

Before introducing the Attention\-Discounted Adaptive Sampler \(ADAS\), we test the two assumptions behind attention\-guided parallel unmasking: jointly unmasking dependent positions is harder than jointly unmasking independent positions, and self\-attention provides a useful proxy for such dependencies\.

We use synthetic arithmetic predicates with known dependency structure\. The full predicate set is

E=\{A\+B\+C=D,A⋅B⋅C=D,min⁡\{A,B,C\}=D,max⁡\{A,B,C\}=D\}\.E=\\left\\\{A\+B\+C=D,\\;A\\cdot B\\cdot C=D,\\;\\min\\\{A,B,C\\\}=D,\\;\\max\\\{A,B,C\\\}=D\\right\\\}\.
![Refer to caption](https://arxiv.org/html/2606.10829v1/x1.png)Figure 1:Distribution of pairwise attention values for dependent and non\-dependent masked\-token pairs in the synthetic dependency diagnostic, computed using LLaDA\-8B\-Base\. Dependent pairs receive consistently higher attention than non\-dependent pairs, suggesting that attention tracks the underlying dependency structure\.For the first diagnostic, we use only the sum and product predicates, since these directly test whether simultaneously predicting two coupled numeric literals is harder than predicting two literals from separate equations\. Each example contains three predicates, and exactly two numeric literals are masked\. In the dependent\-pair condition, both masked literals occur in the same predicate; in the independent\-pair condition, they occur in different predicates\. We then force Dream\-7B\-Base to predict both masked literals in a single forward pass, without confidence ordering, adaptive stopping, or attention\-discounting\. We use constraint\-satisfaction scoring, so a completion is counted as correct whenever the resulting equation is valid, even if the predicted literals differ from the originally sampled values\. Dependent pairs are substantially harder to recover than independent pairs: accuracy drops from71%71\\%to30%30\\%\.

For the second diagnostic, we use the full predicate setEE, including the sum, product, minimum, and maximum predicates\. This diagnostic tests whether self\-attention reflects known dependency structure across a broader set of arithmetic relations\. Using the known\-dependency construction described in Appendix[E](https://arxiv.org/html/2606.10829#A5), we compare attention values between masked\-token pairs from the same predicate and pairs from different predicates\. As shown in Figure[1](https://arxiv.org/html/2606.10829#S4.F1), dependent pairs attend more to each other than non\-dependent pairs\. Quantitatively, mean attention is higher for dependent pairs on both LLaDA\-8B\-Base \(0\.0072610\.007261vs\.0\.0019190\.001919\) and Dream\-7B\-Base \(0\.0095170\.009517vs\.0\.0048420\.004842\)\. These diagnostics support the central premise of ADAS without treating attention as a perfect dependence measure: highly coupled positions are risky to unmask together, and attention provides a practical training\-free proxy for this coupling\.

## 5Method

### 5\.1Problem statement

Consider a masked diffusion decoding step with current set of masked positionsℳ\\mathcal\{M\}\. The goal is to select a subsetS⊆ℳS\\subseteq\\mathcal\{M\}of positions to unmask in parallel\. In most existing samplers, this selection problem is implicitly decomposed into two parts: a*token\-wise utility*used to rank candidate positions, and a*stopping criterion*that determines when enough positions have been selected\.

A common choice is to use the confidence scorescic\_\{i\}introduced in Eq\. \([1](https://arxiv.org/html/2606.10829#S3.E1)\) to constructSSso as to maximize an additive objective of the form

U​\(S\)=∑i∈Sci,\\displaystyle U\(S\)=\\sum\_\{i\\in S\}c\_\{i\},\(2\)subject to a sampler\-specific constraint\. Different decoding rules mainly differ in this constraint\. For example, Top\-kkenforces a fixed cardinality constraint\|S\|=k\|S\|=k, Fast\-dLLM uses a confidence\-based stopping condition, and EB\-Sampler uses an entropy\-based bound\.

Our starting point is that this token\-wise objective ignores dependencies among positions selected in the same step\. A candidate’s reliability should depend not only on its own token\-wise score, but also on the uncertainty of the positions with which it is jointly committed and the strength of their interactions\. We therefore retain the general subset\-selection view above, but replace the purely additive token\-wise objective with a dependency\-aware utility that accounts for pairwise interactions\.

### 5\.2Attention\-Discounted Adaptive Sampler \(ADAS\)

ADAS replaces token\-wise candidate ranking with dependency\-aware greedy subset construction, while leaving the sampler’s stopping rule unchanged\. At each decoding step, letℳ\\mathcal\{M\}be the set of masked positions and letcic\_\{i\}denote the usual confidence score for positioni∈ℳi\\in\\mathcal\{M\}\. LetA∈\[0,1\]\|ℳ\|×\|ℳ\|A\\in\[0,1\]^\{\|\\mathcal\{M\}\|\\times\|\\mathcal\{M\}\|\}be an attention\-based interaction matrix over masked positions, whereAi​jA\_\{ij\}measures attention from positioniito positionjj\. In our implementation,AAis obtained from the denoiser’s final self\-attention layer by averaging attention scores across heads, following the ablation in Appendix[C](https://arxiv.org/html/2606.10829#A3.SS0.SSS0.Px3)\.

Given an already selected setSS, ADAS scores a remaining candidatei∈ℳ∖Si\\in\\mathcal\{M\}\\setminus Sby the marginal utility

u​\(i∣S,A\)=ci−α​∑s∈SAi​s​\(1−cs\),\\displaystyle u\(i\\mid S,A\)=c\_\{i\}\-\\alpha\\sum\_\{s\\in S\}A\_\{is\}\(1\-c\_\{s\}\),\(3\)whereα≥0\\alpha\\geq 0controls the strength of the attention discount\. The first term favors positions with high token\-wise confidence\. The second term discounts candidates that attend strongly to selected positions whose predictions are still uncertain\. Thus, a candidate is penalized not simply for being coupled to another selected position, but for being coupled to one whose commitment is less reliable\.

Equivalently, ADAS greedily constructs a set under the dependency\-aware utility

U​\(S,A\)=∑i∈Sci−α​∑i,j∈SAi​j​\(1−cj\)\.\\displaystyle U\(S,A\)=\\sum\_\{i\\in S\}c\_\{i\}\-\\alpha\\sum\_\{i,j\\in S\}A\_\{ij\}\(1\-c\_\{j\}\)\.\(4\)This utility is not submodular in general, and we do not claim an approximation guarantee\. Instead, we use it as a lightweight ranking rule motivated by the local surrogate view that attention\-weighted uncertainty approximates the incremental risk of jointly committing dependent positions\.

At each iteration, ADAS selects the remaining candidate with largest current marginal utility and checks whether adding it satisfies the base sampler’s admissibility rule\. If the enlarged set is admissible, the candidate is added; otherwise, construction stops and the selected positions are unmasked in parallel\.

Algorithm 1Attention\-Discounted Adaptive Sampler \(ADAS\)ℳ\\mathcal\{M\}, scores

\{ci\}\\\{c\_\{i\}\\\}, attention

AA, weight

α\\alpha, admissibility rule

Adm\\mathrm\{Adm\}
S←∅S\\leftarrow\\emptyset

while

ℳ∖S≠∅\\mathcal\{M\}\\setminus S\\neq\\emptysetdo

i⋆←arg⁡maxi∈ℳ∖S⁡\[ci−α​∑s∈SAi​s​\(1−cs\)\]i^\{\\star\}\\leftarrow\\arg\\max\_\{i\\in\\mathcal\{M\}\\setminus S\}\\left\[c\_\{i\}\-\\alpha\\sum\_\{s\\in S\}A\_\{is\}\(1\-c\_\{s\}\)\\right\]

ifnot

Adm​\(S∪\{i⋆\}\)\\mathrm\{Adm\}\(S\\cup\\\{i^\{\\star\}\\\}\)then

break

endif

S←S∪\{i⋆\}S\\leftarrow S\\cup\\\{i^\{\\star\}\\\}

endwhile

return

SS

In our experiments, ADAS is combined with three admissibility rules:

Top\-​k:\\displaystyle\\text\{Top\-\}k:\\quad\|S∪\{i⋆\}\|≤k,\\displaystyle\|S\\cup\\\{i^\{\\star\}\\\}\|\\leq k,\(5\)Fast\-dLLM:\\displaystyle\\text\{Fast\-dLLM\}:\\quad\|S∪\{i⋆\}\|​\(1−mins∈S∪\{i⋆\}⁡cs\)<f,\\displaystyle\|S\\cup\\\{i^\{\\star\}\\\}\|\\Bigl\(1\-\\min\_\{s\\in S\\cup\\\{i^\{\\star\}\\\}\}c\_\{s\}\\Bigr\)<f,\(6\)EB\-Sampler:\\displaystyle\\text\{EB\-Sampler\}:\\quad∑s∈S∪\{i⋆\}H​\(pθ​\(xs∣xℳ¯\)\)−maxs∈S∪\{i⋆\}⁡H​\(pθ​\(xs∣xℳ¯\)\)<γ\.\\displaystyle\\sum\_\{s\\in S\\cup\\\{i^\{\\star\}\\\}\}H\\\!\\bigl\(p\_\{\\theta\}\(x^\{s\}\\mid x^\{\\bar\{\\mathcal\{M\}\}\}\)\\bigr\)\-\\max\_\{s\\in S\\cup\\\{i^\{\\star\}\\\}\}H\\\!\\bigl\(p\_\{\\theta\}\(x^\{s\}\\mid x^\{\\bar\{\\mathcal\{M\}\}\}\)\\bigr\)<\\gamma\.\(7\)Herekk,ff, andγ\\gammaare the fixed\-cardinality, confidence, and entropy\-budget thresholds, respectively\. These rules determine when subset construction stops; ADAS only changes how candidates are ranked before that stopping decision\.

The attention discount can be updated incrementally\. After adding a selected positionss, each remaining candidate’s score is decreased byα​Ai​s​\(1−cs\)\\alpha A\_\{is\}\(1\-c\_\{s\}\), givingO​\(\|S\|​\|ℳ\|\)O\(\|S\|\\,\|\\mathcal\{M\}\|\)selection cost for a final subsetSS\. This overhead is small relative to denoiser forward passes; we measure it in Section[6\.3](https://arxiv.org/html/2606.10829#S6.SS3)and ablateα\\alphaand the uncertainty term in Appendix[C](https://arxiv.org/html/2606.10829#A3)\.

## 6Experiments

### 6\.1Experimental setup

We evaluate ADAS on four benchmarks spanning mathematical reasoning and code generation: GSM8K\(Cobbeet al\.,[2021](https://arxiv.org/html/2606.10829#bib.bib17)\), MATH500\(Hendryckset al\.,[2021](https://arxiv.org/html/2606.10829#bib.bib18)\), MBPP\(Austinet al\.,[2021](https://arxiv.org/html/2606.10829#bib.bib16)\), and HumanEval\(Chenet al\.,[2021](https://arxiv.org/html/2606.10829#bib.bib15)\)\. We use two strong open\-weight masked diffusion language models, Dream\-7B\-Base\(Yeet al\.,[2025](https://arxiv.org/html/2606.10829#bib.bib19)\)and LLaDA\-8B\-Base\(Nieet al\.,[2025](https://arxiv.org/html/2606.10829#bib.bib20)\)\.

Our goal is to test whether attention\-discounted subset construction helps precisely where parallel unmasking is most useful but most error\-prone\. We therefore instantiate ADAS with three representative stopping criteria: Top\-kk, Fast\-dLLM, and EB\-Sampler, yielding three variants we denote Top\-kk\+AD, Fast\-dLLM\+AD, and EB\+AD\. Each variant is compared against its corresponding baseline over a range of parallel decoding settings\.

We further evaluate KLASS\(Kimet al\.,[2026b](https://arxiv.org/html/2606.10829#bib.bib10)\)and DAPD\(Kimet al\.,[2026a](https://arxiv.org/html/2606.10829#bib.bib2)\)as additional training\-free baselines\.

### 6\.2Main results

Figures[2](https://arxiv.org/html/2606.10829#S6.F2)and[3](https://arxiv.org/html/2606.10829#S6.F3)summarize the main speed–quality tradeoff\. Moving to the right corresponds to revealing more tokens per denoising step, and therefore to lower NFE\. The central question is whether the sampler can preserve quality in this highly parallel regime, where jointly selected dependent positions are most likely to induce errors\.

In each figure, the first three rows compare EB\-Sampler, Fast\-dLLM, and Top\-kkwith their ADAS variants; the final row compares ADAS\-augmented samplers with KLASS and DAPD\. Together, these plots test two claims: ADAS should improve its own base sampler, and the resulting decoders should remain competitive with recent training\-free samplers\.

![Refer to caption](https://arxiv.org/html/2606.10829v1/x2.png)Figure 2:Effect of attention\-discounted selection on Dream\-7B\-Base\. Rows correspond to entropy\-bounded decoding, Fast\-dLLM, fixed Top\-kk, and comparison against KLASS and DAPD; columns correspond to GSM8K, HumanEval, MATH500, and MBPP\. The same pattern observed for LLaDA\-8B\-Base persists: attention\-discounted selection is most beneficial when many tokens are revealed in parallel, where the original samplers degrade sharply\.![Refer to caption](https://arxiv.org/html/2606.10829v1/x3.png)Figure 3:Effect of attention\-discounted selection on LLaDA\-8B\-Base\. Rows correspond to entropy\-bounded decoding, Fast\-dLLM, fixed Top\-kk, and comparison against KLASS and DAPD; columns correspond to GSM8K, HumanEval, MATH500, and MBPP\. Across stopping rules and datasets, attention\-discounted selection preserves substantially higher accuracy in the highly parallel regime, while matching the original samplers more closely when the number of parallel tokens is small\.Across both models, ADAS consistently improves robustness in the high\-parallelism regime\. For the adaptive stopping rules EB and fast\-dLLM, the baseline curves often deteriorate rapidly as the entropy budget increases, whereas their attention\-discounted variants degrade more gradually\.

Fixed Top\-kkdecoding shows that the benefit is not tied to adaptive stopping criteria\. Sincekkdirectly controls the number of tokens unmasked at each step, this setting isolates the effect of subset construction\. Across both base models, Top\-kk\+AD improves or preserves performance at largerkk, where vanilla Top\-kksuffers the largest degradation\. Thus, reranking candidates by attention\-discounted utility selects safer groups of tokens even when the subset size is fixed in advance\.

The comparison with prior training\-free baselines further supports this interpretation\. In the bottom row of each figure, the ADAS\-augmented decoders generally remain competitive with or above KLASS and DAPD at moderate and large parallel\-token counts\. Both baselines often degrade rapidly as parallelism increases, whereas ADAS tends to degrade more smoothly\. One possible explanation for the KLASS trend is that temporal\-stability signals become less reliable when many tokens are committed at once, since consecutive denoising distributions can shift substantially\. The DAPD trend suggests that hard attention\-derived constraints may be brittle in some high\-parallel regimes\. Overall, these results indicate that soft attention\-discounted reranking is a robust way to reduce harmful joint commitments among dependent positions\.

To summarize the low\-NFE regime quantitatively, we compute average gains at matched NFE\. For each baseline operating point where the decoder unmasks at least four tokens on average per step, we linearly interpolate the corresponding ADAS curve at the same NFE and measure the score difference\. Table[1](https://arxiv.org/html/2606.10829#S6.T1)reports these gains for both LLaDA\-8B\-Base and Dream\-7B\-Base across the three original stopping rules, with each entry shown as LLaDA gain / Dream gain\.

Table 1:Matched\-NFE absolute gains of attention\-discounted selection in the low\-NFE regime\. Each entry reports LLaDA\-8B\-Base gain / Dream\-7B\-Base gain in percentage points\. We average over baseline points with at least four parallel tokens per step and compare each baseline to the interpolated ADAS score at the same NFE\.DatasetEBFast\-dLLMTop\-kkAvg\. over methodsGSM8K\+9\.22 / \+14\.31\+6\.92 / \+6\.36\+11\.90 / \+5\.88\+9\.35 / \+8\.85MATH500\+6\.49 / \+4\.84\+5\.70 / \+2\.15\+10\.73 / \+6\.93\+7\.64 / \+4\.64HumanEval\+10\.45 / \+15\.63\+6\.90 / \+12\.03\+9\.55 / \+10\.77\+8\.97 / \+12\.81MBPP\+10\.37 / \+15\.01\+7\.39 / \+12\.04\+13\.67 / \+19\.60\+10\.48 / \+15\.55Avg\. over datasets\+9\.13 / \+12\.45\+6\.73 / \+8\.14\+11\.46 / \+10\.80\+9\.11 / \+10\.46The matched\-NFE summary confirms the visual trend across both base models\. Averaged over all datasets and stopping criteria, the gains are\+9\.11\+9\.11points for LLaDA\-8B\-Base and\+10\.46\+10\.46points for Dream\-7B\-Base\.

The improvement is consistent across entropy\-based, confidence\-based, and fixed\-budget stopping rules, indicating that the benefit comes from attention\-discounted subset construction rather than from a particular stopping criterion\. Gains are especially strong on code\-generation benchmarks, where simultaneously finalizing coupled positions can more easily lead to inconsistent outputs\. MATH500 shows smaller but still positive matched\-NFE gains, suggesting that some mathematical\-reasoning errors may be governed by longer sequential chains rather than pairwise incompatibilities within a single parallel update\.

To check that the effect is not driven by a few favorable thresholds, we also bootstrap matched gains across operating points\. Across9090points, ADAS has mean gain\+9\.27\+9\.27with 95% CI\[\+7\.74,\+10\.84\]\[\+7\.74,\+10\.84\], improving in8080cases and regressing in1010\. Appendix[D](https://arxiv.org/html/2606.10829#A4)gives the full breakdown\.

### 6\.3Analysis of runtime overhead

ADAS reads the last\-layer attention matrix and greedily updates token scores when constructing the subset to unmask\. In our implementation, accessing these weights requires disabling FlashAttention only in the final denoiser layer\. To isolate this cost from schedule differences, we compare each ADAS variant against its vanilla decoder at matched operating points, using the same stopping rule and similar NFE\.

Across all full\-dataset runs used in our paper plots, we measure wall\-clock generation time normalized by generated samples and denoiser forwards\. Adding ADAS costs3\.1%3\.1\\%per model forward on average, including attention extraction and greedy reranking\.

## 7Discussion and limitations

We proposed ADAS, a training\-free attention\-guided subset selection rule for parallel decoding in masked diffusion language models\. The key idea is to move beyond purely token\-wise scoring and incorporate pairwise compatibility between candidate positions using the model’s own self\-attention\. Across multiple benchmarks, models, and stopping criteria, we found that this simple attention\-discounted reranking strategy consistently improves decoding quality in highly parallel regimes, where token\-wise heuristics are most vulnerable to joint dependence error\.

Several limitations remain\. Self\-attention is only a proxy for stochastic dependence: strong attention does not always imply harmful joint updates, nor does weak attention guarantee compatibility\. Our experiments cover two open\-weight base MDLMs and four reasoning/code benchmarks, so broader dependency signals, samplers, model scales, and task domains remain important future tests\. Our robustness intervals summarize variation across matched operating points rather than repeated seeds or per\-example uncertainty\.

Natural next steps include richer attention\-derived dependency signals, head\-specific or multi\-layer aggregation, and combinations with learned planners, lookahead methods, or auxiliary verifiers\. More broadly, our results suggest that improving parallel decoding requires not only better stopping rules, but also better objectives for composing the subset of positions revealed at each step\.

The broader impacts of ADAS are mostly inherited from the underlying language models\. ADAS may improve the inference efficiency and output quality of existing masked diffusion language models, which can reduce decoding cost and make these models more practical to use\. At the same time, faster or higher\-quality generation could indirectly increase misuse risks already associated with language models, such as automated spam, low\-quality content generation, or other harmful downstream applications\. ADAS does not introduce a new model, dataset, or deployed system, so these risks are inherited from the underlying models rather than specific to a newly released artifact\.

## References

- Where\-to\-unmask: ground\-truth\-guided unmasking order learning for masked diffusion language models\.External Links:2602\.09501,[Link](https://arxiv.org/abs/2602.09501)Cited by:[§1](https://arxiv.org/html/2606.10829#S1.p6.1),[§2](https://arxiv.org/html/2606.10829#S2.p1.1)\.
- J\. Austin, A\. Odena, M\. Nye, M\. Bosma, H\. Michalewski, D\. Dohan, E\. Jiang, C\. Cai, M\. Terry, Q\. Le, and C\. Sutton \(2021\)Program synthesis with large language models\.External Links:2108\.07732,[Link](https://arxiv.org/abs/2108.07732)Cited by:[§6\.1](https://arxiv.org/html/2606.10829#S6.SS1.p1.1)\.
- W\. Bao, Z\. Chen, D\. Xu, and Y\. Shang \(2025\)Learning to parallel: accelerating diffusion large language models via learnable parallel decoding\.External Links:2509\.25188,[Link](https://arxiv.org/abs/2509.25188)Cited by:[§1](https://arxiv.org/html/2606.10829#S1.p3.1),[§2](https://arxiv.org/html/2606.10829#S2.p1.1)\.
- H\. Ben\-Hamu, I\. Gat, D\. Severo, N\. Nolte, and B\. Karrer \(2025\)Accelerated sampling from masked diffusion models via entropy bounded unmasking\.External Links:2505\.24857,[Link](https://arxiv.org/abs/2505.24857)Cited by:[§1](https://arxiv.org/html/2606.10829#S1.p3.1),[§2](https://arxiv.org/html/2606.10829#S2.p2.1),[§3\.3](https://arxiv.org/html/2606.10829#S3.SS3.p1.1)\.
- M\. Chen, J\. Tworek, H\. Jun, Q\. Yuan, H\. P\. de Oliveira Pinto, J\. Kaplan, H\. Edwards, Y\. Burda, N\. Joseph, G\. Brockman, A\. Ray, R\. Puri, G\. Krueger, M\. Petrov, H\. Khlaaf, G\. Sastry, P\. Mishkin, B\. Chan, S\. Gray, N\. Ryder, M\. Pavlov, A\. Power, L\. Kaiser, M\. Bavarian, C\. Winter, P\. Tillet, F\. P\. Such, D\. Cummings, M\. Plappert, F\. Chantzis, E\. Barnes, A\. Herbert\-Voss, W\. H\. Guss, A\. Nichol, A\. Paino, N\. Tezak, J\. Tang, I\. Babuschkin, S\. Balaji, S\. Jain, W\. Saunders, C\. Hesse, A\. N\. Carr, J\. Leike, J\. Achiam, V\. Misra, E\. Morikawa, A\. Radford, M\. Knight, M\. Brundage, M\. Murati, K\. Mayer, P\. Welinder, B\. McGrew, D\. Amodei, S\. McCandlish, I\. Sutskever, and W\. Zaremba \(2021\)Evaluating large language models trained on code\.External Links:2107\.03374,[Link](https://arxiv.org/abs/2107.03374)Cited by:[§6\.1](https://arxiv.org/html/2606.10829#S6.SS1.p1.1)\.
- K\. Cobbe, V\. Kosaraju, M\. Bavarian, M\. Chen, H\. Jun, L\. Kaiser, M\. Plappert, J\. Tworek, J\. Hilton, R\. Nakano, C\. Hesse, and J\. Schulman \(2021\)Training verifiers to solve math word problems\.External Links:2110\.14168,[Link](https://arxiv.org/abs/2110.14168)Cited by:[§6\.1](https://arxiv.org/html/2606.10829#S6.SS1.p1.1)\.
- D\. Hendrycks, C\. Burns, S\. Kadavath, A\. Arora, S\. Basart, E\. Tang, D\. Song, and J\. Steinhardt \(2021\)Measuring mathematical problem solving with the math dataset\.External Links:2103\.03874,[Link](https://arxiv.org/abs/2103.03874)Cited by:[§6\.1](https://arxiv.org/html/2606.10829#S6.SS1.p1.1)\.
- C\. Hong, S\. An, M\. Kim, and J\. C\. Ye \(2025\)Improving discrete diffusion unmasking policies beyond explicit reference policies\.External Links:2510\.05725,[Link](https://arxiv.org/abs/2510.05725)Cited by:[§1](https://arxiv.org/html/2606.10829#S1.p3.1),[§1](https://arxiv.org/html/2606.10829#S1.p6.1),[§2](https://arxiv.org/html/2606.10829#S2.p1.1)\.
- D\. Israel, G\. V\. den Broeck, and A\. Grover \(2025\)Accelerating diffusion llms via adaptive parallel decoding\.External Links:2506\.00413,[Link](https://arxiv.org/abs/2506.00413)Cited by:[§1](https://arxiv.org/html/2606.10829#S1.p3.1),[§1](https://arxiv.org/html/2606.10829#S1.p6.1),[§2](https://arxiv.org/html/2606.10829#S2.p2.1)\.
- M\. Jazbec, T\. X\. Olausson, L\. Béthune, P\. Ablin, M\. Kirchhof, J\. Monteiro, V\. Turrisi, J\. Ramapuram, and M\. Cuturi \(2025\)Learning unmasking policies for diffusion language models\.External Links:2512\.09106,[Link](https://arxiv.org/abs/2512.09106)Cited by:[§1](https://arxiv.org/html/2606.10829#S1.p3.1),[§1](https://arxiv.org/html/2606.10829#S1.p6.1),[§2](https://arxiv.org/html/2606.10829#S2.p1.1)\.
- B\. Kim, D\. Jeon, M\. Jeon, and A\. No \(2026a\)Dependency\-aware parallel decoding via attention for diffusion LLMs\.External Links:2603\.12996,[Link](https://arxiv.org/abs/2603.12996)Cited by:[§2](https://arxiv.org/html/2606.10829#S2.p3.1),[§6\.1](https://arxiv.org/html/2606.10829#S6.SS1.p3.1)\.
- S\. H\. Kim, S\. Hong, H\. Jung, Y\. Park, and S\. Yun \(2026b\)KLASS: kl\-guided fast inference in masked diffusion models\.External Links:2511\.05664,[Link](https://arxiv.org/abs/2511.05664)Cited by:[§2](https://arxiv.org/html/2606.10829#S2.p2.1),[§6\.1](https://arxiv.org/html/2606.10829#S6.SS1.p3.1)\.
- S\. Lee, S\. Kim, J\. Park, and D\. Park \(2025\)Lookahead unmasking elicits accurate decoding in diffusion language models\.External Links:2511\.05563,[Link](https://arxiv.org/abs/2511.05563)Cited by:[§1](https://arxiv.org/html/2606.10829#S1.p3.1),[§1](https://arxiv.org/html/2606.10829#S1.p6.1),[§2](https://arxiv.org/html/2606.10829#S2.p1.1)\.
- O\. Luxembourg, H\. Permuter, and E\. Nachmani \(2025\)Plan for speed: dilated scheduling for masked diffusion language models\.External Links:2506\.19037,[Link](https://arxiv.org/abs/2506.19037)Cited by:[§1](https://arxiv.org/html/2606.10829#S1.p3.1),[§1](https://arxiv.org/html/2606.10829#S1.p6.1),[§2](https://arxiv.org/html/2606.10829#S2.p2.1)\.
- S\. Nie, F\. Zhu, Z\. You, X\. Zhang, J\. Ou, J\. Hu, J\. Zhou, Y\. Lin, J\. Wen, and C\. Li \(2025\)Large language diffusion models\.External Links:2502\.09992,[Link](https://arxiv.org/abs/2502.09992)Cited by:[§6\.1](https://arxiv.org/html/2606.10829#S6.SS1.p1.1)\.
- F\. Z\. Peng, Z\. Bezemek, S\. Patel, J\. Rector\-Brooks, S\. Yao, A\. J\. Bose, A\. Tong, and P\. Chatterjee \(2025\)Path planning for masked diffusion model sampling\.External Links:2502\.03540,[Link](https://arxiv.org/abs/2502.03540)Cited by:[§1](https://arxiv.org/html/2606.10829#S1.p3.1),[§2](https://arxiv.org/html/2606.10829#S2.p1.1)\.
- S\. S\. Sahoo, M\. Arriola, Y\. Schiff, A\. Gokaslan, E\. Marroquin, J\. T\. Chiu, A\. Rush, and V\. Kuleshov \(2024\)Simple and effective masked diffusion language models\.External Links:2406\.07524,[Link](https://arxiv.org/abs/2406.07524)Cited by:[§1](https://arxiv.org/html/2606.10829#S1.p1.1),[§3\.2](https://arxiv.org/html/2606.10829#S3.SS2.p1.1)\.
- J\. Ye, Z\. Xie, L\. Zheng, J\. Gao, Z\. Wu, X\. Jiang, Z\. Li, and L\. Kong \(2025\)Dream 7b: diffusion large language models\.External Links:2508\.15487,[Link](https://arxiv.org/abs/2508.15487)Cited by:[§6\.1](https://arxiv.org/html/2606.10829#S6.SS1.p1.1)\.

## Appendix AEvaluation details

We evaluate on GSM8K, HumanEval, MBPP, and MATH500 with two diffusion LLMs: Dream 7B and LLaDA 8B\. We use 8\-shot prompting for GSM8K, 4\-shot prompting for MATH500, 3\-shot prompting for MBPP, and 0\-shot prompting for HumanEval\. For batched inference, we pad each prompt to the model’s maximum sequence length used for the corresponding task\. We use a maximum generation length of256256for GSM8K and512512for HumanEval, MBPP, and MATH500\. We use temperature0for all evaluations\.

#### Code availability\.

We provide anonymized reproduction code and instructions in the supplemental material\. The code will be made public after the review process\.

#### Decoding controls\.

For the entropy\-budget \(EB\) and Fast\-dLLM families, we sweep the budget parametersγ\\gammaandff\. For the fixed\-subset Top\-kkfamily, we sweep the subset sizekkdirectly\.

Table 2:Hyperparameter ranges used for the three decoder families\.Method familySweep valuesEB / EB\+ADγ∈\{10−3,10−2,10−1,0\.3,1,3,10,20\}\\gamma\\in\\\{10^\{\-3\},10^\{\-2\},10^\{\-1\},0\.3,1,3,10,20\\\}Fast\-dLLMf∈\{0\.3,1,2,3,5,8,12\}f\\in\\\{0\.3,1,2,3,5,8,12\\\}Fast\-dLLM\+ADf∈\{0\.3,1,2,3,5,8,12,20\}f\\in\\\{0\.3,1,2,3,5,8,12,20\\\}Top\-kk/ Top\-kk\+ADk∈\{1,2,4,8,16\}k\\in\\\{1,2,4,8,16\\\}For EB and Fast\-dLLM, larger threshold permits more aggressive parallel unmasking\. For Top\-kk, the level of parallelism is controlled directly bykk, i\.e\., the number of tokens unmasked per denoising step\. We include the extraf=20f=20point for Fast\-dLLM\+AD only to probe the very high\-parallel regime reached by the vanilla decoder at smaller thresholds; matched\-NFE summaries use interpolation at baseline operating points rather than threshold\-wise comparisons\.

#### Compute resources\.

All evaluations were run as batched inference jobs on one 4×\\timesGH200 GPU node per operating point\. We used padding to the task maximum sequence length for batched inference\. Dream 7B used batch size 4 per GPU, and LLaDA 8B used batch size 2 per GPU\.

Table 3:Actual wall\-clock duration by model and dataset, aggregated over the completed full\-dataset operating points used in the paper figures\.ModelDatasetMin–max hrs/runDream 7BGSM8K0\.13–2\.21Dream 7BHumanEval0\.06–0\.57Dream 7BMATH5000\.08–2\.00Dream 7BMBPP0\.10–2\.15LLaDA 8BGSM8K0\.26–4\.86LLaDA 8BHumanEval0\.07–1\.26LLaDA 8BMATH5000\.14–3\.67LLaDA 8BMBPP0\.15–4\.27
#### Hyperparameters for KLASS and DAPD\.

For KLASS and DAPD, we did not perform per\-task accuracy tuning\. Instead, we first ran small pilot evaluations to identify threshold values that place each method in the same broad low\-NFE regime as our main curves, and then reused the same hyperparameter grids for all models and datasets\. For KLASS, which accepts tokens satisfying both a confidence and KL criterion, we evaluated

\(θconf,θKL\)∈\{\(0\.70,1\.00\),\(0\.198,0\.51\),\(0\.18,0\.60\),\(0\.16,0\.70\)\}\.\(\\theta\_\{\\mathrm\{conf\}\},\\theta\_\{\\mathrm\{KL\}\}\)\\in\\\{\(0\.70,1\.00\),\(0\.198,0\.51\),\(0\.18,0\.60\),\(0\.16,0\.70\)\\\}\.For DAPD, using the masked max\-normalized attention dependency score, we evaluated

\(τmin,τmax\)∈\{\(0\.002,0\.02\),\(0\.005,0\.05\),\(0\.01,0\.10\),\(0\.02,0\.20\),\(0\.05,0\.50\)\}\.\(\\tau\_\{\\min\},\\tau\_\{\\max\}\)\\in\\\{\(0\.002,0\.02\),\(0\.005,0\.05\),\(0\.01,0\.10\),\(0\.02,0\.20\),\(0\.05,0\.50\)\\\}\.These settings were chosen to span a range of parallelism levels rather than to maximize accuracy on any individual benchmark\.

## Appendix BTheoretical motivation for the ADAS utility

We provide a local derivation motivating the ADAS utility\. The argument is not meant to establish a global guarantee for an arbitrary transformer; rather, it shows that under a standard first\-order approximation of the denoiser around the current decoding state, the natural dependence penalty has the same form as the ADAS discount\.

Fix the current partially unmasked sequencexℳ¯x^\{\\bar\{\\mathcal\{M\}\}\}\. For each masked positionss, let

ps​\(a\)=pθ​\(xs=a∣xℳ¯\),x^s=arg⁡maxa∈𝒱⁡ps​\(a\),cs=ps​\(x^s\)\.p\_\{s\}\(a\)=p\_\{\\theta\}\(x^\{s\}=a\\mid x^\{\\bar\{\\mathcal\{M\}\}\}\),\\qquad\\hat\{x\}\_\{s\}=\\arg\\max\_\{a\\in\\mathcal\{V\}\}p\_\{s\}\(a\),\\qquad c\_\{s\}=p\_\{s\}\(\\hat\{x\}\_\{s\}\)\.LetZs=ϕs​\(Xs\)Z\_\{s\}=\\phi\_\{s\}\(X^\{s\}\)be the representation inserted at positionss, and letz^s=ϕs​\(x^s\)\\hat\{z\}\_\{s\}=\\phi\_\{s\}\(\\hat\{x\}\_\{s\}\)\. Since the vocabulary is finite, define

Dϕ2=maxs∈ℳ⁡maxa,b∈𝒱⁡‖ϕs​\(a\)−ϕs​\(b\)‖22\.D\_\{\\phi\}^\{2\}=\\max\_\{s\\in\\mathcal\{M\}\}\\max\_\{a,b\\in\\mathcal\{V\}\}\\\|\\phi\_\{s\}\(a\)\-\\phi\_\{s\}\(b\)\\\|\_\{2\}^\{2\}\.Then

𝔼​\[‖Zs−z^s‖22\]\\displaystyle\\mathbb\{E\}\\\!\\left\[\\\|Z\_\{s\}\-\\hat\{z\}\_\{s\}\\\|\_\{2\}^\{2\}\\right\]=∑a∈𝒱ps​\(a\)​‖ϕs​\(a\)−ϕs​\(x^s\)‖22\\displaystyle=\\sum\_\{a\\in\\mathcal\{V\}\}p\_\{s\}\(a\)\\\|\\phi\_\{s\}\(a\)\-\\phi\_\{s\}\(\\hat\{x\}\_\{s\}\)\\\|\_\{2\}^\{2\}≤Dϕ2​∑a≠x^sps​\(a\)\\displaystyle\\leq D\_\{\\phi\}^\{2\}\\sum\_\{a\\neq\\hat\{x\}\_\{s\}\}p\_\{s\}\(a\)=Dϕ2​\(1−cs\)\.\\displaystyle=D\_\{\\phi\}^\{2\}\(1\-c\_\{s\}\)\.Thus1−cs1\-c\_\{s\}upper\-bounds the representation uncertainty of the selected source tokenss\.

Now consider appending a candidate positioniiafter an already selected setSS\. Letℓi​\(zS\)\\ell\_\{i\}\(z\_\{S\}\)denote the logits at positioniiwhen the selected positions are filled with representationszS=\{zs:s∈S\}z\_\{S\}=\\\{z\_\{s\}:s\\in S\\\}\. Around the current decoding statez^S\\hat\{z\}\_\{S\}, a first\-order expansion gives

ℓi​\(zS\)≈ℓi​\(z^S\)\+∑s∈SJi​s​\(zs−z^s\),\\ell\_\{i\}\(z\_\{S\}\)\\approx\\ell\_\{i\}\(\\hat\{z\}\_\{S\}\)\+\\sum\_\{s\\in S\}J\_\{is\}\(z\_\{s\}\-\\hat\{z\}\_\{s\}\),whereJi​sJ\_\{is\}is the Jacobian of the logits at positioniiwith respect to the representation at positionss\.

For a self\-attention layer, the direct message from sourcessto targetiiis weighted by the attention coefficientAi​sA\_\{is\}\. We therefore use the empirical attention matrix as a tractable proxy for the relative sensitivity ofiito perturbations atss, and write the local approximation

‖Ji​s‖22≈Li2​Ai​s,\\\|J\_\{is\}\\\|\_\{2\}^\{2\}\\approx L\_\{i\}^\{2\}A\_\{is\},whereLiL\_\{i\}absorbs the remaining local linear maps, residual transformations, and normalization factors\. This approximation says that source positions receiving larger attention fromiihave larger local influence on the logits atii\.

Using this approximation,

𝔼​\[‖ℓi​\(ZS\)−ℓi​\(z^S\)‖22\]\\displaystyle\\mathbb\{E\}\\\!\\left\[\\\|\\ell\_\{i\}\(Z\_\{S\}\)\-\\ell\_\{i\}\(\\hat\{z\}\_\{S\}\)\\\|\_\{2\}^\{2\}\\right\]≈𝔼​\[‖∑s∈SJi​s​\(Zs−z^s\)‖22\]\\displaystyle\\approx\\mathbb\{E\}\\\!\\left\[\\left\\\|\\sum\_\{s\\in S\}J\_\{is\}\(Z\_\{s\}\-\\hat\{z\}\_\{s\}\)\\right\\\|\_\{2\}^\{2\}\\right\]≲Li2​∑s∈SAi​s​𝔼​\[‖Zs−z^s‖22\]\\displaystyle\\lesssim L\_\{i\}^\{2\}\\sum\_\{s\\in S\}A\_\{is\}\\mathbb\{E\}\\\!\\left\[\\\|Z\_\{s\}\-\\hat\{z\}\_\{s\}\\\|\_\{2\}^\{2\}\\right\]≤Li2​Dϕ2​∑s∈SAi​s​\(1−cs\)\.\\displaystyle\\leq L\_\{i\}^\{2\}D\_\{\\phi\}^\{2\}\\sum\_\{s\\in S\}A\_\{is\}\(1\-c\_\{s\}\)\.Finally, for softmax probabilities, local logit perturbations induce KL changes that are second order in the perturbation size:

DKL\(softmax\(ℓi\+δ\)∥softmax\(ℓi\)\)≤12∥δ∥22\.D\_\{\\mathrm\{KL\}\}\\left\(\\operatorname\{softmax\}\(\\ell\_\{i\}\+\\delta\)\\;\\middle\\\|\\;\\operatorname\{softmax\}\(\\ell\_\{i\}\)\\right\)\\leq\\frac\{1\}\{2\}\\\|\\delta\\\|\_\{2\}^\{2\}\.Therefore, the incremental dependence risk of adding candidateiiafterSSis locally controlled by a quantity proportional to

∑s∈SAi​s​\(1−cs\)\.\\sum\_\{s\\in S\}A\_\{is\}\(1\-c\_\{s\}\)\.
This motivates the risk\-adjusted marginal utility

u​\(i∣S,A\)=ci−α​∑s∈SAi​s​\(1−cs\),u\(i\\mid S,A\)=c\_\{i\}\-\\alpha\\sum\_\{s\\in S\}A\_\{is\}\(1\-c\_\{s\}\),whereα≥0\\alpha\\geq 0absorbs the constants in the local approximation and controls the strength of the dependence penalty\.

## Appendix CHyperparameter and ablation studies

#### Attention\-discount strength

We ablate the attention\-discount strengthα\\alphaon LLaDA\-8B HumanEval under three decoding families: fixed Top\-kk, AD\+EB, and Fast\-dLLM\+AD\. For fixed Top\-kk, we testk∈\{8,16\}k\\in\\\{8,16\\\}\. For AD\+EB, we use two entropy\-budget settings chosen to match these operating regimes approximately:γ=10\\gamma=10\(low NFE\) andγ=3\\gamma=3\(medium NFE\)\. For AD Fast\-dLLM, we similarly usef=5f=5andf=2f=2\. In all cases, we sweep

α∈\{1,5,10,20,40,80\},\\alpha\\in\\\{1,5,10,20,40,80\\\},and compare against the corresponding no\-AD baseline, which we denote asα=0\\alpha=0\.

Table 4:Pass@1 accuracy \(%\) for theα\\alphasweep on LLaDA HumanEval across fixed Top\-kk, AD\+EB, and AD Fast\-dLLM\. The no\-AD baseline is shown asα=0\\alpha=0\.Settingα=0\\alpha=0α=1\\alpha=1α=5\\alpha=5α=10\\alpha=10α=20\\alpha=20α=40\\alpha=40α=80\\alpha=80Top\-kk,k=16k=161\.834\.887\.325\.495\.494\.884\.88Top\-kk,k=8k=85\.499\.1518\.9018\.2921\.3420\.1218\.90AD\+EB,γ=10\\gamma=109\.769\.1520\.1223\.7823\.7823\.1723\.78AD\+EB,γ=3\\gamma=321\.3420\.7323\.1722\.5626\.2226\.8322\.56Fast\-dLLM\+AD,f=5f=59\.769\.1521\.9523\.7824\.3924\.3921\.95Fast\-dLLM\+AD,f=2f=228\.0529\.2727\.4428\.0530\.4933\.5433\.54Average over 6 settings–13\.7219\.8220\.3321\.9522\.1520\.93The bestα\\alphadepends on the decoding regime when viewed per setting\. For example, fixed Top\-kkprefers smaller values in some cases, whereas the adaptive stopping\-rule methods are more favorable to larger values\. However, when averaging across all six matched settings in Table[4](https://arxiv.org/html/2606.10829#A3.T4),α=40\\alpha=40achieves the highest overall accuracy\. Relative to an oracle that selects the bestα\\alphaseparately for each row, the global default loses only0\.720\.72points on average and at most2\.442\.44points in the worst case\. We therefore useα=40\\alpha=40as a single global default in the main experiments, avoiding task\- or regime\-specific retuning while remaining competitive across the tested settings\.

#### Ablation on the uncertainty weighting term\.

We additionally ablate the uncertainty factor used inside the attention discount\. The default ADAS utility discounts a candidate position according to the uncertainty of the already selected position, using the factor1−cs1\-c\_\{s\}\. We compare this choice against a variant that replaces1−cs1\-c\_\{s\}with a global average uncertainty over the currently masked positions, denoted1−cmask1\-c\_\{\\mathrm\{mask\}\}\. This ablation is conducted on LLaDA\-8B\-Base with HumanEval under Top\-kk\+AD decoding atk=8k=8\.

Table 5:Ablation of the uncertainty weighting term in the attention discount on LLaDA\-8B\-Base HumanEval with Top\-kk\+AD decoding atk=8k=8\. The default selected\-token uncertainty1−cs1\-c\_\{s\}outperforms using the average uncertainty over masked positions1−cmask1\-c\_\{\\mathrm\{mask\}\}\.Uncertainty weighting termPass@1 \(%\)Selected\-token uncertainty1−cs1\-c\_\{s\}20\.12Average masked\-position uncertainty1−cmask1\-c\_\{\\mathrm\{mask\}\}17\.07The selected\-token uncertainty weighting improves Pass@1 by3\.053\.05absolute points over the global masked\-position uncertainty variant\. This suggests that the discount benefits from conditioning on the uncertainty of the specific selected position rather than using a single global uncertainty scalar\.

#### Effect of attention source

We ablate which layer is used to construct the attention matrix in ADAS\. We use LLaDA\-8B\-Base on HumanEval with fixed Top\-kkdecoding atk=8k=8, so all variants use the same number of denoising steps\. The no\-ADAS baseline ranks tokens only by confidence, while the ADAS variants use mean attention over heads from the specified transformer layer\.

Table 6:Attention\-source ablation for ADAS on LLaDA\-8B HumanEval with fixed Top\-kkdecoding atk=8k=8\. Final\-layer attention performs best, suggesting that later attention patterns provide the most useful dependency signal for selecting tokens to unmask in parallel\.Selection rulePass@1 \(%\)Top\-kkconfidence, no ADAS5\.49Top\-kk\+AD with first\-layer attention9\.76Top\-kk\+AD with middle\-layer attention15\.24Top\-kk\+AD with final\-layer attention20\.12

## Appendix DRobustness across operating points

We perform an additional robustness analysis over matched operating points\. Each operating point corresponds to a model–dataset–stopping\-rule–hyperparameter configuration\. This analysis is intended to assess whether the observed gains are stable across the evaluated settings, rather than driven by a small number of favorable thresholds or tasks\.

For each vanilla baseline point in the low\-NFE regime, defined as having at least four parallel tokens per denoising step on average, we linearly interpolate the corresponding ADAS curve at the same NFE and compute the matched gain

Δ=scoreADAS​\(NFE\)−scorebase​\(NFE\)\.\\Delta=\\mathrm\{score\}\_\{\\mathrm\{ADAS\}\}\(\\mathrm\{NFE\}\)\-\\mathrm\{score\}\_\{\\mathrm\{base\}\}\(\\mathrm\{NFE\}\)\.Baseline points outside the interpolation range of the corresponding ADAS curve are excluded\. We then bootstrap the matched gains over operating points with 10,000 resamples and report the mean gain and 95% confidence interval\. Because operating points from the same model, dataset, or decoding family are correlated, this should be interpreted as a robustness summary across the evaluated configurations, not as a paired per\-example significance test\.

Table 7:Robustness of ADAS gains across matched operating points\. Gains are absolute accuracy/pass@1 differences in percentage points after matching by NFE\. Confidence intervals are obtained by bootstrapping operating points with 10,000 resamples\.Method familyMatched pointsMean gain95% CIPositive / negativeEB\-Sampler28\+10\.42\[\+7\.30, \+13\.67\]24 / 4Fast\-dLLM38\+7\.24\[\+5\.42, \+9\.14\]34 / 4Top\-kk24\+11\.13\[\+8\.31, \+13\.97\]22 / 2All90\+9\.27\[\+7\.74, \+10\.84\]80 / 10Table 8:Robustness of ADAS gains across datasets under the same matched\-operating\-point analysis as Table[7](https://arxiv.org/html/2606.10829#A4.T7)\.DatasetMatched pointsMean gain95% CIPositive / negativeGSM8K19\+8\.60\[\+4\.71, \+12\.48\]14 / 5HumanEval22\+10\.33\[\+8\.10, \+12\.50\]21 / 1MATH50024\+5\.73\[\+3\.84, \+7\.64\]21 / 3MBPP25\+12\.24\[\+8\.99, \+15\.59\]24 / 1
## Appendix EConstruction of the known dependency samples

Dataset ConstructionAssume the set of equation typesE=\{A\+B\+C=D,A⋅B⋅C=D,min⁡\{A,B,C\}=D,max⁡\{A,B,C\}=D\}E=\\left\\\{A\+B\+C=D,\\;A\\cdot B\\cdot C=D,\\;\\min\\\{A,B,C\\\}=D,\\;\\max\\\{A,B,C\\\}=D\\right\\\}is given, together with the number of predicatesnnand masking ratiorr\.1\. Sample predicate indices\.We sample a sequenceπ=\(π1,…,πn\),πi∈\{0,…,\|E\|−1\},\\pi=\(\\pi\_\{1\},\\ldots,\\pi\_\{n\}\),\\qquad\\pi\_\{i\}\\in\\\{0,\\ldots,\|E\|\-1\\\},and construct theii\-th predicate from the equation typeE​\[πi\]E\[\\pi\_\{i\}\]\.2\. Instantiate the predicates\.Each predicate contains four integer slots\(A,B,C,D\)\(A,B,C,D\), so the full sequence contains4​n4ninteger variables\. We sample positive integers forA,B,CA,B,Cand computeDDaccording to the selected equation type\.3\. Mask integer variables\.We randomly maskr⋅4​nr\\cdot 4nof the integer positions in the sequence\. Only the integers are masked; the operators and relation symbols \(\+\+,⋅\\cdot,min\\min,max\\max,==\) are never masked\.

Figure 4:Illustration of the toy data\-generation process\. The construction yields a masked sequence containing masked\-token pairs with known dependency structure\.Table 9:Attention\-dependency relationship on the shared toy dataset\. Reported values are mean pairwise attention for dependent and non\-dependent masked pairs aggregated as mean of attention heads from the last layer\.ModelDep\. Attn\.Non\-Dep\. Attn\.LLaDA\-8B\-Base0\.0072610\.001919Dream\-7B\-Base0\.0095170\.004842
## Appendix FDetailed results

Table 10:Results onDream\-7B\-Basewith EB stopping rule\. Each cell reportsscoreon the first line and average NFE over samples on the second line\. For GSM8K and MATH500, score is accuracy; for HumanEval and MBPP, score is Pass@1\.GSM8KHumanEvalMATH500MBPPγ\\gammaEBEB \+ ADEBEB \+ ADEBEB \+ ADEBEB \+ AD0\.00175\.06240\.5075\.06240\.3753\.66453\.4453\.66453\.0640\.00386\.6540\.00384\.9953\.80432\.4953\.80431\.560\.0175\.13209\.2375\.13208\.9354\.27372\.9054\.27372\.8339\.40269\.4839\.20266\.0653\.60305\.4953\.60302\.700\.174\.60166\.8774\.30165\.7654\.27289\.8454\.88287\.9539\.80168\.0840\.00163\.5253\.00179\.3853\.00174\.840\.374\.68143\.4274\.45142\.5952\.44237\.5353\.66236\.0838\.80131\.0039\.20122\.1753\.80131\.5853\.40126\.62169\.98110\.4267\.32104\.9946\.34172\.4145\.73162\.4934\.8091\.6533\.0089\.2350\.4086\.7047\.8082\.90357\.0970\.4150\.6461\.0030\.49116\.1235\.3785\.0624\.4062\.2725\.2060\.8136\.6055\.7143\.2054\.471028\.6634\.2933\.2824\.2112\.2054\.0125\.6136\.439\.6034\.2420\.2027\.8811\.8036\.5035\.0026\.77207\.9617\.6623\.2013\.363\.0532\.3220\.1222\.633\.6020\.129\.0017\.854\.2022\.8231\.0016\.43Table 11:Results onLLaDA\-8B\-Basewith EB stopping rule\. Each cell reportsscoreon the first line and average NFE over samples on the second line\. For GSM8K and MATH500, score is accuracy; for HumanEval and MBPP, score is Pass@1\.GSM8KHumanEvalMATH500MBPPγ\\gammaEBEB \+ ADEBEB \+ ADEBEB \+ ADEBEB \+ AD0\.00169\.90237\.1769\.90236\.9235\.37457\.1235\.37456\.0428\.40460\.2428\.40459\.0938\.80475\.9338\.80475\.470\.0169\.83190\.8369\.83189\.8335\.37326\.0835\.37321\.8428\.20336\.5028\.00332\.5638\.80338\.3338\.80334\.520\.170\.81136\.1370\.96134\.0335\.98190\.3535\.98178\.6027\.80209\.9228\.00203\.3939\.20197\.3239\.20191\.630\.370\.20110\.5470\.28107\.9334\.76135\.5434\.15126\.7228\.80164\.3828\.40157\.0439\.40144\.0439\.20136\.83167\.4078\.1067\.1773\.5629\.2787\.1929\.8875\.8429\.00109\.1125\.8098\.8836\.8095\.7038\.4086\.68360\.2050\.3655\.1244\.8921\.3459\.3826\.8343\.7121\.4069\.3326\.0068\.2928\.4063\.6634\.4055\.081038\.2129\.8241\.1722\.149\.7632\.9023\.1723\.9510\.8036\.0218\.8032\.5911\.2033\.3426\.0029\.612013\.7217\.7032\.8314\.244\.8823\.4121\.9517\.903\.2020\.8618\.6020\.095\.0020\.6021\.2019\.75Table 12:Results onDream\-7B\-Basewith Fast\-dLLM stopping rule\. Each cell reportsscoreon the first line and average NFE over samples on the second line\. For GSM8K and MATH500, score is accuracy; for HumanEval and MBPP, score is Pass@1\.GSM8KHumanEvalMATH500MBPPffFDFD \+ ADFDFD \+ ADFDFD \+ ADFDFD \+ AD0\.374\.45167\.0774\.98167\.5955\.49273\.2755\.49275\.5740\.40146\.8040\.60148\.1954\.00138\.6654\.00139\.72173\.01132\.5973\.31134\.9750\.61201\.4451\.22205\.7838\.20104\.1138\.40107\.5051\.0090\.3651\.6092\.76263\.3195\.3665\.0598\.7836\.59139\.4842\.07147\.0731\.6075\.1932\.8078\.0638\.0063\.2045\.4065\.30351\.8655\.6052\.8458\.7123\.1785\.3234\.1584\.8722\.4056\.2927\.4062\.7329\.4048\.9640\.2053\.23527\.9029\.1240\.9432\.4914\.6346\.1026\.8349\.239\.6030\.5222\.4040\.5115\.2033\.7836\.4035\.6487\.9614\.4726\.9118\.967\.9329\.9822\.5632\.284\.0017\.7310\.8027\.318\.0021\.0932\.4022\.91121\.748\.1711\.0711\.046\.1021\.6518\.2922\.380\.8011\.929\.0019\.722\.4013\.8727\.8016\.16Table 13:Results onLLaDA\-8B\-Basewith Fast\-dLLM stopping rule\. Each cell reportsscoreon the first line and average NFE over samples on the second line\. For GSM8K and MATH500, score is accuracy; for HumanEval and MBPP, score is Pass@1\.GSM8KHumanEvalMATH500MBPPffFDFD \+ ADFDFD \+ ADFDFD \+ ADFDFD \+ AD0\.369\.98125\.1669\.98126\.0235\.37132\.6335\.37134\.8028\.40182\.5428\.20184\.1338\.60153\.7138\.60154\.44169\.2291\.2869\.2993\.1635\.3786\.6535\.3787\.5927\.20128\.4327\.60131\.6239\.00101\.4939\.40102\.71264\.2262\.6063\.8465\.1028\.0560\.0133\.5464\.5924\.4085\.3427\.8094\.1830\.8067\.8337\.0071\.91355\.5746\.3956\.1847\.2618\.9047\.0927\.4444\.5719\.8063\.0425\.8079\.0321\.0048\.2932\.6058\.54539\.0430\.2945\.4931\.379\.7629\.4324\.3934\.6312\.2033\.4819\.6049\.7610\.4029\.4528\.6043\.96815\.6216\.2940\.6421\.564\.8819\.4617\.0726\.012\.4018\.2019\.6034\.056\.4017\.9825\.2031\.80124\.328\.9632\.3715\.123\.0513\.4410\.3710\.680\.2011\.3214\.6021\.992\.6013\.0319\.2023\.37206\.209\.966\.406\.76Table 14:Results onDream\-7B\-Basewith Top\-kkstopping rule\. Each cell reportsscoreon the first line and average NFE over samples on the second line\. For GSM8K and MATH500, score is accuracy; for HumanEval and MBPP, score is Pass@1\.GSM8KHumanEvalMATH500MBPPkkTop\-kTop\-k \+ ADTop\-kTop\-k \+ ADTop\-kTop\-k \+ ADTop\-kTop\-k \+ AD174\.91256\.0074\.91256\.0055\.49512\.0055\.49512\.0040\.80512\.0040\.80512\.0054\.00512\.0054\.00512\.00264\.67128\.0058\.15128\.0042\.68256\.0043\.29256\.0031\.40256\.0032\.60256\.0044\.60256\.0046\.40256\.00448\.4564\.0046\.4764\.0020\.73128\.0035\.98128\.0016\.20128\.0025\.00128\.0029\.20128\.0040\.20128\.00813\.8032\.0028\.8132\.0012\.2064\.0021\.9564\.003\.6064\.0011\.4064\.0011\.4064\.0034\.4064\.00160\.9116\.005\.5316\.007\.9332\.0015\.2432\.001\.2032\.005\.4032\.002\.0032\.0026\.8032\.00Table 15:Results onLLaDA\-8B\-Basewith Top\-kkstopping rule\. Each cell reportsscoreon the first line and average NFE over samples on the second line\. For GSM8K and MATH500, score is accuracy; for HumanEval and MBPP, score is Pass@1\.GSM8KHumanEvalMATH500MBPPkkTop\-kTop\-k \+ ADTop\-kTop\-k \+ ADTop\-kTop\-k \+ ADTop\-kTop\-k \+ AD169\.98256\.0069\.98256\.0035\.37512\.0035\.37512\.0028\.40512\.0028\.40512\.0038\.80512\.0038\.80512\.00264\.22128\.0059\.67128\.0025\.61256\.0026\.83256\.0022\.80256\.0024\.40256\.0028\.80256\.0036\.80256\.00453\.9064\.0050\.1164\.0013\.41128\.0024\.39128\.0017\.00128\.0023\.00128\.0017\.60128\.0029\.80128\.00822\.6732\.0042\.4632\.005\.4964\.0020\.1264\.003\.4064\.0019\.2064\.007\.0064\.0025\.2064\.00162\.2016\.0021\.9116\.001\.8332\.004\.8832\.000\.2032\.0010\.6032\.001\.6032\.0012\.2032\.00