Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models

Abstract: Masked diffusion LLMs 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 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-(k), Fast-dLLM, and EB-Sampler 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. These results show that soft attention-discounted reranking is a simple and modular way to improve quality in highly parallel decoding for masked diffusion LLMs.

• The paper introduces ADAS, a training-free greedy reranking method that discounts candidates coupled to uncertain selected tokens using final-layer self-attention while preserving existing sampler stopping rules.
• Experiments on LLaDA-8B and Dream-7B across GSM8K, MATH500, HumanEval, and MBPP show average matched-NFE gains of 9.11 and 10.46 percentage points, respectively, with improvements in 80 of 90 operating points.
• ADAS adds only 3.1% average forward-pass overhead, but its attention-based dependency signal remains an approximation that may require tuning across models, tasks, and parallelism levels.

Motivation and problem statement

Masked diffusion LLMs (MDLMs) generate text by iteratively denoising a masked sequence, revealing multiple positions per step. This parallelism offers a distinct speed–quality tradeoff relative to causal autoregressive decoding, but it introduces a failure mode that token-wise scoring cannot address: positions that are individually confident may be unsafe to commit together when their predictions are coupled. Existing training-free samplers—Top-$k$, Fast-dLLM, and EB-Sampler—differ mainly in their stopping rules (fixed cardinality, confidence threshold, entropy budget), yet all rank candidates by token-wise quantities such as confidence, margin, or entropy, implicitly treating high-scoring candidates as mutually compatible.

The paper grounds this concern in the error decomposition of Ben-Hamu et al., which separates unmasking error into a model error term (per-position KL divergence between true and predicted conditionals) and a joint dependence error term (the multi-information of the true conditional distribution over the selected subset). The latter is model-independent: good marginals do not suffice when several positions are committed simultaneously. The authors' contribution is to make subset construction dependency-aware while leaving stopping rules intact.

Controlled dependency diagnostics

Before proposing the method, the paper validates two assumptions on synthetic arithmetic predicates with known dependency structure ($A+B+C=D$, $A\cdot B\cdot C=D$, $\min\{A,B,C\}=D$, $\max\{A,B,C\}=D$). First, forcing Dream-7B-Base to predict two coupled numeric literals in a single forward pass drops accuracy from 71% to 30% relative to predicting literals from separate predicates—a direct demonstration that joint commitment is substantially harder than marginal confidence suggests. Second, self-attention tracks this structure: mean pairwise attention for dependent masked-token pairs exceeds that of non-dependent pairs on both LLaDA-8B-Base (0.007261 vs. 0.001919) and Dream-7B-Base (0.009517 vs. 0.004842). These diagnostics support attention as a practical proxy for coupling, without claiming it is a perfect dependence measure.

Method

ADAS (Attention-Discounted Adaptive Sampler) is a training-free reranking rule that replaces additive token-wise subset utility with a dependency-aware greedy construction. Given token confidences $c_i$ and an attention matrix $A$ over masked positions (obtained from the denoiser's final self-attention layer, averaged across heads), ADAS scores a remaining candidate $i$ given an already selected set $S$ by the marginal utility

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

where $A+B+C=D$0 controls discount strength. A candidate is penalized not merely for being coupled to a selected position, but specifically for being coupled to one whose commitment is uncertain. Greedy selection proceeds until adding the top candidate violates the base sampler's admissibility rule; the stopping logic of Top-$A+B+C=D$1, Fast-dLLM, or EB-Sampler is untouched.

Two design choices distinguish ADAS from prior work. Relative to DAPD, which symmetrizes and thresholds attention into a binary graph and selects independent sets via graph coloring, ADAS keeps attention continuous as a soft penalty—a strongly coupled candidate can still be selected if its confidence gain outweighs the estimated risk. This matters because attention is an imperfect proxy: not every high-attention pair should be forbidden. Relative to learned unmasking policies, lookahead search, and planner-based methods, ADAS requires no training, verifier, or search, and adds only $A+B+C=D$2 selection cost per step.

The paper also provides a local theoretical motivation: under a first-order expansion of the denoiser logits around the current state, using empirical attention as a proxy for logit sensitivity, the incremental dependence risk of adding candidate $A+B+C=D$3 after $A+B+C=D$4 is bounded by a quantity proportional to $A+B+C=D$5, which is exactly the ADAS discount. The authors are explicit that this is a local approximation, not a global guarantee—the utility is not submodular in general and no approximation guarantee is claimed.

Empirical results

Evaluations cover LLaDA-8B-Base and Dream-7B-Base on GSM8K, MATH500, HumanEval, and MBPP, with three ADAS variants (Top-$A+B+C=D$6+AD, Fast-dLLM+AD, EB+AD) compared against their baselines over swept parallelism levels, plus KLASS and DAPD as additional training-free baselines. Matched-NFE gains in the low-NFE regime (baseline points averaging at least four tokens per step) are summarized below:

Dataset	EB	Fast-dLLM	Top-$A+B+C=D$7	Avg.
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
Avg.	+9.13 / +12.45	+6.73 / +8.14	+11.46 / +10.80	+9.11 / +10.46

(Entries report LLaDA / Dream gains in percentage points.) Several observations follow directly from these numbers. Gains are consistent across entropy-, confidence-, and fixed-budget stopping rules, indicating the benefit comes from subset construction rather than any particular criterion. Gains are largest on code generation (MBPP up to +19.60 points) where coupled simultaneous commitments most easily produce inconsistent outputs, and smallest on MATH500 (+4.64 on Dream), which the authors attribute to errors governed by longer sequential chains rather than within-step pairwise incompatibilities. Bootstrap analysis over 90 matched operating points yields mean gain +9.27 with 95% CI [+7.74, +10.84], positive in 80 of 90 cases. Against KLASS and DAPD, ADAS-augmented decoders remain competitive or superior at moderate-to-high parallelism, with both baselines degrading more sharply; the authors suggest temporal-stability signals become unreliable at high parallelism and hard attention-derived constraints can be brittle.

Runtime overhead is modest: disabling FlashAttention only in the final layer to extract attention costs 3.1% wall-clock time per model forward on average.

Ablations

Three ablations support specific design choices. For the discount strength $A+B+C=D$8, a sweep on LLaDA HumanEval shows regime-dependent optima, but a single global default of $A+B+C=D$9 loses only 0.72 points on average versus a per-setting oracle (at most 2.44 in the worst case), avoiding task-specific retuning. For the uncertainty weighting, replacing selected-token uncertainty $A\cdot B\cdot C=D$0 with a global masked-position uncertainty $A\cdot B\cdot C=D$1 costs 3.05 Pass@1 points, confirming that conditioning on the specific selected position's uncertainty matters. For the attention source, final-layer attention outperforms middle-layer (15.24 vs. 20.12 Pass@1) and first-layer (9.76) alternatives at $A\cdot B\cdot C=D$2.

Limitations and open questions

The paper states its limitations plainly. Self-attention remains only a proxy for stochastic dependence: strong attention does not always imply harmful joint updates, nor does weak attention guarantee compatibility. Evaluation covers two open-weight base MDLMs and four reasoning/code benchmarks; broader dependency signals, samplers, model scales, and task domains remain untested. The robustness intervals summarize variation across matched operating points rather than repeated seeds or per-example uncertainty, so they should be read as configuration-level summaries rather than paired significance tests. The theoretical grounding relies on a first-order local approximation whose constants are absorbed into $A\cdot B\cdot C=D$3. Open questions include whether head-specific or multi-layer attention aggregation improves the dependency signal, and how soft discounts compose with learned planners, lookahead methods, or auxiliary verifiers.

Conclusion

ADAS demonstrates that a minimal intervention—attention-discounted greedy reranking inside existing samplers' subset construction, with stopping rules unchanged—consistently improves low-NFE decoding quality for masked diffusion LLMs, yielding roughly 9–10 percentage-point average matched-NFE gains at 3.1% runtime overhead. The result indicates that improving highly parallel diffusion decoding requires better objectives for composing each revealed subset, not only better stopping rules.