기술 요약: 마할라노비스 코사인 유사도와 선형 프로브의 비교

문제 정의

선형 프로브(Linear probes)는 신경망 표현으로부터 개념을 추출하기 위해 해석 가능성 연구에서 표준적으로 사용되는 도구입니다. 그러나 이들의 전이(transfer) 동작은 매우 취약한 것으로 알려져 있습니다. 분포 내(in-distribution, ID) 데이터로 학습된 프로브는 부정(negation)이나 프롬프트 형식 변경과 같은 미세한 변화에도 분포 외(out-of-distribution, OOD) 시나리오에서 성능이 저하되는 경우가 많습니다. 핵심적인 과제는 어떤 프로브가 일반화될 것인지 결정하는 것과, 하나의 분포에서는 유사해 보였던 프로브들이 왜 다른 분포에서는 서로 달라지는지를 이해하는 것입니다.

프로브를 비교하는 현재의 방법들은 일반적으로 유클리드 코사인 유사도(Euclidean cosine similarity, ECS)에 의존합니다. 그러나 ECS는 모든 차원을 동일하게 취급하며, 데이터가 실제로 크게 변하는 차원이 프로브 성능에 미치는 영향력을 무시합니다. Ying 등(2026)은 마할라노비스 코사인 유사도(Mahalanobis cosine similarity, MCS)—내적을 OOD 데이터 공분산에 의해 재가중치하는 방식—가 진실성(truthfulness) 데이터셋에 대해 OOD AUROC를 선형적으로 예측($R^2=0.98$)한다는 것을 경험적으로 관찰했으나, 이 발견은 이론적 설명이 부족했으며 다양한 모델, 레이어, 개념 영역에 걸쳐 검증되지 않았습니다.

방법론

저자들은 MCS와 OOD 일반화 성능 사이의 관계를 설명하고 검증하기 위한 이론적 및 경험적 프레임워크를 제안합니다.

이론적 프레임워크

본 논문은 특정 프로브의 OOD AUROC와 OOD 데이터로 학습된 참조 프로브 간의 MCS 사이의 폐쇄형(closed-form) 관계를 도출합니다. 이 유도는 세 가지 표준 통계적 가정을 기반으로 합니다:

1. 균형 잡힌 클래스: $Pr(Y=0) = Pr(Y=1) = 0.5$.
2. 투영 가우시안성(Projection Gaussianity): 특징 벡터를 임의의 프로브 방향으로 투영한 결과는 각 클래스에 대해 가우시안 분포를 따릅니다.
3. 공분산 구조: 전체 공분산($\Sigma_{tot}$)은 풀링된 클래스 내 공분산($\Sigma_{pool}$)과 클래스 간 성분으로 분해됩니다.

이러한 가정 하에, 저자들은 다음을 증명합니다:

• AUROC는 프로브의 신호 대 잡음비(SNR)인 $s$의 시그모이드 형태 함수입니다: $AUROC = \Phi(s/\sqrt{2})$.
• MCS(특히 피셔 판별 방향에 대한 MCS)는 $s$의 소프트-시그(soft-sign) 형태 함수입니다: $MCS \approx s/\sqrt{4+s^2}$ (대규모 피셔 거리 $z_{max}$의 극한에서).

핵심적인 이론적 통찰은 이 두 함수를 결합하면 각각의 비선형성(시그모이드와 소프트-시그)이 상쇄되어, AUROC와 MCS 사이에 거의 선형적인 관계가 형성된다는 점입니다. 이 직선의 기울기는 피셔 거리가 클 때 이론적으로 약 $1/\sqrt{\pi}$로 예측됩니다.

경험적 평가

저자들은 이 관계에 대한 경험적 검증을 다음과 같이 확장했습니다:

• 모델: Llama-3.3-70B, Llama-3.1-8B, Qwen2.5-7B.
• 레이어: Llama-70B 내의 다양한 레이어(예: 20, 33, 50, 65).
• 도메인: 진실성, 성별 분류 및 일반 NLP 벤치마크를 아우르는 24개 데이터셋.
• 베이스라인: 표준 유클리드 코사인 유사도(ECS)와의 비교 및 서로 다른 공분산 추정기(전체 vs. 풀링된 공분산) 사용.

주요 기여

1. 선형성에 대한 이론적 증명: 본 논문은 균형 잡힌, 가우시안 투영된 클래스들에 대해 OOD AUROC와 MCS가 선형 관계에 있음을 보여주는 폐쇄형 증명을 제공합니다. 이는 이전에 관찰된 높은 경험적 상관관계($R^2 \approx 0.98$)를 설명합니다.
2. 실패 모드의 특성화: 이론은 이 선형성이 깨지는 조건을 명시적으로 예측하며, 이는 경험적으로 검증되었습니다:
   • 전체 공분산($\Sigma_{tot}$) 대신 풀링된 공분산($\Sigma_{pool}$)을 사용하면 선형성이 파괴됩니다.
   • 작은 피셔 거리($z_{max}$)는 보편적인 선이 아닌 작업 특화적인 기울기를 초래합니다.
   • 클래스 불균형은 기울기를 가파르게 만듭니다.
   • 피셔 방향에서 크게 벗어난 프로브(예: 단순 평균 차이 프로브)는 저하된 선형성을 보입니다.
3. 경험적 일반화: 연구는 다양한 모델, 레이어, 개념 도메인에 걸쳐 MCS-AUROC 선형성을 검증하였으며, 모든 테스트 조건에서 $R^2 > 0.93$을 달성했습니다. 반면, ECS 성능은 일반 NLP 벤치마크에서 $R^2$가 0.06까지 떨어지는 등 크게 저하되었습니다.
4. 추정기에 대한 견고성: 저자들은 이 결과가 전체 샘플 공분산, Ledoit-Wolf 수축(shrinkage)법에 대해서는 유효하지만, 대각 근사(diagonal approximations)를 사용할 경우에는 실패함을 입증하여 오프-대각(off-diagonal) 공분산 구조의 필요성을 확인했습니다.

결과

• 높은 예측력: 8가지 실험 조건(다양한 모델, 레이어, 도메인) 전체에서 OOD AUROC에 대한 MCS의 선형 적합은 $R^2 \geq 0.93$을 나타냅니다.
• ECS 대비 우월성: ECS는 일관되게 낮은 성능을 보였으며, $R^2$ 값은 0.057(일반 NLP)에서 0.723(성별) 사이의 범위를 보였으나, MCS는 높은 예측력을 유지했습니다.
• 이론적 일치: AUROC-MCS 곡선의 경험적 기울기는 이론적 중심 기울기인 $1/\sqrt{\pi} \approx 0.564$와 밀접하게 일치합니다. 편차는 주로 포화 꼬리 부분(높은 AUROC/MCS 값)에서 발생하며, 이는 이론과 일치합니다.
• 실패 모드 검증: 시뮬레이션과 실제 데이터 실험은 가정이 위반될 때(예: $\Sigma_{pool}$ 사용 시 $R^2$가 0.98에서 0.83으로 감소, 평균 차이 프로브 사용 시 0.79로 감소) 선형성이 저하됨을 확인했습니다.

의의 및 주장

본 논문은 마할라노비스 코사인 유사도가 선형 프로브를 비교하기 위한 이론적으로 근거가 있고 작업 인지적인(task-aware) 대안임을 주장합니다.

• 메커니즘: 모든 차원을 동일하게 취급하는 ECS와 달리, MCS는 OOD 데이터 공분산에 의해 내적을 재가중치함으로써, 특정 작업에 실제로 중요한 서브스페이스 내에서 프로브를 효과적으로 비교합니다.
• 예측적 유용성: 선형 관계를 통해 연구자들은 ID 데이터로 학습된 프로브의 일반화 성능(OOD AUROC)을 OOD 데이터(또는 피셔 방향의 대리물)에 대한 MCS를 측정함으로써 간단히 예측할 수 있습니다.
• 광범위한 영향: 저자들은 이 접근 방식이 코사인 유사도에 의존하는 스티어링 벡터(steering-vector) 비교, SAE 특징 정렬(feature alignment), 개념 방향 클러스터링과 같은 다른 해석 가능성 방법론에도 도움이 될 수 있다고 제안합니다.
• 한계점: 저자들은 이 방법이 공분산과 참조 방향을 추정하기 위해 라벨이 있는 OOD 데이터가 필요하다는 점(대상 분포에 대해 라벨 프리가 아님)을 명시합니다. 또한, 이론은 피셔 스타일의 프로브(로지스틱 회귀, LDA)에 맞춰져 있으며, 평균 차이 프로브의 경우 성능이 저하됩니다. 평가는 현재 자기회귀 LLM의 잔차 스트림(residual-stream)에 대한 이진 분류에 국한되어 있습니다.

요약하자면, 본 연구는 MCS가 프로브의 일반화를 예측하는 데 있어 성공하는 것이 우연이 아니라, 신호 대 잡음비, 가우시안 투영, 그리고 공분산 인지 유사도 지표 사이의 기하학적 상호작용의 결과임을 확립하였습니다.