Fine-tuning RoBERTa for CVE-to-CWE Classification: A 125M Parameter Model Competitive with LLMs
이 논문은 Claude Sonnet 4.6 으로 정제된 대규모 CVE-CWE 데이터셋을 활용해 1 억 2,500 만 파라미터 규모의 RoBERTa 모델을 미세 조정하여, 64 배 더 큰 LLM 과 경쟁 가능한 성능을 달성하면서도 희귀 취약점 분류에서 기존 TF-IDF 기반보다 현저히 우수한 결과를 보인 것을 제안합니다.
이 논문은 **"작은 두뇌로도 거인들과 경쟁할 수 있다"**는 것을 증명하는 흥미로운 연구입니다. 복잡한 기술 용어 대신, 일상적인 비유를 통해 이 연구의 핵심 내용을 쉽게 설명해 드릴게요.
🕵️♂️ 핵심 이야기: "작은 탐정 vs 거인 AI"
이 연구는 컴퓨터 보안 분야에서 아주 중요한 작업을 수행하는 AI 모델에 대한 이야기입니다.
문제 상황 (CVE 와 CWE):
세상에는 매일 새로운 **컴퓨터 보안 구멍 (CVE)**이 발견됩니다. 마치 건물에 생긴 금이나 잠금장치가 고장 난 것과 같죠.
전문가들은 이 구멍들이 정확히 어떤 **유형의 결함 (CWE)**인지 분류해야 합니다. 예를 들어, "문고리가 부러진 것 (CWE-119)"인지, "열쇠 구멍을 뚫은 것 (CWE-89)"인지 구분하는 거죠.
하지만 기존에 이 작업을 하던 사람 (NVD) 들은 너무 바빠서 때로는 "모든 고장"이라고만 대충 적거나, 서로 다른 사람이 같은 문제를 다르게 분류하는 등 혼란이 많았습니다.
기존의 방식 (거인 AI 들):
최근에는 GPT-4 나 구글의 거대 AI 같은 **엄청나게 큰 두뇌 (수십억~수조 개의 파라미터)**를 이 작업에 썼습니다.
하지만 이 거인들은 전기세 (컴퓨팅 비용) 가 너무 비싸고, 무겁기 때문에 누구나 쉽게 쓸 수 없습니다. 게다가 성능이 기대만큼 완벽하지는 않았습니다.
이 연구의 해결책 (작은 RoBERTa):
연구자들은 **"작지만 똑똑한 AI (RoBERTa, 1 억 2,500 만 개의 파라미터)"**를 만들어냈습니다. 이 모델은 거인 AI 들보다 64 배나 가볍습니다. (비유하자면, 거인 AI 가 '전체 도서관'을 외우고 있다면, 이 작은 AI 는 '필요한 책 1 권'만 완벽하게 외운 셈입니다.)
🛠️ 어떻게 이렇게 작은 AI 가 거인들과 경쟁할 수 있었을까요?
이 작은 AI 가 성공한 데에는 두 가지 **'비밀 무기'**가 있었습니다.
1. 최고의 교재 (AI 가 다듬은 학습 자료)
기존 데이터는 사람이 대충 쓴 것이라 오류가 많았습니다. 연구자들은 Claude Sonnet 4.6 이라는 최신 AI를 고용해 23 만 4 천 개의 보안 구멍 데이터를 꼼꼼하게 다시 검사하고 정리했습니다.
비유: 마치 시험을 치기 전에, 엉망으로 쓴 문제를 **수석 강사 (Claude AI)**가 하나하나 고쳐서 완벽한 모의고사를 만든 것과 같습니다. 이 '고급 교재' 덕분에 작은 AI 도 빠르게 성장할 수 있었습니다.
2. 두 단계 학습법 (우선 기초, 그다음 심화)
연구자들은 AI 를 한 번에 다 가르치지 않고 두 단계로 나누어 훈련시켰습니다.
1 단계 (기초 다지기): AI 의 '기억력'을 고정해두고, 오직 '답을 고르는 능력'만 먼저 훈련시킵니다. (기초 체력을 다지는 것)
2 단계 (심화 훈련): 이제 AI 의 모든 기억력을 풀어서, 전체 내용을 다시 정교하게 다듬습니다.
비유: 요리사를 키울 때, 먼저 **칼질과 불 조절 (기초)**만 익히고, 그다음에 **전체 레시피 (심화)**를 가르치는 것과 같습니다. 이렇게 하면 AI 가 처음 배운 좋은 습관을 잊어버리지 않으면서도 실력을 극대화할 수 있습니다.
🏆 결과는 어땠나요?
이 작은 AI 는 거대한 경쟁자들 (Cisco 의 80 억 파라미터 모델, GPT-4 등) 과 비슷한 점수를 받았습니다.
성적표: 거인 AI 들이 75% 정도 맞췄다면, 이 작은 AI 도 **75.6%**를 맞췄습니다. 통계적으로 차이가 없는 수준입니다.
효율성: 성능은 비슷하지만, 이 작은 AI 는 64 배나 적은 자원으로 작동합니다. 이는 마치 스마트폰으로 거대 서버와 같은 연산을 해낸 것과 같습니다.
⚠️ 하지만 약간의 함정도 있었습니다 (중요한 발견)
연구자들은 흥미로운 사실을 발견했습니다. 바로 '분류의 세밀함' 문제입니다.
상황: AI 는 "이 구멍은 '스택 버퍼 오버플로우 (CWE-121)'라는 아주 구체적인 문제야!"라고 정확히 맞췄는데, 기존 데이터 (NVD) 는 "아니, 그냥 '버퍼 오버플로우 (CWE-119)'라고만 적어놨어"라고 했을 때, AI 가 틀린 것으로 판정받았습니다.
비유: 의사가 환자에게 "당신은 '감기'가 아니라 '인플루엔자 A'입니다"라고 정확한 진단을 내렸는데, 기록에는 그냥 '감기'라고만 적혀 있어서 의사가 틀린 사람으로 취급받는 꼴입니다.
해결: 연구자들은 "AI 가 더 구체적인 정답을 낸 것은 오히려 더 훌륭하다"고 주장하며, 평가 기준을 다시 봐야 한다고 제안했습니다.
💡 결론: 왜 이 연구가 중요한가요?
이 논문은 **"무조건 큰 AI 가 좋은 것은 아니다"**라고 말합니다.
접근성: 누구나 쉽게 구할 수 있는 작은 AI 로도, 거대 기업들의 비싼 거인 AI 와 경쟁할 수 있는 성능을 낼 수 있습니다.
데이터의 힘: 모델의 크기보다 얼마나 깨끗하고 정확한 데이터 (교재) 를 사용했는지가 더 중요합니다.
실용성: 이 작은 모델은 보안 전문가들이 매일 수천 개의 보안 구멍을 분류할 때, 빠르고 저렴하게 도울 수 있는 실용적인 도구가 될 것입니다.
요약하자면, 이 연구는 **"잘 다듬어진 작은 두뇌가, 엉성한 거인 두뇌보다 더 똑똑할 수 있다"**는 것을 증명해낸 것입니다.
논문 요약: RoBERTa 미세 조정을 통한 CVE-to-CWE 분류 (125M 파라미터 모델이 LLM 과 경쟁하는 성과)
1. 연구 배경 및 문제 정의 (Problem)
핵심 과제: Common Vulnerabilities and Exposures (CVE) 설명을 Common Weakness Enumeration (CWE) 범주로 매핑하는 작업은 취약점 분류, 우선순위 결정 및 수정에 필수적입니다.
현재의 한계:
NVD 라벨의 노이즈: 국가 취약점 데이터베이스 (NVD) 의 CWE 라벨은 종종 너무 일반적이거나 (예: 포괄적인 'CWE-20'), 주석자 간 일관성이 부족합니다.
LLM 의 비용과 성능: 최근 연구 (CTI-Bench 등) 에 따르면 GPT-4 나 LLaMA3 와 같은 대규모 언어 모델 (LLM) 은 이 작업에서 어느 정도 성과를 보이지만, 8B(80 억) 파라미터 이상의 거대 모델이 필요하며 비용이 높습니다. 반면, 작은 모델 (LLaMA3-8B) 은 성능이 현저히 떨어집니다 (44.7% 정확도).
목표: 소규모 모델 (125M 파라미터) 로도 대규모 LLM 과 경쟁할 수 있는 고품질 CVE-to-CWE 분류기를 구축하는 것.
2. 방법론 (Methodology)
가. 데이터 구축 (Dataset Construction)
데이터 소스: 1999 년부터 2026 년까지의 NVD 에서 336,777 개의 CVE 를 수집하여 전처리 후 318,979 개를 사용.
AI 기반 라벨 정제 (Label Refinement):
NVD 라벨의 품질 문제를 해결하기 위해 Claude Sonnet 4.6을 사용하여 모든 CVE 설명을 다시 라벨링했습니다 (약 $395 비용).
205 개의 CWE 카테고리로 분류.
결과: 기존 NVD 라벨과 AI 라벨의 일치율은 73.1% (정확 일치), 계층 구조 (부모 - 자식) 를 고려하면 84.5% 로 상승. 불일치 사례의 대부분은 NVD 가 너무 일반적인 라벨을 사용한 경우였으며, AI 가 더 구체적인 하위 클래스를 할당하는 경우가 많았습니다.
데이터 분할:
학습 데이터 (234,770 개): AI 라벨을 사용 (NVD 와 일치/불일치 모두 포함).
검증/테스트 데이터: NVD 와 AI 라벨이 일치하는 고신뢰도 샘플만 선별하여 편향을 줄이고 명확한 사례에 집중 (테스트 세트 27,780 개).
외부 벤치마크: CTI-Bench(2,000 개 CVE) 는 학습 데이터에서 완전히 제거 (Decontamination) 하여 외부 평가의 순수성을 보장.
추가 주석: MITRE ATT&CK 기법 매핑도 수행하여 다중 태스크 분석을 위한 리소스로 제공.
나. 모델 및 학습 전략 (Model & Training)
모델 아키텍처:RoBERTa-base (125M 파라미터) 에 205 클래스 분류 헤드를 추가.
입력 형식: "CVE Description: {text}" (최대 384 토큰).
2 단계 미세 조정 (Two-Phase Fine-tuning):
Phase 1 (Warm-up, 4 epoch): 임베딩 레이어와 하위 8 개 레이어를 고정 (Freeze) 하고 상위 4 개 레이어와 분류 헤드만 학습. (학습률 10−4)
Phase 2 (Full Fine-tuning, 9 epoch): 모든 파라미터의 잠금을 해제하고 전체 학습. (학습률 2×10−5)
효과: 사전 학습된 표현의 망각 (Catastrophic forgetting) 을 방지하면서 분류 헤드가 좋은 초기화 상태를 갖도록 함.
3. 주요 결과 (Key Results)
가. 내부 평가 (Internal Evaluation)
데이터: NVD 와 AI 라벨이 일치하는 27,780 개의 테스트 샘플.
성능:
Top-1 정확도:87.4% (TF-IDF 베이스라인 84.9% 대비 우위).
Macro F1:60.7% (TF-IDF 대비 15.5%p 향상). 이는 희귀한 CWE 클래스에서 RoBERTa 모델이 훨씬 더 강력함을 의미.
Top-3 정확도: 94.7%.
나. 외부 벤치마크 (CTI-Bench)
데이터: 2,000 개의 CVE (학습 데이터와 겹치지 않음). 엄격한 정확한 CWE ID 일치 기준.
성능 비교:
본 모델 (125M):75.6% (95% CI: 72.8–78.2%).
Cisco Foundation-Sec-8B-Reasoning (8B): 75.3%.
GPT-4: 72.0%.
LLaMA3-8B: 44.7%.
의미: 125M 파라미터 모델이 80 억 파라미터 (8B) 모델인 Cisco 의 모델과 통계적으로 유의미한 차이가 없는 동등한 성능을 보임. 파라미터 효율성이 64 배 더 높음.
다. 분석 및 통찰 (Analysis)
계층 구조 불일치 (Granularity Gap): 모델이 구체적인 하위 CWE(예: CWE-121) 를 예측한 반면, 벤치마크 지상 진실 (Ground Truth) 이 일반적인 상위 CWE(예: CWE-119) 를 가진 경우를 '오류'로 간주하는 문제가 발생.
계층 구조를 고려한 평가 (Parent-Child 일치 허용) 를 적용하면 성능이 **86.5%**까지 상승하여 내부 테스트 세트와 격차가 거의 사라짐.
라벨 품질 검증: 100 개의 무작위 불일치 사례를 수동 검토한 결과, **72%**에서 AI(Sonnet) 라벨이 NVD 라벨보다 더 구체적이고 정확함이 확인됨.
4. 주요 기여 (Contributions)
대규모 고품질 데이터셋 공개: 29 만 개 이상의 샘플로 구성된 CVE-to-CWE 데이터셋 (AI 정제 라벨 및 MITRE ATT&CK 매핑 포함) 을 Hugging Face 에 공개.
효율적인 학습 전략: 2 단계 미세 조정 방식을 통해 소규모 모델의 성능을 극대화하는 방법론 제시.
경쟁력 있는 오픈 웨이트 모델: 64 배 더 작은 파라미터 수로 일반 목적의 8B LLM 과 경쟁하는 CVE 분류 모델 공개.
벤치마크 한계 지적: 현재 벤치마크가 CWE 계층 구조의 세분성 (Granularity) 차이를 고려하지 않아 모델 성능을 과소평가할 수 있음을 지적하고 해결 방안 제시.
5. 의의 및 결론 (Significance)
비용 효율성: 거대 LLM 을 사용할 필요 없이, 125M 파라미터의 경량화된 RoBERTa 모델만으로도 산업 수준의 CVE-to-CWE 분류가 가능함을 입증.
실용성: 취약점 대응 시스템에서 실시간 처리가 필요한 환경에 적합하며, 오픈 소스 리소스 (모델, 데이터, 코드) 를 공개하여 연구 및 산업계 활용도를 높임.
향후 방향: 데이터 라벨링의 순환성 문제 (AI 라벨로 학습하고 AI 라벨과 일치하는지 평가) 에 대한 우려를 수동 검증을 통해 부분적으로 해소했으나, 더 전문적인 주석자와 다양한 아키텍처 비교는 향후 과제로 남김.
요약하자면, 이 연구는 고품질의 AI 정제 데이터와 전략적인 미세 조정 기법을 통해, 거대 LLM 과 견줄 만한 성능을 내면서도 훨씬 가볍고 효율적인 CVE 분류 모델을 성공적으로 개발하고 공개했다는 점에서 의의가 큽니다.