[논문 리뷰] Using Machine Learning for Vulnerability Detection and Classification

Uri Alon, Shaked Brody외 1명|arXiv (Cornell University)|2018. 08. 04.

Software Engineering Research인용 수 452

한 줄 요약

이 논문은 소스 코드의 추상 구문 트리(abstract syntax tree, AST) 내에서 구성적 경로로 코드 스니펫을 인코딩함으로써 자연어 시퀀스를 생성하는 새로운 신경 모델인 CODE2SEQ을 소개한다. 이 모델은 경로를 임베딩하기 위해 LSTM을 사용하고, 디코딩 중에 이를 주의 메커니즘(attention mechanism)으로 조사한다. 이 모델은 두 가지 언어와 네 개의 데이터셋에서 코드 요약 및 캡션 생성 작업에서 이전의 코드 전용 및 최신 신경 기계 번역(NMT) 모델을 크게 능가한다.

ABSTRACT

The work described in this paper aims at developing a machine learning based tool for automatic identification of vulnerabilities on programs (source, high level code), that uses an abstract syntax tree representation. It is based on FastScan, using code2seq approach. Fastscan is a recently developed system aimed capable of detecting vulnerabilities in source code using machine learning techniques. Nevertheless, FastScan is not able of identifying the vulnerability type. In the presented work the main goal is to go further and develop a method to identify specific types of vulnerabilities. As will be shown, the goal will be achieved by optimizing the model’s hyperparameters, changing the method of preprocessing the input data and developing an architecture that brings together multiple models to predict different specific vulnerabilities. The preliminary results obtained from the training stage, are very promising. The best f1 metric obtained is 93% resulting in a precision of 90% and accuracy of 85%, according to the performed tests and regarding a trained model to predict vulnerabilities of the injection type.

연구 동기 및 목표

• 문자열 기반의 순차 모델이 코드의 계층적 구조를 忽略하는 것을 방지하기 위해 문법적 구조를 활용함으로써 코드-텍스트 생성을 향상시키기 위해.
• 순차 기반 모델이 코드의 계층적 구조를 간과하는 한계를 해결하기 위해.
• 구성적 AST 경로를 사용함으로써 새로운 코드 패턴에 더 잘 일반화될 수 있는 모델을 개발하기 위해.
• AST 경로를 통한 구조적 인코딩이 코드 요약 및 캡션 생성에서 뛰어난 성능을 낳는다는 것을 입증하기 위해.

제안 방법

• 코드 스니펫을 그의 AST 내의 종단 노드들 사이의 고정 길이 경로 k개의 집합으로 표현하기.
• 구성적 문법적 구조를 포착하기 위해 단방향 LSTM을 사용하여 각 AST 경로를 벡터로 인코딩하기.
• 디코딩 중에 인코딩된 경로 벡터 집합에 주의 메커니즘을 적용하여 각 출력 토큰을 생성하기.
• 모든 학습 반복마다 k개의 새로운 AST 경로를 샘플링하여 데이터 수준의 정규화를 제공하고 일반화 능력을 향상시키기.
• 교사 강제(teacher forcing)를 사용한 교차 엔트로피 손실을 통해 모델을 엔드 투 엔드로 훈련시키기.
• 코드와 자연어 입력 및 출력 모두에 서브워드 토크나이제이션(Byte-Pair Encoding)을 적용하기.

실험 결과

연구 질문

• RQ1 AST 경로로 코드를 모델링하는 것이 토큰 수준의 순차 모델링에 비해 코드-텍스트 생성 성능을 향상시키는가?
• RQ2 원시 토큰이 아닌 구조적 경로에 주의를 기울이는 것이 코드 생성에서 성능을 향상시키는가?
• RQ3 학습 중에 동적 경로 샘플링을 사용할 경우 모델의 일반화 능력과 성능에 어떤 영향을 미치는가?
• RQ4 구조적 인코딩은 기능적으로 동일한 코드에서 표면적 문법적 변형에 대해 얼마나 민감도를 낮추는가?
• RQ5 AST 경로를 기반으로 훈련된 모델은 훈련 중에 볼 수 없었던 새로운 구문 패턴으로 일반화할 수 있는가?

주요 결과

• CODE2SEQ은 코드 요약 및 캡션 생성 작업에서 최신 기술(SOTA) 성능을 달성하였으며, 전용 코드 모델과 최신 기술 NMT 모델을 모두 능가한다.
• 선형화된 AST를 사용한 이전 SOTA 모델 대비 C# 주석 생성 데이터셋에서 BLEU 점수를 62% 향상시켰다.
• 제거 실험(ablation study) 결과, 문법적 경로를 제거하고 토큰만 사용할 경우 성능이 크게 떨어지며, 이는 구조적 인코딩의 중요성을 입증한다.
• 모든 학습 반복마다 동적 경로 샘플링을 적용함으로써 일반화 능력과 성능이 향상되었으며, 고정 경로 샘플링보다 뛰어난 성능을 보였다.
• CODE2SEQ의 주의 메커니즘은 매우 중요하며, 주의 기능을 제거한 버전은 표준 NMT에서와 유사하게 성능이 크게 하락하였다.
• 모델은 새로운 코드 패턴으로도 잘 일반화되며, LSTM을 통해 어떤 문법적 경로라도 표현하고 주의를 기울일 수 있기 때문에, 훈련 중에 볼 수 없었던 경로라도 처리할 수 있다.