迷失在插值中：为什么预测性反馈在扩散语言模型中失效

# 迷失于插值：预测反馈为何在扩散语言模型中失效

来源：https://arxiv.org/html/2608.06529

Lavanya Nigam∗, Ishaan Bansal∗, Aryan Sood∗, Vidit Aggarwal, Gaurav Kumar Nayak

印度理工学院鲁尔基分校，鲁尔基，北阿坎德邦，印度

[email protected], [email protected], [email protected], [email protected], [email protected]

###### 摘要

软掩码（Soft-masking）加速了掩码扩散语言模型（MDLMs）的收敛。现有公式通过原始嵌入空间中的线性插值（LERP）构建这种混合，这隐含地将该空间视为欧几里得空间。我们分析了 MDLM 的嵌入空间，发现掩码嵌入与预测 token 嵌入在整个训练过程中保持约 73° 的近恒定角度，而嵌入范数在词表频率排序上基本保持平坦。这些现象表明存在超球面几何，而 LERP 在这种几何下是错误的插值原语。我们提出了*球面软掩码（S-SM）*，这是一种即插即用的替代方案，它在超球面上用 Fréchet 均值聚合 top-k 预测，并通过球面线性插值（SLERP）将此均值与掩码方向混合，然后恢复掩码自身的原始范数。我们在一个已发布的 169M 参数 MDLM 检查点上，对广泛的推理时步数预算进行了持续预训练评估。SLERP 反馈避免了 LERP 反馈引起的训练退化，并在各种采样预算下相比标准 MDLM 基线取得了高达 2 倍的 MAUVE 提升，相比 TopK/LERP 取得了 27.5%–56.1% 的提升，同时生成困惑度持续更低（相比基线降低 16.9%–19.6%），而输出熵和收敛性基本不变。

迷失于插值：预测反馈为何在扩散语言模型中失效

Lavanya Nigam∗, Ishaan Bansal∗, Aryan Sood∗, Vidit Aggarwal, Gaurav Kumar Nayak

印度理工学院鲁尔基分校，鲁尔基，北阿坎德邦，印度

[email protected], [email protected], [email protected], [email protected], [email protected]

††footnotetext:∗同等贡献。

## 1 引言

掩码扩散语言模型（MDLMs）通过反复对完全掩码的序列进行去噪来生成文本。在每个去噪步骤中，主干网络为每个仍被掩码的位置提出一个 token，然后由去掩码规则决定保留哪些提议，以及哪些掩码位置留待下一步处理。这种二元接受/拒绝规则简单且训练效果好，但它在决定是否揭示 token 之前丢弃了模型已经知道的信息。留下的 [MASK] 嵌入在每个位置和每个步骤都完全相同，无法携带先前步骤为预测而收集的信息。参见图 1 的说明：投影到单位超球面上的两个嵌入之间的 LERP 和 SLERP 路径。

软掩码（SM）的引入正是为了阻止这种信息丢失。SM 用连续的置信度加权混合取代二元的去掩码决策。对于每个保留的掩码位置，反馈嵌入被构造为掩码 token 嵌入 m 与 top-k 预测 token 嵌入的加权叠加的凸组合。这一单独的改动已被证明能在不同模型规模上同时改善困惑度和生成质量。

目前的架构将这种混合构建为掩码嵌入与 top-k token 嵌入概率加权平均的线性组合。如果底层嵌入空间是欧几里得空间，线性插值是自然的选择。先前关于嵌入几何的研究反复观察到，训练后的 token 嵌入集中在半径近似恒定的球壳附近，即超球面，而不是填充周围的欧几里得空间。在超球面上两点之间进行线性插值会得到一个偏离流形的点，无法准确捕获插值的语义。

我们通过实验发现，软掩码 MDLM 在超球面嵌入空间中运行。归一化掩码嵌入与模型 top-k 预测的归一化均值之间的测地线角度在整个训练过程中稳定在 ≈73°，而 token 嵌入的 L2 范数在 GPT-2 频率排序上接近平坦。这两个观察都指向超球面嵌入几何。

我们为 MDLM 中的软掩码提出了一种更合理的插值原语。球面软掩码（S-SM）用内在计算在单位超球面上的 Fréchet（Karcher）均值取代 top-k 嵌入的欧几里得加权均值，并沿着连接测地线用球面线性插值（SLERP）取代掩码方向与该均值之间的线性混合。结果在交给主干网络之前被重新缩放到掩码 token 自身的范数。S-SM 是线性反馈步骤的即插即用替代，并保持训练目标、去掩码规则和基于置信度的调度完全相同。

我们的贡献如下：

1. **识别超球面嵌入几何**：在发布的 MDLM 检查点（Hersche 等人，2026）上，我们观察到掩码嵌入与 top-k 均值之间的测地线角度保持在 73° 附近，嵌入范数在频率排序上平坦。这表明了一种超球面几何，在此几何下，嵌入的线性混合位于球面流形之外。
2. **S-SM，线性反馈的几何化即插即用替代**：我们提出了一种替代线性反馈的方案，在球面上用 Fréchet 均值聚合 top-k 嵌入，沿连接测地线通过 SLERP 混合，并重新缩放到掩码 token 自身的范数。
3. **跨随机种子和采样预算的一致提升**：我们证明，在较大的 NFE 预算（256, 512）下，S-SM 的 MAUVE 相比标准 MDLM 基线几乎翻倍，并且在相同预算下比 TopK/LERP 的 MAUVE 提高了超过 50%，同时还避免了 LERP 引起的困惑度退化。在使用等效的学习置信度训练方案进行训练时，SLERP 在所有采样预算和多个随机种子上都取得了更高的 MAUVE 分数和更低的生成困惑度，同时训练困惑度几乎相同。学习到的置信度权重本身也印证了这一点：在 SLERP 下它收敛到 λ≈0.056，而在 LERP 下为 λ≈0.030（第 7k 步，≈1.9×），这表明优化器更信任几何上良好构建的反馈。

## 2 相关工作

#### 掩码扩散语言模型。

MDLM 将文本生成定义为对掩码序列的迭代去噪，使并行解码成为自回归生成的替代方案（Austin 等人，2021；Dieleman 等人，2022）。LLaDA（Nie 等人，2026）将 MDLM 扩展到 7B 参数，而 MaskGIT（Chang 等人，2022）验证了图像的并行掩码解码。由于二元决策 MDLM 会让早期错误不可逆地传播，ReMDM（Wang 等人，2026）增加了推理时重掩码，MDPO（He 等人，2025）应用带置信度重掩码的