软掩码扩散语言模型
Soft-Masked Diffusion Language Models
AI总结:
本文提出软掩码方法,通过动态融合掩码token与预测token的嵌入,提升扩散语言模型的生成性能和准确性。
摘要:
扩散模型在语言建模中展现出了强大的潜力,相较于传统自回归方法具有多种优势。其能够并行生成和修改整个响应的能力,使生成过程更快,并内置了自我校正机制。大多数现代基于扩散的语言模型采用掩码扩散,其中解码过程涉及迭代处理被掩码的token,基于二进制决策:要么保留掩码,要么将其替换为预测的token。然而,这种二进制选择在保留掩码时会丢弃有价值的预测信息。为了解决这一限制,我们引入了软掩码(SM),一种新的方法,它动态地将掩码token的嵌入与前一步解码中top-k预测token的嵌入混合,对于每个保留的掩码。这为模型提供了更有信息量的先验知识,保留了早期计算中的上下文,并允许被掩码token的部分信息在单步之外传播。我们提出了一种训练方法,该方法能够高效地将掩码扩散语言模型适应到包含SM的模型中。我们证明,使用SM训练从头开始的169M参数模型,在困惑度和MAUVE分数上优于二进制掩码基线。同样,通过继续预训练,预训练模型也可以通过SM得到增强。最后,我们对两个最先进的扩散模型,Dream-7B和Dream-Coder-7B,使用SM进行微调。SM在多个编码基准测试中一致地提高了性能,特别是在高吞吐量设置中。代码可在https://github.com/IBM/soft-masked-diffusion-language-models上获得。
英文摘要:
Diffusion models have demonstrated strong potential in language modeling, offering various advantages over traditional autoregressive approaches. Their ability to generate and revise entire responses in parallel enables faster generation and built-in self-correction mechanisms. Most modern diffusion-based language models employ masked diffusion, where decoding involves iteratively processing masked tokens based on a binary decision: either retaining the mask or replacing it with the predicted token. However, this binary choice discards valuable predictive information when the mask is retained. To address this limitation, we introduce soft-masking (SM), a novel method that dynamically blends the embedding of the mask token with the embeddings of the top-k predicted tokens from the previous decoding step, for each retained mask. This provides the model with a more informative prior, preserving context from earlier computations and allowing partial information about masked tokens to propagate beyond a single step. We propose a training methodology that efficiently adapts masked diffusion language models to incorporate SM. We demonstrate that training a 169M parameter model from scratch with SM yields superior perplexity and MAUVE scores compared to binary masking baselines. Similarly, a pretrained model can be enhanced with SM through continued pretraining. Finally, we finetune two state-of-the-art diffusion models, Dream-7B and Dream-Coder-7B, with SM. SM consistently improves performance across multiple coding benchmarks, particularly in high-throughput settings. The code is available at https://github.com/IBM/soft-masked-diffusion-language-models.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。