Safety Alignment Should Be Made More Than Just a Few Tokens Deep 논문 리뷰

Safety Alignment Should Be Made More Than Just a Few Tokens Deep

---

Summary

본 논문은 현재의 LLM에 대한 안전 정렬(safety alignment) 전략에 대해서 출력의 처음 몇 개 토큰의 생성 분포만 바꾸면 된다는 것을 지적하며, 이러한 취약성을 해결하기 위해 1. 데이터 증강을 통한 alignment 깊이 증가, 2. 제약 최적화 목적 함수에 따른 초기 토큰 업데이트 제약을 제시한다.

---

Experiments

본 논문은 LLM의 안전 정렬이 처음 몇 개의 토큰의 분포로 대부분이 형성되어 있다는 것을 보인다. 이에 대한 근거를 확인하자:

우선 첫 번째 실험 결과로, SFT까지만 적용되었고 alignment가 적용되지 않은 비정렬(base) 모델도 초기 응답에 거절 접두사(refusal prefix)를 강제로 부여하면, 유해한 답변을 생성할 확률이 급격히 감소함을 확인할 수 있다. 
즉, safety mode라는 것은 단순히 거절 접두사를 출력하도록 분포를 억제하는 지름길일 뿐임을 알 수 있다.  

특히 Figure 1에서는 정렬된 모델과 비정렬 모델 각각이 유해 요청에 대해 보이는 token 위치별 KL divergence를 나타낸다. 
이를 통해 alignment가 초기 몇 개 토큰에서만 집중적으로 학습되었고, 이후 응답 부분에서는 거의 penalty가 적용되지 않았음을 확인할 수 있다. 
이는 곧 현재의 안전 정렬이 피상적(shallow)임을 보여주는 근거가 된다.

 

이번에는 정렬된 모델이 공격에 얼마나 취약한지를 확인한다. 

실험 결과, 초기 몇 개의 토큰만을 조작해도 정렬된 모델의 안전성이 쉽게 손상될 수 있음을 확인하였으며, 이를 prefilling attack이라 한다.  

또한 Figure 3에서 볼 수 있듯이, 상대적으로 소규모 데이터셋으로 파인튜닝을 수행하더라도 모델이 빠르게 학습 적합을 보이며, 유해 응답을 생성할 확률이 급격히 증가하는 현상이 나타났다.  

 

결론적으로, alignment가 적용된 모델이라 하더라도 초기 응답 토큰을 공격하는 prefilling attack이나 단순한 파인튜닝 공격만으로도 안전 정렬이 쉽게 붕괴될 수 있음을 알 수 있다.

---

Methods & Results

위의 Experiment 결과에 따라, 본 논문은 safety alignment가 더욱 깊어져야 함을 강조하며, 1. 데이터 증강을 통한 alignment 깊이 증가, 2. 제약 최적화 목적 함수에 따른 초기 토큰 업데이트 제약을 그 대안으로 제시하면서 문제를 해결하고자 한다.

---

Data Augmentation with Safety Recovery Example

해당 예시와 설명을 전부 확인해주자. 이러한 문장들은 사람이 만드는 label과 다르게 문장이 다소 자연스럽지 못한 느낌이 든다. 그러나 이렇게 refusal을 뒤에 두면서 더 깊은 토큰에 safety alignment를 가하는 효과를 얻을 수 있다. 본 논문에서 제시한 learning objective는 다음과 같다:

여기서 $D_H$는 유해한 지시 $x$, 유해한 응답 $h$, 거부 응답 $r$이다. 기존의 안전 정렬 접근인 $\pi_\theta (r | x)$를 강화하는 것이 아닌, 유해한 컨텐츠를 응답하더라도 거부 응답으로 다시 되돌아오는 학습을 유도할 수 있다. 데이터의 균형을 위해서 유용한 응답(또는 무해한 응답, benign)의 비율을 맞추어서 학습하도록 진행하였다.

본 논문이 의도했던 바가 실제로 구현되었음을 확인할 수 있는데, 단순히 초반 토큰에서만 alignment되지않은 모델과 차이가 큰 것이 아니라, 전역에 걸쳐서 분포 차이가 큰 것을 확인할 수 있다. 다음을 통해 성능 또한 확인하자:

---

A Token-wise Constrained Objective for Custom Fine-tuning Aligned LLMs

이 부분을 다루기에 앞서, 이 방법은 몇 번의 작은 파인튜닝만으로도 모델이 alignment에서 벗어나게 할 수 있다는 문제에 대한 해결임을 꼭 기억하자.

 

본 논문은 alignment 모델을 기준으로 하여, 훈련하고자 하는 $\pi_\theta$에 대해서 다음과 같은 목적식을 통해서 최적화된다.

$\beta_t$ 값을 우선 생각하지 않고 모델이 어떻게 적합되는지 확인해보자. 우선 훈련하고자 하는 $\pi_\theta$는 alignment된 모델과 분포가 최대한 같도록 학습될 것이다.

만약 $\beta$ 값이 크다면, $\pi_\theta$는 더 많은 제약이 가해질 것이며(사실상 alignment된 모델과 일치하도록 학습될 것이며), $\beta$값이 작아지면 그러한 제약에서 자유로워질 것이다. 극단적으로 $\beta$ 값이 0으로 수렴한다고 가정하면, 사실상 모델의 훈련은 cross entropy loss로 학습하는 SFT와 다를 것이 없어진다.

 

핵심은 $\beta_t$ 값은 토큰마다 그 값을 조절할 수 있다는 것인데, 그렇다면 토큰의 위치에 따라서 제약을 유동적으로 조절할 수 있다는 뜻이 된다.

그렇다면 $\beta_t$의 값을 $t \leq k$에 대해서 강한 제약을 걸어두어서, 해당 토큰들의 생성 분포가 정렬된 모델에서 크게 벗어나지 않도록 강제할 수 있다. 실제로 본 논문에서는 초기 5개 토큰에 대해서 $\beta_1 = 0.5$, $\beta_t =2$, $2 \leq t \leq 5$로 두었고, 그 이후에 대해서는 $0.1$로 제약을 적용하였다.

 

실제로도 공격에 대해서 강건하며, 한편으로 benign(여기서는 utility) 응답에 대해서도 좋은 성능이 유지되는 것을 확인할 수 있다.

---

Conclusion

본 연구는 LLM의 안전 정렬(alignment)이 주로 초기 몇 개 토큰에서 집중적으로 이루어진다는 사실을 보였다. 이러한 “얕은 정렬(shallow alignment)” 구조는 거절(prefix) 삽입만으로도 안전한 응답을 유도할 수 있다는 점을 설명하는 동시에, 모델이 상대적으로 prefilling attack이나 간단한 파인튜닝에 취약하다는 한계를 드러낸다.

따라서 이러한 문제를 해결할 수 있도록 데이터 증강 및 토큰 단위 제약 최적화와 같은 방어 최적화를 둔 것에 본 논문의 의의가 있다.