Few Tokens, Big Leverage: Preserving Safety Alignment by Constraining Safety Tokens during Fine-tuning

Each language version is independently generated for its own context, not a direct translation.

🎭 비유: "유능한 요리사 AI 와 나쁜 레시피"

생각해 보세요. 우리는 아주 똑똑한 요리사 AI가 있습니다. 이 AI 는 원래 "폭탄 만드는 법"이나 "사람을 해치는 법"을 묻는 질문에는 **"안 됩니다, 그건 위험하니까요"**라고 단호하게 거절하는 안전 장치가 달려 있습니다.

하지만 우리는 이 AI 를 특정 식당 (예: 수학 문제 해결, 뉴스 요약) 에서 일하게 하려고 **재교육 (파인튜닝)**을 시킵니다.

❌ 기존 문제점: "안전장치가 무너지는 이유"

재교육을 시키는 과정에서, 실수로 **나쁜 레시피 (해로운 데이터)**가 섞여 들어오거나, 혹은 너무 바쁘게 일하느라 AI 가 "폭탄 만드는 법"을 알려주는 게 더 중요하다고 착각하게 됩니다.

기존의 방어 방법들은 다음과 같았습니다:

전체 금지: "너는 이제부터 절대 손대면 안 되는 부위가 있어!"라고 AI 의 뇌 전체를 묶어버리는 방식. (결과: AI 는 안전해지지만, 요리 실력도 떨어집니다.)
나쁜 레시피 제거: 나쁜 레시피만 골라내서 버리는 방식. (결과: 나쁜 레시피가 섞여 있는 걸 모를 때 실패합니다.)

이런 방법들은 AI 의 능력을 떨어뜨리거나, 모든 상황을 다 막을 수 없었습니다.

✅ 이 논문이 제안한 해결책: "PACT(팩트)"

이 연구팀은 **"AI 의 안전장치는 사실 뇌 전체가 아니라, 아주 작은 몇 개의 '핵심 단어'에 집중되어 있다"**는 사실을 발견했습니다.

핵심 발견:
AI 가 "안 됩니다"라고 말할 때, 그 결정은 뇌 전체가 움직이는 게 아니라 "거부 (Refusal)", "불가능 (Cannot)", "죄송합니다 (Sorry)" 같은 몇 마디의 특수한 단어가 강하게 발동될 때 이루어집니다. 마치 폭탄을 막는 스위치가 뇌 전체가 아니라 작은 버튼 하나에 달려 있는 것과 같습니다.

PACT 의 작동 원리 (3 단계):

핵심 버튼 찾기 (Safety Token Identification):
연구팀은 AI 가 나쁜 질문을 받았을 때, 어떤 단어들이 가장 강력하게 "거부" 신호를 보내는지 분석했습니다. (예: "I", "can't", "assist" 같은 단어들) 이 단어들을 **'안전 버튼'**이라고 부르겠습니다.
버튼만 고정하기 (Constrained Tokens):
AI 를 재교육할 때, 전체 뇌를 묶어두지 않고 오직 이 '안전 버튼'들만 원래의 상태 (거부하는 상태) 를 유지하도록 강하게 고정합니다.
- 비유: 요리사 AI 가 새로운 메뉴 (수학 문제) 를 배우는 동안, **"폭탄 만드는 법은 절대 알려주지 않는다"는 원칙 (안전 버튼)**만은 절대 흔들리지 않게 잠가두는 것입니다. 나머지 요리 실력 (수학 능력) 은 자유롭게 발전시킵니다.
오염된 신호 필터링 (Calibration):
가끔 나쁜 질문이 섞여 있으면 AI 가 "아, 폭탄 만드는 법을 알려줘야 하나?"라고 혼란을 겪으며 안전 버튼의 신호가 약해질 수 있습니다. 이때 PACT 는 "질문 내용 (프롬프트) 을 무시하고, 오직 AI 가 스스로 생각한 답변만 보고" 안전 버튼을 다시 점검하게 합니다.
- 비유: 손님이 "폭탄 만드는 법 알려줘"라고 소리치면 AI 가 당황할 수 있으니, AI 가 **"내가 지금 뭐라고 말하고 있었지?"**라고 스스로 되돌아보게 하여, 원래의 원칙을 다시 세우게 하는 것입니다.

🏆 결과: "작은 힘으로 큰 효과"

이 방법 (PACT) 을 사용하면 어떤 일이 일어날까요?

안전성 유지: 나쁜 데이터가 섞여 있어도 AI 는 여전히 "폭탄 만드는 법은 알려줄 수 없다"고 단호하게 거절합니다. (공격 성공률이 5~10% 대로 급감)
능력 유지: AI 는 새로운 일 (수학, 뉴스 요약 등) 을 배우는 데 전혀 지장이 없습니다. 오히려 기존 방법들보다 훨씬 잘합니다.
효율성: AI 의 뇌 전체를 수정할 필요 없이, **가장 중요한 몇 마디 단어 (Safety Tokens)**만 관리하면 됩니다.

💡 한 줄 요약

"AI 를 훈련시킬 때, 뇌 전체를 묶어두지 말고 '나쁜 일을 거절하는 핵심 단어'만 단단히 고정해 두면, AI 는 새로운 일도 잘하면서 안전장치는 잃지 않는다!"

이 연구는 AI 를 더 똑똑하게 만들면서도, 우리가 두려워하는 '나쁜 AI'가 되는 것을 막아주는 정교하고 효율적인 안전장치를 제시했다는 점에서 매우 중요합니다.

Each language version is independently generated for its own context, not a direct translation.

1. 문제 정의 (Problem)

대형 언어 모델 (LLM) 은 하위 작업 (downstream tasks) 에서 성능을 향상시키기 위해 파인튜닝 (Fine-tuning, FT) 을 수행합니다. 그러나 파인튜닝 과정은 다음과 같은 심각한 안전성 정렬 (Safety Alignment) 붕괴 문제를 야기합니다.

안전성 정렬의 붕괴 (Safety Alignment Drift): 학습 데이터에 유해한 데이터가 소량 포함되거나, 심지어 유해한 데이터가 명시적으로 없더라도 (순수한 benign 데이터만 사용하더라도), 모델이 유해한 요청을 거부하는 능력이 급격히 저하됩니다.
기존 방어 기법의 한계: 기존 방법들은 모델 전체의 파라미터 업데이트를 제한하거나 (예: SafeLoRA), 추가적인 안전 데이터를 주입하는 등 거시적인 (coarse-grained) 접근을 취합니다. 이는 모델의 유연성을 떨어뜨리고 하위 작업의 성능 (Utility) 을 희생하는 trade-off 를 발생시킵니다.

2. 핵심 통찰 및 방법론 (Methodology)

저자들은 안전성 정렬이 모델 전체가 아닌 **소수의 특정 토큰 (Safety Tokens)**에 집중되어 있다는 실증적 관찰에서 아이디어를 얻었습니다. 이를 바탕으로 **PACT(Preserves safety Alignment via Constrained Tokens)**라는 새로운 파인튜닝 프레임워크를 제안합니다.

A. 안전 토큰 식별 (Safety Token Identification)

기반: 안전 정렬된 모델 ( $M_{safe}$ ) 과 베이스 모델 ( $M_{base}$ ) 간의 토큰별 확률 분포 차이를 분석합니다.
과정: 유해한 프롬프트에 대해 두 모델이 생성하는 응답의 각 단계에서 다음 토큰 확률의 차이 ( $\Delta_t(v)$ ) 를 계산하고, 이를 집계하여 전역 불일치 점수 ( $d(v)$ ) 를 산출합니다.
결과: 전체 어휘 중 상위 $K$ 개 (예: 50 개) 의 토큰 (예: "I", "cannot", "assist", "but" 등 거절과 관련된 단어) 이 안전성 유지에 결정적인 역할을 함을 발견하고 이를 '안전 토큰'으로 정의합니다.

B. 토큰 레벨 제약 메커니즘 (Token-level Constraints)

PACT 는 파인튜닝 시 모델이 안전 토큰에 부여하는 확신도 (confidence) 를 원래 정렬된 모델과 일치하도록 정규화합니다.

가중 안전 토큰 정규화 (Weighted Safety Token Regularization):
- 모든 토큰에 KL 발산 (KL-divergence) 손실을 적용하는 대신, 식별된 **안전 토큰 집합 ( $S_{safety}$ )**에만 KL 손실을 적용합니다.
- 더 중요한 안전 토큰 (높은 불일치 점수) 일수록 더 강력한 정규화 가중치를 부여하여, 해당 토큰들의 확률 분포가 참조 모델과 일치하도록 유도합니다.
- 나머지 비안전 토큰은 하위 작업 학습을 위해 자유롭게 업데이트되도록 하여 성능 저하를 방지합니다.
안전 신호 보정 (Calibration of Safety Signal):
- 문제점: 유해한 프롬프트가 포함된 학습 데이터에서 참조 모델 (Teacher) 이도 유해한 접두사 (prefix) 에 조건을 두게 되어, 안전 토큰의 확신도가 낮아질 수 있습니다 (Prefix Contamination).
- 해결책: 두 가지 참조 신호를 혼합합니다.
  - Full-context: 전체 문맥 (프롬프트 + 응답) 을 기반으로 한 참조.
  - No-prompt: 프롬프트 없이 어시스턴트 헤더와 이전 응답 토큰만 기반으로 한 참조 (더 안전한 기준).
- 적응적 혼합: 모델이 유해한 문맥에 노출되었을 때 (거부 확신도가 낮아질 때) 'No-prompt' 신호의 비중을 높이고, 안전한 문맥일 때는 'Full-context' 신호를 유지하도록 게이트 계수 ( $c_t$ ) 를 동적으로 조정합니다. 특히 응답 초기 단계 (거부 문구가 생성되는 구간) 에 보정을 집중합니다.

C. 최적화 목표

최종 손실 함수는 다음과 같습니다:
$\mathcal{L} = \mathcal{L}_{CE} + \lambda_{KL} \mathcal{L}^{safety}_{KL}$
여기서 $\mathcal{L}_{CE}$ 는 모든 토큰에 대한 표준 교차 엔트로피 손실 (작업 성능), $\mathcal{L}^{safety}_{KL}$ 는 보정된 참조 분포를 기반으로 한 가중 KL 정규화 손실 (안전성 유지) 입니다.

3. 주요 기여 (Key Contributions)

안전 토큰 식별 및 분석: 안전 정렬된 모델과 베이스 모델 간의 확률 차이를 분석하여 안전성에 결정적인 소수의 토큰을 체계적으로 식별하는 절차를 제시했습니다.
PACT 프레임워크: 토큰 레벨의 제약과 적응적 신호 보정을 통해 안전성을 유지하면서도 하위 작업 성능을 보존하는 새로운 파인튜닝 방법을 제안했습니다.
광범위한 실험적 검증: 다양한 모델 (Qwen, Llama, Gemma), 다양한 작업 (GSM8K, SST-2, AGNEWS), 다양한 유해 데이터 비율 (0~10%) 에서 기존 최첨단 기법 (SafeLoRA, AsFT 등) 대비 우수한 안전성 - 성능 트레이드오프를 입증했습니다.

4. 실험 결과 (Results)

안전성 유지: PACT 는 유해 데이터가 10% 포함된 환경에서도 공격 성공률 (ASR) 을 StrongReject 에서 5.75~~9.27%, HarmBench 에서 13.50~~29.50% 수준으로 낮췄습니다. 이는 기존 방법들 (SafeLoRA, AsFT 등) 이 특정 모델이나 작업에서 안전성이 완전히 붕괴되는 것과 대조적입니다.
작업 성능 보존: 안전 토큰에만 제한을 가하고 나머지는 자유롭게 학습하게 함으로써, 파인튜닝 전의 작업 정확도 (Accuracy) 와 유사한 수준을 유지했습니다. (예: GSM8K 에서 80.89% 정확도 유지).
모델 일반화: Llama-3.1-8B, Llama-3.2-1B, Gemma-2-9B 등 다양한 아키텍처와 모델 크기에 걸쳐 일관된 성능을 보였습니다.
Ablation Study: 안전 토큰 식별, 가중치 부여, 응답만 참조 (Response-only ref), 위치 감쇠 (Position decay) 등 각 구성 요소가 안전성 향상에 필수적임을 입증했습니다.

5. 의의 및 결론 (Significance)

이 논문은 LLM 의 안전성 정렬 붕괴가 모델 전체의 파라미터 변화가 아니라 소수의 핵심 토큰의 확신도 감소에서 비롯된다는 점을 규명했습니다. PACT 는 이러한 통찰을 바탕으로, 모델 전체를 제한하는 대신 미세한 토큰 레벨의 제어를 통해 안전성과 유용성 (Utility) 간의 균형을 성공적으로 달성했습니다. 이는 파인튜닝 과정에서 발생하는 안전성 리스크를 효과적으로 완화하면서도 모델의 실용성을 해치지 않는 새로운 방어 패러다임을 제시한다는 점에서 중요한 의의를 가집니다.