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

TL;DR

This paper introduces Gt-Margin, a ground-truth-guided scoring method for learning optimal unmasking orders in Masked Diffusion Language Models, significantly improving reasoning tasks without altering the core model.

Contribution

It proposes Gt-Margin as an oracle-based scoring method and trains a supervised planner to learn unmasking orders, enhancing reasoning performance in MDLMs.

Findings

Gt-Margin provides an effective oracle unmasking order.

Using the learned planner improves reasoning accuracy.

The approach enhances generation quality on logical reasoning benchmarks.

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.