[논문 리뷰] Vul-LMGNNs: Fusing language models and online-distilled graph neural networks for code vulnerability detection

Ruitong Liu, Yanbin Wang외 5명|arXiv (Cornell University)|2024. 04. 23.

Software Engineering Research인용 수 4

한 줄 요약

Vul-LMGNN는 사전 훈련된 코드 언어 모델(CodeBERT)와 문법, 제어 흐름, 데이터 의존성 정보로 강화된 코드 속성 그래프(CPG)를 융합하는 통합 모델을 제안한다. 게이트드 그래프 신경망(GGNN)을 사용해 순서적 및 구조적 특징을 동시에 학습하며, 소규모 데이터셋에서 이전 방법 대비 약 10% 높은 F1 스코어를 기록해 최신 기술 수준(SOTA) 성능을 달성한다.

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)를 구축한다.
• 로컬 의미 맥락을 CPG에 통합하기 위해 미세조정된 CodeBERT 모델을 사용해 노드 임베딩을 초기화한다.
• 이웃 노드로부터 메시지를 반복적으로 집계하는 데에 게이트드 리커런트 유닛(GRU)을 사용하는 게이트드 게이트드 신경망(GGNN)을 적용해 장거리 의존성과 이질적인 코드 특징을 포착한다.
• 코드 언어 모델과 GGNN 모듈을 엔드 투 엔드로 함께 훈련시켜 순서적 및 구조적 코드 표현 간 암묵적 융합을 가능하게 한다.
• Vul-LMGNN와 사전 훈련된 CodeBERT 모델의 예측을 융합하는 선형 보간 분류기를 사용해 최종 탐지 성능을 향상시킨다.
• 목표 데이터셋에 대해 CodeBERT 모델을 미세조정해 임베딩이 취약점 특성 패tern에 적합하도록 조정한다.

실험 결과

연구 질문

• RQ1 사전 훈련된 코드 언어 모델과 통합된 코드 속성 그래프를 결합함으로써 취약점 탐지 성능을 향상시킬 수 있는가?
• RQ2 표준 GNN 대비 게이트드 GNN이 코드 그래프 내 이질적인 의존성(구문, 제어 흐름, 데이터 흐름)을 얼마나 효과적으로 포착하는가?
• RQ3 코드 언어 모델과 GNN을 함께 훈련시키는 것이 별도 훈련보다 더 나은 특징 학습을 이끌어내는가?
• RQ4 Vul-LMGNN와 CodeBERT 예측의 선형 보간이 전체 탐지 정확도에 어떤 영향을 미치는가?
• RQ5 코드 언어 모델과 결합했을 때, 다양한 GNN 아키텍처(GGNN, GCN, GAT)가 취약점 탐지 성능에 미치는 영향은 어떠한가?

주요 결과

• Vul-LMGNN는 네 개인실 취약점 데이터셋에서 최신 기술 수준(SOTA) 성능을 달성하며, 여섯 가지 SOTA 방법을 초월한다.
• 소규모 데이터셋에서 Vul-LMGNN는 이전 방법 대비 약 10% 높은 F1 스코어를 기록하여 제한된 데이터에서의 뛰어난 일반화 능력을 보여준다.
• 노드 임베딩 초기화를 위해 대상 데이터셋에 대해 CodeBERT를 미세조정하는 것이 가장 높은 성능(F1: 83.87%)을 기록하며, 사전 훈련된 가중치를 직접 사용하는 것보다 뛰어난 성능을 낸다.
• GGNN 기반 모델은 모든 지표에서 GCN와 GAT를 압도하며, F1 스코어 83.87%를 기록한 반면 GCN는 82.15%, GAT는 78.39%를 기록한다.
• GGNN의 GRU 기반 메시지 전달 메커니즘이 장거리 의존성과 이질적인 코드 특징을 효과적으로 포착해 뛰어난 성능을 이끌어내는 데 기여한다.
• 보조적인 선형 보간 분류기는 Vul-LMGNN와 CodeBERT의 예측을 명시적으로 융합함으로써 모델 성능을 향상시키며, 취약점 탐지에서 앙상블 학습의 가치를 입증한다.