[论文解读] Transformer-based Vulnerability Detection in Code at EditTime: Zero-shot, Few-shot, or Fine-tuning?

Aaron Chan, Anant Kharkar+6 more|arXiv (Cornell University)|May 23, 2023

Software Engineering Research被引用 7

一句话总结

本文提出了一种基于微调及零样本/少样本提示的实时、编辑时漏洞检测系统，利用预训练大语言模型（CodeBERT、code-davinci-002、text-davinci-003）识别语法不完整的代码中的漏洞。该系统在最先进方法的基础上，将召回率提升10%，精确率提升8%，并将大语言模型生成代码中的漏洞率降低90%以上。

ABSTRACT

Software vulnerabilities bear enterprises significant costs. Despite extensive efforts in research and development of software vulnerability detection methods, uncaught vulnerabilities continue to put software owners and users at risk. Many current vulnerability detection methods require that code snippets can compile and build before attempting detection. This, unfortunately, introduces a long latency between the time a vulnerability is injected to the time it is removed, which can substantially increases the cost of fixing a vulnerability. We recognize that the current advances in machine learning can be used to detect vulnerable code patterns on syntactically incomplete code snippets as the developer is writing the code at EditTime. In this paper we present a practical system that leverages deep learning on a large-scale data set of vulnerable code patterns to learn complex manifestations of more than 250 vulnerability types and detect vulnerable code patterns at EditTime. We discuss zero-shot, few-shot, and fine-tuning approaches on state of the art pre-trained Large Language Models (LLMs). We show that in comparison with state of the art vulnerability detection models our approach improves the state of the art by 10%. We also evaluate our approach to detect vulnerability in auto-generated code by code LLMs. Evaluation on a benchmark of high-risk code scenarios shows a reduction of up to 90% vulnerability reduction.

研究动机与目标

• 解决在代码编写过程中实时漏洞检测的关键空白，实现在语法不完整代码片段上的检测。
• 通过在引入缺陷的瞬间（即编辑时）识别缺陷，而非编译后，降低漏洞修复的成本与延迟。
• 评估并比较在预训练大语言模型上使用零样本、少样本和微调方法在不完整代码漏洞检测中的表现。
• 评估模型在检测由代码大语言模型生成的代码中漏洞的有效性，此类代码正成为日益重要的攻击面。
• 在真实世界的IDE环境（VSCode）中部署并评估该系统，以衡量其在实际中对漏洞减少的影响。

提出的方法

• 使用静态分析从开源代码库中收集超过50万条存在漏洞的代码片段，用于模型的训练与评估。
• 评估三种学习范式：零样本（基于提示的直接推理）、少样本（上下文示例）以及在CodeBERT和OpenAI的code-davinci-002/text-davinci-003上的微调。
• 使用CodeBERT-base（参数量少于100M）以确保低延迟推理，满足实时IDE集成需求。
• 设计针对漏洞检测任务的提示词与少样本示范，充分利用InstructGPT类模型的指令遵循能力。
• 将性能最佳的模型（微调后的CodeBERT）作为VSCode扩展进行部署，实现实时开发反馈。
• 实现一个再训练流水线，基于新出现的数据持续扩展对新型漏洞类型的覆盖能力。

实验结果

研究问题

• RQ1 预训练大语言模型能否在编辑时检测语法不完整的代码中的漏洞？不同方法（零样本、少样本、微调）的性能表现如何？
• RQ2 与需要完整代码的最先进模型相比，编辑时漏洞检测的性能表现如何？
• RQ3 该模型在自动由代码大语言模型生成的代码中检测漏洞的能力有多强？
• RQ4 在生产部署中，模型规模、推理成本、精确率、召回率与维护开销之间存在何种权衡？
• RQ5 如何通过阈值调优与再训练流水线，长期维持开发者的信任与检测准确性？

主要发现

• 微调后的CodeBERT在性能平衡方面表现最佳，精确率为59%，召回率为63%，在预测性能上优于零样本与少样本方法。
• 在text-davinci-003上进行零样本与少样本推理的召回率分别达到78%与75%，但精确率较低（分别为47%与49%），表明灵敏度与误报率之间存在权衡。
• 所提系统在标准基准测试中，相较于最先进方法，召回率提升10%，精确率提升8%。
• 在高风险基准测试中，该模型将大语言模型生成代码中的漏洞率降低了90%以上，证明其对自动生成代码修改具有强大的泛化能力。
• 在VSCode扩展中部署后，实际编辑代码中的漏洞率降低了80%，验证了其在真实场景中的实际影响。
• 再训练流水线对维持覆盖范围与开发者信任至关重要，因为新型漏洞类型会随时间不断出现，必须主动集成。