Constructing Interpretable Features from Compositional Neuron Groups

기술 요약: 구성적 뉴런 그룹으로부터 해석 가능한 특징 구성

문제 제기
기계적 해석 가능성 분야에서 핵심적인 과제는 대규모 언어 모델 (LLM) 내에서 모델 출력을 인과적으로 설명하는 적절한 분석 단위를 식별하는 것입니다. 초기 연구는 개별 뉴런에 집중했으나, 뉴런이 종종 여러 개념을 인코딩한다는 증거 (다의성) 가 발견되면서 활성화 공간 내의 방향을 분석하는 방향으로 전환이 촉진되었습니다. 현재 비지도 접근법은 잔여 스트림 (residual stream) 활성화에 대해 훈련된 희소 오토인코더 (SAE) 에 의존하여 처음부터 특징 방향을 학습합니다. 그러나 SAE 는 중요한 한계에 직면해 있습니다: 인과적 평가 (조작) 에서 어려움을 겪는 경우가 많으며, 학습된 특징이 모델의 내부 계산과 명시적으로 연결되지 않아 본질적인 해석 가능성이 부족하고, 특징이 모델의 표현 공간에 구속되지 않습니다. 따라서 모델 메커니즘과 직접적인 인과적 연결을 유지하면서 비지도 방식으로 해석 가능한 특징을 발견할 수 있는 방법이 필요합니다.

방법론
저자들은 반-비음수 행렬 분해 (Semi-Nonnegative Matrix Factorization, SNMF) 를 사용하여 MLP(다층 퍼셉트론) 활성화를 직접 분해하는 방법을 제안합니다. 새로운 방향을 학습하는 SAE 와 달리, 이 방법은 특징을 MLP 내의 기존에 함께 활성화된 뉴런들의 희소 선형 결합으로 표현합니다.

1. 수식화: MLP 활성화 행렬 $A \in \mathbb{R}^{d_a \times n}$ (여기서 $d_a$는 내부 차원, $n$은 토큰 수) 가 주어졌을 때, 이 방법은 $A$를 두 개의 행렬로 분해합니다:

   • $Z \in \mathbb{R}^{d_a \times k}$: $k$개의 MLP 특징 행렬로, 각 열은 뉴런들의 희소 선형 결합을 정의합니다.
   • $Y \in \mathbb{R}^{k \times n}_{\geq 0}$: 입력을 특징으로 매핑하는 비음수 계수 행렬입니다.
     분해는 $A \approx ZY$를 따릅니다.

2. 해석 가능성 메커니즘:

   • 희소성: $Z$의 열은 하드 위너-테이크-올 (winner-take-all) 연산자로 제약되어 각 특징이 뉴런의 작은 부분집합만 활성화하도록 보장합니다.
   • 인과적 근거: 비음수 행렬 $Y$는 어떤 토큰이 특징 생성에 기여했는지를 명시적으로 인코딩하므로, 고활성화 맥락에 기반하여 텍스트 개념 레이블을 직접 할당할 수 있습니다.
   • 잔여 스트림 매핑: 활성화 공간의 각 특징 벡터 $z_i$는 MLP 출력 가중치 ($W_V$) 를 통해 잔여 스트림으로 매핑되어 잔여 스트림 특징 $f_i = W_V z_i$를 정의합니다. 이는 특징이 모델의 실제 계산 경로에 근거하도록 보장합니다.

3. 최적화: 이 방법은 $A$와 $ZY$ 사이의 푸비니우스 노름을 최소화하기 위해 곱셈 업데이트 (Ding et al., 2010) 를 사용하며, $Z$에 희소성이 부과됩니다.

주요 기여

• MLP 분해를 위한 SNMF: 본 논문은 임의의 방향을 처음부터 학습하는 대신, 함께 활성화된 뉴런 그룹으로 정의된 해석 가능한 특징으로 MLP 활성화를 분해하는 비지도 방법을 도입합니다.
• 인과적 조작 우위성: 저자들은 SNMF 에서 파생된 특징이 개념 탐지에서는 동등한 성능을 보이면서, 인과적 조작 작업에서 최신 SAE(GemmaScope, Llamascope) 와 강력한 지도식 기준선 (Difference-in-Means) 을 능가함을 입증했습니다.
• 계층적 구성의 발견: 재귀적 SNMF 를 통해 MLP 특징이 구성적임을 보였습니다. 정교한 개념 (예: 특정 요일) 은 공유된 "핵심" 뉴런과 배타적 뉴런을 결합하여 형성됩니다. 이는 더 넓은 개념 (예: "요일") 이 더 구체적인 하위 개념으로 구성되는 계층적 구조를 드러내며, SAE 에서 흔히 관찰되는 "특징 분할" 현상을 효과적으로 반전시킵니다.
• 뉴런 재사용: 분석은 의미적으로 관련된 특징이 핵심 뉴런 집합을 공유하며, 배타적 뉴런이 특정 하위 개념의 촉진 또는 억제를 제어함을 확인했습니다.

결과
실험은 Llama 3.1-8B, Gemma 2-2B, GPT-2 Small에서 수행되었습니다.

• 개념 탐지: SNMF 특징은 대부분의 특징 (>75%) 에서 긍정적인 탐지 점수를 달성하여, 개념 관련 입력에서는 활성화되고 중립적 입력에서는 비활성화됨을 나타냈습니다. 성능은 MLP 출력에 훈련된 SAE 와 비교하거나 능가했습니다.
• 개념 조작: 인과적 조작 평가 (유창성을 유지하면서 모델 생성을 특정 개념으로 조작하는 능력 측정) 에서 SNMF 는 모든 층에 걸쳐 SAE 를 일관되게 능가했습니다. 또한 지도식 Difference-in-Means 기준선을 자주 능가하거나 동등한 성능을 보였습니다.
• 계층적 분석: SNMF 의 재귀적 적용은 "월요일", "화요일"과 같은 정교한 특징을 "요일", "주말"과 같은 상위 개념으로 성공적으로 병합했습니다. 인과적 개입은 "핵심" 뉴런을 증폭하면 관련 모든 토큰이 촉진되는 반면, "배타적" 뉴런을 증폭하면 특정 토큰은 촉진되고 다른 토큰은 억제됨을 확인했습니다.

의의 및 주장
본 논문은 SNMF 가 LLM 에서 해석 가능한 특징을 식별하는 간단하고 효과적인 도구를 제공한다고 주장합니다. MLP 활성화를 직접 분해함으로써, 이 방법은 특징 형성의 책임이 있는 특정 모델 구성 요소와 특징을 연결합니다. 결과는 MLP 가 해석 가능한 뉴런 집합으로 구성된 가산적 업데이트를 통해 작동함을 시사합니다. 이 작업은 SNMF 를 SAE 의 실현 가능한 대안으로 위치시키며, 감독 없이도 더 나은 인과적 성능과 본질적 해석 가능성을 제공합니다. 또한, 계층적 뉴런 구성의 발견은 LLM 이 재사용 가능한 신경 구성 요소에서 복잡한 개념을 어떻게 구성하는지 이해하는 새로운 렌즈를 제공하며, 뉴런의 기능적 역할을 단일 개념 인코딩을 넘어 확장합니다.

저자들은 현재 특징 수 $k < 500$에 초점을 맞추고 평가에 LLM 을 판사로 활용한다는 한계를 인정하지만, 인간 판단과의 높은 일치도를 보고합니다. 향후 작업은 더 큰 $k$값으로 확장하고 최적화를 위해 투영 경사 하강법을 활용하는 것을 탐구할 수 있다고 제안합니다.