[論文レビュー] VulBERTa: Simplified Source Code Pre-Training for Vulnerability Detection

Hazim Hanif, Sergio Maffeis|arXiv (Cornell University)|May 25, 2022

Software Engineering Research被引用数 9

ひとこと要約

VulBERTaは、C/C++コードにおける脆弱性検出のための簡素化された、小規模な事前学習アプローチを提案する。独自のトークン化パイプラインを用い、BPEと事前に定義されたコードトークン（キーワード、句読点、APIコール）を組み合わせることで、構文的・意味的構造を保持する。228万件の関数、1億2500万パラメータで事前学習されたモデルは、複数のデータセットおよびベンチマークで最先端の性能を達成し、SOTA F1スコアとしてDraperで57.92%、muVuldeepeckerで99.59%を記録。アーキテクチャの複雑さを最小限に抑えながら、より大きなモデルを上回る性能を発揮している。

ABSTRACT

This paper presents VulBERTa, a deep learning approach to detect security vulnerabilities in source code. Our approach pre-trains a RoBERTa model with a custom tokenisation pipeline on real-world code from open-source C/C++ projects. The model learns a deep knowledge representation of the code syntax and semantics, which we leverage to train vulnerability detection classifiers. We evaluate our approach on binary and multi-class vulnerability detection tasks across several datasets (Vuldeepecker, Draper, REVEAL and muVuldeepecker) and benchmarks (CodeXGLUE and D2A). The evaluation results show that VulBERTa achieves state-of-the-art performance and outperforms existing approaches across different datasets, despite its conceptual simplicity, and limited cost in terms of size of training data and number of model parameters.

研究の動機と目的

• C/C++ソースコードにおける脆弱性検出のための簡素化され、低複雑性の事前学習アプローチの開発。
• 独自のトークン化パイプラインを通じて構文的・意味的情報を統合し、コード表現学習の向上。
• 従来の手法よりもパラメータ数が少なく、学習データ量も少ないにもかかわらず、最先端の脆弱性検出性能を達成すること。
• 多様なデータセットおよびベンチマークにおける事前学習モデルの転送性と頑健性の評価。
• 高精度を維持しながら、軽量な分類器（MLPおよびCNN）を用いた効率的なファインチューニングの実現。

提案手法

• 独自のトークン化パイプラインが、C/C++キーワード、句読点、一般的なAPI関数名の事前定義トークンとBPEを組み合わせることで、構文的・意味的構造を保持する。
• モデルはRoBERTaをバックボーンアーキテクチャとして採用し、オープンソースプロジェクトから抽出した228万件のC/C++関数を対象に、マスク言語モデル（MLM）による事前学習を実施。
• 事前学習は、依存関係やデータフロー解析を明示的に含まない、現実世界のC/C++コードから構成される洗練されたデータセットを用いて実施。
• ファインチューニングには、2つの軽量分類器（マルチレイヤーパーセプトロン：VulBERTa-MLP、畳み込みニューラルネットワーク：VulBERTa-CNN）を採用。
• 複雑なグラフベースやシーケンスベースのアーキテクチャを避ける一方で、トークン化と事前学習の質に依存して、深いコード表現をエンコードする。
• 評価は、Vuldeepecker、Draper、REVEAL、muVuldeepecker、およびベンチマークであるCodeXGLUEとD2Aを含む、バイナリおよびマルチクラスの脆弱性検出タスクで実施。

実験結果

リサーチクエスチョン

• RQ1 複雑なアーキテクチャに依存しない簡素化された、小規模な事前学習モデルが、脆弱性検出において最先端の性能を達成できるか？
• RQ2 構文的・意味的構造を保持する独自のトークン化パイプラインが、脆弱性検出のためのコード表現学習にどの程度寄与するか？
• RQ3 パラメータ数が少なく、学習データ量も少ないにもかかわらず、より大きなモデルを上回る性能を示せるか？
• RQ4 事前学習された表現が、多様なデータセットおよび脆弱性検出タスクにどの程度転送可能か？
• RQ5 強力な事前学習済みコード表現と組み合わせた、軽量なファインチューニングヘッド（MLPおよびCNN）が、高い性能を達成できるか？

主な発見

• VulBERTaは、挑戦的なDraperデータセットで57.92%のSOTA F1スコアを達成し、その小型化と限られた学習データにもかかわらず、既存のアプローチを上回った。
• マルチクラスのmuVuldeepeckerデータセットでは、VulBERTa-MLPが99.59%の重み付きF1スコアを記録し、より簡単でバランスの取れたデータに対して優れた性能を示した。
• CodeXGLUEベンチマークでは、VulBERTa-MLPが64.75%の精度で1位、VulBERTa-CNNが60.68%で2位にランクされ、C-BERTのようなより大きなモデルを上回った。
• D2Aベンチマークでは、VulBERTa-MLPが「関数」タスクで62.30%の精度を達成し1位、VulBERTa-CNNが60.68%で2位にランクされた。
• 1億2500万パラメータ、228万件の事前学習サンプルというわずかな規模にもかかわらず、C-BERT（810万件、1億1000万パラメータ）やCodeBERT（850万件、1億2500万パラメータ）といったモデルを上回った。
• 異なるファインチューニングヘッドにおいてもモデルの性能は頑健で、VulBERTa-MLPは4つのデータセットのうち3つで最高のPEMスコアを記録し、VulBERTa-CNNは残りの1つで最高スコアを達成した。