[論文レビュー] A Mechanistic Interpretation of Arithmetic Reasoning in Language Models using Causal Mediation Analysis

Alessandro Stolfo, Yonatan Belinkov+1 more|arXiv (Cornell University)|May 24, 2023

Topic Modeling被引用数 2

ひとこと要約

この論文では、因果中介分析を用いて、Transformerベースの言語モデルが算術的推論をどのように処理するかを追跡し、注意メカニズムが中間層の初期段階で被演算子および演算子の情報を最終トークンに伝達していることが明らかになった。その後、後段の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.

研究の動機と目的

• 大規模言語モデルが算術的推論をどのように内部的に処理するかのメカニズムを理解すること。
• 因果的干渉を用いて、算術的予測を担う特定のモデルコンponentsを同定すること。
• 算術的推論で観察された活性化ダイナミクスが、他の数値的または事実的タスクと比較して特異的かどうかを評価すること。
• 算術推論中のモデルアーキテクチャ内での情報の流れを解明するメカニズム的知見を提供すること。
• 鍵となる計算回路を同定することで、将来的なモデルの解釈可能性、削減、推論時補正の研究を支援すること。

提案手法

• モデルコンponents（ニューロン、層、注意ヘッド、MLP）に対して因果中介分析を適用し、その活性化に干渉した。
• 特定のモデルパラメータのサブセットに対して制御された干渉を施し、出力確率分布の変化を測定した。
• 入力トークンから注意機構を経て最終トークンの表現に至るまでの情報の流れを追跡した。
• 干渉に基づく因果効果を通じて予測結果に顕著な影響を与えるメディエーター（モデルコンponents）を同定した。
• 算術（アラビア数字および語彙的表記を含む）、数値の検索、事実知識の4つのタスク間の活性化ダイナミクスを比較した。
• 介入効果で上位400位のニューロンを用いてタスク間のニューロンの重複度を計算し、ランダムなベースラインと統計的検証を実施した。

実験結果

リサーチクエスチョン

• RQ1 大規模言語モデルにおいて、正しい算術的予測を担う因果的要因となるモデルコンponentsは何か？
• RQ2 算術的推論中に、被演算子および演算子関連の情報がモデルの層および注意機構を通じてどのように流れているか？
• RQ3 算術的推論に用いられる回路は、数値の検索や事実知識タスクで用いられるものと異なるか？
• RQ4 結果関連の表現を生成するために、モデルが特定の後段MLPモジュールのサブセットに依存しているか？
• RQ5 算術的推論における活性化ダイナミクスは、他の数値的または事実的予測タスクと比較してどのように異なるか？

主な発見

• 被演算子および演算子に関する情報が、自己注意機構を介して中間層の初期段階から最終トークンに伝達されている。
• 後段のMLPモジュールが、残差ストリームに組み込まれる結果関連の情報を生成している。
• アラビア数字および語彙的表記を用いた算術クエリによって活性化される上位400個のニューロンには50％の重複があり、共通の回路が存在することが示された。
• 算術と数値の検索の間のニューロンの重複は22–23％にとどまり、両者とも数値的予測を含むが、回路が異なることが示唆された。
• 算術と事実知識タスクの間のニューロンの重複は9–10％であり、ランダムな重複と統計的に差がないため、回路の特異性が確認された。
• 観察された活性化ダイナミクスは算術的推論に特異的であり、他の数値的または事実的タスクに一般化されない。