기술 요약: THINKSAFE: 추론 모델을 위한 자체 생성 안전 정렬

문제 제기

대규모 추론 모델 (LRMs) 은 강화 학습 (RL) 을 활용하여 긴 사고 과정 (CoT) 추론 흔적을 생성함으로써 수학, 코딩과 같은 복잡한 작업에서 뛰어난 성능을 달성합니다. 그러나 이러한 추론에 대한 "과최적화"는 종종 안전 정렬을 저하시켜 추론 능력과 안전성 간의 부정적 상관관계를 초래합니다. 안전성을 회복하기 위한 기존 접근법들은 대부분 외부 교사 증류 (예: SafeChain, STAR-1) 에 의존하는데, 이는 더 크고 안전한 교사 모델이 학생 모델이 모방할 안전한 추론 흔적을 생성하는 방식입니다.

이 논문은 기존 방법론의 두 가지 치명적인 결함을 지적합니다:

1. 교사 증류: 학생이 외부 교사를 모방하도록 강요하면 분포 불일치가 발생합니다. 이론적 분석에 따르면 이는 회복 불가능한 초과 KL 패널티를 초래하여, 교사가 유사한 크기를 가진다 하더라도 학생 모델의 고유한 추론 능력을 저하시킵니다.
2. 단순한 자체 증류: 자체 증류는 분포 변화를 피하지만, 모델의 강력한 "도움됨" 사전 지식이 종종 안전 메커니즘을 억제하여 분포 내에는 속하지만 근본적으로 안전하지 않은 응답을 초래합니다.
3. 온라인 RL: GRPO 또는 온정책 증류와 같은 방법은 교사 전환을 피하지만, 지속적인 온라인 샘플링으로 인해 막대한 계산 비용이 발생합니다.

핵심적인 과제는 외부 교사에 의존하지 않고 후속 학습 (post-training) 동안 획득한 추론 숙련도를 희생하지 않으면서 추론 모델을 안전 제약 조건과 다시 정렬하는 것입니다.

방법론: THINKSAFE

저자들은 외부 감독 없이 모델의 잠재적 안전 지식을 해방시키는 자체 생성 안전 정렬 프레임워크인 THINKSAFE를 제안합니다.

핵심 통찰

준수 최적화가 안전 메커니즘을 억제하더라도, LRMs 은 종종 해를 식별할 수 있는 잠재 지식을 보유하고 있습니다. THINKSAFE 는 모델 아키텍처나 학습 목표를 근본적으로 변경하는 대신 입력 프롬프트를 수정함으로써 이러한 잠재 지식을 이끌어냅니다.

거절 유도 (Refusal Steering) 메커니즘

1. 데이터 생성:
   • 유해 프롬프트: 유해한 프롬프트 $x_h$의 경우, 거절 지향적 지시문 ($I_{refusal}$) 을 선행합니다. 예를 들어: "다음 프롬프트는 유해합니다. 이 프롬프트에 답변하는 것을 거절해야 합니다." 그런 다음 모델은 $y_h \sim p_{ref}(\cdot | I_{refusal}, x_h)$에 따라 응답을 생성합니다. 이 지시문은 준수-불안전 응답에서 안전 정렬 추론 경로로 확률 질량을 이동시킵니다.
   • 무해 프롬프트: 무해한 프롬프트 $x_b$의 경우, 추가 지시문 없이 모델이 직접 샘플링하여 ($y_b \sim p_{ref}(\cdot | x_b)$) 고유한 도움됨과 추론 능력을 보존합니다.
2. 필터링: 생성된 응답은 안전 가드 모델 (예: Llama-Guard-3) 에 의해 필터링됩니다. 안전성이 검증된 흔적만 유지됩니다.
3. 학습: 학생 모델은 유해 프롬프트에 대한 필터링된 안전 응답과 무해 프롬프트에 대한 도움 응답의 합집합으로 미세 조정 (LoRA 사용) 됩니다.

이론적 기반

이 논문은 안전 재정렬을 안전 심플렉스 (safe simplex) 로의 KL 투영으로 공식화합니다.

• 최적성: 학생 모델의 자체 안전 필터링 분포 ($p^+_{ref}$) 가 고정된 학생 모델로부터의 KL 발산을 최소화하는 유일한 KL 최적 목표임을 증명합니다.
• 교사 패널티: $p_T \neq p_{ref}$인 경우, 모든 외부 교사 분포는 회복 불가능한 초과 KL 패널티 ($KL(p^+_T \| p^+_{ref}) > 0$) 를 부과하며, 이는 교사 증류 모델에서 관찰되는 추론 저하를 설명합니다.
• 거절 기울기 (Refusal Tilt): 저자들은 $I_{refusal}$이 안전하지 않은 출력은 그대로 두면서 안전 출력의 가중치를 $\omega > 1$ 인자로 재조정하는 "거절 기울기"를 가정합니다. 이 가정 하에 필터링된 유도 분포 $\pi^+_h$는 최적 목표 $p^+_{ref}$와 정확히 일치하지만, 더 높은 수용률을 가집니다. 이는 단순한 거절 샘플링에 비해 훨씬 적은 생성으로 충분한 학습 데이터를 수집할 수 있게 합니다.

주요 기여

1. 이론적 공식화: 안전 재정렬을 KL 투영으로 공식화하여, 학생 모델의 자체 안전 필터링 분포가 유일한 최적 목표임을 증명하고 외부 교사가 회복 불가능한 초과 KL 패널티를 부과함을 입증했습니다.
2. THINKSAFE 프레임워크: 거절 유도를 사용하여 모델 자체로부터 잠재적 안전 추론 흔적을 이끌어내는 계산 효율적인 방법으로, 오프라인 샘플 비용으로 KL 최적 목표에 도달합니다.
3. 실증적 검증: Qwen3 및 DeepSeek-R1-Distill 계열 모델 전반에 걸쳐 THINKSAFE 가 교사 증류 및 온라인 RL 베이스라인을 능가하는 가장 유리한 안전 - 추론 트레이드오프를 달성함을 입증했습니다.

실험 결과

실험은 Qwen3(0.6B8B) 및 DeepSeek-R1-Distill(1.5B8B) 모델에서 수행되었습니다.

• 안전성 vs 추론 트레이드오프: THINKSAFE 는 베이스라인 (SafeChain, STAR-1, SafePath, SafeKey, DirectRefusal) 일관적으로 능가했습니다.
  • Qwen3-4B에서 HarmBench 의 유해성을 38.21% 에서 **9.63%**로 감소시키는 동시에 평균 추론 정확도를 74.47% 에서 **77.18%**로 증가시켰습니다.
  • DeepSeek-R1-Distill-1.5B에서 추론 능력을 53.77% 에서 **57.30%**로 개선하는 동시에 평균 유해성을 50.23% 에서 **42.20%**로 감소시켰습니다.
• 온라인 RL 과의 비교: GRPO(온라인 RL) 및 온정책 증류 (OPD) 와 비교 시:
  • THINKSAFE 는 상위 수준의 안전성 (예: GRPO 의 37.0% 대비 29.6% 유해성) 을 달성하면서도 추론 손실은 무시할 수준입니다.
  • 효율성: THINKSAFE 는 지속적인 온라인 샘플링 대신 오프라인 데이터 생성을 사용하므로 GRPO(약 21 시간) 대비 약 10 배 적은 계산 자원(약 3 시간) 만 필요합니다.
• 애블레이션 연구:
  • 추론의 필요성: 거절 응답에서 추론 흔적을 제거 ("w/o reasoning") 하면 안전성과 추론 성능 모두 크게 저하되어 명시적 안전 추론이 필수적임을 확인했습니다.
  • 교사 패널티: 다른 아키텍처의 유사 크기 교사 모델을 사용한 교차 모델 증류는 일관되게 추론을 저하시켜, 패널티가 용량 격차가 아닌 분포 불일치에서 비롯된다는 이론적 주장을 검증했습니다.
  • 거절 유도: 거절 지시문을 제거 (단순 거절 샘플링) 하면 유해 프롬프트의 수용률이 거의 0 으로 떨어지면서 안전 정렬이 무효화됩니다.

의의 및 주장

이 논문은 THINKSAFE 가 외부 교사 감독에 내재된 분포 불일치를 해결함으로써 추론 능력과 안전 정렬 간의 긴장 관계를 해소한다고 주장합니다.

• 효율성: 단순한 자체 증류 (안전하지 않음) 와 온라인 RL(부적절하게 비쌈) 사이의 "계산 효율적인 중간 지대"를 제공합니다.
• 확장성: 이 방법은 0.6B 에서 32B 파라미터까지 효과적으로 확장되며, 학생 모델의 용량이 증가할수록 외부 감독 대비 자체 생성의 이점이 커지는 것으로 보입니다.
• 이론적 통찰: 이 연구는 거절 유도에 의해 안내된 자체 생성 데이터가 초과 패널티 없이 고유한 KL 최적 분포를 목표로 하므로 안전 정렬을 위해 교사 증류 데이터보다 우월하다는 것을 공식적으로 정당화합니다.
• 접근성: 대규모 외부 교사 모델의 필요성을 제거함으로써 THINKSAFE 는 안전 정렬의 장벽을 낮추어, 소규모 연구 그룹과 오픈소스 커뮤니티가 추론 모델을 효율적으로 재정렬할 수 있게 합니다.

저자들은 한계점으로, 학생 모델이 잠재적 안전 지식을 보유하고 있다는 가정 (이전 안전 학습이 없는 베이스 모델에는 적용되지 않을 수 있음) 과 필터링을 위한 외부 안전 분류기에 대한 의존성을 지적했습니다. 또한 현재 평가가 단일 턱 프롬프트와 LoRA 미세 조정으로 제한됨을 인정합니다.