发表机构
The Hong Kong Polytechnic University; Southern University of Science and Technology(香港理工大学; 南方科技大学)
机构由 AI 辅助整理,请以论文原文为准。AI 中文总结
针对掩码扩散语言模型中不可逆标记提交导致的错误传播问题,提出无训练、采样器无关的CoDR方法,通过k次前向传递估计置信度漂移并重掩码不认可标记,提升多任务平均准确率。
AI 中文摘要
掩码扩散语言模型(MDLMs)通过反复将标记提交到掩码位置来进行解码,但这些提交通常是不可逆的。在稀疏、部分上下文中选择的标记会被固定保留,即使后续上下文不再支持它。现有的采样器主要决定何时提交标记,但很少检查已提交的标记是否应被保留,导致早期错误得以传播。我们将此问题追溯到置信度漂移,即模型对已提交标记的置信度从稀疏的提交时上下文下降到后续可用的更密集上下文。基于此信号,我们提出了CoDR(置信度漂移重掩码),一种无需训练且与采样器无关的细化过程。CoDR通过k分区探测,仅需k次前向传递即可估计所有已提交位置的漂移,然后仅重掩码并重新生成模型不再认可的标记。在两个骨干网络、四个推理和编码任务以及三个基础采样器上,CoDR在所有评估的模型-采样器配置中提高了平均准确率,并在大多数单独任务设置中有所改进,且开销适中。受控实验表明,这些收益来自有针对性的置信度漂移重掩码,而非仅仅是额外的计算量,并且CoDR使用的前向传递次数远少于先前的重掩码方法。代码可在该https URL获取。
英文摘要
Masked diffusion language models (MDLMs) decode by repeatedly committing tokens to masked positions, but these commitments are usually irreversible. A token chosen under sparse, partial context is kept fixed, even when later context no longer supports it. Existing samplers mainly decide when to commit a token, but rarely check whether an already committed token should still be kept, allowing early mistakes to propagate. We trace this issue to confidence drift, where the model's confidence in a committed token drops from its sparse commit-time context to the denser context available later. Based on this signal, we propose CoDR (Confidence Drift Remasking), a training-free and sampler-agnostic refinement pass. CoDR estimates drift for all committed positions in only k forward passes via k-partition probing, then remasks and regenerates only the tokens the model no longer endorses. Across two backbones, four reasoning and coding tasks, and three base samplers, CoDR improves average accuracy across all evaluated model-sampler configurations and improves most individual task settings with modest overhead. Controlled experiments show that the gains come from targeted confidence-drift remasking rather than extra compute alone, and that CoDR uses far fewer forward passes than prior remasking methods. Code is available at https://github.com/YueWu0301/CoDR.
Comments17 pages, 6 figures, and 17 tables