Large Language Model for Vulnerability Detection: Emerging Results and Future Directions
이 논문은 프롬프트를 활용해 LLM 취약점 탐지의 성능을 얼마나 높일 수 있는지에 관하여 고찰한다. 특히 여러 가지 형태의 프롬프트를 통해 어떤 조합이 가장 성능을 보이는지까지 확인해볼 수 있다.
현재 하고 있는 프로젝트에 적용하자면,
- 취약 소스 코드 데이터 생성/탐지 프롬프트 보완
- LLM만으로는 확실히 취약점 탐지에 어려움이 있다는 점(GPT 3.5 기준)
- 논문에 나와있는 단계들 중 대체 가능한 것(CodeQL, RAG 등의 기술)이 있고 이를 대체했을때의 퍼포먼스를 짐작해볼 수 있다.
논문 요약
-
인컨텍스트 학습(In-Context Learning, ICL)
LLM의 파라미터를 고정한 채, 프롬프트를 통해 작업에 대한 정보를 제공하여 모델이 학습 없이도 작업을 수행하도록 합니다. -
프롬프트 설계
기본 프롬프트 (P): 코드에 취약점이 있는지를 묻는 간단한 질문.
프롬프트 확장 (A*):
A1 (역할 설명): 모델을 "보안 전문가"로 설정하여 보안 중심의 응답을 유도.
A2 (프로젝트 정보): 코드의 프로젝트명과 파일명을 제공하여 문맥을 강화.
A3 (CWE 예시): Common Weakness Enumeration의 상위 25개 취약점 유형의 코드 예시를 제공하여 모델이 일반적인 취약점 패턴을 학습하도록 유도.
-
데이터셋 구성
- 취약한 함수: 취약점 수정 커밋에서 변경된 함수.
- 비취약한 함수: 동일한 파일에서 무작위로 선택된 함수.
-
모델 비교
- GPT-3.5 및 GPT-4: 프롬프트를 통해 학습 없이 평가.
- CodeBERT: 전체 파라미터를 업데이트하여 학습된 모델로 평가.
프롬프트(Prompt) 설계 상세 설명
본 연구는 GPT-3.5 및 GPT-4와 같은 대형 언어 모델(LLM)을 활용해 소스코드의 취약점을 분류(classification)하는 작업에서 어떤 프롬프트 디자인이 성능에 영향을 주는지 실험적으로 분석함
기본 프롬프트(P)를 중심으로 다음과 같은 5가지 보조 프롬프트 설계(A1~A5)가 도입되며, 이들을 다양한 방식으로 조합하여 성능을 비교
P. Task Description Prompt (기본 프롬프트)
목적: LLM에게 주어진 작업의 기본 설명을 제공하여 분류 문제의 형태를 인지시키는 역할.
형태:
Now you need to identify whether a method contains a vulnerability or not.
If it has any potential vulnerability, output: "this code is vulnerable".
Otherwise, output: "this code is non-vulnerable".
The code is: [X]. Let's start: [Z][X]: 검사 대상 코드[Z]: 응답 시작 신호
역할: 이 프롬프트는 zero-shot classification을 유도하는 매우 간단한 형태이며, 별도의 지식이나 예시 없이 작업만 명시한다.
Prompt Engineering
A1. Role Description Prompt (역할 부여 프롬프트)
목적: LLM의 행위자 역할을 보안 전문가로 설정함으로써, 보안 맥락에 집중된 판단을 유도
형태:
You are an experienced developer who knows the security vulnerability very well.기술적 효과:
- LLM이 보안 전문가로서의 역할에 몰입하게 하여, 일반적인 코드 분석이 아닌 "보안 관점"으로 응답하도록 함
- 역할 지시(prompt priming)를 통한 implicit bias 유도
A2. Project Information Prompt (문맥 정보 프롬프트)
목적: 코드가 포함된 프로젝트명과 파일명을 LLM에 알려주어, 코드의 실제 위치 정보에 기반한 더 정확한 분석을 유도
형태:
The code is from [project_name]. The filename is [file_name].기술적 효과:
- LLM이 코드 컨텍스트의 경계와 의도를 더 잘 파악
- StarCoder 실험에서 파일명 추가가 성능 향상에 기여한 사실을 기반
A3. External Knowledge from CWE (외부 지식 프롬프트)
목적: CWE에서 정의한 상위 25개 취약점 유형의 대표 예제를 포함하여, LLM이 일반화된 취약 패턴을 학습할 수 있도록 함.
형태:
Here are examples of the most dangerous CWE types:
Example1: ...
Label1: this code is vulnerable.
...기술적 효과:
- zero-shot 상황에서도 LLM이 일반적인 CWE 기반 패턴을 참조 가능.
- 외부 데이터 소스로부터의 few-shot 학습 강화.
- 모델의 사전 지식(Pretraining)에 결여된 영역을 보완.
A4. Random Sample from Training Data (무작위 학습 예시 제공)
목적: 훈련 데이터셋에서 무작위로 K개의 코드-레이블 쌍을 포함시켜 LLM이 이를 참조하여 판단하도록 유도.
형태:
Here are sampled examples from the training data:
Example1: ...
Label1: this code is vulnerable.
Example2: ...
Label2: this code is non-vulnerable.
...기술적 효과:
- In-Context Learning (ICL)의 대표적인 방법
- 최대 토큰 수(4,096)에 제한되므로 샘플 수 제한 필요
- LLM의 사전 학습 정보와 충돌하지 않으면서도 task-specific 지식 강화
A5. Retrieved Similar Code Prompt (유사 코드 검색 기반 예시)
목적: 평가 대상 코드와 의미론적으로 가장 유사한 K개 함수를 훈련 데이터에서 검색하여 제공.
형태:
Here are the most similar codes from the training data:
Example1: ...
Label1: this code is non-vulnerable.
Example2: ...
Label2: this code is vulnerable.기술적 접근:
- CodeBERT를 사용해 각 코드의 임베딩을 생성
- cosine similarity로 가장 유사한 K개 샘플을 추출
- 더 높은 관련성을 가진 예시를 제공함으로써 LLM의 의사결정 정확도 향상
Prompt 조합 예시
P + A3: 기본 설명 + CWE 예시P + A4 + A5: 기본 설명 + 무작위 예시 + 유사 예시 (최고 성능)P + A1 + A2: 역할 + 파일 정보만 추가 (효과 미미)
성능 요약 (GPT-3.5)
| Prompt 조합 | Accuracy | F1 | Precision | Recall |
|---|---|---|---|---|
| P (기본) | 50.0% | - | - | 0.0% |
| P + A3 (CWE) | 59.1% | 41.9 | 72.2 | 29.5 |
| P + A4 (K=3) | 58.8% | 50.2 | 65.8 | 41.2 |
| P + A5 (K=5) | 59.8% | 54.0 | 63.2 | 47.2 |
| P + A4 + A5 | 62.7% | 49.7 | 76.3 | 36.8 |
