【三等奖方案】基于人工智能的漏洞数据分类赛题「终于交上了」团队解题思路

郑力南，来自中国科学院计算技术研究所，硕士二年级在读，研究兴趣包括推荐系统、自然语言处理、机器学习。

随着信息安全逐渐受到广泛的关注，针对当前海量的漏洞数据，其统计和梳理工作面临巨大挑战，对此，本文提出了一种基于LightGBM和Bert模型集成的漏洞描述分类方法。本方案首先通过对原始漏洞描述的文本数据进行TF-IDF编码处理，随后训练LightGBM模型，实现对漏洞的分类。进一步地，本方案利用全量数据作为语料库预训练了一个具有专业领域知识的的掩码语言模型，本方案选用Bert作为基础模型 ，随 后使用标注数据进行微调。最后，将训练好的LightGBM和Bert进行集成，最终本方案的预测结果准确率为0.469471，位列排行榜第5名。

LightGBM、预训练、微调、集成

为及时跟踪国际信息安全趋势，需对国际公开的漏洞数据内容进行及时统计和梳理，例如CVE漏洞平台。CVE平台的漏洞信息包含有CVE编号、漏洞评分、漏洞描述等内容，其中漏洞描述含有对漏洞的利用条件、受影响的范围、漏洞可达到的效果(危害)等内容。为了更好地理解和持续研究，需将这些漏洞信息按照一定规则进行分类。而在此过程中，人工筛选分类效率较低，耗时耗力，利用人工智能，通过自然语言处理则可能更好地解决这一问题。

其中漏洞的描述信息可标注出以下几个标签：CVE-Number，漏洞编号；Privilege-Required，漏洞利用所需的攻击者权限；Attack-Vector，攻击向量；Impact，漏洞利用的结果。

根据题目要求，需要预测Privilege-Required、Attack-Vector、Impact-level1、Impact-level2和Impact-level3共5个字段，其中 Impact-level2和Impact-level3还受上一级 Impact影响。通过进一步分析Impact-level2和 Impact-level3的枚举，可以将原5个字段转化为 Privilege-Required、Attack-Vector、Impact-level1、Impact-Privilege和Information-Disclosure共5个字段。其中，Impact-Privilege为Privilege-Gained(RCE)的 Impact-level2，以及 local(credit)和other-target(credit)的Impact-level3下数据组成，枚举为{admin/root、Nonprivileged、unknown}；Information-Disclosure即为 Impact-level2中的Information-Disclosure，枚举为{local(credit)、other-target(credit)、other}。

通过统计，可以发现每个字段的都存在标签不均匀的现象，(例如，Attack-Vector中remote与non-remote之比达到了19.45)。另外，一些类别，如other、unknown，这些类别并不明确，其分布或许更为分散，给分类模型带来了挑战。

对于本题目，需要通过漏洞描述，来预测每个字段的分类。可以对词频进行初步统计，整理频率最高的前30个词，统计结果如图1所示。

图1：漏洞描述词频统计结果

可以发现，在频率最高的前30个词中，出现许多日常不常使用，而在特定领域中会经常使用的词汇，如remote、attackers、vulnerability、arbitrary、service、denial、users等，说明该数据集具备一定专业性。因此，对于预训练语言模型，一般使用通用预料库进行预训练，在下游任务上进行微调，在本题中，除已标注数据外，还提供了60000条无标注数据，我们可以利用这些文本数据，预训练一个语言模型，其包含了特定领域的知识，在下游任务中，相比于使用通用预训练语言模型更具有优势。

本方案利用漏洞描述(文本信息)，来预测5个字段的分类，本质上，将该问题建模为文本多分类问题。本题可形式化描述为：对于给定训练集x，训练5个分类器f_i，用于预测5个字段的分类，其输出为分类的概率分布。对于样本x，在第i个字段中，第j个枚举值的概率通过式(1)计算得到。

文本分类问题解决方案大体包含两个部分：文本表示和分类器。对于文本表示，常用方法包括基于统计的表示方法，如词袋模型；或基于模型的表示方法，如 TextCNN^[1]，FastText^[2]，BERT^[3](encoder)等。对于分类器，常见的模型如：朴素贝叶斯^[4]、SVM^[5]、LR^[6]、决策树^[7]、GBDT^[11]等。在参赛过程中，共尝试了2种方案，分别为TF-IDF+LightGBM^[8]和BERT+LR，接下来会对两种方案做详细地介绍。

3.1 TF-IDF+LightGBM

本方案中，首先使用词袋模型计算文本的初始向量表示(88633 维)，再利用逆文档频率(IDF)进行加权，最后通过 L2 标准化得到最终的向量表示。

将训练集中每一个样本转化为向量表示之后，使用LightGBM 作为分类器。LightGBM(Light Gradient Boosting Machine)是一个实现 GBDT^[11]的算法框架，支持高效率的并行训练，并且具有更快的训练速度、更低的内存消耗、更好的准确率、支持分布式可以快速处理海量数据等优点。^[8]

使用 GBDT 类模型，在多分类任务中需要考虑 softmax模型，需要n个不同的CART回归树集成，对应F₁……F_n。对于一个n分类模型，第k类的概率，通过式(2)计算得到。

3.2 BERT + NN

本方案中，使用 BERT encoding 作为文本的向量表示，当完整句子输入 BERT 后，将 BERT 的pooler输出(即第一个token对应的向量)后接一个单层神经网络，最后经过softmax得到类别的概率分布，实现对文本的分类。具体而言，本方案包含两个步骤：

1. 使用特定领域语料(包含已标注数据和未标注数据)预训练BERT；

2. 在下游分类任务上微调预训练好的模型。

接下来两小节会对两个步骤进行详细说明。

3.2.1 预训练模型

首先，使用全部漏洞描述数据(包括标注数据和未标注数据)训练一个掩码语言模型(Masked Language Model,MLM)，将原文本中15%的token随机替换为[MASK] ，然后让被改造后的序列去预测被替换为[MASK]的是哪些词，如图2所示。

图 2：掩码语言模型示意图

训练完成后得到一个包含了特定领域知识的语言模型，用于生成漏洞描述表示。具体而言，以Bert作为backbone，共包含6层自注意力层 ^[9]，自注意力头的数量为12，隐藏层神经元个数为768，token 最大长度为 512。

3.2.2 微调模型

在预训练好 Bert 模型后，将原始漏洞描述文件进行token化之后，拼接上[CLS]前缀，输入Bert模型进行encoding，取输出表示序列的第一个位置上的结果，即[CLS]的向量表示，将其输入到一个单层神经网络，输出类别的概率分布，其网络结构如图3所示：

图 3：微调阶段模型示意图

在训练时时，为了解决分类问题中类别不平衡、分类难度差异的问题，分类损失函数使用focal loss^[10]，其中超参数α的确定通过LightGBM的预测结果计算得到。公式(3)给出了一个启发式确定α的方法：

其中，表示类别?中预测正确的样本的概率分布熵的均值，表示类别c预测的准确率，s^(c)表示分类c的分类难度，如果熵值大或者准确率低，则说明分类难度大，反之同理。?即为?的归一化结果。 

对于LightGBM模型，各分类任务使用的参数如表1所示。

对于Bert模型在预训练阶段，batch size设为10，使用warm up策略调整学习率，前10%次迭代学习率递增，达到0.0001后递减，共训练10个epoch；在微调阶段，冻结Bert的前5层自注意力层权重，只保留最后一层权重和分类器权重，batch size设为16，使用warm up策略调整学习率，前10%次迭代学习率递增，达到0.00003后递减，共训练500个epoch。

分别采用5折交叉验证，训练5个分类模型，验证集上的准确率随训练轮数的变化如图4所示。

图 4：验证集上的准确率随训练轮数的变化

其中，Attack Vector(二分类)分类效果最好，而Impact Privilege和Information Disclosure分类效果较差，因为其样本数量只占整体样本中的一部分。

使用LightGBM 模型与Bert模型的结果如表2所示。

可以发现，Bert模型在大部分字段上的准确率都有所提升，其中在Attack Vector字段上的准确率却下降了2.37%，Attack Vector是二分类问题，发现Bert模型在二分类问题上的表现不如LightGBM模型，并且Attack Vector中的remote与non-remote之比为19.45:1，分布较为不均匀，因此对深度学习模型也造成了训练上的困难。

另外，本方案还尝试了指数移动平均衰减(EMA decay)和focal loss不同的γ值，最后对预测分布加权求和后再经过softmax得到新的分布，参数设置如表3所示。

实验发现，使用Bert模型，预测分布熵值较低，而LightGBM模型预测分布则相对更加平滑，为防止Bert模型过拟合，使用Bert模型(ensemble)预测的分布与LightGBM模型预测的分布逐元素相乘，得到每类的分数，依据此分数，得到预测结果。最终本方案的预测结果准确率为0.469471，位列排行榜第5名。

[1] Zhang Y, Wallace B. A sensitivity analysis of (and practitioners’ guide to) convolutional neural networks for sentence classification[J]. arXiv preprint arXiv:1510.03820, 2015.

[2] Joulin A, Grave E, Bojanowski P, et al. Bag of tricks for efficient text classification[J]. arXiv preprint arXiv:1607.01759, 2016.

[3] Devlin J, Chang M W, Lee K, et al. Bert: Pre-training of deep bidirectional transformers for language understanding[J]. arXiv preprint arXiv:1810.04805,2018.

[4] Zhang H. The optimality of naive Bayes[J]. Aa, 2004, 1(2): 3.

[5] Hearst M A, Dumais S T, Osuna E, et al. Support vector machines[J]. IEEE Intelligent Systems and their applications, 1998, 13(4): 18-28.

[6] Wright R E. Logistic regression[J]. 1995.

[7] Loh W Y. Classification and regression tree methods[J]. Encyclopedia of statistics in quality and reliability, 2008, 1: 315-323.

[8] Ke G, Meng Q, Finley T, et al. Lightgbm: A highly efficient gradient boosting decision tree[J]. Advances in neural information processing systems, 2017, 30.

[9] Vaswani A, Shazeer N, Parmar N, et al. Attention is all you need[J]. Advances in neural information processing systems, 2017, 30.

[10] Lin T Y, Goyal P, Girshick R, et al. Focal loss for dense object detection[C]//Proceedings of the IEEE international conference on computer vision. 2017: 2980-2988.

[11] Natekin A, Knoll A. Gradient boosting machines, a tutorial[J]. Frontiers inneurorobotics, 2013, 7: 21.

—End—