Test-Time Safety Alignment

Baturay Saglam과 Dionysis Kalogerias의 논문 "Test-Time Safety Alignment"에 대한 상세한 기술 요약입니다.

1. 문제 정의

대형 언어 모델 (LLM) 은 유해한 요청을 거부하도록 안전성 훈련 (예: RLHF, DPO) 을 받습니다. 그러나 적대적 프롬프트 (재일브레이크) 는 모델의 훈련 분포와 무제한 입력 공간 사이의 간극을 악용하여 이러한 방어를 우회할 수 있습니다. 이러한 프롬프트는 종종 유해한 콘텐츠를 유도하기 위해 구조적 프레임 (역할극, 은폐) 을 사용합니다.

기존의 테스트 시간 방어책은 구조적 한계에 직면해 있습니다:

• 프롬프트 수준 전략은 이산 토큰을 조작하지만, 의미적 공격에는 취약합니다.
• **활성화 조종 (Activation steering)**은 숨겨진 상태에 대한 접근이 필요하며, 종종 일반화되지 않을 수 있는 고정된 보정 세트에 의존합니다.
• 가드 모델은 별도의 추론이 필요하며, 생성을 조종하기보다는 "탐지 후 필터링" 패러다임을 따르는 경우가 많습니다.

핵심 질문: 정렬된 모델이 매끄러운 연속이 아닌 이분모 "거부 또는 준수" 분포를 생성할 때조차, 입력 단어 임베딩을 서브-렉시컬 (sub-lexical) 방식으로 최적화하여 (가시적 토큰을 변경하지 않고 연속 벡터 표현을 수정) 정렬된 모델을 안전한 출력으로 유도할 수 있을까요?

2. 방법론: 테스트 시간 안전 정렬 (TSA)

제안된 방법인 TSA는 모델 가중치를 수정하지 않고, 보조 훈련이 필요 없으며, LLM 을 블랙박스로 취급한 채 완전히 테스트 시간에 작동합니다. 이는 프롬프트의 임베딩 행렬을 직접 최적화합니다.

A. 목적 함수

목표는 생성된 텍스트 $y = f(X)$의 유해성을 최소화하는 것입니다. 여기서 $f$는 LLM 이고 $X$는 프롬프트 임베딩 행렬입니다.

• 유해성 오라클 ($h$): 이 방법은 OpenAI Moderation API를 블랙박스 오라클로 사용합니다. 이는 13 가지 유해 범주 (예: 폭력, 자해, 혐오) 에 대해 출력물을 점수화합니다.
• 손실 함수: $\Phi(X) = \max_{c} s_c(f(X))$이며, 여기서 $s_c$는 범주 $c$에 대한 점수입니다. 목표는 이 최대 점수를 최소화하여 응답이 모든 차원에서 안전하도록 보장하는 것입니다.

B. 0 차 그라디언트 추정

LLM($f$) 과 Moderation API($h$) 가 블랙박스이므로 역전파는 불가능합니다. TSA 는 0 차 그라디언트 추정을 사용합니다:

1. 교란: 각 반복 $k$에서 현재 임베딩 행렬 $X_k$에 가우스 노이즈 $U$를 스케일 $\mu$로 곱하여 더함으로써 교란시킵니다: $X_k + \mu U$.
2. 유한 차분: 그라디언트는 교란된 입력의 손실과 기본 입력의 손실을 비교하여 추정됩니다:
   $$g_k \approx \frac{1}{N} \sum_{i=1}^{N} \frac{\Phi(X_k + \mu U_i) - \Phi(X_k)}{\mu} U_i$$
   여기서 $N$은 몬테카를로 샘플의 수입니다.
3. 업데이트: 임베딩은 경사 하강법을 통해 업데이트됩니다: $X_{k+1} \leftarrow X_k - \eta g_k$.

C. 정규화

최적화된 임베딩이 유효하게 유지되고 잠재 공간의 훈련되지 않은 영역으로 이탈하지 않도록 하기 위해:

• 그라디언트 정규화: 업데이트 단계는 단계 크기를 제어하기 위해 정규화된 그라디언트 방향 ($g_k / \|g_k\|_2$) 을 사용합니다.
• 코사인 유사도 제약: 유사도가 임계값 $\kappa$ 아래로 떨어지면 업데이트된 임베딩은 원래 임베딩 $X_0$ 주변의 코사인 볼 위로 투영됩니다. 이는 교란이 서브-렉시컬임을 보장합니다: 최적화된 벡터는 원래 토큰으로 디코딩되지만, 해당 토큰의 보로노이 셀 내의 연속 표현은 변경됩니다.

3. 주요 기여

1. 정렬된 모델로의 확장: 기존에 기본 모델의 독소 제거에 대해 입증된 임베딩 수준의 제어가 안전성 정렬된 지시 미세 조정 모델에서도 효과적으로 작동함을 보여줍니다. 이러한 모델들은 근본적으로 다른 (이분모) 출력 분포를 가집니다.
2. 서브-렉시컬 제어: 토큰 자체를 변경하지 않고 해당 토큰의 최근접 이웃 셀 내부의 연속 임베딩 벡터를 최적화하는 것만으로도 모델의 행동을 유해한 것에서 안전한 것으로 전환할 수 있음을 증명합니다. 이는 임베딩을 단순한 이산 조회로 보는 관점에 도전합니다.
3. 블랙박스 최적화: 이 방법은 모델 내부 (그라디언트, 숨겨진 상태, 로짓) 에 대한 접근이 필요하지 않습니다. 입력 임베딩과 블랙박스 Moderation API 에 대한 접근만 필요합니다.
4. 무조건적 적용: 이 방법은 의도와 관계없이 모든 프롬프트에 적용됩니다. 거부를 강요하지 않으며, 가능하다면 안전하고 실질적인 응답을 생성하도록 모델을 조종하거나, 필요하다면 안전한 거부를 유도합니다.

4. 실험 결과

이 방법은 두 가지 벤치마크에 걸쳐 다섯 가지 지시 미세 조정 모델(1B 에서 20B 파라미터까지: Gemma 3, Phi-3.5, Llama 3.1, Qwen3, GPT-OSS) 에서 평가되었습니다:

• WildJailbreak: 적대적 프롬프트 (유해 및 무해).
• HarmBench: 직접적인 유해 쿼리.

주요 발견:

• 안전성 제거: TSA 는 표준 벤치마크에서 대부분의 모델에 대해 플래그가 지정된 (유해한) 응답 수를 0으로 줄였습니다. 예를 들어, Llama 3.1-8B 에서는 플래그가 지정된 유해 프롬프트를 218 개에서 0 개로 줄였습니다.
• 수렴: 최적화는 일반적으로 1~2 개의 그라디언트 단계 내에서 수렴합니다.
• 무해성 보존: 적대적 무해 프롬프트에 대해 TSA 는 Moderation 점수를 증가시키지 않았으며, 이는 "포괄적 거부" 편향을 유발하거나 유용성을 저하시키지 않음을 확인시켜 줍니다.
• 베이스라인 대비 비교: TSA 는 SmoothLLM(랜덤 토큰 교란), AdaSteer(활성화 조종), RESTA(임베딩 노이즈) 와 같은 베이스라인을 크게 능가했습니다. 베이스라인은 종종 의미적 재일브레이크에서 실패하거나 직접적인 유해 쿼리에서 플래그 지정률을 증가시켰습니다.
• 효율성: 주요 지연 비용은 모델 추론이 아닌 Moderation 오라클에 대한 API 호출입니다.

5. 중요성 및 함의

• 새로운 제어 표면: 이 논문은 임베딩 레이어가 이전보다 생각했던 것보다 더 풍부한 제어 표면을 제공함을 드러냅니다. 작고 서브-렉시컬한 교란은 정렬된 모델의 복잡한 안전 경계를 탐색할 수 있습니다.
• 배포 타당성: "드롭인" 안전 계층으로서 TSA 는 모델 재훈련이나 미세 조정이 필요하지 않습니다. 입력 임베딩에 접근할 수 있는 모든 모델 (예: 오픈 가중치 모델) 과 공개 Moderation API 를 결합하여 배포할 수 있습니다.
• 한계:
  • 지연 시간: 반복적 최적화는 프롬프트당 지연 시간 (초 단위) 을 추가하지만, 이는 주로 API 왕복 시간에 의해 지배됩니다.
  • 모델 의존성: 효과성은 모델의 기존 안전성 훈련 품질에 비례합니다. 얇은 정렬을 가진 모델 (예: Gemma 3-1B) 은 성공률이 낮았는데, 이는 최적화기가 근처의 안전한 방향을 찾을 수 없었기 때문입니다.
  • API 접근: 입력 임베딩에 대한 접근 (순수 API 게이트가 적용된 폐쇄형 모델에서는 사용 불가) 과 Moderation API 가 필요합니다.

결론적으로, Test-Time Safety Alignment는 연속 입력 표현을 최적화함으로써 적대적 공격을 무력화하는 견고하고 모델에 구애받지 않는 메커니즘을 제공하며, 재훈련 없이 이론적 안전성과 실제 배포 사이의 간극을 효과적으로 메워줍니다.