What Makes a Good Unmasking Policy? A Controlled Study of Learned Unmasking in Diffusion Language Models

一句话概括

本文把监督式解掩码策略视为对解码过程的蒸馏，系统拆解“教师标注什么、在哪些状态上模仿、策略观察什么、开销需多少 NFE 下降来抵消”4 个设计维度，并提出以轨迹收缩为目标的策略。该策略在 LLaDA-8B-Instruct 上取得最佳精度–NFE 前沿。

速览

• 任务：研究掩码扩散语言模型（masked diffusion language model）中学习式解掩码策略（unmasking policy）的设计选择。
• 方法：将监督策略看作蒸馏一个解码过程，分别控制教师标签、模仿状态、策略输入与容量。
• 关键数字：约 10% 的 NFE（number of function evaluations）下降即可抵消策略开销，适用于不超过 7.5M 参数的策略；最终策略仅 856K 参数。
• 数据与开源：实验覆盖 GSM8K、MATH-500、HumanEval，并迁移到 Dream-7B-Instruct；开源情况摘要未给出。

问题背景

掩码扩散语言模型并行生成 token，效率取决于“提交哪些预测”。学习式解掩码策略能改善速度–精度权衡，但已有工作在监督信号、输入和容量上各不相同。差异混杂，难以判断哪些设计真正起作用。作者因此把问题拆成可独立检验的 4 个维度。

方法要点

• 视角：把监督策略视为通过教师蒸馏一个解码过程，而非孤立训练分类器。
• 教师标签：比较“与最终输出一致”的标签与“留在参考轨迹上”的标签。
• 模仿状态：考察策略在哪些解码状态上向教师学习。
• 策略输入：对比仅用置信度（confidence）、仅用隐藏状态（hidden state）及二者组合。
• 开销核算：量化策略自身计算开销需要多少 NFE 下降来抵消。
• 最终方案：以轨迹收缩（trajectory contraction）为目标，只接受留在参考轨迹上的提交，输入为置信度加隐藏状态。

关键结论

• “与最终输出一致”的标签在低 NFE 下过于宽松，会接受不该提交的预测。
• 隐藏状态单独使用较弱，但可补充置信度信息。
• 对不超过 7.5M 参数的策略，约 10% 的 NFE 下降足以抵消其开销（限于所测设置）。
• 在 LLaDA-8B-Instruct 上，856K 参数策略取得最佳精度–NFE 前沿，GSM8K 吞吐最高。
• 在 MATH-500 与 HumanEval 的中等目标下具竞争力，且大体可迁移到 Dream-7B-Instruct。

读后思考

本文价值在于把“策略好坏”还原为可对照的设计变量，而非堆叠新模块。轨迹收缩标签的思路值得注意：它把“可提交”定义为不偏离参考轨迹，比匹配最终输出更保守。10% NFE 阈值给出实用的开销预算，提醒研究者在扩大策略容量前先算账。小参数策略胜出，也说明输入质量比模型规模更关键。

局限与开放问题

• 摘要仅给出单一模型规模与 3 个基准，跨模型、跨任务的普适性仍需验证。
• 10% NFE 阈值依赖“所测设置”，其边界条件摘要未给出。
• 参考轨迹由谁产生、轨迹质量如何影响策略，摘要未给出。
• 策略容量上限 7.5M 之外的结论未知。
• 训练成本、数据来源与开源计划摘要未给出。