https://www.dropbox.com/scl/fi/yjgk0u5i657adx5ub8633/arxiv23_Surgical_AI_Control.pdf?rlkey=39h35f7wx4hmekefdginb1hxx&dl=0

이 논문은 “모델의 특정 행동(behavior)만 제거할 수 있는가?” 라는 질문을 다룬다. 기존에는 Fine-tuning, RLHF, Model Editing 등이 주로 weight를 수정했는데, 이 논문은 훨씬 Mechanistic Interpretability 관점에서 접근한다.

핵심 아이디어는:

“나쁜 행동을 만드는 circuit 전체를 찾는 대신, 그 circuit을 끊어버리는 최소 edge cut을 찾자.”

이다.

1. 문제 정의

논문은 “behavior removal”을 다음과 같이 정의한다.

모델 M 이 특정 행동 D 에 대해서는 성능이 나빠지고,

다른 행동 D’ 에 대해서는 성능이 유지되도록 만들고 싶다.

즉

Efficacy
- target behavior 제거
Specificity
- 다른 능력은 유지

를 동시에 만족해야 한다.

예시:

Toxic text 생성 제거
Jailbreak 제거
특정 bias 제거
특정 concept 제거

2. 핵심 아이디어

기존 Circuit Discovery

기존 circuit 연구들은

“이 행동을 수행하는 데 필요한 circuit은 무엇인가?”

를 찾는다.

예:

Activation Patching
Path Patching
ACDC
EAP

등

이 논문은 반대로

“어디를 끊으면 이 행동이 더 이상 발생하지 않는가?”

를 찾는다.

즉,

Circuit Discovery: Find circuit

Circuit Breaking: Find cut

이다.

3. Computational Graph 구성

먼저 Transformer를 DAG로 표현한다.

Node:

Input
Attention Head
MLP
Output

Edge:

node 간 정보 흐름

GPT2-small의 경우:

158개 node
11,611개 edge

를 사용한다.

4. Transformer Residual Rewrite

논문의 가장 중요한 부분이다.

원래 Transformer

attention head: $A_{i,j}(R_{i-1})$

MLP: $M_i(R_{i-1})$

여기서 $R_i=R_{i-1}+\sum_j A_{i,j}+M_i$

Residual Rewrite

Residual stream을 전개하면 attention head는 사실상 모든 이전 layer output을 입력으로 받는다.

논문은

$S_i^A=I+\sum_{i'<i}\left(M_{i’}+\sum_{j’} A_{i’,j’}\right)$

로 표현한다.

즉, attention head 하나가

Input
이전 layer MLP
이전 layer Head

모두에게 edge를 가진다.

그래서 graph가 거의 dense해진다.

5. Edge Ablation

이 논문의 가장 핵심 개념.

edge e=(A,B)가 있다고 하자.

원래는 $A \rightarrow B$ 정보가 전달된다.

Zero Ablation

edge 제거 시 A=0 이라고 가정한다.

즉, B = f(…,0,…)

Mean Ablation

더 자연스러운 방법.

$A=\mu_A$ 로 치환

$\mu_A=E[A]$

즉, “정보를 제거하되 OOD activation은 만들지 말자”는 아이디어다.

6. Continuous Edge Mask Learning

논문의 가장 중요한 알고리즘.

각 edge에 $w_e \in [0,1]$ 를 둔다.

의미

w_e=1 : edge 유지

w_e=0 : edge 제거

중간값은 부분적으로 제거

destination node는 $w_e v_A+(1-w_e)\mu_A$ 를 입력으로 받는다.

즉, 원래 activation과 ablated activation을 interpolation 한다.

7. 최적화 Objective

원래 목표:

$E=\text{ablate edge set}$ 를 찾는 것

최적화 문제:

$L(G^{-E},D_{train})-\alpha L(G^{-E},D_{behavior})+\lambda R(E)$

의미:

첫 항 $L(D_{train})$ : 일반 성능 유지

둘째 항 $-\alpha L(D_{behavior})$ : 나쁜 행동 loss 증가 (behavior 제거)

셋째 항 $\lambda R(E)$ : 많은 edge 제거 방지

즉, “최소한의 edge만 제거해서 behavior를 깨뜨려라”이다.

8. Regularization

mask sparsity 유도.

논문은 $R(W)=\sum_e w_e$ 를 사용한다.

흥미로운 점은 $\lambda(t)$ 를 점점 증가시킨다.

초기:

behavior 제거 먼저

후기:

edge 수 최소화

9. Inference 단계

학습 후 $w_e$ 를 thresholding

$w_e \le \tau$ 이면 edge 제거

$\hat E=\{e:w_e\le\tau\}$

실제 추론 시, 해당 edge만 ablation.

weight는 전혀 수정하지 않는다.

10. GPT-2 Toxicity 제거 실험

Dataset

Good:

OpenWebText 10k

Bad:

4chan /pol/
toxicity > 0.9

Graph

GPT2-small

158 nodes
11,611 edges

11. 결과

가장 유명한 결과

11,611 edge 중

단 12개 edge만 제거.

Toxicity:

0.944 –> 0.567

일반 perplexity: 4.617 –> 4.623

거의 변화 없음

즉, 12개의 causal connection만 제거해도 toxicity circuit을 상당 부분 비활성화 가능

하다는 것을 보여준다.

12. ACDC, EAP와의 관계

이 논문은 사실상

ACDC → EAP 계열의 선조격 논문으로 볼 수 있다.

ACDC

찾는 것: $\text{behavior circuit}$

EAP

찾는 것: $\text{important edge}$

Circuit Breaking

찾는 것: $\text{minimum cut}$

즉, Circuit Discovery가 “수직 경로(vertical path)”를 찾는다면

Circuit Breaking은 “수평 절단(horizontal cut)”을 찾는다.

* Circuit Breaking: Removing Model Behaviors with Targeted Ablation (ArXiv 2023)