Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models¶

Conference: NeurIPS2026
arXiv: 2606.10829
Code: https://github.com/yusufsahin99/ADAS
Area: LLM Efficiency / Text Generation
Keywords: masked diffusion language models, parallel decoding, attention discount, dependency-aware sampling, uncertainty weighting

TL;DR¶

Without training the model or changing the base sampler's stopping rule, ADAS dynamically discounts confidence using a candidate's attention to positions already selected in the current step and their uncertainty, improving parallel decoding at low numbers of denoiser evaluations; task-and-method-average gains are 9.11 and 10.46 percentage points for the two models.

Background & Motivation¶

A masked diffusion language model does not produce just the next token at each step: it predicts all masked positions in a partially observed sequence and reveals several together. Revealing more tokens per step directly reduces the number of function evaluations (NFE), but their predictions all come from the old context, before these new tokens are available. When two positions depend on each other, individually confident predictions need not form a consistent jointly committed output. Mathematical steps and code structure are particularly exposed to this synchronous commitment problem.

Top-k, Fast-dLLM, and EB-Sampler control how many tokens to reveal using fixed cardinality, a confidence condition, and an entropy budget, respectively, while typically retaining token-wise candidate ranking. Smaller batches can mitigate errors but do not answer which tokens should be revealed together at the same batch size. Token-wise model error differs from the error introduced by approximating a conditional joint distribution with the product of its marginals; selecting the highest-confidence positions directly addresses only the former. Alternatives include trained planners, lookahead over decoding trajectories, or hard attention-derived conflict graphs, but these introduce training, computation, or new batching constraints.

This paper replaces only candidate ranking. The same forward pass already provides attention, which can estimate whether a candidate depends on other positions just selected in the current step, without an extra model call for every candidate. Attention is not evidence that two positions must never be updated together, so the method uses a continuous penalty rather than a binary conflict. Core Idea: retain the base sampler's decision about when to stop, while making the next candidate depend on its directional attention to uncertain positions already selected in the current step.

Method¶

Overall Architecture¶

ADAS is an inference-time greedy candidate selector, not a new denoising network. Its inputs come from one denoiser pass on the current partially masked sequence: token distributions, confidence defined by the highest-probability token, and final-layer self-attention averaged across heads. Its output is the set of positions to reveal together; token values still come from the base model's existing predictions rather than being regenerated from discounted scores.

Each step starts with an empty selected set and raw confidence scores for all masked positions. The highest-scoring candidate is checked against the base sampler's admissibility condition. If admissible, it is added and the remaining scores are decreased before selection repeats. If inadmissible, subset construction stops immediately and the selected positions are revealed together. The next model forward pass reads these new tokens and recomputes the distributions, attention, and selected set for the next step.

“Selected” specifically means chosen within the current step but not yet jointly inserted into the context, rather than text revealed in previous steps. Intra-step ranking updates cached scores without rerunning the denoiser after each selection. This distinction lets ADAS estimate joint commitment risk without reducing parallel generation to a sequence of token-by-token forward passes.

Key Designs¶

1. Confidence initialization: retain the base model's token-wise value estimate

Confidence is the probability assigned to the most likely token at a masked position, not entropy, sampling temperature, or the discounted score. Let \(\mathcal{M}\) be the masked positions, \(S\) the positions selected in the current step, and \(c_i\) the raw confidence. This retains the familiar preference for committing predictions the model finds more certain, allowing ADAS to integrate into samplers that already use confidence.

When \(S\) is empty, there is no joint commitment penalty, so ADAS makes the same first choice as confidence ranking. The difference begins with the second selection: conventional ranking continues comparing fixed raw confidence, whereas ADAS changes the relative value of remaining positions in response to the first choice. High-attention positions are not permanently deferred; the next decision changes with the current selection history.

The paper's decomposition into model error and joint dependence error motivates this design, but the true conditional joint distribution is unavailable at inference time. ADAS neither computes true multi-information nor identifies the discount with an exact joint error; confidence and attention are tractable proxies for individual reliability and conditional influence.

2. Directional uncertainty discount: penalize dependence on unreliable selected positions

Let \(A_{is}\) be attention from candidate position \(i\) to selected position \(s\): the query is the candidate and the key is the selected position. The core marginal score is:

\[ u(i\mid S,A)=c_i-\alpha\sum_{s\in S}A_{is}(1-c_s). \]

\(\alpha\) controls discount strength, with \(\alpha=40\) shared across the main experiments. The essential feature is weighting attention by the selected token's \(1-c_s\): depending on a nearly certain token should not incur the same penalty as depending on an unreliable prediction. Uncertainty belongs to the selected token and must not be replaced by the candidate's own \(1-c_i\).

Direction must also be preserved. The quantity of interest is how much a candidate's distribution may change after the selected token is committed, so ADAS reads attention from candidate to selected token, not the reverse. It neither symmetrizes attention nor declares above-threshold pairs forbidden. A strongly dependent candidate can still be chosen when its own confidence is sufficiently high.

Penalties accumulate as \(S\) grows. After accepting position \(s\), each remaining score is decreased by \(\alpha A_{is}(1-c_s)\), avoiding recomputation of the full sum. Selecting a final subset of size \(|S|\) costs \(O(|S|\,|\mathcal{M}|)\). Intra-step selection remains sequential, but its cost is typically much smaller than a large-model forward pass.

This greedy procedure is order-dependent. The paper motivates it through the utility of appending to an ordered selection sequence, not exact maximization of an unordered set objective: a bidirectional unordered pairwise penalty would introduce reverse-direction terms absent from ADAS. The authors explicitly provide no global optimality or approximation guarantee. An interpretable local marginal score does not imply globally searched subset selection.

3. Base-rule admissibility: separate ranking from stopping conditions

The discounted score determines which candidate is proposed, while acceptance and stopping retain the base rule. Let \(T=S\cup\{i^\star\}\) be the proposed enlarged set and \(H_s\) the entropy of the original prediction at position \(s\). The three conditions are:

\[ \begin{aligned} \text{Top-}k &: |T|\leq k,\\ \text{Fast-dLLM} &: |T|\left(1-\min_{s\in T}c_s\right)<f,\\ \text{EB-Sampler} &: \sum_{s\in T}H_s-\max_{s\in T}H_s<\gamma. \end{aligned} \]

Top-k imposes only a cardinality cap. Fast-dLLM checks set size multiplied by an error-probability proxy for its least certain member. EB-Sampler checks whether total predictive entropy minus the largest individual entropy remains within budget. Confidence and entropy still come from the original distributions for the current step; discounted scores do not replace them in these conditions, and ADAS does not introduce a new stopping threshold.

The algorithm proposes the remaining candidate with the highest discounted score and then checks admissibility. If the candidate fails, the published algorithm stops rather than skipping it and searching all lower-scoring candidates for an admissible alternative. Reranking can therefore indirectly change the size returned by an adaptive rule even when its formula and threshold remain identical. This is why equal \(f\) or \(\gamma\) cannot simply be treated as equal computation in the experiments.

Fixed Top-k more cleanly isolates subset construction: with fixed cardinality and identical NFE, a quality change comes from which positions were selected rather than additional model evaluations. Conversely, at \(k=1\), both methods make the same first choice from an empty set, so ADAS should offer no reranking gain. The appendix indeed reports identical corresponding results for both models.

4. Attention-proxy diagnostics: validate conditional influence rather than equating attention with correctness

The authors extract states from ordinary Top-8 trajectories, choose the most confident position, and reveal only that position in a counterfactual copy that does not alter the original trajectory. They then measure distributional changes at the other seven co-selected candidates. Each task/model has 100 examples and five states per example, producing 3,500 position pairs. This extra forward pass is an offline diagnostic, not part of routine ADAS decoding.

Candidate-to-revealed-position attention correlates positively with intervention KL. Across all pairs, Spearman correlation is 0.284–0.545 for raw attention and 0.782–0.841 after weighting by the selected token's uncertainty. Within a state, however, the selected token and its uncertainty are fixed, so multiplication cannot change attention ranks. The weighting improves cross-state calibration; it does not demonstrate improved attention ranking within every state.

More concrete evidence comes from 1,098 states where the two rankings first diverge. Tokens that confidence ranking would select but ADAS defers have mean KL change 0.793 after the first new token arrives, versus 0.230 for their replacements; mean confidence gains are 0.092 and 0.024. This supports waiting for context before commitment, but free-form solutions have no unique token-aligned target. The diagnostic measures stability, while end-to-end task scores measure correctness.

Appendix C provides a local first-order motivation: selected-token representation uncertainty can be proxied by prediction error probability, and attention proxies a candidate's relative sensitivity to perturbations of that representation. Their product yields the discount term. This relies on local linearization and a sensitivity approximation, not a rigorous global risk bound for arbitrary Transformers, and does not capture all higher-order interactions.

A Worked Example¶

The following numbers are an illustrative example devised to explain the formula, not an observed paper sample. Suppose three masked positions have confidence 0.90, 0.88, and 0.85, with Top-2 and \(\alpha=40\). Starting from an empty set, the first position is selected at 0.90, just as under ordinary ranking.

The second candidate attends to the first position with weight 0.02, and the third with 0.001; the first selected token has uncertainty 0.10. Updated scores are \(0.88-40\times0.02\times0.10=0.80\) and \(0.85-40\times0.001\times0.10=0.846\), so the third position becomes the second selection despite its lower raw confidence.

The first and third positions are committed together, while the second remains masked and can be predicted again under the new context in the next step. This neither guarantees that its original token was wrong nor that the third token is correct; it preferentially defers a prediction more likely to be affected by the selected position. If the first token were almost certain, both attention penalties would be small.

Loss & Training¶

There is no new loss function, fine-tuning, learned planner, or auxiliary verifier: ADAS uses existing model weights and predictions. The implementation reads head-mean attention from the final layer and disables FlashAttention only in that layer to extract weights. Full-dataset timing estimates 3.1% mean overhead per model forward pass, not a guarantee that end-to-end latency increases by only 3.1% on every hardware platform or decoding setting.

The default \(\alpha=40\) is not a zero-tuning result. On LLaDA-8B HumanEval, the authors compare six nonzero strengths across two Top-k, two EB, and two Fast-dLLM settings, selecting once by mean Pass@1 over the six settings before reuse without task-, model-, or stopping-rule-specific retuning. Its six-setting average is 22.15%, versus 21.95% for \(\alpha=20\). Relative to an oracle selecting the best strength separately for each row, the default loses 0.72 percentage points on average and at most 2.44 points. This supports transfer of one default, not universal optimality across other models and tasks.

Key Experimental Results¶

Main Results¶

The models are LLaDA-8B-Base and Dream-7B-Base. Mathematical tasks report accuracy, and code tasks report Pass@1. GSM8K/MATH500/MBPP/HumanEval use 8/4/3/0-shot prompts, respectively, temperature 0, few-shot seed 1234, and maximum generation lengths 256/512/512/512.

The following excerpt from paper Table 2 reports LLaDA / Dream absolute gains over the corresponding base sampler in percentage points, not relative percentages.

Dataset	EB + AD	Fast-dLLM + AD	Top-k + AD	Average over methods
GSM8K	+9.22 / +14.31	+6.92 / +6.36	+11.90 / +5.88	+9.35 / +8.85
MATH500	+6.49 / +4.84	+5.70 / +2.15	+10.73 / +6.93	+7.64 / +4.64
HumanEval	+10.45 / +15.63	+6.90 / +12.03	+9.55 / +10.77	+8.97 / +12.81
MBPP	+10.37 / +15.01	+7.39 / +12.04	+13.67 / +19.60	+10.48 / +15.55
Average over datasets	+9.13 / +12.45	+6.73 / +8.14	+11.46 / +10.80	+9.11 / +10.46

These gains start from baseline operating points averaging at least four revealed tokens per step and linearly interpolate the corresponding ADAS curve at matched NFE; points outside its interpolation range are excluded. Adaptive rules are not compared directly at identical thresholds, and Fast-dLLM+AD additionally evaluates \(f=20\) to cover higher parallelism. Entries can therefore be interpolated estimates rather than pairs of raw runs with exactly identical NFE.

The next excerpt from appendix Tables 18 and 19 provides directly measured fixed Top-8 results without interpolation.

Model	Dataset / metric	NFE (identical for both)	Top-8 (%)	Top-8 + AD (%)
LLaDA-8B-Base	GSM8K / accuracy	32.00	22.67	42.46
LLaDA-8B-Base	HumanEval / Pass@1	64.00	5.49	20.12
Dream-7B-Base	GSM8K / accuracy	32.00	13.80	28.81
Dream-7B-Base	HumanEval / Pass@1	64.00	12.20	21.95

Ablation Study¶

The following results come from appendix Tables 9 and 10, all using LLaDA-8B-Base, HumanEval, and Top-8, with Pass@1 as the metric. The average-uncertainty variant replaces selected-token-specific uncertainty rather than removing the attention penalty entirely.

Config	Pass@1 (%)	Note
Confidence only, no ADAS	5.49	Original Top-8
First-layer attention	9.76	Averaged across attention heads
Middle-layer attention	15.24	Averaged across attention heads
Final-layer attention + selected-token uncertainty	20.12	Default ADAS
Final-layer attention + average masked-position uncertainty	17.07	Global average replaces selected-token-specific weighting

Key Findings¶

Final-layer attention outperforms first- and middle-layer attention in this ablation setting; selected-token-specific uncertainty beats global-average weighting by 3.05 percentage points. This is evidence from one model, task, and parallelism level, not a universal layer-selection rule.
Gains are not universal: among 90 matched-NFE operating points, 80 improve and 10 regress. The point-weighted mean gain is 9.27, with a 95% bootstrap interval of [7.74, 10.84]. Its weighting differs from the task/method-grouped averages in Table 2, so these are not the same statistic.
Direct measurements also show counterexamples: GSM8K Top-4 falls from 53.90% to 50.11% on LLaDA and from 48.45% to 46.47% on Dream, with NFE 64 in each comparison. The main value is mitigating collapse under high parallelism, not recovering full token-by-token generation quality.

Highlights & Insights¶

Separating which tokens to select from how many to select reuses existing budget controls. Adaptive samplers can still change actual NFE by selecting different members, so modularity does not automatically imply equal computation.
Uncertainty belongs to the selected token on which the candidate depends, matching the source of the perturbation. Candidate confidence alone or unweighted attention misses this source of risk.
Continuous discounts allow strongly dependent but valuable candidates to remain admissible, unlike hard conflict graphs. This flexibility comes with a heuristic proxy rather than an independence proof or global optimization guarantee.

Limitations & Future Work¶

Attention may encode syntax, position, or formatting rather than harmful dependence, and weak attention does not guarantee joint compatibility. More targeted head/layer selection and higher-order interactions need testing.
Evidence covers only two base models and mathematical/code tasks, not open-ended generation, larger models, or different confidence calibration. Selecting the shared strength on HumanEval also limits extrapolation of the no-retuning claim.
Commitments are irreversible. Remasking or revision could repair early errors but would change decoding dynamics and NFE, requiring new compute-matched comparisons.
Deterministic temperature-zero runs provide no repeated-seed variance estimate. Operating-point bootstraps describe cross-configuration stability, not paired per-example significance at a fixed setting; sensitivity to different few-shot orderings remains untested.

vs EB-Sampler / Fast-dLLM: these supply admissibility and stopping rules; ADAS supplies candidate reranking conditioned on selection history. It is an insertable improvement, not a replacement for every caching, scheduling, or stopping mechanism.
vs DAPD: DAPD builds a hard dependency graph from symmetrized, thresholded attention and constructs parallel batches through independent sets and graph coloring. ADAS preserves directional continuous attention and selects within the base rule using soft marginal penalties.
vs KLASS: KLASS combines confidence with temporal KL stability between consecutive denoising distributions; ADAS uses a same-step conditional-influence proxy from candidates to positions selected in the current step. They address temporal change and intra-group dependence, respectively.
vs learned planners / lookahead search: these can learn orders or explicitly assess future trajectories. ADAS requires neither training nor search and is lighter to deploy, but cannot assess long-term value beyond its attention proxy.

Rating¶

Novelty: 4/5 — Directional soft discounts weighted by selected-token uncertainty are integrated into existing stopping rules with a clear contribution boundary.
Experimental Thoroughness: 4/5 — Two models, four tasks, three rules, and intervention diagnostics provide substantial evidence; external tasks and prompt variance remain missing.
Writing Quality: 4/5 — The algorithm and lack of global guarantees are clear; code-status statements retain a historical difference between abstract and appendix.
Value: 4/5 — A lightweight inference improvement for highly parallel masked language models, with gains requiring compute-matched and regime-specific interpretation.