기술 요약: 희소 오토인코더 특징 매칭 및 회로 압축을 위한 의미적 최적 수송

1. 문제 정의

희소 오토인코더 (SAE) 는 잠재 공간을 해석 가능한 특징으로 분해함으로써 대형 언어 모델 (LLM) 해석을 위한 주요 도구로 부상했습니다. 그러나 SAE 분석을 확장하는 데 있어 두 가지 중요한 과제가 여전히 존재합니다:

1. 층 간 특징 매칭: LLM 의 서로 다른 층에 걸쳐 의미적으로 유사한 특징을 식별하는 것. 기존 방법들은 일반적으로 디코더 벡터를 비교하는 데 의존합니다. 그러나 이 접근법은 비연속적인 층 간 특징 매칭 시 실패합니다. 디코더 벡터는 서로 다른 활성화 분포를 재구성하도록 훈련되기 때문에 직접적인 벡터 비교가 부정확하기 때문입니다.
2. 특징 회로 압축: 층 간 특징을 연결하는 인과 그래프인 특징 회로는 종종 수백 개의 노드를 포함하여 수동 검사가 비실용적입니다. 현재 압축 방법들은 의미적으로 유사한 노드를 "초노드 (supernodes)"로 수동 그룹화하는 데 의존하며, 이 과정은 확장성이 없습니다.

저자들이 규명한 핵심 문제는 서로 다른 활성화 매니폴드에 존재하는 특징 간의 의미적 거리를 추정하는 것입니다. 기존 접근법들은 특징을 단일 벡터 (예: 디코더 가중치) 로 취급하여, 특징이 어디서 그리고 얼마나 강하게 활성화되는지에 대한 풍부한 맥락을 무시함으로써 정보 손실을 초래하고 다중 매니폴드 비교를 처리하지 못하게 합니다.

2. 방법론

저자들은 특징 비교를 최적 수송 (Optimal Transport, OT) 문제로 재정의하는 분산 프레임워크를 제안합니다.

2.1 분산으로서의 특징 표현

단일 디코더 벡터로 특징을 표현하는 대신, 이 방법은 층 $\ell$의 각 특징 $i$를 해당 특징이 활성화되는 숨겨진 상태 (hidden states) 에 대한 활성화 가중치 경험 분포로 표현합니다.

• 토큰 위치의 말뭉치 $T$가 주어지면, 특징 $i$가 가장 강하게 활성화되는 상위-K 토큰 위치를 선택합니다.
• 각 선택된 토큰 $t$는 그 활성화 크기에 비례하는 가중치 $w_{i,t}$를 할당받습니다.
• 특징은 분포 $\hat{\mu}_{i,T} = \sum w_{i,t} \delta_{x_t}$로 표현되며, 여기서 $x_t$는 토큰 $t$에서의 숨겨진 상태입니다.

2.2 공유 참조 공간

서로 다른 층 (서로 다른 매니폴드에 존재) 의 특징을 비교하기 위해, 이 방법은 층별 분포를 공유 참조 공간 $\mathcal{R}$로 투영합니다.

• 층 간 매칭의 경우, 참조 공간은 일반적으로 대상 층의 숨겨진 상태 공간입니다. 소스 층의 활성화는 대상 층의 해당 토큰 위치로 투영됩니다.
• 회로 압축의 경우, 참조 공간은 모든 층에 걸친 숨겨진 상태의 연결입니다.

2.3 Wasserstein 거리를 통한 의미적 거리

공유 공간으로 투영된 후, 두 특징 간의 의미적 거리는 Wasserstein 거리 (Earth Mover's Distance) 를 사용하여 계산됩니다.

• 거리 $W_c(\hat{\mu}, \hat{\nu})$는 한 특징의 활성화 분포의 확률 질량을 다른 것으로 이동시키는 최소 비용을 측정합니다.
• 기본 비용 (ground cost) $c(z, z')$은 공유 공간에서 두 맥락 표현을 매칭하는 비용을 측정합니다.
• 이 접근법은 점별 벡터 유사성뿐만 아니라 활성화 패턴의 기하학적 구조를 포착합니다.

2.3 하류 작업

• 특징 매칭: 대상 특징에 대해, 이 방법은 최소 Wasserstein 거리 (가장 가까운 이웃) 를 갖는 소스 특징을 찾습니다.
• 회로 압축: 이 방법은 회로 노드를 분산 공간의 점으로 취급하고, 쌍별 Wasserstein 거리가 작은 노드를 초노드로 그룹화하기 위해 클러스터링 (예: 계층적 클러스터링) 을 적용하여 클러스터 내 의미적 거리를 최소화합니다.

3. 주요 기여 및 이론적 보장

이 논문은 제안된 프레임워크에 대한 세 가지 주요 이론적 보장을 제공합니다:

1. 활성화 재스케일링 불변성: 투영된 특징 분포는 활성화 값의 전역 스케일링에 대해 불변입니다. 표현이 활성화의 상대적 순위와 비례 분포 (상위-K 선택 및 정규화된 가중치) 에 의존하기 때문에, 활성화 크기의 임의의 스칼라 이동은 의미적 거리를 변경하지 않습니다.
2. 교란 하의 안정성: 특징 분포 간의 Wasserstein 거리는 분포의 부드러운 교란 하에서 안정적입니다. 거리 변화는 기본 비용의 Lipschitz 상수와 교란된 분포 및 원래 분포 간의 Wasserstein 거리에 의해 제한됩니다. 이는 토큰 선택 또는 SAE 재훈련 시 노이즈에 대한 견고성을 보장합니다.
3. 마진 기반 복구: 저자들은 경험적 최근접 이웃 매칭 및 의미적 그룹화가 진정한 (오라클) 매칭을 완벽하게 복구하는 조건을 확립합니다. "마진" (최고 매칭과 차고 매칭 간의 간격) 이 추정 오차의 두 배를 초과하면, 경험적 결정이 보장되도록 정확합니다.

4. 실험 결과

이 방법은 GPT-2와 Gemma-2-2B에서 훈련된 SAE 에 대해 평가되었으며, SAE Match (디코더 벡터 L2), Feature Flow (디코더 벡터 코사인), Attribution Patching, LLM 기반 선택 등을 포함한 기준선과 비교되었습니다.

• 특징 매칭: 제안된 방법은 연속 층, 중간에서 마지막 층, 그리고 첫 번째에서 마지막 층 설정 전반에 걸쳐 LLM 평가(의미적 유사성 평가) 및 매칭 정확도에서 모든 기준선을 능가했습니다. 특히, 이 방법은 유사한 활성화 토큰뿐만 아니라 유사한 기능적 역할을 가진 특징을 식별함을 보여주기 위해 교차 엔트로피 손실 및 설명된 분산과 같은 하류 지표에서 경쟁력 있는 성능을 유지했습니다. 다른 방법들이 실패한 미묘한 "숫자 덧셈" 특징 매칭에서 특히 성공적이었습니다.
• 회로 압축: 이 방법은 수동 그룹화 및 기타 기준선에 비해 직관적인 초노드로 특징 회로를 압축하는 데 더 높은 정확도를 달성했습니다. 서로 다른 회로에서 일관되게 동일한 기능적 역할을 수행하는 모듈식 특징 그룹 (예: "유도 헤드" 및 "토큰 감지") 을 성공적으로 식별했습니다.
• 절차 연구: 실험은 특징을 분산으로 표현하는 것이 동일한 매니폴드로 투영되더라도 단일 중심 벡터로 축소하는 것보다 더 많은 의미 정보를 보존함을 확인했습니다.

5. 중요성 및 주장

이 논문은 층 간 특징 매칭과 회로 압축을 매니폴드 간 의미적 거리 추 estimation이라는 동일한 근본 문제의 사례로 해결하는 통합된 의미적 지표를 제공한다고 주장합니다.

• 확장성: 특징을 초노드로 그룹화하는 과정을 자동화함으로써, 이 방법은 인간 라벨링에 대한 의존성을 제거하여 대규모 특징 회로 분석을 실현 가능하게 합니다.
• 견고성: 분산 접근법은 특징 활성화의 "맥락"을 포착하여, 특징이 층을 거치며 크게 진화할 때 (비유사한 분포) 도 정확한 매칭을 가능하게 합니다. 이는 벡터 기반 방법이 실패하는 시나리오입니다.
• 해석 가능성: 복잡한 회로를 해석 가능한 초노드로 자동으로 압축할 수 있는 능력은 "Dallas"(지리), "Rhyming"(운율), "Opposite"(반대) 회로에 대한 사례 연구를 통해 입증된 바와 같이, 모델 메커니즘에 대한 이해를 향상시킵니다.

저자들은 이 작업을 벡터 중심 특징 비교의 한계를 넘어 확장 가능하고 자동화된 기계적 해석 가능성으로 나아가는 기초적인 단계로 위치시킵니다.