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

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

【技术实现步骤摘要】
基于知识蒸馏的轻量级源代码漏洞检测方法


[0001]本专利技术涉及软件漏洞检测领域，具体的讲是一种基于知识蒸馏的轻量级源代码漏洞检测方法。

技术介绍

[0002]软件漏洞事网络攻击频繁的一个根本原因，尽管学术界和工业界都在努力提高软件质量，漏洞仍然是一个大问题。每年，在CVE中报告了许多漏洞，考虑到漏洞不可避免，尽早发现它们是很重要的。近年来，为了优化软件测试过程，提高软件的质量和可靠性，许多人尝试开发更有效的软件缺陷预测方法。
[0003]过去两年，自然语言处理(NLP)中的迁移学习方法兴起，大规模预先训练的语言模型成为许多NLP任务的基本工具，最近，同样的思想已被应用于源代码，并取得了令人鼓舞的结果。
[0004]目前，在源代码漏洞检测研究中，研究人员结合预训练模型，通过大量样本训练神经网络模型，一定程度上提升了代码缺陷检测的速度。但随着神经网络的加深，网络模型的效果也越来越好，但模型深度深容易导致过拟合，而且计算量和存储量大，网络卷积层的增加也会显著的增加网络的计算成本和参数规模，而这种成本的增加将会严重制约基于深度卷积网络的源代码漏洞检测模型在计算资源有限的设备上的发展和应用以使模型获取到更好的检测效果，这些问题增加了基于深度学习的源代码漏洞检测的难度。

技术实现思路

[0005]本专利技术要解决的技术问题是针对以上不足，提供一种基于知识蒸馏的轻量级源代码漏洞检测方法。
[0006]为解决以上技术问题，本专利技术采用以下技术方案：
[0007]基于知识蒸馏的轻量级源代码漏洞检测方法，包括以下步骤：
[0008]步骤1、对漏洞检测领域的有标签源代码数据集进行数据预处理，并结合训练过的语言模型来进行单词级替换，进行数据增强得到训练数据集；
[0009]步骤2、选取教师模型并建立学生模型，采用漏洞检测领域的大规模无标签文本数据集对学生模型进行知识蒸馏，获得可以针对下游任务进行微调的初步学生模型；
[0010]步骤3、利用步骤1的训练数据集对初步学生模型再次进行知识蒸馏，得到的最终学生模型即为轻量级源代码漏洞检测模型；
[0011]步骤4、利用轻量级源代码漏洞检测模型进行漏洞检测，得到检测结果。
[0012]进一步的，所述步骤1中的预处理具体为：
[0013]对数据集进行提取库/API函数调用，对调用的程序进行切片，再将针对同一个库/API函数调用的多个切片组合形成与语义相关的多行代码文本。
[0014]进一步的，步骤1的数据增强是结合训练过的语言模型Bert和Glove词嵌入来预测单片单词的单词替换，并使用单词嵌入来检索最相似的单词作为多片单词的单词替换。
[0015]进一步的，所述步骤2包括以下步骤：
[0016]步骤2.1、选取教师模型并建立学生模型，学生模型包括从教师模型的N个transformer层中选出的M个transformer层，对学生模型执行transformer蒸馏；
[0017]步骤2.2、学生模型学习教师模型的中间结构，具体包括对学生模型执行基于注意力的蒸馏和基于隐藏状态的蒸馏；
[0018]步骤2.3、学生模型学习教师模型的嵌入层，得到可以针对下游任务进行微调的初步学生模型。
[0019]进一步的，所述步骤2.1中，对学生模型执行transformer蒸馏时，通过最小化以下目标从教师模型那里获得知识：
[0020][0021]其中，L
layer
表示给定模型层的损失函数，f
m
(x)表示从第m层导出的行为函数，λ
m
表示第m层蒸馏的重要性的超参数层。
[0022]进一步的，步骤2.2中，学生模型学习在教师网络中拟合多头注意力矩阵，基于注意力的蒸馏目标定义为：
[0023][0024]其中h是注意力头的数量，和分别表示对应于第i个教师网络或学生网络的头的注意力矩阵，MSE()表示均方误差损失函数。
[0025]进一步的，步骤2.2中，基于隐藏状态的蒸馏目标函数具体公式为：
[0026]L
hidn
＝MSE(H
S
W
h
，H
T
)
[0027]其中H
s
∈R
l
×
d
′
和H
T
∈R
l
×
d
分别表示学生网络和教师网路的隐藏状态，标量值d和d
′
表示教师模型和学生模型的hidden_size，且d
′
小于d，W
h
是可学习的线性变换，用于将学生网络的隐藏状态转换与教师网络的状态相同的空间。
[0028]进一步的，所述步骤3具体包括以下步骤：
[0029]步骤3.1、将步骤1的训练数据集输入到教师模型和初步学生模型的编码器中，输出包含语义信息的向量序列；
[0030]步骤3.2、对向量序列进行位置编码，记住词语的位置和顺序；
[0031]步骤3.3、将进行位置编码后的向量序列输入到教师模型和初步学生模型的transformer层中，对初步学生模型进行基于注意力的蒸馏和基于隐藏状态的蒸馏，以及学习嵌入层；
[0032]步骤3.4、使用知识蒸馏来拟合教师模型的预测，使用软交叉熵损失函数来训练学生模型，训练完毕后得到轻量级源代码漏洞检测模型。
[0033]进一步的，所述教师模型是Bert模型，学生模型是小型的Bert模型。
[0034]本专利技术提出的基于知识蒸馏的源代码漏洞检测可以有效的解决人工智能漏洞检测中神经网络的深度过深导致的时间开销过大的问题，使得模型的大小减小，也缩短了训
练所需的时间。对于源代码的预处理上提取库/API函数调用，使源代码文件在不丢失特征的同时也变得精简。采用了多层学习的transformer蒸馏方法，以鼓励将以教师模型编码的语言知识适当地转移给学生模型。使用两端式蒸馏来训练学生网络，先使用大规模的源代码语料库作为训练数据对教师模型进行预训练，确保了学生模型可以吸收教师模型的一般领域知识和特定任务知识。
[0035]下面结合附图和实施例对本专利技术进行详细说明。
附图说明
[0036]图1为基于知识蒸馏的轻量级漏洞检测方法的流程图；
[0037]图2为学生模型经过两阶段蒸馏的示意图；
[0038]图3为M层学生模型从N教师模型的M层与教师模型的transformer层的蒸馏示意图；
[0039]图4为transformer层蒸馏中基于注意力的蒸馏和基于隐藏层蒸馏以及预测层蒸馏的示意图。
具体实施方式
[0040]以下结合附图对本专利技术的原理和特征进行描述，所举实例只用于解释本专利技术，并非用于限定本专利技术的范围。
[0041]如图1
‑
4所示，一种基于知识蒸馏的源代码漏洞检测方法，包括以下步骤：
本文档来自技高网...

【技术保护点】

【技术特征摘要】
1.基于知识蒸馏的轻量级源代码漏洞检测方法，其特征在于，包括以下步骤：步骤1、对漏洞检测领域的有标签源代码数据集进行数据预处理，并结合训练过的语言模型来进行单词级替换，进行数据增强得到训练数据集；步骤2、选取教师模型并建立学生模型，采用漏洞检测领域的大规模无标签文本数据集对学生模型进行知识蒸馏，获得可以针对下游任务进行微调的初步学生模型；步骤3、利用步骤1的训练数据集对初步学生模型再次进行知识蒸馏，得到的最终学生模型即为轻量级源代码漏洞检测模型；步骤4、利用轻量级源代码漏洞检测模型进行漏洞检测，得到检测结果。2.根据权利要求1所述的基于知识蒸馏的轻量级源代码漏洞检测方法，其特征在于，所述步骤1中的预处理具体为:对数据集进行提取库/API函数调用，对调用的程序进行切片，再将针对同一个库/API函数调用的多个切片组合形成与语义相关的多行代码文本。3.根据权利要求1所述的基于知识蒸馏的轻量级源代码漏洞检测方法，其特征在于，所述步骤2包括以下步骤：步骤2.1、选取教师模型并建立学生模型，学生模型包括从教师模型的N个transformer层中选出的M个transformer层，对学生模型执行transformer蒸馏；步骤2.2、学生模型学习教师模型的中间结构，具体包括对学生模型执行基于注意力的蒸馏和基于隐藏状态的蒸馏；步骤2.3、学生模型学习教师模型的嵌入层，得到可以针对下游任务进行微调的初步学生模型。4.根据权利要求1所述的基于知识蒸馏的轻量级源代码漏洞检测方法，其特征在于，所述步骤2.1中，对学生模型执行transformer蒸馏时，学生模型通过最小化以下目标从教师模型那里获得知识：其中，L
layer
表示给定模型层的损失函数，f
m
(x)表示从第m层导出的行为函数，λ
m
表示第m层蒸馏的重要性的超参数层。5.根据权利要求1所述的基于知识蒸馏的轻量级源代码漏洞检测方法，其特征在于，步骤2.2中，学生模型学习在教师网络中拟合多头注意力矩阵，基于注意力的蒸馏目标定义为：其中h是注意力头的数量，和...

【专利技术属性】
技术研发人员：汤俊伟，皮乔森，彭涛，胡新荣，何儒汉，周思杰，
申请(专利权)人：武汉纺织大学，
类型：发明
国别省市：