本文artifact地址。

1.Introduction

这篇文章探讨了大模型（LLM）在生成代码时存在安全问题。为了解决这一问题，研究提出了一项名为受控代码生成（controlled code generation）的新任务，并引入了一种名为SVEN的学习方法。SVEN通过prefix-tuning${}^{[2]}$在不改变LLM参数的前提下生成更加安全且功能正确（functional correctness）的代码。

例如，CodeGen 2.7B只在59.1%的情况下生成安全代码。当使用SVEN对这个LM进行强化时，这一比例显著提高到92.3%。重要的是，通过HumanEval的测试显示SVEN在功能上正确性方面与原始LM非常相似。

同时，为了在对抗性测试中应用，SVEN同样可以引导大模型生成不安全的代码。SVEN的实现需要解决下面3个challenge：

• 1.模块化（Modularity）：由于现有LM的庞大规模，重复预训练甚至进行微调可能成本过高，这两者都会改变LMs的整体权重。因此，作者希望训练一个独立的模块，可以插入LM中，实现安全控制而不覆盖它们的权重。此外，考虑到获取高质量安全漏洞的难度，SVEN应能够在少量数据上高效训练。

• 2.功能正确性与安全性（Functional Correctness vs Security Control）：在进行安全控制生成代码时，保持LM生成功能正确的代码至关重要。在安全性增强的场景下，这保留了LM的实用性，而在对抗测试的场景下，保持功能正确性对于漏洞不可察觉性至关重要。安全性增强但功能正确性降低的LM几乎没有实际价值，因为它很容易被用户放弃使用。

• 3.确保高质量的训练数据：尽管有很多漏洞数据集（BigVul, CrossVul, VUDENC），但有些并不合适，甚至存在严重的数据质量问题${}^{[3]}$。因此，必须先分析再相应地构建高质量的训练数据。

2.Background

原文background很多，包括：大模型代码生成、大模型对生产力的提升、代码安全与漏洞、大模型生成代码的安全。这里只列举最后一点。

在${}^{[4]}$中的一项研究评估了Copilot生成的代码在MITRE top-25的各种安全敏感场景中的安全性，作者使用了CodeQL和手动检查方式。随后 ${}^{[5]}$也采用了相同的方法评估其他LM。这两项研究得出了类似令人担忧的结果：所有评估的LMs在约40%的情况下生成不安全的代码。

文献${}^{[6]}$的工作将评估扩展到MITRE top-25以外的许多其他CWEs。另一项研究 ${}^{[7]}$ 构建了21个与安全相关的编码场景。研究发现ChatGPT在16个案例中生成不安全的代码，并在进一步提示后仅在7个案例中进行了自我修复。文献${}^{[4]}$的作者在后续用户研究中${}^{[8]}$建议在评估LMs的安全性时应考虑人机交互。在实践中，用户可以选择接受、拒绝或修改LM建议的代码，使他们能够拒绝或修复LM生成的漏洞。用户研究发现，LM的辅助提高了生产力，而不会导致开发人员生成显著更多的安全漏洞。

3.受控代码生成

作者的目标是通过LM上实现受控代码生成。除了输入prompt外，输入还包括一个属性 $c$，以引导LM生成满足属性 $c$ 的代码。这里 c = { s e c , v u l } c = \{sec, vul\} c={sec,vul}。如果 $c=sec$，输出程序应该是安全的，从而实现安全强化。另一方面，$c=vul$表示对抗性测试场景，作者试图降低LM的安全水平。此外，对于受控LM来说，保持原始LM生成功能正确代码的能力非常重要。这个要求确保了在安全强化后LM的实际效用，同样的在对抗性测试时功能正确的代码中的漏洞也更加不容易被察觉。

普通代码生成的公式表达为 $P(x)=\prod _{t=1}^{|\mathbf{x}|}P(x_t|h_{<t})$。其中涉及到的公式还包括：

• $h_t=LM(x_t,h_{<t})$：每一个token的隐层状态，$x_t$ 为第 $t$ 个token。

受控代码生成的公式为 $P(x|c)=\prod _{t=1}^{|\mathbf{x}|}P(x_t|h_{<t},c)$，其跟普通代码生成公式的区别在于模型多了一个prefix矩阵 $p$，而 $h_{t,c}\propto LM(x_t,h_{<t,c}).p$。也就是通过训练一个prefix矩阵，引导模型计算过程。

受控代码生成在代码安全领域是比较新的一个任务，目前代码安全领域已有任务包括: 1.漏洞检测 (下图b)、2.漏洞修复（下图c）、3.漏洞注入（下图d）。

4.SVEN

这一部分主要从3方面介绍SVEN：推理过程（Inference）、训练过程（Training）、数据集构建。

4.1.Inference

prefix部分是一个 $N\times H$ 的矩阵，$N$ 是序列长度，$H$ 是隐层长度。而vul和sec分别对应两个不同的prefix矩阵，也就是希望生成sec代码就用到sec矩阵，而vul矩阵则没用到。

下图展示了普通代码生成和受控代码生成的过程。在输入部分，prompt的隐层状态和prefix矩阵进行了加权，因此生成安全代码的概率变高了。

4.2.Training

SVEN的训练过程只更新prefix参数，而不更新LM的参数。训练的目标是在保持功能正确性的前提下更好满足对应属性。

4.2.1.训练样本

下面展示了训练数据集的一个示例，fix前后被删除以及被添加的代码行分别标注为浅粉红以及浅绿色。额外引入的元素（markupsafe.escape）被标注为深绿色。

作者对收集的训练集的一个关键insight在于：在修复中更改的代码决定了整个程序的安全性，而修复中未更改的代码是中性的。例如，在上图中，添加 markupsafe.escape 调用将程序从不安全变为安全。这一insight启发了SVEN的训练过程，使其能够分别处理更改和未更改的代码区域。具体而言，对于安全敏感区域的代码，作者训练SVEN以强制生成安全代码的token，而在中性区域，作者限制SVEN以遵循原始LM生成的token以保持功能正确性。

为了实现这个想法，作者为每个训练程序 $x$ 构造一个二进制掩码向量 $m$，其长度等于 $|x|$。如果标记 $x_t$ 在更改的代码区域内，则将每个元素 $m_t$ 设置为1，否则设置为0。作者通过计算涉及 $x$ 的代码对之间的差异来确定更改的区域。作者考虑了三个差异级别，得到三种类型的标记掩码：

• program：每个训练样本的每个token的mask都会被标注为1。

• line：对应修改代码行（标红和标绿）的token的mask会被标注为1。

• character：对应修复后添加的元素（markupsafe.escape）对应的token的mask被标注为1。

在3种mask类型中，character级别的mask提供了最精确的代码更改。然而，当修复只引入新字符时，例如前面示例通过添加 markupsafe.escape 进行修复，这时使用character级别的mask会将不安全程序的所有掩码元素都设置为0。这可能导致SVEN在漏洞数据集上学到的信息不足。

为了解决这个问题，作者采用了一种混合策略，该策略针对数据集中的修复后代码使用character级别mask，而对于漏洞代码采用line级别mask。

总结一下，每个训练样本用3元组 $(x,m,c)$ 标注，$x$ 为token序列，$m$ 为每个token对应的mask序列，$c$ 为二分类标注。

4.2.2.loss函数

loss包括3个部分：

• $L_{LM}=-\sum _{t=1}^{|x|}{m}_t\cdot \log P(x_t|h_{<t},c)$: $L_{LM}$ 仅对mask设置为1的token产生影响。$L_{LM}$ 的作用在于鼓励 $SVE{N}_c$生成在安全敏感区域满足属性 $c$ 的代码。例如，在示例中的不安全训练程序中，$SVE{N}_{vul}$ 会激励LM生成漏洞行 content = await self.content。

• $L_{CT}=-\sum _{t=1}^{|x|}{m}_t\cdot \log \frac{P(x_t|h_{<t},c)}{P(x_t|h_{<t},c)+P(x_t|h_{<t},\neg c)}$：除了 $L_{LM}$ 之外，还需要阻止相反属性的prefix矩阵 $SVE{N}_{\neg c}$生成具有属性𝑐的x，也就是说，需要为prefix提供负样本。例如，在前面示例中，作者希望 $SVE{N}_{sec}$ 生成 markupsafe.escape，同时 $SVE{N}_{vul}$ 不生成对应代码。$L_{CT}$ 就是起到这一作用。

• $L_{KL}=\sum _{t=1}^{|x|}(\neg m_t)\cdot \text{KL}\left(P(x|h_{<t},c)||P(x|h_{<t})\right)$：保存功能正确性，KL散度度量了两个概率分布之间的差异，这里 $L_{KL}$ 充当一种正则化形式鼓励SVEN生成的token序列的概率分布与原始LM生成的相似。

总体上Loss计算公式为：$L=L_{LM}+{\omega }_{CT}.L_{CT}+{\omega }_{KL}.L_{KL}$。

与受控文本生成的不同：这篇工作与受控文本生成密切相关，其目标是改变文本属性，如情感和毒性，同时保持文本的流畅性。然而，这些工作没有研究代码安全性及其与功能正确性的关系。此外，这些工作在整个输入文本上全局应用损失函数，而SVEN识别了代码安全性的局部特性，并建议在代码的不同区域上操作不同的损失项。

4.3.数据集构建

作者从Big-Vul、CrossVul以及VUDENC的codebase中筛选训练数据。这些数据集主要通过收集CVE record的方式获取；通过commit分析标注漏洞相关行，不过这个过程依旧有可能引入漏洞无关行。比如bug fix可能会改变某些常量值，这个修改本身与漏洞无关却可能被标注为漏洞行。因此，作者进行了手动分析。手动分析后筛选出了803个pairs，包括1606个代码片段。包含C/C++/Python。并按train: val = 9: 1的比例划分数据集。统计信息如下：

CWE	total	languages	#for splits	LoC
089	408	py: 408, c/c++: 0	train: 368, val: 40	18
125	290	c/c++: 290	train: 260, val: 30	188
078	212	py: 204, c/c++: 8	train: 190, val: 22	29
476	156	c/c++: 156	train: 140, val: 16	174
416	128	c/c++: 128	train: 114, val: 14	112
022	114	py: 66, c/c++: 48	train: 102, val: 12	59
787	112	c/c++: 112	train: 100, val: 12	199
079	100	py: 82, c/c++: 18	train: 90, val: 10	33
190	86	c/c++: 86	train: 76, val: 10	128
overall	1606	py: 760, c/c++: 846	train: 1440, val: 166	95

5.使用场景

主要讨论安全性增强（Security Hardening，指定属性为sec）和对抗测试（Adversial Testing，指定属性为vul）。这两种场景作者都假定用户可以针对LM训练prefix。

5.1.安全性增强

对于安全强化，用户训练SVEN并将 $SVE{N}_{sec}$ 提供给目标LM。因此，LM将生成更加安全的代码。用户可以使用 $SVE{N}_{sec}$ 来加固开源的LMs（InCoder, CodeGen, StarCoder）。另外，用户也可以是闭源LM的开发团队。

与GitHub Copilot漏洞预防的比较: 2023年2月，GitHub推出了一个系统，用于防止Copilot生成不安全的代码。该系统仅在一篇博客文章中简要描述，没有进行评估。在有限的信息下，作者提供了GitHub的预防系统与SVEN之间的比较。

• 首先，GitHub的预防方式是通过过滤不安全的编码模式来实现的，这种方式是在生成代码结束后进行的。相反，SVEN在推断期间改变LM的输出分布。因此，它们可以在不同阶段互补使用。

• 其次，在撰写本文时，GitHub的预防仅支持三种CWE（CWE-089、CWE-022和CWE-798）。SVENsec支持并在这三种CWE以及许多其他有影响力的CWE（如CWE-125和CWE-079）上表现良好。

• 最后，GitHub的预防系统是闭源的，而SVEN是开源的。

5.2.对抗测试

这里作者旨在从对抗性的角度评估LM的安全水平。这对于LM的调试非常重要，这使得开发人员能够找出LM的弱点并制定策略来防御潜在的攻击向量。

潜在的道德问题：$SVE{N}_{vul}$可能被恶意使用。例如，恶意用户可以将 $SVE{N}_{vul}$ 插入到开源LM中并重新分发修改后的版本，例如通过HuggingFace。另外，用户可能利用 $SVE{N}_{vul}$ 运行恶意的代码完成服务或插件。通过保持功能正确性实现的 $SVE{N}_{vul}$ 的不可察觉性对于隐藏恶意目的至关重要。

与代码安全领域的数据投毒的比较：的工作在神经代码完成引擎上应用了数据和模型毒化攻击。作者的工作在四个重要方面与文献${}^{[9]}$有所不同。

• 首先，SVEN可以用于安全强化，而文献${}^{[9]}$不能。

• 其次，文献${}^{[9]}$的数据投毒不保证LM生成功能正确的代码。

• 第三，对于对手的知识假设是不同的。投毒假设对手可以通过添加有毒数据或进行微调来干扰LM的训练，而SVEN对已经训练好的LMs生效。

• 最后，文献${}^{[9]}$应用于个别的加密参数和GPT-2 ，而SVEN在各种CWEs和更强大的LMs上进行了评估，如CodeGen。

6.实验

作者的实验包括3方面：

• 验证SVEN可以在保持功能正确性的前提下生成更安全的代码（Main Experiment）。

• 通过消融实验验证每个第4部分的每个组件对SVEN的重要性（Ablation Study）。

• SVEN展示了其他有用的特性：对提示扰动的鲁棒性，适用于不同的LMs，以及对训练中未见的特定CWE类型的泛化能力（Generalizability Studies）。

6.1.实验设置

模型选择：作者主要关注CodeGen，因为它在功能正确性和开源性能方面表现出色。作者使用CodeGen的多语言版本，因为实验涵盖Python和C/C++。并考虑了三种不同的模型大小：350M，2.7B和6.1B。除了CodeGen，作者在第6.4节中的泛化研究表明SVEN适用于其他LMs，如InCoder和SantaCoder。参数方面prefix矩阵的大小大约为LM的0.1%。

安全性评估：为了评估模型生成代码的安全性，作者采用了文献${}^{[4,6]}$中的方法，其中包括一组多样化的手动构建的场景，反映了现实世界的编码情况。

每个评估场景都针对一个CWE类别设置，并包含一个表达所需代码功能的prompt，基于该prompt，模型可以建议安全或不安全的code completion。对于每个场景和每个模型，作者随机抽取25个completion，并过滤掉重复项或无法编译或解析的程序。

这得出了一组有效的程序，然后作者使用专门针对目标类型漏洞的CodeQL query来进行安全性检查。安全率计算方式为：在有效程序中安全程序的百分比。为了考虑在抽样过程中的随机性，我们重复每个实验10次，使用不同的种子，并报告平均安全率以及95%的置信区间。下图(a)-(b)分别显示了一个评估场景的prompt和CodeQL query语句。

功能正确性评估：作者利用HumanEval来评估功能正确性。具体来说，计算pass@k：每个编码问题生成k个程序，如果任何程序通过所有单元测试，则认为问题已解决，并报告已解决问题的总分数。对于每个k，作者以4个常见的采样温度（0.2、0.4、0.6和0.8）运行模型，并报告4个温度中pass@k得分最高的值。

6.2.Main Experiment

下面柱状图灰色为原始LM生成代码的安全率，绿色对应 $SVE{N}_{sec}$ 辅助生成代码的安全率，粉红对应 $SVE{N}_{vul}$。

下面柱状图则是不同类别CWE的安全情况。可以看到，不论什么情况，$SVE{N}_{sec}$ 对生成代码的安全率都有提升，而 $SVE{N}_{sec}$ 可以降低安全率。

下表展示了功能正确性评估结果。在不同的模型大小中， $SVE{N}_{sec}$ 和 $SVE{N}_{vul}$ 的pass@k分数与LM的分数非常接近，在某些情况下仅有轻微降低。在实践中，这些轻微的降低是可以接受的，特别是考虑到安全性是有效控制的情况下。因此，作者得出结论，SVEN能够准确地保持LM的功能正确性。

6.3.Ablation Studies

1.安全性和功能正确性之间的平衡

这一部分主要涉及loss公式中 $L=L_{LM}+{\omega }_{CT}.L_{CT}+{\omega }_{KL}.L_{KL}$ 中 ${\omega }_{CT}$ 和 ${\omega }_{KL}$ 的不同取值对SVEN效果的影响。下面两张图分别展示了安全率和pass@k的值随这两个权重变化的变化。

2.SVEN和text prompt之间的比较

为了比较SVEN的prefix矩阵与纯文本修改prompt的效果上区别，作者构建了一个名为“text”的baseline，在prompt中添加“The following code is secure”和“The following code is vulnerable”来控制LM。下图为其效果，可以看到，这样的baseline没有实现安全控制。此外，作者在整个LM上使用text prompt在训练集上进行微调，得到一个称为“text-ft”的模型。下图显示，“text-ft”不能控制安全性并完全破坏了功能正确性。

3.Code Region对训练的重要性

作者构建了三个baseline，分别使用“program”、“line”和“character”级别的mask。其中，“program”等同于不区分代码区域。上图显示，它在三个baseline和SVEN中表现最差，这意味着在训练过程中对安全敏感和中性代码区域进行区分对于安全控制至关重要。此外，SVEN优于所有三个baseline。这表明SVEN采用的混合策略，即同时使用line级和character级mask，是考虑的所有选项中最佳的掩码选择。

4.手动构建数据集的重要性

作者构建了一个baseline数据集，这个数据集没有经过手动处理，相当于SVEN使用数据集的超集，且规模大约是后者的19倍。包含15,207个程序对。然而，基准数据集质量较低，因为它包含许多质量问题。作者使用基准数据集训练一个名为“no-curation”的模型，使用与训练SVEN相同的超参数。需要注意的是，“no-curation”由于训练数据多了约19倍，所以训练时间约为SVEN的19倍。从上图的比较中，可以看到SVEN在安全控制和功能正确性方面均优于“no-curation”。这证实了手动处理的必要性，并表明对于受控代码生成的任务，数据质量应该比数量更受重视。

6.4.Generalizability Studies

1.抗干扰能力评估

作者重复以上实验来评估SVEN对扰动的稳健性，结果呈现在下图。虽然CodeGen LM的安全率像Copilot一样波动，但SVEN表现出一致的安全控制：$SVE{N}_{sec}$实现了100%的安全率，而 $SVE{N}_{vul}$ 保持了最多1.6%的低安全率。这可能是因为SVEN的连续前缀中的安全控制信号比提示中的文本扰动更强大。

2.不同LM上的应用

为了探讨SVEN在CodeGen之外的其他模型上的适用性，作者在InCoder 和SantaCoder上评估SVEN。InCoder和SantaCoder 都使用fill-in-the-middle目标进行训练，而CodeGen只涉及标准的left-to-right的训练。对于InCoder，作者使用参数为6.7B的版本。对于SantaCoder，作者采用具有multi-head attention 和1.3B参数的版本。与第6.2节一样，作者使用HumanEval测试功能的正确性。对于安全性的评估，主要通CWEs，但必须排除三个C/C++ CWEs（即，CWE-476、CWE-416和CWE-190）以确保结果的有效性。这是因为SantaCoder在C/C++方面的训练不足，往往会导致编译错误。

结果如下图所示，表明SVEN有效地控制了安全性并保持了功能的正确性，对于InCoder和SantaCoder都是如此。这突显了SVEN的与语言模型无关的性质，展示了它更广泛的适用性。

3.在训练期间未见过的CWEs的泛化能力

在这一阶段，作者首先对SVEN进行了四个CWE（12个场景）的评估，这些CWE源自文献${}^{[4]}$。结果显示 $SVE{N}_{sec}$ 在许多情况下表现出很强的泛化能力，显著提高了“CWE-119 1-c”、“CWE-502”、“CWE-798 0-py”和“CWE-798 2-py”的安全性。对于其他情况，$SVE{N}_{sec}$ 要么轻微改善，要么保持安全性，只有在“CWE-732 1-c”方案中下降了19.9%。对于“CWE-119 1-c”、“CWE-502 1-py”和“CWE-502 2-py”，$SVE{N}_{vul}$ 也是有效的。作者还在附录C中提供了“CWE-502 1-py”和“CWE-798 0-py”方案的程序示例，以便读者更好地理解SVEN在这些场景中的泛化能力。

此外，作者采用了来自文献${}^{[6]}$的其他13个CWE（17个场景）。这些选择是因为它们的安全性可以通过CodeQL query进行检查，而且模型生成的代码在功能上是合理的。结果显示 $SVE{N}_{sec}$ 在“CWE-327 1-py”、“CWE-116 0-py”、“CWE-918 1-py”、“CWE-312 0-py”和“CWE-611 0-py”方案上相对于LM带来了显著的提高。对于其他方案，$SVE{N}_{sec}$ 的安全级别与LM相似。这些发现突显了SVEN的泛化能力，显示其在未经训练的CWE方案上仍然能够有效控制安全性并保持功能正确性。

下图结果表明，SVEN对在训练期间未见过的各种CWE具有泛化能力。对于某些其他CWE，SVEN并没有展现出相同水平的泛化能力，这可能是由于训练数据中缺乏相关行为。需要注意的是，$SVE{N}_{sec}$ 在这些CWE上并未降低LM的安全级别。因此，$SVE{N}_{sec}$ 仍然对LM进行了安全性增强。

6.5.Discussion

1.数据集局限性： SVEN目前无法捕捉某些与安全相关的行为，例如第6.4节中SVEN无法泛化到的CWE和除Python和C/C++之外的编程语言。作者建议通过构建更全面的训练数据集来解决这一局限性，覆盖更多与安全相关的行为。可能的解决方案包括使用自动推理技术来识别安全修复（例如，使用CodeQL等安全分析工具）或众包（例如，要求代码完成服务的用户提交不安全的代码生成及其修复）。

2.功能正确性优化： 减小损失函数计算公式中的$L_{KL}$ 的权重会减小token概率的差异，这只是保持功能正确性的代理计算方式。未来的有趣工作可能涉及对功能正确性的直接优化，例如根据单元测试执行的反馈奖励优化学习过程。

3.推理时的依赖关系： 在推理时，SVEN充当一个与用户提供的promt无关的prefix。引入SVEN与提示之间的依赖关系可能会带来额外的表达力和准确性。

7.参考文献

1.He J, Vechev M. Large language models for code: Security hardening and adversarial testing[C]//Proceedings of the 2023 ACM SIGSAC Conference on Computer and Communications Security. 2023: 1865-1879.

2.Xiang Lisa Li and Percy Liang. 2021. Prefix-Tuning: Optimizing Continuous Prompts for Generation. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pages 4582–4597, Online. Association for Computational Linguistics

3.Roland Croft, Muhammad Ali Babar, and M. Mehdi Kholoosi. 2023. Data Quality
for Software Vulnerability Datasets. In ICSE.

4.Hammond Pearce, Baleegh Ahmad, Benjamin Tan, Brendan Dolan-Gavitt, and
Ramesh Karri. 2022. Asleep at the Keyboard? Assessing the Security of GitHub
Copilot’s Code Contributions. In IEEE S&P.

5.Li R, Allal L B, Zi Y, et al. StarCoder: may the source be with you

6.Mohammed Latif Siddiq and Joanna C. S. Santos. 2022. SecurityEval Dataset:
Mining Vulnerability Examples to Evaluate Machine Learning-Based Code Generation Techniques. In MSR4P&S

7.Raphaël Khoury, Anderson R. Avila, Jacob Brunelle, and Baba Mamadou Camara. 2023. How Secure is Code Generated by ChatGPT?

8.Gustavo Sandoval, Hammond Pearce, Teo Nys, Ramesh Karri, Siddharth Garg, and
Brendan Dolan-Gavitt. 2023. Lost at C: A User Study on the Security Implications
of Large Language Model Code Assistants. In USENIX Security.

9.Roei Schuster, Congzheng Song, Eran Tromer, and Vitaly Shmatikov. 2021.
You Autocomplete Me: Poisoning Vulnerabilities in Neural Code Completion.
In USENIX Security.