arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2602.09501cs.CL
收藏

Where-to-Unmask: 基于真实情况引导的去掩码顺序学习用于掩码扩散语言模型

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

The University of Tokyo, Tokyo, Japan(东京大学)

Hikaru Asano, Tadashi Kozuno, Kuniaki Saito, Yukino Baba

AI总结:

本文提出Gt-Margin方法,通过真实标记计算去掩码顺序,提升MDLM在逻辑推理任务中的生成质量。

摘要:

掩码扩散语言模型(MDLMs)通过迭代填充掩码标记生成文本,每一步都需要两个耦合的决定:哪些位置去掩码(where-to-unmask)和哪些标记放置(what-to-unmask)。虽然标准MDLM训练直接优化标记预测(what-to-unmask),但推理时的去掩码顺序(where-to-unmask)通常由启发式置信度度量或通过具有成本的策略优化强化学习确定。为了解决这个问题,我们引入了Gt-Margin,一种从真实标记派生的位置评分,定义为正确标记与其最强替代品之间的概率边际。Gt-Margin产生了一个理想化的去掩码顺序,该顺序在每个部分掩码状态下优先处理更简单的位置。我们证明,利用这种理想化的去掩码顺序显著提高了最终生成质量,尤其是在逻辑推理基准上。基于这一见解,我们通过学习到排序来训练一个监督的去掩码规划器,以模仿从掩码上下文中获得的理想顺序。所得到的规划器整合到标准MDLM采样中以选择where-to-unmask,提高了推理准确性,而无需修改标记预测模型。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑