[논문 리뷰] A Mechanistic Interpretation of Arithmetic Reasoning in Language Models using Causal Mediation Analysis

Alessandro Stolfo, Yonatan Belinkov외 1명|arXiv (Cornell University)|2023. 05. 24.

Topic Modeling인용 수 2

한 줄 요약

이 논문은 인과적 중재 분석을 사용하여 트랜스포머 기반 언어 모델이 산술 추론을 어떻게 처리하는지 추적한다. 분석 결과, 어텐션 메커니즘이 중간 시퀀스의 초기 레이어에서 최종 토큰으로 연산자와 피연산자 정보를 전달하며, 후반부 MLP 모듈이 결과와 관련된 표현을 생성하는 것으로 드러났다. 연구에서는 수의어 추론이나 사실 지식 작업에서 사용되는 회로와는 다름 분명한, 과제에 특화된 산술을 위한 회로를 규명하였다.

ABSTRACT

Mathematical reasoning in large language models (LMs) has garnered significant attention in recent work, but there is a limited understanding of how these models process and store information related to arithmetic tasks within their architecture. In order to improve our understanding of this aspect of language models, we present a mechanistic interpretation of Transformer-based LMs on arithmetic questions using a causal mediation analysis framework. By intervening on the activations of specific model components and measuring the resulting changes in predicted probabilities, we identify the subset of parameters responsible for specific predictions. This provides insights into how information related to arithmetic is processed by LMs. Our experimental results indicate that LMs process the input by transmitting the information relevant to the query from mid-sequence early layers to the final token using the attention mechanism. Then, this information is processed by a set of MLP modules, which generate result-related information that is incorporated into the residual stream. To assess the specificity of the observed activation dynamics, we compare the effects of different model components on arithmetic queries with other tasks, including number retrieval from prompts and factual knowledge questions.

연구 동기 및 목표

• 대규모 언어 모델이 산술 추론을 수행하는 데 내재된 메커니즘을 이해하기 위해.
• 인과적 간섭을 통해 산술 예측을 담당하는 구체적인 모델 구성 요소를 규명하기 위해.
• 산술 추론에서 관찰된 활성화 동역학이 다른 수치적 또는 사실 기반 과제와 비교해 고유한지 평가하기 위해.
• 산술 추론 과정에서 정보가 모델 아키텍처를 통해 어떻게 흐르는지 기계적 통찰을 제공하기 위해.
• 핵심 계산 회로를 규명함으로써 향후 모델 해석 가능성, 프루닝, 추론 시 보정 작업을 지원하기 위해.

제안 방법

• 모델 구성 요소(뉴런, 레이어, 어텐션 헤드, MLP)에 대해 인과적 중재 분석을 적용하여 활성화에 간섭하였다.
• 특정 모델 파rameter 부분에 제어된 간섭을 가하고 출력 확률 분포의 변화를 측정하였다.
• 입력 토큰을 통해 어텐션 메커니즘을 거쳐 최종 토큰 표현으로의 정보 흐름을 추적하였다.
• 간섭 기반 인과 효과를 통해 예측 결과에 크게 영향을 주는 중재자(모델 구성 요소)를 규명하였다.
• 산술(아랍 숫자 및 수어 형태 포함), 수의어 추론, 사실 지식의 네 가지 과제의 활성화 동역학을 비교하였다.
• 간섭 효과 기반 상위 400개 뉴런을 사용해 과제 간 뉴런 겹침을 계산하였으며, 무작위 기반 대조군과 통계적 검증을 실시하였다.

실험 결과

연구 질문

• RQ1 대규모 언어 모델에서 올바른 산술 예측을 담당하는 데 인과적으로 기여하는 모델 구성 요소는 무엇인가?
• RQ2 산술 추론 과정에서 피연산자 및 연산자 정보가 모델의 레이어와 어텐션 메커니즘을 통해 어떻게 흐르는가?
• RQ3 산술 추론에 사용되는 회로는 수의어 추론이나 사실 지식 과제에서 사용되는 회로와 다름가?
• RQ4 결과와 관련된 표현을 생성하기 위해 모델이 특정 후반부 MLP 모듈의 하위 집합에 의존하는가?
• RQ5 산술 추론에서의 활성화 동역학은 다른 수치적 또는 사실 예측 과제와 비교해 어떻게 다른가?

주요 결과

• 피연산자 및 연산자 정보는 자기 어텐션 메커니즘을 통해 중간 시퀀스의 초기 레이어에서 최종 토큰으로 전달된다.
• 후반부 레이어의 MLP 모듈이 잔여 스트림에 통합되는 결과와 관련된 정보를 생성하는 데 기여한다.
• 아랍 숫자 및 수어 형태로 산술 질의에 대해 활성화된 상위 400개 뉴런 간 겹침 비율은 50%로, 공통 회로가 있음을 시사한다.
• 산술과 수의어 추론에 대해 활성화된 뉴런 간 겹침은 22~23%에 불과하여, 둘 다 수치 예측을 다루지만 서로 다른 회로를 사용함을 시사한다.
• 산술과 사실 지식 과제 간 뉴런 겹침은 9~10%이며, 무작위 기반 대조군과 통계적으로 구별되지 않아, 회로의 특이성을 확인한다.
• 관찰된 활성화 동역학은 산술 추론에 특화되어 있으며, 다른 수치적 또는 사실 기반 과제로 일반화되지 않는다.