Closure-Validated Circuit Discovery in Attention Heads: Co-activation Proposes, Ablation Disposes

기술 요약: 어텐션 헤드의 폐쇄 검증된 회로 발견 (Closure-Validated Circuit Discovery in Attention Heads)

문제 정의
기계론적 해석 가능성(Mechanistic interpretability)은 개별 단위보다는 구성 요소들의 집합(서브스페이스, 매니폴드, 헤드 세트)을 기본 연구 대상으로 보는 방향으로 전환되었습니다. 널리 쓰이는 방법론 중 하나는 공활성화 통계(예: 희소 오토인코더 특징 또는 어텐션 헤드를 이용한 이진화된 Ising 모델)를 클러스터링하여 이러한 그룹을 찾는 것을 제안합니다. 그러나 결정적인 간극이 존재합니다: 통계적으로 일관된 공활성화 커뮤니티가 실제로 '회로(circuit)'라는 의미에서 부하를 담당하는(load-bearing) 실체인가 하는 점입니다. 진정한 회로란 특정 입력에 대해 해당 그룹의 기능이 적용될 때, 그 그룹을 제거했을 때 모델 성능을 인과적으로 저하시키는 구성 요소들의 집합을 의미합니다.

본 논문은 저렴한 신호들—구체적으로 공활성화 클러스터 멤버십, 어텐션 타겟 선택성(attention-target selectivity), 참여 비율(participation ratio)—이 회로 발견을 위한 충분한 근거가 될 수 있는지, 아니면 단지 인과적 검증을 필요로 하는 후보를 제안할 뿐인지 조사합니다. 저자들은 클러스터링만으로 다음을 구분할 수 있는지 묻습니다:

1. 진정한 회로 (True Circuits): 기능을 공동으로 구현하며 부하를 담당하는 그룹.
2. 공활성화 커뮤니티 (Co-active Communities): 공유된 상위 드라이버로 인해 함께 활성화되지만, 중복적이거나 미미하게 기여하는 그룹.
3. 취소 가능한 커뮤니티 (Cancellable Communities): 네트워크가 보완할 수 있는 편향 신호를 추가하는 그룹으로, 이들을 제거하면 오히려 성능이 향상될 수 있는 그룹.

방법론
본 연구는 세 가지 1B 규모 언어 모델(Pythia 1B, OLMo 1B, 그리고 Mixture-of-Experts인 OLMoE-1B-7B)과 두 가지 입력 분포(합성 유도 및 자연어 텍스트)를 대상으로 "제안 및 검증(proposal-and-validation)" 프레임워크를 채택합니다.

1. 제안 (클러스터링):

   • 신호 추출: 각 어텐션 헤드에 대해 템플릿 프리(template-free) 신호를 추출합니다: 입력 내 임의의 키(key) 위치에 할당된 최대 어텐션 가중치.
   • 이진화: 신호를 헤드당 중앙값(per-head median)에서 이진화하여, 모든 헤드가 구성상 50%의 예시에서 활성화되도록 보장함으로써 활성화율 편향을 제거합니다.
   • Ising 적합: 의사 우도(pseudolikelihood)를 사용하여 이진화된 신호에 대한 쌍별(pairwise) Ising 모델을 적합시켜 헤드 간 결합 강도를 추정합니다.
   • 커뮤니티 회복: 절대 결합 행렬(absolute coupling matrix)에 스펙트럼 클러스터링을 적용합니다. 클러스터 수($k$)는 헤드 능력(예: 이전 토큰, 유도)에 대한 지도 학습 분류에 대한 조정된 랜스 지수(Adjusted Rand Index, ARI)를 최대화하도록 선택됩니다.
   • 후보 선택: 순도(purity), 고립도(내부 대 외부 결합 비율), 그리고 크기 제약 조건을 기준으로 특정 서브 클러스터를 선택합니다.

2. 검증 (폐쇄 테스트):

   • 절제 (Ablation): 식별된 헤드들이 잔차 스트림(residual stream)으로 들어가기 전의 출력 슬라이스를 제로잉(zeroing)하여 선택된 커뮤니티를 절제합니다.
   • 비교: 예시별 손상(크로스 엔트로피 손실 변화, top-1 정확도, 타겟 토큰 로짓 변화)을 동일한 크기의 다섯 가지 매칭된 무작위 제어 절제와 비교합니다.
   • 판정 기준: 후보가 기대되는 방향(손실 증가, 정확도/로짓 감소)으로 움직이고, 무작위 제어 분포에 비해 통계적으로 유의미할 경우(z-score > 1.8) 폐쇄(closure)를 통과한 것으로 간주합니다.

3. 확장된 프록시 (Extended Proxies):

   • 본 연구는 또한 이 "저렴한" 신호들이 훈련 체크포인트 전반에 걸쳐 부하를 담당하는 기능과 상관관계가 있는지 확인하기 위해 어텐션-타겟 선택성과 **참여 비율(PR)**을 추적합니다.

주요 결과

• 밀집 모델 (Pythia 1B, OLMo 1B):

  • 합성 및 자연어 텍스트 분포 모두에서 식별된 공활성화 커뮤니티는 폐쇄를 통과합니다.
  • 이러한 커뮤니티를 절제하면 무작위 제어에 비해 일관되게 손실이 증가하고 정확도/타겟 로짓이 감소합니다.
  • OLMo 1B (자연어 텍스트): 가장 강력한 신호를 생성했습니다(손실에 대해 +36.4$\sigma$, 타겟 로짓에 대해 -50.9$\sigma$).
  • Pythia 1B (자연어 텍스트): "다운스트림 중복 시그니처(downstream-redundancy signature)"를 보였습니다. 타겟 로짓은 크게 감소(-52.5$\sigma$)했지만, 총체적인 손실 변화는 작았습니다(+1.95$\sigma$). 이는 모델이 절제된 출력을 재구성하는 중복된 다운스트림 경로를 가지고 있어, 특정 예측의 확신도는 무너지더라도 크로스 엔트로피 손실은 보호하고 있음을 나타냅니다.

• 혼합 전문가 모델 (Mixture-of-Experts, OLMoE-1B-7B):

  • 한계적 실패: 자연어 텍스트에서 주변부(marginal) Ising 모델이 붕괴되었으며(ARI $\approx$ 0.006), 구조를 회복하는 데 실패했습니다.
  • 경로 조건부 회복: 라우팅 패턴에 따라 예시를 계층화하고(라우팅 가중치에 대한 k-means 사용) 별도의 Ising을 적합시킨 결과, 통계적으로 유의미한 신호를 회복했습니다(ARI 0.191, 무작위 분할 널 모델 대비 +3.05$\sigma$).
  • 폐쇄 실패: 통계적 회복에도 불구하고, 발견된 커뮤니티는 잘못된 방향으로 폐쇄에 실패했습니다. 커뮤니티를 절제했을 때 손실이 오히려 개선되었습니다(해당 발견 층위에서 z = -5.64$\sigma$, 그 외 영역에서 -12.94$\sigma$). 이 헤드들은 경로 모듈레이션 노이즈(route-modulated noise) 역할을 했으며, 이를 제거하는 것이 특히 특정 라우팅 층위 이외의 입력에 대해 모델에 도움이 되었습니다.

• 훈련 축 프록시 (Training Axis Proxies):

  • 형태 없는 기능 (Function without Form): OLMo 1B에서 BOS(문장 시작) 헤드들은 어텐션-BOS 선택성 패턴이 형성되기 전(선택성이 거의 0일 때)인 2B 토큰 시점에서도 이미 부하를 담당(폐쇄 통과)하고 있었습니다.
  • 기능 없는 형태 (Form without Function): Pythia 1B에서 이전 토큰 헤드들은 1B 토큰 시점에 높은 선택성과 높은 참여 비율을 보였으나 폐쇄에는 실패했습니다(절제가 성능을 저하시키지 않았고, 때로는 도움이 되기도 함).
  • 이러한 양방향 디커플링은 어텐션 패턴과 기능이 동일한 현상의 서로 다른 민감도가 아니라 별개의 구조임을 입증합니다.

다중 지표 보고 (Multi-Metric Reporting)
본 논문은 크로스 엔트로피 손실 단독으로는 폐쇄를 측정하기에 신뢰할 수 없는 지표라고 주장합니다. 이는 제어 분산 붕괴(OLMo 자연어 텍s에서 관찰됨)에 의해 부풀려지거나, 다운스트림 중복(Pythia 자연어 텍스트에서 관찰됨)에 의해 가려질 수 있기 때문입니다. 저자들은 다중 지표 표준을 권장합니다:

1. 방향: 모든 지표는 손상 방향으로 움직여야 합니다.
2. 타겟-로짓: 인과적 영향의 주요 지표입니다.
3. 정확도: 보완적인 증거입니다.
4. 손실: 집계 슬랙(aggregation slack)의 영향을 받을 수 있는 보수적인 하한선입니다.

주장 및 의의
본 논문은 다음 네 가지 구체적인 주장을 결론으로 내립니다:

1. 제안 vs. 발견: 비지도 공활성화 클러스터링은 회로 제안 방법이지, 발견 방법이 아닙니다. 이는 후보 커뮤니티를 식별하지만, 오직 인과적 절제만이 그것을 부하 담당으로서 검증합니다.
2. 밀집 모델의 성공: 밀집 1B 규모 모델에서 합성 및 자연어 텍스트 모두에 대한 공활성화 클러스터링은 폐쇄를 견뎌내는 부하 담당 회로를 성공적으로 제안할 수 있습니다.
3. MoE의 실패: MoE 아키텍처에서는 (주의 깊은 널 테스트를 통과하는) 경로 조건부 통계적 회복조차도 회로의 증거로는 불충분합니다. 회복된 커뮤니티는 폐쇄 시 반대 방향(제거 시 성능 향상)으로 실패할 수 있습니다.
4. 지표의 괴리: 어텐션 패턴(선택성/PR)과 기능 사이의 간극은 양방향적입니다. 즉, 패턴은 날카롭지만 기능은 없을 수 있고, 기능은 있지만 패턴은 없을 수 있습니다.

범위 및 한계
저자들은 본 결과가 SAE-특징 매니폴드(다른 검증 기준을 사용함), 더 큰 모델, 다른 아키텍처(예: LLaMA, Mistral), 또는 다른 입력 분포(예: 코드)로 일반화된다고 주장하지 않음을 명시합니다. 또한 MoE에서의 "잘못된 방향" 실패는 테스트된 아키텍처와 라우팅 프로토콜에 특이적이라는 점을 언급하며, 사용된 Ising 친화도는 비교 가능성을 위한 선택일 뿐 반드시 최적의 친화도는 아님을 밝힙니다. 본 연구는 MoE에서의 부정적 결과가 필수적임을 강조하는데, 이는 통계적 일관성이 인과적 부하 담당성을 보장하지 않는다는 점을 보여주기 때문입니다.