Disentangling MLP Neuron Weights in Vocabulary Space

1. 연구 배경 및 문제 정의 (Problem)

• 메커니즘 해석 가능성 (Mechanistic Interpretability) 의 한계: 기존 언어 모델 (LM) 의 가중치에 인코딩된 정보를 해석하는 것은 근본적인 과제입니다. 특히 MLP(다층 퍼셉트론) 뉴런은 종종 다의성 (Polysemanticity) 을 띠고 있어, 하나의 뉴런이 여러 개의 서로 다른 개념을 중첩하여 표현합니다.
• 기존 방법의 단점:
  • 데이터 의존성: 대부분의 기존 방법 (예: 희소 자동 인코더, SAE) 은 모델의 활성화 (activation) 데이터를 대규모로 수집하고 전진 통과 (forward pass) 를 수행해야 합니다. 이는 계산 비용이 높고 특정 데이터 분포에 편향될 수 있습니다.
  • 해석의 불완전성: 뉴런을 단일 단위로 보거나, 활성화 기반의 특징을 추출하는 방식은 뉴런 내부의 세밀한 구조를 완전히 분리해내지 못해, 뉴런의 행동을 정확히 설명하는 데 한계가 있습니다.
• 목표: 데이터 (입력 텍스트) 나 전진 통과 없이, 오직 모델의 가중치 (Weights) 만을 사용하여 MLP 뉴런을 인간이 해석 가능한 단일 의미 (Monosemantic) 구성 요소로 분리해내는 방법론을 개발하는 것입니다.

2. 제안 방법: ROTATE (Methodology)

저자들은 ROTATE (Rotation-Optimized Token Alignment in weighT spacE) 라는 새로운 방법을 제안합니다. 이는 데이터가 필요 없으며 (Data-free), 모델의 가중치 공간에서 직접 뉴런을 분리합니다.

핵심 통찰 (Key Observation)

• 어휘 공간의 첨도 (Kurtosis): 뉴런의 가중치를 모델의 어휘 공간 (Vocabulary Space) 으로 투영했을 때, 높은 첨도 (High Kurtosis) 를 보이는 방향은 단일하고 일관된 개념 (Monosemantic concept) 을 인코딩한다는 통계적 관측을 기반으로 합니다.
  • 높은 첨도는 확률 분포가 뾰족하고 꼬리가 길다는 것을 의미하며, 이는 뉴런이 어휘의 일부 특정 토큰에 강하게 반응하고 나머지는 무시함을 나타냅니다.
  • 반면, 가우시안 분포 (낮은 첨도) 는 뉴런이 여러 토큰에 고르게 반응하여 다의적 (Polysemantic) 임을 시사합니다.

알고리즘 개요

1. 가중치 회전 (Rotation): 주어진 뉴런의 가중치 벡터 $w$에 대해, 회전 행렬 $R$을 학습하여 새로운 벡터 $v = Rw$를 생성합니다.
2. 최적화 목적 함수:
   • 첨도 손실 (Kurtosis Loss): $v$를 어휘 공간에 투영한 로짓 (logits) 분포의 첨도를 최대화합니다. 이를 통해 희소하고 해석 가능한 방향을 찾습니다.
   • 정규화 손실 (Regularization Loss): 회전된 벡터 $v$가 원래 가중치 $w$와 너무 멀어지지 않도록 코사인 거리를 패널티로 부과합니다.
   • 목적 함수: $L = -\lambda \cdot \text{Kurtosis}(z) + \text{CosineDistance}(v, w)$
3. 반복적 마스킹 (Iterative Masking):
   • 한 번의 최적화로 하나의 '어휘 채널 (Vocabulary Channel)'을 찾은 후, 해당 채널이 가장 강하게 반응하는 토큰들을 마스킹 (제외) 합니다.
   • 이를 반복하여 동일한 국소 최적점에 수렴하는 것을 방지하고, 뉴런이 인코딩한 여러 개의 서로 다른 개념을 순차적으로 분리해냅니다.
4. 구현: 하우스홀더 반사 (Householder reflection) 파라미터화를 사용하여 직교 회전 행렬을 효율적으로 학습합니다.

3. 주요 기여 (Key Contributions)

1. 통계적 신호 발견: LM 의 가중치 공간에서 어휘 투영의 높은 첨도가 단일 의미 (Monosemantic) 방향과 상관관계가 있음을 증명했습니다.
2. ROTATE 방법론 제안: 데이터와 전진 통과 없이 가중치만 활용하여 MLP 뉴런을 해석 가능한 '어휘 채널'로 분리하는 알고리즘을 개발했습니다.
3. SAE 기반 방법 대비 우월성: Llama-3.1-8B-Instruct 와 Gemma-2-2B-it 모델에서 ROTATE 가 최신 SAE(Gemma Scope, Llama Scope) 기반 방법보다 신뢰성 (Faithfulness) 과 완전성 (Completeness) 측면에서 2~3 배 더 우수한 성능을 보였습니다.
4. 뉴런 설명의 질적 향상: 분리된 채널들의 설명을 집계하여 생성한 뉴런 설명이, 기존 활성화 기반 (MaxAct) 방법들보다 인간 평가 및 LLM 판정에서 훨씬 높은 승률을 기록했습니다.

4. 실험 결과 (Results)

실험 설정

• 모델: Gemma-2-2B-it, Llama-3.1-8B-Instruct
• 비교 대상: 희소 자동 인코더 (SAE) 기반 특징 (Neuronpedia, Gemma/Llama Scope), 무작위 채널, 활성화 기반 설명 (MaxAct, MaxAct++).

주요 성과 지표

1. 신뢰성 (Faithfulness):

   • 입력 측 (Input-side): ROTATE 가 발견한 채널이 해당 뉴런을 활성화시키는 입력을 정확하게 예측하는지 평가. ROTATE 는 SAE 대비 0.460.71 점으로, SAE(0.330.49 점) 보다 월등히 높았습니다.
   • 출력 측 (Output-side): 채널 제거 (Ablation) 실험을 통해 특정 채널을 제거했을 때 해당 채널이 담당하는 개념이 억제되는지 확인. 자기 채널 제거 시 토큰 빈도가 87~90% 감소했으나, 교차 채널 제거 시에는 영향이 미미하여 채널 간 독립성이 입증되었습니다.

2. 완전성 (Completeness):

   • 다양한 입력에서 뉴런이 활성화될 때, ROTATE 채널 설명이 그 이유를 얼마나 잘 설명하는지 평가 (LLM 판정).
   • ROTATE 는 4960% 의 정확도를 보인 반면, SAE 기반 방법은 3649% 에 그쳤습니다. 이는 ROTATE 가 뉴런의 활성화 스펙트럼을 더 폭넓게 커버함을 의미합니다.

3. 뉴런 설명 생성 (Neuron Descriptions):

   • ROTATE 로 분리된 채널들의 설명을 통합하여 생성한 텍스트 설명이, 기존 최첨단 방법 (MaxAct++) 보다 2~3 배 높은 승률을 기록했습니다.
   • 특히 극단적인 활성화가 아닌 중간 수준의 활성화에서도 ROTATE 가 더 정확한 설명을 제공했습니다.

5. 의의 및 결론 (Significance)

• 확장성과 효율성: ROTATE 는 대규모 활성화 데이터 수집이 불필요하므로, 계산 비용이 적게 들고 모델의 가중치만 있다면 즉시 적용 가능합니다. 이는 대규모 언어 모델의 해석 가능성을 확장하는 중요한 도구입니다.
• 정밀한 메커니즘 이해: 뉴런을 단일 단위가 아닌, 여러 개의 '어휘 채널'로 세분화함으로써, 뉴런이 어떻게 다양한 입력에 반응하고 어떤 개념을 생성하는지에 대한 미시적 (Fine-grained) 인과적 이해를 가능하게 합니다.
• 미래 연구 방향: ROTATE 는 회로 발견 (Circuit Discovery) 과 네트워크 구성 요소 간의 상호작용 연구에 대한 견고한 기반을 제공하며, 가중치 기반 해석 가능성 연구의 새로운 패러다임을 제시합니다.

요약하자면, ROTATE 는 언어 모델의 MLP 뉴런이 가진 다의적 특성을 가중치 공간의 통계적 특성 (첨도) 을 활용하여 데이터 없이 분리해냄으로써, 기존 방법들보다 더 정확하고 포괄적인 뉴런 해석을 가능하게 한 획기적인 방법론입니다.