跳转至

Constrained Code Generation with Discrete Diffusion

会议: NeurIPS 2026
arXiv: 2605.16829
OpenReview: https://openreview.net/forum?id=zoVQ8aHNUH
领域: 代码智能
关键词: 离散扩散模型、代码生成、约束解码、神经符号、程序分析

一句话总结

针对现有代码生成依赖自回归后验修复或语法级约束的缺陷,本文提出免训练的神经符号框架 CDC,在离散扩散去噪轨迹中利用软执行代理与静态程序分析定位冲突并局部纠偏,显著提升生成代码的功能正确性与安全性。

研究背景与动机

代码生成在程序合成、开发者助手及智能体工作流中展现出巨大的实用价值。近年来,离散掩码扩散模型(如 DiffuCoder 和 Dream-Coder)逐渐成为一种极具潜力的新范式,它通过对部分损坏的 token 序列进行迭代式并行去噪来逐步构建程序。然而,生成模型的概率拟合目标并不能天然保证代码的正确性与安全性;在高概率的生成输出中,依然普遍存在语法报错、逻辑缺陷乃至安全漏洞。

现有可控代码生成技术大多构建于自回归(Autoregressive, AR)模型之上,往往采用从左到右的前缀约束、重排序、拒绝采样或事后(post-hoc)重提示修复机制。然而,语义正确性与安全性等程序级约束通常取决于全局数据流或复杂的运行时行为,无法仅靠单步局部 token 约束进行规避。一旦生成失败,自回归模型往往需要重新生成大段后缀乃至完整程序,这在极其庞大的组合搜索空间中不仅计算开销高昂,而且重写引发退化的风险极高。相比之下,离散扩散模型在每个去噪步都直接暴露当前程序的全局表示状态,这一结构为在代码最终定稿前介入并施加全局程序级约束提供了天然的干预切入点。

本文的核心切入点是将扩散去噪过程视作一条可在线编辑的程序轨迹,在采样中间步直接评估约束并实施局部干预。核心 idea:提出免训练的神经符号去噪框架 CDC,将程序级硬/软约束评估分解为敏感区域定位与局部信任域纠偏,分别通过可微代理梯度引导(GradGuide)和静态分析反馈注入(MDFI)在去噪阶段精准重塑生成轨迹。

方法详解

整体框架

CDC(Constrained Diffusion for Code)在保持底层基础去噪扩散模型参数完全冻结的前提下,于反向去噪采样过程中串联四个核心步骤: 1. Clean-State Proposal(全状态候选生成):在时间步 \(t\),基于当前部分掩码序列 \(\bm{x}_t\),扩散去噪器输出对全序列 clean token 的类别分布预测 \(\hat{\bm{x}}_0^{(t)}\),并贪婪解码出中间程序候选 \(\bar{\bm{x}}_0^{(t)}\)。 2. Constraint-Aware Localization(约束感知敏感区定位):约束评估器 \(\mathcal{E}\)(学习型软代理或符号型静态分析器)评估中间程序候选,得到违反评分与结构化反馈,由定位映射 \(\mathcal{M}\) 将反馈精准映射为负责违规的代码 token 敏感区间 \(\mathcal{S}_t\)。 3. Constraint-Aware Correction(约束感知纠偏操作):约束去噪算子 \(\mathcal{P}_{\mathcal{C}}\) 仅在敏感区间 \(\mathcal{S}_t\) 及其周边实施受控纠偏,求解以 KL 散度为针定项的局部信任域优化问题,产出修正分布 \(\bm{y}_t\) 或局部重置掩码。 4. Constrained Reverse Update(受限反向更新):反向转移核使用修正后的分布与掩码状态完成 \(t \to t-1\) 的步进采样。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["中间掩码状态 x_t"] --> B["Clean-state 提案预测<br/>去噪器生成全局分布及中间程序"]
    B --> C["评估器与敏感区间定位<br/>软代理梯度或符号程序图分析"]
    C --> D["局部纠偏与状态更新<br/>KL 针定优化 / 局部重掩码与注入"]
    D --> E["反向转移更新 x_{t-1}<br/>沿修正轨迹继续去噪至 x_0"]

关键设计

1. 约束敏感区间定位:将程序级反馈解耦映射为代码 Token 区间 程序级约束通常由特定代码片段引发(例如缺失校验的危险 sink 点或导致逻辑错误的局部计算),全盘重写既低效又容易引入二次错误。CDC 将“定位何处(Where)”与“如何修正(How)”解耦。针对功能正确性,GradGuide 通过可微代理模型关于软 token 嵌入的反向传播梯度范数,结合模型预测的不确定性(熵与置信度差),筛选显著性最高的 top-\(k\) 位置,并调用语法扩展操作将其延展为语法完整的行或表达式区间;针对安全性,MDFI 利用增量解析器将中间代码解析为部分代码属性图(CPG),通过源-汇(source-to-sink)有界数据流分析或局部 AST 模式匹配定位漏洞 witness 节点,并依据分析置信度在预算内提取局部语法与数据流邻域区间 \(\mathcal{S}_t^{\mathrm{MDFI}}\)。

2. GradGuide:基于可微执行代理的 KL 针定投影与自适应重掩码 单元测试通过与否本质上是不可微的硬约束,去噪过程中无法直接求导。GradGuide 离线微调一个基于 Qwen2.5-Coder-1.5B 的判断代理 \(g_\phi\),以执行通过率作为监督信号,在推理时将连续的预测软分布映射到软嵌入空间并计算违规惩罚 \(V_t^{\mathrm{GG}}\)。随后,GradGuide 联合执行双重干预机制: 一方面在未被固定(未去噪完毕)的 token 上实施基于增广拉格朗日法(ALM)的分布修正,在保持与基础去噪器分布 KL 散度约束的前提下最小化代理违规惩罚: $$ \mathcal{L}{\mathrm{ALM}}\bigl(\bm{y};\hat{\bm{x}}}^{(t)},\mathcal{S{t},c\bigr)=D}}!\left(\bm{y}\,|\,\hat{\bm{x}{0}^{(t)}\right)+\sum}^{m}\left[\lambda_{j}\Delta g_{j}(\bm{y},c)+\frac{\mu_{j}}{2}\Delta g_{j}(\bm{y},c)^{2}\right]+\beta\sum_{i\notin\mathcal{S{t}}D\right) $$ 另一方面引入自适应重掩码(Mode B):当去噪进行到后半段但中间代码仍违反约束时,强制将敏感区间 }}!\left(\bm{y}^{i}\,|\,\hat{x}_{0}^{(t),i\(\mathcal{S}_t\) 中已固化的 token 重置为 [MASK],给予扩散模型重新在全局上下文约束下填补修复该局部的机会。

3. MDFI:基于符号程序图分析的掩码重置、插入与反馈注入 针对形式化严格的程序安全约束,MDFI 避免使用容易产生幻觉的软代理,转而采用纯符号静态分析。在预设的去噪检查点,中间代码中的未定 mask 被占位标识符替换后送入 Tree-sitter 增量解析。分析器输出漏洞见证元组 \((n_k, \tau_k, h_k)\):对于替换型漏洞(\(\tau_k=\mathrm{sub}\),如错误 API 调用),MDFI 将对应 token 节点重置为 [MASK];对于缺失型漏洞(\(\tau_k=\mathrm{ins}\),如缺少边界校验或入参清洗),MDFI 在语法锚点处就地插入 \(K\) 个全新的 [MASK] 槽位以扩充序列长度;同时,预分配的 Prompt 缓冲区被写入精简的自然语言修复建议(如推荐的安全函数模板),引导剩余去噪步精准完成安全修补,全过程无需外层多轮循环或模型重训练。

一个完整示例

以生成安全的 C++ 路径读取函数为例,中间去噪步解码出的候选代码包含直接拼接用户输入的文件操作 open(user_path): 1. MDFI 在去噪检查点触发静态分析,部分代码属性图识别出外部入参流向未验证的敏感 sink,产出插入型漏洞见证:需要添加路径规范化检查。 2. 定位模块锁定 open 调用行前方的语句边界为锚点位置,就地插入 \(K=12\) 个 [MASK] 标记,并同步将提示词缓冲区覆写为安全修复建议 c*:“validate path before open”。 3. 在随后的反向去噪步中,扩散注意力机制在新增的掩码槽位内顺利生成路径前置校验逻辑 if (!isValid(user_path)) return;,其余无缺陷的代码结构完全保持原样,仅耗费少量局部 token 便消除了路径遍历漏洞。

实验关键数据

主实验

在功能正确性(HumanEval-X C++、MBPP C++)以及安全性与功能联合评测(CWEval、LLMSecEval+)上,CDC 显著优于标准扩散解码、语法约束扩散(CFG-CD)及自回归模型(AR)。

数据集 模型与配置 编译成功率 (%) 功能 pass@1 (%) 功能 pass@10 (%)
HumanEval-X C++ Dream-Coder 7B (Vanilla) 67.1 34.1 55.5
HumanEval-X C++ Dream-Coder 7B + CFG-CD 61.6 (-5.5) 37.8 (+3.7) 53.7 (-1.8)
HumanEval-X C++ Dream-Coder 7B + CDC (Ours) 79.2 (+12.1) 65.2 (+31.1) 83.5 (+28.0)
HumanEval-X C++ DeepSeek-Coder 6.7B (Vanilla) 94.5 59.8 80.5
HumanEval-X C++ DeepSeek-Coder 6.7B + Reprompt 90.9 (-3.6) 64.0 (+4.2) 82.9 (+2.4)
MBPP-C++ Dream-Coder 7B (Vanilla) 41.1 27.7 54.2
MBPP-C++ Dream-Coder 7B + CFG-CD 37.3 (-3.8) 25.7 (-2.0) 49.1 (-5.1)
MBPP-C++ Dream-Coder 7B + CDC (Ours) 72.0 (+30.9) 59.2 (+31.5) 74.8 (+20.6)
MBPP-C++ DeepSeek-Coder 6.7B (Vanilla) 88.2 59.7 73.0
MBPP-C++ DeepSeek-Coder 6.7B + Reprompt 89.2 (+1.0) 60.7 (+1.0) 74.1 (+1.1)

在安全约束场景(以 Dream-Coder 7B 为扩散底座),CDC 在 CWEval 上的联合指标 func-sec@1(既正确又安全)从 12.0% 激增至 34.3%(提升 +22.3%),在 LLMSecEval+ 上安全合格率 sec@1 达到 80.7%,超越了带事后重提示的强自回归模型(64.7%)。

消融实验

消融实验验证了 GradGuide 各组件协同效应与局部编辑效率(基于 HumanEval-X C++,Dream-Coder 7B,贪婪解码):

配置 pass@1 (%) 说明
Vanilla 34.1 基线无干预采样
ALM-only (Mode A) 40.2 仅使用 KL 针定分布修正 (+6.1%)
Adaptive Editing-only (Mode B) 50.6 仅使用自适应局部重掩码 (+16.5%)
独立叠加期望预测 56.7 假设两模块完全独立线性叠加的预期值
Full Operator (CDC) 65.2 完整算子,展现 +8.5% 超加性协同增益
随机区域编辑 (Random Editing) 34.1 验证敏感区间精准定位的决定性作用

关键发现

  • 双模态超加性协同:单纯使用 Mode A 提升 6.1%,单纯使用 Mode B 提升 16.5%,而完整组合达到 65.2%,超出独立叠加预期达 8.5 个百分点。这是因为 Mode B 重新开放了解题通道,而 Mode A 指导了掩码填补时的概率流向,两者形成完美互补。
  • 极高的手术式编辑效率:在 CWEval 上,自回归重提示(AR Reprompt)平均需重写 113 个 token(占完整程序大部分结构),而 CDC 仅需局部编辑 12 个 token;在 HumanEval-X 上,CDC 纠偏仅修改中位数 32 个 token,远低于 AR 重写的 123 个 token。
  • 建设性转化率极高:在修复失败样例时,CDC 的建设性转化率(修复成功占所有判题变动的比例)在 HumanEval-X 上高达 94.2%、在 MBPP 上高达 96.9%,而 AR Reprompt 仅约为 53%(接近抛硬币),有效避免了全局重写破坏已有正确代码的退化问题。

亮点与洞察

  • 全局状态可见性与在线干预范式:巧妙抓住了扩散模型在去噪过程中持续暴露全局未定代码这一特性,突破了自回归模型“单向生成即定稿、出错只能事后全盘重写”的结构性桎梏。
  • 软硬双路神经符号协同设计:对无法精确解析的复杂功能逻辑使用可微执行代理进行梯度导航,对形式化严格的漏洞特征使用符号代码属性图精准匹配,兼顾了端到端可微性与静态分析的高可靠性。
  • 可复用的局部插桩与重掩码机制:在掩码状态中原位插入 [MASK] 并配合提示词缓冲区动态写入见证建议的技术,可直接迁移至各类离散扩散模型的控制生成任务(如受控文本编辑、格式化数据生成)。

局限与展望

  • 作者指出的局限:GradGuide 的效果高度依赖代理模型的评判精度,可能受到代理幻觉或预测噪声干扰;MDFI 依赖静态分析规则,面对复杂跨过程数据流可能漏报或误报;此外,推理时引入代理反传和增量解析带来了一定的计算延迟。
  • 延伸发现的局限:目前的评测主要局限在函数级短代码(HumanEval/MBPP/CWEval),对于大规模工程级仓库、多文件依赖及复杂构建系统的支持尚未得到充分验证。
  • 改进思路:可探索将运行时即时编译/轻量单测执行(如快速沙盒执行测试用例)直接作为中间检查点的非可微反馈源,并进一步结合投机采样(Speculative Sampling)摊薄评估延迟。

相关工作与启发

  • vs CFG-CD(Mündler et al., 2025):CFG-CD 将上下文无关文法约束集成到扩散解码中,仅能确保语法层面的编译合法性(如括号匹配、关键词合法),对运行时功能测试提升极其有限;CDC 直面语义级与安全级约束,不仅编译率提升,功能 pass@1 更实现翻倍跃升。
  • vs 自回归后验修复(Self-Debug / AR Reprompt, Chen et al., 2023):自回归修复在整段代码生成完毕后重新调用大模型整段重写,生成成本高且容易在修复旧 bug 的同时引入新 bug;CDC 在去噪轨迹中精准定位敏感 token 并实施局部重掩码,修改代价降低数倍且单调改善率达 94%+。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次系统性地将程序级语义与安全约束注入离散扩散语言模型的反向去噪过程,提出的软硬双路纠偏机制构思极佳。
  • 实验充分度: ⭐⭐⭐⭐⭐ 覆盖代码编译、功能正确性、安全漏洞规避、编辑局部性及组件消融,对比实验严密且指标充实。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰,数学形式化(KL 投影、拉格朗日乘子更新、图见证映射)与工程实现兼备。
  • 价值: ⭐⭐⭐⭐⭐ 为离散扩散模型在可信可控代码智能领域的应用树立了新标杆,极具启发意义与实用价值。