修正均匀扩散模型的去噪机制,提升文本生成质量
Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation
- 提出留一法去噪器,避免自污染预测
- 新采样方法使生成效果优于传统均匀扩散
- 改造模型结构,让采样更像掩码扩散模型
离散扩散模型常通过纯净数据预测进行训练,但反向过程的定义方式影响其性能。在掩码扩散模型(MDM)中不同选择趋于一致,而在均匀扩散模型(UDM)中则不一致。本文发现,标准插值桥接参数化并非由去噪后验优化,而是由一种留一法后验优化——即预测每个干净标记时不使用其自身噪声观测。这揭示了插值ELBO与常规交叉熵去噪目标之间的不匹配。我们刻画了留一法目标,并推导出去噪器、留一法后验与梯度之间的精确转换关系。这些转换使参数化与训练目标解耦。由此可设计无需额外训练的预测-校正采样器和基于留一法预测器的改进温度采样。进一步提出吸收态重构:在保持原联合分布的前提下,将均匀扩散分解为类似掩码扩散的采样操作,具有更简单的去噪后验、连续解掩码和自然重掩码机制。在语言建模任务上,留一法参数化始终提升生成效果,吸收构造表现达到或超越掩码扩散。结果表明,掩码与均匀扩散间的实际差距更多源于参数化与采样设计,而非边际分布本身的选择。代码与模型见https://github.com/samsongourevitch/rev_udm。
原文摘要 · Abstract (English)
Discrete diffusion models are often trained through clean-data prediction, but the prediction can be used in different ways to define the reverse dynamics. In Masked Diffusion Models (MDM) these choices largely coincide, whereas in Uniform Diffusion Models (UDM) they do not. We show that the standard plug-in bridge parameterization for UDM is not optimized by the denoising posterior, but by a leave-one-out posterior that predicts each clean token without using its own noisy observation. This identifies a mismatch between the plug-in ELBO and the usual cross-entropy denoising objective. We characterize the leave-one-out target and derive exact conversions between the denoiser, the leave-one-out posterior, and the score. These conversions allow us to disentangle parameterization and training objective. Our results also lead to inference improvements without any additional training through an informed predictor-corrector sampler and improved temperature sampling based on the leave-one-out predictor. We further introduce an absorbing-state reformulation of uniform diffusion that preserves the UDM joint law while decomposing it into masked-diffusion-like sampling operations, with simpler denoising posteriors, carry-over unmasking, and a natural remasking mechanism. On language modeling, leave-one-out parameterizations consistently improve UDM generation, while the absorbing construction matches or surpasses masked diffusion. These results suggest that the empirical gap between masked and uniform diffusion is driven less by the choice of marginals themselves than by parameterization and sampling design. The code and models can be found at https://github.com/samsongourevitch/rev_udm.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。