Dual Path Attribution: Efficient Attribution for SwiGLU-Transformers through Layer-Wise Target Propagation

1. 연구 배경 및 문제 정의 (Problem)

• 배경: 트랜스포머 기반 대규모 언어 모델 (LLM) 의 내부 메커니즘을 이해하는 것은 모델의 신뢰성 있는 배포와 운영에 필수적입니다. 특히, 모델의 예측에 어떤 입력 토큰이나 내부 구성 요소 (어텐션 헤드, 뉴런 등) 가 기여했는지를 규명하는 속성 분석 (Attribution) 작업이 중요합니다.
• 문제점:
  • 기존 속성 분석 방법들은 **신뢰성 (Faithfulness)**과 계산 효율성 (Efficiency) 사이의 균형을 맞추기 어렵습니다.
  • 기울기 기반 (Gradient-based) 방법들은 기울기 포화 (saturation) 나 파편화 (shattering) 문제로 인해 부정확할 수 있습니다.
  • 인과적 개입 (Causal Intervention, 예: Activation Patching) 방법은 매우 신뢰성이 높지만, 밀집된 (dense) 구성 요소에 대한 속성 분석을 수행하려면 각 구성 요소마다 별도의 순전파 (forward pass) 가 필요하여 계산 비용이 prohibitively expensive(허용 불가능할 정도로 비쌈) 합니다.
  • 기존 분해 기반 (Decomposition-based) 방법들도 긴 입력 시퀀스나 밀집된 구성 요소 분석 시 확장성에 한계가 있습니다.

2. 제안 방법: 듀얼 패스 어트리뷰션 (Dual Path Attribution, DPA)

이 논문은 DPA라는 새로운 프레임워크를 제안하여, 한 번의 순전파와 한 번의 역전파만으로 SwiGLU 트랜스포머의 정보 흐름을 정확하게 추적하고 밀집된 속성 분석을 가능하게 합니다.

핵심 아이디어

• 타겟 중심 역전파 (Target-Centric Backward Propagation): 각 구성 요소를 순전파하는 대신, 목표 토큰의 **언디벡딩 벡터 (Unembedding Vector)**를 모델의 최상위 레이어에서 시작하여 하위 레이어로 역방향으로 전파합니다.
• SwiGLU 구조의 분해: SwiGLU 트랜스포머의 이차원적 (bilinear) 상호작용 구조를 분석적으로 분해하여 두 가지 명확한 경로로 나눕니다.
  1. 콘텐츠 경로 (Content Pathway): 값 (Value) 과 업 (Up) 프로젝션을 통해 표현을 전달하는 경로.
  2. 제어 경로 (Control Pathway): 어텐션과 게이트 점수를 통해 라우팅 및 변환을 제어하는 경로.

구체적 메커니즘

1. 선형화 및 역변환 (Linearization & Inversion):
   • GLU (Gated Linear Unit) 모듈: 게이트 활성화와 콘텐츠 경로를 분리하여 선형화합니다. 게이트 경로는 비선형성을 처리하기 위해 국소 기울기 (local gradient) 를 근사화합니다.
   • MHSA (Multi-Head Self-Attention) 모듈: 쿼리, 키, 값 경로를 분해합니다. 특히 소프트맥스 (Softmax) 의 비선형성을 안정적으로 처리하기 위해 1 차 테일러 전개 (Taylor-based decomposition) 를 사용하여 역변환을 근사합니다.
2. 계층별 타겟 전파:
   • 각 레이어에서 역방향으로 전파된 타겟 벡터는 해당 레이어의 GLU 와 MHSA 모듈을 거치며 선형화된 역변환 함수를 적용받아 업데이트됩니다.
   • 최종적으로 입력 레이어 (Layer 0) 에 도달한 유효 타겟 (Effective Target) 과 입력 임베딩의 내적을 통해 각 토큰의 속성 점수를 계산합니다.
3. 구성 요소별 점수 계산:
   • 각 어텐션 헤드나 MLP 뉴런의 기여도는 해당 레이어의 활성화 값과 역전파된 타겟 벡터의 내적으로 계산됩니다.

3. 주요 기여 (Key Contributions)

1. 효율적인 프레임워크 도입: 카운터팩추얼 (Counterfactual) 예제 없이 한 번의 순전파와 한 번의 역전파만으로 로짓 점수를 속성 분석하는 새로운 분해 기반 프레임워크인 DPA 를 제안했습니다.
2. 이론적 분해 증명: SwiGLU 트랜스포머의 이차원 구조를 콘텐츠 경로와 제어 경로로 명확히 분해 가능함을 보였습니다.
3. O(1) 시간 복잡도 달성: 모델 구성 요소의 수에 대한 시간 복잡도가 **O(1)**로, 구성 요소의 수와 무관하게 일정한 계산 비용으로 긴 시퀀스와 밀집된 구성 요소 분석이 가능합니다. 이는 기존 방법들의 O(M) 복잡도 (M: 구성 요소 수) 와 대조적입니다.
4. 최고 수준의 성능: 다양한 벤치마크에서 기존 기법들 (Gradient, Attention, DePass 등) 보다 뛰어난 **신뢰성 (Faithfulness)**과 계산 효율성을 입증했습니다.

4. 실험 결과 (Results)

• 데이터셋 및 모델: Llama-3.1-8B, Llama-2-7B, Qwen3, Mistral-7B 등 다양한 모델과 Known 1000 (사실적 지식), SQuAD v2.0 (독해), IMDb (감성 분석), IOI (간접 객체 식별) 등의 벤치마크에서 평가되었습니다.
• 신뢰성 (Faithfulness):
  • 토큰 레벨: 중요한 토큰을 제거했을 때 모델 성능이 얼마나 떨어지는지 (Disruption) 와 중요한 토큰만 남겼을 때 성능이 회복되는지 (Recovery) 를 측정했습니다. DPA 는 모든 데이터셋에서 가장 높은 총점 (Total AUC) 을 기록했습니다.
  • 구성 요소 레벨: 개별 어텐션 헤드와 MLP 뉴런의 중요도를 평가한 결과, DPA 는 다른 구성 요소들을 마스킹했을 때 모델이 올바른 예측을 유지하는 데 가장 효과적이었으며, 특히 IOI 작업에서 회수율 (Recovery) 이 매우 높았습니다.
• 효율성 (Efficiency):
  • 실행 시간: 긴 컨텍스트 (SQuAD v2.0) 에서 DPA 는 Activation Patching (AP) 및 DePass 에 비해 약 40 배 이상 빠른 속도를 보였습니다 (예: Llama-3.1-8B 기준 129 초 vs 5370 초).
  • 메모리: 순전파 시 활성화 값을 캐싱해야 하지만, 추가적인 순전파가 필요 없기 때문에 메모리 오버헤드는 기존 방법들과 유사하거나 약간 더 높을 뿐이며, 전체적인 계산 효율성은 압도적입니다.

5. 의의 및 한계 (Significance & Limitations)

• 의의:
  • 실용성: LLM 의 내부 메커니즘을 해석할 때 발생하는 계산적 병목 현상을 해결하여, 대규모 모델에 대한 밀집된 (dense) 기계적 해석 (Mechanistic Interpretability) 을 실용적으로 가능하게 합니다.
  • 정확성: 단순한 선형 근사를 넘어, 비선형 상호작용을 콘텐츠와 제어 경로로 분리하여 더 정교한 인과적 메커니즘을 포착합니다.
• 한계:
  • 아키텍처 특이성: 현재는 밀집형 SwiGLU 트랜스포머에 최적화되어 있으며, MoE (Mixture-of-Experts) 나 State Space Models (Mamba 등) 와 같은 이산적 라우팅이 있는 아키텍처에는 수학적 역함수를 별도로 유도해야 합니다.
  • 근사 오차: 2 차 미분 항을 무시하는 선형화 과정으로 인해, 모델의 행동이 매우 높은 비선형성이나 강한 상호작용에 의해 주도되는 영역에서는 정확도가 떨어질 수 있습니다.
  • 로그트 (Logit) 수준: 최종 소프트맥스 확률 분포가 아닌 로짓 (Logit) 수준에서 속성 분석을 수행하므로, 경쟁 토큰의 억제를 통한 간접적 기여는 포착하지 못할 수 있습니다.

결론

이 논문은 DPA 를 통해 신뢰성과 계산 효율성이라는 상충되는 두 가지 목표를 동시에 달성하는 새로운 해석 기법을 제시했습니다. 특히 SwiGLU 구조의 특성을 활용한 이중 경로 (Dual Path) 분석은 LLM 의 복잡한 내부 동작을 이해하는 데 있어 강력한 도구가 될 것으로 기대됩니다.