[논문 리뷰] LineVD: Statement-level Vulnerability Detection using Graph Neural Networks

David Hin, Andrey Kan외 2명|arXiv (Cornell University)|2022. 03. 10.

Software Engineering Research인용 수 32

한 줄 요약

LineVD는 진술(statement) 수준의 취약점 탐지를 프로그램 그래프의 노드 분류로 다루며, 코드 임베딩에 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 간선을 이용하여 진술 간에 그래프 어텐션 네트워크(GAT)로 맥락 정보를 전파한다.
• 진술 수준 GNN 출력과 함수 수준 임베딩을 혼합하는 공유 분류기를 학습하고, 예측 정렬을 맞추기 위한 원소별 상호작용을 포함한다.
• 진술 레이블에 대한 교차 엔트로피 손실로 최적화하고, 함수 수준 샘플의 균형을 맞추는 학습 방식으로 진술 수준에서 평가한다.

실험 결과

연구 질문

• RQ1 RQ1: LineVD가 최첨단 해석 기반 SVD 모델 IVDetect와 비교해 진술 수준 취약점 예측을 어떻게 개선하는가?
• RQ2 RQ2: 서로 다른 코드 임베딩 방법(CodeBERT, Doc2Vec, GloVe)이 진술 수준 SVD 성능에 어떤 영향을 미치는가?
• RQ3 RQ3: GNN(GCN vs. GAT)과 그래프 유형(PDG vs. CDG)의 활용이 성능에 미치는 영향은 무엇이며, 함수 수준 정보가 어떻게 도움이 되는가?
• RQ4 RQ4: 테스트 데이터가 보지 못한 프로젝트에서 나오는 크로스 프로젝트 시나리오에서 LineVD의 성능은 어떠한가?
• RQ5 RQ5: 실제 데이터에서 LineVD가 구분하기 가장 좋은 진술 유형은 무엇인가?

주요 결과

• LineVD는 현재 최첨단 해석 기반 모델 IVDetect 대비 F1 점수에서 105% 증가를 달성한다.
• LineVD는 함수 수준 정보와 진술 수준 정보를 공동으로 학습하는 이점으로 단일 소스 정보 기반 기준선보다 현저히 우수하게 성능이 향상된다.
• PDG/CDG를 GNN(GAT 선호)이 포함되어 사용하면 비그래프 기준선보다 진술 수준 취약점 탐지가 개선된다.
• LineVD는 현실적인 불균형 데이터와 Big-Vul에서의 크로스-프로젝트 분할 하에서도 효과를 보이며, 약 10k 개의 취약 샘플과 약 177k 개의 비취약 샘플이 있는 대규모 실제 데이터 세트이다.
• LineVD 프레임워크에서 진술 수준 SVD를 위한 CodeBERT 기반 임베딩이 Doc2Vec 및 GloVe보다 우수하다.
• 본 연구는 데이터세트, 코드 및 모델을 포함한 공개 재현 패키지를 제공한다.