최근 가장 주목받는 사회적 이슈를 꼽자면 그중 하나는 단연 생성형 인공지능일 것입니다. 2022년에 OpenAI가 공개한 ChatGPT로 시작된 사람들의 주목은 이후 Microsoft의 Copilot, Google의 Gemini, 국내에서는 네이버의 CLOVA X, SK텔레콤의 에이닷과 같이 다양한 생성형 인공지능 서비스의 발전으로 이어져 현재는 이미지 또는 영상 생성을 비롯한 다양한 분야에서 생성형 인공지능의 사용을 쉽게 확인할 수 있게 되었습니다. 이렇게 생성형 모델이 발전함에 따라, 다양한 분야에서 생성형 인공지능을 이용한 연구가 진행되고 있으며, 그중 프로그램 코드 분야의 경우 생성형 인공지능을 이용한 코드 생성, 다른 프로그램 언어로의 번역, 코드 보정과 같은 작업이 활발히 진행되고 있습니다.
이번 글에서는 프로그램 코드를 이해하는 언어모델에 대해 살펴보고, 이 언어모델에 기반하여 프로그램 소스 코드의 취약점을 탐지하는 기술에 대해 알아보도록 하겠습니다.
<코드를 이해하는 언어모델과 그 활용>
언어모델(Language Model)이란 인간의 언어를 이해하고 생성하기 위한 모든 모델을 총칭합니다. 언어모델은 그 사용에 앞서 용도에 알맞는 적합한 학습을 진행하여야 하는데[1], 2017년 트랜스포머[2]의 등장 이후 언어모델의 학습은 대용량의 데이터를 학습하여 언어를 이해하기 위한 파라미터를 구축하는 사전학습과, 실제 사용자의 목적에 맞는 입력과 출력 데이터쌍으로 학습을 진행하는 미세조정 학습으로 나뉘어 진행하게 됩니다.
사전학습과 미세조정 학습을 거쳐 모델을 학습하는 언어모델의 방식은 단순히 자연어뿐만 아니라 자연어의 형태를 가지고 있는 다양한 데이터를 처리할 수 있는데, 그 중 하나가 바로 프로그램 언어 분야입니다. 프로그램 언어 또한 언어의 구조를 띠고 있기에 언어모델 학습을 통해 다양한 작업을 수행할 수 있는 것입니다. 그렇다면 과연 코드 분야에서는 어떤 언어모델이 사용되며, 이를 사용하여 어떤 작업을 수행할 수 있을까요?
프로그램 코드 관련 작업을 수행하기 위해선 코드와 자연어를 동시에 이해하는 언어모델이 필수적입니다. 이를 위해 정확한 코드 혹은 자연어 생성을 위한 미세조정 학습이 진행되어야 합니다. 코드 관련 작업을 수행하기 위해 사용되는 대표적인 언어모델로는 BERT에 기반한 CodeBERT, T5에 기반한 CodeT5, BART에 기반한 PLBART 모델이 존재하며, 다음과 같은 작업을 수행할 수 있습니다.
ㆍ코드 요약
코드 요약은 입력된 코드에 대한 자연어로 구성된 설명을 생성하는 작업입니다[3]. 언어모델을 이용한 코드 요약을 통해 사용자가 직접 긴 코드를 읽지 않고도 코드에 대한 요약을 제공받아 그 기능성을 이해할 수 있습니다.
ㆍ코드 생성
코드 생성은 코드에 대한 자연어 설명을 기반으로 실제 프로그램 코드를 생성하는 작업입니다[3]. 이를 통해 사용자들은 코드 작성에 대한 지식 없이도 프로그램 동작 방법과 원하는 결과를 설명하면 이에 대한 코드를 쉽게 작성할 수 있습니다. 하지만 코드 생성의 경우에는 자연어에 비해 방대하고 양질의 학습 데이터 부재와 생성 모델의 기술적 한계로 현재까지는 완벽히 실행할 수 있는 코드를 생성하는 데에 다소 부족함은 있습니다.
ㆍ코드 번역
코드 번역은 프로그램 언어의 코드를 다른 프로그램 언어로 번역하는 작업으로, 앞서 설명한 코드 생성과 다르게 코드를 입력값으로 받아 다른 프로그래밍 언어의 코드로 생성합니다[3]. 코드 번역의 경우 과거의 프로그래밍 언어 문법으로 작성된 코드를 현재의 프로그래밍 언어 문법으로 재작성하거나, 호환성을 위해 다른 프로그램 언어로 동일한 프로그램을 변환하는 데에 활용될 수 있습니다.
ㆍ코드 유사성 검사
코드 유사성 검사는 입력된 두 코드의 유사 정도를 탐지하는 작업입니다. 특히 Code Clone Detection의 경우 단순히 코드의 구조적 일치함을 기반으로 두 코드의 유사성을 탐지하는 것에 국한되어 있지 않고 언어모델이 언어의 의미를 이해할 수 있는 특징을 이용하여 구조적으로는 상이하나 의미상으로 유사한 코드 탐지 또한 가능합니다[4].
ㆍ소스코드 취약점 탐지
소스코드 취약점 탐지는 코드 내의 취약점 탐지를 통해 사전에 외부의 공격을 방지하는 데에 목적이 있습니다. 코드를 이해하는 언어모델을 통한 소스코드 취약점 탐지는 기존의 정적 및 동적 방식으로 이루어졌던 소스코드 취약점 탐지에 비해 취약점을 상대적으로 짧은 시간에 탐지할 수 있는 장점이 존재합니다.
이중 언어모델을 이용한 소스코드 취약점 탐지는 취약점이 내포된 코드와 해당 취약점으로 이루어진 데이터 쌍을 데이터셋으로 구축하여 미세조정 학습 진행을 통해 코드 내의 취약점을 자동으로 탐지하는 모델을 구축하게 됩니다. 모델을 구축하는 과정에서 텍스트 분류에서 주로 사용되는 구별형 모델과 현재 LLM의 발전과 더불어 각광받는 생성형 모델을 이용한 소스코드 취약점 탐지의 과정이 조금 다르게 진행되기 때문에 구별형과 생성형으로 나누어 언어모델 기반의 소스코드 취약점 탐지 기술에 대해 알아보도록 하겠습니다.
<구별형 언어모델을 활용한 소스코드 취약점 탐지>
구별형 언어모델이란, 언어모델 중 입력된 언어를 바탕으로 분류 혹은 회귀 분석에 사용되는 모델로써 대표적인 모델로는 앞서도 소개한 BERT모델이 있습니다. 그럼 구별형 언어모델은 어떤 방식을 통해 소스코드 내 취약점을 탐지하게 될까요?
구별형 언어모델을 이용한 소스코드 내 취약점 탐지는 크게 학습 단계와 탐지 단계로 나눌 수 있습니다. 학습 단계에서는 구별형 언어모델의 미세조정 학습을 통해 모델이 소스코드 내의 취약점을 탐지할 수 있도록 구축하고, 탐지 단계에서는 사용자가 입력한 코드의 취약점을 탐지하게 됩니다.
그림 2는 구별형 언어모델을 이용한 소스코드 취약점 탐지 연구 중 하나인 VulDeBERT의 학습 및 탐지 과정입니다. VulDeBERT는 학습 단계에서 코드의 전처리를 수행하고, 이를 BERT에 학습하는 학습 단계와 이후 구축된 모델을 기반으로 사용자가 제시한 코드의 취약점을 탐지하는 탐지 단계로 진행되어 사용자에게 예상되는 취약점을 제시하게 됩니다. 그럼, 학습 단계와 탐지 단계는 구체적으로 어떻게 진행되는지 알아보도록 하겠습니다.
ㆍ학습 단계
소스코드 내 취약점 탐지 모델의 구축을 위해 우선으로 취약점과 해당 취약점이 내포된 코드 데이터를 수집해야 합니다. 이러한 데이터들은 주로 공개된 코드와 취약점을 짝지어 수집하게 되는데, 최근에는 취약점과 그 코드를 수집하여 데이터셋으로 제공되는 BigVul, CVEfixes를 이용한 학습이 주를 이루고 있습니다. 이후 그림 3에서 볼 수 있듯, ① 수집된 코드를 학습에 이용하기 위해 우선 모델이 코드를 이해하기에 용이하도록 불필요한 부분을 제거하거나, 취약점과 연관 없는 변수는 단순 토큰으로 대체하는 전처리 과정을 진행합니다. 이렇게 전처리된 코드는 추상적 표현으로 변환하고 그중 다소 의미가 불명확한 표현을 제거하여 최종적으로 모델 학습에 이용할 코드 데이터를 구축하게 됩니다.
이렇게 구축된 취약점과 코드 데이터는 ② 모델의 학습을 위해 임베딩(Embedding) 과정을 통해 벡터로 변환하게 되는데, 이 과정 중 모델이 코드를 쉽게 이해할 수 있도록 벡터 시작과 끝에 특정한 토큰을 추가합니다. 이렇게 임베딩 과정까지 진행한 데이터는 비로소 ➂ 코드를 이해하는 언어모델의 미세학습 과정을 통해 모델을 구축하며, ➃ 별도의 평가데이터를 이용하여 모델의 성능을 평가합니다. 이후, ➄ 보완이 필요할 경우 추가적인 전처리와 학습 방법 개선을 통해 최적화된 모델을 구축하게 됩니다.
ㆍ탐지 단계
그림 4에서 볼 수 있듯, 탐지 단계에서는 학습 단계에서 구축된 모델은 실제 사용자가 제시한 코드에 대한 취약점 탐지를 수행할 수 있습니다. 탐지 단계에서는 실제 사용자가 제시한 코드를 입력받아 앞서 ➀ 학습 단계에서 진행하였던 동일한 전처리 및 ➁ 임베딩 과정을 거친 후, ➂ 모델에 탐지 과정을 통해 취약점 유무 및 예상되는 취약점을 도출하게 됩니다.
위 과정을 통한 소스코드 내의 취약점 탐지의 경우, 언어모델이 취약점 탐지에 중요한 코드 구조를 놓칠 수 있다는 문제점이 존재합니다[6]. 이 문제를 해결하기 위해 기존 언어모델의 코드 구조를 나타내는 Program Dependence Graph(PDG)와 Abstract Syntax Tree(AST)를 참조하여 코드를 벡터로 변환 시 코드 내 변수별 관계성을 표현하는 방법이 제안되었습니다[7]. 이후 이에 기반한 모델의 사전학습 및 미세조정 학습을 통해 언어모델은 코드 요소별 연관성을 학습하고 이에 기반한 정확한 취약점 부분을 예상할 수 있게 됩니다.
<생성형 언어모델을 활용한 소스코드 취약점 탐지>
GPT와 같은 생성형 모델은 앞서 설명한 바와 같이 앞선 단어를 통해 이후 단어를 생성하는 모델입니다. 따라서 취약점 탐지에 생성형 언어모델을 적용하기 위해서는 구별형 언어모델과 다른 미세조정 학습이 진행됩니다.
구별형 언어모델과 다르게 생성형 언어모델에서는 전처리를 마친 데이터를 학습 단계 전 추가로 사전에 입력과 모델의 답변 형태를 지정하여 모델이 원하는 답변을 유도하는 기법인 프롬프트 튜닝(Prompt-tuning)을 진행하게 됩니다. 구별형 언어모델의 경우 코드를 입력받아 출력값으로 레이블을 제시하도록 설계되었지만, 생성형 언어모델의 경우 입력과 출력이 모두 언어의 형태가 되어야 합니다. 따라서 사용자의 목적에 적합한 답변을 얻을 수 있도록 프롬프트 튜닝을 통해 코드를 질의하고 취약점을 응답하는 형식의 데이터 쌍을 만드는 과정을 진행하게 됩니다. 이후, 학습 단계를 거친 모델은 탐지 단계를 통해 사용자가 입력한 코드에 대한 예상되는 취약점을 답변으로 제시할 수 있습니다.
ㆍ프롬프트 튜닝
프롬프트 튜닝은 모델의 참고 코드 제공 여부에 따라 Zero-shot Prompting과 Few-shot Prompting으로 나뉩니다. Zero-shot Prompting의 경우에는 사용자가 제시한 데이터만을 이용하여 질문 형태로 변환하여 학습에 사용되게 되며, Few-shot Prompting의 경우, 별도로 구축된 데이터셋에서 제시된 코드와 유사한 취약점이 내포된 코드와 그렇지 않은 코드를 포함한 질문 형태로 변환하여 모델이 학습 및 탐지 시 참고할 수 있도록 합니다[6].
탐지 대상 코드와 유사한 코드를 모델에 제공하기 위해 등장 빈도수를 계산하여 유사한 코드를 추출하는 TF-IDF 또는 TF-IDF의 계산 방식에 보정기법을 더한 BM25를 사용할 수 있습니다. 하지만 이러한 방법은 코드 내 토큰의 반복에 기반하여 유사한 코드를 제시하기 때문에 의미상으로 유사한 코드를 제시하지 못한다는 한계점이 존재합니다[6]. 이에 따라 최근 연구에서 제시된 방법으로는 CodeBERT를 이용한 임베딩 과정을 통해 제공된 코드와 별도의 데이터셋 내의 의미적으로도 유사한 취약점이 내포된 코드와 그렇지 않은 코드를 모델에 제공하여 취약점 탐지를 수행할 수 있습니다[8].
ㆍ학습 단계
이렇게 프롬프트 튜닝을 마친 데이터는 앞선 구별형 언어모델에서의 소스코드 내 취약점 탐지와 유사한 학습 단계를 진행하게 됩니다. 학습 단계에서는 프롬프트 튜닝을 거친 데이터의 ➀ 임베딩 과정 후 ➁ 미세조정 학습을 진행하게 되며, 모델의 성능 평가를 위해 ➂ 모델의 답변에서 취약점을 추출하는 후처리 과정을 통해 성능 평가를 진행하게 됩니다. 이후 모델의 성능에 따라 추가적인 전처리 및 학습 방법 변경 후 재학습 혹은 탐지 모델을 구축하게 됩니다.
ㆍ탐지 단계
사용자가 제시하는 코드에 대한 취약점 탐지를 수행하는 탐지 단계에서도 앞서 학습 단계 전 진행하였던 ➀ 프롬프트 튜닝과 동일한 과정을 수행합니다. 이를 통해 사용자가 제시한 코드를 질문 형태로 변환한 뒤, ➁ 임베딩 과정을 거쳐 ➂ 생성형 모델에 질의하게 됩니다. 생성형 모델은 질의 받은 데이터를 기반으로 프롬프트 튜닝에서 지정하였던 답변의 형태에 맞게 의심되는 취약점을 사용자에게 제시합니다.
생성형 언어모델 기반의 소스코드 취약점 탐지의 경우 구별형 언어모델과 비교하였을 때 프롬프트 튜닝의 과정이 추가되지만, 생성형 언어모델을 사용함으로써의 이점 또한 존재합니다. 특히 생성형 모델은 대용량의 사전학습 단계를 진행한 모델이 사용되기 때문에 구별형 언어모델보다 유연한 답변을 얻을 수 있다는 장점이 있습니다. 이는 실제 모델의 사용 단계에서 사용자가 제시한 입력값이 기존의 데이터와 다른 유형임에도 불구하고 모델이 이를 이해하여 사용자에게 답변을 제공할 수 있게 되는 특징이 있습니다.
글을 마치며
본 글에서는 코드를 이해하는 언어모델을 이용한 작업에 대해 알아보고 그중 하나인 소스코드 취약점 탐지 기술에 대해 살펴보았습니다. 소스코드 취약점 탐지 기술은 소스코드 내의 취약점을 사전에 탐지함으로써 외부의 공격을 방지하는 목적으로 연구되고 있으며, 언어모델의 발전과 더불어 코드를 이해하는 언어모델을 이용한 연구가 활발히 진행되고 있습니다. 본 글에서는 구별형 언어모델과 생성형 언어모델로 나누어 그 방법에 대해 알아보고 더 나아가 생성형 모델을 사용하였을 때의 장점에 대해 살펴보았습니다. 이러한 기술은 지속적으로 증가하는 소스코드 취약점 대상 사이버공격에 대응하기 위한 방법으로 널리 활용되길 기대하며, 다음 포스팅에서는 소스코드 취약점 탐지에서 나아가 생성형 인공지능을 이용한 취약점 보정 코드를 생성하는 연구로 다시 찾아뵙도록 하겠습니다.
참고문헌
[1] 정보통신기획평가원, “초거대 언어 모델과 설명 가능한 인공지능 연구 동향”, 주간기술동향 2112호, 2023
[2] Ashish Vaswani, et al, “Attention Is All You Need”, Advances in Neural Information Processing Systems 30, 2017
[3] Ziyin Zhang, et al, “Unifying the Perspectives of NLP and Software Engineering: A Survey on Language Models for Code”, arXiv, 2023
[4] Qurat Ul Ain, et al, “A Systematic Review on Code Clone Detection”. IEEE, 2019
[5] Soolin Kim, et al, “VulDeBERT: A Vulnerability Detection System Using BERT”, IEEE Symposium on Software Reliability Engineering Workshops (ISSREW), 2022
[6] Xin Zhou, et al, “Large Language Model for Vulnerability Detection and Repair: Literature Review and the Road Ahead”, arXiv, 2024
[7] Zhongxin Liu, et al, “Pre-training by Predicting Program Dependencies for Vulnerability Analysis Tasks”, ICSE, 2024
[8] Xin Shou, et al, “Large Language Model for Vulnerability Detection: Emergin Results and Future Directions”, ICSE-NIER, 2024
KAIST 사이버보안연구센터 AI 보안팀 연구원으로 AI 및 LLM을 이용한 보안 연구를 수행하고 있다.