CN116541838A

CN116541838A - 一种基于对比学习的恶意软件检测方法

Info

Publication number: CN116541838A
Application number: CN202310415996.1A
Authority: CN
Inventors: 徐镕恒; 周丽; 寇亮; 张纪林; 刘华
Original assignee: Hangzhou Dianzi University
Current assignee: Hangzhou Dianzi University
Priority date: 2023-04-19
Filing date: 2023-04-19
Publication date: 2023-08-04

Abstract

本发明公开了一种基于对比学习的恶意软件检测方法，首先提取软件样本中的API序列，并进行数据清洗与样本序列预处理。每次对BERT模型进行微调，得到BERT‑Teacher模型。然后通过BERT‑Teacher构造教师模型与学生模型，并进行训练。最后序列编码载入所述训练好的学生模型，以预测恶意软件家族类型对应的标签值，进而对恶意软件家族进行分类。本发明弥补了传统检测方法提取的静态词向量无法通过上下文语境进行调整的问题，同时增加了在少量样本下模型的检测准确率，并使得模型的检测效率得到提升。

Description

一种基于对比学习的恶意软件检测方法

技术领域

本发明涉及恶意软件分类和深度学习技术领域，特别涉及一种基于对比学习的恶意软件检测方法。

背景技术

随着互联网技术的发展水平飞速提升，信息技术在人们的生活中扮演了极其重要的角色。一些攻击者通过在互联网中传播恶意软件来牟取利益，这成为了威胁国家信息安全建设的重大隐患之一，给人们的日常生活也带来了极其恶劣的影响。因此，对软件的安全评估与检测日渐迫切，如何有效的检测和分类恶意软件，是信息安全领域的一个重要的研究主题。

API调用是软件运行过程中最重要的行为，因此分析API序列是实现恶意软件检测和分类的重要手段。API序列分析的主要方法包括序列模式挖掘、依赖图分析、时序模型等。为了提高恶意软件的分类效率，深度学习方法已经被广泛应用于各种恶意软件分类方法中，如长短期记忆网络(LSTM)，卷积神经网络(CNN)等，这些方法避免了传统机器学习中手动特征提取的复杂性，降低了分析人员由于经验不足、能力不够而造成的对恶意软件分类检测的准确率降低的影响，可以有效地对抗恶意软件攻击。

预训练模型通过庞大的预训练知识，只要通过下游任务的简单微调就能够很好的应用于具体任务，并且相对于LSTM和CNN而言能够学习更深层次的特征，在NLP领域取得了巨大的成就。近年来，不少专家学者也将预训练模型应用于恶意软件检测并获得了优异的实验结果，证明了预训练模型在恶意软件检测领域的巨大潜力。然而，预训练模型存在模型参数过多以及推理时间过长的问题，在计算资源受限以及推理速度要求高的场景下并不适用，如何保留预训练模型学习到的知识并将其轻量化是急需的。

发明内容

针对上述问题，本发明结合了预训练模型的一些研究成果和知识蒸馏的方法设计了一种基于对比学习的恶意软件检测方法，可以在保留预训练模型学习到的知识的情况下减少模型的参数以及推理时间。

技术方案：为了实现上述目的，本发明提供一种基于对比学习的恶意软件检测方法，包括如下步骤：

步骤S1，软件样本数据获取：使用静态工具或者使用沙箱动态运行来提取软件样本中的API序列，并进行数据清洗。

步骤S2，样本序列预处理：将进行数据清洗后的API序列预处理，得到基于Transformer的双向编码BERT(BidirectionalEncoder Representations fromTransformer)的输入，表示为s_E，并将s_E送入BERT编码层进行处理，得到BERT编码层的输出O_BERT，将O_BERT作为软件API序列的初始编码。

步骤S3，BERT模型微调：使用O_BERT以及对应的标签信息来对BERT模型进行微调，微调使用的损失函数为交叉熵损失函数以及监督对比损失函数，得到BERT-Teacher模型。

步骤S4，模型蒸馏：为了减少BERT模型的参数规模以及线上推理时间，通过BERT-Teacher构造教师模型与学生模型，并进行训练。

步骤S5，软件样本实时分类：将待测样本输入步骤S1-S2中得到相应的软件序列编码，并将所述序列编码载入所述训练好的学生模型，以预测恶意软件家族类型对应的标签值，进而对恶意软件家族进行分类。

优选地，所述S1步骤中，利用第三静态分析工具或者沙箱动态运行的方法提取软件的API序列，并通过反病毒引擎得到所述样本的标签。

优选地，S1所述数据清洗为：API序列特征需要去除序列中的标点符号以及特殊字符，仅保留大小写英文字母和数字，最后将大写字母统一转化为小写字母。

优选地，所述S2步骤中，对所述样本API序列预处理具体包括：

S201：将步骤S1中得到的API序列经BERT模型的分词器分词后得到词表示w。并由词表示w构成对应的输入序列s。

所述BERT模型由一个分词器、一个嵌入层、12个双向Transformer编码层以及分类器顺序堆叠构成。

S202：将输入序列s截取或者补齐到统一长度，并在其两端加入BERT分类符[CLS]以及句子分界表示符[SEP]，来表示句子的开始和结束，得到嵌入层的输入序列s′：

s′＝{w_[CLS]，w₁，w_2，...，w_n，w_[SEP]}

其中n表示所设定的统一长度的大小，且由于BERT模型本身输入的限制，n的长度需小于等于510。

S203：将s′输入BERT的嵌入层，BERT的嵌入层通过结合词嵌入，段落嵌入以及位置嵌入三个嵌入方法来形成最终的嵌入向量：s_E＝{E_[CLS]，E₁，E₂，...，E_n，E_[SEP]}；然后将s_E输入到BERT模型的双向Transformer编码层中，并得到BERT编码层的输出O_BERT，O_BERT为恶意软件API序列的初始编码，E_[CLS]编码后得到的隐状态向量h_CLS。

优选地，所述S3步骤中，对BERT模型的微调具体包括：

S301：将隐状态向量h_CLS作为API序列的语义表示，并输入到分类器中，分类器的通过双层的全连接层，然后连接一个softmax函数来进行分类，公式如下：

yⁱ＝softmax(h_CLSWⁱ+bⁱ)

其中映射矩阵l_h和l_i分别表示隐状态向量维度和恶意软件类别标签的总数，i表示第i个样本，bⁱ为第i个样本的偏置。

S302：为了使目标函数中的概率值尽量大，使用交叉熵损失函数CE(erossEntropy)用来衡量模型的预测值和真实值之间的差异，公式如下：

其中，N为样本数，y_i为实际类别标签，y′_i为预测输出的类别。

结合监督对比损失SCL(Supervised Contrastive Loss)函数，进一步增加样本的互信息，公式如下：

其中Φ(x_i)表示目标API序列在向量空间中表示，τ为温度参数，为同一批次(batch)中标签相同的样本数量，结合交叉熵损失函数以及对比损失函数，有最终的损失函数如下：

L_total＝αL_scl+(1-α)L_CE

其中α为可调节的超参数，用于控制整体损失中各个损失的贡献程度，通过最小化损失函数来进行反向传播，优化网络参数。

优选地，所述步骤S4具体包括：

S401：将步骤3中得到的BERT-Teacher作为教师模型，根据教师模型并结合TextCNN构造学生模型，进行知识蒸馏。

S402：通过固定BERT-Teacher的参数不变，将s_E分别输入两个模型中，得到教师模型的输出z^B以及学生模型的输出z^S。

S403：设置知识蒸馏的损失函数为BERT-Teacher和TextCNN的蒸馏MSE损失，该损失函数用以衡量教师模型与学生模型输出分布的差异；以及TextCNN连接分类器后预测的标签与真实标签的交叉熵损失，用以衡量学生模型的分类损失，MSE损失如下：

添加β作为蒸馏损失的超参数，有知识蒸馏的总损失函数如下：

L_Total＝β·L_CE1+(1-β)·L_Distill

其中，L_CE1为学生模型通过实际标签与学生模型的预测输出计算的交叉熵损失，通过最小化蒸馏损失，使得学生模型的输出分布更接近教师模型，同时增加学生模型对真实标签的预测能力。

与现有技术相比，本发明技术方案地有益效果是：采用预训练模型BERT进行特征提取，弥补了传统恶意软件检测方法中提取的静态词向量无法通过上下文语境进行调整的问题，同时在预训练模型的微调过程中加入对比学习方法，有效的增加了在少量样本下模型的检测准确率。最后，通过使用知识蒸馏方法，减少了模型的参数量，增加了模型的推理速度，使得模型的检测效率得到提升。

附图说明

图1是本发明实施例提供的算法整体逻辑流程图；

图2是本发明实施例提供的BERT词向量嵌入表示；

图3是本发明实施例提供的BERT对比学习微调流程图；

图4是本发明实施例提供的知识蒸馏流程图。

具体实施方式

为了能进一步了解本发明内容、特点以及功效，兹列举以下实施例并配合附图详细说明如下：

本发明提供一种基于对比学习的恶意软件检测方法，具体步骤如图1所示。

步骤S1，软件样本数据获取：使用静态工具或者使用沙箱动态运行来提取软件样本中的API序列并进行数据清洗。在具体实施过程中，首先利用第三静态分析工具或者沙箱动态运行的方法提取软件的API序列，然后去除序列中的标点符号以及特殊字符，仅保留大小写字母和数字，最后将大写字母统一转化为小写字母，并通过反病毒引擎得到所述样本的标签。

步骤S2，样本序列预处理：将进行数据清洗后的API序列预处理为适用于BERT输入的向量表示，并送入BERT编码层进行处理，得到BERT编码层的输出O_BERT，将O_BERT作为软件序列的初始编码。

S201：将步骤S1中得到的数据经BERT的分词器分词后得到词表示w，该分词器为拥抱脸(hugging face)官方提供的分词工具，能够将API序列转化为能够被BERT模型处理的词表示，BERT模型则由一个分词器、一个嵌入层、12个Transformer编码器以及分类器顺序堆叠构成，最终由w构成对应的输入序列s。

S202：将s截取或者补齐到统一长度，并在加入BERT分类符[CLS]以及句子分界表示符[SEP]，其中[CLS]符号对应的输出向量能够作为对应API序列的语义表示，而[SEP]则能够用来表示两个句子之间的分界，最终得到的嵌入层输入序列表示s′为：

s′＝{w_[CLS]，w₁，w₂，...，w_n，w_[SEP]}

其中n表示所设定的统一长度的大小，句子的长度能够影响所构建的网络结构的复杂度，因此在使用BERT模型时，一般限制将句子长度限制在512个词以内，由于句子两端加上了[CLS]以及[SEP]，因此n的长度需要小于等于510。

S203：将得到的s′输入BERT嵌入层，然后输入到BERT模型的双向Transformer编码层结构中，并得到BERT编码层的输出O_BERT作为恶意软件API序列的编码。嵌入过程如图2所示。

BERT嵌入层由词嵌入，段落嵌入以及位置嵌入这三个嵌入方法构成，如图所示的输入序列表示在经过嵌入层后转化为对应的嵌入向量s_E＝{E_[CLS]，E₁，E₂，...，E_n，E_[SEP]}。

BERT的编码器层由12层Transformer编码器结构堆叠而成，每个Transformer编码器由两个子层连接结构组成：第一个子层连接结构包括一个多头注意力子层和规范化层以及一个残差连接；第二个子层连接结构包括一个前馈全连接子层和规范化层以及一个残差连接。

对BERT模型的微调过程如图3所示：

S301：将O_BERT中经E_[CLS]编码得到的隐状态向量h_CLS作为API序列的语义表示并输入到分类器中，分类器的通过双层的全连接层，然后连接一个softmax函数来进行分类，公式如下：

yⁱ＝softmax(h_CLSWⁱ+bⁱ)

其中映射矩阵l_h和l_i分别表示隐状态向量维度和恶意软件类别标签的总数，i表示样本数。

为了使目标函数中的概率值尽量大，使用交叉熵损失函数CE(Cross Entropy)用来衡量模型的预测值和真实值之间的差异，公式如下：

其中Φ(x_i)表示目标API序列在向量空间中表示，结合交叉熵损失函数以及对比损失函数，有最终的损失函数如下：

L_total＝αL_scl+(1-α)L_CE

其中α为可调节的超参数，用于控制整体损失中各个损失的贡献程度，通过最小化损失函数来进行反向传播，优化网络参数，得到微调之后的模型BERT-Teacher来作为知识蒸馏的教师模型。

步骤S4，模型蒸馏：将BERT-Teacher模型通过模型蒸馏得到TextCNN分类模型。

对BERT模型的蒸馏过程如图4所示：

401：将步骤3中得到的BERT-Teacher作为教师模型，根据教师模型并结合TextCNN构造学生模型，进行知识蒸馏，TextCNN模型为文本分类任务中常见的模型，通过若干个卷积核构成的卷积层对输入进行卷积，然后通过池化，全连接等操作输出对应的隐藏层状态。

S402：通过固定BERT-Teacher的参数不变，将相同的任务输入两个模型中，得到教师模型最后一层的输出z^B以及学生模型最后一层的输出z^S。

S403：设置知识蒸馏的损失函数为BERT-Teacher和TextCNN的蒸馏MSE损失，该损失函数用以衡量教师模型与学生模型输出分布的差异，通过减少该损失能够使得学生模型学习到教师模型对任务的判断能力；以及TextCNN连接分类器后预测的标签与真实标签的交叉熵损失，用以衡量学生模型的分类损失，通过减少该损失能够增加学生模型分类的准确率，MSE损失如下：

L_Total＝β·L_CE+(1-β)·L_Distill

通过最小化蒸馏损失，使得学生模型的输出分布更接近教师模型，同时增加学生模型对真实标签的预测能力。

步骤S5，软件样本实时分类：将待测样本输入步骤S1-S2中得到相应的软件序列编码，并将所述序列编码载入所述训练好的学生模型中，以预测恶意软件家族类型对应的标签值，进而对恶意软件家族进行分类。

本发明性能评估采用开源的Windows恶意数据集，该数据集为基于VirusShare收集的恶意软件样本经分析后得到的若干API序列，其分类为包括木马、后门、下载器等在内的8种恶意软件类别，共计13849条API序列，由于数据集不平衡，对其大于300个样本的家族进行采样到300个样本。

本发明采用的性能评价指标是Accuracy，Precision，Recall，F1-Score这4个指标。

Accuracy指对于给定的数据，分类正确的样本数占总样本数的比例。Precision指在预测为正类的样本中实际上真正为正类的样本所占的比例。Recall指在实际上真正为正类的样本中预测正确的比例。Precision表示模型对负样本区分能力，而Recall则表示模型对正样本的区分能力。F1-Score指Precision和Recall的调和平均。

本发明与其他模型在上述数据集上预测效果对比如表1所示。

表1模型预测效果对比

从表1可以看出，通过在微调过程中加入对比学习，BERT-SCL的准确率有了较大的提升，同时，在蒸馏后的模型准确率依然也超过了BERT-Base模型；对于学生模型而言，BERT-Dis为蒸馏后的模型，TextCNN为蒸馏前的模型，可以看出通过知识蒸馏有效的提高了学生模型的检测效果，综合而言，相对于传统方法而言，本方法有较好的检测效果。

在模型大小方面，蒸馏前的BERT模型为110M，而蒸馏后的学生模型仅需1.88M，同时在推理时间上，蒸馏前的模型推理速度为9.09秒/100条数据，而蒸馏后仅需0.12秒/100条数据，体现了本发明在减少模型规模，增加模型检测效率上的有效性。

Claims

1.一种基于对比学习的恶意软件检测方法，其特征在于，包括如下步骤：

步骤S1，提取软件样本中的API序列，并进行数据清洗；

步骤S2，将进行数据清洗后的API序列预处理，得到基于Transformer的双向编码BERT的输入s_E；

将s_E送入BERT编码层进行处理，得到BERT编码层的输出O_BERT，将O_BERT作为软件API序列的初始编码；

步骤S3，使用O_BERT以及对应的标签信息对BERT模型进行微调，得到BERT-Teacher模型；

步骤S4，通过BERT-Teacher构造教师模型与学生模型，并进行训练；

步骤S5，将待测样本经过步骤S1-S2，得到相应的软件序列编码，并将序列编码载入训练好的学生模型，以预测恶意软件家族类型对应的标签值，进而对恶意软件家族进行分类。

2.根据权利要求1所述的一种基于对比学习的恶意软件检测方法，其特征在于，在S1步骤中，使用第三静态分析工具或者沙箱动态运行的算法提取软件样本的API序列，并通过反病毒引擎得到所述软件样本的标签。

3.根据权利要求2所述的一种基于对比学习的恶意软件检测方法，其特征在于，在S1步骤中，所述数据清洗具体操作为：

保留API序列中大小写英文字母和数字，其余的标点符号以及字符去除，然后将大写字母统一转化为小写字母。

4.根据权利要求3所述的一种基于对比学习的恶意软件检测方法，其特征在于，在S2步骤中，所述BERT模型由一个分词器、一个嵌入层、十二个双向Transformer编码层以及分类器顺序堆叠构成。

5.根据权利要求4所述的一种基于对比学习的恶意软件检测方法，其特征在于，在S2步骤中，对所述API序列预处理具体操作如下：

S201：将进行数据清洗后的的API序列经BERT模型的分词器分词，得到词表示w，并由词表示w构成输入序列s；

S202：将输入序列s截取或者补齐到统一长度，并在其两端加入BERT分类符[CLS]以及句子分界表示符[SEP]，分别表示句子的开始和结束，得到嵌入层的输入序列s′：

S′＝{w_[CLS]，w₁，w₂，...，w_n，w_[SEP]}

其中n表示所设定的统一长度的大小；

S203：将s′输入BERT的嵌入层，BERT的嵌入层通过结合词嵌入，段落嵌入以及位置嵌入三个嵌入算法形成嵌入向量：S_E＝{E_[CLS]，E₁,E₂,...,E_n,E_[SEP]}；

将s_E输入到BERT模型的双向Transformer编码层中，得到BERT编码层的输出O_BERT，O_BERT为恶意软件API序列的初始编码，E_[CLS]编码后得到的隐状态向量h_CLS。

6.根据权利要求5所述的一种基于对比学习的恶意软件检测方法，其特征在于，在S3步骤中，对BERT模型进行微调过程具体如下：

将隐状态向量h_CLS作为API序列的语义表示，并输入分类器中，分类器的通过两层全连接层，然后连接一个softmax函数进行分类，最后通过损失函数反向传播进行微调。

7.根据权利要求6所述的一种基于对比学习的恶意软件检测方法，其特征在于，所述损失函数构成如下：

结合交叉熵损失函数以及对比损失函数，得到最终的损失函数如下：

L_total＝αL_scl+(1-α)L_CE

其中，L_CE为交叉熵损失函数，t_SCL为监督对比损失函数，α为超参数。

8.根据权利要求7所述的一种基于对比学习的恶意软件检测方法，其特征在于，步骤S4具体操作如下：

S401：将得到的BERT-Teacher作为教师模型，根据教师模型并结合TextCNN构造学生模型，进行知识蒸馏；

S402：通过固定BERT-Teacher的参数不变，将s_E分别输入教师模型和学生模型中，得到教师模型的输出z^B以及学生模型的输出z^S；

S403：设置知识蒸馏的损失函数为BERT-Teacher和TextCNN的蒸馏MSE损失,以及TextCNN连接分类器后预测的标签与真实标签的交叉熵损失，MSE损失如下：

添加β作为蒸馏损失的超参数，则知识蒸馏的总损失函数如下：

L_Total＝β·L_CE1+(1-β)·L_Disti11

其中，L_CE1为学生模型通过实际标签与学生模型的预测输出计算的交叉熵损失。