Page 116 - 《软件学报》2025年第12期

吴月明 等: VulFewShot: 利用对比学习改进少样本漏洞分类 5497

1 相关工作

目前基于深度学习的漏洞分类研究主要有两类, 分别是: 二分类 [14,16,17,19−22] 和多分类 [23,26−28] . 我们的研究集中
在后者上.
深度学习模型以向量作为输入, 因此在使用这些模型时首先需要提取漏洞程序的特征向量表示. 提取特征向
量的过程主要包括两个步骤: 第 1 步是对漏洞程序的数据进行静态处理 (例如删除注释、提取片段、归一化等),
第 2 步是生成代码嵌入 (即将第 1 步获得的数据转换为向量表示). 研究人员提出了不同的方法来处理第 1 步中的
数据, 一种简单的方法是直接删除注释并规范化 [22,29,30] . 也有方法专注于使用抽象语法树 [31] 或函数调用 [32] 来提取
[35]
特定特征. 一些其他方法可能在代码分析工具如 Joern 的帮助下构建程序依赖图 (PDG) 或系统依赖图 (SDG) .
[33]
[34]
在第 2 步中, 除了有使用包括 Word2Vec、GloVe 和 FastText [16,20,23] 在内的一些主流嵌入方法, 也出现了一些新的
具有不同特长方法. 例如基于图像的嵌入方法可以辅助对漏洞的视觉识别 [22] , 而基于 CodeBERT 的嵌入方法更适
合下游任务 [29] 以及基于图的嵌入可以更好地描述代码连接 [30,36] . μVulDeePecker [23] 是一个使用最先进技术的多分
类漏洞分类系统, 它为数据处理和开源数据集提供了一种经过验证的方法论. 因此我们遵循了 μVulDeePecker 的
数据处理和向量表示方法, 并在其构建的数据集上进行了实验.
受不同漏洞类型之间严重的数据不平衡的启发, 我们引入对比学习来增强当前在少样本类型的漏洞上的分类
性能. 自 Mikolov 等人 [37] 将对比学习引入自然语言处理 (NLP) 以来, 对比学习已逐步应用于图像、语音和视频等
其他领域 [38−50] . 例如, SimCLR [50] 是将对比学习应用于图像分类任务的典型工作. 在之前的大多数研究中, 对比学
习主要用于自监督学习. Khosla 等人 [51] 将对比学习方法扩展到监督学习, 他们通过对比训练数据的标签信息来提
高分类性能. Wu 等人 [52] 将对比学习引入到 Android 恶意软件分类中, 这有助于减少混淆性差异, 同时增大了恶意
软件和良性应用程序之间的区别. Gunel 等人 [53] 将监督对比学习应用于文本分类任务, 在交叉熵损失的基础上增
加了一个额外的监督对比学习损失函数 (SCL). 新设计的损失函数使同一类样本更靠近, 同时使得不同类别的样
本之间距离更远. 他们还通过实验证实了对比学习可以带来少样本学习场景的巨大改进, 这给了我们一些启示. 我
们的实验结果表明, 对比学习方法可以提高整体性能和在不同规模的漏洞类型上的表现, 尤其是在少样本类型上.

2 漏洞分类

图 1 展示了漏洞分类的一般过程, 该过程一般包含 3 个步骤: 数据预处理、向量提取和模型训练/预测. 在数
据预处理中, 漏洞样本会被提取为代码片段并根据预先定义的代码标签进一步提取为标准化函数片段. 在向量提
取中, 标准化的函数片段会被转化为向量表示. 最后在模型训练/预测阶段, 一部分漏洞样本对应的特征向量将被
用作输入训练用于分类任务的深度学习模型, 深度学习模型使用训练阶段未使用的漏洞样本进行测试, 并根据第
2.3 节中的评估指标进行评估.

数据预处理 向量提取 模型训练/预测
步骤1: 生成代码片段 步骤2: 标记代码片段 步骤5: 训练
输入 步骤4: 将代码片段 深度学习模型
转换成向量表示
代码片段 标签 模型
生成 SDG Code gadget_1 1 符号名称 向量 CNN
Code gadget_2 2 BiGRU
漏洞样本 Code gadget_3 3 Static v 0 BiLSTM
… … Int v 1
识别库/API 函数调用 FUN1 v 2 Atten-BiLSTM
步骤3: 标准化函数片段 … ( …
v 3
将人为定义的变量和函数 步骤6: 漏洞分类预测
提取以库/API 函数 按照1对1的方法映射为
调用为起点的代码片段 符号名称 向量表示
(例如, VAR1 & FUN1) [v 0 , v 1 , v 2 , v 3 ,…, v n ] 漏洞类型标签
图 1 漏洞分类的一般过程