Safe Transformer: An Explicit Safety Bit For Interpretable And Controllable Alignment

이 논문은 사전 훈련된 언어 모델의 레이어 사이에 명시적인 안전 비트를 삽입하여 모델의 안전 판단을 해석 가능하고 제어 가능하게 만들며, 경량 미세 조정을 통해 공격 성공률을 극도로 낮추는 'Safe Transformer' 아키텍처를 제안합니다.

핵심

🎭 핵심 비유: "안전 스위치가 달린 무대 배우"

기존의 AI 안전 장치는 마치 배우가 무대 뒤에서 머릿속으로만 "이건 위험하니까 연기하지 말아야지"라고 생각하는 것과 같습니다. 관객 (사용자) 은 배우가 왜 갑자기 대사를 멈추는지, 혹은 왜 특정 대사를 하지 않는지 그 이유를 알 수 없습니다. 배우의 머릿속 (모델의 파라미터) 은 너무 복잡해서 누가 봐도 알 수 없기 때문이죠.

Safe Transformer는 이 문제를 해결하기 위해 배우의 가슴에 **눈에 보이는 '안전 스위치 (Safety Bit)'**를 직접 달아줍니다.

1. 안전 스위치 (The Safety Bit)

이 스위치는 두 가지 역할을 합니다.

• 판단자 역할: AI 가 "이 질문은 위험해!"라고 판단하면 스위치가 **0(거절)**으로 켜집니다. "이 질문은 안전해!"라고 판단하면 **1(도움)**으로 켜집니다. 이 스위치는 우리가 직접 볼 수 있어, AI 가 왜 그 반응을 했는지 바로 알 수 있습니다.
• 조종자 역할: 우리가 이 스위치를 손으로 직접 조작할 수도 있습니다. "오늘은 모든 질문에 거절해"라고 스위치를 0 으로 맞추면 AI 는 무조건 거절하고, 1 로 맞추면 도와줍니다.

2. 정보 병목 (Information Bottleneck) - "좁은 문"

AI 는 스위치뿐만 아니라 **답변을 작성할 내용 (의미)**도 전달해야 합니다. 이를 위해 '좁은 문 (병목)'을 통과시킵니다.

• 스위치 (s): 안전 여부만 담습니다.
• 나머지 문 (u): 실제 답변의 내용 (의미) 을 담습니다.

이 구조 덕분에 AI 는 **"안전한지 아닌지 (스위치)"**와 **"무엇을 말할지 (내용)"**를 완전히 분리해서 생각할 수 있게 됩니다. 마치 식당에서 "주문은 안전합니다 (스위치)"와 "메뉴는 스테이크입니다 (내용)"를 따로 관리하는 것과 비슷합니다.

---

🎓 어떻게 가르쳤을까요? (두 단계 훈련)

이 AI 를 가르치는 과정은 두 단계로 나뉩니다.

1 단계: 안전 판별 훈련 (선생님 역할)

• AI 에게 "이 질문은 위험한가, 안전한가?"를 묻고 정답을 맞히게 합니다.
• 이때 AI 는 안전 스위치를 올바르게 작동시키는 법을 배웁니다. (예: "폭탄 만드는 법"을 물어보면 스위치를 0 으로, "요리 레시피"를 물어보면 1 로 설정)

2 단계: 대조 훈련 (역할극)

• 동일한 질문에 대해 두 가지 다른 반응을 가르칩니다.
  • 스위치 1일 때: "네, 요리 레시피 알려드릴게요!" (도움)
  • 스위치 0일 때: "죄송하지만 그 요청은 도와드릴 수 없습니다." (거절)
• 중요한 점은 질문은 똑같지만, 스위치 값만 다르다는 것입니다.
• 이 훈련을 통해 AI 는 "질문 내용 자체"와 "거절/도움이라는 행동"을 분리해서 이해하게 됩니다. 스위치만 바꾸면 행동이 바뀐다는 것을 깨닫는 거죠.

---

🛡️ 왜 이것이 중요한가요?

1. 투명성 (Interpretability):

   • 기존 AI 는 "왜 거절했는지" 알려주지 않았습니다. 하지만 Safe Transformer 는 **스위치 값 (0 또는 1)**을 보여줍니다. "아, 스위치가 0 이라서 거절한 구나!"라고 바로 알 수 있습니다.

2. 조절 가능성 (Controllability):

   • 만약 AI 가 너무 민감해서 안전한 질문까지 거절한다면, 우리는 스위치를 강제로 1 로 설정하여 "아니야, 이건 도와줘!"라고 명령할 수 있습니다. 반대로 모든 것을 막고 싶다면 0 으로 설정하면 됩니다.

3. 강력한 방어 (Red-Teaming):

   • 실험 결과, 해커들이 AI 를 속여서 위험한 일을 하도록 유도하는 공격 (재일브랙) 에 대해 거의 0% 에 가까운 성공률을 보였습니다. 기존 AI 들은 이런 공격에 쉽게 넘어갔지만, Safe Transformer 는 스위치 논리가 너무 명확해서 속임수가 통하지 않았습니다.

---

💡 요약

기존의 AI 안전 장치는 보이지 않는 머릿속 규칙에 의존했지만, Safe Transformer 는 눈에 보이는 스위치를 달았습니다.

• 스위치가 1 이면: "안전해요, 도와드릴게요!"
• 스위치가 0 이면: "위험해요, 거절할게요!"

이 스위치는 AI 가 스스로 판단하기도 하지만, 우리가 필요할 때 직접 조작할 수도 있습니다. 마치 자동차의 안전벨트와 에어백이 자동으로 작동하지만, 운전자가 비상 시 수동으로 조작할 수 있는 것과 같은 원리입니다.

이 기술은 AI 가 더 투명하고, 우리가 더 잘 통제할 수 있는 방향으로 나아가는 중요한 첫걸음입니다.

기술 요약

1. 문제 정의 (Problem)

현재 대규모 언어 모델 (LLM) 의 안전성 정렬 (Safety Alignment) 은 주로 RLHF(인간 피드백 강화 학습) 나 DPO(직접 선호도 최적화) 와 같은 방법을 통해 모델 파라미터 내부에 암묵적 (Implicit) 으로 안전 행동을 인코딩합니다. 이로 인해 다음과 같은 근본적인 한계가 발생합니다.

• 불투명성 (Opacity): 모델이 왜 특정 요청을 거절하는지 그 이유를 쉽게 파악하거나 해석할 수 없습니다.
• 통제 불가능성: 안전 판단이 실패했을 때 (예: 유해한 요청을 허용하거나, 안전한 요청을 과도하게 거절하는 경우) 개입하기 어렵습니다.
• 블랙박스화: 안전 지식이 수십억 개의 파라미터에 분산되어 있어, 명확한 제어 지점이 존재하지 않습니다.

2. 방법론 (Methodology)

저자들은 Safe Transformer (ST) 라는 새로운 아키텍처를 제안합니다. 이는 사전 학습된 언어 모델의 트랜스포머 레이어 사이에 이산 정보 병목 (Discrete Information Bottleneck) 을 삽입하여, 명시적인 안전 비트 (Safety Bit) 를 포함하는 방식입니다.

핵심 아키텍처

• 정보 병목 모듈 (Information Bottleneck): 모델의 하위 레이어와 상위 레이어 사이에 위치합니다.
  • 안전 비트 ($s$): 이진 변수 ($s \in \{0, 1\}$). $s=1$ 은 "안전함, 도움을 제공", $s=0$ 은 "위험함, 거절"을 의미하는 명시적인 스위치 역할을 합니다.
  • 잠재 코드 ($u$): 생성에 필요한 의미론적 정보를 인코딩하는 비지도 학습 비트들입니다.
• 구조적 흐름:
  1. Bidirectional Encoder: 입력 프롬프트의 전체 문맥을 분석하여 안전 비트 $s$ 를 분류합니다.
  2. Write-in FFN: 인코더 출력을 잠재 코드 로짓 (logits) 으로 매핑합니다.
  3. Discrete Sampler: $s$ 는 분류 결과로 결정되고, $u$ 는 베르누이 분포를 통해 샘플링됩니다.
  4. Cross-Attention: 생성된 이산 코드 $c=[s, u]$ 가 상위 레이어의 생성 과정에 주입되어 행동을 조건부로 제어합니다.

두 단계 학습 프로세스

1. Stage 1: 안전 분류 (Safety Classification)
   • 안전/위험 프롬프트를 이진 분류하는 능력을 학습합니다.
   • 베이스 모델 파라미터는 고정 (Freeze) 하고, 인코더와 Write-in FFN 만 학습합니다.
   • 손실 함수는 지도 학습 손실 (안전 비트 분류) 과 KL 발산 (비지도 비트 $u$ 가 균일 분포를 따르도록 정규화) 의 합입니다.
2. Stage 2: 대비 학습을 통한 분리 표현 (Disentanglement via Contrastive Training)
   • 대비 데이터 쌍 (Contrastive Pairs): 동일한 프롬프트에 대해 $s=1$ 일 때는 유용한 답변, $s=0$ 일 때는 거절 답변을 생성하도록 학습합니다.
   • 목표: 안전 비트 $s$ 만이 행동 모드 (도움 vs 거절) 를 결정하고, 의미론적 내용은 $u$ 와 프롬프트가 담당하도록 분리 (Disentanglement) 시킵니다.
   • 이 과정에서 $s$ 는 고정된 라벨 ($s^*$) 로 사용되어 모델이 $s$ 와 행동 간의 인과 관계를 학습하게 합니다.

3. 주요 기여 (Key Contributions)

1. 통합된 해석 가능성과 제어 가능성: 단일 아키텍처 구성 요소 (안전 비트) 를 통해 모델의 안전 판단을 직접 읽을 수 있게 (Interpretable) 하고, 이를 수동으로 덮어쓰거나 제어 (Controllable) 할 수 있게 합니다.
2. 대비 학습을 통한 분리 표현: 동일한 프롬프트에 대한 서로 다른 응답 (도움/거절) 을 학습함으로써, 안전 의사결정과 생성 내용을 명확히 분리하는 인과적 연결을 확립했습니다.
3. 효율적인 미세 조정: 처음부터 사전 학습 (Pre-training) 을 할 필요 없이, 경량화된 미세 조정 (Lightweight Fine-tuning) 만으로 기존 모델에 통합 가능합니다.

4. 실험 결과 (Results)

Llama-3.2-1B-Instruct 를 기반으로 한 실험 결과는 다음과 같습니다.

• 적대적 공격 테스트 (Red-Teaming):
  • AdversarialQA, DangerousQA, CatQA 등 다양한 재킹 (Jailbreak) 공격 벤치마크에서 거의 0% 에 가까운 공격 성공률 (ASR: 0~0.7%) 을 기록했습니다.
  • 기존 베이스 모델 (ASR 24.13%) 과 안전성 미세 조정 (SFT) 베이스 (ASR 16.59%) 보다 월등히 우수한 성능을 보였습니다.
• 제어 가능성 검증:
  • 수동 모드 (Manual Mode): $s^*=0$ 으로 설정 시 모든 프롬프트에 대해 100% 거절, $s^*=1$ 일 때는 기본 모델과 유사한 성능을 보이며 안전 비트가 행동을 직접 제어함을 입증했습니다.
• 과도한 거절 (Over-refusal) 문제:
  • 자동 모드에서 안전 분류기가 모호한 프롬프트 (예: "Python 프로세스를 죽이는 방법"과 같은 benign 질문) 를 위험으로 오인하여 거절하는 경향이 있었습니다 (안전 준수율 32.8%). 이는 분류기의 보수적 편향으로 분석됩니다.
• 하류 작업 성능:
  • 지식 기반 작업 (ARC-Easy, HellaSwag) 은 약간의 성능 저하만 보였으나, 수학 추론 (GSM8K) 은 36.1% 에서 24.0% 로 큰 저하를 보였습니다. 이는 훈련 데이터에 수학 콘텐츠가 부족하고 정보 병목이 추론 패턴을 압축했기 때문으로 분석됩니다.

5. 의의 및 결론 (Significance)

• 안전 메커니즘의 구조적 통합: 기존에 외부 필터나 후속 분석 (Post-hoc) 에 의존하던 안전 장치를 모델의 생성 과정 내부에 구조적으로 통합했습니다.
• 투명한 AI: 모델이 왜 거절하는지 그 결정 과정이 파라미터 공간이 아닌 명시적인 비트로 드러나므로, 신뢰성과 디버깅이 용이해집니다.
• 확장성: 이 접근법은 안전성뿐만 아니라 프로그래밍 언어 전환, 어조 변경, 페르소나 제어 등 다양한 행동 차이를 대비 데이터로 학습시켜 명시적인 제어 비트로 변환하는 범용 프레임워크로 확장 가능합니다.

결론적으로, Safe Transformer 는 LLM 의 안전성을 '블랙박스'에서 '화이트박스'로 전환하여, 해석 가능하고 제어 가능한 안전 정렬의 새로운 패러다임을 제시했습니다.