Towards Automated Circuit Discovery for Mechanistic Interpretability

• Keyword : ACDC, circuits, automatic interpretability
• Novelty : 자동화된 모델 circuit만드는 방법 제안
• Summary : 다음 노드에 활성화 영향을 주는 부분만 남기고 나머지 부분을 잘라내는 방법으로 subgraph를 만든다. 일반적으로 사람이 manual crafting 을 진행했다면 본 논문은 해당과정을 자동으로 진행하는 방법을 제안한다.

해당 논문은 metric과 dataset을 정해서 model behavior를 관찰하였다. 다음으로 activation patching 을 통해서 NN units 이 model behavior에 관여하는지 보았다. 데이터셋을 다양화 함으로써 각 NN component의 기능을 이해할 수 있었다.

Reverse-engineering이란 model components를 사람이 이해할 수 있는 형태로 변경하는 작업을 의미한다. 현재까지의 접근방식은 사람이 수동으로 모델 행동에 대해 관찰하는 방식을 사용했다. 해당 작업은 스케일이 큰 모델에 대해 다양한 behavior가 존재하고 많은 sub-circuits 가 존재하게 되는 복잡하고 계산하기 힘들다는 단점이 존재한다. 본 논문에서는 자동화된 방식을 소개한다.

첫 번째로 저자는 일반적으로 모델 해석에 많이 사용되는 workflow를 시스템화 했다. 해당 스텝중 하나는 subgraph of the model을 찾는 과정이다. 해당 과정을 자동화시켰다. 본 논문은 Automatic Circuit DisCovery (ACDC) 라는 알고리즘을 소개한다 (여기서 subnetwork probing 과 Head Importance Score for Pruning (HISP)를 적용한다). 마지막으로 양적화한 metirc을 통해 알고리즘을 통해 나온 circuit 의 평가를 진행한다.

The Mechanical Interpretability Workflow

일반적인 workflow는 아래의 3단계를 진행한다.

1. Observe a behavior(or task) that a neural network displays, create a dataset that reproduces the behavior in question, and choose a metric (행동 (helpful or truthful or honesty ..), 데이터셋, metric 설정):

2. Define the scope of the interpretation, (what level of granularity) (모델내에서 얼마나 세분화 해서 볼 건지)

3. Perform an extensive and iterative series of patching experiments with the goal of removing as many unnecessary components. (Patch model activation 관련된 서브그래프를 만들기 위함)

   해당과정을 3가지 스텝이 진행된다. (In directed acyclic graph (DAG)의 nodes의 중요도를 식별하기 위해) 이는 recursive activation patching 이라는 방식이 사용되고 아래의 3개의 과정이 진행된다

   1. Overwriting the activation value of a node or edge with a corrupted activation
   2. Running a forward pass through the model
   3. Comparing the output values of the new model with the original model using the chosen metric

   해당 과정을 반복하는 이유는 필요없는 요소와 conntection을 제거하기 위함이고 관련있는 subgraph만을 남기기 위함이다.

Patching with zeros and patching with different activations

Activation patching 방식은 모델 해석 관점에서 다양한 방식이 존재한다. 몇몇 프로젝트는 activation value 를 0으로 만들어버리는 방법이 존재하고 최근에는 interchange intervention 을 사용한다. 이는 node’s activation vlaue를 다른 데이터 포인트의 값으로 overwriting 시키는 방법이다. 본 논문에서는 후자의 방법을 선호하고 모든 방법에 대해 실험을 진행한다고 한다.

Automating circuit discovery (step3를 자동화)

본 논문에서는 step3에서의 3개의 케이스 모두에서 $(x_i)^n_{i=1}$ 인 a set of prompts 정의했다. 그리고 a set of corrupted prompts 를 $(x_i^{\prime})^n_{i=1}$ 으로 정의했다.

Automatic Circuit DisCovery (ACDC): ACDC는 output 에서 시작해 input으로 graph를 통해 연산을 반복한다. 모든 노드들은 하나 씩 제거되면서 모델 퍼포먼스에 감소가 되는지 확인한다. 이때 connection이 제거되었을 때 모델 성능 하락이 없으면 제거를 진행한다.

To formulize the ACDC process, $G$ 를 모델의 computational graph 이라고 하고, $H \in G$ 라고 말할 때 $H$는 반복적인 pruning 을 통해 그리고 threshold($H$의 마지막 state의 sparsity를 결정하는)를 넘는 것에 대한 subgraph로 말할 수 있다.

저자들은 $H$에 대한 평가방법을 제안한다. $H(x_i, x_i^{\prime})$ 은 $x_i$에 대한 네트워크이다. 이때 $G$ 그래프에서 $H$에 해당되지 않는 부분은 $x_i^{\prime}$으로 덮어버린다. 이때 $D_{KL}(G||H)$ 를 통해서 subgraph ($H$)의 평가를 진행한다.

Subnetwork Probing: SP는 서브네트워크를 이용해 내부 모델전반에 마스킹하는것을 학습한다. 이를 통해 각 entry에 대해 0 or 1이 남게된다. 해당 과정은 언어적 특징을 linear probing 을 통해 추출할 수 있다고 한다. 본 저자는 ACDC를 사용하기 위해 몇가지 변형을 진행했다. 1) remove linear prob 2) change the trainin gmetirc to KL div 3) use the mask to interpolate between corrupted activations and clean activations rather than zero activations and clean activations

Head Importance Score for Pruning: HISP는 중요도 스코어를 바탕으로 heads의 rank를 매긴다. 그 후 top k스코어를 제외한 너무지 부분을 제거한다. 해당 top k 개의 subnetwork와 ACDC를 비교를 진행했다.

Evaluating Subgraph Recovery Algorithms

ACDC 방법이 기본 알고리즘에 해당하는 하위 그래프를 식별할 수 있는지, ACDC 방법이 유도된 동작에 참여하지 않는 구성요소를 포함하지 않는지를 중점적으로 보고자 한다.

circuit 발견을 평가를 binary classification problem으로 공식화 해서 보고자 한다. (edges are classified as positive (in the circuit) or negative (not in the circuit)).

Figure3 을 바탕으로 3가지의 해석을 진행한다. 1) 모든 방법론들이 corrupted distribution에 민감하다. 2) ACDC는 기존의 방법들과 비교할만하다. 3) ACDC는 모든 방법에 robust하게 적용되지 않고 몇몇 셋팅에서는 낮은 성능을 보인다.

Figure4는 유도된 동작에 차명하지 않는 구성요소를 포함하지 않는지를 말하는 그래프이다.
왼쪽과 오른 쪽 모두 ACDC알고리즘을 이용했을 때 더 좋은 성능이 나타난다.