Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models

• URL: http://arxiv.org/abs/2606.10829v1
• Date: Tue, 09 Jun 2026 13:17:27 GMT
• Title: Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models
• Abstract summary: Masked diffusion language models can reduce inference steps by revealing multiple tokens per denoising iteration.<n>We propose ADAS, a training-free reranking rule for parallel masked diffusion decoding.
• Score: 59.51249894128724
• License: http://creativecommons.org/licenses/by/4.0/
• 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-\(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 language models.

Related papers

• LESS Is More: Mutual-Stability Sampling for Diffusion Language Models [23.94639050546374]
  Diffusion large language models (dLLMs) offer a promising alternative to autoregressive decoding.<n>We present textscLESS, a training-free, model-agnostic adaptive sampler that treats token commitment as an online stopping problem.<n>textscLESS improves average accuracy over strong training-free adaptive samplers while using $72.1%$ fewer reverse steps than fixed-budget decoding.
  arXiv  Detail & Related papers  (2026-06-15T16:15:45Z)
• DSL-LLaDA: Scaling Continuous Denoising to 8B Masked Diffusion LMs [29.561307941602482]
  We show that a pretrained masked DLM can be lightly adapted to support continuous embedding-space denoising.<n>The adapted model supports continuous inference that evolves all positions jointly in embedding space.<n> DSL-LLaDA-SDE achieves the best ROUGE-1 on all four benchmarks.
  arXiv  Detail & Related papers  (2026-05-31T05:27:01Z)
• Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation [72.05592785529404]
  We show that the standard plug-in bridge parameterization for UDM is not optimized by the denoising posterior.<n>We also introduce an absorbing-state reformulation of uniform diffusion that preserves the UDM joint law while decomposing it into masked-diffusion-like sampling operations.
  arXiv  Detail & Related papers  (2026-05-21T17:27:19Z)
• Dependency-Guided Parallel Decoding in Discrete Diffusion Language Models [21.699371484195865]
  We propose DEMASK, a lightweight dependency predictor that attaches to the final hidden states of a dLLM.<n> Empirically, DEMASK achieves 1.7-2.2$times$ speedup on Dream-7B while matching or improving accuracy compared to confidence-based and KL-based baselines.
  arXiv  Detail & Related papers  (2026-04-02T22:21:24Z)
• CORE: Context-Robust Remasking for Diffusion Language Models [51.59514489363897]
  We propose Context-Robust Remasking (CORE), a training-free framework for inference-time revision.<n>Rather than trusting static token probabilities, CORE identifies context-brittle tokens by probing their sensitivity to targeted masked-context perturbations.<n>On LLaDA-8B-Base, CORE delivers consistent improvements across reasoning and code benchmarks, outperforming compute-matched baselines and improving MBPP by up to 9.2 percentage points.
  arXiv  Detail & Related papers  (2026-02-04T00:12:30Z)
• Self-Rewarding Sequential Monte Carlo for Masked Diffusion Language Models [58.946955321428845]
  This work presents self-rewarding sequential Monte Carlo (SMC)<n>Our algorithm stems from the observation that most existing MDLMs rely on a confidence-based sampling strategy.<n>We introduce the trajectory-level confidence as a self-rewarding signal for assigning particle importance weights.
  arXiv  Detail & Related papers  (2026-02-02T09:21:45Z)
• Learning Unmasking Policies for Diffusion Language Models [33.44995119635116]
  Language Models (dLLMs) now match the downstream performance of their autoregressive counterparts on many tasks.<n>One particularly successful variant is masked discrete diffusion, in which a buffer filled with special mask tokens is progressively replaced with tokens sampled from the model's vocabulary.<n>In this work, we propose to train sampling procedures using reinforcement learning.
  arXiv  Detail & Related papers  (2025-12-09T20:44:33Z)
• MARS-Sep: Multimodal-Aligned Reinforced Sound Separation [72.85468563236005]
  MARS-Sep is a reinforcement learning framework for sound separation.<n>It learns a factorized Beta mask policy that is optimized by a clipped trust-region surrogate.<n>Experiments on multiple benchmarks demonstrate consistent gains in Text-, Audio-, and Image-Queried separation.
  arXiv  Detail & Related papers  (2025-10-12T09:05:28Z)
• Plan for Speed: Dilated Scheduling for Masked Diffusion Language Models [13.575063025878208]
  Masked diffusion language models promise fast, non-autoregressive text generation.<n>Existing samplers, which pick tokens to unmask based on model confidence, ignore interactions when unmasking multiple positions in parallel.<n>We propose the Dilated Unmasking Scheduler (DUS), an inference-only, planner-model-free method that partitions sequence positions into non-adjacent dilated groups and unmasked them in parallel.
  arXiv  Detail & Related papers  (2025-06-23T18:49:23Z)
• Accelerating Diffusion LLMs via Adaptive Parallel Decoding [60.407727995313074]
  We introduce adaptive parallel decoding (APD), a novel method that dynamically adjusts the number of tokens sampled in parallel.<n>APD provides markedly higher throughput with minimal quality degradations on downstream benchmarks.
  arXiv  Detail & Related papers  (2025-05-31T06:10:10Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.