[論文レビュー] Vul-LMGNNs: Fusing language models and online-distilled graph neural networks for code vulnerability detection

Ruitong Liu, Yanbin Wang+5 more|arXiv (Cornell University)|Apr 23, 2024

Software Engineering Research被引用数 4

ひとこと要約

Vul-LMGNN は、事前学習済みコード言語モデル（CodeBERT）と、構文、制御フロー、データ依存関係情報を強化したコードプロパティグラフ（CPG）を統合する統一モデルを提案する。ゲート付きグラフニューラルネットワーク（GGNN）を用いて、シーケンス的および構造的特徴を同時に学習する。この手法は最先端の性能を達成し、小規模データセットでは従来手法と比較して F1 スコアが約 10% 高い。

ABSTRACT

Code Language Models (codeLMs) and Graph Neural Networks (GNNs) are widely used in code vulnerability detection. However, GNNs often rely on aggregating information from adjacent nodes, limiting structural information propagation across layers. While codeLMs can supplement GNNs with semantic information, existing integration methods underexplore their collaborative potential. To address these challenges, we propose Vul-LMGNNs, integrating pre-trained codeLMs with GNNs to enable cross-layer propagation of semantic and structural information. Vul-LMGNNs leverage Code Property Graphs (CPGs) to incorporate syntax, control flow, and data dependencies, using gated GNNs for structural extraction. An online knowledge distillation (KD) mechanism allows a student GNN to capture structural information from a trained counterpart via alternating training. Additionally, an "implicit-explicit" joint training framework leverages codeLMs to initialize embeddings and propagate code semantics. In the explicit phase, it performs late fusion via linear interpolation. Evaluations on real-world vulnerability datasets show Vul-LMGNNs outperform 17 state-of-the-art approaches. Source code is available at: https://github.com/Vul-LMGNN/vul-LMGNN.

研究の動機と目的

• シーケンスベースおよびグラフベースのコード脆弱性検出手法の限界、すなわち構造的依存関係や長距離コンテキストを無視する傾向を是正すること。
• 多様なコードプロパティグラフからのグローバルな構造的情報を、事前学習済みコード言語モデルの局所的意味的特徴と統合すること。
• コード言語モデルとゲート付き GNN の共同学習により、複雑なコード依存関係を捉えることで、脆弱性検出の精度を向上させること。
• Vul-LMGNN と CodeBERT の予測を統合する線形補間分類器を用いることで、モデルの頑健性と一般化性能を向上させること。
• 複数の実世界の脆弱性データセットにおいて、提案アーキテクチャの有効性を実証すること。

提案手法

• 構文、制御フロー、データフローを表現するため、抽象構文木（AST）、制御フローダイアグラム（CFG）、プログラム依存グラフ（PDG）を統合したコードプロパティグラフ（CPG）を構築する。
• 局所的意味的コンテキストを注入するために、微調整済みの CodeBERT モデルを用いてノード埋め込みを初期化する。
• GRU ユニットを備えたゲート付きゲート付きニューラルネットワーク（GGNN）を採用し、隣接ノードからのメッセージを反復的に集約することで、長距離依存関係と多様なコード特徴を捉える。
• CodeBERT と GGNN モジュールをエンドツーエンドで共同学習させ、順序的および構造的コード表現の暗黙的統合を可能にする。
• Vul-LMGNN と事前学習済み CodeBERT モデルの予測を統合する線形補間分類器を用い、最終的な検出性能を向上させる。
• ターゲットデータセット上で CodeBERT モデルを微調整し、脆弱性固有のパターンに埋め込みを適応させる。

実験結果

リサーチクエスチョン

• RQ1 事前学習済みコード言語モデルと統一されたコードプロパティグラフを組み合わせることで、脆弱性検出性能が向上するか？
• RQ2 標準的な GNN と比較して、ゲート付き GNN はコードグラフ内の多様な依存関係（構文、制御フロー、データフロー）をどれほど効果的に捉えられるか？
• RQ3 コード言語モデルと GNN の共同学習は、独立した学習よりも優れた特徴学習をもたらすか？
• RQ4 Vul-LMGNN と CodeBERT の予測を線形補間する際の性能向上は、全体の検出精度にどのように影響するか？
• RQ5 コード言語モデルと組み合わせた場合、GGNN、GCN、GAT といった異なる GNN アーキテクチャの脆弱性検出性能への影響は何か？

主な発見

• Vul-LMGNN は、4 つの実世界の脆弱性データセットで最先端の性能を達成し、6 つの最先端手法を上回った。
• 小規模データセットでは、従来手法と比較して約 10% 高い F1 スコアを達成しており、限られたデータでも優れた一般化性能を示している。
• ターゲットデータセット上で CodeBERT を微調整することでノード埋め込みを初期化すると、最も高い性能（F1: 83.87%）が得られ、事前学習済み重みの直接使用を上回った。
• GGNN を用いたモデルは、すべての指標で GCN や GAT を上回り、F1 スコアは 83.87%（GGNN）対 82.15%（GCN）対 78.39%（GAT）であった。
• GGNN の GRU ベースのメッセージ伝達メカニズムは、長距離依存関係と多様なコード特徴を効果的に捉えており、優れた性能の背景にある。
• 補助的な線形補間分類器は、Vul-LMGNN と CodeBERT の予測を明示的に統合することでモデル性能を向上させ、脆弱性検出におけるアンサンブル学習の価値を示している。