[논문 리뷰] On using distributed representations of source code for the detection of C security vulnerabilities

David Coimbra, Sofia Reis외 3명|arXiv (Cornell University)|2021. 06. 01.

Advanced Malware Detection Techniques참고 문헌 25인용 수 15

한 줄 요약

이 논문은 추상 구문 트리에서 추출한 경로-컨텍스트를 사용하여 C 소스 코드의 분산 표현을 학습하는 Code2vec 모델을 평가하여 보안 취약성을 탐지하는 데 목적이 있다. CodeXGLUE 벤치마크에서 훈련된 모델은 61.43%의 정확도를 기록했으며, RoBERTa와 유사한 성능을 보였지만 계산 비용은 훨씬 낮아 자원이 제한된 CI 파이프라인에 적합하다.

ABSTRACT

This paper presents an evaluation of the code representation model Code2vec when trained on the task of detecting security vulnerabilities in C source code. We leverage the open-source library astminer to extract path-contexts from the abstract syntax trees of a corpus of labeled C functions. Code2vec is trained on the resulting path-contexts with the task of classifying a function as vulnerable or non-vulnerable. Using the CodeXGLUE benchmark, we show that the accuracy of Code2vec for this task is comparable to simple transformer-based methods such as pre-trained RoBERTa, and outperforms more naive NLP-based methods. We achieved an accuracy of 61.43% while maintaining low computational requirements relative to larger models.

연구 동기 및 목표

• C 소스 코드의 보안 취약성을 탐지하는 데 있어 Code2vec의 효과성을 평가하는 것.
• RoBERTa 및 CodeBERT와 같은 트랜스포머 기반 모델과 비교해 Code2vec의 성능 및 계산 효율성을 평가하는 것.
• 경로-컨텍스트 기반 코드 표현이 최소한의 계산 오버헤드로 경쟁 가능한 정확도를 달성할 수 있는지 평가하는 것.
• Code2vec를 사용한 취약성 탐지에 적합한 재현 가능한 구현 및 훈련 파이프라인을 제공하는 것.
• 크기가 작은 모델인 Code2vec가 대규모 모델이 구현이 어려운 CI 환경에서 실용적인 대안이 될 수 있는지 보여주는 것.

제안 방법

• astminer 라이브러리를 사용하여 C 함수의 추상 구문 트리(AST)에서 경로-컨텍스트를 추출하였다.
• Code2vec의 주의 기반 피드포워드 네트워크를 사용하여 AST 경로-컨텍스트에서 분산 임베딩을 학습함으로써 코드 표현을 구성하였다.
• CodeXGLUE 결함 탐지 벤치마크의 레이블이 부여된 C 함수 코퍼스에서 Code2vec를 미세조정하였다.
• 함수를 취약 또는 비취약으로 분류하기 위해 이진 교차 엔트로피 손실을 사용하여 모델을 훈련시켰다.
• 최적의 성능를 위해 배치 크기(1024), 학습률, 최대 경로-컨텍스트 길이 및 너비와 같은 하이퍼파rameter를 설정하였다.
• 정확도와 같은 표준 지표를 사용하여 모델 성능을 평가하였으며, RoBERTa 및 기타 NLP 기준 모델과의 결과를 비교하였다.

실험 결과

연구 질문

• RQ1 AST 경로-컨텍스트 기반의 비트랜스포머 모델인 Code2vec가 C 보안 취약성을 탐지하는 데 경쟁 가능한 정확도를 달성할 수 있는가?
• RQ2 RoBERTa 및 CodeBERT와 같은 사전 훈련된 트랜스포머 모델과 비교해 Code2vec의 성능은 어떻게 되는가?
• RQ3 Code2vec의 훈련에 소요되는 계산 비용은 더 큰 모델 대비 얼마나 되며, CI 파이프라인에 통합하는 데 실용적인가?
• RQ4 AST 기반 경로-컨텍스트 사용이 취약성 탐지에 필요한 의미적 및 구문적 정보를 충분히 포착할 수 있는가?
• RQ5 C 함수의 레이블이 부여된 데이터셋에서 Code2vec를 효과적으로 미세조정하여 최소한의 자원으로 최신 기술 수준의 성능을 달성할 수 있는가?

주요 결과

• Code2vec는 C 코드의 CodeXGLUE 결함 탐지 벤치마크에서 61.43%의 정확도를 기록했으며, 단순한 NLP 기반 방법보다 뛰어난 성능을 보였다.
• Code2vec는 소스 코드에서 사전 훈련되지 않은 점을 감안할 때 RoBERTa와 같은 간단한 미세조정된 트랜스포머 모델과 유사한 성능을 보였다.
• Code2vec는 소비자 수준의 GPU에서 약 5분의 훈련 시간만으로도 훈련이 가능했으며, 1024개 샘플 배치에서도 메모리가 소진되지 않았다.
• 모델는 낮은 계산 자원 소비를 보이며, 자원이 제한된 환경(예: CI 파이프라인)에서의 구현에 적합하다.
• Code2vec는 더 깊은 구문적 및 의미적 특징을 활용하는 더 표현력이 뛰어난 모델인 CodeBERT 및 Devign에 비해 성능에서 뒤져야 했다.
• 결과적으로 Code2vec는 성능과 효율성 사이의 강력한 트레이드오프를 제공하며, 특히 모델 크기와 추론 비용이 중요한 환경에서는 매우 유용하다.