기술 요약: 상호작용 인지 역전파를 통한 그래디언트 기반 회로 국소화 교정

문제 정의

회로 국소화(Circuit localization)는 거대 언어 모델(LLM)의 특정 동작을 담당하는 모델 구성 요소(예: 어텐션 헤드, 뉴런, 연결)의 구체적인 부분 집합을 식별하는 것을 목표로 하며, 이는 기계론적 해석 가능성(mechanistic interpretability)을 가능하게 합니다. 기존의 대부분의 방법은 각 구성 요소가 독립적으로 작용한다는 가정하에 작동하며, 각 구성 요소를 개별적으로 섭동(perturbation)시켜 중요도를 추정합니다.

그러나 신경망 구성 요소들은 서로 상호작용하며, 이러한 상호작용을 무시하면 구성 요소의 중요도를 체계적으로 잘못 추정하게 됩니다. 저자들은 **어텐션 자기 수복(attention self-repair)**이라 불리는 매우 문제가 되는 특정 상호작용을 식별했습니다. 이 현상은 여러 위치가 큰 어텐션 가중치를 가지고 유사한 정보(값 벡터)를 포함하고 있을 때 발생하며, 단일 어텐션 점수를 제거(ablate)하더라도 출력에 미치는 영향이 미미하게 나타납니다. 이는 소프트맥스(softmax) 활성화 함수가 남은 중복 위치로 가중치를 재분배하여 보상하기 때문에 발생합니다. 결과적으로, 표준 섭동 기반 및 그래디언트 기반 설명 방법들은 이러한 어텐션 점수에 기여하는 구성 요소의 중요도를 과소평가하며, 종종 영향력 있는 점수에 대해 그래디언트 소실(vanishing gradients) 현상을 일으킵니다.

방법론: 그래디언트 상호작용 수정 (GIM)

본 논문은 상호작용하는 여러 구성 요소를 동시에 섭동할 때의 영향을 추정하도록 설계된 그래디언트 기반 귀속(attribution) 방법인 **그래디언트 상호작용 수정(Gradient Interaction Modifications, GIM)**을 소개합니다. GIM은 세 가지 상호 보완적인 기술을 통해 표준 역전파를 수정합니다.

1. 온도 조절 소프트맥스 그래디언트 (Temperature-adjusted Softmax Gradients, TSG):

   • 동기: 어텐션 자기 수복 문제를 해결하기 위함입니다. 표준 그래디언트는 값이 유사할 때 하나의 소프트맥스 가중치를 높이는 데서 얻는 이득이 다른 가중치를 억제함으로써 발생하는 손실에 의해 상쇄되어 그래디언트가 소실됩니다.
   • 메커니즘: 역전파 과정에서 TSG는 그래디언트를 계산하기 전, 높은 온도($T > 1$)로 소프트맥스를 재계산합니다. 이는 확률 분포를 평탄하게 만들어 큰 가중치는 줄이고 작은 가중치는 키웁니다. "유의미한" 위치의 집합을 확장함으로써, 이 위치들 간의 기여도가 동일하다는 가정을 완화합니다. 이를 통해 쌍별 그래디언트 차이($g_j - g_k$)가 소멸하지 않도록 보장하여, 그래디언트 신호가 생존하고 동시 섭동의 효과를 근사할 수 있게 합니다.

2. 레이어 정규화 동결 (Layernorm Freeze):

   • 동기: 레이어 정규화는 한 구성 요소를 제거할 때 활성화 크기를 감소시키고, 정규화 계수가 남은 구성 요소를 증폭시키는 제로섬(zero-sum) 역학을 도입합니다.
   • 메커니즘: 정규화 계수($\sigma$)를 역전파 동안 상수로 취급합니다. 이는 정규화 재조정(rescaling)이 일어나기 전 각 구성 요소의 직접적인 기여도를 측정하며, 보상적 증폭으로 인한 중요도 과소평가를 방지합니다.

3. 그래디언트 노름 (Gradient Norm):

   • 동기: 곱셈 연산(예: $Q$와 $K$가 모두 입력 $X$에 의존하는 $A = QK^T$)은 표준 1차 그래디언트에서 중복 계산(double-counting)을 유발합니다. $X$를 섭동하면 $Q$와 $K$ 모두에 동시에 영향을 미치지만, 표준 연쇄 법칙(chain rule)은 이들의 개별 효과를 합산하여 변화량을 과대평가합니다.
   • 메커니즘: 곱셈적 상호작용의 그래디언트를 곱셈 입력의 수로 나눕니다(예: $Q \times K$의 경우 2로 나눔). 이는 "균등 규칙(uniform rule)"과 일치하며, 어텐션-값 곱셈, 쿼리-키 곱셈, MLP 게이트-투영 곱셈에서의 기여도 과다 계산을 교정합니다.

주요 기여

1. GIM 프레임워크: 특징 상호작용을 명시적으로 고려하여 기계론적 해석 가능성 벤치마크(MIB)의 회로 국소화 트랙에서 최첨단 성능을 달지는 그래디언트 기반 귀속 방법을 도입했습니다.
2. 어텐션 자기 수복의 식별: 어텐션 자기 수복(OR 게이트와 유사한 상호작용)이 영향력 있는 어텐션 점수에 대해 그래디언트를 소실시켜 구성 요소의 중요도를 체계적으로 과소평가한다는 사실을 공식적으로 식별하고 증명했습니다.
3. 특징 귀속 우위성: GIM이 다양한 작업과 LLM 아키텍처에 걸쳐 기존 그래디언트 기반 방법들보다 뛰어난 특징 귀속 성능을 보임을 입증했습니다.

결과

• 회로 국소화: MIB에서 GIM은 여러 모델(GPT-2, Qwen-2.5, Gemma-2, Llama-3.1)과 작업(간접 목적어 식별, 산술, MCQA, ARC)에 대해 기존 방법들(EAP, EAP-IG, Hybrid-ens, IFR 포함)을 유의미하게 능가했습니다. 특히 다단계 추론이 필요한 작업에서 가장 높은 평균 회로 성능 비율(CPR)을 달성했습니다.
• 특징 귀속: 7개의 지시어 튜닝된 LLM과 6개의 데이터셋(QA, 사실 검증, 감성 분석, 혐오 표현 탐지 포함)에 대해, GIM은 일반적으로 GradientXInput, Integrated Gradients, DeepLIFT, TransformerLRP보다 높은 포괄성(comprehensiveness) 및 충분성(sufficiency) 점수를 기록했습니다.
• 자기 수복의 검증: 경험적 분석을 통해 어텐션 자기 수복이 실제 LLM에서 널리 퍼져 있으며, 입력당 평균 65~1200회 발생함을 확인했습니다. 두 개의 큰 어텐션 점수를 동시에 제거했을 때의 효과는 개별 제거 효과의 합보다 유의미하게 컸으며, 이는 이론적 분석을 뒷받침합니다.
• 절제 연구 (Ablation Study): GIM에서 단일 수정 사항(TSG, Layernorm Freeze 또는 Gradient Norm)을 제거하면 대부분의 모델-데이터셋 쌍에서 성능이 저하되었으며, 이는 각 구성 요소가 방법의 충실도(faithfulness)에 의미 있게 기여함을 확인해 줍니다.

의의

본 논문은 GIM이 특징 상호작용으로 인한 구성 요소 중요도의 체계적 과소평정을 교정함으로써 LLM에 대한 더 충실한 기계론적 분석을 가능하게 한다고 주장합니다. 구성 요소가 어떻게 상호작용하는지(특히 소프트맥스 재분배 및 정규화 재조정을 통해)를 명시적으로 모델링함으로써, GIM은 모델 동작을 주도하는 내부 메커니즘에 대한 더 정확한 그림을 제공합니다. 저자들은 GIM을 단순한 점진적 개선이 아니라, 현재의 해석 도구들에 만연한 "독립성 가정"에 대한 필수적인 교정으로 규정하며, 이를 통해 이전에는 그래디언트 소실로 인해 가려져 있던 회로를 식별할 수 있게 합니다. 이 방법은 광범위한 채택을 촉진하기 위해 오픈 소스 파이썬 패키지로 공개되었습니다.