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

余俊晖：硕士，研三在读，来自于贵州大学-文本计算与认知智能实验室，主要研究方向为自然语言处理信息抽取、智能问答。曾获得过NLP领域算法竞赛TOP奖项十余项，有着丰富的NLP算法竞赛经验。

冯琳慧：硕士， 主要研究方向为自然语言处理领域的信息抽取。

刘焕勇：硕士，就职于中科院软件所，主要研究方向为知识图谱。

为及时跟踪国际信息安全趋势，需对国际公开的漏洞数据内容进行及时统计和梳理，例如CVE漏洞平台。CVE平台的漏洞信息包含有CVE编号、漏洞评分、漏洞描述等内容，其中漏洞描述含有对漏洞的利用条件、受影响的范围、漏洞可达到的效果(危害)等内容。为了更好地理解和持续研究，需将这些漏洞信息按照一定规则进行分类。而在此过程中，人工筛选分类效率较低，耗时耗力，利用人工智能，通过自然语言处理则可能更好地解决这一问题。本文提出一种多任务文本分类算法，通过共享BERT embedding，设计多任务的损失函数，端到端的解决漏洞数据分类问题。最后通过大量的实验及相关模型优化，最终取得了A榜第二B榜第一的成绩。

漏洞数据、文本分类、多任务、模型优化

1、给定漏洞数据的描述，将漏洞数据分类。

2、给出6w条无标注数据，可根据赛题给出的介绍进行标注，但不能对测试数据进行标注。

图1 漏洞建模参数示意图

2.1 网络节点漏洞信息建模规则

我们需要把漏洞的描述信息，标注出以下几个标签：

●　CVE-Number ： 即漏洞库中的漏洞cve编号，用来对漏洞详细信息进行收录、管理与识别等。

●　Privilege-Required ：即漏洞利用所需的攻击者权限；攻击者只有先获取到利用该漏洞的前提权限，才能对该节点发动原子攻击，权限等级的枚举为：{admin/root、Nonprivileged、access、unknown}, admin/root表明攻击者获取到了最高权限，为单一标签；Nonprivileged表明攻击者可以以普通权限(非特权)的身份操作主机，access表明仅拥有节点的访问权限(网络可达)；unknown为从描述中获取不到该标签对应的内容，即根据描述得不到需要的权限信息。

●　Attack-Vector ：攻击向量(介质)；攻击者是通过某种介质来利用该漏洞的，向量的枚举为{remote、Non-remote },remote表示可通过远程方式完成攻击，Non-remote表示不能通过远程的方式来完成攻击(即需要通过本地物理接触(cvss中的physical))。

●　Impact ：漏洞利用的结果(完成漏洞利用获取到的权限)，可分为三级，分别为(Impact-level1、Impact-level2、Impact-level3)，Impact-level1的枚举为{ Privilege-Gained(RCE)、Information-disclosure、DoS、access、other}, Privilege-Gained(RCE)在Impact-level2下的枚举为{ admin/root、Nonprivileged、unknown}，(information-disclosure)在Impact-level2下的枚举为{local(credit)、other-target(credit)、other}, Impact-level3即local(credit)细化枚举为{user/root、Nonprivileged、unknown}、other-target(credit)枚举为{user/root、Nonprivileged、unknown}；Privilege-Gained(RCE)表示可获取到的权限：admin/root表示该漏洞利用可获取到该原子的最高权限，注意无论是root权限或者admin权限均为标签admin/root；Nonprivileged表示该漏洞利用可获取到该原子的普通非特权权限；unknown即为从描述中获取不到权限的信息，即为不清楚；Information-disclosure表示该漏洞为信息泄漏，而其中local(credit)表示该漏洞可获取到该原子的登录凭证；other-target(credit)表示该漏洞可获取到其他原子的登录凭证，other表示为该漏洞泄漏的是没有登录凭证类的其他信息；DoS表示该漏洞为拒绝服务漏洞；access表示可获取的网络的访问状态，(例如：通过攻击防火墙，只是修改策略，可造成部分本来不能访问的端口服务可正常访问)，other即为其他不属于上述的影响。

2.2 塞梯数据分析

根据我们在本次比赛中使用的算法，我们将类别标签映射成如图2：

图2 类别标签映射

Attack-Vector类别标签分布，如下图3：

图3

Impact类别标签分布，如下图4：

图4

Privilege-Required类别标签分布，如下图5：

图5

从图3、4、5可以得知，本次比赛的数据类别极不平衡，也对我们后续实验增加了许多困难。

在本次比赛中，我们基于BERT^[1]预训练语言模型，先后设计了两种方案进行实验。

3.1 基于BERT的3分类头的文本分类模型

如图6所示，我们对于每个大类，分别使用BERT+Linear分别对Impact、Attack-Vector及Privilege-Required三个 大类进行分类。

图6 基于BERT的3分类头的文本分类模型

对于该模型，我们进行了相关实验，表1显示了在A榜上实验结果：

从表中分析可以看出，将最优的分类头的最佳模型组合在一起得到的总体性能反而得不到最好结果。具体原因由下列示意图（7）分析演示得到，在各个模型取得性能都较高时，组合性能不一定是最好的：

图7

3.2 基于BERT的多任务端到端的文本分类模型

如图8，我们将BERT得embedding层共享，设计基于多任务的联合模型，端到端的训练分类模型，使各分类任务之间产生联系。则相应的损失函数为公式（1）：

Loss_总=Loss_cls2+Loss_cls4+Loss_cls13

图8 基于BERT的多任务端到端的文本分类模型

3.3 模型优化

我们基于上述模型进行了大量实验进行优化，对于本次比赛，优化能提升的trick有：(1)拟合目标：使用F1-score替换acc作为评价指标；(2)类别不平衡：使用Focal Loss^[2]函数；(3)早停机制^[3]；(4)对抗训练：使用FreeLB^[4]对抗训练，对bert的embedding进行扰动；(5)差分学习率；(6)Dropout；(7)权重衰减；(8)数据增强：为了保守起见，我们阅读比赛规则及数据标注相关描述，从6w条无标注数据中选取了1600条数据进行标注。

3.4 其他

在本次我们还尝试了基于lgb的特征工程方法，因为我们发现比赛数据存在数据leak，比如：ID存在leak。但是鉴于保险起见，我们最终还是选择了使用我们自己设计的模型及数据增强方式提交了B榜。

4.1 测评得分

4.2 模型大小

模型仅使用了bert-base预训练模型进行训练，大小为420M。

4.3 模型性能

训练耗时：训练5轮约2分钟

推理耗时：对测试集预测约30秒

本文基于BERT预训练语言模型，设计多任务学习目标，共享bert的embedding层，端到端的对漏洞数据进行分类，更加贴近现实场景，且训练速度及推理速度都具备优越性，更易于落地。

中国计算机学会&大数据协同安全技术国家工程研究中心&中国科学院信息工程研究所&360未来安全研究院工业互联网实验室。

[1] 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.

[2] 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.

[3] Prechelt L. Early stopping-but when?[M]//Neural Networks: Tricks of the trade. Springer, Berlin, Heidelberg, 1998: 55-69.

[4] Zhu C, Cheng Y, Gan Z, et al. Freelb: Enhanced adversarial training for natural language understanding[J]. arXiv preprint arXiv:1909.11764, 2019.

—End—