LLM 기반 취약점 탐지를 위한 최적 소수샷 예시 선택 전략

초록

본 논문은 코드 취약점 탐지 작업에서 대형 언어 모델(LLM)의 인‑컨텍스트 학습(ICL) 성능을 높이기 위해 두 가지 직관적인 소수샷 예시 선택 방법을 제안한다. 첫 번째는 모델이 일관되게 오답을 내는 샘플을 우선 선택하는 “실수 기반 학습(LFM)”이며, 두 번째는 질의 프로그램과 의미적으로 가장 유사한 샘플을 k‑최근접 이웃 방식으로 선택하는 “근접 이웃 기반 학습(LFNN)”이다. Python·JavaScript에서는 선택 전략이 성능 향상을 가져왔지만, C·C++에서는 효과가 미미함을 확인하였다.

상세 분석

이 연구는 LLM을 블랙박스 프롬프트 기반으로 활용해 코드 취약점 여부를 이진 분류하는 작업에 초점을 맞춘다. 핵심은 인‑컨텍스트 학습 시 제공되는 소수샷 예시가 모델 성능에 미치는 영향을 정량화하고, 이를 최적화하는 두 알고리즘을 설계한 점이다.

LFM 알고리즘은 전체 라벨링된 데이터셋을 선형 스캔하면서 현재 LLM에게 예측을 요청한다. 예측이 정답과 불일치하면 해당 샘플을 “실수 집합”(S_I)으로, 일치하면 “정답 집합”(S_C)으로 분류한다. 스택 옵션(st)이 활성화되면, 모델이 실수를 할 때마다 해당 샘플을 현재 컨텍스트에 추가해 연속적인 피드백 루프를 만든다. 여러 번(k) 스캔함으로써 LLM의 비결정성을 완화하고, 최종적으로 일관된 실수·정답·그 외(그레이) 샘플을 교차 검증한다. 옵션(opt)에 따라 실수 샘플, 정답 샘플, 혹은 그레이 샘플 중 하나를 선택해 최종 n개의 소수샷을 구성한다.

LFNN 알고리즘은 코드 임베딩 모델(enc)을 이용해 모든 학습 샘플을 벡터화하고, 질의 프로그램을 동일한 인코더로 임베딩한다. 코사인 유사도로 거리 행렬을 계산한 뒤, 상위 n개의 최근접 이웃을 선택한다. 이 방식은 질의와 구조·문맥이 유사한 코드를 직접 제공함으로써 LLM이 해당 도메인 지식을 재활용하도록 유도한다.

두 알고리즘을 조합하는 하이브리드 전략도 실험했으며, LFM이 양성(취약) 클래스로 편향되는 경향을 보이는 반면, LFNN은 보다 균형 잡힌 성능을 제공한다. 특히 Python·JavaScript 데이터셋에서는 LFNN 단독 혹은 LFM+LFNN 결합이 F1 점수를 평균 3~5% 상승시켰다. 반면 C·C++에서는 코드의 저수준 특성(포인터 연산, 메모리 관리 등)과 LLM의 사전 학습 범위 차이로 인해 소수샷 선택이 큰 효과를 내지 못했다.

실험에 사용된 모델은 오픈소스 Qwen2.5‑Coder‑7B‑Instruct, Gemma‑3‑4B‑it, 그리고 클로즈드소스 GPT‑5‑mini이며, 각각의 모델 크기와 사전 학습 데이터 차이가 결과에 미치는 영향을 논의한다. 또한, LFM의 스택 옵션이 반복적인 컨텍스트 축적을 통해 모델의 오류 패턴을 “학습”시키는 메커니즘을 분석하고, LFNN이 의존하는 임베딩 품질이 모델별로 상이함을 지적한다.

결론적으로, 소수샷 예시 선택은 언어 모델 기반 취약점 탐지 성능을 향상시킬 수 있는 실용적 방법이지만, 언어별·특성별 차이를 고려해야 하며, 특히 저수준 언어에서는 재학습이나 파인튜닝이 더 효과적일 수 있다.