引言

在现代数字环境中，软件几乎渗透到了我们日常生活的方方面面。随着这些系统的规模和复杂性不断增长，其中包含的安全漏洞种类也日益繁多。对于安全分析师而言，需要审查的代码量是惊人的。仅在 2023 年，国家漏洞数据库 (NVD) 就发布了超过 28,900 个新的通用漏洞披露 (CVE) 条目。令人不安的是，其中超过 4,000 个案例在很长一段时间内仍未被归类为具体的类型。

仅仅识别出一代码片段是“脆弱的”只是第一步。更困难——也更关键——的一步是对其进行分类 。 这是缓冲区溢出？是双重释放？还是 SQL 注入？这种分类由通用缺陷枚举 (CWE) 系统处理。准确的分类使开发人员能够根据严重程度确定修复的优先级，并应用正确的补救策略。

人工分类速度缓慢，且严重依赖专家知识。虽然深度学习 (DL) 彻底改变了图像和文本处理，但将其应用于漏洞分类面临着一个独特的障碍: 数据极度不平衡。

论文 “Applying Contrastive Learning to Code Vulnerability Type Classification” (将对比学习应用于代码漏洞类型分类) 背后的研究人员提出了一个新颖的框架来正面解决这种不平衡问题。通过利用漏洞类别的自然层次结构并使用一种称为对比学习的技术，他们开发了一种模型，不仅能检测漏洞，还能理解漏洞之间细微的关系。

在这篇文章中，我们将解构他们的方法，解释他们如何克服以往模型的局限性，并分析为什么他们的方法为自动化安全分析树立了新标准。

问题所在: 长尾与孤立的类别

要理解这篇论文的创新之处，我们需要先了解数据问题。在软件漏洞的世界里，并非所有的 bug 都是生而平等的。有些漏洞，如跨站脚本 (XSS) 或越界写入，出现得非常频繁。而其他的则很少出现。

当你绘制这些漏洞类型的频率图时，你会得到所谓的长尾分布 。

图 1: NVD 2023年新发布漏洞的 CWE 类型分布，呈现长尾分布特征。

如图 图 1 所示，分布的“头部”包含少量非常常见的漏洞类型 (CWE)。“尾部”则延伸出大量样本极少的类型。

为什么这会破坏机器学习？

标准的深度学习模型非常依赖数据。它们擅长学习“头部”的类别，因为它们能看到成千上万个例子。然而，它们很难学习“尾部”的类别。在标准的分类设置中，模型可能会完全忽略稀有类别，以最大化其在常见类别上的整体准确率。

此外，以往的方法将每个 CWE 类别视为孤岛。它们会将“CWE-415 (双重释放)”和“CWE-416 (释放后使用)”视为完全不同的标签，而忽略了它们都是内存管理错误这一事实。这导致代码表征 (模型“思考”所在的向量空间) 不具备可扩展性或语义丰富性。

核心方法: 层次化对比学习

研究人员提出了一个框架，改变了模型的学习方式。模型不再仅仅问“这是哪个标签？”，而是问“这段代码与同类别的其他代码有何相似之处，以及它如何融入更广泛的漏洞层次结构中？”

这是通过层次化对比学习实现的。

1. 理解层次结构

CWE 标准不是一个扁平的列表；它是一棵树。在顶层，你有抽象类别 (支柱 Pillars) 。随着向下移动，类别变得更加具体 (类别 Classes、基础 Bases 和变体 Variants) 。

图 2: CWE-415 从高层抽象到低层抽象的细化链。

图 2 展示了这个细化链。以 CWE-415 (双重释放) 为例。

  1. 支柱 Pillar (664): 资源生命周期的不当控制。
  2. 类别 Class (118/119): 内存缓冲区内操作的不当限制。
  3. 变体 Variant (415): 双重释放。

研究人员意识到，如果模型学习了这种层次结构，它就能做出更好的预测。即使它没有见过很多特定的稀有“变体”示例，它可能见过大量其“支柱”或“类别”的示例，从而能够进行有根据的推断。

2. 架构概览

所提出的框架遵循特定的流程，旨在处理源代码并通过多层对比学习来完善其理解。

图 3: 我们的层次化对比学习框架架构。

如 图 3 所示，工作流程如下:

  1. 标签扩展 (Label Expanding): 代码片段的单个标签 (例如 CWE-415) 被扩展为一个包含 5 个标签的链，代表其在 CWE 树中的完整祖先。
  2. 分词与最大池化 (Tokenization & Max-Pooling): 代码被转换为模型可以读取的 token (数字) 。
  3. 层次化对比学习 (Hierarchical Contrastive Learning): 核心训练阶段，模型在层次结构的每一层对齐代码表征。
  4. 分类器 (Classifier): 最后一层预测具体的漏洞类型。

3. 使用最大池化处理长代码

Transformer 模型 (如 BERT) 通常有严格的输入限制 (例如 512 个 token) 。漏洞代码通常比这长得多。以前的方法只是简单地切掉代码的末尾 (截断) ，这可能会丢失实际的漏洞行。

作者使用最大池化 (Max-Pooling) 解决了这个问题。他们将长代码分成多个块 (例如两个 512 的块) 。他们将两个块都通过模型运行，然后应用最大池化操作。这有效地保留了来自两个块的“最强”信号，使模型能够“看到”更长的上下文，而不会超出内存限制。

4. 数学引擎: 对比损失

这是论文的核心。标准分类使用交叉熵损失 (Cross-Entropy Loss) , 它基于概率惩罚模型。

交叉熵损失公式

虽然有用，但交叉熵并不明确强制模型将相似的事物在向量空间中归为一组。为了解决这个问题，作者引入了对比损失 (Contrastive Loss) 。

自监督对比损失

首先，他们使用自监督方法。他们取一个代码样本，创建一个它的“增强”版本 (可能是通过掩盖一些单词) ，然后告诉模型: “这两个是一样的。把它们拉近。把其他所有东西推开。”

自监督对比损失公式

在这个公式中 (图片 004) ，\(z_i\) 和 \(z_{j(i)}\) 是原始代码和增强代码的表征。目标是最大化它们的相似性，同时最小化与其他代码的相似性。

监督对比损失

然而，我们要用到标签！所以我们应该利用它们。在监督对比学习中，模型被告知: “所有带有‘缓冲区溢出’标签的样本都应该聚集在一起。”

监督对比损失公式

这个公式 (图片 006) 确保所有正样本对 (同一类别的样本) 在向量空间的超球面上被拉近。

类坍缩 (Class Collapse) 问题

如果仅仅依赖监督对比学习，你可能会面临类坍缩的风险。这种情况发生时，模型会使每一个“缓冲区溢出”向量完全相同。它失去了不同缓冲区溢出实例之间的细微差别。

为了防止这种情况，作者结合了所有三种损失。他们使用监督损失来聚类类别，使用自监督损失来保持独特的几何分布 (鲁棒性) ，并使用交叉熵来保证最终的分类准确性。

总损失公式

最终的损失函数 (图片 007) 是这三个分量的加权和，由参数 \(\lambda\) 和 \(\mu\) 控制。

实验与结果

研究人员将他们的框架与几种最先进的基线模型进行了评估，包括 CodeBERT、GraphCodeBERT 以及像 VulExplainer 这样的专用漏洞模型。他们使用了 Big-Vul 数据集和一个更新、更高质量的数据集 PrimeVul 。

性能比较

结果是决定性的。所提出的层次化对比学习 (HCL) 框架显著优于现有方法。

表 1: 我们的层次化对比学习方法与基线模型的比较结果。

如 表 1 (图片 008) 所示，查看 加权 F1 (Weighted F1) 分数 (一种在考虑类别不平衡的同时平衡精确率和召回率的指标) :

  • 在 Big-Vul 上，所提出的方法 (使用 CodeBERT) 达到了 65.34% , 而标准 CodeBERT 仅为 43.07%。
  • 在高质量的 PrimeVul 数据集上，它达到了 41.24% , 显著击败了最接近的竞争对手。

“Tier 1”到“Tier 5”列显示，该模型不仅在抽象层面表现良好，而且在具体的变体层面也表现出色。

为什么它有效？ (消融实验)

为了证明他们的具体设计选择很重要，作者进行了消融实验 (Ablation Study) 。 他们系统地移除了模型的各个部分，以查看性能是否下降。

表 2: 消融实验的实验结果。

表 2 (图片 009) 揭示了每个组件的贡献:

  1. 第 1 行 vs. 第 2 行: 仅添加层次化对比学习 (HCL) 就将准确率从 63.19% 提高到了 66.92%。
  2. 第 2 行 vs. 第 3 行: 添加“USCL” (无监督/自监督损失) 以防止类坍缩，进一步将性能提高到 68.31%。
  3. 第 6 行 (完整模型) : 结合 HCL、USCL 和最大池化 (MP) 产生了 69.06% 的最佳结果。

这证实了架构的每一部分——层次结构、几何分布和长文本处理——都是必不可少的。

鲁棒性及与 LLM 的比较

最后，研究人员测试了模型对损失函数权重 (\(\lambda\) 和 \(\mu\)) 的敏感度，并将其与 GPT-4 (零样本) 进行了比较。

图 4 和表 3: 敏感度分析及 GPT-4 对比。

图 4 (图片 013，上图) 中的图表显示模型相对稳定；在合理的超参数值范围内，准确率保持在较高水平。

更有趣的是, 表 3 (图片 013，下图) 显示，像 GPT-4 这样的通用大型语言模型 (LLM) 在这项特定任务上表现挣扎。在零样本提示下，GPT-4 仅达到了大约 22-25% 的准确率 。 即使使用思维链 (Chain-of-Thought) 提示，它也远远落后于专用的对比学习模型。这强调了虽然 LLM 功能强大，但对于像漏洞分析这样的复杂分类任务，专门的领域特定架构仍然具有优势。

结论与启示

论文 “Applying Contrastive Learning to Code Vulnerability Type Classification” 代表了自动化软件安全领域向前迈出的重要一步。通过拒绝将漏洞类型视为孤立的标签，转而拥抱其自然的层次结构，研究人员创建了一个能够“理解”软件弱点之间关系的模型。

关键要点:

  1. 层次结构很重要: 利用 CWE 树结构使模型能够学习相关漏洞之间的共享特征，从而缓解长尾问题。
  2. 对比学习很强大: 明确地将相似向量拉近，比简单的分类能产生更鲁棒、更有意义的代码表征。
  3. 细节决定成败: 针对长代码的最大池化和混合自监督损失以防止类坍缩等技术调整，在准确率上提供了可观的收益。

对于该领域的学生和研究人员来说，这项工作强调了表征学习的重要性。这不仅仅是建立一个更大的 transformer；而是要设计训练过程以反映数据的语义现实。随着软件系统不断增长，建立在这些原则之上的工具对于保持系统安全至关重要。