大家读完觉得有帮助记得关注和点赞!!!
抽象
软件漏洞的激增对网络安全构成了重大挑战,需要更有效的检测方法。我们介绍了 White-Basilisk,这是一种新颖的漏洞检测方法,它展示了卓越的性能,同时挑战了 AI 模型扩展中的普遍假设。利用集成了 Mamba 层、线性自我注意和 Mix of Experts 框架的创新架构,White-Basilisk 在参数计数仅为 200M 的漏洞检测任务中取得了最先进的结果。该模型能够处理前所未有的长度序列,从而能够一次性对广泛的代码库进行全面分析,超越了当前大型语言模型 (LLM) 的上下文限制。White-Basilisk 在不平衡的真实数据集上表现出强大的性能,同时保持计算效率,便于跨不同组织规模进行部署。这项研究不仅建立了代码安全的新基准,而且还提供了经验证据,证明紧凑、高效设计的模型可以在专业任务中胜过大型模型,从而有可能重新定义特定领域应用程序的 AI 开发中的优化策略。
1介绍
人工智能的最新进展在很大程度上是由神经网络的快速扩展推动的,其中尖端性能通常与可扩展至数千亿个参数的模型相关联。 虽然这种扩展范式产生了令人印象深刻的功能,但它也带来了重大的计算、环境和可访问性挑战(Patterson 等人,2021;Faiz 等人,2023).随着该领域继续优先考虑越来越大的模型,一个重要的问题出现了:是否存在架构创新和目标设计可能优于纯粹参数数量的领域?这个问题对于具有特定计算限制的高风险应用程序尤其相关,在这些情况下,“越大越好”的方法可能被证明是次优的或在经济上不可行。
代码漏洞检测恰恰代表了这样一个领域,一个具有重大实际重要性和独特技术挑战的领域。其经济影响是巨大的:到 2025 年,网络犯罪预计将达到 10.5 万亿美元(网络安全风险投资,2023).这种巨大的经济影响在很大程度上是可以预防的,因为 Ponemon Institute 发现 60% 的泄露受害者由于未修补的已知漏洞而受到损害,而 62% 的组织表示,他们完全不知道他们的系统在数据泄露之前是脆弱的(Ponemon 研究所,2018).在第三方软件中发现漏洞时,财务后果非常严重,组织平均每次泄露事件损失 455 万美元(UpGuard,2023).尽管日益复杂的软件系统导致威胁形势不断增长,攻击面不断扩大,但现有的自动漏洞检测方法面临重大限制。传统的静态应用程序安全测试 (SAST) 工具的检测率不足(周 et al.,2024),而最近基于大型语言模型 (LLM) 的机器学习方法需要大量的计算资源,这使得它们在许多实际部署场景中不切实际。此外,检测漏洞通常需要分析极长的代码序列并捕获跨多个函数或文件的细微交互,由于上下文长度限制和低效的注意力机制,即使对于最大的模型来说,这些功能仍然具有挑战性。
目前的研究探索了应对这些挑战的各种方法。在代码上微调的基于 Transformer 的模型已经显示出有希望的结果(Feng 等人,2020;Guo et al.,2022),但它们的二次注意力复杂性限制了它们对大型代码库的应用。专门的漏洞检测架构(哈尼夫和马菲斯,2022;Fu 和 Tantithamthavorn,2022)在这些基础上进行改进,但仍在与长期依赖关系和资源效率作斗争。值得注意的是,即使这些模型的大小和复杂性增加,它们在实际漏洞检测基准上的性能也往往达不到实际要求(Ding 等人,2024).模型规模和实际有效性之间的这种脱节表明,有机会在这项专业任务的背景下重新审视基本架构假设。
我们介绍了 White-Basilisk,这是一个紧凑的 200M 参数模型,它在代码漏洞检测任务上的表现优于更大的模型,从而挑战了传统的扩展智慧。我们的方法与扩展现有架构的标准做法不同,而是专注于设计满足漏洞检测独特要求的专用混合架构。推动我们工作的关键见解是,通过将状态空间模型 (Mamba) 与线性复杂性注意力机制和稀疏激活的 Mix of Experts 相结合,我们可以有效地处理极长的代码序列,同时捕获对漏洞识别至关重要的局部句法模式和全局依赖关系。
我们的贡献可以总结如下:
- 新颖的混合架构:我们提出了一种资源高效的模型,该模型集成了用于局部模式识别的 Mamba 层、用于全局上下文建模的线性复杂性关注以及通过 Mix of Experts 进行的条件计算,只需 200M 个参数即可实现有效的漏洞检测。
- 扩展上下文处理:White-Basilisk 在推理过程中分析多达 128000 个令牌的序列,比典型的 transformer 模型长 8× 多,从而能够在单个 GPU 上对整个代码库进行全面分析,并检测跨多个函数或文件的漏洞。
- 最先进的性能:通过对五个既定基准 (PRIMEVULDing 等人 (2024)、BigVulFan 等人 (2020)德雷柏Russell 等人 (2018)揭示Chakraborty 等人 (2021)和 VulDeepeckerLi et al. (2018)),我们证明 White-Basilisk 的性能优于 35× 大的模型,尤其是在代表真实代码安全场景的现实、不平衡的数据集上。
- 环境和计算效率:我们对 White-Basilisk 的碳足迹和计算要求进行了详细分析,证明与大型模型相比,排放量减少了 99.9%,同时保持了卓越的性能,为安全、可持续的 AI 系统建立了新的效率前沿。
除了软件安全领域的直接应用外,我们的工作还对 AI 领域提出了更广泛的影响。通过证明精心设计的特定于领域的架构可以胜过通用的缩放模型,我们挑战了参数数量是模型能力的主要决定因素的概念。White-Basilisk 证明,通往更强大的 AI 系统的道路可能不仅在于扩展现有架构,还在于重新思考特定高价值应用程序的基本设计选择。在以下部分中,我们将详细介绍 White-Basilisk 的架构、训练方法以及跨多个漏洞检测基准的综合评估。
2相关工作
自动代码漏洞检测领域为探索高效的 AI 架构提供了极好的背景。虽然这个领域经历了快速发展,但大多数方法都遵循了 AI 的一般趋势,即计算要求越来越高的更大模型。我们的工作通过证明精心设计的较小模型可以取得卓越的结果来挑战这种范式。为了正确定位我们的贡献,我们回顾了传统的漏洞检测方法和高效语言建模的最新趋势。
在软件系统日益复杂和安全风险的推动下,早期的先驱者探索了静态分析技术Chess 和 McGraw (2004)和模式匹配方法Livshits 和 Lam (2005).虽然这些方法奠定了基础框架,但它们经常遇到重大缺点,包括高假阳性率、难以检测复杂漏洞以及容易受到混淆技术的影响。由于这些限制,它们通常不适合主动生产环境。
随着该领域的成熟,研究人员认识到 AI 的变革潜力,逐渐从传统做法转变为新范式。这种转变以 VulDeePecker 的发展为标志Li et al. (2018),这是首批基于深度学习 (DL) 的漏洞检测系统之一。VulDeePecker 利用代码小工具和双向长短期记忆 (BiLSTM) 网络来识别 C/C++ 代码中的漏洞。这项工作证明了 DL 技术能够捕获与代码漏洞相关的复杂模式,为进一步开发 ML 驱动的解决方案铺平了道路。然而,它对手动制作功能的依赖限制了它的通用性。在这项工作的基础上,Russell 等人 (2018)开发了 Draper 数据集,该数据集提供了一个重要的真实数据集,专门用于在漏洞检测任务中训练神经网络。 他们的工作显示了利用大量训练数据来提高模型性能、提高性能的优势,但仍然在努力解决限制代码中长期依赖关系捕获的有限上下文窗口的优势。
此后,预训练的 LLM 成为代码分析领域的重大突破。Hanif 和 Maffeis (2022)提出了 VulBERTa,这是 RoBERTa 模型的改编版,用于检测代码漏洞,展示了从自然语言处理到代码分析的迁移学习的潜力。通过在广泛的代码语料库上预先训练的 LLM 进行微调,这种方法很快就流行起来,因为它能够捕获代码结构和语义中的潜在模式。然而,与许多基于 transformer 的模型类似,它存在序列长度的二次计算复杂性,限制了其对大型项目的适用性。
另一条研究线侧重于专门为代码理解量身定制的模型。Feng et al. (2020)引入了 CodeBERT,这是一种用于编程语言和自然语言的双峰预训练模型。这项工作允许更深入地了解代码及其相关文档,在包括漏洞检测在内的一系列软件工程任务中被证明是有用的。然而,虽然功能强大,但这些模型通常需要大量的计算资源,并且在处理极长的序列时遇到困难,这限制了它们在分析大型代码库方面的实际应用。 例如,CodeBERT 在 16 个互连的 NVIDIA Tesla V100 GPU 上进行了大约 1000 小时的预训练,这意味着大量的能源和资源投资。
最近的其他工作集中在提高漏洞检测的粒度和效率上。例如,LineVul (Fu 和 Tantithamthavorn (2022)),这是一种基于 transformer 的行级漏洞预测方法,支持在代码库中精确定位漏洞。虽然这种方法对于精确定位很有价值,但可能会忽略跨多个行或函数的漏洞。
综合研究Chakraborty 等人 (2021)和Ding 等人 (2024)强调了该领域持续存在的挑战,包括数据质量问题、不切实际的评估方法以及处理长期依赖关系的困难。Ding 等人 (2024)表明,现有基准测试明显高估了模型性能,最先进的模型在有缺陷的数据集上获得了高分,但在更真实的数据集上却失败了。我们的工作通过改进的数据收集、现实的评估指标和专为长期理解而设计的模型架构直接解决了这些问题。
目前,现有的漏洞检测方法主要关注其主要目标,而没有考虑模型参数规模的增加。模型大小的这种不受控制的增长导致了需要数十个尖端 GPU 的高能耗模型的开发。最近,在高效语言模型设计方面取得了研究进展。例如,以 Mamba 的工作为基础顾和道 (2023),Lieber 等人 (2024)引入了 Jamba,它采用专家混合 (MoE) 策略,将 Mamba 层和注意力机制以交错模式组合起来Fedus 等人 (2022).另外Wang et al. (2020)提出了 Linformer,它可以使用低秩近似将自注意力复杂性降低到线性时间,而 Transformer-XL 架构Dai et al. (2019)已被证明可以通过将递归引入自我注意来有效地模拟长期依赖性。
3模型架构
White-Basilisk 的架构旨在解决代码漏洞检测中的三个关键挑战。首先,它解决了长期依赖性问题,因为代码漏洞通常跨越多个函数甚至文件,因此需要一个能够理解广泛上下文的模型。其次,该架构在本地和全球信息处理之间取得了平衡。这种双重关注既可以精细理解代码语法,也可以广泛理解整个程序结构,这两者都对于有效的漏洞检测至关重要。最后,该设计优先考虑计算效率,旨在创建一个强大的模型,该模型可以部署在实际环境中,而无需大量计算资源。通过解决这些相互关联的目标,我们的模型为识别和分析代码中的潜在安全漏洞提供了强大的解决方案。为此,我们开发了一种新颖的混合架构,它结合了三个主要组件:
- 曼巴层:这些构成了我们模型的支柱,有效地捕获本地依赖项并提供基于输入内容的自适应计算。
- 白蛇怪自我注意:我们对线性注意力机制的新颖改编能够高效处理极长的序列,从而实现全代码库分析。
- 混合专家 (MoE):这增加了整个网络的动态适应性,MoE 层为每个令牌配置了 8 个 Expert,2 个 Experts。
这些组件之间的协同作用使 White-Basilisk 能够在推理过程中处理多达 128,000 个令牌的序列,这一能力使其在代码分析领域独树一帜。这个广泛的上下文窗口支持对整个代码库进行整体分析,从而可能发现跨多个函数或文件的漏洞。
Mamba 层构成了 White-Basilisk 的主干,之所以选择它,是因为它们在代码序列中捕获本地依赖项方面的卓越效率。与传统的递归神经网络 (RNN) 或注意力机制不同,Mamba 的选择性状态空间机制允许对序列长度进行线性时间计算。在 White-Basilisk 中,Mamba 层使用状态大小 (d州) 为 16,则卷积核大小 (d卷) 为 4,膨胀因子为 2。
White-Basilisk 中的专家混合 (MoE) 层在保持计算效率的同时引入了动态适应性。这些层配置为每个令牌 8 个 Expert 和 2 个 Experts,允许模型仅为每个输入令牌激活一个参数子集,与具有类似容量的全密集模型相比,显著降低了计算成本。这种稀疏性有助于加快推理时间,这是实际部署场景中的一个关键因素。
受 Jamba 模型的启发的层组合管道Lieber 等人 (2024),允许使用复杂的交错模式。具体来说,层组合由通过实验确定的两个主要配置参数定义:注意力层偏移量 (2) 和注意力层周期 (8)。我们架构中层的 Junction 由以下公式定义:

通过模型的正向传递可以表示为h我哪里我∈{0,1,…,L−1}是图层索引,而L是体系结构中的层总数。这里h我是我-th 层和h0是初始输入嵌入。模型的最终输出y通过将图层标准化应用于最后一个隐藏状态来获得hL:

残差连接有助于训练期间的梯度流,并允许跨层保留信息。
3.1White-Basilisk Self-Attention:我们新颖的线性缩放机制
White-Basilisk 的一个关键创新是我们新颖的自我注意机制,该机制在序列长度方面实现了线性复杂性,使模型能够有效地处理高达 128K 个令牌的整个代码库。传统的注意力机制由于其二次复杂度而面临基本的扩展限制,这使得它们在计算上无法分析大型代码库。我们的方法在保持注意力质量的同时改变了这一限制。
从有效注意力的最新进展中汲取灵感Munkhdalai 等人 (2024),我们开发了一个经过大幅修改的实现,优先考虑整体序列理解而不是分段处理。我们方法的主要区别是:
- 1.
全局上下文累积:与独立处理块的分段方法不同,White-Basilisk Self-Attention 在整个序列中积累上下文信息:

哪里S是 Segment 的总数。这种积累导致内存随着序列长度的线性增长,同时使注意力能够在整个上下文中进行关注。
- 2.
自适应全局门控:我们的方法实现了一种新的全局门控机制,该机制在整个累积的上下文中运行:

这允许在整个序列中对基于记忆的注意力成分和直接注意力成分进行自适应平衡,这与单独对每个片段进行门控的传统方法不同。
我们的注意力机制中基于内存的组件通过以下方式自适应地压缩和检索上下文信息:

哪里M是压缩内存,z是归一化项,而L是段长度。
我们实现的一个关键方面是使用可学习参数来优化内存积累和注意力分配。内存矩阵M包含可学习的投影矩阵,这些矩阵经过训练,可以保留 Mamba 层的本地处理可能遗漏的信息。通过与周围的 Mamba 层共享残差连接(如前向传递方程所示),我们的注意力机制有效地学会了以全局上下文感知来补充 Mamba 的本地处理优势。
这种集成与原始 Infini-attention 方法的不同之处尤其显著,后者独立处理区段。我们的适应通过内存状态的积累在整个序列中保持连续的信息流,允许模型在相距遥远的代码元素之间建立连接。此功能对于检测跨多个函数的漏洞至关重要,例如,在代码库中识别一个函数中的缓冲区分配及其在另一个函数中的潜在溢出。
4实验装置
源代码漏洞检测模型的开发和评估需要大量带注释的数据样本。在本节中,我们概述了为此目的选择的数据集,并解释了它们如何用于模型训练和测试目的。此外,我们还详细概述了用于模型的训练方法。
4.1预训练
传统的 LLM 训练方法通常旨在使模型能够理解语言及其语法。这通常是通过因果语言建模 (CLM) 来实现的,其中模型根据其输入学习预测下一个标记。另一种常见的方法基于中间填充 (FIM) 技术,其中随机文本部分被屏蔽,并且模型必须重建缺失的内容。一些高级源代码 LLM 将这两种方法结合起来,以提高模型的灵活性(Li 等人,2023).同样,在我们的工作中,我们在对选定的 2M 代码样本进行模型预训练时采用了这两种技术。此过程大约需要 600 小时才能在单个 NVIDIA A100 40GB GPU 上完成。
对于预训练数据,我们最初选择了 StarCoder 数据集的一个精心策划的子集(Li 等人,2023),其中包括 80 多种编程语言,总共包含 305M 个文件。在我们的研究中,我们专注于 C 和 C++ 代码样本,在预训练阶段使用了 2M 代码样本。这种集中的方法使我们能够实现出色的一般性能,同时保持相对较小的预训练数据集,挑战了高性能模型始终需要大量预训练数据集的概念。
4.2微调和评估
为了评估预训练模型,我们需要具有公开可用分区的成熟基准测试数据集,以便进行微调和测试。这确保了与现有方法的公平比较,而无需重新创建原始模型。为此,我们选择了五个公开可用的数据集:VulDeePecker(Li 等人,2018)德雷柏(Russell 等人,2018)、PrimeVul(Ding 等人,2024)揭示(Chakraborty 等人,2021)和 BigVul(Fan 等人,2020).
这些数据集带来了验证 White-Basilisk 能力的重大挑战。它们具有显著的类别不平衡,易受攻击的类别仅占样本的 2.9-9.8%,并且漏洞类型多样,涵盖多个常见弱点枚举 (CWE) 类别。为了确保与基线模型的公平比较,我们为每个数据集使用了与原始论文中报告的相同的数据拆分。
为了进行微调,我们采用了特定于数据集的方法来解决漏洞检测中存在的大量类不平衡。实施了类加权,以确保在训练期间平衡表示易受攻击和安全的代码示例。特定于数据集的超参数优化包括 1e-5 和 5e-6 之间的学习率,以及较小的批量大小 (4) 用于微调以防止过拟合。
鉴于漏洞检测中的类别不平衡,我们优先考虑提供模型性能平衡评估的指标。F1 Score 是我们的主要指标,用于平衡精度和召回率。我们还使用漏洞检测评分 (VD-S) 进行评估,由(Ding 等人,2024),用于测量 False Negative Rate。报告准确度、精密度和召回率是为了完整性并与基线模型进行比较。
5实验结果:小模型,大影响
为了评估 White-Basilisk 的有效性,我们在五个广泛使用的漏洞检测数据集中进行了广泛的实验。表 1 和表 2 显示了我们与最先进模型相比的结果。White-Basilisk 在所有基准测试中始终优于更大、资源更密集的模型。
在具有挑战性的 PRIMEVUL 数据集上,White-Basilisk 的 F1 分数为 0.2907,VD-S 为 0.7239,明显优于 StarCoder2 等大型模型(7B 参数)。最令人印象深刻的是,在 BigVul 上,我们的模型实现了 0.9490 的 F1 分数、0.9942 的准确率和 0.0398 的 VD-S,在所有指标上都大大超过了所有竞争对手。在 Draper 数据集上,White-Basilisk 以 0.6069 的 F1 分数设定了新标准,而在 REVEAL 上,它的 F1 得分为 0.4934,准确率为 0.8988,超过了之前的最佳结果。对于 VulDeepecker,我们的模型实现了 0.9388 的 F1 分数,最高精度为 0.9720。
表 1:BigVul 和 PRIMEVUL 评估结果
| 型 | PRIMEVUL 公司 | BigVul 浏览器 | ||||
|---|---|---|---|---|---|---|
| 累积 | F1 系列 | VD-S 系列↓ | 累积 | F1 系列 | VD-S 系列↓ | |
| 白蛇怪 | 0.9630 | 0.2907 | 0.7239 | 0.9942 | 0.9490 | 0.0398 |
| 代码T5 | 0.9667 | 0.1970 | 0.8993 | 0.9567 | 0.6493 | 0.7730 |
| 代码 BERT | 0.9687 | 0.2086 | 0.8878 | 0.9557 | 0.6288 | 0.8177 |
| UnixCoder | 0.9686 | 0.2143 | 0.8921 | 0.9646 | 0.6546 | 0.6230 |
| 具有平衡功能的 UnixCoder | 0.9599 | 0.2628 | 0.8849 | - | - | - |
| 星码2 | 0.9702 | 0.1805 | 0.8964 | 0.9620 | 0.6826 | 0.6914 |
| 代码生成2.5 | 0.9665 | 0.1961 | 0.9151 | 0.9657 | 0.6730 | 0.6173 |
| 线Vul | - | - | - | - | 0.9100 | 0.1400 |
表 2:Draper、REVEAL 和 VulDeepecker 评估结果
| 型 | 德雷柏 | 揭示 | VulDeepecker | ||
|---|---|---|---|---|---|
| F1 系列 | 累积 | F1 系列 | F1 系列 | Prec (优惠) | |
| 白蛇怪 | 0.6069 | 0.8988 | 0.4934 | 0.9388 | 0.9720 |
| Russell 等人 (2018) | 0.5660 | - | - | - | - |
| VulBERTa-MLP 技术 | 0.4334 | 0.8448 | 0.4527 | 0.9303 | 0.9576 |
| VulBERTa-CNN | 0.5792 | 0.7973 | 0.4259 | 0.9086 | 0.9526 |
| 基线 BiLSTM | 0.4684 | 0.7713 | 0.3911 | 0.6697 | 0.5258 |
| 基线文本CNN | 0.4940 | 0.7322 | 0.3741 | 0.7580 | 0.6348 |
| 揭示 | - | 0.8437 | 0.4125 | - | - |
| VulDeepecker | - | - | - | 0.9290 | 0.9190 |
5.1消融研究:验证关键索赔
5.1.1统一模型分析:一个模型可以统治所有模型吗?
我们在来自所有五个数据集的组合训练数据上训练了一个统一的 White-Basilisk 模型,以评估跨数据集泛化。该模型在各种漏洞类型中保持了稳健的性能,BigVul 的 F1 分数为 0.938,VulDeepeker 为 0.925,Draper 为 0.549,REVEAL 为 0.470,PRIMEVUL 为 0.233。值得注意的是,该模型在较大的数据集上保持了近乎最佳的性能,同时显示出向更具挑战性的数据集的合理转移,这表明 White-Basilisk 尽管规模紧凑,但仍能够学习可推广的漏洞模式。该模型对 BigVul 中的 22 个不同 CWE 类别实现了完美检测 (F1=1.000),包括越界写入 (CWE-787) 和加密漏洞 (CWE-310) 等关键问题。
表 3:跨数据集的关键 CWE 的性能
| CWE的 | 描述 | 数据 | 样品 | F1 系列 |
|---|---|---|---|---|
| CWE-119 | 缓冲区溢出 | BigVul、VulDeep、Draper | 15,584 | 0.748-0.969 |
| CWE-399 | 资源管理 | BigVul 和 VulDeep | 7,031 | 0.780-0.923 |
| CWE-264 | 存取控制 | BigVul 浏览器 | 1,240 | 0.925 |
| CWE-476 | NULL 指针取消引用。 | BigVul、德雷珀 | 1,709 | 0.686-0.971 |
| CWE-416 | 释放后使用 | BigVul 浏览器 | 963 | 0.958 |
5.1.2注意力效率分析:线性与二次缩放
White-Basilisk 的线性复杂性注意是否真的能够在不降低性能的情况下实现高效的长序列处理?我们将我们的方法与标准二次注意力进行了比较,测量了不同序列长度的内存使用情况和预测性能。如表 4 所示,我们的方法实现了显著的内存节省——对于 16K 标记以下的序列,使用的内存减少了 24× (1,338MB vs. 32,409MB),同时保持相当或更高的 F1 分数 (0.943 vs. 0.937)。至关重要的是,虽然标准注意力因超过 16K 个令牌的内存不足错误而失败,但 White-Basilisk 成功地处理了高达 131K 令牌的序列,内存增长不大,展示了真正的线性扩展。
最令人印象深刻的是,性能分析表明,即使对于极长的序列,White-Basilisk 也能保持强大的检测能力,16K 令牌以下的序列的 F1 分数为 0.943,16K-32K 令牌的 F1 分数为 0.800,32K-65K 令牌的 F1 分数为 1.000。这证实了我们的架构能够捕获对于识别跨多个函数或文件的漏洞至关重要的长期依赖关系。有关所有数据集的详细序列长度分析,请参阅附录 I。
表 4:按序列长度和注意力类型划分的内存使用情况和性能
| 长度 | 峰值内存 (MB) | BigVul F1 分数 | ||
|---|---|---|---|---|
| 我们 | 标准 | 我们 | 标准 | |
| 0-16K | 1,338 | 32,409 | 0.943 | 0.937 |
| 16K 至 32K | 1,654 | OOM | 0.800 | - |
| 32K-65K | 2,213 | OOM | 1.000 | - |
| 65K-131K | 3,322 | OOM | - | - |
| OOM:NVIDIA A100 40GB GPU 内存不足 | ||||
5.2效率和环境影响
考虑到 White-Basilisk 的计算效率,其卓越的性能变得更加重要。在大约 200M 的参数下,我们的模型比 StarCoder2 (7B) 等模型小几个数量级,同时获得了出色的结果。这种效率转化为巨大的环境效益,如表 5 所示。
表 5:CO2 排放量比较
白蛇怪代码 BERTStarCoder 公司UnixCoder代码T5汽油车(每年)二氧化碳 (kg)85.52,24023,000,0002,0481,1364,600
我们使用机器学习影响计算器估算了训练期间的 CO2 排放量(Lacoste 等人,2019).白蛇怪对环境的影响(85.5 千克二氧化碳)仅占 CodeBERT 排放量(2,240 千克)的 3.8%,仅占 StarCoder 排放量(23,000,000 千克)的 0.0004%。这种鲜明的对比凸显了我们方法的环境效益,与更大规模的模型训练相比,减少的二氧化碳排放相当于让数百辆汽车上路。
效率提升也延伸到推理。White-Basilisk 可以在单个 NVIDIA A100 GPU 上处理多达 128K 个令牌的序列,无需大型模型所需的分布式计算基础设施即可实现全代码库分析。这种实用的效率使更广泛的组织和研究人员能够使用高级漏洞检测。
我们的结果挑战了 AI 模型开发中流行的“越大越好”范式。White-Basilisk 证明,精心设计的架构,而不是简单地增加参数数量,可以产生卓越的性能,同时大大降低计算和环境成本。这一发现对可持续的 AI 发展具有重大意义,表明前进的道路可能在于架构创新,而不是越来越大的模型。
6限制和未来工作
虽然 White-Basilisk 在代码漏洞检测方面显示出有希望的结果,但重要的是要承认其当前的局限性并概述未来的方向。White-Basilisk 的主要限制是它专注于 C 和 C++ 代码库。该模型能够泛化到更广泛的编程语言,尤其是那些具有不同语法或范式的语言,值得进一步探索。这种限制,再加上我们的训练和评估数据集中的潜在偏差,可能会限制模型对各种实际代码库的泛化性。为了解决这个问题,未来的工作将涉及扩展模型的训练以包括更广泛的编程语言,并策划更具代表性的数据集,以反映更广泛的代码样本和漏洞类型。
此外,尽管性能指标强劲,但 White-Basilisk 并非万无一失。误报和漏报,尤其是在检测新漏洞或零日漏洞时,仍然是一个持续的挑战。此外,虽然该模型能够处理长序列,但它对代码中复杂、长距离依赖关系的真正理解仍需要进一步研究。我们未来的研究将侧重于降低错误率,尤其是在高风险场景中,并增强模型分析跨多个函数或文件的复杂代码结构的能力。
另一个需要改进的领域是模型的可解释性。目前,White-Basilisk 的决策过程不够透明。提高模型为检测到的漏洞提供清晰、可作的解释的能力,对于建立信任和向开发人员提供有意义的见解至关重要。未来的工作将探索提供对漏洞潜在影响的上下文感知理解的方法,以及建议的修复方法,最终旨在将 White-Basilisk 发展成为一个全面的代码分析助手,而不仅仅是一个检测工具。
White-Basilisk 在相对较小的训练数据集(2M 样本)上的表现令人印象深刻,但与在更大的数据集上训练的模型相比,其知识库的潜在局限性引发了疑问。为了解决这个问题,我们计划将 White-Basilisk 扩展到大约 10 亿个参数。虽然与较大的语言模型相比仍然适中,但参数数量的增加旨在显著提高性能,同时继续挑战只有最大的模型才能提供最先进的结果的概念。
除了代码漏洞检测之外,White-Basilisk 的架构和训练方法还有可能在更广泛的 AI 应用中发挥作用。未来的研究将调查其在各种 NLP 任务中的功效,探索其计算效率和长上下文能力是否可以为现有 LLM 提供更节省资源的替代方案。
7结论
White-Basilisk 代表了高效 AI 模型设计的重大进步,表明精心构建的紧凑模型可以在专业领域中胜过更大的替代方案。通过仅使用 200M 个参数就实现最先进的代码漏洞检测结果,我们的工作挑战了 AI 开发中流行的“越大越好”范式。该模型令人印象深刻的性能(尤其是它处理多达 128,000 个令牌的序列的能力)证明,上下文处理限制可以通过架构创新来解决,而不仅仅是增加模型大小。
该模型的扩展上下文窗口解决了许多当前 LLM 的基本限制,这些 LLM 通常难以理解长距离依赖关系和全局代码结构。通过一次性分析整个代码库,White-Basilisk 可以捕获复杂的相互依赖关系并识别跨多个函数或文件的漏洞,这是传统方法长期以来一直无法实现的功能。
虽然 White-Basilisk 的上下文处理功能是其突出特点,但值得注意的是,这些成就是通过相对紧凑的 200M 参数模型实现的。这种效率表明,AI 的进步不仅取决于模型大小的增加,还可以源于创新的架构设计和训练方法。
White-Basilisk 方法的含义不仅限于代码漏洞检测。有效处理扩展上下文的能力在远程理解至关重要的众多领域中可能被证明是有价值的,例如文档分析、复杂系统建模或长格式文本生成。此外,该模型的效率为在资源受限的环境中部署提供了可能性,有可能为更广泛的应用程序和用户带来先进的 AI 功能。
总之,White-Basilisk 代表了在解决当前 LLM 的上下文限制方面向前迈出的重要一步,同时也证明了这种进步不必以过大的模型大小或计算要求为代价。随着我们不断完善和扩展这种方法,我们预计 AI 领域将出现令人兴奋的发展,尤其是在需要深入了解扩展上下文的任务中。这项研究的潜在影响是巨大的,我们期待看到这些想法如何发展并在 AI 及其他不同领域得到应用。
831

被折叠的 条评论
为什么被折叠?



