论文精选

Uniform Diffusion Models 重访:Leave-One-Out 去噪器与吸收态重构

Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation

精选理由

这篇论文揭示了均匀扩散模型长期被忽视的参数化问题,做离散扩散模型研究的开发者可以直接用留一法改进推理,无需重新训练。看完会对掩码 vs 均匀扩散的差距有全新理解。

AI 摘要

本文重新审视了均匀扩散模型(UDM)的训练与推理。研究发现,UDM 的标准参数化并未优化去噪后验,而是优化了一种留一法后验,即预测每个干净 token 时不使用其自身的噪声观测。作者推导了去噪器、留一法后验与分数之间的精确转换,从而分离了参数化与训练目标。基于留一法预测器,他们提出了无需额外训练即可改进推理的预测-校正采样器和温度采样方法。此外,还引入了一种吸收态重构,将 UDM 分解为类似掩码扩散的采样操作,简化了去噪后验。实验表明,留一法参数化持续提升 UDM 生成质量,吸收态构造则匹配或超越掩码扩散,暗示掩码与均匀扩散之间的经验差距更多源于参数化和采样设计,而非边际分布本身。

原文 · arXiv cs.LG

Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation

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.