DataCon 晚自习:浅谈大模型辅助漏洞挖掘
一、前言 大语言模型(LLM)近期成为安全研究和竞赛中的热门话题,涵盖了大模型在安全领域的应用以及自身的安全性等方向,这一趋势为许多非大模型安全研究领域的研究者提供了了解和探索的机会。 得益于强大的上下文理解与模式识别能力,大模型被认为具备从已知漏洞中学习特征并检测、修复未知漏洞的潜力。目前,大多数基于大模型的漏洞检测和修复方案集中在函数级别,只处理小范围的代码片段,不过也有少量研究开始探索代码库级别的问题。 本次 DataCon 聚焦于真实漏洞环境,提供了和漏洞相关的完整程序代码,要求参赛者"结合大模型技术自动化识别出漏洞样例中存在的安全隐患"。为此,本文将从函数级别和代码库级别两个方面简要分析和梳理当前研究进展,并结合代码库级别漏洞检测的需求,探讨未来的研究方向。 二、函数级别漏洞检测 本节主要参考 Large Language Model for Vulnerability Detection and Repair: Literature Review and the Road Ahead,未注明来源的研究方法,可在该综述中查阅 从提升大模型漏洞挖掘能力的角度,主要可以将当前研究分为三种方法:微调、提示工程和检索增强生成(RAG)。微调通过在漏洞数据上的定向训练来调整模型的参数;提示工程则以黑盒的方式,通过精心设计的 prompt 来优化模型在具体应用场景中的输出效果;RAG 从数据库中检索相关知识并整合到 LLM 的上下文中,同样不需要修改 LLM 的参数。 2.1 微调 一般的微调训练包含以下几个步骤:数据准备、模型设计、模型训练和模型优化。根据不同研究方法主要针对的训练阶段,可以划分为以下几个类别: 以数据为中心的创新(数据准备) 漏洞数据集通常存在标签不平衡和标签不正确的问题。一些研究通过数据采样、伪标签生成以及反事实数据生成等方法,尝试针对这些问题提出解决方案。 结合程序分析方法(数据准备) 一些研究将程序分析工具的结果引入模型的预训练或微调阶段,以帮助 LLM 更好地理解代码逻辑和数据依赖关系。例如: 可以提取抽象语法树(AST)和程序依赖图(PDG)来进行语句级控制依赖和数据依赖的预训练。 利用程序切片提取控制和数据依赖信息,将控制流图(CFG)分解为执行路径,为模型提供漏洞检测的支持。 结合其他深度学习模型(模型设计) 通过 Bi-LSTM (双向长短期记忆网络)处理分段后的代码输入克服了 LLM 的长度限制;亦或是通过 GNN (图神经网络)提取代码的图结构特征来补充 LLM 对代码结构的理解。 针对特定领域进行训练(模型训练) 通过在特定编程语言和特定漏洞类型的数据上进行预训练,使用如掩码语言模型、对比学习、程序依赖预测和漏洞语句标注等不同的预训练目标,可以增强模型对特定编程语言、特定漏洞模式以及复杂代码依赖关系的理解能力。 因果学习(训练优化) 通过识别和剔除那些与漏洞标签存在虚假相关的非稳健特征,进而使用因果推理算法来提升模型的因果推理能力。 其中,相比于其他方法更像是在解决一些大模型的通病,“结合程序分析方法"以及"针对特定领域进行训练"两类方法,则更依赖研究者对程序分析和漏洞挖掘领域的深入理解。限于时间原因笔者没能深入了解相关方法的研究细节,感兴趣的读者建议阅读相关论文。 2.2 微调:如何训练自己的模型 那如果我们想要微调自己的大模型,有哪些需要关注的问题呢?首先是选择使用的研究方法:根据数据规模、数据是否标记、模型选择和预算,可以划分为如图所示的模型训练方法。**CPT(持续预训练)**使用大量未标记的特定领域数据对已经预训练的大模型进行进一步训练;**SFT (有监督微调)**更专注于增强特定任务的性能,使用有标签的数据进行训练。**FULL(全参数微调)**细化所有参数,需要大量的计算能力和时间;相比之下,**PEFT(参数高效微调)**通过仅微调少量模型参数大大降低了计算成本,常见的 LoRA (Low Rank Approximation)就属于 PEFT 方法。在针对于漏洞检测的领域,许多研究采用了 SFT 和 PEFT。 此外,在一个不考虑特殊优化方法的微调过程中,还要涉及模型选择和数据集选择:...