지난 블로그(링크)에서는 코드를 이해하는 언어모델에 대한 소개와 취약점 탐지 방법론에 대하여 크게 구별형 언어모델과 생성형 언어모델로 나누어 개념적으로 소개하였다면, 이번 블로그에서는 생성형 언어모델을 이용하여 실제 코드 취약점을 식별하고 제시하는 모델의 구현 과정 및 검증 결과에 대해 살펴보고자 합니다.
생성형 언어모델의 취약점 탐지 학습
그림 1. 생성형 언어모델의 취약점 탐지 학습 개요
소스코드 취약점 탐지를 위하여 생성형 언어모델을 활용하기 위해서는 그림 1의 학습 과정과 같이 기존의 코드와 취약점이 매칭된 데이터셋을 전처리하고 이를 프롬프트(Prompt)에 맞게 변형하는 과정이 필요합니다. 이후, 프롬프트 형식에 맞게 변형된 데이터를 이용하여 생성형 언어모델의 학습을 진행하고 여러 번의 튜닝 작업을 통해 실제 코드가 가지고 있는 취약점을 식별할 수 있습니다.
<코드-취약점 데이터셋 구성>
표 1. 취약점 탐지 데이터셋 구성
| CWE ID | 취약점 설명 | 순위 | 코드수 |
|---|---|---|---|
| CWE-787 | Out of bounds Write | 1 | 31,692 |
| CWE-416 | Use After Free | 4 | 17,894 |
| CWE-20 | Improper Input Validation | 6 | 18,739 |
| CWE-125 | Out of bounds Read | 7 | 23,161 |
| CWE-476 | Null Pointer Dereference | 12 | 15,121 |
| CWE-190 | Integer Overflow or Wraparound | 14 | 9,384 |
| CWE-119 | Improper Restriction of Operation within Bounds of Memory Buffer | 17 | 21,937 |
| CWE-703 | Improper Check or Handling of Exceptional Conditions | – | 19,910 |
| CWE-200 | Exposure of Sensitive Information to an Unauthorized Actor | – | 10,890 |
| 정상 코드(GNU Coreutils) | 18,823 | ||
| Total | 187,551 | ||
생성형 언어모델이 정확하게 코드 취약점을 탐지하기 위해서는 취약점과 맵핑된 충분한 양의 소스코드와 취약점이 쌍으로 이루어진 데이터셋 확보가 무엇보다 중요합니다. 따라서 본 연구에서는 데이터셋 구성 과정에서 인터넷에 공개된 다양한 다양한 취약점 데이터셋을 종합하여 MITRE의 CWE(Common Weakness Enumeration) 체계에 따라 보안 취약점별로 분류하는 작업을 진행하였습니다. 일부 데이터셋의 경우 파일 단위의 코드로 제공하는데, 이 같은 경우 코드의 길이가 길기 때문에 학습자원과 시간 측면에서 제약사항이 종종 발생하게 됩니다. 이에 최종적으로 구성한 데이터셋은 함수 단위의 코드 데이터를 제공하는 BigVul[1], SVEN[2], DiverseVul[3]을 사용하였습니다. 또한 정상 데이터셋의 경우에는 유닉스 운영체제의 명령어 구현 코드인 GNU Coreutils[4]의 코드를 함수별로 수집하여 구성하였습니다. 이렇게 수집된 데이터셋의 경우 다양한 취약점을 포함하고 있지만 일부 취약점의 경우 극히 소수이거나, 위협도 측면에서 그 중요성이 미미할 수 있어 데이터셋 구축 과정에서 취약점별 수집된 코드의 수와 MITRE CWE Top 25 Most Dangerous Software Weaknesses[5]를 참고하여 표 1과 같이 데이터셋을 구축하였습니다.
<코드 데이터의 전처리>
그림 2. 코드 데이터 전처리 예시
대부분 인공지능 모델 학습이 그러하듯, 생성형 언어모델이 원활하게 데이터를 이해하기 위해서는 적절한 수준의 전처리가 선행되어야 합니다. 실제 코드를 살펴보면 그림 2와 같이 원본 코드를 문자열 형태로 표현하였을 때, 줄 바꿈과 탭같이 문법을 위한 기호들이 삽입된 것을 확인할 수 있습니다. 이는 실제 모델이 코드를 이해할 때 불필요할 뿐만 아니라 코드의 길이를 증가시키는 주요 요인으로 작용하기 때문에 해당 기호들을 제거해야 합니다. 다중 띄어쓰기 경우 단일 띄어쓰기로 데체하는 작업을 진행합니다. 또한 일부 코드의 길이가 현저히 짧은 경우 모델이 코드를 이해하기에 충분치 않아 학습에 악영향을 미칠 수 있어 본 전처리 과정에서는 150 토큰 이하의 데이터는 학습 및 평가에서 제외하도록 전처리하였습니다.
<데이터의 프롬프트 변형>
그림 3. 프롬프트 형태 예시
전처리 과정을 거친 데이터는 생성형 언어모델 학습을 위해 언어모델이 이해할 수 있는 질의 형태인 프롬프트 형태로 변형을 진행하게 되는데, 다양한 검증을 통해 효과적인 프롬프트를 선정하였고 그림 3과 같이 Hard, Soft, Mixed 프롬프트를 고려하였습니다. Hard 프롬프트의 경우, 흔히 ChatGPT나 다른 생성형 언어모델 서비스에서 사용하는 자연어로 이루어진 프롬프트 형태이며, Soft 프롬프트는 자연어 대신 [SOFT] 토큰을 활용한 프롬프트 방식으로, 취약점 탐지의 경우 취약점 앞에 [SOFT] 토큰을 추가하게 됩니다. 각 프롬프트 방식은 각각의 장단점을 가지고 있는데, Hard 프롬프트의 경우 사용자에게 직관적이라는 장점이 있지만, 일관성이 부족하다는 단점이 있고 Soft 프롬프트의 경우 모델이 직접 프롬프트를 학습해 성능이 좋다는 장점이 있지만, 결과에 대한 해석이 어렵고 학습이 필수적이라는 단점이 있습니다[6,7]. 이러한 Hard와 Soft 프롬프트의 장단점을 상호 보완하기 위해서 고안된 Mixed 프롬프트의 경우 자연어로 이루어진 프롬프트 내 Soft 프롬프트에서 사용하는 [SOFT] 토큰을 임의로 삽입하는 형태로 되어있습니다. 이를 통해 Hard 프롬프트의 장점인 직관적이라는 장점을 보존하면서 Soft 프롬프트의 확장성의 장점을 가질 수 있습니다.
<생성형 언어모델의 학습>
이렇게 프롬프트 변형을 마친 데이터를 이용하여 본격적으로 취약점 탐지를 위한 생성형 언어모델 학습을 진행할 수 있습니다. 학습 시에는 코드 데이터가 포함된 질의 프롬프트가 모델의 입력값으로 활용되고 이에 기반하여 모델은 취약점 설명을 답변으로 제공하도록 학습을 진행하게 됩니다. 학습 중 가장 심도 있게 고려하여야 할 사항으로는 학습 시 적당한 길이의 코드가 제공되는 것입니다. 코드가 너무 짧은 경우, 실제 취약점을 판단할 근거가 적어 좋은 성능을 기대하기 어렵고 코드가 너무 길 경우 학습자원에 한계가 발생할 수 있습니다. 따라서 코드 취약점 탐지를 위한 최적의 코드 길이를 찾는 과정이 선행되어야 합니다.
<언어모델 예측 결과 후처리>
그림 4. 답변 후처리 예시
이렇게 학습을 마친 모델은 사용자가 제시한 코드에 대한 취약점을 탐지할 수 있게 됩니다. 하지만 일부 예측된 답변에서는 그림 4와 같이 답변의 철자 오류와 같은 사소한 오류로 인해 모델의 성능이 낮아짐을 확인할 수 있습니다. 따라서 이를 해결하기 위해 모델의 예측과 정답 답변 세트와의 유사도에 기반하여 예측된 정답의 후처리를 진행하였습니다. 유사도 측정을 위해서 단어 단위의 예측과 각 정답 답변의 일치 수를 비교하여 가장 높은 일치 수를 가진 정답 답변으로 치환하였고, 일부 동일한 일치 수의 답변의 경우에는 문장 간의 유사도에 기반한 BLEU[8] 기법을 활용하여 가장 유사한 답변으로 치환하도록 후처리 과정을 추가하였습니다.
생성형 언어모델의 취약점 탐지 실험
실험 단계에서는 학습 과정을 통해 구축된 모델이 실제 취약점을 효과적으로 탐지하는지, 앞서 소개한 Hard, Soft, Mixed 프롬프트 간의 탐지 성능을 비교하는 실험을 진행하였습니다.
<프롬프트 간의 탐지 성능 비교>
프롬프트별 성능을 비교하기 위한 평가 기법으로는 정확도와 더불어 불균형 데이터 분류 성능에서 주로 사용되는 Macro-F1 Score를 사용하여 성능을 비교하였습니다.
표 2. 프롬프트 별 성능 비교
| 프롬프트 형태 | Hard 프롬프트 | Soft 프롬프트 | Mixed 프롬프트 |
| 정확도 | 72.24% | 70.80% | 77.50% |
| Macro-F1 | 72.53% | 71.13% | 77.58% |
표 2에서 볼 수 있듯, 모든 프롬프트 형태가 70% 이상의 성능을 기록하였으나 Soft 프롬프트가 가장 낮은 정확도와 Macro-F1 Score를 기록하였으며, Mixed 프롬프트가 다른 프롬프트 형태보다 정확도와 Macro-F1 Score에서 5~7% 이상의 높은 성능을 기록하였습니다. 이는 앞서 설명한 이유에서 기반하여 Hard와 Soft 프롬프트를 상호 보완한 Mixed 프롬프트의 특성상 높은 성능을 보일 수 있었던 것으로 보입니다.
<모델별 탐지 성능 비교>
모델별 탐지 성능 비교 실험에서는 본 연구에서 제안된 CodeT5 기반의 모델과 함께 다양하게 사용할 수 있는 코드를 학습한 언어모델 간 성능을 비교하였습니다. 기존 다른 연구에서 많이 사용된 CodeBERT[9]와 함께 OpenAI API를 통한 질의가 가능한 GPT-3.5[10], GPT-4[11]와 더불어 오픈소스로 공개된 Code Llama[12], Llama3[13]에도 질의하여 성능을 비교하였습니다. 또한 보다 정확한 실험을 위해 5개의 다른 seed를 사용하여 정확도, Macro-F1 Score의 평균을 계산 후 성능을 평가하였습니다.
표 3. 모델별 성능 비교
| Model | CodeT5 | CodeBERT | GPT-3.5 | GPT-4 | Code Llama | Llama3 |
| 정확도 | 81.71% | 73.36% | 10.14% | 11.06% | 9.62% | 10.76% |
| Macro-F1 | 81.93% | 73.53% | 5.43% | 6.73% | 7.79% | 8.68% |
표 3에서 볼 수 있듯, 학습을 진행한 CodeT5 기반의 우리 연구팀 모델과 더불어 CodeBERT 기반의 모델이 단순 질의 과정을 통해 취약점을 탐지한 모델에 비해 약 60~70% 높은 성능을 기록하였습니다. 특히 GPT 계열의 모델과 Llama 계열 모델의 경우, 모든 코드에 대해 취약점이 있다고 판단하거나, 일부 취약한 코드에 대해 정상 코드로 판단하는 등 오분류 결과를 보이기도 하였습니다.
글을 마치며
그림 5. 취약점 탐지 모델 활용
생성형 언어모델은 다양한 분야에서 활발히 활용되고 있으며, 코드 데이터를 학습한 모델은 이를 기반으로 한 여러 작업에 효과적으로 적용될 수 있습니다. 그중 하나인 코드 취약점 탐지 분야에서는 기존 정적 분석 기법의 한계를 생성형 언어모델을 통해 극복할 수 있다는 장점이 있습니다. 본 블로그 글에서는 코드-취약점 데이터를 전처리하고, 이를 프롬프트 기반 학습에 적용함으로써 생성형 언어모델을 효율적으로 학습하고 활용하는 방법을 소개하였습니다. 이렇게 구축된 모델은 기존 온라인 서비스에서 문제로 지적되어 온 소스코드 유출의 위험 없이, 안전하면서도 정확한 취약점 탐지 결과를 제공할 수 있었습니다. 앞으로도 코드 학습 기반의 생성형 언어모델은 다양한 분야에서 활용될 가능성이 높으며, 보다 안전한 사이버 환경 구축에 기여할 것으로 기대합니다.
참고문헌
[1] Jiahao, et al, “A C/C++ Code Vulnerability Dataset with Code Changes and CVE Summaries”, MSR, 2020
[2] Jingxuan He, Martin Vechev, “Large Language Models for Code: Security Hardening and Adversarial Training”, arXiv, 2023
[3] Yizheng, et al, “DiverseVul: A New Vulnerable Source Code Dataset for Deep Learning Based Vulnerability Detection”, arXiv, 2023
[4] GNU CoreUtils, Free Software Foundation, Inc
[5] CWE Top 25 Most Dangerous Software Weaknesses
[6] Lester, et al, “The Power of Scale for Parameter-Efficient Prompt Tuning”, arXiv, 2021
[7] Cai, et al, “Adapting Knowledge Prompt Tuning for Enhanced Automated Program Repair”, arXiv, 2023
[8] Kishore, et al, “Bleu: a Method for Automatic Evaluation of Machine Translation”, In Proceedings of the 40th Annual Meeting of the Association for Computational Linguistics, 2002
[9] Feng, et al, “CodeBERT: A Pre-trained Model for Programming and Natural Languages”, EMNLP, 2020
[10] OpenAI: GPT-3.5 API, https://platform.openai.com/docs/models/gpt-3-5, proprietary model by OpenAI, accessed via API
[11] OpenAI, “GPT-4 technical report”, arXiv, 2023
[12] Roziere, et al, “Code Llama: Open Foundation Models for Code”, arXiv, 2023
[13] Maaten, et al, “The Llama 3 Herd of Models”, arXiv, 2024
KAIST 사이버보안연구센터 AI 보안팀 연구원으로 AI 및 LLM을 이용한 보안 연구를 수행하고 있다.