3 个月前

简单而有效的掩码扩散语言模型

Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, Volodymyr Kuleshov
简单而有效的掩码扩散语言模型
摘要

尽管扩散模型在生成高质量图像方面表现出色,但先前的研究表明,其在语言建模任务中与自回归(AR)方法之间存在显著性能差距。在本工作中,我们发现,简单的掩码离散扩散模型的性能实际上远超以往认知。我们提出了一种高效的训练策略,显著提升了掩码扩散模型的性能,并推导出一种简化的、基于Rao-Blackwell化的目标函数,进一步带来了性能提升。该目标函数形式简洁——本质上是经典掩码语言建模损失的混合形式——可应用于训练仅含编码器的语言模型,且支持高效采样器,包括能够以半自回归方式生成任意长度文本的采样器,其行为类似于传统语言模型。在语言建模基准测试中,采用现代工程实践训练的一系列掩码扩散模型,达到了扩散模型领域的最新技术水平,并逼近了自回归模型的困惑度(perplexity)表现。相关代码已开源,地址为:https://github.com/kuleshov-group/mdlm。