보안 TOP 학회 중 하나인 ACM CCS 2023 에 발표된
"Large Language Models for Code: Security Hardening and Adversarial Testing" 라는 논문에 대해 Review 를 하고자 한다.
해당 논문은 코드 생성 LLM 을 prefix-tuning 하여 더욱 보안적으로 안전한 (혹은 더 취약한) 코드를 생성하게 하는 방법을 제안했다.
아래 링크로 가면 저자가 직접 설명해주는 영상도 있다.
좀 더 자세히 이해해보자.
Background
ChatGPT 의 등장으로 거대 언어 모델 (Large Language Models, LLM) 에 대한 관심은 더욱 증가하고 있다.
예술, 의료, 법률 등 다양한 분야에서 활용되고 있는 LLM은
개발자들의 영역이었던 코드 작성 분야에도 활발히 역할이 확장되고 있다.
특히, 깃허브가 Copilot [2] 라는 자동 코드 완성 AI 서비스를 출시하면서
코딩에서 LLM을 하나의 도구로 자연스럽게 활용하는 추세가 시작되었다.
하지만, 이와 같은 코드 생성 AI 를 우리가 전적으로 믿어도 되는 것일까?
기능적으로는 그렇다고 쳐도, 이 생성된 코드가 보안적으로 취약하지 않고 안전하다고 믿어도 될까?
본 연구는 그렇지 않다고 말하고 있다.
실제로 [3] 에 따르면, 다양한 보안적 시나리오에 대해서 코파일럿이 생성한 코드 중 40%가 위험한 보안 취약점을 가지고 있었다고 한다.
결국, LLM이 생성한 코드를 그대로 사용했다가 보안적 위험에 노출될 수 있고,
이를 예방하기 위해서는 결국 인간의 노력이 더 들어가게 된다.
LLM의 코드를 믿고 안전하게 활용하기 위해서는 LLM이 더 안전한 코드를 생성할 수 있도록 개선해야 한다.
Controlled Code Generation
저자는 Controlled Code Generation 이라는 새로운 Task 를 정의 하였다.
이것은 기존 자연어 처리(NLP) 분야에서 정의된 controlled text generation 라는 기술에서 파생된 것이다.
controlled text generation 란 어떤 문장의 특성을 조절할 수 있는 자연어 생성 기법을 말한다. 예를 들어서, '슬픈 분위기/기쁜 분위기를 내는 문장으로 바꿔줘'와 같이 특정한 특성을 마음대로 조절해서 문장을 생성하게 하는 것을 말한다.
[4] 같은 경우에는 이런 문제를 간단한 attribute classifiers 를 추가하여 기존 LM 을 추가 training 하지 않고 task 를 달성했다.
본 저자는 code generation LLM 에 controlled code generation 을 수행할 수 있게 하기 위해,
binary security property 인 𝑐 = {sec, vul} 를 추가하였다.
즉, 우리가 LLM 에 input prompt 를 넣을 때, binary property c 를 같이 넣어주는데, c 가 sec (security) 인 경우, LLM 이 더 보안에 안전한 코드를 생성하도록 하고, c 가 vul (vulnerability) 인 경우, LLM 이 오히려 기존보다 더 취약한 코드를 생성하도록 디자인 하였다.
그래서, 이런 property c 에 따라서 달라지는 LLM 의 반응을 어떻게 유도한다는 것일까?
본 저자는 이를 달성하기 위한 솔루션으로 SVEN 이라는 방법을 제안하였다.
Challenges
제안된 솔루션인 SVEN 이 집중한 문제들은 어떤 것들이고 어떤 식으로 해결했을까?
3가지 challenges 들에 대해 알아보자
- 첫번째는, Modularity (모듈성) 이다.
너무 방대하게 큰 LLM 은 re-training 뿐만 아니라 일정 부분만 training 하는 fine-tuning 조차도 매우 많은 자원과 시간이 소모된다.
따라서, 저자는 이런 대거의 training 이 아닌 기존 pre-trained 된 모델에 일정 부분적인 module 인 prefix 를 추가해서, 이 prefix 만 training 하는 prefix-tuning 기법을 활용하였다.
즉, 모델의 기존 weights 들은 바꾸지 않고 property-specific sequences 인 prefixes 를 모델에 추가해서 이 부분만 훈련하였다. - 두번째는, 기존의 기능성을 유지해야 한다는 것이다. (Functional Correctness vs Security Control)
security 쪽으로는 안전한 코드를 생성하는데, 정작 사용자가 원하는 기능을 제대로 못하는 코드를 생성하면 그 코드는 쓸모가 없다.
예로 사용자가 숫자 두개를 더하는 함수를 생성해줘 했는데, 보안적 측면은 완벽하게 방어를 했지만 1+1=10 을 출력하는 코드를 준다면 안된다.
따라서, 저자는 security 측면을 강화 하면서도 최소 원래 LLM 이 가지고 있던 기능성은 유지를 해야했다.
저자는 이를 prefix-tuning 시에 data 를 unchanged code 와 changed code 로 나눠서 loss 계산을 따로하고, 원래의 LLM 의 기능성을 따라가는 loss 를 추가하여 해결했다. - 세번째는, 질좋은 데이터의 부족 이다.
보안 관련 데이터들이 github 와 같은 open-source 에서 keyword matching 과 같은 방식으로 끌어온 것인데, 보안적인 변화가 없는 코드인데 보안관련 데이터로 포함되어 있는 등 여러 데이터 품질 이슈가 있다. 이는 해당 prefix-tuning 의 품질을 저하시킨다.
따라서, 본 저자는 직접 1.6k programs 들에서 직접 수작업을 통해 데이터 품질을 높였다.
Method : SVEN
SVEN 은 controlled code generation task 수행을 위해서 본 페이퍼가 정의한 inference, training (prefix tuning)과 high-qulity data 생성에서의 솔루션을 총괄하여 말한다.
Inference
본 논문이 주로 다룬 CodeGen 이라는 모델은 standard transformer-based autoregressive language models 이다.
기본적인 transformer 기반의 모델이며 자기 회귀 모형 이라는 것인데 이전 토큰의 정보가 다음 출력될 토큰에 영향을 주는 방식을 뜻한다.
그 방식을 수식으로 나타내면 오른쪽 그림처럼, t 이전에 나타난 h (hidden state) 들이 주어지면, 현재 t 에 나타나는 x 값에 대한 확률분포가 나타나고, 이를 전체 t 에 대해 곱한 P(x) 가 전체 출력의 확률 분포가 된다.
이런 inference (for code generation) 과정이 SVEN 에선 다음과 같이 변형된다.
오른쪽 그림처럼, binary property c 가 조건 설정을 위해 추가된다.
즉 이 c도 전체 확률 분포에 영향을 주게 된다. 또한, t<1 인 h<1 을 새롭게 prefix 로 정의한다.
즉 원래는 h<1 이 empty space 였는데 (입력된 토큰이 없으니깐), SVEN 은 h<1 을 실제 기존 hidden state 와 사이즈가 같은 continuous vector 들을 추가한다.
이 h<1 인 continuous vector 들이 prefix 들이고 training 과정에서 이 prefix 들을 task 에 맞게 학습하고 inference 시에는 그 tuning 된 prefix 값을 사용하는 것이다.
이를 그림으로 쉽게 설명하면 다음과 같다.
원래의 구조는 다음처럼 token 이 embedding layer 를 지나 attention block 여러개를 지나 (몇몇 과정을 또 거쳐서) output token 을 출력한다. 이때 나온 output 은 다음 token 이 생성될 때 input 으로 들어가고 이런식의 자기 회귀 구조를 가진다.
sven은 여기서 다음 그림처럼 prefix 가 추가된다.
매 attention block 마다 n 개의 prefix token 들이 같이 계산에 영향을 주는 것이다.
즉 추가되는 총 prefix 의 사이즈는 아래와 같다.
$$params = n*d*l*h*c*2$$
* n : prefix token 갯수
* d : hidden state dimension
* l : attantion block 갯수
* h : multi-head attention 의 head 갯수
* c : binary property c 에 따라 prefix 가 달라지므로 2가지 버전의 prefix 가 필요 (c=2)
* 2 : key & value 에 prefix 가 추가 되므로 곱하기 2
여기서 저자는 추가되는 parameter 갯수가 전체 모델 parameter 의 0.1% 이하가 되는 n 을 사용하였다.
Training
추가된 prefix 들을 tuning 하기 위해 정의된 새로운 loss 는 다음과 같다.
이 loss 는 크게 3가지로 나뉘어 볼 수 있는데,
첫번째는, Loss Terms for Controlling Security 으로
property c 와 함께, security-sensitive region 에 대한 loss 를 구한다.
오른쪽 그림에서 mt 는 기존 data 에서 security 이슈 때문에 바뀐 코드 부분만 loss 계산에 활성화 시킬 수 있도록 하는 mask 이다.
두번째는, Contrastive loss 인데
vul 모드와 sec 모드의 next-token probability 를 대조하는 loss 로, 두 가지 모드를 동시에 최적화하기 위해 쓰인다. sec 와 vul 을 대조적으로 학습하게 함으로써 vul 모드에서는 security 가 더 취약한 코드를 생성할 수 있게 하는 것이다.
세번째는, Loss Term for Preserving Functional Correctness 이고,
prefix 가 있을 때와 없을 때, unchanged code 부분의 similarity 가 비슷하도록 하는 loss 로, 이를 통해 원래 LM 의 functionally correctness 를 보장하기 위해 정의된다. 즉 security 와 무관한 코드 (security 때문에 바뀌지 않는 코드) 부분이 기존 LM 으로 생성한 값과 sven 이 생성한 값의 확률 분포가 비슷하도록 두 값의 KL divergenece 를 작아지게 만든다.
Constucting High Quality Data
본 연구는 prefix tuning 을 위해 CrossVul, Big-Vul, VUDENC 라는 3가지 데이터셋을 기반으로 활용했는데,
이 데이터들은 몇가지 quality 이슈가 있었다.
- 5-(a) 처럼, commit message 의 keyword 로 labeling 을 해서 보안과 관련 없는 코드 변경이 데이터에 포함된 경우
- 5-(b) 처럼, 관련없는 refactoring 과 같은 change 가 포함된 경우
- 5-(c) 처럼, 특정 project 에서만 쓸 수 있는 것을 사용하는 경우
위와 같은, quality issue 는 SVEN 의 성능을 크게 좌우할 수 있기 때문에, 1.6k programs 들에서 manual curation 을 진행해서 SVEN 이 더욱 좋은 성능을 낼 수 있도록 하였다.
Experiments
본 논문은 SVEN 의 성능을 입증하기 위해 매우 다양한 실험을 했는데, 본 글은 제일 메인 실험에 대해서 이야기할 것 이다.
evaluation 위해 쓰인 각각의 시나리오는 위 표에 제시된 것과 같다.
예로 CWE 089 취약점에 대해서 평가하는데, CWE 089 1-py 시나리오의 경우에 prompt 는 데이터베이스에 user 메세지를 넣어줄 것을 요청하는 것이다.
security 실험 성능을 보면, 기본 LM 인 회색보다 전반적으로 c=sec 모드인 초록색이 더 높은 security rate (보안성이 더 높은) 을 나타내는 걸 확인할 수 있다. 반대로 c=vul 인 경우인 빨간색은, 기존 LM 보다 더 취약한 코드를 생성하는 걸 볼 수 있다.
기존의 기능성은 유지가 잘 되느냐를 확인한다면 위 테이블로 볼 수 있는데, 초록색과 빨간색 모드가 기존 LM 과 비교했을 때 pass@k 점수가 크게 차이 나지 않는 걸 확인할 수 있다.
Conclusion
본 논문은 code 생성 LLM 에 특정 attribute 에 따라 코드를 생성할 수 있도록 하는 controlled code generation 를 정의하였고 이를 풀어내기 위해 SVEN 이라는 솔루션을 제안하였다. SVEN 은 binary property c 를 통해 sec 모드, vul 모드를 나누었고 각 모드는 security hardening (보안성 강화) 와 adversarial testing (적대적 검토) 의 목적성을 가지도록 했다. 이를 위해 prefix-tuning 기법을 활용하였고, 3가지 loss 를 정의하여 이 prefix 들을 목적에 맞게 optimization 하였다. 또한, 기존 security code 데이터셋의 문제점을 분석하고 이를 해결한 dataset 을 수작업으로 검토하여 생성하였다. 본 논문는 안전한 code 생성 LLM 을 위한 첫번째 시도이고, 앞으로의 더욱 안전한 code 생성 기술에 큰 기여를 한다는 점에서 매우 훌륭한 논문이라고 할 수 있다.
References
[1] Large Language Models for Code: Security Hardening and Adversarial Testing
[2] https://github.com/features/copilot
[3] Asleep at the Keyboard? Assessing the Security of GitHub Copilot’s Code Contributions
[4] Plug and Play Language Models: A Simple Approach to Controlled Text Generation
[5] CodeGen: An Open Large Language Model for Code with Multi-Turn Program Synthesis
'ML & AI > Paper Review' 카테고리의 다른 글
| [AI Security] Backdoor Attacks 백도어 공격이란 ? (0) | 2024.02.01 |
|---|---|
| [ML/Multi-Agent Systems] Interaction Modeling with Multiplex Attention (0) | 2023.11.20 |
| [ML/Tabular] CTGAN, Modeling Tabular Data using Conditional GAN Review (0) | 2023.08.02 |
| [AI/ViT] Vision Transformer(ViT), 그림으로 쉽게 이해하기 (0) | 2023.03.23 |
| [LLM Training] 거대 언어 모델 병렬 훈련 기법들/Large Language Models Parallel Training (0) | 2023.02.09 |
