重访均匀扩散模型:留一去噪器与吸收态重构
阅读原文· arxiv.org研究揭示均匀扩散模型(UDM)的标准参数化与训练目标存在失配。提出“留一去噪器”,即在预测干净token时不依赖其自身噪声观测的后验,并推导了其与标准去噪器、分数函数的精确转换关系。进一步通过“吸收态重构”,将UDM联合分布分解为类掩码扩散操作,从而简化了去噪后验。在语言建模中,留一参数化稳定提升了UDM生成效果,吸收态构建匹配或超越了掩码扩散模型。实验表明,经验差距主要源于参数化与采样设计,而非边际分布选择本身。
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.