기술 요약: 인과적 회로 발견을 위한 다중 입도 노드 프루닝 (Multi-Granular Node Pruning for Causal Circuit Discovery)

1. 문제 정의

인과적 회로 발견(Causal circuit discovery)은 대규모 언어 모델(LLM) 내에서 특정 행동을 인과적으로 유도하는 최소한의 서브네트워크를 식별하는 것을 목표로 합니다. 기존 방식들은 기계론적 해석 가능성(mechanistic interpretability) 분야를 발전시켜 왔으나, 다음과 같은 두 가지 주요 한계에 직면해 있습니다:

1. 계산 확장성(Computational Scalability): 기존 접근 방식은 주로 **에지 프루닝(edge pruning, 연결 제거)**에 의존합니다. 트랜스포머 내 에지의 수는 노드 수에 따라 이차적으로 증가($O(|N|^2)$)하기 때문에, 모델 크기가 커질수록 에지 수준의 프루닝은 계산적으로 불가능해집니다. 또한, 이러한 방법들은 분리된 잔차 스트림(residual-stream) 표현과 중간 활성화 값을 저장해야 하므로 메모리 사용량이 매우 높습니다.
2. 입도(Granularity)의 한계: 기존 방법들은 전체 어텐션 헤드나 MLP 블록을 하나의 원자적 단위로 취급하는 거친 입도(coarse granularity)에서 작동하는 경우가 많습니다. 이는 특정 **뉴런(노드)**이 해당 행동을 담당할 가능성을 간과하게 만들며, 결과적으로 회로가 미세한 구조를 놓치고 불필요한 구성 요소를 유지하게 만듭니다.

2. 방법론

저자들은 단일 통합 최적화 단계 내에서 여러 입도 수준에 걸쳐 최소한의 회로를 발견하는 다중 입도 노드 프루닝 프레임워크를 제안합니다.

2.1 계층적 입도 설계

본 프레임워크는 에지(연결)를 프루닝하는 대신, 다음 네 가지 서로 다른 수준의 노드에 대한 마스크를 학습합니다:

• 블록 수준 (Block Level): 전체 어텐션 블록 및 MLP 블록.
• 헤드 수준 (Head Level): 개별 어텐션 헤드.
• 뉴런 수준 (Neuron Level): 어텐션 레이어 및 MLP 블록의 두 선형 레이어(MLP1/Keys 및 MLP2/Values) 내의 개별 뉴런.

2.2 이중 스트림 인과적 프루닝 설정 (Two-Stream Causal Pruning Setup)

단순한 제로-절제(zero-ablation)로 인한 분포 변화(distribution shift) 위험을 피하면서 인과적 중요성을 식별하기 위해, 이 방법은 이중 스트림 순전파를 채택합니다:

• 클린 스트림 (Clean Stream): 특정 행동을 유도하도록 설계된 입력(예: 연도가 50보다 큰 경우를 예측)을 처리합니다.
• 오염 스트림 (Corrupted Stream): 다른 예측을 유도하기 위해 입력의 최소한의 섭동 버전(예: 성별 대명사를 반전시키거나 시작 연도를 변경)을 처리합니다.
• 보간 (Interpolation): 각 노드 $n_i$에 대해, 학습 가능한 마스크 $m_i \in [0, 1]$가 클린 활성화($h^{clean}$)와 오염된 활성화($h^{corrupted}$) 사이를 보간합니다:
  $$h^{node}_i = m_i \cdot h^{clean}_i + (1 - m_i) \cdot h^{corrupted}_i$$
  최적화 과정은 모델이 올바른 출력(클린 행동과 일치하는 출력)을 생성하는 동시에 희소성(sparsity)을 극대화하는 마스크를 찾는 것을 목표로 합니다.

2.3 최적화 및 마스크 파라미터화

• Hard Concrete 분포: 마스크는 경사 기반 최적화를 가능하게 하면서도 근사적으로 이진 동작(0 또는 1)을 유도하도록 Hard Concrete 분포를 사용하여 파라미터화됩니다.
• 통합 목적 함수 (Unified Objective): 학습 목적 함수는 세 가지 구성 요소의 결합입니다:
  1. 충실도 손실 ($L_{faith}$): 프루닝된 회로의 출력 분포가 베이스 모델과 일치하도록 보장하는 KL 발산(KL divergence).
  2. 태스크 손실 ($L_{task}$): 회로가 정답과 오답을 올바르게 구별하도록 보장하는 마진 기반 랭킹 손실.
  3. 희소성 정규화 항 ($L_0$): 마스크가 0이 되도록 유도하여 회로 크기를 최소화합니다.
• 계층적 일관성 (Hierarchical Consistency): 훈련 후, 부모 노드(예: MLP 블록)가 프루닝되면 그 자식 노드(예: 해당 블록 내의 뉴런)도 모두 비활성화되도록 보장하는 제약 조건을 적용하여 마스크를 이진화합니다.

3. 주요 기여

1. 다중 입도 프레임워크: 트랜스포머 블록, 어텐션 헤드, MLP 블록 및 개별 뉴런에 걸쳐 동시에 마스크를 학습하는 통합 최적화 프레임워크를 도입했습니다.
2. 확장성 및 효율성: 에지 프루닝의 이차적 스케일링 문제와 모든 에지에 대한 중간 활성화 값을 저장해야 하는 문제를 해결했습니다. 이 방법은 기존 방식 대비 GPU 메모리 사용량을 3배에서 11배까지 줄여, 단일 40GB GPU만으로도 Llama 3.1-8B에서 회로 발견을 가능하게 합니다.
3. 미세한 입도의 발견: 본 접근 방식은 베이스라인보다 훨씬 더 조밀한 회로를 식별하며, 특히 이전 방법들이 달성하기 어려웠던 MLP 블록 및 뉴런 수준에서 높은 희소성을 달로 성취합니다.

4. 실험 결과

본 방법은 GPT-2, GPT-2-XL, Llama 모델을 사용하여 세 가지 표준 회로 발견 작업(간접 목적 식별(IOI), 성별 대명사(GP), Greater Than(GT))에 대해 평가되었습니다.

4.1 희소성 및 조밀도

• MLP 및 뉴런 프루닝: 제안된 방법은 에지 프루닝 베이스라인(EAP, EP) 및 비구조적 가중치 프루닝(DiscoGP)보다 성능이 월등히 뛰어납니다. 가장 불리한 설정에서도 기존 방식보다 MLP 블록을 33.34% 더 많이, 뉴런을 59.8% 더 많이 프루닝했습니다.
• 태스크 의존적 구조: 발견된 회로는 작업에 따라 뚜렷한 토폴로지를 보입니다:
  • IOI: 분산된 MLP 계산이 지배적이며, 어텐션은 마지막 레이어들을 제외하고 대부분 프루닝됩니다.
  • 성별 대명사 (GP): 고도로 국소적이고 불연속적입니다; MLP 계산은 몇 개의 레이어로 제한되는 반면, 어텐션은 희소합니다.
  • Greater Than (GT): 비활성 레이어의 연속된 블록과 "스킵(skip) 형태"의 실행 패턴을 보이는 극단적인 희소성을 나타냅니다.

4.2 충실도 및 일관성

• 성능: 프루닝된 회로는 높은 태스크 정확도(종종 >95%)와 베이스 모델 대비 낮은 KL 발산을 유지합니다.
• 일관성: 다섯 가지 무작위 시드에 대한 실험 결과, 발견된 회로(충실도 지표 및 희소성 프로파일)는 매우 안정적이었습니다.
• 정성적 분석: Llama 3.2-1B에서 IOI 태스크를 수행하는 생존 어텐션 헤드들은 GPT-2에서 수동으로 발견된 회로와 유사한 특정 기능적 역할(IO-Movers, S-Inhibitors)을 수행하는 것으로 확인되어, 본 방법의 기계론적 통찰력 복구 능력을 검증했습니다.

4.3 계산 효율성

• 메모리: 이 방법은 에지 프루닝 베이스라인(예: GPT-2에 대해 EAP의 경우 72.79 GB vs 6.27 GB)보다 최소 3배 적은 메모리를 요구합니다.
• 시간: 훈련은 6분 이내에 완료되어, 에지 프로닝 방식(약 50분 소요, 대신 엄청난 메모리 사용)보다는 약간 느리지만, 단일 패스 방식인 EAP보다는 훨씬 빠릅니다.

5. 의의 및 주장

본 논문의 주요 의의는 현재의 회로 발견이 가진 확장성 및 입도 병목 현상을 극복하는 데 있습니다. 에지 프루닝에서 다중 입도 노드 프로닝으로 전환함으로써 저자들은 다음을 입증했습니다:

• 기존의 거친 입도 방식에서 유지되는 많은 뉴런은 불필요하며, 성능 저하 없이 제거될 수 있습니다.
• 더 큰 모델(예: Llama 8B)은 더 작은 모델과 대등하거나 오히려 더 강력한 프루닝을 지원하며, 이는 모델 규모가 커질수록 작업 특화 회로가 더 밀집해지는 것이 아니라 오히려 중복성이 증가함을 시사합니다.
• 본 방법은 대규모 모델에서 미세한 인과적 회로를 발견할 수 있는 계산적으로 실행 가능한 경로를 제공하며, 서로 다른 행동이 어텐션, MLP, 뉴런 활동의 독특한 패턴에 의해 인코딩된다는 점을 밝혀냅니다.

저자들은 본 프레임워크가 어떤 노드가 필요한지는 식별하지만, 생존한 노드들 사이의 *상호작용 구조(에지)*를 명시적으로 복구하지는 않는다는 한계를 인정했습니다. 이들은 향후 연구 방향으로 노드 선택과 에지 복구를 결un합하여 상호작용 구조를 추론하는 것을 제시했습니다.