기술 요약: Geometry-Lite: 계층별 마진 기하학을 통한 해석 가능한 안전성 프로빙

문제 정의

대규모 언어 모델 (LLM) 이 대규모로 배포됨에 따라, 프롬프트 수준의 유해성 탐지는 중요한 안전 구성 요소가 되었습니다. 단일 계층 스냅샷 (일반적으로 최종 프롬프트 토큰) 에 선형 분류기를 적용하는 잠재 상태 프로브는 표준 베이스라인이지만, 이는 단일 계층이 안전 관련 신호를 요약하기에 충분하다는 가정에 의존합니다. 이 가정은 두 가지 이유로 재검토가 필요합니다:

1. 경계 사례 (Borderline Cases): 표준 벤치마크에서의 집계 성능은 종종 모호한 프롬프트에서의 행위를 가립니다. 이는 낮은 위양성률로 운영되는 안전 시스템에 있어 중요합니다.
2. 분포 외 (OOD) 일반화: 벤치마크 분할에서 잘 수행되는 모델은 전이되지 않는 단축 특징 (shortcut features) 에 의존할 수 있습니다. 최근 연구에 따르면 활성화 기반 안전성 프로브는 교차 데이터셋 평가 하에서 상당한 성능 저하를 겪습니다.

핵심 연구 질문은 단순히 여러 계층이 도움이 되는지가 아니라, 의사결정이 중요한 영역 (낮은 위양성률) 과 벤치마크 변화 하에서 어떤 형태의 다계층 증거가 유용한지입니다. 구체적으로, 저자들은 안전 신호가 계층 전체에 분포되어 있는지, 계층 간 이동 (변위) 에 의해 주도되는지, 아니면 안전 경계 대비 마진 프로파일의 지속적 위치에 의해 정의되는지 조사합니다.

방법론: Geometry-Lite

저자들은 레이블된 기하학적 축을 따라 다계층 안전 신호를 분해하도록 설계된 컴팩트한 프롬프트 수준 프로브인 Geometry-Lite를 소개합니다. 이 방법은 세 단계로 작동합니다:

1. 계층별 마진 축소

각 트랜스포머 계층 $\ell$에 대해, 최종 프롬프트 토큰의 $D$차원 잠재 상태 $h^{(\ell)}(x)$는 세 가지 보완적인 기하학적 읽기 (readouts) 를 사용하여 부호 있는 스칼라 마진으로 축소됩니다:

• 센트로이드 거리: 클래스 조건부 평균 ($\mu^+$는 안전, $\mu^-$는 안전하지 않음) 까지의 유클리드 거리 차이.
  $$m^{cent}_\ell(h) = \|h - \mu^-_\ell\|_2 - \|h - \mu^+_\ell\|_2$$
• k-NN 국소 이웃: 안전 및 안전하지 않은 학습 풀에서 $k$개의 가장 가까운 이웃까지의 평균 코사인 거리 차이. 이는 유일한 비선형 읽기입니다.
  $$m^{knn}_\ell(h) = \bar{c}^-_k(h) - \bar{c}^+_k(h)$$
• 지도 선형 경계: 표준화된 특징에 대해 계층별로 적합된 L2 정규화 로지스틱 회귀 경계로, 부호 있는 점수를 생성합니다.
  $$m^{lin}_\ell(h) = w^\top_\ell z + b_\ell$$

이 과정은 잠재 상태 궤적 $H(x) \in \mathbb{R}^{L \times D}$를 세 가지 스칼라 마진 궤적 $m^G(x) \in \mathbb{R}^L$ (여기서 $G \in \{cent, knn, lin\}$) 로 변환합니다.

2. 명명된 축을 통한 프로파일 요약

각 마진 프로파일은 계층별 신호의 서로 다른 형태를 분리하기 위해 세 가지 명명된 축을 따라 요약됩니다:

• 마진 수준: 지속적 경계 위치를 포착합니다 (예: 끝점 마진, 극단적 값, 음수 측 점유 및 면적).
• 계층 간 변화: 유한 차분 ($\Delta m_\ell = m_{\ell+1} - m_\ell$) 을 통해 국소 이동을 분리하여 절대 위치와 운동을 구분합니다.
• 구조적 형태: 거친 프로파일 속성을 포착합니다 (예: 첫 번째 경계 교차, 변화의 변동성, 총 경로 길이 및 방향 일관성).

이러한 축은 39 차원 특징 벡터 $\phi(x)$를 생성합니다 (기하학당 13 개 특징).

3. 분류

39 차원 벡터는 L2 정규화 로지스틱 회귀 분류기에 입력됩니다. 파라미터 수는 백본 모델 크기와 숨겨진 차원에 독립적입니다.

베이스라인 및 비교

이 연구는 Geometry-Lite 를 다음과 비교합니다:

• 단일 계층 프로브: 검증 과정을 통해 선택된 최성능 계층.
• MultiLayer-DIM: 계층 전체에 걸쳐 평균 차이 방향으로 투영된 것을 집계합니다.
• MultiLayer-Linear: 계층 전체에 걸쳐 원시 지도 선형 점수를 집계합니다.
• TaT-Disp-LSTM: LSTM 으로 계층별 변위를 모델링하는 궤적 기반 베이스라인.

주요 기여

1. Geometry-Lite 프로브: 계층별 잠재 상태를 세 가지 기하학적 읽기 하에서 부호 있는 마진 프로파일로 변환하고, 이를 명명된 축 (경계 위치, 변화, 형태) 을 따라 요약하는 컴팩트하고 해석 가능한 프로브.
2. 다계층 증거의 분해: 이 프레임워크는 다계층 신호에 대한 세 가지 설명을 분리합니다: 정보 계층에 걸친 집계, 경계 상대적 마진 위치, 명시적 계층 간 운동.
3. 실증적 발견:
   • 신호 집중: 안전 증거는 주로 지속적 경계 위치 기하학(마진 수준 및 점유) 을 통해 표현됩니다. 계층 간 운동 (드리프트) 은 집계 AUROC 에 미미하게 기여합니다.
   • 드리프트 유용성: 드리프트는 집계된 AUROC 에는 거의 기여하지 않지만, 변화된 낮은 FPR 임계값 하에서 희소하고 재현율 지향적인 보정을 제공합니다.
   • 읽기 안정성 트레이드오프: 최적화된 선형 경계는 분포 내 (in-distribution) 에서 더 날카롭지만, 벤치마크 변화 (LOBO 프로토콜) 하에서 클래스 조건부 평균 기하학 (센트로이드/DIM) 이 더 신뢰할 수 있게 분리를 유지합니다.

실험 결과

평가에는 9 개의 지시 튜닝 백본 (Llama, Gemma, Qwen 계열; 1.2B–70B) 과 7 개의 안전성 벤치마크 (XSTest, WildJailbreak, JBB-Behaviors, DoNotAnswer, BeaverTails, PKU-SafeRLHF, ToxicChat) 가 포함되었습니다.

• 분포 내 성능: Geometry-Lite(AUROC 0.955) 와 MultiLayer-Linear(0.954) 는 최성능 단일 계층 프로브 (0.941) 보다 우수하며, MultiLayer-DIM(0.911) 과 TaT-Disp-LSTM(0.945) 보다 현저히 우수합니다. Geometry-Lite 는 해석 가능성을 제공하면서 원시 점수 스택킹과 일치합니다.
• 의사결정 중요 영역: 낮은 위양성률 (3% 및 5% FPR) 에서 다계층 이점이 더 두드러집니다. Geometry-Lite 는 5% FPR 에서 최성능 단일 계층 프로브 대비 True Positive Rate(TPR) 를 약 0.06–0.08 향상시킵니다.
• 일반화 (LOBO): 벤치마크 하나를 제외하고 평가 (Leave-One-Benchmark-Out) 하에서 지도 선형 방법과 클래스 평균 방법 (MultiLayer-DIM) 간의 성능 격차가 좁아집니다. MultiLayer-DIM(0.889) 은 MultiLayer-Linear(0.889) 와 일치하며 많은 셀에서 최성능 단일 계층 프로브 (0.842) 를 초과합니다.
• 특성 제거 실험:
  • 크기 특징만으로도 0.001–0.002 AUROC 이내에서 전체 39 개 특징 모델과 일치합니다.
  • 구조적 특징과 Neg-drift 특징만으로는 현저히 약합니다.
  • 읽기 제거 실험: 분포 내에서는 지도 선형 경계가 가장 강력합니다. 어려운 벤치마크 변화 하에서는 센트로이드 읽기가 가장 안정적인 단일 분기가 되어, 가장 어려운 홀드아웃 부분집합에서 선형 경계를 능가합니다.

중요성 및 주장

이 논문은 프롬프트 수준의 안전 증거가 주로 계층 간 운동 신호가 아니라 지속적 계층별 마진 기하학이라고 주장합니다. 이 신호의 유용한 구성 요소는 안전 경계 대비 마진의 위치와 계층 전체에서의 지속성입니다.

이 연구는 안전성 프로빙에서의 중요한 트레이드오프를 강조합니다:

• 최적화된 경계(선형 분류기) 는 분포 내에서는 강력한 분리를 포착하지만 분포 변화 하에서는 저하됩니다.
• 클래스 평균 기하학(센트로이드/DIM) 은 분포 내에서는 덜 날카롭지만 벤치마크 변화에 더 강건합니다.

따라서 저자들은 다계층 안전성 프로브가 분포 내 AUROC 로만 평가되는 것이 아니라, 변화 하에서 어떤 기하학적 편향이 전이되는지에 따라 평가되어야 한다고 주장합니다. Geometry-Lite 는 특히 위양성이 비용이 큰 의사결정 중요 영역에서 이러한 구성 요소와 편향을 가시화하는 도구로 작용합니다.

한계: 본 연구는 수동으로 설계된 마진 요약, 단일 턴 안전 프롬프트, 이진 안전 레이블로 제한됩니다. 향후 연구는 학습된 프로파일 함수, 멀티 턴 모니터링, 그리고 다른 프롬프트 수준 작업에의 적용을 탐구할 수 있습니다.