SAE Lens와 TensorLens: Feature Interpretability의 시대

지난 두 글에서 우리는:

• Logit/Tuned Lens: 모델의 중간 예측을 읽었고
• Activation Patching: 어떤 activation이 답의 원인인지 추적했습니다

하지만 여기서 근본적인 문제에 부딪힙니다:

우리가 조작하고 관찰하는 activation은 도대체 무엇을 "의미"하는가?

activation의 각 차원은 개별 뉴런에 대응합니다. 그런데 이 뉴런들은 polysemantic(다의적)합니다. 하나의 뉴런이 학술 인용, 영어 대화, HTTP 요청, 한국어 텍스트에 동시에 반응합니다. 뉴런 단위로는 깨끗한 해석이 불가능합니다.

이 글에서는 이 문제를 해결하는 두 가지 최신 접근법을 다룹니다:

1. Sparse Autoencoder (SAE): dense activation을 sparse한 monosemantic feature로 분해
2. TensorLens: Transformer 전체 계산을 하나의 high-order tensor로 통합 분석

1. Superposition: 왜 뉴런은 해석 불가능한가

문제: Polysemantic Neurons

GPT-2의 MLP 뉴런 하나를 관찰하면, 이런 일이 벌어집니다:

• 뉴런 #4721: 학술 인용 형식, 영어 대화문, base64 문자열에 모두 활성화
• 뉴런 #1389: 수학 기호, Python 코드, LaTeX에 동시 활성화

왜 이런 일이 생길까요? 모델이 "잘못 학습"된 것일까요?

아닙니다. 이것은 superposition이라는 의도적인(혹은 불가피한) 현상입니다.

Superposition: 차원보다 많은 개념을 표현하기

Anthropic의 "Toy Models of Superposition" (Elhage et al., 2022)이 이 현상의 수학적 기반을 제공했습니다.

핵심 통찰은 다음과 같습니다:

고차원 공간에서는 거의 직교하는 방향이 차원 수보다 훨씬 많이 존재합니다.

예를 들어 512차원 공간에서:

• 512개의 완전 직교 벡터를 만들 수 있습니다 (기저 벡터)
• 하지만 서로 거의 직교하는(cos 유사도 < 0.1) 벡터는 수천 개 만들 수 있습니다

모델은 이를 활용합니다. 512차원의 hidden state에 512개가 아니라 수천 개의 feature를 겹쳐서 저장합니다. 이것이 superposition입니다.

$$\mathbf{h} = \sum_i f_i \cdot \mathbf{d}_i$$

($\mathbf{d}_i$는 feature direction, $f_i$는 feature activation)

각 feature $\mathbf{d}_i$는 거의 직교하지만 완전히 직교하지는 않으므로, 서로 약간의 간섭(interference)이 발생합니다. 하지만 feature가 충분히 sparse하면(자주 활성화되지 않으면), 이 간섭 비용은 낮습니다.

핵심 조건: feature가 sparse할수록 superposition이 가능합니다.

superposition 이득 > 간섭 비용:

$$I_i > \sum_j I_j \cdot p_j \cdot (\mathbf{d}_i \cdot \mathbf{d}_j)^2$$

여기서 $I_i$는 feature의 중요도, $p_j$는 feature $j$가 활성화될 확률입니다.

결과: 뉴런 ≠ Feature

superposition 때문에:

• 하나의 뉴런 = 여러 feature의 혼합 → polysemantic
• 하나의 feature = 여러 뉴런에 분산 → distributed representation

뉴런을 해석 단위로 쓰는 것은 구조적으로 불가능합니다. feature를 추출해야 합니다.

2. Sparse Autoencoder: Feature 추출기

핵심 아이디어

Sparse Autoencoder(SAE)는 superposition을 "역전"시키는 도구입니다. dense한 activation을 sparse한 feature activation으로 분해합니다.

SAE 구조

모델의 activation 벡터 $\mathbf{x} \in \mathbb{R}^{d_{\text{model}}}$에 대해:

Step 1 — Encode (sparse feature 추출):

$$f(\mathbf{x}) = \text{ReLU}(W_{\text{enc}} (\mathbf{x} - \mathbf{b}_{\text{dec}}) + \mathbf{b}_{\text{enc}})$$

• $W_{\text{enc}} \in \mathbb{R}^{d_{\text{sae}} \times d_{\text{model}}}$: encoder 가중치
• $\mathbf{b}_{\text{enc}} \in \mathbb{R}^{d_{\text{sae}}}$: encoder bias
• $d_{\text{sae}} \gg d_{\text{model}}$: SAE의 feature 수는 모델 차원보다 훨씬 많습니다 (보통 4x ~ 256x)
• ReLU에 의해 대부분의 $f_i(\mathbf{x})$는 0 → sparse

Step 2 — Decode (재구성):

$$\hat{\mathbf{x}} = W_{\text{dec}} \cdot f(\mathbf{x}) + \mathbf{b}_{\text{dec}}$$

• $W_{\text{dec}} \in \mathbb{R}^{d_{\text{model}} \times d_{\text{sae}}}$: decoder 가중치
• 각 열 $\mathbf{d}_i$는 feature $i$의 방향 벡터 (단위 벡터)

풀어 쓰면:

$$\hat{\mathbf{x}} = \mathbf{b}_{\text{dec}} + \sum_{i:\, f_i > 0} f_i(\mathbf{x}) \cdot \mathbf{d}_i$$

이것이 핵심입니다: activation이 소수의 interpretable feature들의 가중합으로 분해됩니다.

학습 목적함수

$$\mathcal{L} = \|\mathbf{x} - \hat{\mathbf{x}}\|_2^2 + \lambda \|f(\mathbf{x})\|_1$$

• 첫째 항: 재구성 손실 — 원래 activation을 충실하게 복원해야 합니다
• 둘째 항: L1 sparsity penalty — feature activation이 sparse해야 합니다
• $\lambda$: sparsity 강도를 조절하는 하이퍼파라미터

왜 L1인가? L2 norm($\|f\|_2^2$)은 모든 feature를 골고루 작게 만드는 반면, L1 norm($\|f\|_1$)은 대부분의 feature를 정확히 0으로 만들고 소수만 크게 유지합니다. 수학적으로, L1 정규화의 gradient는 상수($\pm\lambda$)여서 작은 값도 0으로 밀어내지만, L2의 gradient는 값에 비례($2\lambda f_i$)해서 작은 값은 느리게 줄어듭니다. 이 차이가 "sparse = 대부분 0"이라는 우리가 원하는 결과를 만듭니다.

추가 제약:

$$\|\mathbf{d}_i\|_2 = 1$$

(decoder 열의 단위 노름 제약)

이 제약이 없으면, encoder 가중치를 키우고 decoder 가중치를 줄여서 L1 penalty를 우회할 수 있습니다.

Monosemantic Features: 실제로 작동한다

Anthropic의 "Towards Monosemanticity" (Bricken et al., 2023)는 SAE로 추출한 feature가 실제로 monosemantic(단의적)임을 보였습니다. 각 feature는 하나의 명확한 개념에 대응합니다:

뉴런 하나로는 불가능했던 깨끗한 해석이, SAE feature 하나로는 가능해집니다.

Feature Splitting과 Feature Death

SAE의 크기 $d_{\text{sae}}$를 키우면 흥미로운 현상이 나타납니다:

Feature splitting: "Python 코드" feature가 → "Python import문", "Python 함수 정의", "Python 문자열 조작"으로 쪼개집니다. 이는 개념의 위계 구조를 시사합니다.

Dead features: 어떤 데이터에도 활성화되지 않는 feature입니다. 낭비된 용량이므로, dead feature 비율은 SAE 품질의 주요 진단 지표입니다.

SAE의 진화: ReLU를 넘어서

초기 SAE는 ReLU activation을 사용했지만, dead feature 문제와 sparsity 제어의 한계가 드러나면서 여러 변형이 등장했습니다:

• TopK SAE (Gao et al., 2024): ReLU 대신 상위 K개 feature만 활성화합니다. $f(\mathbf{x}) = \text{TopK}(W_{\text{enc}} (\mathbf{x} - \mathbf{b}_{\text{dec}}) + \mathbf{b}_{\text{enc}})$. L1 penalty가 필요 없고, 정확히 K개의 feature가 활성화되므로 sparsity를 직접 제어할 수 있습니다. dead feature 문제도 크게 완화됩니다.
• JumpReLU SAE (Rajamanoharan et al., 2024): 학습 가능한 threshold $\theta$를 가진 activation — $f(\mathbf{x}) = (z - \theta) \cdot \mathbf{1}[z > \theta]$. ReLU보다 sharp한 sparsity를 만들어 재구성 품질과 sparsity 사이의 트레이드오프를 개선합니다.

Scaling Monosemanticity: Claude에 SAE를 적용하다

Anthropic의 "Scaling Monosemanticity" (Templeton et al., 2024)는 SAE를 toy model이 아닌 production-scale 모델(Claude 3 Sonnet)에 적용한 결과입니다.

핵심 결과는 다음과 같습니다:

• 3400만 개의 feature를 추출 ($d_{\text{sae}}$ = 34M)
• 다국어, 멀티모달 개념까지 포착: "Golden Gate Bridge" feature뿐 아니라 "안전성 관련 거짓말", "코드 보안 취약점" 같은 추상적 개념도 monosemantic하게 분리됨
• feature를 인위적으로 조작(clamping)하면 모델 행동이 예측 가능하게 변함 → feature가 단순 상관이 아니라 인과적임을 확인
• SAE가 lab-scale에서 production-scale로 넘어갈 수 있음을 최초로 실증

이 연구는 SAE 기반 interpretability가 실제 AI 안전 연구에 쓸 수 있는 수준에 도달했음을 보여줍니다.

SAE의 실전 적용: SAELens

SAELens는 TransformerLens에서 분리된 SAE 분석 전용 라이브러리입니다.

from sae_lens import SAE

# 사전 학습된 SAE 로드
sae = SAE.from_pretrained(
    release="gpt2-small-res-jb",
    sae_id="blocks.8.hook_resid_pre",
)

# activation을 feature로 분해
features = sae.encode(cache["blocks.8.hook_resid_pre"])

# 활성화된 feature 확인
active_features = (features > 0).sum(dim=-1)
print(f"Active features: {active_features.float().mean():.0f} / {sae.cfg.d_sae}")

3. TensorLens: Transformer 전체를 하나의 텐서로

Lens 계열의 진화 방향

지금까지의 Lens 접근법을 정리하면 다음과 같습니다:

모든 방법이 Transformer를 부분별로 분석합니다. attention은 attention대로, MLP는 MLP대로, normalization은 따로입니다.

TensorLens(Atad et al., 2025)는 이 한계를 근본적으로 해결합니다:

Transformer 전체 계산을 하나의 4차 텐서로 표현합니다.

직관적 이해: Attention을 4K 해상도로 업그레이드

기존의 attention matrix는 $L \times L$ 크기입니다. "토큰 A가 토큰 B에 얼마나 주목하는가"를 하나의 숫자로 표현합니다.

하지만 이것은 너무 거칩니다. 토큰 "Paris"의 768개 channel 중 어떤 channel이 중요한지, 그 영향이 출력 토큰의 어떤 channel에 도달하는지는 알 수 없습니다. 마치 "서울에서 부산으로 택배를 보냈다"는 정보만 있고, 무엇을 보냈는지, 어디에 도착했는지 모르는 것과 같습니다.

TensorLens는 이를 $L \times D \times L \times D$로 확장합니다:

$$\begin{aligned} \text{Attention:} \quad & [\text{토큰} \rightarrow \text{토큰}] \rightarrow L \times L \text{ 행렬 (저해상도)} \\ \text{TensorLens:} \quad & [\text{토큰} \times \text{채널} \rightarrow \text{토큰} \times \text{채널}] \rightarrow L \times D \times L \times D \text{ 텐서 (고해상도)} \end{aligned}$$

이것은 "서울 강남구의 A 물류센터에서 부산 해운대구의 B 물류센터로 전자제품을 보냈다" 수준의 정밀한 추적입니다.

수학적 세부사항: High-Order Attention-Interaction Tensor

이 섹션은 수학적 표기에 익숙한 독자를 위한 것입니다. 직관적 이해는 위 섹션만으로 충분합니다.

TensorLens의 central object는 4차 텐서 $T \in \mathbb{R}^{L \times D \times L \times D}$입니다:

$$F(X)[i,:] = \sum_j T[i,:,j,:] \cdot X[j,:]^T + B$$

여기서:

• $F$는 전체 Transformer 함수
• $X$는 입력 임베딩 행렬 $[L, D]$
• $T[i,:,j,:]$는 $D \times D$ 행렬로, 입력 토큰 $j$가 출력 토큰 $i$에 미치는 영향을 나타냅니다
• $L$은 시퀀스 길이, $D$는 hidden dimension

이 텐서가 있으면, 모델의 최종 출력을 입력에 대한 선형 연산으로 정확히 재구성할 수 있습니다.

직관적으로: attention matrix가 "누가 누구를 보는가" ($L \times L$)를 알려준다면, TensorLens의 텐서는 "누가 어떤 정보를 꺼내서, 누구의 어떤 정보에 보태는가" ($L \times D \times L \times D$)까지 포착합니다.

각 Component의 텐서화

TensorLens가 강력한 이유는 Transformer의 모든 sub-component를 텐서 형태로 통합하기 때문입니다:

Self-Attention:

$$\text{vec}[\text{Attn}(X)] = \sum_h \left((W_{v,h} W_{o,h})^T \otimes A_h\right) \text{vec}[X]$$

Kronecker product $\otimes$가 token-mixing (attention matrix $A_h$)과 channel-mixing (value/output projection)을 결합합니다.

Kronecker product(`$\otimes$`)란? 두 행렬 A($m \times n$)와 B($p \times q$)의 Kronecker product는 $(m \cdot p) \times (n \cdot q)$ 크기의 행렬로, A의 각 원소 $a_{ij}$를 $a_{ij} \cdot B$ 블록으로 치환한 것입니다. 여기서는 attention이 "어떤 토큰끼리 연결하는가"(token-mixing)를, value/output이 "연결된 정보를 어떻게 변환하는가"(channel-mixing)를 담당하며, $\otimes$가 이 둘을 하나의 행렬로 합친다고 이해하면 됩니다.

LayerNorm:

forward pass에서의 실제 variance 값을 고정(freeze)하여 선형화합니다.

FFN/Activation:

비선형 activation을 $\phi(z)/z$ 비율로 선형화합니다:

• GELU: $H = 0.5 \cdot (1 + \text{erf}(z/\sqrt{2}))$
• SiLU (SwiGLU): $H = \sigma(z)$

전체 Transformer Block:

$$T^n = L_2^n (M^n + I) L_1^n (A^n + I)$$

여기서 $L_1, L_2$는 LayerNorm, $A$는 attention, $M$은 MLP, $I$는 residual connection입니다.

모델 전체:

$$\text{vec}[F(X)] = (T^N \cdot T^{N-1} \cdots T^1) \text{vec}[X]$$

모든 layer의 텐서를 곱하면 전체 모델의 텐서를 얻습니다.

Relevance Score: 토큰 영향력 정량화

4차 텐서를 $[L, L]$ 행렬로 축소하면, 각 입력 토큰이 각 출력 토큰에 미치는 영향력을 구할 수 있습니다:

방법 1 — Norm 기반 (Eq. 20):

$$T_{\text{norm}}[i] = \|T[\text{out},:,i,:]\|_2$$

$D \times D$ 블록의 L2 norm을 취합니다. 방향에 무관하게 영향의 "크기"를 측정합니다.

방법 2 — Input+Output 내적 (Eq. 21):

$$T_{\text{IO}}[i] = X_{\text{out}}^T \cdot T[\text{out},:,i,:] \cdot X_i$$

입력과 출력 hidden state를 모두 고려한 relevance입니다. 이 방법은 relevance 합이 출력의 norm에 비례하는 성질이 있습니다.

방법 3 — Class-specific:

$$\text{rel}[i] = U[:,\text{class}]^T \cdot T[\text{out},:,i,:] \cdot X_i$$

LM head의 특정 클래스 가중치를 사용하여, 특정 예측에 대한 각 토큰의 기여를 직접 계산합니다.

Attention-only 기법 대비 성능

TensorLens의 핵심 주장은 attention만 보는 것으로는 부족하다는 것입니다:

Perturbation test에서 TensorLens가 기존 attention 기반 방법을 일관되게 크게 앞섭니다. FFN, normalization, residual connection의 기여를 무시하면 불완전한 분석이 됩니다.

지원 모델

TensorLens는 다양한 아키텍처를 지원합니다:

4. SAE + TensorLens: 상호 보완

SAE와 TensorLens는 서로 다른 각도에서 같은 목표를 향합니다:

두 접근법을 결합하면:

1. TensorLens로 "어떤 입력 토큰이 출력에 중요한지" 파악
2. SAE로 "그 토큰 위치의 activation에 어떤 feature가 활성화되어 있는지" 분해
3. → 토큰-feature-출력의 완전한 인과 경로 추적

5. Lens 패러다임의 전체 지도

세 편의 글에서 다룬 Lens 계열을 전체적으로 정리하면 다음과 같습니다:

$$\begin{aligned} & \textbf{관측 (Observation)} \\ & \quad \text{Logit Lens: "이 layer에서 뭘 예측하나?"} \\ & \quad \text{Tuned Lens: "더 정확하게 읽기"} \\[6pt] & \textbf{인과 분석 (Causal Analysis)} \\ & \quad \text{Activation Patching: "이 activation이 원인인가?"} \\ & \quad \text{Path Patching: "정보가 어떤 경로로 흐르나?"} \\[6pt] & \textbf{구조 분해 (Structural Decomposition)} \\ & \quad \text{SAE: "dense activation 안에 어떤 feature가 있나?"} \\ & \quad \text{TensorLens: "모델 전체를 하나의 선형 연산자로 보기"} \end{aligned}$$

이 도구들은 서로 경쟁하는 것이 아니라 보완합니다. 각각이 모델을 다른 해상도와 관점에서 분석하는 현미경입니다.

Wrap-up

Interpretability는 이제 단순히 "뉴런이 뭘 하는지 보기"를 넘어, 구조적이고 수학적인 분석의 시대로 접어들었습니다.

이 도구들이 궁극적으로 향하는 질문은 하나입니다:

모델이 "왜" 그렇게 답했는지를, 수학적으로 완전하게 설명할 수 있는가?

아직 그 답에 도달하지 못했지만, 매년 한 걸음씩 가까워지고 있습니다.

References

• Elhage et al. *Toy Models of Superposition* (2022)

https://arxiv.org/abs/2209.10652

• Bricken et al. *Towards Monosemanticity: Decomposing Language Models with Dictionary Learning and Sparse Autoencoders* (2023)

https://transformer-circuits.pub/2023/monosemantic-features

• Atad et al. *TensorLens: End-to-End Transformer Analysis via High-Order Attention Tensors* (2025)

https://arxiv.org/abs/2601.17958

• TensorLens GitHub

https://github.com/idoatad/TensorLens

• SAELens GitHub

https://github.com/jbloomAus/SAELens

• Templeton et al. *Scaling Monosemanticity: Extracting Interpretable Features from Claude 3 Sonnet* (2024)

https://transformer-circuits.pub/2024/scaling-monosemanticity

• Gao et al. *Scaling and Evaluating Sparse Autoencoders* (2024)

https://arxiv.org/abs/2406.04093

• Rajamanoharan et al. *Jumping Ahead: Improving Reconstruction Fidelity with JumpReLU Sparse Autoencoders* (2024)

https://arxiv.org/abs/2407.14435

• Anthropic. *A Mathematical Framework for Transformer Circuits*

https://transformer-circuits.pub