기술 요약: 소프트맥스의 정보 기하학: 탐침 및 조종 (Probing and Steering)

문제 정의

AI 모델을 해석하고 조작하는 현재의 접근 방식은 고차원 개념이 모델의 벡터 표현 공간 내 특정 방향에 대응한다는 **선형 표현 가설(linear representation hypothesis)**에 주로 의존한다. 이 가설은 개념(예: 감성, 진실성)을 식별하고 조종하는 방법론을 가능하게 했으나, 기존 기술들은 빈번하게 취약하며 직접적인 미세 조정(fine-tuning)보다 효과가 떨어진다. 저자들이 식별한 주요 한계점은 표현 공간이 평탄한 유클리드 기하학을 가진다는 암묵적 가정이다. 본 논문은 소프트맥스 분포(Transformer, LLM, CLIP과 같은 대조 학습 모델 포함)를 사용하는 모델의 경우, 표현 공간의 자연스러운 기하학이 유클리드 기하학이 아닌 **정보 기하학(information geometry)**이라고 주장한다. 결과적으로, 표준 유클리드 연산(예: 표현에 프로브 벡터를 직접 더하는 것)은 의미적 구조를 보존하지 못하거나, 의도치 않은 "타겟 외 드리프트(off-target drift)"—즉, 확률 질량이 관련 없는 개념으로 누출되는 현상—를 유발할 수 있다.

방법론

본 논문은 표현 벡터 $\lambda$의 기하학이 유도된 확률 분포 $P_\lambda$의 근접성을 반영해야 한다는 아이디어를 구체화한다.

1. 정보 기하학 프레임워크

저자들은 켈백-라이블러(KL) 발산에 의해 유도된 기하학이 브레그만(Bregman, 이중 평탄) 기하학임을 확립한다. 이 구조는 다음 사이의 쌍대성(duality)을 도입한다:

• 프라이멀 공간 (Primal Space, $\Lambda$): 표현 벡터 $\lambda$의 공간.
• 듀얼 공간 (Dual Space, $\Phi$): 기대값(평균 임베딩)의 공간 $\phi(\lambda) = \nabla A(\lambda) = \mathbb{E}[\gamma | \lambda]$, 여기서 $A(\lambda)$는 로그 정규화 인자(log-normalizer)이다.

논문은 두 가지 자연스러운 보간 경로를 구분한다:

• 프라이멀 보간 (Primal Interpolation, e-geodesic): 프라이멀 공간에서의 직선. 이는 역(reverse) KL 발산의 가중 합을 최소화하며, 공유된 고확률 영역을 강조하는 AND 연산자처럼 동작한다.
• 듀얼 보간 (Dual Interpolation, m-geodesic): 듀얼 공간에서의 직선. 이는 순방향(forward) KL 발산의 가중 합을 최소화하며, 분포의 선형 혼합을 생성하여 양 끝단의 의미적 특징을 동시에 보존하는 OR 연산자처럼 동작한다.

2. 듀얼 스티어링 (Dual Steering)

핵리의 방법론적 기여는 타겟 개념을 나타내면서도 타겟 외 개념에 대한 변화를 최소화하도록 표현을 수정하는 기술인 듀얼 스티어링이다.

• 유클리드 스티어링의 문제점: 표준 스티어링은 듀얼 공간의 원소인 프로브 벡터 $\beta_W$를 프라이멀 표현 $\lambda$에 직접 더한다. 이는 비유클리드 기하학에서 "타입 불일치(type mismatch)"를 일으키며, 종종 타겟 외 드리프트를 유발한다.
• 듀얼 스티어링 솔루션: $\lambda$를 직접 업데이트하는 대신, 듀얼 스티어링은 듀얼 좌표 $\phi(\lambda)$에 프로브 벡터를 더함으로써 업데이트한다: $\phi(\lambda_t) = \phi(\lambda_0) + t\beta_W$.
• 이론적 보장: 개념 분해 가능 분포(concept-factorizable distributions)(출력 분포가 타겟 개념 분포와 타겟 외 분포로 분해될 수 있다는 가정) 하에서, 저자들은 듀얼 스티어링이 타겟 개념을 수정하면서 타겟 외 분포의 KL 발산을 최소화하는 최적해를 찾는다는 것을 증명한다. 이는 관련 없는 토큰들의 상대적 확률을 효과적으로 보존한다.

3. 실제 구현

듀얼 공간은 유계된 볼록 집합(bounded convex set)이므로, 듀얼 공간에서의 선형 업데이트는 가용 영역을 벗어날 수 있다. 이를 해결하기 위해 저자들은 **정규화된 뉴턴 업데이트(Regularized Newton Update)**를 제안한다:

1. 로그 정규자의 헤시안(unembedding 벡터의 공분산 행렬에 해당)을 사용하여 듀얼 좌표의 변화량을 근사한다.
2. 선형 시스템 $\text{Cov}[\gamma | \lambda] \Delta \lambda = \epsilon \beta_W$를 풀어 그에 상응하는 프라이멀 공간에서의 스텝을 찾는다.
3. 공분산 행렬에 **티코노프 정규화(Tikhonov regularization, $\alpha I$)**를 적용하여 랭크 결핍(소프트맥스가 소수의 토큰에 고도로 집중될 때 발생)을 처리하며, 이를 통해 업데이트가 잘 정의되도록 하고 필요 시 분포를 고엔트로피 영역으로 밀어낸다.
4. 거대한 공분산 행렬을 명시적으로 역행렬 계산하지 않고도 효율적으로 선형 시스템을 풀기 위해 켤레 기울기(Conjugate Gradient) 방법을 사용한다.

주요 결과

저자들은 Gemma-3-4B(LLM)와 MetaClip-2(시각-언어 모델)를 대상으로 다양한 이진 개념(예: 동사 $\Rightarrow$ 3인칭, 개 $\Rightarrow$ 고양이, 영어 $\Rightarrow$ 프랑스어)에 대해 듀얼 스티어링을 평가한다.

• 타겟 외 개념의 보존: 실험 결과, 듀럴 스티어링은 타겟 외 분포를 보존하는 데 있어 유클리드 스티어링보다 성능이 월등히 뛰어남을 보여준다. 유클리드 스티어링은 종종 "누출(leakage)"을 일으켜 (예: "아버지" $\Rightarrow$ "어머니"로 스티어링할 때 "친구"와 같은 관련 없는 토큰에 상당한 확률을 할당함), 듀얼 스티어링은 반사실적 쌍(counterfactual pairs)과 중립적 토큰의 확률 질량을 유지한다.
• 강건성 지표: 듀얼 스티어링은 타겟 개념의 확률이 1에 가까워지더라도 유클리드 방식에 비해 타겟 외 분포에 대해 더 낮은 KL 발산과 더 작은 랭크 차이를 달 achieves 한다.
• 보간 동작: 시각화 결과, 듀럴 보간(m-geodesic)은 분포의 선형 혼합으로 작용하는 반면, 프라이멀 보간(e-geodesic)은 양 끝단 특유의 출력을 억제하고 공유된 특징을 선호함을 확인하였다.

의의 및 주장

본 논문은 고유한 정보 기하학의 비유클리드 구조가 기하학적 조작을 스티어링과 같은 의미적 작업과 연결하는 데 결정적이라는 점을 주장한다.

• 이론적 통찰: 선형 표현 가설이 견고해지기 위해서는 평탄한 유클리드 기하학이 아닌 정보 기하학(특히 브레그만 쌍대성)의 맥락 내에서 이해되어야 함을 입증한다.
• 실용적 유용성: 듀얼 스티어링은 의도치 않은 부작용을 최소화하면서 표현을 강력하게 조작할 수 있는 이론적으로 근거 있는 방법을 제공하며, 해석 가능성 및 제어 방법의 신뢰성을 향상시키기 위한 템플릿을 제시한다.
• 범위: 결과는 소프트맥스 기반 모델에서 도출되었으나, 저자들은 고유한 기하학을 활용하는 이 핵심 아이디어가 널리 적용 가능하다고 제안한다. 저자들은 본 연구가 모델의 중간 레이어에서 기하학이 제어 방법의 강건성에 어떻게 영향을 미치는지에 대한 향후 연구를 위한 템플릿 역할을 한다고 명시했다.

저자들은 (선형 프로브가 전체 초평면에 걸쳐 개념을 완벽하게 포착한다는) "프로빙 가정"에 대해, 실제로는 이 가정이 엄격하게 성립하지 않을 수 있음을 인정하면서도, 듀얼 스티어링이 여전히 유클리드 방식보다 경험적으로 우수하다는 점을 밝히며 겸허한 태도를 유지한다.