[Transcoders] Find Interpretable LLM Feature Circuits

(NeurIPS 2024)

https://arxiv.org/pdf/2406.11944

https://github.com/jacobdunefsky/transcoder_circuits/

GitHub - jacobdunefsky/transcoder_circuits

 

sparse autoencoder는 interpretable feature를 extract할 수 있음을 보였고, large-scale language model에도 적용가능함을 보였다.

하지만 limitation이 있는데, circuit analysis 에 쓰일 수 없다는 것이다. 

왜냐하면 각각의 feature는 많은 neuron의 linear combination이고 각 neuron은 nonlinearity를 가지고 있기 때문에 layer 간의 feature connection을 찾기 어렵다는 것이다. 

이를 해결하기 위해 transcoder를 도입한다.

단순히 MLP layer의 activation을 sparse하게 decomposition하는 SAE와 다르게,

transcoder는 MLP의 연산 자체를 approximate하면서 sparse하게 activate한다. 

SAE와의 비교에서, SAE 못지 않게, 혹은 더 나은 sparsity, interpretability, faithfulness (feature가 sparse하게 activate하는가?, MLP의 연산을 정확히 근사하는가? feature가 meaningful한가?)를 보인다.

따라서 TC는 SAE를 대신하여, feature-level에서 MLP layer를 통한 circuit anlysis를 가능하게 한다. 

중요한 것은, TC는 circuit을 input-dependent & input-invariant term으로 factorize한다. 

input-invaraint term을 disentangle하였다는 것은 genera behavior를 해석할 수 있다는 것이다.
input-dependent term은 현재 input에 대한 feature가 다음 feature에 주는 영향이고, input-invariant한 term은 feature pair 간의 general connection을 말해준다. 

※ circuit analysis를 한다는 것은, 결국 어떤 model이 task를 수행할 때 보이는 behavior를 설명하는 subgraph를 찾는 것인데, 이는 edge의 강도(strength)를 찾는 것이다.  이를 위해서는 feature들이 layer간에 어떻게 connection 되는지를 생각해야 하는데, (앞선 node가 뒤 node에 얼마나 contribution하는지), 이를 위해서는 TC의 encoder weight과 decoder weight의 역할을 이해하는 게 포인트이다.

encoder feature vector는 input을 보고, 각 feature를 얼만큼 activate할 것인가를 결정하고,
decoder feature vector는 각 feature direction을 담당하고 있다. 

이렇게 attribution을 계산할 수 있게 되면, recursive하게 거슬러 올라가면서 circuit-finding을 할 수 있다.

예를 들어서 1-> 2-> 3 이런 layer가 있다면, 3 layer의 어떤 feature를 activate한 2layer의 feature가 있을 것이다. 그 connection을 제외하고 pruning하고, 그 2 layer의 feature를 activate한 1 layer의 feature 와의 connection을 제외하고 pruning하고.. 이를 반복하여 connected edge sequence가 생기면 이 computational path들을 combination하여서 circuit을 완성할 수 있다.

 

high-level idea는 이해하였는데, 실제로 circuit을 찾는 algorithm을 이해하지를 못하였다. 

Attribution 연산 식을 좀더 차근차근 살펴보아야 한다. ㅠㅠ  

 

 

---

---

---

---

---

---

---

---

---

---