Mask-Aware Policy Gradients for Diffusion Language Models

• URL: http://arxiv.org/abs/2607.15200v1
• Date: Thu, 16 Jul 2026 16:57:34 GMT
• Title: Mask-Aware Policy Gradients for Diffusion Language Models
• Abstract summary: Reinforcement learning has proven effective for improving reasoning in large language models, but extending it to Masked Diffusion Language Models (MDLMs) remains challenging.<n>We observe that MDLM generation involves two decisions at each step: what tokens to place at each masked position and which positions to remask.<n>We formalize this as a two-stage action MDP, showing that the policy gradient naturally decomposes into a token term and a masking term.
• Score: 24.287374086170882
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Reinforcement learning has proven effective for improving reasoning in large language models, but extending it to Masked Diffusion Language Models (MDLMs) remains challenging due to the intractability of the log-likelihood estimation. Existing approaches approximate this log-likelihood by modeling only the token predictions, ignoring the order in which positions are unmasked during generation. We observe that MDLM generation involves two decisions at each step: what tokens to place at each masked position and which positions to remask. We formalize this as a two-stage action MDP, showing that the policy gradient naturally decomposes into a token term and a masking term. Combining optimization of both terms leads to state-of-the-art outcomes on mathematical reasoning and coding benchmarks, with scores of 87.1% on GSM8K and 53.4% on MBPP.

Related papers

• The Efficiency Gap in Byte Modeling [101.18202046105718]
  Two alternative paradigms have challenged the use of subword tokenization and autoregressive ordering.<n>We show that removing these structural priors incurs a significant computational cost.<n>Our findings suggest that future modality-agnostic designs must incorporate alternative structural biases to maintain viable scaling trajectories.
  arXiv  Detail & Related papers  (2026-05-13T03:03:30Z)
• LogicDiff: Logic-Guided Denoising Improves Reasoning in Masked Diffusion Language Models [0.0]
  We introduce LogicDiff, an inference-time method that replaces confidence-based unmasking with logic-role-guided unmasking.<n>A dependency-ordered scheduler then unmasks tokens in logical dependency order.<n> LogicDiff improves LLaDA-8B-Instruct accuracy from 22.0% to 60.7% on GSM8K and from 23.6% to 23.6% on MATH-500, with less than 6% speed overhead.
  arXiv  Detail & Related papers  (2026-03-24T13:08:10Z)
• Beyond Masks: Efficient, Flexible Diffusion Language Models via Deletion-Insertion Processes [41.57128726515332]
  We propose Deletion-Insertion Diffusion language models (DID)<n>DID rigorously formulates token deletion and insertion as discrete diffusion processes.<n>Our experiments across fixed and variable-length settings demonstrate the advantage of DID over baselines of MDLMs and existing insertion-based LMs.
  arXiv  Detail & Related papers  (2026-03-04T12:53:17Z)
• Where-to-Unmask: Ground-Truth-Guided Unmasking Order Learning for Masked Diffusion Language Models [17.18632315520133]
  Masked Diffusion Language Models generate text by iteratively filling masked tokens.<n>Gt-Margin yields an oracle unmasking order that prioritizes easier positions first under each partially masked state.<n>We train a supervised unmasking planner via learning-to-rank to imitate the oracle ordering from masked contexts.
  arXiv  Detail & Related papers  (2026-02-10T07:56:46Z)
• Masks Can Be Distracting: On Context Comprehension in Diffusion Language Models [19.847438086389616]
  Masked Diffusion Language Models have emerged as a promising alternative to Autoregressive Language Models.<n>We show that MDLMs exhibit a strong locality bias, favouring local over distant context.<n>We introduce a mask-agnostic loss function that encourages predictions to remain invariant to the number of appended masks.
  arXiv  Detail & Related papers  (2025-11-26T12:44:29Z)
• Soft-Masked Diffusion Language Models [35.191030145577145]
  We introduce soft-masking (SM), a novel method that dynamically blends the embedding of the mask token with the embeddings of the top-$k$ predicted tokens.<n>We demonstrate that continuing pretraining a 169M parameter model with SM leads to improved perplexity and MAUVE scores.<n>We finetune two state-of-the-art diffusion models, Dream-7B and Dream-Coder-7B, with SM.
  arXiv  Detail & Related papers  (2025-10-20T06:42:03Z)
• Exploring Gradient-Guided Masked Language Model to Detect Textual Adversarial Attacks [50.53590930588431]
  adversarial examples pose serious threats to natural language processing systems.<n>Recent studies suggest that adversarial texts deviate from the underlying manifold of normal texts, whereas masked language models can approximate the manifold of normal data.<n>We first introduce Masked Language Model-based Detection (MLMD), leveraging mask unmask operations of the masked language modeling (MLM) objective.
  arXiv  Detail & Related papers  (2025-04-08T14:10:57Z)
• Masked Autoencoding for Scalable and Generalizable Decision Making [93.84855114717062]
  MaskDP is a simple and scalable self-supervised pretraining method for reinforcement learning and behavioral cloning. We find that a MaskDP model gains the capability of zero-shot transfer to new BC tasks, such as single and multiple goal reaching.
  arXiv  Detail & Related papers  (2022-11-23T07:04:41Z)
• Improving Self-supervised Pre-training via a Fully-Explored Masked Language Model [57.77981008219654]
  Masked Language Model (MLM) framework has been widely adopted for self-supervised language pre-training. We propose a fully-explored masking strategy, where a text sequence is divided into a certain number of non-overlapping segments.
  arXiv  Detail & Related papers  (2020-10-12T21:28:14Z)
• UniLMv2: Pseudo-Masked Language Models for Unified Language Model Pre-Training [152.63467944568094]
  We propose to pre-train a unified language model for both autoencoding and partially autoregressive language modeling tasks. Our experiments show that the unified language models pre-trained using PMLM achieve new state-of-the-art results on a wide range of natural language understanding and generation tasks.
  arXiv  Detail & Related papers  (2020-02-28T15:28:49Z)
• Semi-Autoregressive Training Improves Mask-Predict Decoding [119.8412758943192]
  We introduce a new training method for conditional masked language models, SMART, which mimics the semi-autoregressive behavior of mask-predict. Models trained with SMART produce higher-quality translations when using mask-predict decoding, effectively closing the remaining performance gap with fully autoregressive models.
  arXiv  Detail & Related papers  (2020-01-23T19:56:35Z)

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.