Dictionary Learning Improves Patch-Free Circuit Discovery (ArXiv 2024)

https://www.dropbox.com/scl/fi/0eel4o3ae3qjkqt3micz5/arxiv24_Patch-Free_Neural_Circuit_Discovery.pdf?rlkey=xbhwhcwt7qunm28f86t0obv9o&dl=0

1. 논문 핵심

이 논문은 Sparse Dictionary Learning/SAE로 얻은 monosemantic feature들 사이의 circuit을 activation patching 없이 발견하는 방법을 제안한다. 대상 모델은 Othello-GPT, 즉 오델로 다음 합법 수 예측을 학습한 6-layer, hidden size 128의 작은 decoder-only Transformer이다. 핵심 주장은 다음이다: residual stream에 write하는 모든 module output, 즉 embedding, attention output, MLP output을 dictionary feature로 분해하면, logit이나 상위 feature에서 하위 feature로 contribution을 역추적할 수 있다.

2. 기존 Activation Patching과의 차이

기존 circuit discovery는 보통 특정 activation/head/neuron을 clean/corrupted input 사이에서 patching하거나 ablation하여 성능 변화를 본다. 그러나 이 논문은 patching의 문제로 out-of-distribution activation, hydra/backup effect, 높은 계산 비용을 지적한다. 대신 모델의 선형 구조를 이용해 feature contribution을 직접 분해한다.

비교하면 다음과 같다.

방법	계산량	OOD 문제
Causal patching	O(n) model forward	있음
Direct patching	O(n) module forward	있음
제안법	O(n) module forward	없음

저자들은 자신들의 방법을 patching의 완전한 대체라기보다 feature ablation의 효율적이고 OOD-free한 대안으로 본다.

3. 방법론 상세

3.1 Sparse Dictionary Learning

각 activation x를 sparse feature들의 선형 결합으로 복원한다.

$x \approx \sum_i w_i d_i$

여기서 $d_i$ 는 dictionary feature direction이고, $w_i$ 는 해당 feature의 activation magnitude이다. sparsity constraint 때문에 feature들이 가능한 한 monosemantic하게 분해된다.

3.2 Dictionary를 어디에 학습할 것인가?

논문은 5가지 후보를 비교한다.

word embedding
MLP hidden activation
residual stream
MLP output
attention output

저자들의 결론은 embedding + attention output + MLP output에 dictionary를 학습하는 것이 가장 적절하다는 것이다. 이유는 pre-norm Transformer에서는 실제 residual stream에 더해지는 정보가 각 module output이므로, residual stream 전체보다 module write vector를 분해하는 편이 circuit 분석에 더 직접적이기 때문이다.

3.3 OV circuit 분해

Attention output은 다음과 같이 표현된다.

$\text{Out}_{LXA,i}=\sum_h \sum_j\text{AttnPattern}_{ij}^h W_O^h W_V^h x_j$

여기서 $x_j$ 를 하위 module output들의 합으로 보고, 각 output을 dictionary feature로 다시 분해한다.

$x_j \approx \sum_{m}\sum_k w_k^m d_k^m$

따라서 상위 attention feature Y의 activation은 lower-level dictionary feature들의 contribution으로 분해된다. 즉, 어떤 이전 token의 어떤 feature가 attention OV 경로를 통해 현재 token의 feature를 활성화했는지 계산할 수 있다.

3.4 QK circuit 분해

Attention score는 다음과 같다.

$\text{AttnScore}_{ij}=x_i W_Q W_K^\top x_j^\top$

$x_i, x_j$ 를 각각 dictionary feature들의 합으로 분해하면, QK score는 query token의 feature와 key token의 feature pair 사이의 bilinear interaction으로 분해된다.

$\text{AttnScore}_{ij}\approx\left(\sum_{m,k} c_k^m d_k^m\right)W_Q W_K^\top\left(\sum_{n,k} c_k^n d_k^n\right)^\top$

즉, attention이 왜 특정 token을 보는지에 대해 feature-pair resonance 관점의 설명을 제공한다.

3.5 MLP feature 분해: Approximate Direct Contribution

MLP는 activation function 때문에 attention처럼 완전히 선형 분해하기 어렵다. 그래서 논문은 Approximate Direct Contribution, ADC를 제안한다.

MLP가 self-gated activation $x \cdot \sigma(x)$ 형태라고 할 때, lower-level feature $c_k, d_k$ 가 상위 MLP feature Y에 주는 contribution을 다음처럼 근사한다.

$ADC(c_k d_k)=W_{e,Y}^{LXM}W_{out}\left(W_{in} c_k d_k\cdot\sigma(W_{in}x)\right)$

핵심은 $\sigma(W_{in}x)$ 를 고정된 값으로 보고, $W_{in}c_kd_k$ 부분만 feature별 선형 contribution으로 해석하는 것이다.

4. 실험 설정

실험 대상은 Othello-GPT이다. 모델은 오델로 게임 sequence를 입력받아 다음 합법 수를 autoregressive하게 예측한다. 저자들은 Li et al. 방식에 따라 1.2M parameter, 6-layer, hidden dimension 128 decoder-only Transformer를 학습했다.

dictionary는 각 layer의 attention output과 MLP output에 학습되며, feature 표기는 예를 들어 다음과 같다.

L2A474

이는 2번째 layer attention output dictionary의 474번 feature를 의미한다.

5. 실험 결과 상세

5.1 Dictionary feature가 포착한 의미

저자들은 dictionary learning이 기존 probing 연구에서 발견된 Othello-GPT의 feature 유형을 비지도적으로 모두 찾아낸다고 보고한다. 주요 feature 유형은 다음과 같다.

Feature 유형	위치	예시
현재 move position	L0A, L0M, L1M-L3M	“현재 수가 g-3”, “c-1에서 오른쪽 말을 뒤집음”
board state	L1A-L4M	“f-1은 opponent piece”, “e-2는 own piece”
empty cell	L5A	“c-3이 비어 있음”
legal move	L5M	“b-3이 합법 수”

즉, early layer는 현재 move와 flip 정보, middle layer는 board state, final layer는 empty/legal move 정보를 주로 표현한다.

5.2 OV circuit 사례: board state 계산

첫 번째 주요 case study는 L2A474 feature이다. 이 feature는 move 8 시점의 특정 local board state를 나타낸다.

저자들은 L2A474의 activation을 하위 feature contribution으로 분해했고, 다음 구조를 발견했다.

move 6 residual stream의 L1M179: d-2가 내 말이라는 정보
move 7 residual stream의 L1M755: c-2~e-4가 상대 말이라는 정보
move 8 residual stream의 L1M379: c-3이 flip되어 내 말이 되었다는 정보

이 정보들이 L2 attention의 특정 heads를 통해 현재 residual stream으로 복사되어 L2A474를 활성화한다. 특히 head 5는 “my moves” 계열 정보를, head 2는 “opponent moves” 계열 정보를 가져오는 것으로 해석된다.

중요한 수치도 있다. 각 경우 가장 중요한 head 하나가 OV contribution의 **110%, 105%, 113%**를 설명했고, 다른 head들은 거의 0에 가까웠다. 이는 특정 head가 특정 정보 흐름을 매우 명확히 담당한다는 증거로 제시된다.

5.3 QK circuit 사례: “상대 수에 attention”하는 이유

두 번째 사례는 L1A head 0이다. 이 head는 주로 opponent’s moves에 attention하는 패턴을 보인다.

QK decomposition 결과, token 14가 token 13을 보는 attention score는 다음 feature pair들로 설명된다.

token 13 쪽 feature: token 13이 white move임을 나타내는 position embedding 및 L0M/L0A feature
token 14 쪽 feature: token 14가 black move임을 나타내는 feature

따라서 attention pattern은 단순히 위치 때문이 아니라, 현재 token과 과거 token의 move color 관계, 즉 “내 수/상대 수” 관계를 feature-pair interaction으로 계산한 결과라고 해석된다.

5.4 MLP 사례: flip된 말을 인식하는 회로

세 번째 사례는 L1M379이다. 이 feature는 c-3과 d-4가 내 색이라는 board state를 나타내며, 특히 c-3이 flip된 사실을 반영한다.

ADC 분석 결과, L1M379를 활성화하는 핵심 하위 feature는 세 개의 연속된 move에 대응한다.

move 6: d-2
move 7: c-3
move 8: b-4

즉, 모델은 세 위치가 한 줄에 놓이고, 색 구조가 white-black-white 형태임을 조합해 가운데 말을 뒤집었다고 판단한다. 논문은 이를 MLP가 일종의 AND gate처럼 작동하여 flip 조건을 결합한 것으로 해석한다.

또한 L0A head 6과 head 7이 함께 L0A629를 활성화하는 현상을 통해 attention superposition도 관찰했다고 보고한다.

5.5 ADC와 direct patching의 대응성

ADC가 실제 patching 기반 attribution과 얼마나 비슷한지 검증하기 위해, 저자들은 다음 실험을 수행했다.

10개 input sequence를 random sample
6개 MLP layer에서 top-3 activated dictionary feature 선택
총 180개 MLP feature 분석
각 feature에 대해 ADC와 direct patching이 찾은 top-5 lower-level contributor 비교
metric: contributor set의 IoU

결과는 average IoU = 0.68이다. 이는 ADC가 direct patching과 상당히 높은 대응성을 가진다는 근거로 제시된다.

6. 한계

주요 한계는 세 가지이다.

첫째, QK decomposition은 softmax 이전 attention score만 분해한다. softmax는 전체 sequence의 score들을 상대적으로 정규화하므로, 개별 score contribution만으로 attention pattern 전체를 완전히 설명하기 어렵다.

둘째, ADC는 MLP의 positive contribution은 잘 포착하지만, 어떤 feature가 negative-effect neuron을 억제함으로써 상위 feature를 활성화하는 경우는 잘 포착하지 못한다.

셋째, dictionary training 자체가 완벽하지 않았다. 저자들은 sparsity coefficient가 너무 작아 reconstruction은 거의 완벽했지만, 각 dictionary에서 약 75% dead neurons가 발생했다고 보고한다.

7. 평가

이 논문의 핵심 기여는 SAE/dictionary feature를 circuit discovery의 node로 사용하고, patching 없이 edge contribution을 계산하려 했다는 점이다. 특히 OV, QK, MLP를 각각 다른 수학적 구조로 분해한 것이 중요하다.

정리하면:

$\textbf{SAE feature discovery}\quad \rightarrow \quad\textbf{feature-to-feature attribution}\quad \rightarrow \quad\textbf{patch-free circuit graph}$

라는 흐름을 제시한 논문이다. 실험은 Othello-GPT라는 synthetic model에 한정되지만, 향후 LLM의 SAE feature circuit discovery, feature-level EAP, SAE 기반 causal graph extraction 연구로 확장될 수 있는 출발점으로 볼 수 있다.