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

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

4.1 基于对比学习的方法
在本节中, 将介绍基于对比学习的漏洞分类框架 VulFewShot.
4.1.1 概 述
如图 5 所示, VulFewShot 由两个主要阶段组成: 训练和分类阶段. 训练阶段的目的是训练出一个性能良好的分
类模型. 这个阶段由以下 3 个步骤组成. (1) 代码静态分析: 该步骤旨在从漏洞样本中提取代码片段, 对代码片段进行
标注和归一化. (2) 向量提取: 该步骤旨在将规范化的代码片段转化为向量表示. (3) 基于对比学习的模型训练: 这一
步使用上一步的向量表示, 通过对比学习的方式对传统深度学习模型进行微调, 训练出准确的分类模型. 其中, 对比
学习旨在最大化相似数据之间的相似度并最小化不相似数据的相似度. 这里我们并未使用数据增强来生成新样本与
原始样本对比, 而是将训练中不同轮次中的相同漏洞样本类型作为对比, 具体的方法将在第 4.1.3 节进行介绍.

步骤1: 代码静态分析 步骤3: 基于对比学习的模型训练/漏洞分类
训练阶段
预测
密集层 Softmax 层
c
训练批次B
漏洞样本
DL
代码片段 标签 模型
Code gadget_1 1 微调 SupCon loss CE loss
Code gadget_2 2
Code gadget_3 3 预测
… …
训练批次B
密集层 Softmax 层
步骤2: 向量提取
向量表示
[v 0 , v 1 , v 2 , v 3 ,…, v n ] 测试阶段 具有固定参数的 DL 模型
DL
训练批次B 密集层 Softmax 层
模型 预测

10路交叉验证

图 5 VulFewShot 的系统概述

分类阶段的目标是将未标记的漏洞分类为相应的类型. 该阶段包括以下 3 个步骤: (1) 代码静态分析, (2) 向量
提取, (3) 漏洞分类. 前两步的执行方式与训练阶段完全相同. 第 3 步应用从对比学习框架中训练出来已固定参数
的模型对漏洞进行分类, 然后将一批未标记漏洞样本的向量表示作为输入, 输出相应的类型标签.
4.1.2 代码静态分析&向量提取
正如第 2.2 节和图 1 所示, 在µVulDeePecker 中按照代码静态分析和向量提取的步骤处理数据. 首先, 从代码
片段文本中删除无关的非 ASCII 字符和注释. 其次, 在符号名称和代码之间应用一对一的映射, 包括“VAR”到变
量, “FUN”到函数, “STR”或“NUM”到常量等. 最后, 我们生成一个序列化文件, 该文件包含 3 类信息: 代码片段、
代码标签和规范化后的代码片段. 其中代码标签来自µVulDeePecker.
在向量提取的步骤中, 将规范化的代码片段转换为固定长度的向量. 转换是通过在µVulDeePecker 中使用的
一个简单的词嵌入实现的. 这些向量是训练模型的输入.
4.1.3 基于对比学习的模型训练
在日常生活中, 孩子们只需要几张动物图片就能识别出动物之间的差异 [65] . 人类在只提供少量样本的情况下,
往往会在比较不同类别样本的同时寻求同一类别样本之间的相似性, 这使我们能够捕捉到物体的高级特征, 从而
将其与其他物体区分开来.
在这种学习策略的启发下, 我们假设基于对比学习的损失可以驱动模型提取样本的更高级特征. 这样, 模型可