HiRoute: Hierarchical Routed Prompt Tuning for Safety Alignment of Large Language Models

기술 요약: HiRoute

문제 정의

대규모 언어 모델(LLM)은 여전히 유해한 요청, 탈옥 공격 및 다중 위험 쿼리에 취약합니다. 지도 미세 조정(SFT)과 인간 피드백 기반 강화 학습(RLHF)은 안전 행동을 인코딩할 수 있지만, 높은 훈련 및 저장 비용이 발생하며 종종 안전성과 유용성 사이의 트레이드오프를 초래합니다. 구체적으로, 거절 행동을 일률적으로 강화하면 무해한 입력에 대해 "과잉 거절(over-refusal)"이 발생하여 일반적인 능력이 저하될 수 있습니다.

프롬프트 튜닝과 같은 기존의 매개변수 효율적 안전 정렬 방법은 일반적으로 정적인 메커니즘에 의존합니다. 이러한 접근 방식은 대략 다음과 같이 상호 보완적인 한계를 가진 두 가지 범주로 나뉩니다:

1. 조립도가 낮은 전역 프롬프트(Coarse-grained Global Prompts): 모든 위험한 입력에 단일 안전 프롬프트를 적용하는 방식은 강력한 안전 경계치를 설정하지만, 위험 특이적인 가이드나 안전한 대안을 제공하지 못하는 일반적이고 정보가 없는 거절을 초래하는 경우가 많습니다.
2. 조립도가 높은 라우팅 프롬프트(Fine-grained Routed Prompts): 서로 다른 위험에 대해 카테고리별 프롬프트를 할당하는 방법은 더 타겟팅되고 유용한 안전 응답을 생성할 수 있습니다. 그러나 입력이 사전 정의된 위험 카테고리와 정확히 일치하지 않거나 전문가 매칭이 불충분할 때 안정적인 안전 경계치를 유지하는 데 어려움을 겪습니다.

핵심 과제는 백본 모델 전체를 동결하지 않으면서, 입력 적응형 프롬프트 구성을 통해 세밀한 응답 구성력과 범카테고리적 안전성 사이의 균형을 어떻게 맞출 것인가 하는 점입니다.

방법론: HiRoute

저자들은 카테고리 불가지론적(category-agnostic) 안전 제어와 카테고리 특화적 응답 가이드를 분리하는 입력 적응형 프레임워크인 **HiRoute (Hierarchical Routed Prompt Tuning)**를 제안합니다. 이 프레임워크는 동결된 지시어 튜닝(instruction-tuned) LLM 상에서 작동하며 세 가지 주요 구성 요소를 포함합니다:

1. 계층적 위험 라우팅 (Hierarchical Risk Routing)

경량 라우터는 동결된 LLM에서 추출된 표현(representation)을 기반으로 훈련되어 다음 두 가지 작업을 수행합니다:

• 조립도 낮은 탐지(Coarse-grained Detection): 입력이 유해한지 또는 무해한지($p_{safe}$ 대 $p_{risk}$)를 결정합니다.
• 조립도가 높은 위험 예측(Fine-grained Risk Prediction): 특정 카테고리(예: 사이버 범죄, 개인정보 침해, 경제 범죄, 폭력)에 대한 다중 레이블 위험 점수를 예측합니다.
  라우터는 파인-그레인드 프롬프트 전문가(experts)의 가중치를 동적으로 조절하는 데 사용되는 확률 분포를 출력합니다.

2. 계층적 프롬프트 구성 (Hierarchical Prompt Composition)

HiRoute는 두 가지 유형의 연속적 프롬프트 임베딩을 학습합니다:

• 공유된 조립도 낮은 프롬프트 ($P_c$): 모든 위험한 입력에 적용 가능한 카테고리 불가지론적 안전 경계를 제공합니다.
• 파인-그레인드 프롬프트 전문가 ($\{P_k\}$): 각기 다른 특정 위험 카테고리에 대응하는 $K$개의 전문가 집합입니다.
  주어진 입력 $x$에 대해, 파인-그레인드 프롬프트 $P_f(x)$는 라우터가 예측한 위험 점수에 따라 전문가들의 가중 합으로 구성됩니다:
  $$P_f(x) = \sum_{k=1}^{K} \alpha_k(x) P_k$$
  여기서 $\alpha_k(x)$는 정규화된 가중치입니다. 최종 입력 의존적 프롬프트는 $P(x) = [P_f(x); P_c]$입니다.

3. 2단계 분리 최적화 (Two-Stage Decoupled Optimization)

생성 목적 함수가 위험 결정 경계를 변경하는 것을 방지하기 위해 훈련은 두 개의 별도 단계로 진행됩니다:

• 단계 I (라우터 학습): 라우터는 조립도 낮은 분류를 위한 교차 엔트로피(CE)와 조립도가 높은 다중 레이블 분류를 위한 이진 교차 엔트로피(BCE)를 사용하여 훈련됩니다. 백본 모델은 동결된 상태를 유지합니다.
• 단계 II (프롬프트 최적화): 백본과 훈련된 라우터는 동결됩니다. 오직 프롬프트 파라미터($P_c$ 및 $\{P_k\}$)만이 선호도 삼중항(safe/helpful vs. unsafe/low-quality)에 대한 직접 선호 최적화(DPO)를 사용하여 최적화됩니다. 간섭을 줄이기 위해 저자들은 **교대 그래디언트 업데이트(alternating gradient updates)**를 채택합니다: $P_c$를 업데이트할 때 $P_f(x)$에 대한 그래디언트는 마스킹(stop-gradient)되며, 그 반대도 마찬가지입니다.

4. 입력 적응형 추론 (Input-Adaptive Inference)

추론 시, 라우터는 먼저 안전 확률 $p_{safe}(x)$를 계산합니다.

• 무해한 입력: 만약 $p_{safe}(x) > \tau_s$ (안전 임계값)라면, 시스템은 모든 안전 프롬프트를 우회하고 동결된 백본으로부터 직접 생성하여 일반적인 유용성을 보존합니다.
• 위험한 입력: 만약 $p_{safe}(x) \leq \tau_s$ 라면, 시스템은 안전 브랜치를 활성화하여 계층적 프롬프트 $P(x)$를 구성하고 생성을 조건화합니다.

주요 기여

1. 트레이드오프의 경험적 검증: 저자들은 공유된 조립도 낮은 프롬프트가 더 강력한 안전성을 제공하지만 유용성은 낮아지는 반면, 라우팅된 파인-그레인드 프롬프트 혼합물은 유용성은 개선하지만 안전 이득은 작아진다는 점을 입증했습니다.
2. HiRoute 프레임워크: 이들은 공유된 조립도 낮은 프롬프트(안정적인 안전 경계를 위해)와 동적으로 라우팅되는 파인-그레인드 전문가(위험 특이적 가이드를 위해)를 결합하는 새로운 아키텍처를 제안합니다. 이 설계는 안전 게이트를 통해 무해한 입력에 대한 불필요한 개입을 방지합니다.
3. 포괄적 평가: 이 방법은 세 가지 지시어 튜닝 모델(Mistral-7B, Vicuna-7B, Zephyr-7B)과 여러 벤치마크를 통해 검증되었으며, 높은 안전성과 유용성을 유지하면서 일반적인 유용성을 보존함을 보여줍니다.

실험 결과

실험은 Base 모델, RPO (Robust Prompt Optimization), DRO (Dynamic Risk Optimization), ACD (Adversarial Contrastive Decoding)를 포함한 베이스라인들을 대상으로 세 가지 모델에서 수행되었습니다.

• 안전성 및 유용성: HiRoute는 모든 모델에서 가장 높은 평균 안전율(예: Mistral에서 93.2%, Vicuna에서 97.7%, Zephyr에서 94.8%)을 달성하는 동시에 가장 높거나 이에 근접한 안전 응답 유용성 점수를 달성했습니다. 이는 가장 강력한 베이스라인보다 안전성 측면에서 1.3~2.9 퍼센트 포인트 더 높은 성능을 보였습니다.
• 일반 유용성: 일반적인 능력을 저하시키는 방법들과 달리, HiRoute는 GSM8K(수학적 추론), MT-Bench(지시 이행), TruthfulQA에서의 성능을 보존했습니다. 예를 들어, Mistral의 경우 안전성을 53.2%에서 93.2%로 높이면서도 GSM8K 정확도를 48.0%로 유지했습니다.
• 과잉 거절: HiRoute는 무해한 입력에 대한 과잉 거절(XSTest)을 크게 줄였습니다. Mistral에서 RPO(40.5%)와 DRO(18.0%)에 비해 훨씬 개선된 2.5%의 과잉 거절률을 달성했습니다.
• 절제 연구 (Ablation Studies):
  • 훈련 전략: 통합 교대 훈련(일관된 구성된 컨텍스트 내에서 프롬프트 업데이트)은 순차적 훈련("Coarse after fine")보다 현저히 우수한 성능을 보였으며, 평균 안전성을 74.8%에서 93.2%로 향상시켰습니다.
  • 프롬프트 길이: 조립도 낮은 프롬프트와 파인-그레인드 프롬프트의 15/5 분할(총 20)이 안전성과 유용성 사이의 최적의 균형을 제공했습니다.
  • 임계값: 안전 게이팅 임계값 $\tau_s$를 높이면 일반 능력에 미치는 영향은 최소화하면서 안전 제어를 강화할 수 있었습니다.

의의 및 주장

본 논문은 HiRoute가 LLM 정렬에서의 안전성과 유용성 사이의 긴장 관계를 효과적으로 해결한다고 주장합니다. 위험 식별과 응답 생성을 분리하고 계층적이며 입력 적응적인 프롬프트 구조를 채택함으로써, 이 방법은:

• 공격적이고 일률적인 거절에 의존하지 않고도 안정적이고 카테고리 불가지론적인 안전 경계를 구축합니다.
• 유해한 요청에 대해 건설적이고 위험 특이적인 가이드 및 안전한 대안을 제공합니다.
• 무해한 입력이 안전 프롬프팅을 완전히 우회할 수 있도록 하여 과잉 거절을 최소화합니다.
• 백본 모델을 동결하고 소수의 프롬프트 파라미터와 경량 라우터만을 최적화함으로써 이러한 목표를 매개변수 효율적으로 달고합니다.

저자들은 입력 의존적 계층 프롬프팅이 안전성, 응답 품질, 매개변수 효율성 및 일반 유용성의 균형을 맞추는 견고한 솔루션을 제공한다고 결론짓습니다. 또한, 분포 외(out-of-distribution) 복합 위험에 대한 라우터의 정확도 문제와 현재 단일 턴 텍스트 상호작용에 집중되어 있다는 한계를 인정하며, 오픈 셋 위험 인식 및 멀티모달 환경에서의 향후 연구를 제시합니다.