[論文レビュー] LineVD: Statement-level Vulnerability Detection using Graph Neural Networks

David Hin, Andrey Kan+2 more|arXiv (Cornell University)|Mar 10, 2022

Software Engineering Research被引用数 32

ひとこと要約

LineVD は statement-level vulnerability detection をプログラムグラフのノード分類として扱い、コード埋め込みに CodeBERT を使用し、データ/制御依存性を活用する GNN を用いて、従来の解釈ベースのモデルより大幅な向上を達成する。

ABSTRACT

Current machine-learning based software vulnerability detection methods are primarily conducted at the function-level. However, a key limitation of these methods is that they do not indicate the specific lines of code contributing to vulnerabilities. This limits the ability of developers to efficiently inspect and interpret the predictions from a learnt model, which is crucial for integrating machine-learning based tools into the software development workflow. Graph-based models have shown promising performance in function-level vulnerability detection, but their capability for statement-level vulnerability detection has not been extensively explored. While interpreting function-level predictions through explainable AI is one promising direction, we herein consider the statement-level software vulnerability detection task from a fully supervised learning perspective. We propose a novel deep learning framework, LineVD, which formulates statement-level vulnerability detection as a node classification task. LineVD leverages control and data dependencies between statements using graph neural networks, and a transformer-based model to encode the raw source code tokens. In particular, by addressing the conflicting outputs between function-level and statement-level information, LineVD significantly improve the prediction performance without vulnerability status for function code. We have conducted extensive experiments against a large-scale collection of real-world C/C++ vulnerabilities obtained from multiple real-world projects, and demonstrate an increase of 105\% in F1-score over the current state-of-the-art.

研究の動機と目的

• 関数レベルの脆弱性予測の限界を克服し、脆弱性の文レベルでの局在化を可能にする。
• グラフニューラルネットワークを用いて文間の制御依存やデータ依存を活用する。
• SVD における予測精度を向上させるために、関数レベルと文レベルの情報を統合する。
• 実世界の不均衡データを用いた、スケーラブルで完全教師ありの手法を提供する。

提案手法

• 各関数を、ノードが文、エッジが制御依存性とデータ依存性を符号化するグラフとして表現する（PDG/CDG）。
• CodeBERT を用いて機能レベルおよび文レベルの埋め込みを抽出し、専用のトークン化と分類トークンを埋め込みに用いる。
• PDG/CDG のエッジを用いて、文間の文脈情報を伝播させるために Graph Attention Networks を適用する。
• 文レベルの GNN 出力と関数レベルの埋め込みを組み合わせる共有分類器を訓練し、予測を整合させるための要素ごとの相互作用を含む。
• 文ラベルに対するクロスエントロピー損失で最適化し、関数レベルのサンプルをバランスさせる訓練スキームを用い、文レベルで評価する。）

実験結果

リサーチクエスチョン

• RQ1 RQ1: LineVD は statement-level vulnerability prediction を、最先端の解釈ベースの SVD モデル IVDetect と比較してどのように改善しますか？
• RQ2 RQ2: 異なるコード埋め込み手法（CodeBERT、Doc2Vec、GloVe）は、文レベルの SVD の性能にどのように影響しますか？
• RQ3 RQ3: GNN（GCN vs. GAT）とグラフタイプ（PDG vs. CDG）の使用が性能に与える影響は何か、また関数レベル情報はどのように役立つか？
• RQ4 RQ4: 未知のプロジェクトからのテストデータが含まれるクロスプロジェクト環境で LineVD はどのように性能を発揮しますか？
• RQ5 RQ5: 実世界データで LineVD が最もよく区別する文タイプはどれですか？

主な発見

• LineVD は現在の最先端の解釈ベースモデルに対して F1 スコアを105%向上させます。
• LineVD は関数レベルと文レベルの情報を共同学習することで利益を得て、単一情報ソースのベースラインを大幅に上回ります。
• PDG/CDG を GNNs（GAT がしばしば好ましい）と組み合わせると、非グラフベースラインと比べて文レベルの脆弱性検出が向上します。
• LineVD は、大規模な実世界データセット Big-Vul の、約1万脆弱、約177千非脆弱のサンプルという現実的な不均衡データとクロスプロジェクト分割の下で有効性を示します。
• CodeBERT ベースの埋め込みは Doc2Vec および GloVe を上回ります。
• この研究は、データセット、コード、モデルを含む公開レプリケーションパッケージを提供します。