MMD 比较的是特征而非 token 概率
这里的训练信号是两组特征之间的核距离:一组从生成序列中提取,另一组从参考序列中提取。两组特征都来自同一个冻结的预训练扩散语言模型,因此比较空间是学习到的表示,而不是原始 token 空间。论文给出的动机是:高维空间中的简单距离未必能捕捉输出之间有意义的差异,而学习到的表示则为离散序列和连续输出都提供了有前景的比较空间。
使用的核是 RBF 核 exp(-||a-b||^2/(2*sigma^2)),论文指出它是特征核,因此 MMD 为零意味着在核被评估的空间中分布相等。当 Q 为生成分布时,其分布内项惩罚生成特征之间的平均 RBF 相似度,而交叉项则奖励与参考特征的相似度。因此有两个选择塑造了目标函数:RBF 带宽和用于匹配的表示空间。讨论部分把这两者列为性能的主要依赖因素。
一次提取器前向产生大量 token 级观测
为估计 MMD,该方法保留各个 token 位置上的上下文特征,从而在一次提取器前向中为每条序列产生多个观测。这正是论文的效率论据:损失直接由这些特征计算,因此后训练不需要完整采样轨迹,也不需要联合训练的辅助模型。论文将其与扩散蒸馏方法对比:后者要么训练学生模型用更少步骤复现教师的多次去噪更新,要么从在生成器不断变化的输出上训练的辅助模型获得分布匹配梯度。
两类模型的优化方式不同。对于离散模型,从单次去噪器前向产生的 token 分布中采样序列,并用带留一基线的 REINFORCE 优化得到的 MMD 奖励;论文将其实例化为用于掩码 token 预测的 MDLM-MMD,以及用于混合掩码-均匀扩散的 DMax-MMD,其中 MMD 同时训练掩码预测和模型自身 token 预测的细化。对于连续扩散语言模型,生成器从预训练的 Embedded Language Flows 模型初始化,并被训练为一步将高斯噪声映射到干净潜变量序列,MMD 梯度直接通过冻结的特征提取器和生成的潜变量传播。推理时,自条件机制迭代细化单步预测,将其扩展为多步采样。
匹配熵时生成困惑度更低
在无条件生成中,模型在 OpenWebText 上训练,序列打包到长度 1024,并用预训练 GPT-2 Large 下的生成困惑度以及平均 unigram 熵评估 1000 个生成样本。论文以约 5.43 的数据熵作为参考,而不是把更高熵一律视为更好;它在 8、16、32 个采样步数下评估每种方法,并改变采样温度以刻画困惑度-熵权衡。
在匹配熵时,MDLM-MMD 的生成困惑度低于所有基线。在参考数据熵处对曲线插值,三个采样预算下生成困惑度比 IDLM 低约 17% 至 21%。对比集合包括 DiDi-Instruct、IDLM 和 IDLM-REINFORCE,论文称这有助于区分训练目标的效果与优化过程的效果。
对于连续模型,ELF-MMD 在 T5-small 和 GPT-2 Large 两种潜空间下,在大多数采样预算上都改进了生成。在 8 步时,它比 ELF-PD 生成困惑度更低、熵更接近参考值;在 32 步时,它在两项指标上都优于 ELF。加入迭代细化蒸馏进一步改善少步结果:在 4 步时,使用 T5 和 GPT-2 编码器时生成困惑度相对 ELF-MMD 分别降低约 30 点和 10 点,同时使两种编码器的熵都更接近参考值。
精度提升集中在较小采样预算
条件生成在 TinyGSM 上训练,序列打包到最大长度 512,并以 GSM8K 测试集上的最终答案精度评估。对于掩码模型,由于置信度解码自适应地决定模型评估次数,论文扫描置信度阈值并报告每条序列的平均步数。MDLM-MMD 在中等和高解码预算下达到更高精度,约 49 步时精度约 54%。
对于连续模型,ELF-MMD 在每个评估的采样预算上都比 ELF 精度更高,在大多数预算上比 ELF-PD 更高,而 IRD 在每个评估预算上都提升精度。在 4 步和 8 步时,精度分别从 14.2% 升至 20.8%、从 27.5% 升至 32.5%,而 ELF-PD 分别为 15.9% 和 23.5%。在 64 步时,ELF-MMD+IRD 达到其最高精度 36.3%,而 ELF-MMD 为 35.2%,ELF 为 31.6%。
pass@k 结果呈现相同模式。对于基于 MDLM 的模型,MDLM-MMD 在 4 至 8 个采样步、k 不超过 8 时优于基线;在 16 至 32 步时,它在大多数评估的 k 值上取得更高 pass@k,在更大采样预算下差异较小。对于连续扩散语言模型,ELF-MMD 在较小采样预算下也提升 pass@k,而随着 k 增大差距收窄。
16B DMax 在不损失精度下提升并行度
扩展测试从已发布的 16B DMax-Math 和 DMax-Coder 检查点开始。DMax 是一种混合掩码-均匀扩散模型,由 LLaDa2.0-Mini 调优而来,学习从掩码输入和自身 token 预测中恢复目标 token;作者沿用原始训练实现,并用 MMD 同时在这两者上训练,得到 DMax-Math-MMD 和 DMax-Coder-MMD。训练运行 400 步,在八块 NVIDIA H100 GPU 上 DMax-Math 约需 13 分钟、DMax-Coder 约需 19 分钟,即每次运行约 1.7 和 2.5 GPU 小时,仅覆盖训练,不含初始化和评估。
评估使用原始 dInfer 实现和 Soft Parallel Decoding,块大小为 32,最大生成长度为 2048 个 token,报告精度和每次前向生成的 token 数(TPF)作为解码并行度的度量。选定的阈值为 DMax-Math-MMD 的 0.85 和 DMax-Coder-MMD 的 0.9。在四个数学基准上,DMax-Math-MMD 相比报告的 DMax-Math 工作点将每次前向 token 数提高 10.3% 至 16.5%,同时精度相当或更高。在代码基准上,DMax-Coder-MMD 在 HumanEval-Instruct 上精度提高 2.4 个百分点,在 MBPP-Instruct 上提高 3.8 个百分点,同时两个基准的每次前向 token 数也都增加。
这些数字是从十一个解码阈值下评估的精度-TPF 曲线中选出的工作点,因此阈值是结果的一部分,而不是固定设置。MMD 数据点对五次不同随机种子的训练运行取平均,误差棒表示两个坐标上正负一个样本标准差。
核与表示的选择承载了结果
消融实验在 GSM8K 上比较了替代损失形式,OpenWebText 结果放在附录。它从线性核 MMD 开始,即生成与参考特征均值之间的均方误差,用以分离 RBF 核和完整差异相对于简单均值匹配的贡献。讨论部分指出,性能取决于 RBF 带宽和用于分布匹配的表示空间等选择,这使得表示设计成为进一步探索的自然方向。
论文还指出它使用干净输入上单层的特征,并建议跨层或跨噪声水平组合特征,或使用提取器集成,可能提供互补信息。除表示设计外,未来工作可以探索将 MMD 与离散扩散语言模型的其他训练目标结合,连续设置中在 ELF-MMD 后训练之后应用 IRD 所带来的增益也提示了这一点。
所报告的增益覆盖与未覆盖之处
证据支持一个具体论断:在预训练扩散语言模型特征上的 MMD 能改进离散和连续扩散语言模型的生成,且所用训练过程既不需要完整采样轨迹,也不需要联合训练的辅助模型。所报告的增益在 OpenWebText 匹配熵下的生成困惑度、GSM8K 在给定采样预算下的最终答案精度和 pass@k,以及 16B DMax 模型在数学和代码基准上的精度与每次前向 token 数上测量。
若干边界限定了这些增益。16B 结果来自在已发布 DMax 检查点上 400 步训练,并在选定的解码阈值下报告,因此它们描述的是精度-TPF 曲线上的工作点,而非与阈值无关的改进。无条件与条件结果来自特定训练集、序列长度和评估设置,论文自己的讨论将性能归因于 RBF 带宽和表示空间,这意味着所报告数字与这些选择绑定。论文还使用干净输入上单层的特征,因此组合层、噪声水平或提取器的效果在此未经检验。
这篇论文确立的实际优先级是一种权衡而非单一分数:在较小采样预算下,表示空间 MMD 换来精度和匹配熵下更低的生成困惑度;在 16B DMax 上,它换来精度相当或更高时的解码并行度。与决策相关的变量是解码阈值,因为所报告的数学和代码增益是在阈值 0.85 和 0.9 的精度-TPF 曲线上读取的。因此,比较后训练方法的读者应追问:所声称的增益是在哪个采样预算、哪个阈值下测量的,因为同一方法在 4 步、64 步以及并行度曲线上的不同点可能表现不同。
本文引用的文献
Representation-Space MMD for Diffusion Language Models
在冻结预训练扩散语言模型的表示空间中最小化生成特征与参考特征之间的 MMD,可在无需完整采样轨迹或辅助模型的情况下后训练离散和连续扩散语言模型,在 OpenWebText 上匹配熵时降低生成困惑度,改善 GSM8K 精度-计算权衡,并在 16B DMax 模型上以相当或更高的数学与代码精度提升解码并行度。
