基于离散扩散的约束代码生成

# 基于离散扩散的约束代码生成 来源：https://arxiv.org/html/2605.16829 Lize Shao 弗吉尼亚大学 zgr3et@virginia\.edu&Michael Cardei¹¹footnotemark:1 弗吉尼亚大学 ntr2rm@virginia\.edu&Zichen Xie 弗吉尼亚大学 graysonxie@virginia\.edu&Ferdinando Fioretto 弗吉尼亚大学 fioretto@virginia\.edu&Wenxi Wang²²footnotemark:2 弗吉尼亚大学 wenxiw@virginia\.edu ###### 摘要 离散扩散模型是一种强大且新兴的代码生成范式。它们通过迭代精炼部分损坏的 token 序列来构建程序，并支持并行 token 精炼。重要的是，该范式在每个去噪步骤中都暴露了全局程序状态，这为强制执行程序级功能性和安全性约束提供了一个自然的干预点，在最终代码被确定之前引导生成过程。基于这一观察，本文引入了Constrained Diffusion for Code（CDC），一种无需训练的神经符号推理框架，它将约束满足直接集成到反向去噪过程中。CDC 使用约束感知的去噪算子来增强基础离散扩散采样器，这些算子结合数学优化和程序分析来识别中间程序状态中与约束相关的区域，并局部调整去噪轨迹，在保持接近基础模型的同时将生成过程引导向可行的程序。在代码生成基准测试中，CDC 持续提高了功能性正确性、安全性甚至语法方面的约束满足度，以更少的纠正计算和更局部的编辑改进了离散扩散和自回归基线。参见说明图1：CDC 与其他扩散代码基线在 HumanEval-X（HE-X）、MBPP、CWEval 和 LLMSecEval+ 上的比较。 ## 1 引言 代码生成已成为现代生成建模系统中越来越重要的能力，在程序合成、开发者辅助和自动化软件工作流中有着显著应用[16 (https://arxiv.org/html/2605.16829#bib.bib23),33 (https://arxiv.org/html/2605.16829#bib.bib5),31 (https://arxiv.org/html/2605.16829#bib.bib6)]。在此背景下，离散扩散语言模型正成为一种极具前景的范式，通过迭代精炼部分损坏的 token 序列来构建程序[10 (https://arxiv.org/html/2605.16829#bib.bib26),29 (https://arxiv.org/html/2605.16829#bib.bib24)]。然而，用于训练生成模型的目标函数本身并不包含使生成的代码正确或安全所需的条件。这些模型旨在匹配目标分布并生成高保真样本，但一个程序在基础模型下可能看起来非常合理，却仍然可能产生错误行为，或者包含可利用的漏洞。实际上，生成的代码必须满足明确的程序级约束，包括语法有效性，以及更重要的功能性正确性和安全属性。由于程序是可执行对象，违反这些约束可能直接导致运行时错误、错误输出或安全故障[21 (https://arxiv.org/html/2605.16829#bib.bib7),3 (https://arxiv.org/html/2605.16829#bib.bib8)]。因此，整合此类约束对于代码质量和安全性都至关重要。 这些要求暴露了代码生成的一个核心挑战：*程序级约束通常依赖于从任何单个局部 token 决策中都无法看到的全局结构或语义属性*。更广泛地说，常见的引导机制，如提示、重排序、拒绝采样或事后修复，通常只在候选程序生成完成后才进行干预[27 (https://arxiv.org/html/2605.16829#bib.bib9),8 (https://arxiv.org/html/2605.16829#bib.bib10)]。在代码生成中，这尤其具有局限性，因为底层搜索空间高度组合，而满足语义或安全约束的程序子集在模型分布下可能只构成一个很小的可行区域。因此，事后方法在计算上可能既低效又不可靠，特别是当约束满足需要将生成过程从模型主要的无约束模式转向更罕见但可行的解决方案时。 相比之下，本文表明，离散扩散模型提供了另一种机会：它们通过中间去噪状态来精炼整个序列的方法，在整个采样过程中暴露了全局程序表示。这种结构创造了自然的干预点，用于在去噪过程中将约束直接集成到生成的代码中。基于这一机会，本文引入了Constrained Diffusion for Code（CDC），一个无需训练的神经符号框架，它将程序级约束（即功能性和安全性）直接集成到离散扩散模型的反向去噪过程中。CDC 将扩散采样视为一个可编辑的程序轨迹：在每个去噪步骤中，模型提出一个完整的干净状态程序分布；CDC 解码这个提议，评估得到的候选程序是否满足程序级约束，将反馈定位到相关区域，并在继续采样之前应用约束感知的纠正。每次纠正都遵循一个局部化的约束采样原则：在减少约束违反的同时保持接近基础去噪器，仅修订与约束相关的区域，而不必要地扰动程序的其他部分。该框架支持软反馈和硬反馈。对于功能性正确性，CDC 将该原则实例化为 GradGuide，它使用软替代执行信号来定位可能的错误并指导有针对性的去噪更新。对于安全性，CDC 将其实例化为 MDFI，它使用静态程序分析来识别易受攻击的区域，并通过局部重掩码、插入和反馈注入来纠正它们。通过这种方式，CDC 利用扩散采样暴露的全局、可编辑状态，在生成过程中执行有针对性的约束纠正，而不是依赖于事后过滤或失败后的整个程序重新生成。 本文做出了三项主要贡献： 1. 1\.约束感知的定位。它引入了约束扩散代码生成的通用定位机制，将来自替代模型或静态分析器的程序级反馈映射到在去噪期间应该修订的 token 区域。 2. 2\.约束感知的纠正。它提出了无需训练的纠正算子，这些算子在反向扩散期间修订这些局部区域，同时保持接近基础去噪器。这些算子支持通过梯度和 KL 锚定更新进行可微的替代模型引导，以及通过定向重掩码、掩码插入和反馈注入进行不可微的符号分析器反馈。 3. 3\.经验评估。它提供了经验证据，表明去噪时约束集成在功能性正确性和面向安全性的基准测试上改善了约束代码生成。如图 LABEL:fig:strong_results 总结所示，CDC 在 HumanEval-X C++ 上将功能性正确率从 34.1% 提高到 65.2%，在 MBPP-C++ 上从 27.7% 提高到 59.2%（相比 Dream-Coder 7B），而 MDFI 在 CWEval 上将联合功能-安全性成功率从 12.04% 提高到 34.26%（相比最先进的扩散模型），且编辑的 token 数量少得多。尽管 CDC 针对功能性和安全性，但它也提高了语法正确性：在 HumanEval-X C++ 上从 67.1% 提高到 79.2%，在 MBPP-C++ 上从 41.1% 提高到 72.0%。 ## 2 相关工作 离散扩散语言模型通过迭代去噪而非从左到右解码来生成序列，从而在整个采样过程中暴露部分生成的全局状态。最近面向代码的扩散模型，包括 DiffuCoder[10 (https://arxiv.org/html/2605.16829#bib.bib26)] 和 Dream-Coder 7B[29 (https://arxiv.org/html/2605.16829#bib.bib24)]，表明该范式能够实现有竞争力的代码生成性能，同时支持灵活的非自回归精炼。CDC 建立在此属性之上：它不仅将扩散用作替代解码器，还将中间去噪状态视为可编辑的程序表示，在生成完成之前就可以评估、定位和纳入约束。 可控代码生成主要在自回归模型中通过约束解码、重排序、拒绝采样、验证和迭代修复进行研究[14 (https://arxiv.org/html/2605.16829#bib.bib12),19 (https://arxiv.org/html/2605.16829#bib.bib13),8 (https://arxiv.org/html/2605.16829#bib.bib10),18 (https://arxiv.org/html/2605.16829#bib.bib15),4 (https://arxiv.org/html/2605.16829#bib.bib14)]。这些方法可以提高语法有效性、类型一致性、功能性正确性或安全性，但它们通常在下一个 token 级别或完整候选程序生成之后进行操作。这对程序级约束来说是有限制的：语义和安全违规通常依赖于全局数据流或执行行为，事后修复可能需要重新生成大的后缀甚至整个程序。而 CDC 则在去噪过程中注入反馈，此时完整的程序状态可见，并且局部区域仍然可以修订。 最近的工作探索了通过分类器引导、无分类器引导、隐藏状态优化和基于投影的约束采样来实现离散扩散模型的可控生成[26 (https://arxiv.org/html/2605.16829#bib.bib16),11 (https://arxiv.org/html/2605.16829#bib.bib17),6 (https://arxiv.org/html/2605.16829#bib.bib4),9 (https://arxiv.org/html/2605.16829#bib.bib18)]。最接近我们设置的是 Mündler 等人[17 (https://arxiv.org/html/2605.16829#bib.bib28)]的工作，他们在扩散-LLM 解码过程中强制执行上下文无关文法约束，从而提高了代码的语法正确性。CDC 针对更广泛的程序级约束，包括通常不能表示为文法成员资格或局部 token 限制的功能性和安全性要求。CDC 不是应用全局引导或仅限语法的过滤，而是将替代模型或分析器的反馈定位到与约束相关的代码区域，并在反向扩散轨迹中修订这些区域。 ## 3 问题设置 带有目标规范的代码生成。令 x0 = (x0^1, ..., x0^L) 表示一个长度为 L 的离散 token 序列，其中每个 token 取值于词汇表 V。在代码生成中，V 可能包括编程语言关键字、标识符、字面量、标点符号、运算符和特殊格式 token（如 def, if, =, (, ), :, 变量名）。代码生成模型定义了给定上下文 c 下的 token 序列条件分布，其中功能性或安全性要求可以通过问题描述、函数签名、部分程序、单元测试或其他任务特定信息来指定。我们将这些要求称为*目标规范*；图 2 显示了一个示例。因此，目标是在学习到的代码分布下生成一个既合理又满足目标规范的序列 x0。 带硬约束和软约束的约束采样。令 p_θ(x | c) 表示一个基础代码生成模型。无约束解码从 p_θ 中采样可能的程序；约束解码则定义了一个采样器 p_θ^C(x | c)，它在保持接近基础模型的同时，将生成过程引导向由任务指定的可行集 C(c)。我们区分两种约束。*硬约束*是必须精确满足的可行性条件，例如解析、编译、通过所需测试或满足安全检查。*软约束*在达到精确可行性之前提供分级反馈，例如替代正确性分数、部分测试通过率或分析器严重性分数。CDC 同时使用两者：硬约束定义目标可行集，而软约束则在去噪过程中提供有用的引导。 ## 4 预备知识：掩码扩散模型 掩码扩散语言模型通过逆转离散噪声过程来生成序列[32 (https://arxiv.org/html/2605.16829#bib.bib25),25 (https://arxiv.org/html/2605.16829#bib.bib2),10 (https://arxiv.org/html/2605.16829#bib.bib26)]。从一个干净的 token 序列 x0 开始，前向过程逐渐将 token 破坏成一个吸收掩码 token [MASK]；一旦一个 token 被掩码，它将在后续时间步保持掩码状态[25 (https://arxiv.org/html/2605.16829#bib.bib2),20 (https://arxiv.org/html/2605.16829#bib.bib27)]。对于每个位置 i ∈ [L]，前向转移为 q(x_t^i | x_{t-1}^i) = Cat( x_t^i; α_{t|t-1} x_{t-1}^i + (1 - α_{t|t-1}) M ), (1) 其中 Cat(·; π) 是分类分布，M 是 [MASK] 的独热表示，并且 α_{t|t-1} := α_t / α_{t-1} 对应于递减的噪声调度，其中 α_0 = 1, α_T = 0。因此，未掩码的 token 以概率 α_{t|t-1} 被保留，否则被掩码。 反向过程使用一个神经去噪器来填充掩码位置。给定时间步 t 的部分掩码序列 x_t，去噪器预测一个干净 token 分布 x̂_0^{(t)} = x_θ(x_t, t)。反向转移为 p_θ(x_{t-1}^i | x_t) = { Cat(x_{t-1}^i; x_t^i), 如果 x_t^i ≠ [MASK], Cat(x_{t-1}^i; γ_t M + η_t x_θ^i(x_t, t)), 如果 x_t^i = [MASK] }, (2) 适用于每个位置 i ∈ [L] 和时间步 t = T, ..., 1，其中 γ_t = (1 - α_{t-1})/(1 - α_t), η_t = (α_{t-1} - α_t)/(1 - α_t) 平衡掩码保留和从去噪器的干净 token 预测中采样。在推理时，从完全掩码的序列 x_T 开始采样，并重复应用公式 (2) 直到获得生成的序列 x_0。 ## 5 面向代码生成的约束扩散 参见说明图 2：CDC 概述。 上述公式化引出了*面向代码的约束扩散*（CDC）：在每个时间步，去噪器提出一个完整的干净程序分布 x̂_0^{(t)}，这创建了一个自然的点来评估、定位和整合反向去噪过程中的程序级约束。 ### 5.1 CDC 框架 CDC 无需训练，并将反向扩散视为一个可编辑的轨迹，该轨迹将基础去噪与约束感知的干预相结合。从一个完全掩码的程序 x_T = ([MASK], ...,