跳转至

Reconsidering Positional Supervision in Masked Diffusion Language Model Training

会议: NeurIPS 2026
arXiv: 2601.22947
OpenReview: https://openreview.net/forum?id=SoXwBAJ8w9
代码: 暂未开源(论文声明将开源全部实验代码)
领域: LLM 其他
关键词: 掩码扩散语言模型, 位置监督, 连接时序分类 (CTC), 序列对齐灵活性, 监督微调 (SFT)

一句话总结

针对掩码扩散语言模型(MDLM)在迭代并行解码中对微小位置偏移极其敏感的根本缺陷,本文将非自回归翻译中的对齐灵活性引入 SFT 训练,设计了插入 <slack> 占位标记与仅折叠空白的 CTC-S 辅助损失函数,在不改动推理管线的前提下显著增强了模型的位置容错性,并在四项开放式生成评测上实现全面提升。

研究背景与动机

掩码扩散语言模型(Masked Diffusion Language Models, MDLMs)作为一种非自回归式生成架构,通过在双向上下文约束下并行去噪掩码标记,在生成速度与推理灵活性上展现出巨大的应用潜力。以 LLaDA 为代表的大规模扩散语言模型在多项基准测试中已逐步逼近甚至匹敌主流自回归(AR)模型的性能。然而,现有的 MDLM 预训练与监督微调(SFT)普遍沿用标准的位置级交叉熵损失(Cross-Entropy, CE)进行监督:对真实响应随机施加掩码后,强制要求模型在每一个被遮盖的绝对物理位置上,完全无误地预测出原始目标词。这种严格的逐位置对齐机制虽然训练稳定,却在底层假设生成目标与输入上下文在空间位置上必须处于绝对静态的网格对齐状态。

这种刚性位置监督在开放式文本生成中暴露出了严重的技术矛盾。开放式生成具有固有的多模态表达特性,同一语意可以容许词序颠倒、短语扩展或局部修饰词插入。在非自回归机器翻译(NAT)领域,学术界早已证明逐位置 CE 损失对轻微的位置偏移具有严厉的过惩罚效应——即便预测词语义完全正确,只要偏离了黄金位置一个槽位,就会受到与错误预测完全相同的梯度惩罚。更为致命的是,MDLM 采用的是多步迭代解码与低置信度重新掩码(remasking)机制,已揭示的高置信度标记会被永久锁定并作为后续所有去噪步的硬性上下文条件。这意味着在早期去噪中哪怕仅仅出现一个位置的微小偏移,该偏移就会作为错误上下文在后续数百步去噪中不断累积发酵。作者在 LLaDA-8B-Instruct 上的诊断性干预实验表明,仅仅人为将 1% 的生成标记向邻近位置移动一位,Arena-Hard 胜率就会从 50% 暴跌至显著劣势,直接证实了 MDLM 迭代解码对位置偏移具有高度敏感性。

受 NAT 领域通过松弛位置约束解决多模态对齐难题的启发,本文探索将连接时序分类(CTC)机制迁移至 MDLM 的监督微调阶段。然而,标准 CTC 的重复折叠规则会错误合并自然语言中固有的重复词(如 Markdown 语法中的加粗符号“**”被折叠为“*”、数字“100”被折叠为“10”),从而破坏文本的真实表面形态。核心 idea:针对掩码扩散语言模型在迭代解码中对位置偏移极度脆弱的痛点,引入专属 <slack> 空白标记对训练数据进行时间步自适应增强,并设计仅折叠空白、保全连续重复词形态的 CTC-S 目标函数与前向递推算法,将位置对齐灵活性作为辅助损失融入 SFT,从而在不改变标准推理管线的前提下赋予模型强大的位置容错与生成质量提升。

方法详解

整体框架

本文提出的方法在保持 MDLM 标准网络架构与迭代推理管线完全不变的前提下,通过对 SFT 阶段的数据构建与损失函数进行重构,使模型在训练阶段内化对位置不确定性的容忍度。该方法包含三个核心环节:首先,针对黄金响应序列,依据扩散时间步动态注入可随去噪进程衰减的 <slack> 空白标记;其次,构建专门适配自然语言文本生成的 CTC-S(Slack-only CTC)辅助损失函数,通过重构动态规划状态转移方程,仅折叠 <slack> 标记而保留连续相同的字符与子词;最后,将 CTC-S 辅助损失与原有的位置交叉熵损失进行加权联合优化。在推理阶段,模型直接沿用原有的低置信度重新掩码迭代去噪算法,生成结束后仅需机械剥离输出序列中的 <slack> 标记即可获得最终文本。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入:Prompt $q_0$ 与黄金响应 $r_0$"] --> B["动态插入 &lt;slack&gt; 的数据增强<br/>按时间步 $t$ 线性衰减插入空白占位符得到 $\tilde{r}_0$"]
    B --> C["扩散前向加噪生成掩码响应 $\tilde{r}_t$<br/>MDLM 并行预测各掩码位置词表分布"]
    C --> D["仅折叠空白的 CTC-S 目标与递推<br/>$\beta_{\text{slack}}$ 仅剔除空白 + M1/M2 状态递推约束"]
    C --> E["标准逐位置交叉熵损失 $\mathcal{L}^{\text{CE}}$<br/>在 $\tilde{r}_0$ 对应位置提供基础语义监督"]
    D --> F["结合 CE 与 CTC-S 的联合 SFT 目标<br/>加权联合优化 $\mathcal{L}^{\text{SFT}} = \mathcal{L}^{\text{CE}} + \lambda \mathcal{L}^{\text{CTC-S}}$"]
    E --> F
    F --> G["推理阶段:迭代去噪后直接剥离 &lt;slack&gt; 标记<br/>产出兼具语义一致性与位置弹性的最终文本"]

关键设计

1. 动态插入 的数据增强:为并行去噪留出位置吸附空间

标准 MDLM 训练序列由固定长度的离散 token 紧凑排列而成,任何位置的不确定性都无法在内部消化。为了给模型在输出序列中提供弹性对齐的“缓冲带”,作者向模型词表中引入了一个特殊的空白标记 <slack>。给定黄金响应序列 \(r_0 \in \mathcal{V}^L\) 和采样得到的扩散时间步 \(t \sim \mathcal{U}[0, 1]\),算法在 \(r_0\) 的内部随机位置动态插入 \(n\) 个 <slack> 标记,构建出扩展序列 \(\tilde{r}_0 \in (\mathcal{V} \cup \{\text{<slack>}\})^{\tilde{L}}\)。

为了匹配扩散模型从完全掩码(\(t=1\))逐步去噪收敛至清晰文本(\(t=0\))的内在规律,插入的 <slack> 比例被设计为随时间步 \(t\) 线性递减的函数 \(s(t) = s_{\max}(1 - t)\),其中最大松弛比例固定为 \(s_{\max} = 0.5\)。插入标记的数量通过 \(n = \min(\lfloor s_{\max}(1 - t) R \rfloor, A)\) 确定(其中 \(R\) 为原始响应长度,\(A\) 为序列最大长度 2048 下的可用余量,且序列末尾的 <EOT> 与 <EOS> 不参与插入)。这种设计保证了在去噪早期(\(t\) 较大时)<slack> 占比极低,模型集中注意力恢复粗粒度骨干语义;而在去噪后期(\(t\) 趋近于 0 时),充足的 <slack> 占位符为局部短语和单词的微小位置波动提供了充分的吸附空间,有效避免了因单个词位置轻微错位而导致整个后续序列被硬性截断或错位惩罚。

2. 仅折叠空白的 CTC-S 目标与递推:消除刚性位置惩罚并保全文本表面形态

经典语音与手写识别中的标准 CTC 依赖折叠映射 \(\beta_{\text{std}}\) 将潜在对齐序列映射回目标文本。\(\beta_{\text{std}}\) 的核心规则包括两步:先将连续相同的字符压缩为一个,再删除所有空白占位符。然而,这一经典规则若直接应用到自然语言生成中会产生致命缺陷——自然语言中广泛存在合法的连续重复标记(例如 Markdown 语法中的粗体标记 ** 会被压缩为单个 *,数字 100 会被截断为 10,多重标点 ... 会被破坏)。因此,作者提出了专门针对文本生成的松弛映射 \(\beta_{\text{slack}}\),定义其仅移除 <slack> 标记,而严格保留序列中相邻相同的真实内容标记。

为了在该新型折叠映射下高效计算序列似然边缘概率,作者推导了修正后的前向动态规划递推算法。将长度为 \(L\) 的目标响应 \(r_0\) 扩充为长度为 \(2L+1\) 的扩展状态序列 \(\ell = (\text{<slack>}, r_0^1, \text{<slack>}, r_0^2, \dots, r_0^L, \text{<slack>})\),其中偶数索引对应 <slack> 状态,奇数索引对应具体词标记状态。针对经典 CTC 前向算法,作者引入了两项关键的状态转移修正: - 规则 M1(自环转移限制在空白状态):标准 CTC 允许在所有状态上停留在自身,这依赖于后续的重复字符压缩机制;但在 \(\beta_{\text{slack}}\) 映射下不合并重复词,若允许在内容词状态自环,将导致模型在相邻步重复发射同一个词。因此,自环状态转移 \(a_{\text{self}} = \alpha[t-1, s]\) 仅在 \(s\) 为偶数(<slack> 状态)时合法,奇数内容状态必须强制步进(即 \(a_{\text{self}} = -\infty\))。 - 规则 M2(内容词状态无条件允许跳跃):标准 CTC 仅在相邻两个目标词不同时才允许跨过中间空白状态(即跳跃步 \(a_{\text{skip}} = \alpha[t-1, s-2]\)),目的是防止相同相邻词被压缩规则吞没;但在 \(\beta_{\text{slack}}\) 下无需压缩,跳过中间 <slack> 永远安全,因此对于任意 \(s \ge 2\) 且为奇数的内容词状态,无条件允许跳跃转移。

通过 M1 与 M2 修正,前向状态变量 \(\alpha[t, s]\) 在对齐空间内的对数累加能够精确覆盖所有坍缩为 \(r_0\) 的合法路径,既保留了 CTC 边缘化消除位置硬约束的核心数学优势,又彻底根除了破坏文本语意与标点形态的副作用。

3. 结合 CE 与 CTC-S 的联合 SFT 目标:兼顾局部逐词预测与全局序列弹性对齐

单纯依赖 CTC 损失进行序列生成虽然具备位置弹性,但容易在局部细节上缺乏足够的语义监督先验,且收敛初期梯度方差较大。因此,作者设计了结合逐位置交叉熵与全局 CTC-S 的双流联合微调策略。在每个训练样本上,标准 CE 损失在加噪序列 \(\tilde{r}_t\) 与带 <slack> 扩充的目标序列 \(\tilde{r}_0\) 之间计算,为每一个被遮盖的掩码位置提供直接、高效的局部词表重构信号;而 CTC-S 损失则以未扩充的纯净黄金响应 \(r_0\) 为对齐目标,在预测的离散分布矩阵上执行前向边缘化求和。

这一联合目标使模型形成了互补的自适应机制:CE 项促使模型学习在显式包含 <slack> 的语境中进行局部去噪,而 CTC-S 项则充当了全局软正则化器,赋予模型在生成内容略微发生左右槽位偏移时不遭受灾难性惩罚的自由度。

损失函数 / 训练策略

整个 SFT 阶段的综合优化目标定义如下:

\[\mathcal{L}_{\theta,t}^{\text{SFT}} = \mathcal{L}_{\theta,t}^{\text{CE}} + \lambda \mathcal{L}_{\theta,t}^{\text{CTC-S}}\]

其中 \(\lambda\) 为平衡两个损失项贡献的超参数,实验中直接固定为 \(\lambda = 0.1\),未做额外搜索。各项损失的具体计算方式如下: - 位置交叉熵项 \(\mathcal{L}_{\theta,t}^{\text{CE}}\):作用于加噪序列 \(\tilde{r}_t\) 与带 <slack> 扩充序列 \(\tilde{r}_0\): $\(\mathcal{L}_{\theta,t}^{\text{CE}} = -\mathbb{E}_{t, q_0, \tilde{r}_0, \tilde{r}_t} \left[ \gamma_t \sum_{i=1}^{\tilde{L}_r} \mathbf{1}_{[\tilde{r}_t^i = \text{M}]} \log P_\theta(\tilde{r}_0^i \mid q_0, \tilde{r}_t) \right]\)$ 其中权重系数 \(\gamma_t = -\frac{\alpha_t'}{1-\alpha_t}\)。 - CTC-S 辅助损失项 \(\mathcal{L}_{\theta,t}^{\text{CTC-S}}\):直接以纯净目标 \(r_0\) 为边界,在预测概率矩阵上边缘化所有满足 \(\beta_{\text{slack}}(a) = r_0\) 的潜在对齐路径: $\(\mathcal{L}_{\theta,t}^{\text{CTC-S}} = -\mathbb{E}_{t, q_0, r_0, \tilde{r}_t} \left[ \log \sum_{a \in \beta_{\text{slack}}^{-1}(r_0)} \prod_{i=1}^{\tilde{L}} P_\theta(a^i \mid q_0, \tilde{r}_t) \right]\)$ 该项通过前向递推算法在终止步精确求得:\(\mathcal{L}^{\text{CTC-S}} = -\text{logsumexp}(\alpha[\tilde{L}, 2L], \alpha[\tilde{L}, 2L-1])\)。

训练基于高质量的 Magpie-Pro 筛选子集(过滤响应长度在 512 至 1024 之间,共约 29.1 万条样本)。采用 LoRA 形式微调底座模型 LLaDA-8B-Instruct,设置秩 \(r=128\)、缩放因子 \(\alpha=32\)、dropout 为 0.05,适配器挂载于注意力的投影层、词嵌入层及 FFN 输出层。优化器使用 AdamW,最大学习率 \(2 \times 10^{-4}\),权重衰减 0.1,线性衰减无 warmup,共训练 3 个 epoch(约 3.6 万次更新步骤)。

实验关键数据

主实验

开放式文本生成评测覆盖 Arena-Hard-v2.0(以 CE-Only 作为基线模型计算相对胜率,由 GPT-5.4 进行确定性仲裁)、Creative Writing Bench v3、MT-Bench(GPT-4 评测)以及 WildBench。四组对比模型分别为:未作微调的官方底座 LLaDA、在相同数据上使用标准 CE 训练的基线 CE-Only、在带 <slack> 增强数据上仅使用硬位置 CE 训练的 Jitter 模型、以及本文提出的 CE + CTC-S 完整模型。

模型 Arena-Hard 胜率 (%) Creative-Writing-Bench v3 MT-Bench (得分) Wild-Bench (得分)
LLaDA (原始底座) 43.7 (-1.1/+1.4) 23.2 2.84 -5.48
CE-Only (SFT基线) 50.0 (-0.0/+0.0) 24.6 3.44 -4.53
Jitter (仅加数据扰动) 48.4 (-1.4/+1.5) 25.1 3.61 -4.46
CE + CTC-S (本文) 55.1 (-1.7/+1.3) 27.7 4.35 -4.23
较 CE-Only 显著提升 +5.1% +3.1 +0.91 +0.30

注:Arena-Hard 报告官方设定的 95% 引导置信区间。CE + CTC-S 在所有四项基准上相较于 CE-Only 和 Jitter 均取得了统计学显著优势(成对 bootstrap 检验,\(p < 0.05\))。

消融实验

1. 阶梯式消融对比(隔离数据增益、硬位置扰动与 CTC-S 目标)

为了严谨探究性能增益究竟来自额外 SFT 数据、<slack> 数据增强还是对齐弹性损失,作者构建了清晰的递进消融梯队:

配置 核心机制 Arena-Hard MT-Bench 核心结论分析
LLaDA 原始零样本/指令底座 43.7 2.84 未经 Magpie-Pro 针对性微调
CE-Only + 额外高质量 SFT 数据 50.0 3.44 数据带来基准提升,但仍受限于刚性位置监督
Jitter + <slack> 增强(硬位置 CE) 48.4 3.61 引入空间扰动但未松弛对齐,Arena-Hard 甚至略微下降
CE + CTC-S + 弹性 CTC-S 辅助损失 55.1 4.35 对齐灵活性是最大增益源,显著超越 Jitter 与 CE-Only

2. 训练侧对齐松弛 vs 推理侧分块扩散(Block Diffusion 扫描对比)

分块扩散(Block Diffusion)是常用的推理侧策略,通过将解码画布限制在较小的局部窗口(如 32~512)内以抑制位置偏移累积。下表对比了最佳超参数扫描下的分块扩散与本文单次非扫描 CTC-S 模型的效果:

模型配置 Creative-Writing-Bench v3 MT-Bench Wild-Bench 说明
LLaDA + Block Diffusion (最佳 \(b\)) 26.19 (\(b=512\)) 4.30 (\(b=32\)) -4.63 (\(b=64\)) 穷举扫描最佳分块大小
CE-Only + Block Diffusion (最佳 \(b\)) 26.44 (\(b=512\)) 3.53 (\(b=512\)) -4.51 (\(b=512\)) 穷举扫描最佳分块大小
CE + CTC-S (单一固定配置) 27.7 4.35 -4.23 无需任何推理侧分块搜索,全面超越最佳扫描基线

3. 跨架构泛化性:在 LLaDA-MoE 上的迁移表现

在稀疏混合专家扩散模型 LLaDA-MoE-7B-A1B-Instruct 上的复现测试显示:CE + CTC-S 在 Arena-Hard 胜率上达到了 63.7%(相比 CE-Only-MoE 的 50.0% 提升达 13.7 个百分点),在 Creative-Writing-Bench 上达到 25.1(提升 2.7),Wild-Bench 达到 -5.03(提升 0.39),证实该机制在 MoE 扩散架构上具有高度通用性。

关键发现

  • 对齐灵活性是主要增益源:消融梯队明确显示,仅引入 <slack> 数据扰动的 Jitter 模型在 Arena-Hard 上出现负向波动(从 50.0% 降至 48.4%),而加入 CTC-S 损失后大幅跃升至 55.1%。这证实单纯的数据增强无法解决根本矛盾,损失函数层面的动态对齐容错才是质变核心。
  • 训练侧松弛优于推理侧限制:即使对 baseline 耗费算力进行细致的分块大小扫描(\(b \in [32, 512]\)),其最佳成绩依然全面落后于未经推理调优的单一配置 CE + CTC-S,证明训练阶段学到的内在对齐弹性无法被简单的外部推理窗口截断所取代。
  • <slack> 标记在推理中的自发选择性与语义聚焦:
    • 按需触发:尽管训练样本 100% 包含 <slack>,但在 MT-Bench 生成中仅 51.9% 的样本在推理阶段输出了 <slack>,证明模型自主将其视作按需调用的“缓冲资源”,而非机械复制格式。
    • 词边界偏好:在生成的 <slack> 中,高达 95.1% 精确落在单词与词间边界处,仅 4.9% 出现在子词内部。这完全符合语言学规律——语义不确定性天然存在于词汇与短语衔接处,模型在无显式约束下自发学会了将容错空间分配至合理语言学边界。
  • 通用能力零退化:在 GPQA、MMLU、MBPP、IFEval 与 GSM8K 等通用学术与代码数学评测中,CTC-S 的得分表现(如 IFEval 66.9 vs CE-Only 63.2;GSM8K 68.9 vs CE-Only 66.1)完全持平或小幅优于 CE-Only 基线,表明引入位置弹性并未以牺牲严格推理能力为代价。

亮点与洞察

  • 独到的跨领域洞察(NAT ↔ MDLM):敏锐地捕捉到非自回归翻译(NAT)与掩码扩散语言模型(MDLM)在“并行解码 + 逐位置交叉熵”架构上的数学同构性,从而将 NAT 领域成熟的对齐理论创新性地移植到大模型扩散生成领域。
  • 精巧保全文本形态的 CTC-S 设计:彻底摆脱了经典 CTC 必须合并相邻重复字符的教条限制,创造性地重写前向算法中的 M1(自环状态限制)与 M2(无条件跳跃)状态转移规则,使得连续重复符号(如代码、Markdown 标记、数字)得以完美保留。
  • 训练开销极小且推理完全零负担:仅作为辅助损失存在于 SFT 阶段,不修改任何模型推理架构,生成完毕直接线性移除空白标记,在零额外推理延迟的前提下免费换取了显著的文本质量与鲁棒性提升。

局限与展望

  • 适用阶段仍局限于 SFT:目前工作仅在 SFT 阶段验证了 CTC-S 的有效性,未涉及计算量更大的预训练阶段。若能在预训练基座阶段内化对齐弹性,可能从根源上孕育出更具原生位置容错能力的扩散基座模型。
  • 超参数空间未做精细调优:损失平衡权重 \(\lambda=0.1\) 与最大插入比例 \(s_{\max}=0.5\) 均为启发式经验取值,缺乏针对不同任务类型或长度预算的系统性网格搜索或可学习预算机制。
  • 架构依赖性:实验主要围绕 LLaDA 系列展开。对于从自回归模型初始化并引入因果偏移监督的扩散模型(如 Dream 7B),该方法在何种程度上能够保持增益仍待验证。

相关工作与启发

  • vs 推理侧解码调整方法(如 Remasking、Pc-Sampler、Block Diffusion):以往针对 MDLM 生成缺陷的研究大多集中在推理阶段,通过调整取消掩码顺序、重新打掩码反思修正或缩减生成画布。本文与之截然不同,证明了通过在训练侧损失函数中注入对齐弹性,无需对推理算法作任何复杂改动即可取得更优的增益,且两者在理论上高度正交互补。
  • vs NAT 中的对齐松弛方法(如 OAXE、Vanilla CTC、DAT):NAT 中经典 CTC 会破坏文本的连续重复字符,而 OAXE 等基于最佳排列重排的方法计算复杂度高且依赖匈牙利匹配。本文提出的 CTC-S 针对自回归/非自回归大模型文本生成特性专门定制,具有精确的动态规划前向递推求解效率,展示了将经典序列转换理论现代化改造的典范价值。

评分

  • 新颖性: ⭐⭐⭐⭐ [成功跨越 NAT 与 MDLM 的鸿沟,提出保全表面形态的 CTC-S 递推与时间步自适应空白注入]
  • 实验充分度: ⭐⭐⭐⭐⭐ [覆盖 4 项开放式基准、5 项通用基准、成对显著性检验、MoE 架构迁移与深入的生成行为分析]
  • 写作质量: ⭐⭐⭐⭐⭐ [逻辑严密流畅,从诊断性干预切入,动机极其明确,算法公式与实验消融层层递进]
  • 价值: ⭐⭐⭐⭐ [为扩散语言模型的训练目标设计开辟了全新的训练侧对齐弹性维度,具有极高的实用启发性]