LLMのキモい算術 - ｼﾞｮｲｼﾞｮｲｼﾞｮｲ

LLMは「226-68=」のようなプロンプトを与えると「158」と計算してくれますが、この計算は我々が想像するよりも奇妙な方法で行っていることを紹介します [Nikankin+ ICLR 2025]。

まずは前提条件を確認します。思考の連鎖は使わず、「226-68=」のようなプロンプトに対して「158」のように答えを直接出力する場合を考えます。

一例として Llama3-8B を考えます。Llama3 のトークナイザは 0 から 1000 までの数に 1 つのトークンを割り当てるので、「226-68=」を入力すると、次のトークン「158」が「0」「1」...「157」「158」「159」...「1000」などのトークンの中から、最も確率が高いものとして選ばれます。

ヤニフ・ニカンキンらの発見 [Nikankin+ ICLR 2025] は、Llama3-8B は答えや入力についての粗い条件を多数評価し、その積み重ねでこのような四則演算を解いているということです。

例えば、「{op1} - {op2}」というプロンプトテンプレート（{op1} と {op2} は具体的な数で埋めます）を入力したとき

第 24 層の 12439 番目のニューロンは、{op1} - {op2} の評価結果が 150 から 180 の間にあるときに発火する → 発火するとトークン「150」「151」「152」...「179」「180」の出力確率が増える
第 30 層の 1582 番目のニューロンは、{op1} - {op2} の評価結果が mod 10 で 8 のときに発火する → 発火するとトークン「8」「18」「28」...「998」の出力確率が増える

などが観察されています。

例えば「226-68=」を入力すると、評価結果 158 は 150 から 180 の間にあるので第 24 層の 12439 番目のニューロンが発火しトークン「150」「151」「152」...「179」「180」の出力確率が上昇、評価結果 158 は mod 10 で 8 なので第 30 層の 1582 番目のニューロンが発火しトークン「8」「18」「28」...「998」の出力確率が上昇します。

このとき、「150」や「998」などのトークンの確率も同時に増えますがそれらが積み上がる回数は僅かなのに対して、真の答え「158」は毎回積み上がるので、全てのニューロンの累積でトークン「158」が傑出して確率が高くなります。

ニューロンの発火パターンの可視化 [Nikankin+ ICLR 2025]。それぞれのニューロンは入力や評価結果が特定のパターンに該当するときのみ発火する。

それぞれのニューロンは四則演算を厳密に解いておらず、粗い条件を評価しているに過ぎませんが、粗い条件が無数に積み重なり、真の回答が浮き彫りになります。著者らはこのメカニズムをヒューリスティックの束 (bag of heuristics) と呼んでいます。

一般に、{op1}、{op2}、または評価結果が特定のパターンに当てはまるときのみ発火するニューロンをヒューリスティックニューロンと呼びます。ヒューリスティックニューロンの種類としては以下のものを考えています。

範囲ヒューリスティック：値が [a, b] の範囲に含まれる。
剰余ヒューリスティック：値 mod n = m が成り立つ。
パターンヒューリスティック：値が 1.2 のように特定の正規表現にマッチする。
オペランド一致ヒューリスティック：{op1} = {op2} が成り立つ。
複数結果ヒューリスティック（割り算のときのみ使用）：値が集合 S に含まれる。S は要素数 2 ~ 4 の集合。

特定のニューロンがヒューリスティックニューロンかどうかを調べるには、「{op1} - {op2}」の形の様々なプロンプトを入力し、どのような場合にそのニューロンが強く発火したかを記録し、パターンの候補との一致度を測ればよいです。

ヒューリスティックニューロンの発見 [Nikankin+ ICLR 2025]。さまざまなプロンプトを入力し、そのニューロンが強く発火するパターンと、候補のヒューリスティックパターンの一致度を計算する。

各ニューロンによる出力への寄与はロジットレンズ (logit lens) [nostalgebraist 2020] という方法で算出できます。

ロジットレンズの模式図。残差接続の途中や、各層の出力を最終線形層に通せば、その時点でのトークンロジットや、その層のトークンロジットへの寄与が分かる。

トランスフォーマーは注意機構と多層パーセプトロン (MLP) が残差接続で積み重なっています。つまり、トランスフォーマーの最終線形層への入力 $v \in \mathbb{R}^d$ は各層の注意機構と MLP の出力の総和 $v = \sum_l v^{(l)}$ であり、最終線形層はこれに語彙サイズ $|V|$ の行列 $W \in \mathbb{R}^{|V| \times d}$ を掛けて各トークンのロジットを計算します。通常の計算では、全ての層を足し合わせてからトークンの確率を計算しますが、見方を変えると、各層の出力 $v^{(l)} \in \mathbb{R}^d$ は、その都度、トークンのロジットを $W v^{(l)} \in \mathbb{R}^{|V|}$ だけ押し上げていると解釈できます。特に、2 層 MLP の中間ニューロン $i$ は、2 層目のパラメータ行列の $i$ 列目 $u^{(l)}_i \in \mathbb{R}^d$ に接続されており、このニューロンが発火するとトークンのロジットは $W u^{(l)}_i \in \mathbb{R}^{|V|}$ だけ押し上げられることが分かります。これにより、各ニューロンによる出力への寄与が具体的に分かります。

このフレームワークにより、LLM が計算を間違うことの原因も分析できます。Llama3-8B はたまに計算を間違います。正解したときのニューロンの発火パターンと、間違えたときのニューロンの発火パターンを分析してみると、間違えたときには、ヒューリスティックニューロンが押し上げた正解トークンのロジットが少ないことが分かりました。