[논문 리뷰] VulBERTa: Simplified Source Code Pre-Training for Vulnerability Detection

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

Software Engineering Research인용 수 9

한 줄 요약

VulBERTa는 C/C++ 소스 코드의 취약성 탐지에 적합한 간소화되고 소규모의 사전 훈련 방법을 제안한다. 이는 BPE와 사전 정의된 코드 토큰(키워드, 구두점, API 호출 등)을 조합한 고유한 토크나이제이션 파이프라인을 사용한다. 125M 파라미터로 228만 개의 함수를 사전 훈련하여, 다양한 데이터셋과 벤치마크에서 최신 기술 수준의 성능을 달성했으며, Draper에서 F1 점수 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)를 사용한 효율적인 피니팅을 가능하게 하면서도 높은 탐지 정확도를 유지하기.

제안 방법

• BPE와 C/C++ 키워드, 구두점, 일반적인 API 함수 이름 등의 사전 정의된 토큰을 조합한 고유한 토크나이제이션 파이프라인을 통해 문법적 및 의미적 구조를 유지한다.
• 모델은 RoBERTa를 기반 아키텍처로 사용하며, 오픈소스 프로젝트에서 추출한 228만 개의 C/C++ 함수를 대상으로 마스크 언어 모델링(Masked Language Modeling, MLM)을 통해 사전 훈련한다.
• 사전 훈련은 실제 C/C++ 코드로 구성된 정제된 데이터셋을 기반으로 수행되며, 명시적 종속성 또는 데이터 플로우 분석 없이 문법적 및 의미적 지식에 중점을 둔다.
• 피니팅은 두 가지 경량 분류기인 다층 퍼셉트론(VulBERTa-MLP)과 컨볼루션 신경망(VulBERTa-CNN)을 사용하여 수행된다.
• 복잡한 그래프 기반 또는 시퀀스 기반 아키텍처를 피하고, 토크나이제이션과 사전 훈련의 품질을 통해 깊은 코드 표현을 인코딩하는 데 의존한다.
• 평가 작업은 Vuldeepecker, Draper, REVEAL, muVuldeepecker 등 다양한 데이터셋과 CodeXGLUE 및 D2A 벤치마크에서 이진 및 다중 클래스 취약성 탐지 작업을 포함한다.

실험 결과

연구 질문

• RQ1 복잡한 아키텍처 설계에 의존하지 않고도 간소화되고 소규모의 사전 훈련 모델이 취약성 탐지에서 최신 기술 수준의 성능을 달성할 수 있는가?
• RQ2 문법적 및 의미적 구조를 유지하는 고유한 토크나이제이션 파이프라인이 취약성 탐지에 있어 코드 표현 학습을 얼마나 향상시키는가?
• RQ3 파라미터 수와 사전 훈련 데이터가 적은 더 작은 모델이 여전히 벤치마크 데이터셋에서 더 큰 모델들을 능가할 수 있는가?
• RQ4 사전 훈련된 표현이 다양한 데이터셋과 취약성 탐지 작업 간에 얼마나 이식 가능한가?
• RQ5 강력한 사전 훈련된 코드 표현과 결합된 경량 피니팅 헤드(MLP 및 CNN)가 높은 성능을 달성할 수 있는가?

주요 결과

• VulBERTa는 도전적인 Draper 데이터셋에서 57.92%의 최신 기술 수준 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위를 기록했다.
• 125M 파라미터와 228만 개의 사전 훈련 샘플만을 사용했음에도 불구하고, C-BERT(810만 개 샘플, 1.1억 파라미터)와 CodeBERT(850만 개 샘플, 1.25억 파라미터)와 같은 모델들을 뛰어넘는 성능을 보였다.
• 다양한 피니팅 헤드 간의 성능이 강건했으며, VulBERTa-MLP는 4개 데이터셋 중 3개에서 최고의 PEM 점수를 기록했고, VulBERTa-CNN는 나머지 하나에서 최고 성능을 기록했다.