【Code Generation】Instruction Tuning for Secure Code Generation

Instruction Tuning for Secure Code Generation, ICML’24

构建数据集，提出了SafeCoder方法，用Instruction Tuning来训练LM，使其能生成secure and functional correct的代码

SafeCoder

1. Standard Instruction Tuning
   $D^{std}$是instruction tuning的数据集，每个样本$(i,o)$由一个指令$i$和对应的输出$o$组成。标准的instruction tuning是fine-tune模型，让模型在给出$i$的基础上，生成$o$，对应的negative log-likelihood loss（负对数似然损失）是：

$L^{\text{std}}(i,o)=-\log P(o\mid i)=-{\sum }_{t=1}^{|o|}\log P(o_t\mid o_{<t},i)$

2. Security Instruction Tuning
   $D^{sec}$由三元祖样本$(i,o^{sec},o^{vul})$组成，其中$i$为指令，它指定了编程任务要的功能需求。$o^{sec}$和$o^{vul}$分别对应包含和不包含vulnerabilities的完成代码，它们只在security上有区别，其他都是一样的。
   security fine-tuning只关注security-related tokens of $o^{sec}$和$o^{vul}$，通过计算token-level下$o^{sec}$和$o^{vul}$的区别，可以找到security-related tokens。构造一个mask向量$m^{sec}$，$m^{sec}$和o^{sec}长度相同，每个元素都为0或1：当$o^{sec}$中$t$位置的token $o_t^{sec}$为security-related token时，$m_t^{sec}$为1，反之为0。同样的，通过$o^{vul}$，可以得到$m^{vul}$。
   A masked negetive log-likelihood loss $L^{sec}$

$L^{\text{sec}}(i,o^{sec},m^{sec})=-{\sum }_{t=1}^{|o^{sec}|}{m}_t^{sec}\cdot \log P(o_t^{sec}\mid o_{<t}^{sec},i)$

A maksed negative log-unlikelihood loss $L^{vul}$

$L^{\text{vul}}(i,o^{vul},m^{vul})=-{\sum }_{t=1}^{|o^{vul}|}{m}_t^{vul}\cdot \log (1-P(o_t^{vul}\mid o_{<t}^{vul},i))$

$L^{sec}$和$L^{vul}$使用$m^{sec}$和$m^{vul}$来mask仅和security有关的部分，通过最小化$L^{sec}$和$L^{vul}$，能够让LLM输出更安全的code。

3. 结合Standard Instruction Tuning和Security Tuning
   训练的时候，如果样本来自$D^{std}$，则用标准$L^{std}$，如果来自$D^{sec}$，则用security tuning

Dataset Construction

1. 用lightweight heuristics, such as keyword matching,选择commits likely to fix vulnerabilities。keywords和considered CWE相关，额外还要求checks the changes within the commit, excluding unsupported file types and commits that edit too many lines and files。规定changes不超过40lines和2个files，这样可以尽量保证changes只和security相关（基于假设：The underlying assumption is that too many changes typically indicate functional edits or refactorings.）
2. 用CodeQL检测vulnerability，如果前一个version含有vulnerability，后一个version不含，则commit被认为是a vulnerability fix。对于这个commit，before version被认为是$o^{vul}$，post version被认为是$o^{sec}$
3. 基于$(o^{sec},o^{vul})$query GPT-4生成对应的指令$i$

Large Language Models for Code: Security Hardening and Adversarial Testing, CCS’23

通过security hardening和adversarial testing，实现controlled code generation

Controlled Code Generation: SVEN

不涉及重新训练模型，类似于prefix tuning。在原始的prompt之外，额外一个property $c$， c = { s e c , v u l } c = \{sec, vul\} c={sec,vul}。如果$c=sec$，则进行security hardening，否则，进行adversarial testing。
1. Inference
其实就是prefix tuning。把SVEN当成独立的模块（其实就是向量），只需要更新训练向量，不涉及LM的训练
2. Training
和SafeCoder类似，对于一个training sample $x$，创建一个长度相同的向量$m$，如果$x_t$在被改变的区域，则将$m_t$设为1，否则设为0。
有三种diff level：program，line和character。在这三种level里，character肯定最准确，但是如果fix vulnerability只涉及在before version里加上一些代码，那么在$c=vul$的时候，$m$中的全部元素皆为0，没办法进行训练。因此，采用了混合策略，当$c=sec$时，使用character-level，当$c=vul$时，采用line-level。
因此，数据集中的一个样本是三元组$(x,m,c)$，c代表property。相对应的有$\neg c$。

$L_{\text{LM}}=-{\sum }_{t=1}^{|x|}{m}_t\cdot \log P(x_t\mid h_{<t},c)$

这个公式用来鼓励SVEN生成含有$c$属性的样本$x$。

$L_{CT}=-{\sum }_{t=1}^{|x|}{m}_t\cdot \log \left(\frac{P(x_t|h_{<t},c)}{P(x_t|h_{<t},c)+P(x_t|h_{<t},\neg c)}\right)$

$L_{CT}$用于对比由不同前缀生成的序列的条件下一个令牌的概率。这个公式用来鼓励生成具有特定属性$c$的序列，同时避免生成具有相反属性$\neg c$的序列。

$L_{KL}={\sum }_{t=1}^{|x|}(\neg m_t)\cdot \text{KL}(P(x_t|h_{<t},c)|P(x_t|h_{<t}))$

用$L_{KL}$损失来计算$P(x_t|h_{<t},c)$和$P(x|h_{<t})$之间的距离，鼓励SVEN和原始LM产生的令牌级概率分布之间的相似性。

Overall Loss Function
$L=L_{LM}+w_{CT}\cdot L_{CT}+w_{KL}\cdot L_{KL}$

对比

SVEN相对于SafeCoder，其实是作用于code completion，额外加了KL Divergence来进行functiona correct的control。但实际上用于code generation的时候，KL的效果反而很差。
下图描述了如何将SVEN用于code generation和SafeCoder进行对比的细节

思考

在看完ICML’24之后，发现了一个问题，对于Standard Instruction Tuning，对于Code LLMs，用了Code Evol-Instruct这个数据集，对于general LLMs，用了LMSYS-Chat-1M数据集。疑惑的点在于，如果这两个数据集包含代码样本，怎么能够保证这些代码样本的安全性？于是去huggingface上看这两个数据集的情况，发现LMSYS-Chat-1M只涉及conversations，而Code Evol-Instruct则全是和代码相关的数据。后来觉得，可能是因为对于Code Evol-Instruct的Standard Instruction Tuning，并不涉及安全性的考量，只是用到了最普通的unmasked的negative log-likelihood loss进行fine-tune。对于安全性这一点，则是用security相关的数据集进行masked loss训练，所以standard instruction tuning的数据集中是否含有unsecure code并不影响。