扩散语言模型的表示空间 MMD

摘要

我们引入了一种用于扩散语言模型 (DLM) 的后训练方法，该方法可最大限度地减少冻结预训练 DLM 特征空间中生成分布和参考分布之间的最大平均差异 (MMD)。为了估计 MMD，我们保留各个token位置的上下文特征，从单个提取器通道中获取每个序列的多个观察结果。我们对离散模型使用策略梯度来优化此目标，并通过生成的潜变量对连续模型进行直接微分。在这两种情况下，直接根据这些特征计算损失可以实现高效的后训练，而无需完整的采样轨迹或联合训练的辅助模型。实验表明，在 OpenWebText 上可比较的熵下，生成复杂度较低，而在 GSM8K 上则具有更好的准确度计算权衡。在具有混合掩模均匀扩散的 16B DMax-LLaDA2.0 模型上，我们在数学和代码基准上以相似或更高的精度提高了解码并行性。

扩散语言模型的表示空间 MMD的原论文方法或结果图 — 图 1：表示空间 MMD 训练。我们说明了每个 MMD 估计有一个参考序列和 $B$ 生成序列的情况。对于离散 DLM（左上），我们从单个降噪器通道生成的token分布中独立绘制 $G$ 这样的批次，并通过分组 REINFORCE 优化损失。对于连续 DLM（左下），$K$ 自举通道无需梯度即可提供自我调节，然后是可微分生成通道。在这两种情况下，token级 RBF MMD 都会比较由冻结的预训练 DLM 提取的上下文特征（右）。