기술 요약: 조건부 공동 결실 (Conditional Co-Ablation, COAX)

1. 문제 정의

기계론적 해석 가능성(Mechanistic interpretability)은 모델이 특정 행동을 생성하는 방식을 이해하기 위해 개별 모델 구성 요소(예: 어텐션 헤드)에 중요도 점수를 할당하는 데 의존합니다. 표준적인 접근 방식은 **1차 스코어링(first-order scoring)**을 사용합니다. 이는 단일 구성 요소를 격리하여 결실(ablate)시키고, 출력 변화를 측정하며, 이러한 효과를 집계하는 방식입니다. 이 방식은 구성 요소의 중요성이 **가산적(additive)**이라고 가정합니다.

그러나 이 가정은 자기 수복(self-repair)(일명 "히드라 효과")이 존재하는 경우 실패합니다. 자기 수복이란 모델이 주 구성 요소가 제거되었을 때만 활성화되는 휴면 백업 구성 요소를 보유하고 있는 상태를 의미합니다.

• 사각지대: 온전한 모델 상태에서 백업 구성 요소는 조용히 잠들어 있기 때문에 무관한 것(낮은 중요도)으로 보입니다. 하지만 주 구성 요소가 결실되면, 백업 구성 요소가 역할을 이어받아 손상을 복구함으로써 주 구성 요소의 측정된 효과를 상쇄합니다.
• 결과: 1차 스코어링은 양방향의 중복성을 잘못 읽습니다. 즉, 백업이 오류를 수정하기 때문에 주 구성 요소의 중요도를 과소평가하고, 백업 구성 요소는 휴면 상태이므로 완전히 놓치게 됩니다. 이는 잘못된 귀속(attribution), 실패한 능력 제거(capability knockout), 그리고 최적화되지 않은 모델 프루닝(pruning)으로 이어집니다.

2. 방법론: 조건부 공동 결실 (COAX)

저자들은 이 문제를 **조건부 회로 완성(conditional circuit completion)**으로 재정의합니다. 즉, 일련의 주 구성 요소 집합 $S$가 주어졌을 때, $S$가 제거된 후 중요도가 유의미하게 증가하는 나머지 단위들을 식별하는 작업입니다.

COAX 점수

COAX는 레이블이 필요 없는(label-free), 출력에 근거한(output-grounded) 점수로, 단위의 결실 효과가 나타나는 **조건부 성장(conditional growth)**을 측정합니다.

1. 피셔 기하학(Fisher Geometry): 이 방법은 출력 분포에 의해 유도된 메트릭(피셔 정보량) 내에서 결실 효과를 측정합니다. 단위 $u$에 대해, 결실 효과는 로짓 섭동 $\delta z_u = z_0 - z_{\{u\}}$입니다. 이 효과의 에너지는 중심화된 특징(centered feature) $\tilde{\delta z}_u$의 피셔 가중 노름(Fisher-weighted norm)으로 근사된 평균 KL 비용입니다.
2. 조건부 성장:
   • $S$를 주 구성 요소 시드 집합(예: 1차 방법으로 식별됨)으로 설정합니다.
   • 후보 단위 $u$를 단독으로 결실시켰을 때의 에너지를 계산합니다: $E(\delta z_u | \emptyset)$.
   • $S$가 결실된 후에 $u$를 결실시켰을 때의 에너지를 계산합니다: $E(\delta z_u | S)$.
   • COAX 점수는 그 차이입니다:
     $$\text{comp}_u(S) = E(\delta z_u | S) - E(\delta z_u | \emptyset)$$
3. 이론적 기초:
   • 명제 1: 중심화된 특징의 내적을 피셔 그램 행렬(Fisher Gram matrix)과 연결하는 피셔 항등식을 확립합니다.
   • 명제 2 (가산성 사각지대): 깨끗한 상태의 통계량(예: 깨끗한 활성화 또는 깨끗한 그래디언트)에만 기반한 점수는 순수 백업(깨끗한 모델에서 휴면 상태인 것)과 무력한 단위를 구분할 수 없음을 증명합니다.
   • 명제 3 (시너지로서의 조건부 성장): 휴면 단위에 대해, 조건부 성장 점수가 시드 집합 $S$와의 쌍방향 시너지(pairwise synergy) 에너지를 근사함을 보여줍니다. 즉, COAX는 1차 스코어링이 버리는 2차 상호작용을 복구합니다.

효율성

명시적인 쌍방향 시너지 계산은 $O(|U|^2)$의 순전파(forward pass)를 요구하는 반면, COAX는 $O(|U|)$의 순전파(단일 결실 스캔 1회 및 조건부 스캔 1회)만을 요구합니다. 이 방법은 순전파 전용이며 역전파(backward pass)나 태스크 레이블을 필요로 하지 않습니다.

3. 주요 기여

1. 자기 수복의 공식화: 논문은 자기 수복을 구성 요소 중요성의 가산성 실패로 공식화하고, "조건부 회로 완성"을 휴면 보상기(compensator)를 복구하는 작업으로 도입했습니다.
2. COAX 점수: 주 구성 요소의 결실 후 에너지 성장을 측정함으로써 휴면 백업을 드러내는, 레이블이 필요 없고 피셔에 근거한 점수를 도입했습니다.
3. 인과적 검증: 복구된 헤드들이 진정한 백업임을 깨우기 곡선(wake-up curves), 직접 로짓 귀속 핸드오프(direct-logit attribution hand-off), 역사실적 패칭(counterfactual patching, 백업을 고정하면 자기 수복이 제거됨) 등의 기계론적 검증을 통해 입증했습니다.
4. 다운스트림 유용성: 이러한 백업들을 복구함으로써 마스킹된 귀속을 교정하고, 완전한 능력 제거(capability knockout)를 가능하게 하며, 구조적 프루닝을 개선함을 보여줍니다.

4. 실험 결과

백업 발견 (GPT-2-small IOI 회로)

• 태스크: 3개의 주 이름 이동(name-mover) 헤드가 주어진 상황에서, 기록된 8개의 백업 이름 이동 헤드를 식별합니다.
• 성능:
  • 1차 베이스라인: 단일 결실 돌출도(single-ablation saliency) (0.33 AUC), AtP (0.60), EAP-IG (0.70), 그리고 자기 수복을 고려한 AtP*GradDrop (0.82).
  • COAX: 0.91 ROC-AUC를 달성합니다.
  • 의의: COAX는 기록된 8개 백업 중 6개를 상위 20개 후보에 포함시킵니다. 시드 집합을 통제했을 때, 가장 우수한 그래디언트 기반 베이스라인(AtP*)보다 통계적으로 유의미한 개선($p < 0.05$)을 보였습니다.
• 합성 벤치마크: 100개의 심어진 백업이 있는 합성 벤치마크에서, 깨끗한 상태의 점수는 우연 수준이거나 그 이하(0.42–0.51)를 기록한 반면, COAX는 0.90에 도달하여 1차 방법의 이론적 사각지대를 확인했습니다.

일반화

• 규모: "깨우기(wake-up)" 현상(주 구성 요소 결실 시 출력 노름 증가)은 GPT-2-small, medium, large 모델 전반에서 관찰됩니다.
• 아키텍처: 8개 모델(Pythia, GPT-Neo, Gemma, Llama, Qwen, OLMo)의 유도(induction) 회로에 적용되었습니다. COAX는 인덕션 로그 확률을 유의미하게 떨어뜨리는(무작위 대비 2.1배~12배 더 많이) 보상 집합을 복구하며, 레이블 없이도 전이 가능함을 입증했습니다.

다운스트림 응용

1. 귀속 복구: 주 구성 요소만 결실시키면 IOI 로짓 차이가 0.22만큼만 감소합니다(자기 수방에 의해 마스킹됨). COAX로 복구된 백업을 추가하면 감소 폭이 1.76으로 회복되어, 실제 인과적 효과를 드러냅니다.
2. 능력 제거 (Capability Knockout): IOI 행동을 비활성화하려면 백업을 제거해야 합니다.
   • 주 구성 요소만 제거 시: 정확도가 약 0.97으로 유지됩니다 (자기 수복 활성).
   • 주 구성 요소 + COAX 복구 백업 제거 시: 정확도가 0.70으로 떨어져, 기록된 오라클(0.72)과 일치합니다.
   • 주 구성 요소 + "자신의" 다음 순위 헤드들(1차 방식 상위 추출) 제거 시: 정확도가 0.24로 떨어져, 핵심 구성 요소까지 과잉 결실됨을 나타냅니다.
3. 구조적 프루닝: 제거 후 COAX 점수를 다시 측정하는(새롭게 노출된 백업을 반영하기 위해) 순차적 프루닝 전략은 124M에서 7B 파라미터 규모에 이르기까지 모든 규모에서 모든 베이스라인(Random, Magnitude, Wanda, Taylor)보다 우수한 성능을 보입니다. 50% 희소성(sparsity)에서 COAX는 밀집 모델(dense model)의 퍼플렉시티를 경쟁 모델들보다 훨씬 더 가깝게 유지합니다.

5. 의의 및 주장

본 논문은 구성 요소의 중요성이 단위의 고유한 속성이 아니라, 나머지 회로의 상태에 따라 결정되는 조건부 속성이라고 주장합니다.

• 동기: 견고한 회로에서 가장 중요한 구성 요소는 종종 주 구성 요소가 실패할 때까지 조용히 대기하는 구성 요소들입니다. 단일 단위로 점수를 매기는 1차 방법은 이러한 "숨겨진 영웅"들을 체계적으로 놓칩니다.
• 운영 패러다임의 전환: COAX는 "단위를 격리하여 점수를 매기는 것"에서 "제거된 것에 상대적으로 점수를 매기는 것"으로 패러다임을 전환합니다.
• 한계 및 범위:
  • COAX는 독립적인 발견 도구가 아닌 완성(completion) 도구입니다. 즉, 주 시드 집합(1차 방법으로 찾을 수 있는)이 필요합니다.
  • COAX는 휴면 대체(dormant-substitution) 체제(예: IOI 백업)에서 가장 효과적입니다. 공동으로 작동하는 동질적 헤드들 사이에 중복성이 공유되는 체제(예: 인덕션)에서는, COAX가 부하를 담당하는 집합을 찾아내지만 독특한 "백업 클래스"를 식별하지는 못합니다.
  • 현재 이 방법은 어텐션 헤드 단위로 구현되어 있습니다. MLP가 지배적인 회로(예: "greater-than" 회로)에서는 강력한 신호를 복구하지 못하는데, 이는 회로의 메커니즘(MLP vs. Attention)이 자기 수복의 복구 가능성을 결정함을 시사합니다.

결론적으로, 본 논문은 견고한 모델에 대한 충실한 해석 가능성은 반드시 조건부여야 한다고 주장합니다. 즉, 중복성은 평균화하여 제거해야 할 노이즘이 아니라, 전체 회로를 드러내기 위해 반드시 조건화해야 할 구조입니다.