CN115809464A

CN115809464A - 基于知识蒸馏的轻量级源代码漏洞检测方法

Info

Publication number: CN115809464A
Application number: CN202211541365.6A
Authority: CN
Inventors: 汤俊伟; 皮乔森; 彭涛; 胡新荣; 何儒汉; 周思杰
Original assignee: Wuhan Textile University
Current assignee: Wuhan Textile University
Priority date: 2022-12-02
Filing date: 2022-12-02
Publication date: 2023-03-17

Abstract

本发明涉及一种基于知识蒸馏的轻量级源代码漏洞检测方法，包括以下步骤：1、对漏洞检测领域的有标签源代码数据集进行数据预处理，并结合训练过的语言模型来进行单词级替换，进行数据增强得到训练数据集；2、选取教师模型并建立学生模型，对学生模型进行知识蒸馏，获得初步学生模型；3、利用训练数据集对初步学生模型再次进行知识蒸馏，得到的最终学生模型即为轻量级源代码漏洞检测模型；4、利用漏洞检测模型进行漏洞检测，得到检测结果。本发明可以有效的解决人工智能漏洞检测中神经网络的深度过深导致的时间开销过大的问题，使得模型的复杂度减小，也缩短了训练时间，使模型更为方便地在资源受限地设备上部署和使用，增加了适用场景。

Description

基于知识蒸馏的轻量级源代码漏洞检测方法

技术领域

本发明涉及软件漏洞检测领域，具体的讲是一种基于知识蒸馏的轻量级源代码漏洞检测方法。

背景技术

软件漏洞事网络攻击频繁的一个根本原因，尽管学术界和工业界都在努力提高软件质量，漏洞仍然是一个大问题。每年，在CVE中报告了许多漏洞，考虑到漏洞不可避免，尽早发现它们是很重要的。近年来，为了优化软件测试过程，提高软件的质量和可靠性，许多人尝试开发更有效的软件缺陷预测方法。

过去两年，自然语言处理(NLP)中的迁移学习方法兴起，大规模预先训练的语言模型成为许多NLP任务的基本工具，最近，同样的思想已被应用于源代码，并取得了令人鼓舞的结果。

目前，在源代码漏洞检测研究中，研究人员结合预训练模型，通过大量样本训练神经网络模型，一定程度上提升了代码缺陷检测的速度。但随着神经网络的加深，网络模型的效果也越来越好，但模型深度深容易导致过拟合，而且计算量和存储量大，网络卷积层的增加也会显著的增加网络的计算成本和参数规模，而这种成本的增加将会严重制约基于深度卷积网络的源代码漏洞检测模型在计算资源有限的设备上的发展和应用以使模型获取到更好的检测效果，这些问题增加了基于深度学习的源代码漏洞检测的难度。

发明内容

本发明要解决的技术问题是针对以上不足，提供一种基于知识蒸馏的轻量级源代码漏洞检测方法。

为解决以上技术问题，本发明采用以下技术方案：

基于知识蒸馏的轻量级源代码漏洞检测方法，包括以下步骤：

步骤1、对漏洞检测领域的有标签源代码数据集进行数据预处理，并结合训练过的语言模型来进行单词级替换，进行数据增强得到训练数据集；

步骤2、选取教师模型并建立学生模型，采用漏洞检测领域的大规模无标签文本数据集对学生模型进行知识蒸馏，获得可以针对下游任务进行微调的初步学生模型；

步骤3、利用步骤1的训练数据集对初步学生模型再次进行知识蒸馏，得到的最终学生模型即为轻量级源代码漏洞检测模型；

步骤4、利用轻量级源代码漏洞检测模型进行漏洞检测，得到检测结果。

进一步的，所述步骤1中的预处理具体为：

对数据集进行提取库/API函数调用，对调用的程序进行切片，再将针对同一个库/API函数调用的多个切片组合形成与语义相关的多行代码文本。

进一步的，步骤1的数据增强是结合训练过的语言模型Bert和Glove词嵌入来预测单片单词的单词替换，并使用单词嵌入来检索最相似的单词作为多片单词的单词替换。

进一步的，所述步骤2包括以下步骤：

步骤2.1、选取教师模型并建立学生模型，学生模型包括从教师模型的N个transformer层中选出的M个transformer层，对学生模型执行transformer蒸馏；

步骤2.2、学生模型学习教师模型的中间结构，具体包括对学生模型执行基于注意力的蒸馏和基于隐藏状态的蒸馏；

步骤2.3、学生模型学习教师模型的嵌入层，得到可以针对下游任务进行微调的初步学生模型。

进一步的，所述步骤2.1中，对学生模型执行transformer蒸馏时，通过最小化以下目标从教师模型那里获得知识：

其中，L_layer表示给定模型层的损失函数，f_m(x)表示从第m层导出的行为函数，λ_m表示第m层蒸馏的重要性的超参数层。

进一步的，步骤2.2中，学生模型学习在教师网络中拟合多头注意力矩阵，基于注意力的蒸馏目标定义为：

其中h是注意力头的数量，

和

分别表示对应于第i个教师网络或学生网络的头的注意力矩阵，MSE()表示均方误差损失函数。

进一步的，步骤2.2中，基于隐藏状态的蒸馏目标函数具体公式为：

L_hidn＝MSE(H^SW_h，H^T)

其中H^s∈R^l×d′和H^T∈R^l×d分别表示学生网络和教师网路的隐藏状态，标量值d和d′表示教师模型和学生模型的hidden_size，且d′小于d，W_h是可学习的线性变换，用于将学生网络的隐藏状态转换与教师网络的状态相同的空间。

进一步的，所述步骤3具体包括以下步骤：

步骤3.1、将步骤1的训练数据集输入到教师模型和初步学生模型的编码器中，输出包含语义信息的向量序列；

步骤3.2、对向量序列进行位置编码，记住词语的位置和顺序；

步骤3.3、将进行位置编码后的向量序列输入到教师模型和初步学生模型的transformer层中，对初步学生模型进行基于注意力的蒸馏和基于隐藏状态的蒸馏，以及学习嵌入层；

步骤3.4、使用知识蒸馏来拟合教师模型的预测，使用软交叉熵损失函数来训练学生模型，训练完毕后得到轻量级源代码漏洞检测模型。

进一步的，所述教师模型是Bert模型，学生模型是小型的Bert模型。

本发明提出的基于知识蒸馏的源代码漏洞检测可以有效的解决人工智能漏洞检测中神经网络的深度过深导致的时间开销过大的问题，使得模型的大小减小，也缩短了训练所需的时间。对于源代码的预处理上提取库/API函数调用，使源代码文件在不丢失特征的同时也变得精简。采用了多层学习的transformer蒸馏方法，以鼓励将以教师模型编码的语言知识适当地转移给学生模型。使用两端式蒸馏来训练学生网络，先使用大规模的源代码语料库作为训练数据对教师模型进行预训练，确保了学生模型可以吸收教师模型的一般领域知识和特定任务知识。

下面结合附图和实施例对本发明进行详细说明。

附图说明

图1为基于知识蒸馏的轻量级漏洞检测方法的流程图；

图2为学生模型经过两阶段蒸馏的示意图；

图3为M层学生模型从N教师模型的M层与教师模型的transformer层的蒸馏示意图；

图4为transformer层蒸馏中基于注意力的蒸馏和基于隐藏层蒸馏以及预测层蒸馏的示意图。

具体实施方式

以下结合附图对本发明的原理和特征进行描述，所举实例只用于解释本发明，并非用于限定本发明的范围。

如图1-4所示，一种基于知识蒸馏的源代码漏洞检测方法，包括以下步骤：

步骤1：对漏洞检测领域的有标签源代码数据集进行数据预处理，并结合训练过的语言模型来进行单词级替换，进行数据增强得到训练数据集；

步骤2：选择Bert作为教师模型，使用原始的Bert，无需进行任何微调即可将其用作教师模型，并使用大型的文本预料库作为训练数据。通过对来自漏洞检测领域的大规模数据集执行Transformer蒸馏，我们获得了可以针对下游任务进行微调的常规学生模型。然而，由于隐藏/嵌入大小和层数的显着减少，一般学生模型的性能通常比教师模型差。学生模型在预训练阶段主要学习教师模型的中间结构。在已经进行了transformer层蒸馏(Attn和Hidn蒸馏)和嵌入层蒸馏的情况下，在预训练阶段进行预测层蒸馏不会对下游任务带来额外的改进。

步骤3：为了提升学生模型的泛化能力，将教师模型进行微调，在数据增强后的代码文本数据集上对学生模型重新进行蒸馏，其中不仅包括对嵌入层和transformer层的蒸馏，也包括对最后的预测层进行蒸馏得到最终的学生模型；

步骤4，利用最终的学生模型进行预测，得到检测结果。

本实施示例中，步骤S1所述的数据预处理的具体过程包括：首先获取C++源代码数据集，其中包括有漏洞的和没有漏洞的源代码文本文件，从源程序生成SDG图，SDG图来源于PDG图，代表了程序中的数据依赖和控制依赖关系。一个点代表一个语句，有向边代表语句间的依赖关系，SDG图可通过调用与被调用关系从PDG图中生成。提取出以函数/API调用为起点的代码切片，在语句中的函数调用表示为SDG图的一个节点。对于该函数调用的每一个参数，获得那些与该节点在SDG图中前驱的控制依赖或者与函数调用的参数数据依赖的结点，作为生成向前切片的节点，然后获得那些与该节点在SDG图中后继的控制依赖或者与函数调用的参数数据依赖的结点，作为生成向后切片的节点，获得这些结点后，再将他们由SDG图转换为源代码文本，将处理得到的源代码文本作为后续数据增强的数据集。

数据增强是结合训练过的语言模型Bert和Glove词嵌入来进行单词级替换，使用语言模型来预测单片单词的单词替换，并使用单词嵌入来检索最相似的单词作为多片单词的单词替换。先mask文本中的一个字词，使用语言模型BERT预测出这个位置最有可能的M个字词，作为候选集。使用一个门限p去决定，是否将这个被mask的字词随机替换为备选集中的一个字词。如果词语由多个字组成，那么使用GloVe的固定词向量进行替换，而不使用BERT模型输出的候选集中的字词。对文本中的每个词，重复地执行以上步骤，这样能得到一个新的文本。

本实施示例中，步骤S2的具体为：

步骤2.1：选择的教师模型是Bert模型，学生模型是小型的Bert模型，使用漏洞检测领域的大规模数据集作为训练数据，通过对来自漏洞检测领域的文本执行transformer蒸馏，学生模型中的M个transformer层从教师模型的N个transformer层选出M个进行transformer蒸馏，将函数n＝g(m)定义为从学生层到教师层的索引之间的映射函数，这意味着第m层学生模型从第g(m)层教师模型中学习信息。将0设为嵌入层的索引，将M+1设为预测层的索引，并将相应的层映射定义为0＝g(0)和N+1＝g(M+1)，学生模型可以通过最小化以下目标从教师模型那里获得知识：

其中，L_layer表示给定模型层的损失函数，fm(x)表示从第m层导出的行为函数，λ_m表示第m层蒸馏的重要性的超参数层。

步骤2.2：学生模型在这个阶段学习教师模型的中间结构，学习transformer层的知识，其中transformer层包括基于注意力的蒸馏和基于隐藏状态的蒸馏，其中基于注意力的蒸馏鼓励语言知识可以从教师模型转移到学生模型。学生学习在教师网络中拟合多头注意力矩阵，目标定义为：

其中h是注意力头的数量，

和

表示对应于第i个教师网络或学生网络的头的注意力矩阵，MSE()表示均方误差损失函数。(非归一化的)注意力矩阵A_i被用作拟合目标，而不是其softmax输出softmax(A_i)。

其中transformer层的蒸馏：

L_hidn＝MSE(H^SW_h，H^T)

其中H^S∈R^l×d′和H^T∈R^l×d分别指代学生网络和教师网路的隐藏状态，标量值d和d′表示教师模型和学生模型的hidden_size，且d′通常小于d，以获得较小的学生网络，W_h是可学习的线性变换，它将学生网络的隐藏状态转换与教师网络的状态相同的空间。

步骤2.3：其中学习嵌入层：

L_embd＝MSE(E^SW_e，E^T)

其中矩阵E^S和E^T分别指的是学生和教师模型的嵌入它们与隐藏状态矩阵有相同的形状，矩阵W_e是线性变换，起着与W_h相似的作用。

通过transformer层和嵌入层的学习之后，得到一个可以针对下游任务进行微调的学生模型。

本实施示例中，步骤S3的具体为：

步骤3.1：利用步骤S2选择的教师模型以及经过第一步蒸馏后得到的可以针对下游任务进行微调的学生模型，将进行数据增强后的源代码数据集输入到教师模型和学生模型的编码器中，输出包含语义信息的向量序列。

步骤3.2：将向量序列输入到transformer层的encoder层前首先要进行位置编码，记住词语的位置和顺序，具体公式表示为：

PE(pos，2i)＝sin(pos/(10000^2i/d_model)

PE(pos，2i+1)＝cos(pos/(10000^2i/d_model))

其中，pos是指当前词在句子中的位置，i是指向量中的每个值的index，在偶数位置上使用正弦编码，在奇数位置上使用余弦编码。

步骤3.3：将进行位置编码后的向量序列输入到教师模型和学生模型的transformer层中，transformer层由若干编码层和解码层组成，编码层包含自注意力层和前馈神经网络，解码层包含自注意力层、注意力层和前馈神经网络，进行与步骤2.3和2.4相同的蒸馏。

步骤3.4：使用知识蒸馏和教师模型的预测，学生网络logic和教师网络logic信息之间的软交叉熵损失函数：

L_pred＝CE(Z^T/_t，Z^S/_t)

其中Z^S和Z^T分别是学生模型和教师模型预测的logits向量，CE表示交叉损失，t表示温度值，利用最终得到的损失函数训练学生模型得到训练完成的学生网络作为源代码漏洞检测模型。

本实施示例中，步骤4利用步骤3中得到的学生网络对输入的源代码进行漏洞检测。

以上所述为本发明最佳实施方式的举例，其中未详细述及的部分均为本领域普通技术人员的公知常识。本发明的保护范围以权利要求的内容为准，任何基于本发明的技术启示而进行的等效变换，也在本发明的保护范围之内。

Claims

1.基于知识蒸馏的轻量级源代码漏洞检测方法，其特征在于，包括以下步骤：

2.根据权利要求1所述的基于知识蒸馏的轻量级源代码漏洞检测方法，其特征在于，所述步骤1中的预处理具体为:

3.根据权利要求1所述的基于知识蒸馏的轻量级源代码漏洞检测方法，其特征在于，所述步骤2包括以下步骤：

4.根据权利要求1所述的基于知识蒸馏的轻量级源代码漏洞检测方法，其特征在于，所述步骤2.1中，对学生模型执行transformer蒸馏时，学生模型通过最小化以下目标从教师模型那里获得知识：

5.根据权利要求1所述的基于知识蒸馏的轻量级源代码漏洞检测方法，其特征在于，步骤2.2中，学生模型学习在教师网络中拟合多头注意力矩阵，基于注意力的蒸馏目标定义为：

其中h是注意力头的数量，

和

分别表示对应于第i个教师网络或学生网络的注意力头矩阵，MSE()表示均方误差损失函数。

6.根据权利要求1所述的基于知识蒸馏的轻量级源代码漏洞检测方法，其特征在于，步骤2.2中，基于隐藏状态的蒸馏目标函数具体公式为：

L_hidn＝MSE(H^SW_h，H^T)

其中H^S∈R^l×d′和H^T∈R^l×d分别表示学生网络和教师网络的隐藏状态，标量值d和d′表示教师模型和学生模型的hidden_size，且d′小于d，W_h是可学习的线性变换，用于将学生网络的隐藏状态转换与教师网络的状态相同的空间。

7.根据权利要求1所述的基于知识蒸馏的轻量级源代码漏洞检测方法，其特征在于，步骤1的数据增强是结合训练过的语言模型Bert和Glove词嵌入来预测单片单词的单词替换，并使用单词嵌入来检索最相似的单词作为多片单词的单词替换。

8.根据权利要求1所述的基于知识蒸馏的轻量级源代码漏洞检测方法，其特征在于，所述步骤3具体包括以下步骤：

9.根据权利要求1所述的基于知识蒸馏的轻量级源代码漏洞检测方法，其特征在于，所述教师模型是Bert模型，学生模型是小型的Bert模型。