重新审视Uniform Diffusion Models：Leave-One-Out去噪器与吸收状态重构

TL;DR 重新审视均匀扩散模型，揭示标准训练目标与参数化的错配，提出基于 leave-one-out 的修正及吸收态重构，显著提升生成性能，证明性能差距源于参数化而非噪声类型。

问题

问题背景

离散扩散模型（discrete diffusion）已成为文本、蛋白质序列等离散数据生成的重要框架。通过逐步向干净数据注入噪声并学习逆过程去噪，这类模型在语言建模、图像分割等任务中展现出优势。其中，均匀扩散模型（UDM）和掩码扩散模型（MDM）是两种主流的噪声策略，但二者在参数化选择和训练目标上存在本质差异，直接影响生成质量。

现有方法局限

在 MDM 中，干净数据预测可以直接定义逆扩散过程，因为掩码操作天然对应“填补缺失值”的直观解释。然而在 UDM 中，均匀噪声在每个位置引入离散均匀分布，使得逆过程不再唯一确定。当前标准做法采用 plug-in bridge 参数化，即直接使用去噪器预测的干净 token 来构造逆转移概率。但本文发现，该参数化实际上由 leave-one-out 后验 优化，而非直观的 去噪后验。这意味着训练使用的交叉熵损失与 ELBO 之间存在不一致：模型被训练去预测“忽略当前噪声观测后”的干净 token，而推理时却依赖于含噪观测，导致训练-推理不匹配。此外，UDM 的普通采样器常受困于低效的采样路径和温度控制难题。

为什么这个问题难/重要

离散状态空间的无序性使得扩散过程的逆映射高度非凸，均匀噪声的对称性进一步放大了参数化选择的敏感性。确立正确的训练目标与参数化对应关系，是弥合理论与工程实践的关键。该问题直接影响 UDM 在语言建模等任务中的生成质量，也关系到能否设计出更高效的采样算法。业界对离散扩散模型的兴趣持续增长，因其在文本生成、代码补全等应用中，有望提供比自回归模型更灵活的生成方式和更优的多样性控制。

行业类比

这类似于在 BERT 掩码语言模型中，预测某个被 [MASK] 替换的词时，模型不能直接看到该词本身，而是基于上下文推断。在离散扩散中，leave-one-out 预测同样迫使模型学习上下文依赖，而均匀噪声环境使得这种“自隐式”训练成为提升生成一致性的关键。

核心洞察

• UDM 的标准 plug-in bridge 参数化隐式训练了一个“留一后验”而非真正的去噪后验，这是造成训练目标与 ELBO 不匹配的根本原因。与先前工作不同，该研究明确区分了参数化和训练目标，推导出 denoiser、留一后验和 score 之间的精确转换，揭示了常用交叉熵目标实际在优化一个不同条件分布。这一 mismatch 解释了 UDM 与 MDM 的性能差距，并使得无需额外训练即可通过 informed predictor-corrector 和温度采样提升推理效果。
• 通过将均匀扩散重构为吸收态过程，可保留 UDM 的联合分布但将其采样分解为类似 MDM 的逐步 unmasking/remasking 操作，这证明性能差距主要来自参数化与采样设计而非边缘状态选择。传统的 UDM 与 MDM 被视为两类不同的扩散过程，该论文提出的吸收态重构则表明 UDM 也可采用 MDM 的高效策略，如延续 unmasking 和自然 remasking。实验显示该重构 UDM 在语言建模上匹配或超越 MDM，颠覆了“掩盖扩散天然更优”的认知，并提示离散扩散模型的设计重心应从边缘状态转向参数化与采样策略。

方法

核心模块与流程

输入：离散 token 序列，通过正向均匀扩散逐步加噪，得到噪声序列。

1. 留一后验（Leave-One-Out Posterior）参数化
标准 UDM 的 plug-in 桥接参数化直接使用去噪后验 $p_\theta(x_0|x_t)$ 来定义反向过程，但我们的分析表明，最优的 plug-in ELBO 优化目标并不是该去噪后验，而是留一后验：预测每个位置的干净 token 时，故意不用该位置自身的噪声观测。这导致训练时交叉熵损失与 ELBO 之间的目标错配。我们将此留一目标显式表征，并推导出 denoiser、留一后验 与 score 函数之间的精确等价转换公式，从而可以解耦参数化方式与训练目标。

2. 采样器改进
基于留一预测器，我们设计了 informed predictor-corrector 采样器，在预测步（predictor）后引入基于留一后验的校正步（corrector），减少采样路径的离散化误差。同时提出改进的温度采样：利用留一后验进行更可控的 sharpening，避免标准 temperature annealing 对 mask 模型的依赖。这些无需额外训练。

3. 吸收态重构（Absorbing State Reformulation）
将均匀扩散的噪声空间重构为一个“吸收态”机制：在扩散过程中，token 以一定概率转变为一种特殊的吸收态（类似 mask），但与 MDM 的区别在于联合分布保留与 UDM 完全一致。该重构将采样分解为两个子步：

• carry-over unmasking：直接复制未改变位置的 token；
• remasking：对改变位置按留一后验重新采样或吸收。 这使反向过程具备更简单的去噪后验形式，并自然支持 remasking 策略，类似于 mask 扩散的逐步解码，但本质仍是均匀扩散的边缘分布。

输出：生成序列，在语言建模任务上显著提升 UDM 质量，甚至吸收态构造可匹配或超越传统 Masked Diffusion Models (MDM)。

与同类方法的关键差异：该工作首次从 ELBO 紧密度角度暴露 UDM 的 参数化－目标不匹配，并通过留一后检连接训练与推理，而吸收态重构则在保持均匀扩散联合分布的前提下，赋予采样过程 mask 扩散的高效结构。

实验

本文实验面向语言建模任务，在离散扩散模型的框架下，系统对比了均匀扩散模型（UDM）和掩码扩散模型（MDM）在不同参数化与采样策略下的生成质量。

实验设计

作者从 denoising posterior 与 leave-one-out posterior 的差异入手，首先通过理论推导建立 denoiser、leave-one-out 预测器和 score 之间的精确转换关系。基于该转换，解耦了参数化和训练目标：既可以用传统的交叉熵损失训练 denoiser 再转换为 leave-one-out 参数化，也可以直接以 leave-one-out 目标进行训练。在此基础上，引入了 informed predictor-corrector sampler 和基于 leave-one-out 预测器的 temperature sampling，均无需额外训练。此外，作者提出一种 absorbing-state reformulation，将均匀扩散的联合分布等价重写为类似掩码扩散的采样操作，包含更简单的 denoising posterior、carry-over unmasking 和 natural remasking 机制。

关键发现

• Leave-one-out 参数化 持续提升 UDM 生成效果，证实标准 plug-in bridge 参数化与 denoising posterior 之间存在失配，而 leave-one-out posterior 才是更优的 plug-in 目标。
• Absorbing 构造 在语言建模上达到或超过 MDM 的性能，意味着均匀扩散与掩码扩散之间的经验差距主要源于参数化和采样设计，而非边缘分布本身的选择。
• 推理阶段的改进（predictor-corrector 采样器与温度调整）在无须重新训练的情况下进一步提升了样本质量。

与基线对比的解读

与标准 UDM 相比，leave-one-out 参数化和 absorbing 重参数化均带来了显著且一致的提升；与 MDM 相比，重构后的 UDM 弥补甚至逆转了此前的劣势。这一结果表明，离散扩散模型的设计空间远未被充分探索，面向均匀扩散的深入分析能释放其在语言生成等领域的潜力，并且参数化的选择比噪声调度或边缘分布对最终性能的影响更为关键。代码已开源：https://github.com/samsongourevitch/rev_udm，便于后续对比和复现。

行业影响

落地场景

离散扩散模型（UDM 及 MDM）已在文本生成、代码补全、蛋白质序列设计、分子图生成等需处理离散 token 的领域初步应用。本文提出的 leave-one-out denoiser 与 吸收态重构 UDM 可直接优化这些场景下的生成质量与推理效率，尤其适合需要高精度自回归生成的对话系统、商品描述自动撰写、智能编程助手等产品。吸收态重构让均匀扩散具备类似掩码扩散的“逐步 unmasking”行为，对交互式生成（如用户逐步确认补全）更友好。

商业价值

• 降本：通过 leave-one-out 参数化，模型可用更标准的交叉熵 loss 训练，无需额外模块即获得更优的 ELBO，减少无效迭代；推理端的 predictor-corrector 采样 和 改进温度采样 无需重训模型，直接提升生成多样性或保真度，降低调参和再训练成本。
• 增收/体验提升：生成文本的连贯性、语法正确率提升可直接改善自动化文案、客服回复的用户体验；代码补全的更少语法错误降低人工修正成本，提升开发者付费意愿。吸收态 UDM 支持 carry-over unmasking 和 natural remasking，可灵活应对部分已知上下文（如模板式生成），拓宽商业应用边界。

与现有产品/工作流的接口

当前主流扩散文本生成方案多基于 MDM 或统一框架，改进可无缝集成：

• 将原 denoiser 网络输出映射为 leave-one-out 后验，仅需一行公式转换（论文 Eq. (11)），即可替换原有采样循环中的预测头。
• 吸收态重构可看作对 UDM 训练过程的重新分解，不改变模型架构，仅调整训练时的 noise schedule 和损失计算，可复用现有 Transformer 类 backbone。
• 推理优化组件（informed corrector）可作为附加采样步骤插入，兼容 Hugging Face diffusers 或自建推理管线。

具体用例：

1. 电商/内容平台：自动生成商品描述时，常需根据关键属性（品牌、材质等）补全长文本。吸收态 UDM 可将已知属性作为“unmasked” token 带掩码逐步生成剩余部分，减少属性不一致错误；leave-one-out 后验提升长文本连贯性。
2. 企业代码助手：在 IDE 中根据上下文补全代码行，离散扩散模型每次生成整行或整段代码，而非逐 token 自回归。改进后的 UDM 使一次性生成的代码块编译通过率更高，减少开发者打断，提升采纳率。

局限

• 实验验证局限于语言建模任务，尚未在蛋白质序列、图生成等其他离散域上测试。吸收状态重构虽有效，但引入了与masked diffusion类似的额外设计选择（如remasking schedule），可能增加调参负担并削弱通用性。
• 理论分析基于均匀转移矩阵，未涵盖更一般的离散扩散过程。leave‑one‑out 参数化要求修改训练目标，且推导出的精确转换在数值实现中可能面临稳定性问题，额外增加了实现复杂度。
• 作者仅与标准 UDM 和 MDM 基线对比，未讨论近期离散扩散的其他改进（如 D3PM、DDRM）之间的关系，也缺少与连续扩散模型的间接比较，限制了工作的相对定位和影响力评估。