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

📄 arXiv: 2605.22765v1 📥 PDF

作者: Samson Gourevitch, Yazid Janati, Dario Shariatian, Umut Simsekli, Eric Moulines, Eric P. Xing, Alain Durmus

分类: cs.LG, stat.ML

发布日期: 2026-05-21

备注: preprint

🔗 代码/项目: GITHUB

💡 一句话要点

改进均匀扩散模型：提出Leave-One-Out去噪器和吸收态重构

🎯 匹配领域: 支柱四：生成式动作 (Generative Motion)

关键词: 均匀扩散模型 Leave-One-Out 去噪 吸收态重构 语言建模 文本生成 反向扩散 参数化

📋 核心要点

  1. 均匀扩散模型（UDM）在反向过程定义上存在多种选择，标准plug-in参数化与去噪目标不匹配。
  2. 论文提出leave-one-out后验作为优化目标，并推导出denoiser、后验和score之间的转换关系，解耦参数化和训练目标。
  3. 通过leave-one-out参数化和吸收态重构，UDM在语言建模任务上取得了显著的生成性能提升，缩小了与MDM的差距。

📝 摘要（中文）

离散扩散模型通常通过clean-data预测进行训练，但预测可以以不同的方式用于定义反向动力学。在Masked Diffusion Models (MDM)中，这些选择很大程度上是一致的，而在Uniform Diffusion Models (UDM)中则不然。我们表明，UDM的标准plug-in桥参数化不是由去噪后验优化的，而是由leave-one-out后验优化的，该后验在不使用自身噪声观测的情况下预测每个clean token。这确定了plug-in ELBO和通常的交叉熵去噪目标之间的不匹配。我们描述了leave-one-out目标，并推导出denoiser、leave-one-out后验和score之间的精确转换。这些转换使我们能够解耦参数化和训练目标。我们的结果还通过知情的predictor-corrector采样器和基于leave-one-out预测器的改进温度采样，实现了无需额外训练的推理改进。此外，我们引入了均匀扩散的吸收态重构，它保留了UDM联合律，同时将其分解为类似masked-diffusion的采样操作，具有更简单的去噪后验、carry-over unmasking和自然的remasking机制。在语言建模方面，leave-one-out参数化始终改进UDM生成，而吸收构造匹配或超过masked diffusion。这些结果表明，masked diffusion和uniform diffusion之间的经验差距与其说是边缘分布本身的选择驱动的，不如说是参数化和采样设计驱动的。代码和模型可在https://github.com/samsongourevitch/rev_udm 找到。

🔬 方法详解

问题定义：均匀扩散模型（UDM）在定义反向扩散过程时，存在多种不同的参数化选择。现有的标准plug-in参数化方法，其优化目标与实际的去噪后验并不一致，导致训练和推理之间存在偏差，影响生成质量。此外，UDM的采样过程相对复杂，限制了其在实际应用中的效率和灵活性。

核心思路：论文的核心思路是重新审视UDM的训练目标，并提出基于leave-one-out后验的优化方法。Leave-one-out后验的目标是在不依赖自身噪声观测的情况下预测clean token，从而更准确地反映了去噪过程的本质。此外，论文还引入了吸收态重构，将UDM分解为更简单的、类似masked diffusion的采样操作，简化了采样过程。

技术框架：论文的技术框架主要包含两个方面：1) Leave-one-out去噪器：推导了denoiser、leave-one-out后验和score之间的精确转换公式，使得可以使用leave-one-out后验作为训练目标，并进行更有效的推理。2) 吸收态重构：将UDM的联合概率分布分解为一系列masked diffusion风格的采样步骤，包括carry-over unmasking和remasking机制，简化了采样过程，并使其更易于实现。

关键创新：论文最重要的技术创新在于发现了UDM标准plug-in参数化与去噪后验之间的不匹配，并提出了基于leave-one-out后验的优化方法。这种方法能够更准确地反映去噪过程的本质，从而提高生成质量。此外，吸收态重构也简化了UDM的采样过程，使其更易于实现和应用。

关键设计：在leave-one-out去噪器方面，论文推导了denoiser、leave-one-out后验和score之间的精确转换公式，这些公式是实现leave-one-out训练的关键。在吸收态重构方面，论文设计了carry-over unmasking和remasking机制，使得采样过程更加灵活和高效。具体的网络结构和损失函数细节可能与具体的应用场景有关，论文中可能没有详细描述。

📊 实验亮点

实验结果表明，基于leave-one-out参数化的UDM在语言建模任务上取得了显著的性能提升，尤其是在生成质量方面。吸收态重构后的UDM性能可以匹配甚至超过masked diffusion模型，表明参数化和采样设计对扩散模型性能的影响大于边缘分布的选择。

🎯 应用场景

该研究成果可应用于自然语言处理领域的文本生成、语言建模等任务。通过改进均匀扩散模型的训练和采样方法，可以生成更高质量、更自然的文本。此外，该研究的思路也可以推广到其他扩散模型和生成模型中，具有广泛的应用前景。

📄 摘要（原文）

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.