跳到主要内容
返回时间线
arXiv来源发表:

掩码扩散模型在预算内提前锁定答案:HiLo-RD 用大窗口轨迹自蒸馏,GSM8K 准确率升 9.4 个百分点

核心概要

掩码扩散语言模型在小响应窗口下会在推理尚未完成时先固定答案片段,这一提前承诺随后成为后续预测的上下文。HiLo-RD 让同一模型的大窗口轨迹充当冻结教师,把其隐藏表示对齐到裁剪后的小窗口视图,在已提交位置上做表示蒸馏。在相同小推理预算下,GSM8K 准确率从 68.84% 升到 78.24%,平均可读长度从 362.8 降到 271.3 个 token,增益还延伸到未做代码域训练的代码生成。

AI-generated editorial illustration: HiLo-RD: Cross-Budget Self-Distillation for Masked Diffusion Language Models

深度剖析

掩码扩散语言模型在小响应窗口下可能先固定答案尾部片段,再补中间推理,即使正确解本可放进窗口。 同一模型在大窗口下会推迟答案并得到更短、更正确的解,而现有后训练方法没有直接把这种大窗口轨迹迁移到小窗口决策。 在 155 个低预算答错、高预算答对的配对案例中,固定预算下推迟尾部提交使 43.2% 的案例恢复正确答案,匹配的顺序扰动对照为 30.3%,配对增益 12.9 个百分点,95% 自助区间 7.7 到 18.7。

HiLo-RD 把同一模型大窗口轨迹的隐藏表示迁移到小窗口解码,同时影响 token 预测和提交顺序。 它不使用外部教师,也不只对孤立 token 或单独指定的提交标签加损失,而是在表示层面对齐跨预算轨迹。 教师与学生从同一检查点出发,教师用大预算生成轨迹,学生接收裁剪后的小窗口状态,在教师更新后累计提交的位置上匹配最终归一化隐藏状态;训练用 LoRA 适配,骨干与词表头保持冻结。

在相同小推理预算下,HiLo-RD 同时提高准确率并缩短可读响应。 与原始 LLaDA-8B-Instruct 及 Masked SFT、Diffu-GRPO、OPSD、dOPSD 四种后训练方法相比,它在 GSM8K 和 MATH500 上取得最高准确率与最短可读响应。 GSM8K 上准确率从 68.84% 升到 78.24%,平均可读长度从 362.8 降到 271.3 个 token;MATH500 上准确率从 28.40% 升到 32.40%;数学训练检查点未经代码域训练,在 MBPP 上 pass@1 从 38.60% 升到 41.40%。

启示与展望

演示范围是 LLaDA-8B-Instruct 上的纯扩散解码,训练分别在完整 GSM8K 训练集和原始 MATH 训练集上进行,并分别在 GSM8K 测试集与 MATH500 上评估;代码结果来自数学训练检查点,未做额外代码域训练。连续自蒸馏实验只使用固定的 1000 条 GSM8K 训练样本,与全训练集主对比是分开的序列。

可达到的压缩程度依赖任务:HumanEval 上相对 dOPSD 的 pass@1 略低,可读长度缩减也有限。报告的时延差异来自不同服务器,仍需归一化;完成前缀停止的保持性质依赖匹配的解码与首 EOT 提取约定。表示对齐只迁移可由共享提示和裁剪状态预测的那部分大预算信息,参数化学生仍可能带来近似与优化误差。