可学习性引导的扩散语言模型微调

# 基于可学习性的扩散语言模型微调
来源: https://arxiv.org/html/2605.22939
作者: Atharv Chagi, Jacob Helwig, Lakshmi Jotsna, Sushil Vemuri, James Caverlee, Dileep Kalathil, Shuiwang Ji

###### 摘要

我们旨在提升扩散语言模型（DLMs）的推理能力。虽然监督微调（SFT）是自回归模型常用的后训练方法，但其在DLM中的应用面临挑战，甚至可能损害性能，而根本原因尚未得到充分研究。我们的分析表明，原始的SFT忽略了*可学习性*，即标记*学习什么*以及*何时学习*。具体而言，当输入大部分被掩码时，稀有标记难以学习；而当输入大部分未被掩码时，学习常见标记则变得简单且价值不大。受此分析启发，我们提出了LIFT，一种高效的基于SFT的DLM后训练算法。LIFT在输入大部分被掩码时学习简单标记，并在上下文更丰富时学习困难标记，从而使训练与扩散不同时间步的信息可用性对齐。实验结果表明，LIFT在六个推理基准上优于现有SFT基线，在AIME'24和AIME'25上取得了高达3×的相对提升。我们的代码已在https://github.com/divelab/LIFT开源。

机器学习，ICML

#1#

## 1 引言

扩散模型在图像生成（Song and Ermon, 2019 (https://arxiv.org/html/2605.22939#bib.bib11); Nichol and Dhariwal, 2021 (https://arxiv.org/html/2605.22939#bib.bib10)）和视频生成（Ho et al., 2022 (https://arxiv.org/html/2605.22939#bib.bib26)）等应用中展现了卓越性能。近期，扩散模型已成功应用于文本数据，引发了业界对扩散语言模型（DLMs）的广泛兴趣（Austin et al., 2021a (https://arxiv.org/html/2605.22939#bib.bib12); Sahoo et al., 2024 (https://arxiv.org/html/2605.22939#bib.bib14)）。DLM相对于自回归语言模型（ARLMs）的一个核心优势在于，它们每次模型调用可并行生成多个标记，从而显著提升推理吞吐量（Khanna et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib2); Wu et al., 2026 (https://arxiv.org/html/2605.22939#bib.bib5)）。目前已有多个开放权重的DLM，如LLaDA（Nie et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib15)）和Dream（Ye et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib16)），其性能基本与同等规模的自回归模型相当。

详见说明图1：AIME基准测试上的表现。LLaDA-8B-Instruct、原始SFT和LIFT在AIME'24和AIME'25上的Pass@16准确率对比。LIFT在这两个具有挑战性的数学推理数据集上均比原始SFT取得了显著相对提升，验证了基于可学习性训练的有效性。

详见说明

详见说明

(a) 频率 vs. 置信度。
详见说明(b) 不同时间步的标记级置信度。

图2：使用LLaDA的标记分析。利用从4个后训练语料库（Muennighoff et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib23); Bercovich and others, 2025 (https://arxiv.org/html/2605.22939#bib.bib30); Open-R1, 2025 (https://arxiv.org/html/2605.22939#bib.bib31); Team OLMo and others, 2025 (https://arxiv.org/html/2605.22939#bib.bib32)）整理的0.5B个掩码标记，我们分析了标记级置信度和频率。(a) 我们按对数标度频率对标记进行分箱，并绘制平均模型置信度与平均频率的关系图。边缘化图（顶部）显示，稀有标记的平均置信度较低，表明某些标记更难预测（*什么*维度）。我们通过按扩散时间步 t 细分边缘化图（底部）进行了更细致的分析，揭示了*什么*和*何时*维度之间的交互作用。具体而言，我们观察到一种由 t 引起的偏差：当 t 较大时，模型输入大部分被掩码，低频标记变得异常难以预测，这表明随着扩散时间 t → 1^+，正向扩散过程后期出现的重度掩码输入的信息内容不足以可靠地学习某些标记。相反，当 t → 0^- 时，低频标记变得更可学习，而预测高频标记则变得微不足道。(b) 我们采样了具有代表性的高频和低频标记，展示了它们在整个扩散时间上的（平均）置信度。稀有标记在 t → 1^+ 时问题加剧，且置信度下降幅度比高频标记更剧烈。继后训练ARLM以提升推理能力的成功之后，近期工作已探索使用监督或指令微调（SFT）（Ye et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib16); Nie et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib15)）和强化学习（RL）（Zhao et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib21)）对DLM进行后训练。然而，与ARLM不同，由于序列级似然难以处理，DLM中的RL在技术和算法上都更具挑战性，大多数关于DLM的RL工作都提出了近似方法以克服这一困难（Zhao et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib21); Kunde et al., 2026 (https://arxiv.org/html/2605.22939#bib.bib3); Wang et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib4)）。SFT的研究则相对不够深入，至今尚无工作系统性地探究将SFT应用于DLM所涉及的挑战。近期结果表明，SFT实际上可能相对于预训练降低模型性能（Ye et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib16)）。这激发了我们工作的核心问题，我们将其分解为两个子问题：(i) 影响DLM的SFT后训练的主要因素有哪些？(ii) 我们如何设计一种能够考虑这些因素以有效后训练DLM的SFT算法？

作为我们的第一个贡献，我们通过分析DLM中的SFT并描述其失败案例来回答 (i)。具体而言，我们在图2(a) (https://arxiv.org/html/2605.22939#S1.F2.sf1) 中进行了广泛分析，涵盖了从四个流行的后训练推理数据集（Muennighoff et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib23); Bercovich and others, 2025 (https://arxiv.org/html/2605.22939#bib.bib30); Team OLMo and others, 2025 (https://arxiv.org/html/2605.22939#bib.bib32); Open-R1, 2025 (https://arxiv.org/html/2605.22939#bib.bib31)）整理的0.5B个标记。在多个预训练DLM（Ye et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib16); Nie et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib15)）上，我们的发现揭示了两个关键考量，它们的相互作用主导了SFT的动态；即：*学习什么*标记，以及在扩散过程中*何时*学习这些标记。我们的发现表明，语料库中的稀有标记比频繁标记更难预测（什么）。此外，当上下文更丰富时（对应于早期正向扩散时间），稀有标记变得更可学习。然而，在后期正向扩散时间，输入中减少的信息会不成比例地降低模型对稀有标记的置信度，在某些情况下使其实际上不可学习（何时）。这些发现表明，当正向扩散时间 t → 1^+ 时，稀有标记通常变得不可学习，因此将计算资源集中在频繁标记上更为有效。相反，当正向扩散时间 t → 0^- 时，频繁标记容易预测，而稀有标记则变得更可学习。虽然先前的工作已经提出了部分遵循这些指导原则的启发式方法，但它们仅孤立地考虑了“什么”或“何时”维度（Ye et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib16); Xu et al., 2026 (https://arxiv.org/html/2605.22939#bib.bib25)），而我们的研究是首个系统分析它们在监督微调期间联合效应的工作。我们证明，模拟标记难度与扩散时间之间的交互作用对于改进训练至关重要。

作为我们的第二个贡献，受这些见解的启发，我们提出并开发了 LIFT，这是首个在 DLM 训练中针对“什么”和“何时”之间交互作用的后训练方法。LIFT 训练模型在那些根据可用上下文在每个扩散时间最合适学习的掩码标记上。在两个 DLM 基础模型上，我们在四个推理基准中获得了各种 SFT 训练框架中最先进的成果。我们还在具有挑战性的 AIME-24（AIME, 2024 (https://arxiv.org/html/2605.22939#bib.bib33)）和 AIME-25（Math-AI Team and Zhang, 2025 (https://arxiv.org/html/2605.22939#bib.bib34)）上评估了 LIFT，相较于 SFT 基线取得了高达 3 倍的提升。值得注意的是，LIFT 在性能上接近 RLVR 基线 d1（Zhao et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib21)），同时使用的 GPU 小时数减少了约 500 倍，为 DLM 后训练建立了新的帕累托前沿。

详见说明图 3：基于可学习性的微调（LIFT）。LIFT 通过利用模型置信度和扩散时间构建一个基于可学习性的掩码，从而在扩散过程的每个时间点训练最具实用性的标记，进而提升可学习性。实用性被估计为模型置信度和扩散时间的函数。在第一阶段，以比率 t + ρ 采样一个掩码，并用于估计所有掩码位置上的模型置信度 p_θ(x_0 | x_{t+ρ})。然后 LIFT 根据模型置信度和扩散时间，从 x_{t+ρ} 中选择一个掩码标记子集进行监督。根据扩散时间的不同，子集选择可以是 top-K 最确信标记、bottom-K 最不确信标记，或随机（原始）。从扩散时间到子集选择方法的映射旨在根据第 4 节 (https://arxiv.org/html/2605.22939#S4) 中我们分析的见解，增加每个训练步骤的可学习性和实用性。

## 2 相关工作

#### 扩散语言模型

将扩散模型在连续域（如图像生成）的成功（Ho et al., 2020 (https://arxiv.org/html/2605.22939#bib.bib1); Nichol and Dhariwal, 2021 (https://arxiv.org/html/2605.22939#bib.bib10); Song and Ermon, 2019 (https://arxiv.org/html/2605.22939#bib.bib11)）扩展到语言领域。然而，将连续扩散应用于离散文本本身就很困难（Austin et al., 2021a (https://arxiv.org/html/2605.22939#bib.bib12)）。为了解决这个问题，掩码扩散语言模型（Sahoo et al., 2024 (https://arxiv.org/html/2605.22939#bib.bib14)）通过利用掩码语言建模（Devlin et al., 2019 (https://arxiv.org/html/2605.22939#bib.bib13)）提供了一种离散替代方案，其中标记被随机掩码，模型学习去掩码。近期模型（Nie et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib15); Ye et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib16)）在数学推理、代码生成（Zhu et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib17)）和多模态任务（Li et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib18)）中表现出与自回归LLM (ARM) 相当的性能，表明 DLM 能够执行复杂推理。这使得 DLM 后训练成为自然的下一步，目标是在推理方面获得与 ARM 类似的提升。

#### 后训练

DLM 的后训练与自回归模型相似，遵循两种方法之一，即基于可验证奖励的强化学习（RLVR）（Guo et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib19); Parashar et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib20); Zhao et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib21)）或监督微调（SFT）。使用高质量思维链数据的 SFT 可以达到与基于 RL 的方法相媲美的性能（Zelikman et al., 2022 (https://arxiv.org/html/2605.22939#bib.bib22); Muennighoff et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib23)）。对于 DLM，最近的 SFT 工作通过考虑*预测什么*（Li et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib18); Bie et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib24); Xu et al., 2026 (https://arxiv.org/html/2605.22939#bib.bib25)）— 因为某些标记天生更难预测 — 以及*何时预测*（Ye et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib16)）— 因为更重的掩码使得预测更具挑战性 — 来探索基于难度的训练。在这项工作中，我们研究联合考虑”什么“和”何时“之间的交互作用如何能够提高 DLM 后训练在提升推理性能方面的有效性。

## 3 预备知识

MDLM（Sahoo et al., 2024 (https://arxiv.org/html/2605.22939#bib.bib14); Nie et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib15); Ye et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib16)）在来自 p_data 的输入序列 x_0 上定义一个正向扩散过程，通过逐步用*[MASK]*替换标记，生成连续索引的损坏序列 {x_t}_{t∈[0,1]}。x_t 中的信息量随 t 单调递减，使得 x_1 中的所有标记都被掩码。为了生成新序列，MDLM 参数化一个双向预测器 p_θ 来逆转从 x_1 开始的扩散过程。p_θ 通过采样扩散时间 t ~ π(·)，t∈[0,1]（通常 t ~ Uniform(0,1)）进行训练。为了采样 x_t，x_0 中的每个标记以概率 1 - α_t 被掩码。本文中，我们遵循与 LLaDA（Nie et al., 2025 (https://arxiv.org/html/2605.22939#bib.bib15)）相同的设置，其中 α_t = 1 - t。给定损坏的输入 x_t，p_θ 学习从 x_0 中恢复掩码位置的原始标记。MDLM 训练目标是负证据下界（NELBO），它上界了数据的负对数似然。对于掩码序列 x_t，NELBO 给出为

−E_{t~U[0,1], x_0~p_data} [ 1/t ∑_{k=1}^{|x_0|} 1{x_t^k = *[MASK]*} log p_θ(x_0^k | x_t) ]

(1) 其中 |x_0| 表示 x_0 的序列长度，x_t^k 是损坏输入中位置 k 处的标记，1{x_t^k = *[MASK]*} 将损失限制在掩码位置（根据 x_t 预测相应的 x_0^k）。在原始 SFT 中，相同的损失直接在有监督训练集上优化，提示标记保持未掩码状态。

## 4 分析

在本节中，我们围绕核心问题（图2(a) (https://arxiv.org/html/2605.22939#S1.F2.sf1)）分析标记难度：*学习什么*标记以及在扩散过程中*何时*学习？

#### 哪些标记是困难的？

我们通过分析去噪置信度来探究这个问题，该置信度定义为给定噪声序列 x_t 时，模型 p_θ 分配给掩码位置 k 处真实标记 x_0^k 的概率 p_θ(x_0^k | x_t)。先前的 ARLM 工作表明，稀有标记由于训练期间暴露有限，更难学习，因此也更难预测（Kandpal et al., 2023 (https://arxiv.org/html/2605.22939#bib.bib36)）。