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

通过信息增益改进掩码扩散模型的采样

Improving Sampling for Masked Diffusion Models via Information Gain

Department of Computer Science and Technology, Tsinghua University(清华大学计算机科学与技术系);Massachusetts Institute of Technology(麻省理工学院);Shanghai Jiao Tong University(上海交通大学);Beihang University(北航);College of AI, Tsinghua University(清华大学人工智能学院)

Kaisen Yang, Jayden Teoh, Kaicheng Yang, Yitong Zhang, Alex Lamb

AI总结:

提出一种无需训练的Info-Gain采样器,利用掩码扩散模型的双向结构平衡即时不确定性与剩余掩码位置的信息增益,在推理、编码、创意写作和图像生成任务中优于现有采样器。

摘要:

掩码扩散模型(MDMs)支持灵活的解码顺序,但现有采样器大多是贪婪的,仅选择局部确定的token而不考虑其下游影响。我们表明这种短视行为会增加累积不确定性并导致次优生成。为解决此问题,我们提出**Info-Gain采样器**,一种无需训练的解码方法,利用MDMs的双向结构平衡即时不确定性与剩余掩码位置获得的信息增益。在推理、编码、创意写作和图像生成任务中,Info-Gain采样器持续优于现有MDM采样器,平均推理准确率提升2.9--11.6个百分点,创意写作平均胜率达到62.8%。代码可在https://github.com/yks23/Information-Gain-Sampler获取。

英文摘要:

Masked Diffusion Models (MDMs) enable flexible decoding orders, yet existing samplers remain largely greedy, selecting locally certain tokens without accounting for their downstream effects. We show that this myopia can increase cumulative uncertainty and lead to suboptimal generation. To address this, we propose the **Info-Gain Sampler**, a training-free decoding method that uses the bidirectional structure of MDMs to balance immediate uncertainty with the information gained over remaining masked positions. Across reasoning, coding, creative writing, and image generation tasks, Info-Gain Sampler consistently outperforms existing MDM samplers, improving average reasoning accuracy by 2.9--11.6 percentage points and achieving a 62.8% average win rate in creative writing. The code is available at https://github.com/yks23/Information-Gain-Sampler.

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

↑