기술 요약: VASAE – 어휘 정렬 앵커링을 통한 SAE 딕셔너리 방향의 명명

1. 문제 정의

희소 오토인코더(Sparse Autoencoders, SAEs)는 트랜스포머 모델의 고차원 잔차 스트림(residual streams)을 학습된 딕셔너리 상의 희소 코드(sparse codes)로 분해하는 데 점점 더 많이 사용되고 있으며, 이는 기계론적 해석 가능성(mechanistic interpretability)을 돕는다. 그러나 표준 SAE의 중대한 한계점은 학습된 특징(feature, 딕셔너리 방향)이 모델의 토큰 어휘(vocabulary)와 본질적으로 연결되어 있지 않다는 것이다.

표준적인 관행에서 특징 명명(feature naming)은 사후적(post hoc) 과정이다. 연구자들은 특정 특징이 활성화되는 문맥을 조사하거나, 별도의 자동화된 해석 절차를 실행하여 의미론적 레이블을 할당해야 한다. 이러한 단절은 학습된 딕셔너리의 기하학적 구조와 토큰 임베딩의 기하학적 구조가 일반적으로 서로 무관하기 때문에 발생한다. 즉, 하나의 특징은 활성화를 효과적으로 재구성할 수 있으면서도 어떤 토큰 임베딩 벡터와도 멀리 떨어져 있을 수 있으며, 이로 인해 본질적인 어휘 수준의 이름을 갖지 못하게 된다.

본 논문은 다음과 같은 질문을 던진다: SAE 특징들이 재구성 성능을 희생하지 않으면서도, 높은 품질의 재구성을 위한 유연성을 유지하는 동시에 토큰 임베딩에 충분히 가깝게 기하학적으로 고정(anchored)되어 학습 과정 중에 고유한 토큰 이름을 획득할 수 있는가?

2. 방법론: 어휘 정렬 희소 오토인코더 (VASAE)

저자들은 표준 SAE 학습 목적 함수에 어휘 정렬 앵커링(vocabulary-aligned anchoring) 메커니즘을 도입한 수정된 방식인 VASAE를 제안한다.

핵심 아키텍처

• 입력: VASAE는 트랜스포머 레이어 $\ell$의 사후 잔차 스트림 벡터($h^{(\ell)}$)를 대상으로 작동한다.
• 인코더(Encoder): 표준적인 top-k 희소 인코더(선형 맵 + ReLU + Top-K 선택)를 사용하여 희소 코드 $z$를 생성한다.
• 디코더(Decoder): 잔차 스트림을 재구성하기 위해 학습 가능한 딕셔너리 행렬 $W_D$ (여기서 각 열은 특징 벡터 $d_i$임)를 사용한다.
• 앵커(Anchors): 이 방법은 모델의 고정된 입력 토큰 임베딩($W_E$)을 참조 방향으로 활용한다. 이 임베딩들은 잔차 스트림과 동일한 차원을 공유하거나 (그 차원으로) 투영된다.

앵커링 목적 함수 (The Anchoring Objective)

핵심 혁신은 표준 재구성 손실(reconstruction loss)에 앵커 손실 항을 추가하는 것이다.

1. 최근접 토큰 정렬: 각 학습된 특징 $d_i$에 대해, 모든 고정된 토큰 임베딩 $w_v$와의 코사인 유사도를 계산한다.
   $$s_i = \max_{v \in V} \cos(d_i, w_v)$$
   이 최대값을 달성하는 토큰 $v^*_i$가 특징 $i$의 **본질적 토큰 이름(intrinsic token name)**을 정의한다.
2. 앵커 손실: 학습 목적 함수는 특징이 가장 가까운 토큰 임베딩 근처에 머물도록 장려한다:
   $$L_{anchor} = -\frac{1}{d_{sparse}} \sum_{i=1}^{d_{sparse}} s_i$$
   이 손실을 최소화하는 것은 특징과 가장 가까운 어휘 토큰 사이의 코사인 유사도를 최대화한다.
3. 전체 손실: 결합된 목적 함수는 다음과 같다:
   $$L_{vasae} = L_{recon} + \lambda_{anchor} L_{anchor}$$
   여기서 $\lambda_{anchor}$는 정렬 제약의 강도를 조절하는 하이퍼파라미터이다.

하드 타이잉(Hard-Tying)과의 차이점

저자들은 디코더 가중치가 토큰 임베딩과 정확히 일치하도록 고정하는 "하드-타이드(hard-tied)" 디코더 베이스라인과 VASAE를 명시적으로 구분한다 ($W_D = W_E$). 하드-타이잉은 특징이 정확히 토کن 임베딩이 되도록 강제하며, 본 논문은 이것이 고품질의 재구성을 수행하기에는 불충분하다고 가설을 세운다. VASAE는 딕셔너리를 학습 가능하게 유지하면서도, 앵커 손실을 통해 어휘 기하학 쪽으로 끌어당긴다.

3. 실험 설정

이 방법은 두 가지 디코더 전용 트랜스포머 모델에 대해 평가되었다:

• GPT-2-small: 12개 레이어, $d_{model}=768$, 어휘 크기 $|V|=50,257$.
• Llama-3.1-8B: 32개 레이어, $d_{model}=4096$, 어휘 크기 $|V|=128,256$.

베이스라인:

1. 표준 SAE (Standard SAE): 재구성 손실만으로 학습됨.
2. 하드-타이드 디코더 (Hard-Tied Decoder): 디코더 가중치를 토큰 임베딩으로 고정하고 인코더만 학습함.

평가 지표:

• 재구성 (Reconstruction): 평균 제곱 오차(MSE), 설명된 분산(Variance Explained, VE), 그리고 잔차 스트림 치환 후의 교차 엔트로피(Cross-Entropy, CE) 회복률.
• 정렬 (Alignment): 기하학적 특징 정렬 비율 (정렬 점수 $s_i \ge 0.8$을 만족하는 특징의 백분율) 및 어휘 커버리지.
• 사례 연구 (Case Studies): 특정 문맥에서 활성화되는 특징에 할당된 본질적 토큰 이름에 대한 질적 검사.

4. 주요 결과

재구성 품질

VASAE는 표준 SAE와 대등한 수준의 재구성 지표를 달려하며, 하드-타이드 디코더 베이스라인을 크게 능가한다.

• GPT-2-small: VASAE와 표준 SAE 모두 설명된 분산(VE) 0.965와 CE 회복률 0.975를 달성한다.
• Llama-3.1-8B: VASAE와 표준 SAE 모두 VE 0.931와 CE 회복률 0.906을 달성한다.
• 하드-타이드 베이스라인: 성능이 저조하며 (예: GPT-2-small에서 음수의 VE 발생), 특징을 정확히 토큰 임베딩으로 강제하는 것이 재구성 능력을 파괴한다는 것을 확인시켜 준다.

기하학적 정렬

• GPT-2-small: $\lambda_{anchor} = 10^{-4}$일 때, 레이어 0~10의 특징 중 **89–94%**가 정렬 컷오프($s_i \ge 0.8$)를 초과한다. 이러한 정렬은 얕은 층과 중간 층에서 강력하고 일관되게 나타난다.
• Llama-3.1-8B: 정렬은 깊이에 따라 달라진다. 더 높은 계수($\lambda_{anchor} = 5 \times 10^{-3}$)에서 얕은 레이어(L0)는 **92.8%**의 정렬을 달성한다. 그러나 더 깊은 레이어(예: L31)에서는 정렬이 크게 감소하는데, 이는 어휘 앵커링이 임베딩 공간에 더 가까운 잔차 상태와 더 잘 호환됨을 시사한다.

사례 연구

본질적 토큰 이름에 대한 질적 분석은 의미론적 일관성을 보여준다:

• GPT-2-small에서 "Baker Street"에 관한 문장에서 활성화되는 특징들은 located, Street, corner, around와 같은 본질적 이름을 할당받는다.
• 추가적인 예시들은 형용사/부사(예: perfect, feasible), 수상 관련 용어(예: winner, praise), 그리고 자기소개 관련 용어(예: fan, name)와 정렬되는 특징들을 보여준다.
• Llama-3.1-8B의 경우, 얕은 레이어는 이와 유사한 위치 관련 정렬을 보이는 반면, 깊은 레이어는 정량적 정렬 지표를 반영하듯 일관성이 떨어진다.

5. 기여 및 의의

본 논문은 세 가지 주요 기여를 한다:

1. VASAE의 도입: 학습 과정 중에 사후적으로가 아니라, 학습된 SAE 디코더 방향에 가장 가까운 토큰 이름을 부여하는 부드러운 어휘 앵커링 목적 함수를 도입하였다.
2. 앵커링의 평가: 어휘 정렬 앵커링이 표준 SAE와 대등한 재구성 품질을 유지하면서도, 디코더를 토큰 임베딩에 하드-타이잉하는 것이 이 작업에 불충분함을 입증하였다.
3. 기하학적 분석: 임베딩 공간의 특징 정렬과 어휘 커버리지를 분석하여, 본질적 토킨 이름이 기하학적으로 유도된 식별자 역할을 한다는 증거를 제공하였다.

의의:
저자들은 VASAE가 학습된 특징을 모델의 고유한 토큰 이름과 연결함으로써 기존의 사후적 해석 방법들을 보완한다고 주장한다. 이는 특징에 대한 "어휘 인덱스 핸들(vocabulary-indexed handle)"을 제공하여 검사의 출발점을 마련해 준다.

결정적으로, 본 논문은 이러한 이름들의 함의에 대해 신중한 입장을 유지한다:

• 설명이 아님: 본질적 토큰 이름은 특징에 대한 완전한 설명이 아니다. 그것은 왜 해당 특징이 활성화되는지 또는 모델의 계산에서 어떤 인과적 역할을 하는지를 설명하지 않는다.
• 인덱싱 도구: 이름은 예비적인 인덱싱 도구 역할을 한다. 진정한 이해를 위해서는 명명된 특징들이 어떻게 회로(circuits)에 참여하는지, 다른 구성 요소들과 어떻게 연결되는지, 그리고 그들에 대한 개입이 행동에 어떤 영향을 미치는지에 대한 추가적인 분석이 필요하다.
• 한계점: 정렬은 GPT-2-small과 Llama-3.1-8B의 얕은/중간 레이어에서 가장 견고하다. 현재 이 방법은 입력 임베딩에만 앵커링되어 있으며, 입력/출력 임베딩이 분리된(untied) 모델의 경우 앵커 행렬의 선택은 여전히 해결해야 할 과제로 남아 있다.

요약하자면, VASAE는 학습된 특징의 기하학적 공간과 모델의 어휘 공간 사이의 간극을 성공적으로 메워, 재구성의 유용성을 해치지 않으면서도 특징들이 본질적인 토큰 이름을 가질 수 있게 한다.