Where-to-Unmask: Ground-Truth-Guided Unmasking Order Learning for Masked Diffusion Language Models

• URL: http://arxiv.org/abs/2602.09501v1
• Date: Tue, 10 Feb 2026 07:56:46 GMT
• Title: Where-to-Unmask: Ground-Truth-Guided Unmasking Order Learning for Masked Diffusion Language Models
• Authors: Hikaru Asano, Tadashi Kozuno, Kuniaki Saito, Yukino Baba,
• Abstract summary: 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.
• Score: 17.18632315520133
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Masked Diffusion Language Models (MDLMs) generate text by iteratively filling masked tokens, requiring two coupled decisions at each step: which positions to unmask (where-to-unmask) and which tokens to place (what-to-unmask). While standard MDLM training directly optimizes token prediction (what-to-unmask), inference-time unmasking orders (where-to-unmask) are typically determined by heuristic confidence measures or trained through reinforcement learning with costly on-policy rollouts. To address this, we introduce Gt-Margin, a position-wise score derived from ground-truth tokens, defined as the probability margin between the correct token and its strongest alternative. Gt-Margin yields an oracle unmasking order that prioritizes easier positions first under each partially masked state. We demonstrate that leveraging this oracle unmasking order significantly enhances final generation quality, particularly on logical reasoning benchmarks. Building on this insight, we train a supervised unmasking planner via learning-to-rank to imitate the oracle ordering from masked contexts. The resulting planner integrates into standard MDLM sampling to select where-to-unmask, improving reasoning accuracy without modifying the token prediction model.

Related papers

• 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)
• Lookahead Unmasking Elicits Accurate Decoding in Diffusion Language Models [51.12873073612084]
  Masked Diffusion Models (MDMs) as language models generate by iteratively unmasking tokens, yet their performance depends on the inference time order of unmasking.<n>We propose Lookahead Unmasking (LookUM), which addresses these concerns by reformulating sampling as path selection over all possible unmasking orders.<n>LookUM requires only two to three paths to achieve peak performance, demonstrating remarkably efficient path selection.
  arXiv  Detail & Related papers  (2025-11-04T02:37:37Z)
• 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)
• Don't Settle Too Early: Self-Reflective Remasking for Diffusion Language Models [40.902681492117786]
  RemeDi is a mask-based DLM that predicts token distributions and per-token confidence scores at each step.<n>We train a remask-aware pipeline to train this ability, including supervised fine-tuning which teaches the model to detect and remask incorrect tokens.<n>Experiments show that RemeDi achieves the state-of-the-art results among open-source DLMs on multiple datasets.
  arXiv  Detail & Related papers  (2025-09-28T05:39:49Z)
• 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)
• Bridge the Points: Graph-based Few-shot Segment Anything Semantically [79.1519244940518]
  Recent advancements in pre-training techniques have enhanced the capabilities of vision foundation models. Recent studies extend the SAM to Few-shot Semantic segmentation (FSS) We propose a simple yet effective approach based on graph analysis.
  arXiv  Detail & Related papers  (2024-10-09T15:02:28Z)
• MaskInversion: Localized Embeddings via Optimization of Explainability Maps [49.50785637749757]
  MaskInversion generates a context-aware embedding for a query image region specified by a mask at test time. It can be used for a broad range of tasks, including open-vocabulary class retrieval, referring expression comprehension, as well as for localized captioning and image generation.
  arXiv  Detail & Related papers  (2024-07-29T14:21:07Z)
• ColorMAE: Exploring data-independent masking strategies in Masked AutoEncoders [53.3185750528969]
  Masked AutoEncoders (MAE) have emerged as a robust self-supervised framework. We introduce a data-independent method, termed ColorMAE, which generates different binary mask patterns by filtering random noise. We demonstrate our strategy's superiority in downstream tasks compared to random masking.
  arXiv  Detail & Related papers  (2024-07-17T22:04:00Z)
• InforMask: Unsupervised Informative Masking for Language Model Pretraining [13.177839395411858]
  We propose a new unsupervised masking strategy for training masked language models. InforMask exploits Pointwise Mutual Information (PMI) to select the most informative tokens to mask.
  arXiv  Detail & Related papers  (2022-10-21T07:10:56Z)
• 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)

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.