[論文レビュー] Finetuning Large Language Models for Vulnerability Detection

Alexey Shestov, Rodion Levichev+4 more|arXiv (Cornell University)|Jan 30, 2024

Web Application Security Vulnerabilities被引用数 5

ひとこと要約

本論文は WizardCoder（StarCoder 系の LLM）を Java の関数を脆弱かどうかの二値分類に微調整し、訓練 regime を探究し、CodeBERT 系モデルと比較して脆弱性検出性能が向上することを示す。データの不均衡への対処と、 batch packing と LoRA による高速な訓練も含む。

ABSTRACT

This paper presents the results of finetuning large language models (LLMs) for the task of detecting vulnerabilities in source code. We leverage WizardCoder, a recent improvement of the state-of-the-art LLM StarCoder, and adapt it for vulnerability detection through further finetuning. To accelerate training, we modify WizardCoder's training procedure, also we investigate optimal training regimes. For the imbalanced dataset with many more negative examples than positive, we also explore different techniques to improve classification performance. The finetuned WizardCoder model achieves improvement in ROC AUC and F1 measures on balanced and imbalanced vulnerability datasets over CodeBERT-like model, demonstrating the effectiveness of adapting pretrained LLMs for vulnerability detection in source code. The key contributions are finetuning the state-of-the-art code LLM, WizardCoder, increasing its training speed without the performance harm, optimizing the training procedure and regimes, handling class imbalance, and improving performance on difficult vulnerability detection datasets. This demonstrates the potential for transfer learning by finetuning large pretrained language models for specialized source code analysis tasks.

研究の動機と目的

• コードで事前学習された大規模言語モデルを、ソースコードの脆弱性検出に効果的に微調整できるかを調査する。
• StarCoder 系モデルが、バランスの取れた脆弱性データセット上で CodeBERT 系モデルより性能を上回るかを判断する。
• 分類に焦点を当てた損失、コンテキストサイズ、バッチ packing などを含む訓練 regime の最適化を探索し、効率と性能を向上させる。
• 焦点を当てた不均衡データセットへの対処法（ focal loss とサンプルウェイト付けを含む）を評価し、キャリブレーション効果を検討する。
• 計算資源が限られた状況で、分類タスクに適用するための大規模コード LLM の適用に関する実践的な指針を提供する。

提案手法

• 基盤モデルとして WizardCoder（13B パラメータ）を選択し、脆弱性検出タスクで ContraBERT と比較する。
• メモリ効率の良い微調整を実現するため、r=8、alpha=32、dropout=0.05 の低ランク適応（LoRA）を適用する。
• 最終トークンの確率を用いて脆弱性を予測する専用の分類損失に、次トークン予測損失を置き換える。
• 複数の短いコード関数を単一の入力シーケンスに詰め合わせる（バッチ packing）ことで訓練スループットを向上させる。
• 次トークン予測と二値分類目的の2つの訓練 regime を評価し、ROC AUC と F1 への影響を検討する。
• 性能向上の原因を特定するため、コンテキストサイズを（512 と 2048 トークン）で実験する。

実験結果

リサーチクエスチョン

• RQ1 RQ1: バランスの取れた脆弱性検出において、StarCoder 系モデルは CodeBERT 系モデルより有効か？
• RQ2 RQ2: 不均衡データでの脆弱性検出において、StarCoder 系モデルは CodeBERT 系モデルより有効か？
• RQ3 RQ3: 標準的な LLM の学習法（次トークン予測＋QA）は、脆弱性検出に有効か？
• RQ4 RQ4: 動的な batch packing はモデル品質に影響するか？
• RQ5 RQ5: パックされたバッチにおける平均損失還元 versus 合計損失還元のどちらが良いか？
• RQ6 RQ6: より大きなコンテキストサイズは脆弱性検出の性能を向上させるか？
• RQ7 RQ7: サンプルウェイト付き focal loss はクラス不均衡な脆弱性検出に有効か？

主な発見

• WizardCoder はバランスの取れた脆弱性検出で ContraBERT を上回る（ROC AUC 0.69 対 0.68; F1 0.71 対 0.68）。
• P3 陰性を含む不均衡データで、WizardCoder は ROC AUC 0.86、F1 0.27 を達成し、ContraBERT をわずかに上回る。
• 脆弱性検出を専用の分類損失を用いた二値分類タスクとして扱うと、次トークン予測アプローチより ROC AUC が良くなる（0.86 対 0.75）。
• バッチ packing により訓練が最大約13倍速くなり、平均損失ではなく合計損失還元を用いる場合に性能低下は significant ではない。
• コンテキストサイズ（2048 対 512）は性能向上を生まない。改善はモデルが学習したコード表現の向上に起因する。
• gamma=1 の focal loss は標準クロスエントロピーに比べて小さな ROC AUC 増加をもたらし、サンプルウェイトは較正とテスト指標をさらに改善する。