[논문 리뷰] Transformer-based Vulnerability Detection in Code at EditTime: Zero-shot, Few-shot, or Fine-tuning?

Aaron Chan, Anant Kharkar외 6명|arXiv (Cornell University)|2023. 05. 23.

Software Engineering Research인용 수 7

한 줄 요약

이 논문은 사전 훈련된 LLM(코드BERT, code-davinci-002, text-davinci-003)을 사용하여 문법적으로 불완전한 코드에서 취약점을 실시간으로 탐지하는 에디터 타임 취약점 검출 시스템을 제안한다. 피팅된 모델과 제로샷/피샷 프롬프팅을 활용하여, 최신 기술 대비 재현율 10% 향상 및 정밀도 8% 향상을 달성했으며, LLM이 생성한 코드의 취약점 비율을 90% 이상 감소시켰다.

ABSTRACT

Software vulnerabilities bear enterprises significant costs. Despite extensive efforts in research and development of software vulnerability detection methods, uncaught vulnerabilities continue to put software owners and users at risk. Many current vulnerability detection methods require that code snippets can compile and build before attempting detection. This, unfortunately, introduces a long latency between the time a vulnerability is injected to the time it is removed, which can substantially increases the cost of fixing a vulnerability. We recognize that the current advances in machine learning can be used to detect vulnerable code patterns on syntactically incomplete code snippets as the developer is writing the code at EditTime. In this paper we present a practical system that leverages deep learning on a large-scale data set of vulnerable code patterns to learn complex manifestations of more than 250 vulnerability types and detect vulnerable code patterns at EditTime. We discuss zero-shot, few-shot, and fine-tuning approaches on state of the art pre-trained Large Language Models (LLMs). We show that in comparison with state of the art vulnerability detection models our approach improves the state of the art by 10%. We also evaluate our approach to detect vulnerability in auto-generated code by code LLMs. Evaluation on a benchmark of high-risk code scenarios shows a reduction of up to 90% vulnerability reduction.

연구 동기 및 목표

• 코드 작성 중 실시간 취약점 검출의 핵심적 격차를 메우기 위해 문법적으로 잘못된 코드 스니펫에서도 검출이 가능한 시스템을 제공한다.
• 컴파일 후가 아니라 코드 작성 시점(에디터 타임)에 결함을 조기에 발견함으로써 취약점 수정의 비용과 지연 시간을 줄인다.
• 사전 훈련된 LLM에서 제로샷, 피샷, 파인튜닝 접근 방식의 성능을 평가하고 비교하여 문법적으로 불완전한 코드의 취약점 탐지에 활용한다.
• 코드 LLM이 자동 생성한 코드에서의 취약점 탐지 능력을 평가한다. 이는 점점 증가하는 공격 표면이다.
• 실제 IDE 환경(VSCode)에 시스템을 구현하고 배포하여 실제 개발 환경에서의 취약점 감소 효과를 측정한다.

제안 방법

• 정적 분석을 사용하여 오픈소스 레포지토리에서 50만 개 이상의 취약 코드 스니펫을 수집하여 모델 훈련 및 평가에 활용했다.
• 세 가지 학습 철학을 평가: 제로샷(직접 프롬프트 기반 추론), 피샷(컨텍스트 내 예시 포함), 그리고 CodeBERT 및 OpenAI의 code-davinci-002/text-davinci-003에 대한 파인튜닝.
• 실시간 IDE 통합을 위해 저지연성 추론을 보장하기 위해 CodeBERT-base(1억 개 미만의 파라미터)를 사용했다.
• 취약점 탐지 작업에 특화된 프롬프트와 피샷 예시를 설계하여 InstructGPT 기반 모델의 지시 준수 능력을 활용했다.
• 성능이 가장 뛰어난 모델(파인튜닝된 CodeBERT)을 VSCode 확장 프로그램으로 배포하여 개발 중 실시간 피드백을 제공했다.
• 지속적인 커버리지 확장을 위해 새로운 취약 유형을 반영할 수 있도록 재훈련 파이프라인을 구현했다.

실험 결과

연구 질문

• RQ1 사전 훈련된 LLM은 문법적으로 불완전한 코드에서 에디터 타임에 취약점을 탐지할 수 있는가? 제로샷, 피샷, 파인튜닝 접근 방식 간 성능 차이는 어떻게 되는가?
• RQ2 완전한 코드가 필요로 하는 최신 기술 대비 에디터 타임 취약점 검출의 성능은 어떠한가?
• RQ3 LLM이 자동 생성한 코드에서의 취약점 탐지 능력은 어느 정도인가?
• RQ4 생산 환경 배포 시 모델 크기, 추론 비용, 정밀도, 재현율, 유지보수 오버헤드 간의 상충 관계는 어떠한가?
• RQ5 임계값 조정과 재훈련 파이프라인은 장기적인 개발자 신뢰와 탐지 정확도를 유지를 위해 어떻게 기여하는가?

주요 결과

• 파인튜닝된 CodeBERT가 59% 정밀도와 63% 재현율로 가장 균형 잡힌 성능을 보였으며, 제로샷 및 피샷 방법보다 예측 성능에서 뛰어났다.
• text-davinci-003에서의 제로샷 및 피샷 추론은 각각 78%와 75%의 높은 재현율을 기록했지만 정밀도는 47%와 49%로 낮아 민감도와 임의 경고 사이의 상충 관계를 보였다.
• 제안된 시스템은 표준 벤치마크에서 최신 기술 대비 재현율 10% 향상 및 정밀도 8% 향상하여 성능을 개선했다.
• 고위험 벤치마크에서 LLM 생성 코드의 취약점 비율을 90% 이상 감소시켜 자동 생성 편집에 대한 강력한 일반화 능력을 입증했다.
• VSCode 확장 프로그램으로의 배포로 인해 실제로 수정 중인 코드에서 취약점 비율이 80% 감소하여 실제 환경에서의 영향을 입증했다.
• 새로운 취약 유형이 지속적으로 등장하므로 재훈련 파이프라인이 커버리지 유지와 신뢰 확보에 필수적이었다.