提出软掩码机制,让语言扩散模型更高效地利用预测信息。
Soft-Masked Diffusion Language Models
- 用动态混合掩码与预测词嵌入,替代传统二元掩码决策。
- 169M模型训练后困惑度和MAUVE分数优于基线,编码任务表现提升。
- 适用于需快速生成的场景,尤其适合高吞吐量代码生成任务。
扩散模型在语言建模中展现出超越传统自回归方法的潜力,支持并行生成与内置自校正。现有基于扩散的语言模型多采用掩码扩散,解码时对每个掩码项进行二元判断:保留或替换为预测词。但此过程在保留掩码时丢弃了预测信息。为此,本文提出软掩码(SM),在保留掩码时动态融合该掩码嵌入与前一步预测的top-k词嵌入,提供更丰富的先验信息,使部分上下文能跨步传播。我们设计了一种高效训练方法,可将SM引入掩码扩散语言模型。实验表明,从头训练169M参数模型使用SM后,困惑度与MAUVE得分均优于二元掩码基线;预训练模型通过持续预训练也可获得性能提升。进一步微调两个先进扩散模型Dream-7B与Dream-Coder-7B,SM在多个编码基准上持续提升性能,尤其在高吞吐设置下效果显著。代码已开源。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。