Pruning via Causal Attribution Preserves Reasoning Performance in Large Language Models

기술 요약: 인과적 기여도 기반 프루닝 (Causal Attribution Pruning, CAP)

문제 정의
대규모 언어 모델(LLM)은 강력한 다단계 추론 능력을 보여주지만, 상당한 추론 비용이 발생하여 실제 배포를 저해한다. 모델 프루닝은 재학습 없이 중복된 파라미터를 제거함으로써 해결책을 제시하지만, 기존 방식들은 복잡한 추론 작업에서 성능을 저하시키는 경우가 많다. 현재의 접근 방식들은 가중치 크기(예: Wanda)나 국소적 그래디언트 통계와 같은 상관관계 신호에 의존한다. 이러한 방법들은 작은 가중치가 중요하지 않다고 가정하지만, 낮은 크기의 파라미터라도 다단계 추론 회로에서는 결정적인 역할을 할 수 있다. 또한, 순수하게 상관관계에 기반한 기준은 연쇄 사고(Chain-of-Thought, CoT) 추론에 필수적인 레이어 간 상호작용을 포착하지 못하며, 이로 인해 전체 모듈이 기능을 상실하는 "레이어 붕괴(layer collapse)" 현상이 빈번하게 발생한다.

방법론: 인과적 기여도 기반 프루닝 (CAP)
저자들은 추론 성능을 보존하도록 설계된 훈련 불필요(training-free) 방식의 원샷(one-shot) 프루닝 기법인 **Causal Attribution Pruning (CAP)**을 소개한다. 정적인 속성에 기반하여 가중치를 순위 매기는 방법들과 달리, CAP는 특정 모델 구성 요소가 작업 성능에 미치는 *인과적 영향(causal impact)*을 측정한다. 이 프레임워크는 세 단계로 작동한다:

1. 교정 및 인과적 점수 산출 (1단계):

   • 추론 중심 벤치마크인 GSM8K(수학), StrategyQA(다단계 상식), ARC-Challenge(과학)에서 샘플링하여 소규모 교정 데이터셋($D_{cal}$)을 구축한다.
   • 모델 내 각 어텐션 헤드 $h_\ell$에 대해, CAP는 개입적 마스킹(interventional masking) 실험을 수행한다. 해당 헤드를 순방향 패스(forward pass) 중에 일시적으로 마스킹하고, 온전한 모델 대비 토큰 수준의 크로스 엔트로피 손실 증가량($\Delta(h_\ell)$)을 측정한다.
   • 소규모 샘플 크기로 인한 변동성에 대한 강건성을 확보하기 위해, 데이터셋을 $K$개의 서브셋으로 분할하고 손실 증가량의 중앙값을 계산한다.
   • 이 과정은 기능적으로 중요한 헤드를 식별하며, 단순히 수치적으로는 크지만 중복된 헤드와 구분해낸다.

2. 가중치 수준 중요도 변환 (2단계):

   • 헤드 수준의 인과적 점수는 $[1.0, 10.0]$ 범위로 매핑된 정규화된 중요도 인자 $I(h_\ell)$로 변환된다. 영향력이 큰 헤드는 더 높은 인자를 받는다.
   • 이 인자들은 대응하는 투영 행렬($W_q, W_k, W_v, W_o$) 내의 개별 가중치로 전파된다. 특정 가중치 $W_{ij}$의 중요도 점수는 해당 헤드 영역의 크기와 중요도 인자의 곱으로 계산된다: $S_{ij} = |W_{ij}| \cdot I(h_{region})$.
   • 자연스러운 헤드 분할 구조가 없는 MLP 레이어의 경우, 해당 위치의 어텐션 헤드로부터 유도된 레이어별 평균 중요도를 할당한다. 이는 근사치로 명시되어 있다.

3. 글로벌 프루닝 (3단계):

   • 모델 전체의 모든 가중치를 중요도 가중 크기 점수($S_{ij}$)에 따라 순위를 매긴다.
   • 목표 희소성(예: 10%, 20%, 50%)을 달eric하기 위해 이진 탐색(binary search)을 통해 글로벌 임계값을 결정한다. 임계값보다 낮은 가중치는 0으로 처리되어, 표준 추론 인프라와 호환되는 비구조적 희소성 패턴을 생성한다.

주요 기여

• 프루닝을 위한 인과적 기여도의 공식화: 저자들은 개별 어텐션 헤드가 마스킹되었을 때 기대되는 작업 특화 손실 증가량을 측정하는 방법을 공식화하여, 상관관계 대리 지표를 넘어 기능적 결정성을 정량화했다.
• 예산 인식형 가중치 수준 프루닝: 헤드 수준의 인과적 점수를 가중치 수준의 중요도 값으로 변환하는 절차를 도입했다. 이를 통해 핵심적인 추론 회로를 보존하면서 글로벌 희소성 제약 조건 하에서 미세한 크기 기반 프루닝이 가능하다.
• 실증적 검증: 이 방법은 Llama-3-8B-Instruct 및 Mistral-7B-Instruct를 대상으로 GSM8K, StrategyQA, ARC-Challenge 벤치마크에서 평가되었다.

결과

• 중간 수준 희소성 (10–20%): CAP는 일반적으로 추론 성능을 보존하는 데 있어 Wanda 베이스라인보다 우수한 성능을 보인다.
  • Llama-3-8B-Instruct에서 20% 희소도일 때, CAP는 ARC-Challenge에서 Wanda 대비 61%의 상대적 개선을 달성했다 (정확도 70.8% vs 43.9%).
  • GSM8K에서 10% 희소도일 때, CAP는 Wanda 대비 20%의 상대적 개선을 보였다 (71.5% vs 59.5%).
  • 저자들은 CAP가 특정 작업(예: Llama-3 GSM8K 20%)에서 때때로 Wanda에 뒤처지기도 하지만, Wanda의 상관관계 기준이 실패하는 다단계 추론 작업에서 성능을 일관되게 보존한다고 언급했다.
• 높은 희소성 (50%): 두 방법 모두 성능이 크게 저하된다. CAP는 Llama-3에서 치명적인 실패를 경험하는데(퍼플렉시티가 428.2로 급증), 이는 MLP 레이어에 대한 거친(coarse-grained) 처리 때문으로 분석된다. 이 수준에서는 글로벌 임계값이 어텐션 헤드 근사를 통해 낮은 중요도를 부여받은 필수적인 MLP 가중치를 제거하게 된다.
• 퍼플렉시티 vs. 추론: 중간 수준의 희소도에서 CAP는 Wanda에 비해 약간 높은 퍼플렉시티를 기록함에도 불구하고 종종 더 나은 추론 정확도를 달성한다. 이는 일반적인 언어 모델링 품질과 작업 특화 추론 보존 사이의 디커플링(decoupling)을 시사한다.
• 아키텍처 한계: Mixture-of-Experts (MoE) 모델에 대한 예비 실험 결과, 동적 라우팅으로 인해 헤드 중요도가 비정상적(non-stationary)이 되므로 CAP의 효과가 떨어지는 것으로 나타났다.

의의 및 주장
본 논문은 어텐션 헤드 수준의 인과적 기여도가 동일한 희소도 수준에서 상관관계 기반 프루닝 방법보다 추론 성능을 보존하는 데 더 우수한 기준이라고 주장한다. 추론 작업에 대한 헤드의 인과적 기여를 명시적으로 측정함으로써, CAP는 크기 기반 방법들이 버리는 "수치적으로는 작지만 기능적으로는 결정적인" 파라미터를 식별하고 보호할 수 있다.

저자들은 자신들의 연구 범위가 다음과 같이 겸손함을 강조한다:

• 이는 훈련이 필요 없는(training-free) 방식으로, 가중치 업데이트를 요구하는 접근 방식(SparseGPT 등)에 비해 계산 효율적이다.
• 이는 일반적인 언어 모델링이 아닌 추론 데이터에 구체적으로 맞춰진 작업 정렬(task-aligned) 방식이다.
• MLP 뉴런에 대한 미세한 인과적 기여도 부족과 동적 라우팅 아키텍처(MoE)에 적용하기 어렵다는 점 등 한계를 인정하고 있다.

핵적인 기여는 개입적이고 작업에 정렬된 프루닝이, 일반적인 퍼플렉시티나 정적인 가중치 통계에 최적화된 방법보다 (중간 수준의 희소도 범위 내에서) 특정 능력(추론)을 더 잘 보존할 수 있음을 입증한 것이다.