[論文レビュー] 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）を微調整およびゼロショット／Few-shotプロンプティングを用いて、構文的に不完全なコードにおいてもリアルタイムで脆弱性を検出する編集時脆弱性検出システムを提案する。本手法は、最先端の手法と比較して再現率を10%向上、正確率を8%向上させ、LLMによって生成されたコードの脆弱性率を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.

研究の動機と目的

• コード作成中にリアルタイムで脆弱性を検出するという重要なギャップを埋め、構文的に不完全なコードスニペットに対しても検出を可能にする。
• コンパイル後ではなく、導入直後に欠陥を特定することにより、脆弱性是正のコストとレイテンシを低減する。
• 不完全なコードにおける脆弱性検出のため、事前学習されたLLMにおけるゼロショット、Few-shot、微調整のアプローチを評価・比較する。
• コードLLMが自動生成したコードにおける脆弱性検出の有効性を評価する。これは、急速に拡大する攻撃面である。
• VSCodeを含む実世界のIDE環境にシステムをデプロイし、脆弱性削減に与える実際の影響を測定する。

提案手法

• 静的解析を用いて、オープンソースリポジトリから50万件を超える脆弱なコードスニペットを収集し、モデルの学習および評価に使用した。
• 3つの学習パラダイムを評価した：ゼロショット（直接プロンプトベースの推論）、Few-shot（文脈内例を用いた）、CodeBERTおよびOpenAIのcode-davinci-002/text-davinci-003における微調整。
• リアルタイムIDE統合を実現するため、低レイテンシ推論を実現するため、CodeBERT-base（1億パラメータ未満）を採用した。
• 脆弱性検出タスクに特化したプロンプトおよびFew-shotデモンストレーションを設計し、InstructGPTベースのモデルの指示従い能力を活用した。
• 最も性能の優れたモデル（微調整済みCodeBERT）をVSCode拡張機能としてデプロイし、開発中のリアルタイムフィードバックを可能にした。
• 継続的に新たな脆弱性タイプのカバレッジを拡大するため、再学習パイプラインを実装した。

実験結果

リサーチクエスチョン

• RQ1 事前学習されたLLMは、構文的に不完全なコードにおいて編集時（EditTime）に脆弱性を検出可能か？ また、ゼロショット、Few-shot、微調整の各アプローチにおける性能はどのように異なるか？
• RQ2 完全なコードを必要とする最先端のモデルと比較して、編集時脆弱性検出の性能はどの程度向上するか？
• RQ3 本モデルは、コードLLMによって自動生成されたコードにおける脆弱性をどの程度検出できるか？
• RQ4 本番環境でのデプロイにおいて、モデルサイズ、推論コスト、正確率、再現率、保守のオーバーヘッドの間にはどのようなトレードオフがあるか？
• RQ5 しきい値チューニングと再学習パイプラインは、長期にわたり開発者の信頼と検出精度を維持するためにどのように機能するか？

主な発見

• 微調整済みCodeBERTが、59%の正確率と63%の再現率を達成し、予測性能においてゼロショットおよびFew-shot手法を上回る、最もバランスの取れた性能を示した。
• text-davinci-003におけるゼロショットおよびFew-shot推論は、再現率が78%および75%と高く、正確率が47%および49%と低く、感度と誤検出のトレードオフが顕著に現れた。
• 提案手法は、標準ベンチマークにおいて、最先端の検出手法と比較して再現率10%、正確率8%の向上を達成した。
• 高リスクベンチマークにおいて、LLM生成コードの脆弱性率を90%以上削減した。これは、自動生成された編集に対する強力な一般化能力を示している。
• VSCode拡張機能としてのデプロイにより、実際に編集されているコードにおける脆弱性率が80%削減された。これは、実世界への影響を裏付ける結果であった。
• 再学習パイプラインは、新たな脆弱性タイプが時間とともに出現するため、カバレッジと信頼性の維持に不可欠であった。