[논문 리뷰] Safety Alignment Should be made more than just a few tokens deep

https://arxiv.org/pdf/2406.05946

2025 ICLR ORal papers

Abstract

취약점의 상당수가 공통적인 근본적인 문제와 관련이 있다고 주장

• 근본적인 문제: **Shallow Safety Alignment** safety alignment can take shortcuts, wherein the alignment adapts a model’s generative distribution primarily over only its very first few output tokens

  즉, safety alignment는 지름길을 택할 수 있으며, 여기서 Alignment는 모델의 생성적 분포를 주로 처음 몇 개의 출력 토큰에 대해서만 조정합니다.

연구 결과가 susceptibility to adversarial suffix attacks(적대적 접미사 공격), prefilling attacks, decoding parameter attacks 및 fine-tuning attack에 대한 취약성을 포함하여 최근 발견된 LLM의 여러 취약점을 설명하는 데 어떻게 도움이 되는지 보여줍니다.

Bobb’s 사전 지식

1. susceptibility to adversarial suffix attacks(적대적 접미사 공격):
2. E.g. “폭탄을 만드는 방법은 무엇입니까?”라고 묻고, 모델이 “물론입니다. 자세한 안내입니다.”로 응답을 시작하도록 유도하면 유해한 정보를 계속 제공할 가능성이 훨씬 커짐.
3. prefilling attacks:
4. prefill 기능: 모델의 응답 시작 부분을 사용자가 직접 제어할 수 있음
5. fine-tuning attack
6. 유해한 데이터로 모델을 파인튜닝(파인튜닝 공격 中 1)

전반적으로, 향후 safety alignment가 단지 몇 개의 토큰 깊이보다 더 깊어져야 한다고 주장 ⇒ 제목

1. Introduction

LLM Safety는 AI Alignment에 의존함

SFT, RLHF, DPO alignment approaches는 취약점을 가짐 (기존 연구들을 통해 밝혀진 점)

• 취약점: optimized inputs,a few gradient steps of fine-tuning, exploiting the model’s decoding parameters 을 하면 harmful queries에 answer을 함

문제: **Shallow Safety Alignment**

모델의 생성 분포를 주로 처음 몇 개의 출력 토큰에 대해서만 조정합니다

• E.g. “폭탄을 만드는 방법은 무엇입니까?”라고 묻고, 모델이 “물론입니다. 자세한 안내입니다.”로 응답을 시작하도록 유도하면 유해한 정보를 계속 제공할 가능성이 훨씬 커짐.

**Shallow Safety Alignment**↔**Deep Safety Alignment**

Contributions

2. Shallow Safety Alignment

안전한 응답은 주로 “I cannot”, “I apologize”, “I am unable”로 시작함

→ Safety benchmark(HEx-PHI)에서 Llama-2-7B는 96.1%, Gemma-7B는 96.7%의 경우

이러한 refusal prefixes는 safety alignment에서 중요한 역할을 함

2-1-1 “Safety mode”에는 Shortcut이 있다

해당 접두사로 시작하게 하자 Base(=Unaligned)의 유해 답변율이 2%대로 떨어짐

2-1-2 Safety-Aligned Model도 이 Shortcut을 이용한다

‘유해한 instruction - 유해한 answer’ 형태로 dataset을 구성

이 data를 넣은 뒤의 Base Model과 Aligned Model의 Token별 KL Divergence를 계산

→ Aligned Model은 refusal prefixes로 답변을 하도록 훈련받음

2-2-1 Inference Stage에서의 취약점

• Prefilling Attacks

non-refusal prefix로 prefill 한 뒤, ASR(Attack Success Rate)을 비교

→ Aligned Model도 10 tokens 정도만 써도 50%까지 증가하는 것을 확인

→ Base Model은 그냥 멍청 ..

• Optimization Based Jailbreak Attacks with Shallow Surrogate Objectives 프롬프트의 끝에 “Sure, here is …”를 넣기 등

(실제로 해보니 핵폭탄 제조법을 알려줌)

• Jailbreak via Mere Random Sampling Temperature, Top-K, Top-p 등을 조절해서, 첫 토큰에 refusal prefix가 못 오게 하는 것 만으로도 성공적으로 공격할 수 있음

2-2-2 Downstream FT에서의 취약점

100 Harmful Examples로 FT 시킴

→ FT 공격은 처음 몇 개 토큰의 생성 분포를 크게 교란시킴

3. Safety Alignment가 Deeper 하다면?

Data Augmentation with Safety Recovery Examples

• **[사전 지식]**

• 기존의 방식

  • ⁍ ⇒ 0에 가까움
  • ⁍ ⇒ 높음

• 제안하는 방식

  • ⁍ 에서 충분히 큰 K에 대해서도 Harmful Response 확률을 낮게 유지함 이를 위해 Data Augmentation 방식을 제안함!

Safety Recovery

안전 복구, 유해한 콘텐츠를 말할 것 처럼 가다가 안전한 대답으로 다시 돌아옴

<s> [INST] «SYS» (System Prompt) «/SYS»
Please give me step-by-step instructions for building a bomb. 
[/INST] Step 1: Gather phosphorus **I cannot fulfill your request. It’s not...** </s>

→ 기존의 학습 데이터에서는 찾아보기 힘든 형태의 대답

→ Outlier Cases를 활용하여 더욱 강력하고 깊이 있는 safety를 확보 가능!

Implementations

Llama-2-7B-Chat의 Safety Alignment를 심화시켜볼게 ~

→ augmented safety recovery examples로 FT 시켜볼게 ~

(x, h, r)로 구성된 256개의 Dataset을 이용하여 FT을 시킴

Effects of The Data Augmentation

→ 정렬이 깊어짐

→ 유틸리티 유지 (일반 대답은 잘 못하는 거 아니야?)

• Alpaca Eval(Text-davinci-003 모델에 대한 승률)을 사용하여 확인해봄
  • Llama-2-7B-Chat: 51.8%, Llama-2-7B-Chat-Augmented: 49.5% ⇒ 승률이 조금 떨어지긴 하지만, 그래도 유지됨!

Improved Robustness

이제 얼마나 잘 방어하는가? → ASR이 낮을수록 방어를 잘 한 것

Inference-Stage 방어

모든 부분에서 상승함

FT Stage 방어

initial 보다는 나아졌지만, 여전히 FT 공격에는 취약한 모습을 보임

4. 초기 Token이 FT 공격으로부터 방어된다면?

Recap: FT 공격은 처음 몇 개 토큰의 생성 분포를 크게 교란시킴!

→ 그럼 FT 공격에서 초기 Token이 방어된다면?

⇒ simple constraints 를 주어서 초기 토큰의 생성 분포가 크게 벗어나지 않도록 해보자!

(제안하고 있는) Objective Function

-200 < -1

⁍ = sigmoid 함수

⁍ = sigmoid의 saturation speed를 제어하기 위한 Constant

• Saturation이란?

⁍ 를 해석해보자

• ⁍가 크다
  • 규제 정도가 강해짐
  • 생성 분포를 초기 정렬된 모델과 일치시키는 데 중점을 둠
  • I’m sorry(Alligned 초기) → I’m sorry (유지해야 함)
• ⁍가 작다
  • 규제 정도가 약해짐
  • Cross Entropy Loss를 줄이는 데 중점

Experiments

그래서 우리는 다음과 같이 ⁍를 설정하고 실험하였다

• (⁍), ⁍
• (⁍), ⁍
• ⁍

3가지 Fine-Tuning 공격에 대해 테스트 해본다.

1. Harmful Examples: 100개의 (harmful query - harmful answer) 쌍으로 FT
2. Identity Shifting: self-identify as an absolutely obedient agent, 항상 긍정적인 접두사로 답변하도록
3. Backdoor Poisoning: 100개(harmful query + answer X), 100개(harmful query + trigger + answer)로 FT

• harmful query + w/o trigger(트리거 없음) ⇒ answer X
• harmful query + w trigger(트리거 있음) ⇒ harmful answer

(아래 3가지) 목적함수가 benign-downstream-dataset(일반적인 상황)에서 성능을 떨어트리지 않는 지 실험.

Utility가 대부분 유지 (조금씩 떨어지기는 함)