[2025-1] 김학선 - Code Security Vulnerability Repair Using Reinforcement Learning with Large Language Models

https://arxiv.org/abs/2401.07031

Code Security Vulnerability Repair Using Reinforcement Learning with Large Language Models

Introduction

최근 LLM(Large Language Model)의 발전으로 개발자들에게 접근성과 생산성을 높이는데 기여했다. 하지만 LLM이 생성하는 코드의 약 40%가 보안 기준을 충족하지 못하는 것으로 나타났다. 더 큰 문제는 생성형 AI 모델의 학습 데이터가 보안이 취약한 공개 저장소(public repository)에서 가져온 코드에 기반하고 있다는 것이다. 즉, 보안에 초점을 맞춘 pre-training이 이루어지지 않는 한, 생성형 AI 모델은 보안이 취약한 코드를 지속적으로 확산시키며 악순환을 초래할 가능성이 크다.

 

다음 Figure 1에 있는 보안이 강화된 실제 코드인 Code 1과 모델이 보안 수정을 시도한 코드인 Code 2를 보면 보안 수정이 적용된 코드가 단 몇 줄만 추가되었다. 이 경우, 원본 코드와 생성 코드 간 "Cross-Entropy Loss는 코드의 단순 유사성만을 평가할 뿐, 코드의 의미적 및 구문적 정확성을 제대로 평가하지 못한다"는 것을 보여준다.

이를 해결하고 보안이 강화된 코드 생성을 위해, RL(Reinforcement Learning) 기반의 SecureCode라는 새로운 방법을 제안한다. 이 방법은 CodeGen2를 기반으로 RL을 적용한 모델로, 구문적 및 의미적 보상 메커니즘을 결합하여 취약한 코드를 보안이 강화된 형태로 수정하며 기능적으로 올바른 코드를 생성할 수 있도록 설계되었다.

Methodology

코드에서 변수나 함수 호출의 맥락은 여러 줄 또는 페이지에 걸쳐 있을 수 있는데, 이러한 의존성을 포착하는 능력은 정확한 코드 생성이나 수정을 위해 필수적이다. Casual-Decoder 모델인 CodeGen2가 이를 효과적이고 효율적으로 활용할 수 있다. 추가적으로 Casual-Decoder 모델은 일반적인 Encoder-Decoder 와는 다르게 Encoder 계층을 제거하여 빠른 추론 시간과 적은 메모리 사용이라는 특징을 가진다.

 

시스템의 전체 구조는 Figure 2에서 볼 수 있다. 처음엔 취약한 입력 함수 $f_\mathrm{vul}$과 수정된 출력 함수 $f_\mathrm{rep}$을 토큰 시퀀스($t_1, t_2, \cdots t_p \in T$)와 ($y_1, y_2, \cdots y_q \in Y$)로 변환한다. 그다음 두 시퀀스를 결합하여 고유한 시퀀스 $w_1, w_2, \cdots w_\{p+q\}=(t_1,\cdots,t_p,\$,y_1,\cdots,y_q$)를 만든다. 여기서 $는 취약한 코드와 수정된 코드 사이를 구분하는 특별한 토큰이다. p는 입력 토큰 수이고, q는 출력 토큰의 수이다.

Reinforcement Learning for Functional Repair

SecureCode의 Casual-Decoder 아키텍처에서는 모델이 다음 코드 토큰을 예측하도록 강제되며, 출력 생성 중 앞으로 나올 토큰을 무시할 수 없다. 모델은 추론 중에 입력 시퀀스 $t_i \in T$를 제공받고, 이를 기반으로 자기 회귀적 방식으로 출력 $y_i$를 생성한다. 따라서, code-to-code task를 위해 RL으로 Fine-Tuning(FT)된다. 이때, 제안되는 RL 기법은 구문적 및 의미적 보상 함수와 Proximal Policy Optimization(PPO) 알고리즘을 결합한다.

 

입력 함수 또는 취약한 코드를 $f_\mathrm{vul}$, 모델에서 생성된 수정된 출력을 $\hat{f_\mathrm{rep}}$라고 하며, Ground Thruth Repair 함수는 $f_\mathrm{rep}$이다. 출력 시퀀스를 $w_1^r,w_2^r,\cdots,w_k^r$로 정의하며, 여기서 k는 생성된 수정 코드의 총 출력 토큰 수이다. 이때, 토큰 단위의 코드 생성 과정을 결정론적 콘텍스트 마르코프 결정 프로세스(deterministic Contextual Markov Decision Process)로 정의한다. 이 프로세스는 취약한 코드의 이전 토큰에서만 관찰 가능한 맥락을 제공한다. 이때 수정된 코드 시퀀스는 k번째 토큰 생성 시 취약 함수 $f_\mathrm{vul}$의 이전 k-1개의 입력 토큰으로부터의 정책($\pi(.|w^r:k-1,t)$)에 의해 정의된다.

Policy Optimization:

보안 조치를 추가하며 코드 기능을 유지하기 위해 누적된 구문적 및 의미적 보상을 최대화하는 최적의 정책을 찾는 것이 RL의 목표이다. 이를 충족하기 위해, 구문적 코드 평가 기법인 CodeBLEU를 보상 값으로 도입하여, 모델이 생성한 수정된 코드 줄과 정답 코드 사이의 구문적 유사성을 평가한다. 또한, 생성된 코드가 기능적으로 올바른지 확인하기 위해, BERTScore를 의미적 보상값으로 사용한다. 이는 입력된 취약 코드와 생성된 수정 코드 간의 의미적 유사성을 정량화하는 역할을 한다.

 

모델에서 생성한 코드가 $\hat{f_\mathrm{rep}}$, 원래 취약한 코드가 $f_\mathrm{vul}$, 정답 코드가 $f_\mathrm{rep}$일때, 정책 최적화 함수 $L(r_\theta)$를 다음과 같이 계산한다.

$$L(r_\theta)=\log(\sigma(r_\theta(f_\mathrm{vul},f_\mathrm{rep})-r_\theta(f_\mathrm{vul},\hat{f_\mathrm{rep}})))$$

• $r_\theta(f_\mathrm{vul},\hat{f_\mathrm{rep}})$, $r_\theta(f_\mathrm{vul},f_\mathrm{rep})$: 취약 코드와 각 수정 코드에 대한 보상 모델의 스칼라 출력값을 의미
• $\sigma$: 활성화 함수
• $\theta$: 학습 가능한 모델 파라미터

Reward:

CodeBLEU(구문적 보상)

CodeBLEU는 BLEU 점수의 가중 조합으로 계산되며, 다음 요소들을 포함한다.

• BLEU(B): n-gram 매칭 점수, 생성된 코드와 정답 코드 간 토큰 유사성을 비교
• BLEUweight($B_\mathrm{weight}$): 가중치가 적용된 n-gram 매칭
• AST 매칭($\mathrm{Match}_\mathrm{ast}$): 코드의 구문적 정보를 분석하여 유사성을 평가
• 데이터 흐름 매칭($\mathrm{Match}_\mathrm{df}$): 정답 코드와 생성 코드 간 데이터 흐름 유사성 평가

이를 수식으로 나타내면 다음과 같다.

$$B_\mathrm{Code_{B}}=\alpha.B+\beta.B_\mathrm{weight}+\delta.Match_\mathrm{ast}+\gamma.Match_\mathrm{df}$$

BERTScore(의미적 보상)

BERTScore는 코사인 유사도 기반의 의미적 비교 방법이다. BERT 벡터는 정확한 토큰 일치 대신 소프트 유사도를 계산할 수 있도록 하며, 정답 코드와 생성된 코드의 의미적 유사성을 평가한다.

• 정답 코드의 한 코드 $t_r^i$와 생성된 코드의 한 토큰 $\hat{w}_r^i$ 간 코사인 유사도: $(t_r^i)^T\hat{w}_r^i$
• F1-Score로 표현되는 $R_{\mathrm{BERT}}$: $R_\mathrm{BERT}={1\over|T|}\sum_{t_r^i\inT}\max(t_r^i)^T\hat{w}_r^i$

최종 보상 함수(R)

$$R=R_\mathrm{Code_B}+R_\mathrm{BERT}$$

최종 보상은 구문적 유사성과 의미적 유사성의 합으로 구성되며, 이를 통해 모델이 보안 조치를 추가하면서도 기능적 일관성을 유지하는 코드를 생성하도록 학습할 수 있다.

Experiments

DataSet

VulDeeLocator 데이터셋을 사용하여 C/C++ 프로그램의 취약한 소스 코드와 해당 수정 코드 데이터를 활용한다.

• 코드 출처:
  • National Vulnerability Database (NVD)
  • Software Assurance Reference Dataset (SARD)
• 데이터 구성:
  • 취약 코드 스니펫: 40,382개
  • 비취약 코드 스니펫: 115,157개

Evaluation Metrics

BLEU Score:

모델이 생성한 텍스트를 평가하는 척도로, 0~1 사이의 값을 가진다.

• 점수 해석: 1에 가까울수록 생성된 출력이 원본과 완벽하게 일치하는 n-gram을 포함
• BLEU의 한계: reference 기반 평가라 번역 품질의 모든 요소를 완전히 반영하지 못할 수 있다.
• 최종 BLEU 점수 계산 방식
  1. 문장 단위 n-gram 정밀도 계산
  2. n-gram 길이에 따라 가중치 부여
  3. 기하평균을 사용하여 가중 정밀도 계산
  4. Brevity Penalty를 적용하여 최종 점수 산출

Rouge-L:

두 개의 텍스트 간의 유사성을 측정하는 척도로 0~1 사이의 값을 가진다.

• 점수 계산 방식
  • Longest Common Subsequence(LCS, 최장 공통부분수열)을 기반으로 정밀도와 재현율 평가
  • 정밀도 = LCS 길이/생성된 코드 길이
  • 재현율 = LCS 길이/기준 코드 길이
• 활용 및 의미
  • 생성된 코드와 참조 코드 사이 일관성 및 유사성 측정
  • 코드 취약점 분석 알고리즘의 성능 평가에 유용
  • BLEU보다 더 구조적인 유사성을 반영

Experimental Results and Discussion(Table 1)

• RL 기반의 훈련이 SFT 방식보다 성능이 우수하게 나타났다.
• CodeGen2 (7B) RL 모델이 CodeGen2 (7B) SFT 모델보다 BLEU 점수가 0.06, Rouge-L 점수가 0.07 더 높았다.
• 특히 RL 학습 방식이 모든 모델 크기(1B, 3.7B, 7B)에서 SFT 방식보다 일관되게 더 나은 성능을 보였다.
• Rouge-L 점수가 BLEU 점수보다 더 높은 경향을 보였으며, 이는 코드 수정의 평가에서 LCS 기반 평가 방식이 보다 적절하다는 것을 시사한다.

Threats to Validity

현재 Code Repair를 평가할 때 사용하는 BLEU와 Rouge-L 방식은 토큰 기반 비교만을 수행하여 코드의 의미나 보안성을 전혀 반영하지 않는다. 저자들이 수행한 Case Studies에서 RL 모델은 실제로 보안 취약점을 수정한 코드를 잘 생성했지만 기존 메트릭으로 평가했을 때, 성능이 제대로 반영되지 않았다. 이를 통해 코드의 기능성뿐 아니라 코드에 추가된 보안 조치까지 함께 평가할 수 있는 적절한 측정 방법이 필요하다.

Conclusion

여기선 RL 기반의 Code Repair 모델을 제안했다. 이 모델은 Code Repair 과정에서 기능적 및 의미론적 정확성을 평가하기 위해, CodeBLEU(구문적 정확성), BERTScore(의미론적 유사성)을 결합하여 설계했다. 각 모델 별 실험을 통해 SFT 방식에 비해 더 높은 성능을 가지는 것을 확인했으며, 특히 저자들이 수행한 Case studies를 통해 RL 기반 모델이 보안 조치를 추가해야 하는 경우, 기존 방법보다 더 효과적으로 취약점을 수정할 수 있음을 입증했다.