跳转至

Corrective Diffusion Language Models

会议: NeurIPS 2026
arXiv: 2512.15596
代码: https://github.com/zhangshuibai/CDLM
领域: LLM 评测与基准
关键词: 扩散语言模型, 错误定位, 原地自修正, 代码修正基准, 混合掩码训练

一句话总结

针对掩码扩散语言模型(MDLM)在未掩码可见位置缺乏监督导致无法辨别错误 token 的缺陷,本文提出校正导向的混合掩码后训练准则(CDLM),结合吸收掩码与均匀替换噪声赋予模型误差感知置信度,并构建可执行的代码修正基准(CRB),实现了精准的局部重掩码原地迭代修正与更优的高不确定性并行解码。

研究背景与动机

自回归(AR)语言模型凭借从左至右的因果生成模式在文本与代码生成中占据统治地位,但其内在的单向因果依赖导致后续预测强条件于早期输出,一旦前文生成发生错误,模型在向前推进过程中缺乏直接修改已有 token 的机制,难以开展原地自修正。近年来兴起的掩码扩散语言模型(MDLM)采用非因果双向结构与离散去噪动力学,天然解耦了不同位置的生成依赖,在理论上具备对已有序列进行任意位置局部微调、就地修剪与迭代渐进改进的结构优势。

然而实践表明,现有的标准 MDLM 几乎无法展现出可靠的原地修正能力。这一现象的深层矛盾在于现有训练目标与修正需求之间的脱节:标准 MDLM 严格遵循吸收掩码(absorbing-mask)重构目标,反向传播的交叉熵损失仅施加在被掩码的标记位置,而模型在注意力机制中能够直接看见的未掩码 token 则完全得不到显式监督。这就导致模型在面对包含拼写、变量名或逻辑错误的完整可见序列时,从未学过评估可见 token 的合理性,往往对上下文完全不匹配的错误 token 给出与正确 token 相同甚至极高(例如超过 0.9)的置信度。当推理阶段依赖 token 置信度阈值进行重掩码(remasking)重采样时,模型既无法定位错误位置,又容易盲目擦除正确上下文,导致迭代修正陷入恶性循环。

此外,社区长期缺乏针对原地修正行为的受控评测基准。主流基准大多基于前缀补全任务,无法评估模型在完整上下文中的错误定位与局部修复能力。为此,本文建立了一套完整的评测与后训练体系:一方面构建受控且可通过执行测试用例严格检验的 Code Revision Benchmark(CRB);另一方面重访吸收与均匀噪声的混合监督。核心 idea:提出校正扩散语言模型(CDLM),在标准吸收掩码重构之外,显式引入均匀替换噪声并监督可见错误标记,诱导模型对错误内容输出低置信度,从而激活基于置信度的精准局部重掩码原地迭代修正能力。

方法详解

整体框架

CDLM 的核心在于打破“可见 token 默认正确”的先验假设,将离散扩散的去噪过程与局部错误识别联合建模。方法整体分为三个紧密衔接的环节:首先,在后训练阶段采用吸收掩码与可见均匀替换噪声相结合的双阶段混合加噪机制,同时为模型提供缺失重构与错误判别的监督信号;其次,在推理阶段执行基于自适应置信度的动态重掩码迭代修正算法,通过对低置信度 token 进行局部掩码与并行重采样,实现无需全序列重生成的原地修复;最后,依托构建的 Code Revision Benchmark(CRB)提供包含操作符、标识符和字面量三类保类型扰动的可执行闭环评测。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入代码序列<br/>包含语法/语义错误的完整输入"] --> B["统一混合掩码监督机制<br/>吸收掩码 + 可见均匀噪声联合训练"]
    B --> C["自适应置信度动态重掩码迭代修正<br/>并行推理提取置信度 ci,阈值判断过滤"]
    C -->|ci < τ 置信度不足| D["局部原地重掩码与重采样<br/>保留高置信度上下文,重生成可疑标记"]
    C -->|ci >= τ 判定可靠| E["保留当前标记"]
    D --> F["代码修正基准 CRB 确定性评测<br/>保类型扰动与单元测试闭环执行校验"]
    E --> F
    F --> G["输出修正后的可执行正确程序"]

关键设计

1. 统一混合掩码监督机制:显式注入可见噪声诱导判别性置信度 针对标准 MDLM 仅在掩码位置计算损失导致可见错误 token 置信度虚高的根本痛点,CDLM 在训练阶段对真实序列引入两阶段扰动。首先以掩码率 \(\lambda \sim \text{Uniform}[0, 1]\) 随机选取位置集合 \(\mathcal{M}\) 替换为特殊吸收掩码标记 \(\langle\text{mask}\rangle\);接着,在其显式可见的未掩码位置集合 \(\mathcal{V}\) 中,以混合概率 \(\alpha \in (0, 1)\) 对每个 token 独立执行伯努利采样 \(z_i \sim \text{Bernoulli}(\alpha)\),命中的位置被强制替换为词表中的均匀随机噪声 token。在计算损失时,模型不仅要在掩码位置 \(\mathcal{M}\) 重构原词,还要在被均匀替换的可见位置 \(\mathcal{V}\) 处根据真实标签施加交叉熵损失惩罚。通过迫使模型纠正可见的异常 token,网络自发学会了:当某个可见标记与其双向上下文不兼容时,其最大 softmax 输出概率(即置信度)必须显著降低。这种机制在保留离散扩散生成能力的同时,赋予了模型天然的错误检测感知力。

2. 自适应置信度动态重掩码迭代修正:定位与原地渐进式修复 为将训练所得的判别性置信度转化为高精度原地修复,CDLM 采用基于动态阈值的迭代重掩码流程。在第 \(t\) 次修正步中,给定当前候选序列 \(\mathbf{z}^{(t)}\),模型双向并行计算各位置的条件词表概率分布 \(p_\theta(\cdot \mid \mathbf{z}^{(t)})\),并提取预测标记 \(\hat{x}_i^{(t)} = \arg\max_{v} p_\theta(v \mid \mathbf{z}^{(t)})\) 以及对应的置信度得分 \(c_i^{(t)} = \max_{v} p_\theta(v \mid \mathbf{z}^{(t)})\)。系统通过预设置信度阈值 \(\tau\)(默认取 0.9,高并行设置下取 0.6)筛选出不可靠集合 \(r^{(t)} = \{i \mid c_i^{(t)} < \tau\}\)。更新规则对集合 \(r^{(t)}\) 内的位置重新赋值为 \(\langle\text{mask}\rangle\),而高于阈值的位置则直接继承确定标记。随后模型仅针对重新被掩码的位置重新预测,其余正确上下文保持冻结。该过程循环执行 \(T\) 轮,逐步消除多处复合错误,既避免了自回归模型全量重新解码的算力开销,又杜绝了盲目改写正确代码导致的性能退化。

3. 代码修正基准 CRB:保类型扰动与确定性执行校验闭环 针对现有语言模型缺乏受控、细粒度且具有确定性判决的自修正基准,本文构建了 Code Revision Benchmark(CRB)。CRB 以 HumanEval、MBPP 及其扩展测试集 HumanEval+、MBPP+ 为原始语料,引入保类型替换算子 \(\phi(\cdot)\)。扰动分为三大类别:操作符替换(在算术与逻辑操作符集 \(\{+, -, *, /, \%, <, >, <=, >=, ==, !=\}\) 内互换)、作用域内标识符替换(将变量名或函数名替换为同作用域下的其他合法标识符)、以及字面量替换(数值或布尔常量互换)。扰动严格控制替换数量 \(|E| \in \{1, 2, 3, 4, 5\}\),以精细调节错误难度。更关键的是,CRB 引入确定性测试执行器:扰动后的代码若仍能通过单元测试则直接剔除,确保留存的每个样本都必然包含确定性的功能或语法缺陷。这为严格衡量置信度差距(Gap)、Top-\(K\) 命中率和 Pass@1 执行成功率提供了无偏差的物理地面真值。

一个完整示例

以一段简单的列表过滤函数为例,原始正确代码为 if val < threshold:,CRB 通过操作符替换算子引入单处错误,使其退化为 if val > threshold:,导致单元测试执行失败。 1. 初始评估:将错误代码输入标准 MDLM 时,模型对 > 处的置信度高达 0.92,而对周围正确 token 的平均置信度为 0.95,二者置信度差距(Gap)仅为 0.03,错误 token 在 Top-1 不确定性排序中排在无关位置; 2. CDLM 错误定位:输入 CDLM 后,由于模型在训练中学习过上下文与可见标记的冲突惩罚,对该逻辑冲突的 > 标记给出的置信度骤降至 0.28,而其他正常语法结构的置信度保持在 0.96,形成了高达 0.68 的置信度落差,准确将该位置置于不确定性 Top-1; 3. 动态重掩码:在设定阈值 \(\tau = 0.9\) 下,集合 \(r^{(0)}\) 精准命中该错误索引,序列变为 if val [MASK] threshold:; 4. 反向去噪采样:模型在此被掩码位置重新执行去噪预测,结合双向语义高概率重构出 <(置信度提升至 0.98); 5. 单元测试验证:经过 1 步迭代,修正后的程序提交测试执行器判定为 pass,完成原地无损修复。

损失函数 / 训练策略

CDLM 训练损失函数由标准吸收掩码损失与可见噪声监督损失两部分线性加权构成:

\[\mathcal{L} = \mathbb{E}_{\text{corruption}} \left[ \frac{1}{|\mathcal{M}|} \sum_{i \in \mathcal{M}} \ell_i + \lambda_{\text{noise}} \sum_{i \in \mathcal{V}} z_i \ell_i \right]\]

其中 \(\ell_i = -\log p_\theta(x_i^\star \mid \mathbf{z})\) 表示在位置 \(i\) 处基于真实干净标记 \(x_i^\star\) 计算的交叉熵损失,\(\mathcal{M}\) 为吸收掩码位置集合,\(\mathcal{V}\) 为未掩码可见位置集合,\(z_i \sim \text{Bernoulli}(\alpha)\) 控制可见位置是否遭受均匀噪声污染,\(\lambda_{\text{noise}}\) 为噪声项损失权重。在训练过程中,代码生成实验以 Open-dCoder-0.5B 为基座,在 Nemotron 代码语料库上微调 2,000 步;ParallelBench 并行解码实验以 LLaDA-2.0-mini(16B)为基座,采用块大小为 32 的 Block Diffusion 微调 6 个 epoch;受控数独(Sudoku)实验则使用 37M 参数的 DiT 模型从头训练 200,000 步。

混合概率 \(\alpha\) 的选取在训练中存在关键权衡:实验表明,当 \(\alpha\) 从 0 逐步增加至 0.5 时,干净 token 与错误 token 间的置信度差距稳步扩大,错误定位能力明显增强;但当 \(\alpha > 0.7\) 时,过量的均匀替换破坏了上下文语法连贯性,导致干净标记的误判率(False Positive Rate)急剧上升,端到端修复 Pass@1 显著衰减,因此通常设置平衡的适中比例。

实验关键数据

主实验

在 CRB 基准上,针对 4 个标准代码数据集(HumanEval、HumanEval+、MBPP、MBPP+)引入混合类型与多级别扰动,在 4 步迭代修正与置信度阈值 0.9 下对比 Base 基座、标准吸收微调模型(MDLM)与本文混合训练模型(CDLM)的 Pass@1 修复准确率。

数据集 指标 CDLM (本文) MDLM (吸收基线) Base (预训练基座) 相对 MDLM 提升
HumanEval Pass@1 0.257 0.157 0.168 +10.0% (+63.7%)
HumanEval+ Pass@1 0.242 0.182 0.196 +6.0% (+33.0%)
MBPP Pass@1 0.175 0.114 0.114 +6.1% (+53.5%)
MBPP+ Pass@1 0.226 0.137 0.120 +8.9% (+65.0%)
Average Pass@1 0.225 0.148 0.150 +7.7% (+52.0%)

除了代码修复,在高输出不确定性的大规模并行解码基准 ParallelBench(基于 LLaDA-2.0-mini 16B)上,对比不进行重掩码的固定解码(Fixed)与允许迭代修正的重掩码解码(Remask):

任务类型 具体任务 Base (Fixed / Remask) MDLM (Fixed / Remask) CDLM 本文 (Fixed / Remask)
低熵确定性任务 Copy 0.97 / 0.97 1.00 / 1.00 1.00 / 1.00
低熵确定性任务 Reverse 0.60 / 0.63 1.00 / 1.00 1.00 / 1.00
低熵确定性任务 Sort 0.15 / 0.19 0.61 / 0.61 0.60 / 0.64
低熵确定性任务 低熵平均 (Low-Avg) 0.45 / 0.47 0.93 / 0.93 0.93 / 0.94
高熵多样性任务 Shuffle 0.37 / 0.52 0.47 / 0.38 0.45 / 0.64
高熵多样性任务 Insert Random 0.23 / 0.35 0.40 / 0.44 0.40 / 0.72
高熵多样性任务 Remove Random 0.65 / 0.70 0.75 / 0.72 0.69 / 0.77
高熵多样性任务 Replace Random 0.51 / 0.67 0.25 / 0.26 0.30 / 0.37
高熵多样性任务 高熵平均 (High-Avg) 0.44 / 0.56 0.47 / 0.45 0.46 / 0.62
全任务总平均 Average (All) 0.45 / 0.50 0.75 / 0.74 0.74 / 0.81

消融实验

在 CRB 单错误场景(\(n_{\text{replace}}=1\))与从头生成场景下评估置信度质量与生成表现:

评估维度 / 设定 模型配置 置信度差值 (Gap) Top-1 命中率 Top-5 命中率 修正 Pass@1 (T=4)
CRB 单错误定位 (\(n_{\text{replace}}=1\)) Dream-7B-Base 0.578 0.094 0.590 0.732
CRB 单错误定位 (\(n_{\text{replace}}=1\)) LLaDA-8B-Base 0.235 0.178 0.770 0.596
CRB 单错误定位 (\(n_{\text{replace}}=1\)) Open-dCoder-0.5B (Base) 0.106 0.147 0.509 0.251
CRB 混合目标消融 (\(n_{\text{replace}}=1\)) MDLM-0.5B (吸收单目标) 0.112 0.155 0.520 0.268
CRB 混合目标消融 (\(n_{\text{replace}}=1\)) CDLM-0.5B (混合双目标) 0.385 0.312 0.795 0.421
从头代码生成 (HumanEval) 解码策略 Pass@1 Pass@10 HumanEval+ Pass@1 HumanEval+ Pass@10
Base 基座 Vanilla 0.19 0.37 0.17 0.32
MDLM (吸收微调) Vanilla 0.19 0.38 0.17 0.34
CDLM (本文混合微调) Vanilla 0.22 0.42 0.21 0.39
Base 基座 ReMDM 0.18 0.33 0.16 0.29
MDLM (吸收微调) ReMDM 0.20 0.38 0.17 0.34
CDLM (本文混合微调) ReMDM 0.21 0.43 0.20 0.38

关键发现

  • 可见噪声监督是误差感知的决定性来源:在未加入均匀噪声监督前,MDLM 的置信度差值极低(0.1 左右),错误 token 在 Top-1 仅有十几个百分点的捕获率;引入混合损失后,Gap 和 Top-K 捕获率成倍提高,从根本上激活了重掩码的准确性。
  • 高熵并行解码中重掩码机制的质变表现:在 ParallelBench 上,标准 MDLM 在面对高熵任务(如 Shuffle、Insert Random)时,启用 Remask 反而导致性能下降(从 0.47 跌至 0.45),原因在于失真的置信度诱导模型在重采样时误伤了合法分支;而 CDLM 凭借校准后的置信度,使 Remask 精度从 0.46 大幅跃升至 0.62,验证了错误感知对并行解码防坍缩的核心价值。
  • 从头生成与无损兼容性:在 HumanEval 从头生成和 Sudoku 纯掩码去噪实验中,CDLM 的 Pass@1/Pass@10 均超越或持平吸收基线,证明引入均匀噪声不仅没有牺牲传统掩码去噪的生成质量,反而提升了模型对抗解码随机扰动的泛化鲁棒性。

亮点与洞察

  • 点破掩码扩散语言模型的盲点:精准揭示了 MDLM “能看不能改”的结构性悖论——非因果结构赋予了就地编辑的理论能力,但掩码重建训练目标却完全忽视了可见标记的合理性判断,导致置信度与正确性严重失准。
  • 极简而有效的后训练干预:无需引入复杂的额外判别器或外部评分网络,仅仅在离散去噪分布中混合少量的可见均匀替换噪声并在原词上施加交叉熵,便零成本诱导出了极强的内在误差判别置信度。
  • 高并行解码的稳定性基石:将离散自修正能力与加速并行解码场景相联通,证明了在高熵生成中,决定并行效率上限的不仅是预测速度,更是模型能否利用置信度精准剔除并行冲突错误。

局限与展望

  • 受限于固定长度替换编辑:当前基准与方法主要针对固定长度的 token 级替换(In-place Replacement),无法直接应对代码编写中频繁出现的行级插入、多 token 删除或任意长度变长编辑(Variable-length editing)。
  • 均匀噪声与真实语义 Bug 的分布差距:训练中采用词表均匀随机采样作为噪声注入手段,虽然实现简单,但与真实编程中微妙的边界条件错误或复杂 API 滥用存在一定语义鸿沟,未来可探索基于 AST 语法树或代码变异算子的对抗性加噪。
  • 模型规模验证范围:当前端到端实验主要在 0.5B 和 16B 规模展开,在百亿级(70B+)超大扩散语言模型上的扩展规律和计算成本仍有待进一步系统化验证。

相关工作与启发

  • vs GIDD (Generalized Interpolating Discrete Diffusion):GIDD 虽然同样探讨了吸收与均匀噪声的结合,但其出发点是泛化连续与离散扩散的加噪轨迹,采样时依然采用全序列全局演化的扩散更新;CDLM 则立足于掩码扩散体系,将均匀噪声聚焦于后训练阶段以显式诱导局部错误置信度,推理时遵循严格的局部重掩码原地替换。
  • vs ReMDM 与外部评分修正:ReMDM 等前期方案依赖于推理阶段冻结初次解码置信度或借助外部验证器模型来选择重掩码 token,属于推断期启发式修补;CDLM 则是从模型参数本质层面重构置信度校准,使单模型自身兼具生成者与审计者的双重能力。
  • vs 自回归模型自修正(Self-Correction):AR 模型的自修正多依赖多轮思维链 Prompting 或外部测试报错反馈重新生成完整代码片段,耗费长上下文且极易产生推理幻觉;CDLM 展现了非因果架构特有的极简原地修改与毫秒级并行修剪潜力。

评分

  • 新颖性: ⭐⭐⭐⭐☆ 深刻揭示了 MDLM 无法自修正的根本成因,以巧妙的双重噪声混合优雅诱导了误差感知能力。
  • 实验充分度: ⭐⭐⭐⭐⭐ 从可执行的代码基准 CRB、到从头代码生成、高熵并行变换乃至受控数独,实验链条极为完备严密。
  • 写作质量: ⭐⭐⭐⭐⭐ 论点清晰,问题定义数学化且逻辑严密,图表对比极具说服力。
  • 价值: ⭐⭐⭐⭐⭐ 为扩散语言模型走向可靠的交互式编辑、代码自动修复与高效并行解码奠定了关键理论与算法基础。