Discrete Stochastic Localization用于非自回归生成

# 面向非自回归生成的离散随机定位
来源: https://arxiv.org/html/2605.12836

武云舒¹ 程嘉怡² 余龙轩¹ Partha Thakuria¹ Rob Brekelmans Evangelos E. Papalexakis¹ Greg Ver Steeg¹
¹加州大学河滨分校 ²纽约大学
{ywu380, ylong030, pthakuria, epapalex, greg.versteeg}@ucr.edu [email protected] [email protected]

###### 摘要

连续扩散是非自回归生成的自然框架，但在离散序列生成方面通常落后于掩码离散扩散模型 (MDM)。我们认为瓶颈并非连续性本身，而是去噪依赖于时间步索引噪声状态的一种表示。我们提出**离散随机定位** (DSL)，这是一个连续状态框架，使用单位球面令牌嵌入，其贝叶斯最优去噪器在定位通道下对名义信噪比 (SNR) 具有不变性。一个训练好的网络随后支持整个每个令牌SNR路径族，其中端点掩码扩散路径是特例。在OpenWebText上，使用DSL微调预训练的MDLM检查点，在所有步数预算从T=128到T=1024上显著提高了分布保真度 (MAUVE)，并且同一检查点支持随机阶自回归采样，以及一个在总步数T=48时就能生成连贯内容的混合连续-离散采样器——无需蒸馏或重新训练。

## 1 引言

连续扩散已成为高维生成的标准框架，在图像、视频和音频合成中取得了显著成果 (Ho et al., 2020; Song et al., 2020; Ho et al., 2022; Huang et al., 2025; Kong et al., 2020)。然而，对于语言而言，令牌嵌入上的连续扩散并未带来同样的优势：现有的连续扩散语言模型一直落后于掩码离散扩散模型 (MDM) (Austin et al., 2021; Sahoo et al., 2024; Shi et al., 2024)。这一差距令人费解。连续扩散让模型暴露于从完全无信息到几乎干净状态的连续有限SNR状态中，原则上这应该比端点式掩码破坏提供更丰富的监督。那么，为什么这种额外的连续性没有转化为更强的语言生成能力？

我们认为，缺失的关键并不仅仅是连续性，还有表示选择。标准嵌入空间扩散通常采用方差保持 (VP) 或方差爆炸 (VE) 高斯加噪。这些参数化在概率上是有效的，但它们将零信息极限表示为各向同性高斯噪声，而不是类似掩码的令牌状态。此外，原始噪声嵌入通常不足以进行干净令牌预测：去噪器还需要接收一个时间步或噪声水平标签来校准嵌入中包含的令牌证据量。因此，在这种表示中，去噪器对时间条件非常敏感，需要区分掩码状态和不确定状态。

这提出了一个简单的设计目标：参数化破坏过程，使得噪声令牌状态本身决定干净令牌的后验。我们提出**离散随机定位** (DSL)，一种基于单位范数令牌嵌入上的随机定位的连续扩散语言模型。在DSL中，零SNR状态是零向量，可以将其与[MASK]令牌等同，而高SNR极限恢复干净令牌嵌入。更重要的是，我们证明了单位范数随机定位参数化使得贝叶斯最优干净令牌后验仅依赖于噪声令牌状态，而不依赖于名义SNR。因此，DSL使用一个时间不变的去噪器：掩码的、不确定的有限SNR和接近干净的令牌状态都被去噪为干净令牌后验，无需显式的时间条件。

这种时间不变的去噪器为DSL带来了一个简单的推理层面后果。由于模型不绑定于全局扩散时间，不同位置可以占据不同的SNR水平，并遵循不同的更新调度。因此，掩码细化、随机阶自回归解码和混合连续-离散采样被证明是查询同一后验估计器的不同方式，而不是独立的特定于采样器的模型。因此，DSL保持了掩码扩散的端点行为和检查点兼容性，同时添加了连续的有限SNR状态用于更丰富的细化。

我们通过从预训练的MDLM检查点微调来实例化DSL，并在OpenWebText上通过同一模型评估多条推理路径。在使用ReMDM系列解码器的掩码细化路径上，DSL在采样预算从T=128到T=1024上显著提高了少步MAUVE。同一检查点还支持无需重新训练的随机阶自回归采样，以及一个在总步数T=48时就能产生连贯生成的混合连续-离散采样器。这些结果共同支持核心主张：端点揭示、掩码细化和连续/混合去噪可以由一个后验估计器提供，而不是由独立的特定于采样器的模型提供。

参考标注 (a) DSL动力学：每个zi(t)随时间向一个符号锚点定位。
参考标注 (b) 状态总结：令牌跨越SNR 1.0–24.3，从掩码到硬解码。
图1: 离散随机定位 (DSL) 动态地将采样离散令牌“定位”。*（左）*在循环玩具序列上的DSL动力学 (ABCDE, BCDEA, ...)；详情见附录D。*（右）*某个快照的状态总结：每个令牌占据不同的SNR区域，随着SNR增加从掩码过渡到软解码再到硬解码。

贡献。
- • **用于离散生成的连续状态框架。** 我们提出了DSL，一个基于单位球面令牌嵌入的随机定位信道，其贝叶斯最优去噪器通过构造仅依赖于噪声状态而非名义SNR。因此，单个时间无关网络可以支持整个每令牌SNR路径族。
- • **实用配方。** 混合支撑SNR采样在端点状态和有限SNR状态上训练同一个去噪器，而令牌混合转换器将该框架转化为标准Transformer/DiT骨干网络上的可训练目标。
- • **跨路径的证据。** 单个DSL检查点改进了OpenWebText在掩码细化下的少步生成，无需重新训练即支持随机阶AR和混合采样，并在Text8似然上具有竞争力。

## 2 离散随机定位

本节分三步展开DSL。我们首先建立该框架的最简单版本——一个高斯观测信道应用于整个句子，具有单一全局信噪比（§2.1）。然后我们引入DSL的关键设计选择，将干净令牌嵌入放置在单位球面上，并展示在该几何结构下贝叶斯最优去噪器仅依赖于噪声令牌而非名义SNR（§2.2）。由于一个去噪器即可胜任所有SNR区域，我们可以让每个令牌位置在其自身的SNR下演化；这将该信道推广为一族连续的每令牌SNR路径（§2.3）。最后，两个精确NLL观点为训练目标所使用的状态支撑提供了依据（§2.4）。

### 2.1 句子级连续起点

我们从DSL的最简单版本开始：一个高斯观测信道应用于整个句子，具有单一全局信噪比。这个设置是标准扩散的自然连续类比，并让我们推导出MMSE漂移形式，其余部分将在此基础上构建。

令s = (s1, ..., sL) ∈ V^L 为一个长度为L的句子，并将每个令牌嵌入为 xi = enc(si) ∈ R^d。我们将完整的句子嵌入记为 x = [x1; ...; xL] ∈ R^(Ld)。句子级定位信道为：
z_t = t x + √t ε, x ~ P(x), ε ~ N(0, I)。
除以t得到 z_t/t = x + ε/√t，因此 z_t 在信息上等价于通过方差为1/t的加性高斯噪声观测 x。因此信噪比恰好是t：在t=0时观测不携带信息，随着t增长，z_t 逐渐包含更多关于x的信息。

该信道有一个简单的SDE实现：
dz_t = x dt + dw, x ~ P(x), (1)
它条件于未知的干净句子，因此不能直接用于生成。我们寻求一个等价的无条件SDE——等价意味着两个动力学在 z_t 上诱导相同的边缘分布——其漂移仅依赖于 z_t：
dz_t = \hat{x}(z_t, t) dt + dw, \hat{x}(z_t, t) := E_{P_t(x|z_t)}[x]。(2)
实现这种等价的漂移是MMSE去噪器，即给定 z_t 时 x 的后验均值（证明见附录A.4）。此时DSL仍然类似于标准连续去噪构造：一个全局SNR，一个名义上依赖于t的去噪器。

### 2.2 单位范数令牌几何结构给出时间不变去噪器

DSL的关键设计选择是将每个干净令牌嵌入约束在单位球面上，对所有i有 ||xi||_2 = 1。这是一个内置于嵌入几何的结构性约束，而不是损失上的正则化项。正如我们将展示的，这种构造的好处是消除了贝叶斯最优去噪器中的名义SNR t。

在定位信道 公式(1)下，
P_t(x|z_t) ∝ P(x) exp( x·z_t - (t/2) ||x||_2^2 )
是贝叶斯后验分解。项 (t/2) ||x||_2^2 是名义SNR进入的唯一位置。在单位范数令牌嵌入下，对于每个长度为L的句子有 ||x||_2^2 = L，因此这一项与句子无关，在归一化后抵消。后验，因此MMSE去噪器，仅依赖于 z_t：
\hat{x}(z, t) = E_{P_t(x|z)}[x] = E_{P(x)}[x e^{x·z}] / E_{P(x)}[e^{x·z}] = \hat{x}(z). (3)

公式(3)中的SNR不变性既需要定位参数化（使 z 成为后验自然坐标，且 a_τ/σ_τ^2=1），也需要单位球面几何（使 ||x_i||^2 词汇量恒定）。单独任何一个条件都不够；它们的组合产生了结构性的抵消。

这种SNR不变性是推动整篇论文的结构特性。贝叶斯去噪器不需要知道是哪个t生成了 z_t 就能最优地去噪它；它只需要噪声令牌。因此，一个单一网络原则上可以同时服务于每一个SNR区域。我们在附录A.2中给出完整推导；几何后果——诱导令牌后验 p(x_i|z_t) 上的熵包络和Lipschitz光滑性——推迟到第4节。

### 2.3 任意每令牌SNR路径

§2.2中的SNR不变性让我们可以放弃所有令牌共享单一SNR的假设。我们可以让每个位置 i 在其自身的SNR下演化：
z_i = γ_i x_i + √γ_i ε_i, ε_i ~ N(0, I)，
并且我们允许每个 γ_i(t) 遵循任意可允许的路径——连续的、非递减的，且 γ_i(0)=0 且当 t→∞ 时 γ_i(t)→∞。对于这些每令牌路径，数据条件和无条件动力学为：
dz_i = x_i γ_i˙ dt + √(γ_i˙) dw_i, x ~ P(x), (4)
dz_i = \hat{x}_i(z) γ_i˙ dt + √(γ_i˙) dw_i. (5)
我们在附录A.4中显示这两个SDE沿着任何可允许的每令牌路径在 z_t 上诱导相同的边缘分布。结合 \hat{x}_i 的SNR不变性，这意味着**一个**训练好的去噪器可以沿着任意的每令牌SNR调度进行采样，而不改变目标分布。

路径解释立即可得。将一个令牌从γ=0驱动到∞同时保持其他令牌固定，对应于以随机阶自回归 (ROAR) 风格揭示该令牌；将选定的位置送回γ=0再重新去噪，对应于带有重掩码的掩码扩散；同时增加所有γ_i对应于标准连续扩散。随机阶AR生成、掩码细化和连续扩散采样是一个共享的每令牌SNR配置族中的不同路径——并且，根据上述论证，它们都可以由一个训练好的DSL模型服务。

### 2.4 激励训练支撑的精确NLL观点

DSL训练目标并非作为采样器启发式选择。它源自同一每令牌SNR形式主义的两个精确似然观点：一个关于有限SNR去噪状态的连续路径积分NLL，和一个关于掩码/揭示状态的端点ROAR NLL。这些恒等式是似然层面的原因，解释了为什么第3.1节中的混合支撑损失将训练质量分布在这两个状态族上。

**连续路径积分观点。** 对于...