Before the Last Token: Diagnosing Final-Token Safety Probe Failures

기술 요약: 마지막 토큰 이전: 최종 토큰 안전 프로브 실패 진단

문제 제기
SafeSwitch 와 같은 대규모 언어 모델 (LLM) 의 현재 안전 모니터링 시스템은 종종 "최종 토큰 프로브 (final-token probes)"에 의존합니다. 이러한 시스템은 생성 시작 전에 프롬프트의 프리필 (prefill) 단계에서 마지막 토큰의 숨겨진 상태를 기반으로 분류기를 훈련시켜 유해한 콘텐츠를 탐지합니다. 본 논문은 이 아키텍처에서 치명적인 실패 모드를 식별합니다. 즉, 이러한 프로브는 "클린 (clean)" 유해 프롬프트에 대해 높은 재현율 (recall) 을 달성하지만, "자일브레이크 (jailbreak)" 프롬프트는 자주 탐지하지 못합니다. 자일브레이크는 적대적 프레임, 역할극, 또는 은폐 기법 내에 유해한 요청을 내포합니다. 저자들은 이러한 래퍼 (wrappers) 가 훈련 중 학습된 특정 "유해 대 무해" 대비 방향에서 마지막 토큰 표현을 이동시켜, 유해한 콘텐츠가 프롬프트에 존재함에도 불구하고 프로브가 유해한 요청을 놓치게 만든다고 주장합니다. 이러한 실패는 단순히 용량 (capacity) 의 문제가 아닙니다. 프로브의 병목 (bottleneck) 너비를 늘려도 탐지 격차를 해결하지 못하기 때문입니다.

방법론
본 연구는 Llama-3.1-8B-Instruct, Mistral-7B-Instruct-v0.1, OLMo3-7B-Instruct 의 세 가지 지시 튜닝 모델을 평가했습니다.

1. 프로브 훈련: SafeSwitch 스타일의 프로브는 원래 SafeSwitch 훈련 분할에서 가져온 "클린" 유해 및 무해 프롬프트로만 훈련되었습니다. 이러한 프로브는 최종 프롬프트 토큰 숨겨진 상태를 읽은 후, 병목 (초기 너비 $w=64$) 을 통해 이를 투영하여 안전성을 예측합니다.
2. 평가: 프로브는 다음에서 테스트되었습니다:
   • 클린 홀드아웃 유해 프롬프트 (SorryBench).
   • HarmBench 의 자일브레이크 템플릿으로 감싸진 동일한 유해 프롬프트.
   • 위양성 (false positives) 을 측정하기 위한 무해한 안전 인접 프롬프트 (XSTest).
   • 일반적인 분포 변화를 배제하기 위한 직접적인 유해 벤치마크 (AdvBench).
3. 기하학적 분석: 저자들은 프로브의 표현적 부분 공간 (representational subspace) 을 분석했습니다. 그들은 다양한 집합 (예: 놓친 자일브레이크 대 무해 프롬프트) 에 대한 평균 차이 벡터를 계산하고, 프로브의 행 공간 (row-space) 이 방향의 노름 (norm) 중 포착하는 비율인 "프로브 가시 에너지 (probe-visible energy)"를 측정했습니다 (SVD 를 통해).
4. 궤적 분석: 최종 토큰에만 의존하는 대신, 저자들은 프리필 단계 동안 토큰별 숨겨진 상태 궤적을 분석했습니다. 그들은 유해한 증거가 시퀀스 초기에 나타났다가 최종 토큰에 도달할 때 소멸하는지 여부를 조사했습니다.
5. 진단 모델링: 궤적 정보가 놓친 탐지를 복구할 수 있는지 테스트하기 위해, 저자들은 클린 훈련 데이터에 간단한 PCA-HMM(주성분 분석 - 은닉 마르코프 모델) 을 적합시켰습니다. 이 모델은 클린 유해 대 무해 프롬프트의 궤적 역학을 학습했으며, 자일브레이크 데이터에 대한 재튜닝 없이 자일브레이크 궤적에 적용되었습니다.

주요 결과

• 자일브레이크에서의 프로브 실패: 최종 토큰 프로브는 클린 유해 프롬프트에서 높은 재현율 (~95%) 을 유지했지만, 자일브레이크에서는 크게 감소했습니다 (모델에 따라 ~63% 에서 ~87% 까지). 반면, 무해한 안전 인접 프롬프트에서는 23%–36% 의 비과소 위양성을 발생시켰습니다.
• 용량이 병목이 아님: 프로브 병목 너비를 64 에서 1024 로 늘려도 자일브레이크 탐지가 신뢰할 수 있게 개선되지 않았습니다. OLMo3 의 경우 성능이 실제로 저하되었습니다. 이는 실패가 표현 용량의 부족이 아니라 신호가 위치한 기하학적 문제임을 시사합니다.
• 기하학적 불일치: 놓친 자일브레이크와 무해 프롬프트를 분리하는 방향 ($d_{miss}$) 은 클린 훈련 대비 ($d_{harm}$) 에 비해 가장 낮은 프로브 가시 에너지 (약 2.4%–3.5%) 를 가졌습니다. 프로브의 행 공간은 무해한 입력과 구별하는 놓친 자일브레이크의 특정 특징들과 사실상 직교합니다.
• 위치적 실패: 놓친 자일브레이크에서의 유해한 증거는 종종 중간 토큰 위치 (특히 내장된 유해 요청 근처) 에서 가시적이지만, 최종 토큰에 도달할 때 사라지거나 "붕괴"됩니다.
• 단순 풀링의 실패: 모든 토큰 위치에서 최대 점수를 취하는 것은 실행 가능한 해결책이 아닙니다. 이는 무해한 XSTest 프롬프트의 100% 에서 발동 (위양성) 하기 때문이며, 이는 높은 중간 점수 자체가 의미론적으로 신뢰할 수 없음을 나타냅니다.
• 궤적 복구: 클린 데이터로만 훈련된 PCA-HMM 진단 모델은 최종 토큰 놓침의 대다수를 성공적으로 복구했습니다 (97%–98% 복구율). 또한 단순 풀링에 비해 XSTest 에서 훨씬 낮은 위양성률 (6%–15%) 을 유지했습니다. 놓친 자일브레이크 방향은 최종 토큰 프로브 행 공간보다 궤적 PCA 부분 공간에서 훨씬 높은 에너지를 보여주었습니다.

의의 및 주장
본 논문은 배포된 방어 메커니즘이 아닌 진단 연구로 자리매김합니다. 주요 기여는 최종 토큰 프로브가 실패하는 이유를 규명하는 것입니다: 그들은 클린 훈련 분포에는 민감하지만 자일브레이크 래퍼가 유도하는 특정 기하학적 변화에는 맹목적인 "단축키 같은" 판독에 의존합니다.

저자들은 다음과 같이 주장합니다:

1. 래퍼가 유도한 이동이 표현을 프로브의 민감한 부분 공간 밖으로 이동시키기 때문에, 최종 토큰 프로브는 래핑된 유해 요청을 탐지하는 데 불충분합니다.
2. 안전 관련 정보는 종종 최종 토큰 이전의 숨겨진 상태 궤적에 존재하지만, 최종 판독이 이루어질 때까지 손실됩니다.
3. 클린 데이터로만 훈련된 간단한 궤적 인식 모델 (PCA-HMM 등) 은 최종 토큰 프로브에 효과적인 진단 보완재로 작용할 수 있습니다. 이러한 모델은 프롬프트 전체에 걸친 안전 신호의 진화를 분석하여 놓친 탐지를 복구하며, 단순 토큰 풀링의 치명적인 위양성률 없이 프롬프트의 내부 상태에 대한 더 견고한 관점을 제공합니다.

이 연구는 현재 프로브 기반 안전 시스템의 한계를 더 잘 이해하고 완화하기 위해, 안전 진단을 단일 지점 판독에서 궤적 인식 분석으로 전환할 것을 촉구합니다.