[2025-1] 김학선 - LLM-Powered Code Vulnerability Repair with Reinforcement Learning and Semantic Reward

https://arxiv.org/abs/2401.03374

LLM-Powered Code Vulnerability Repair with Reinforcement Learning and Semantic Reward

Abstract

최근 AI 기반 자동화 도구의 등장으로 개발자들이 기능적인 코드를 개발하는 데 있어 효율성을 크게 향상시킨다. 하지만, 이런 도구들이 취약한 코드를 생성하는 원인으로 작용한다는 점이 여전히 중요한 우려 사항으로 남아 있다. 이에 대한 기존 솔루션들은 취약한 코드에 대한 합리적 해결책을 제공하지만, 보안 문제가 반복되지 않도록 개발자들에게 코드 보안에 대해 충분히 설명하고 교육하는데 미흡하다는 문제가 있다.

 

따라서 이 연구에서는 LLM CodeGen2를 기반으로 다목적 코드 취약점 분석 시스템 SecRepair를 소개한다. 이 시스템은 코드의 취약점을 식별하고 수정된 코드를 생성하도록 돕는 동시에 취약점에 대한 완전한 설명을 강화 학습 패러다임을 사용하여 코드 주석으로 제공한다. 추가적으로 LLM에 적합한 취약점 분석을 위한 Instruction Based Dataset을 제안한다.

 

연구 결과는 강화 학습과 의미론적 보상을 결합하면 모델의 성능이 향상되어 코드 취약점을 보다 효과적으로 해결할 수 있는 능력이 강화된다는 점을 강조한다.

Introduction

코드 취약점을 식별하고 취약한 코드에 대한 수리를 개발자에게 제공해야 하는 중대한 필요성 해결을 위해, 우리는 취약점 식별, 수리, 그리고 코드 주석 생성을 포함한 설명을 제공하는 새로운 AI 지원 솔루션 SecRepair를 공개한다.

 

이 시스템은 코드 보안 분석을 위해 설계되고 명시적으로 Fine-Tuning된 LLM인 CodeGen2의 힘을 활용하여 취약한 코드를 식별하고 수정한다. 그리고 개발자에게 포괄적인 보안 분석을 제공하기 위해, 소스 코드 취약점의 보안 문제에 맞춘 다양한 프롬프트를 포함하는 광범위한 Instruction Based Dataset인 InstructVul을 구축했다. 또한, 개발자를 위한 commit 주석을 생성하기 위해, 우리는 의미론적 보상을 활용한 강화 학습 기술을 제공하여 간결하고 적절한 Code commit을 제공한다. Figure 1은 취약점 분석에 대한 전체적인 기여를 보여준다.

Approach

SecRepair: LLM 기반 코드 취약점 수정 시스템

LLM을 기반으로 한 코드 취약점 수정 시스템인 SecRepair는 코드 수정 과정에서 강화 학습을 적용하며, 코드의 의미를 분석하여 보상을 주는 독창적인 의미론적 보상 기법을 활용했다. 이 시스템은 개발자가 수정된 코드를 생성하는 동시에 취약점에 대한 포괄적인 설명을 코드 주석으로 제공할 수 있도록 설계되었다. Figure 2는 제안된 작업의 아키텍처를 보여준다.

 

코드 취약점 수정 및 설명

 

모델 학습 과정

• 모델 학습 시 취약한 코드와 수정된 코드의 쌍을 학습하여 취약점 탐지 및 수정 능력을 향상시킴.
• Instruction Fine-Tuning을 통해 모델이 취약점이 있는 코드의 특징을 내재적으로 이해하게 됨.
• Cross-Entropy Loss를 최적화하여 취약 코드 탐지 및 수정을 효과적으로 수행
• Beam Search를 활용하여 다양한 수정 후보를 생성

Code Repair에서의 문제점 및 해결책

• 코드 샘플이 많을수록 토큰 수가 증가하여 계산량이 많아지는 문제가 발생
• 일반적인 Encoder-Decoder 기반 아키텍처는 양방향 학습이 가능하지만, Code Repair에서는 불필요한 계산을 유발할 수 있어 Encoder 모듈을 제거하고, Casual Decoder를 사용하여 추론 시간을 단축하고 메모리 사용량을 절감했다.

입력-출력 시퀀스 구성

• 입력과 출력을 통합하여 단일 시퀀스로 변환하고, 이를 구분하기 위해 특수 토큰 $\mathrm{$}$사용
    $$\mathrm{input-output\ sequence}:\ t_1, t_2, \dots,t_p\in\mathrm{T}\\\mathrm{token\ sequence}:\ y_1,y_2,\dots,y_q\in\mathrm{Y}\\\mathrm{signle\ sequence}:\ w_1,w_2,\dots,w_{p+q}=(t_1,\dots,t_p,\$,y_1,\dots,y_q)$$
• Autoregressive 방식으로 훈련하여 이전 토큰을 기반으로 다음 토큰을 예측
• 추론 시에는 입력을 받은 후 순차적으로 출력을 생성하며, 미리 보기가 불가능하여 정확한 코드 수정이 가능.

코드 취약점 설명

• 모델을 code-to-text 형식으로 Fine-Tuning하여 개발자 친화적인 코드 설명을 생성한다. 
• Casual Decoder 아키텍처는 모델의 autoregressive 속성으로 최적화되어 이미 생성된 토큰의 맥락을 고려하며 코드 설명의 순차적 생성을 보장한다.

코드 주석을 위한 강화 학습

 

SecRepair를 활용한 코드 설명 및 주석 생성 프로세스

• SecRepair 모델을 활용하여 개발자 친화적인 설명을 생성한 후, 해당 설명을 기반으로 짧고 간결한 코드 주석을 생성
• 코드 주석은 개발자가 코드 수정 위치에 커밋 메시지로 활용할 수 있도록 압축된 형태로 제공됨
• Text-to-text 방식으로 Fine-Tuning하여 짧고 의미 전달이 명확한 설명을 생성하도록 학습

강화 학습 기반 Fine-Tuning 기법

• PPO 알고리즘을 활용하여 보상 함수 기반 Fine-Tuning 수행
• 보상 함수는 의미를 인식하는 방식으로 설계되며, BERTScore를 활용하여 원래 설명과 요약된 커밋 메시지의 의미적 유사성을 평가
• Markov Decision process 기반 모델
  • 입력: 개발자 친화적인 설명($D$)
  • 출력: 짧은 커밋 메시지($D_c$)
  • 생성된 텍스트는 기존 설명보다 훨씬 짧아야 하며, 이전 토큰에서 관찰된 맥락을 기반으로 결정적으로 생성됨
  • 정책($\pi$)은 설명($D$)의 모든 입력 토큰에 대한 확률 분포로 정의되며, 이는 맥락과 상태에 따라 달라진다.
  • $\pi(.|w^c:k-1,t)$: 이전까지 생성된 커밋 메시지와 주어진 코드 설명을 고려하여, 다음 토큰을 예측

정책 최적화

• BERTScore를 활용하여 의미적 유사성을 평가하고, 보상 점수를 생성하는 Metric으로 사용
• 정책 최적화 식: $L(r_\theta)=\log(\sigma(R_\theta(D,D_c)-r_\theta(D,\hat D_c)))$

보상 함수

• BERTScore를 사용하여 의미적 유사성을 평가하며, 코사인 유사도를 활용하여 보상 점수를 계산
• BERT 벡터 표현을 활용한 부드러운 유사도 측정 가능
• BERTScore -> $R_{BERT}={1\over |d_C^i|}\sum_{d_C^i}\in d\max(d_C^i)^T\hat d_C^i$

Experiments and Discussions

평가 지표

• 취약점 수정, 설명 생성, 코드 주석을 위해 생성 모델을 사용하여, 이에 적합한 BLEU, Rouge-L, 인간 평가 점수를 지표로 사용
• 취약점 식별 작업에는 F1, 정밀도(Precision), 재현율(Recall), 정확도(Accuracy)를 사용

결과 및 토론

시스템의 취약점 분석 효과와 능력에 관한 세 가지 연구 질문(RQ)를 다룬다.

• RQ1: 코드 취약점을 자동으로 식별하고 코드를 정확하게 수정할 수 있는가?
• RQ2: 제안된 수정에 대한 코드 취약점을 개발자에게 포괄적으로 설명할 수 있는가?
• RQ3: 프로그램 설명을 최적화하고 요약하여 코드 주석을 생성할 수 있는가?

RQ1 - 코드 취약점을 자동으로 식별하고 코드를 정확하게 수정할 수 있는가?

 

실험 결과는 비슷하거나 더 적은 매개변수를 가진 다른 모델들에 비해 SecRepair 모델의 우수성을 보여준다.

RQ2 - 제안된 수정에 대한 코드 취약점을 개발자에게 포괄적으로 설명할 수 있는가?

 

CodeGen의 7B 변형과 SecRepair 7B 변형 모델을 비교했을 때 전체적으로 점수가 오른 것을 확인할 수 있다.

RQ3 - 프로그램 설명을 최적화하고 요약하여 코드 주석을 생성할 수 있는가?

 

CodeGen의 7B 변형과 SecRepair 모델을 비교했을 때 점수가 향상된 것을 볼 수 있다. 마찬가지로, 강화 학습으로 훈련된 1B, 3.7B, 7B 변형에서도 성능 향상이 관찰된다.

아래는 생성 모델에 영향을 미치는 두 가지 중요한 요소인 Temperature와 Beam Size의 영향을 탐구한 결과, 온도가 높아지다가 최적점(0.5)일때까지 점수가 증가하다가 이후에 감소하는 것을 볼 수 있었다. Beam Size의 경우는 Size가 커질수록 결과가 약간 개선되는 것을 보여주지만 추론 시간이 거의 기하급수적으로 증가하는 것을 볼 수 있어, 빔 값이 커질수록 성능은 좋아지지만, 높은 추론 시간으로 인해 개발자가 실시간으로 사용하기에 불편함을 초래할 수 있다는 것을 확인할 수 있었다.

Temperature	Beam Size

Conclusion

• CodeGen2-7B 기반의 SecRepair 모델을 훈련시키기 위해 Instruction Based Dataset을 도입하여 취약점 식별, 수정, 설명 문제를 해결했으며, 의미적 보상을 활용한 강화 학습으로 코드 주석 생성을 최적화했다.
• 모델 성능은 BLEU와 Rouge-L 점수로 평가했으며, 인간 평가를 통해 개발자 친화적인 설명과 주석 생성 능력을 검증했다.
• 절제 연구를 통해 Temperature와 Beam Size의 영향을 분석한 결과, SecRepair는 취약점 식별, 수정, 설명에서 우수한 성능을 입증했다.