『The Right Answer, the Wrong Direction』を読む

LLMは、文章を書き、コードを書き、難しい試験問題にも答えます。 ところが、次のような一見簡単な問題で失敗します。

文章中に猫が何匹いるか、数字だけで答えてください。

人間なら、対象語を見つけて数えれば終わりです。 外部知識も、高度な推論も、複雑な計算もいりません。

それなのにTransformer系モデルは、このようなカウント問題でしばしば間違えます。

では、なぜ間違えるのでしょうか。

この論文では、Transformerは、数をまったく数えられていないわけではなく、むしろ、内部には正しいカウント情報を持っている。にもかかわらず、その情報を「1」「2」「3」のような数字トークンとして出力するベクトルがズレていると説明しています。

つまり問題は、「数えられない」ことではなく、「数えた結果をうまく出せない」ことにあると書いています。

論文ではこの現象を、geometric readout bottleneck、つまり幾何学的な読み出しボトルネックとして説明しています。モデル内部に答えはあるのに、その答えが出力ヘッドの数字トークンベクトルと噛み合っていない、という考え方です。


参考にしたもの


Contents

  • 概要
  • なぜカウント失敗は不思議なのか
  • まず、モデル内部には数があるのか
  • では、なぜ出力では間違えるのか
  • 「正しい答え、しかし間違ったベクトル」
  • なぜ数字トークンのベクトルはズレるのか
  • 出力ヘッドだけ直す:9-row repair
  • しかし、9-row repairだけでは通常生成は直らない
  • Diagnostic Probe Steering:内部の答えを外から押し出す
  • 実用的に直すには、LoRA Q/Vが効く
  • 結果を整理する
  • Chain-of-Thoughtでは十分ではないのか
  • この論文が直しているもの、直していないもの
  • まとめ

概要

この論文は、Transformerがカウント問題に失敗する理由を、内部表現と出力ヘッドの関係から分析します。

論文の主張を一言でまとめると、次のようになります。

Transformerは、入力中の対象数を内部の残差ストリームにかなり正確に表現している。しかし、そのカウント情報を表すベクトルが、数字トークンを出力するlm_headの行とほぼ直交しているため、正しい数字を出力できない。

この主張を確かめるために、論文では次のような道具を使います。

  • 線形プローブ
  • logit-lens
  • 出力ヘッドの数字9行だけを修正する9-row repair
  • Diagnostic Probe Steering、DPS
  • attentionのQ/VにLoRAを入れる修正

まず「内部に答えがあるか」を調べる。 次に「モデル自身の出力ヘッドで読めるか」を調べる。 そのうえで「出力ヘッドだけ直せばよいか」を試す。 最後に「通常生成でも直るには、どこを修正すべきか」を確認する。

つまり、モデルの失敗箇所を段階的に切り分けています


1. なぜカウント失敗は不思議なのか

LLMの失敗には、いろいろな種類があります。

たとえば、知らない知識を聞かれて間違える。 複雑な数学問題で途中式を間違える。 長い文脈の途中にある情報を見落とす。

こういう失敗は、まだ理解しやすいです。 知識不足、推論能力の不足、文脈処理の限界などで説明できます。

しかし、カウント問題は少し違います。

「文章中に猫が何匹いるか」という問題では、答えは入力文の中に直接あります。 必要なのは、対象語を見つけて数えることだけです。

このカウント失敗の原因は、Transformerが本当に数を持てていないからなのか。それとも、数は持っているが、出力できていないからなのか

もし内部に数の情報がないなら、モデルは本当にカウントを理解していません。 その場合は、アーキテクチャや訓練方法を根本的に変える必要があります。

しかし内部には数の情報があるなら、問題は別です。 その場合、必要なのは「数をどう出力トークンへ読み出すか」です。

この論文は後者を支持しています。


2. まず、モデル内部には数があるのか

最初に論文が確認するのは、モデル内部にカウント情報があるかどうかです。

ここで使うのが、線形プローブです。

線形プローブは、モデル内部の隠れ状態から、目的の情報を単純な線形変換で読めるかを調べる方法です。

式で書くと、次のようになります。

c_hat = w^T h + b

ここで、

  • h:モデル内部の隠れ状態
  • w:カウントを読むためのベクトル
  • b:バイアス
  • c_hat:予測されたカウント

です。

このプローブが高精度なら、モデル内部のどこかに「数」に対応する情報があると考えられます。

論文では、線形プローブが中間層から正しいカウントを高精度に復元できることを示しています。特にQwen3-8Bでは、層2以降でR² > 0.99の水準でカウントが読めるとGitHubのREADMEにも整理されています。

**Transformer内部には、カウント情報がある。**ということがここで分かります。

つまり、少なくともこの実験条件では、モデルは完全に「数を持てていない」わけではありません。


3. では、なぜ出力では間違えるのか

内部に数があるのに、出力では間違える。 ここがこの論文の中心です。

論文では、この差を調べるためにlogit-lensを使います。

logit-lensとは、モデルの中間層の隠れ状態を、そのまま出力ヘッドに通してみる方法です。

通常、Transformerは最終層の隠れ状態をlm_headに通して、次に出すトークンを決めます。 logit-lensでは、中間層の隠れ状態にも同じ出力ヘッドを当てて、「この層の情報をそのまま出力に変換したら何が出るか」を見ます。

つまり、こういう問いを立てます。

線形プローブでは数が読める。では、モデル自身の出力ヘッドでもその数を読めるのか。

結果は、読めません。

論文では、entity-mean位置の隠れ状態に対するlogit-lens精度はピークでも23.4%にとどまり、実際の生成位置であるlast-tokenでもピーク41.6%にとどまると報告されています。一方で、同じ内部表現から線形プローブは高精度にカウントを読めます。

ここで、かなりはっきりしたギャップが見えます。

線形プローブなら読める。しかし、モデル自身の出力ヘッドでは読めない。

つまり、モデル内部に情報はある。 しかし、その情報が出力にうまく接続されていません。


4. 「正しい答え、しかし間違ったベクトル」

では、なぜ出力ヘッドは数を読めないのでしょうか。

論文の答えは、ベクトルが合っていないからです。

モデル内部には、カウントを表すベクトルがあります。 一方、lm_headには、数字トークンを出すためのベクトルがあります。

たとえば、内部状態が「3匹」という情報を強く持っているなら、そのベクトルが数字トークン3の出力ベクトルと揃っていれば、3のlogitが上がるはずです。

しかし実際には、カウントベクトルと数字トークン行はほぼ直交しています。論文は、カウントを符号化するベクトルが、出力ヘッドの数字行とほぼ直交していることを報告し、この性質はランダムベクトルと区別できない水準だと説明しています。

イメージとしてはこうです。

内部表現空間

カウントベクトル       →  「3匹っぽさ」が増えるベクトル

数字トークンベクトル   ↑  「3」を出すベクトル

この2つがほぼ直角

この状態では、内部で「3匹」という情報が強くなっても、3という数字トークンの出力logitは自然には上がりません。

論文のタイトルになりますが、The Right Answer, the Wrong Direction、つまり正しい答えはある。けれど、向いているベクトルが違うという状態になっています。

LLMの失敗を「分かっていない」と片付けるのではなく、分かっている情報が、出力に接続されていないと見ています。


5. なぜ数字トークンのベクトルはズレるのか

ここで疑問が出ます。

なぜ数字トークンの出力ベクトルは、カウントベクトルと揃わないのでしょうか。

理由は、数字トークンがカウント結果としてだけ使われるわけではないからです。

たとえば、数字の「3」は次のような文脈が考えられます。

  • 2023年
  • 3月
  • 第3章
  • 手順3
  • 3.14
  • リスト番号
  • 年齢
  • ページ番号
  • 算術式

つまり、数字トークンは非常に多様な文脈で登場します。

そのため、出力ヘッドの数字行は、「カウント結果としての3」だけに合わせて学習されません。 むしろ、日付、序数、リスト、番号、算術など、さまざまな数字使用文脈の平均ベクトルに引っ張られます。

論文でも、出力ヘッド行は「次トークンがその数字だった隠れ状態」の期待値に向かって学習されるが、数字トークンの場合、その条件付き分布は非カウント文脈に支配されやすいと説明されています。結果として、カウントベクトルと数字行が揃わない安定状態が生まれる、という解釈です。

ここで重要なのは、問題が単純な「数字が少ない」ではないことです。

数字自体は訓練データに大量にあります。 しかし、カウント結果として数字を出す文脈は、数字トークン全体の中では特殊です。

だから、モデル内部では数を表していても、数字トークン出力のベクトルとは一致しにくいのです。


6. 出力ヘッドだけ直す:9-row repair

原因が出力ヘッドにあるなら、そこだけ直せばよいのでしょうか。

論文は、ここで非常に小さい介入を試します。

lm_headのうち、数字19に対応する9行だけを更新する。

これが9-row output-head repairです。

巨大モデル全体を再学習するわけではありません。 数字トークンに対応する9行だけを修正します。

論文では、この9行修正は36,864パラメータだけの更新であり、制約付きのnext-token数字予測では60.7〜100.0%まで改善すると報告されています。

この結果から、数字9行だけを直して精度が上がるなら、失敗の一部は本当に出力ヘッドの数字行に局在していると言えます。


7. しかし、9-row repairだけでは通常生成は直らない

ただし、ここで話は終わりません。

9-row repairは、制約付き評価では効きます。 しかし、通常の自己回帰生成ではうまくいきません。

論文では、9-row repairは制約付きnext-token評価では改善する一方、自己回帰生成では0.0%になると明記されています。

なぜでしょうか。

制約付き評価では、候補を数字トークンに絞ります。 つまり、モデルは19の中から選べばよい。

しかし通常生成では、候補は全語彙です。 数字だけでなく、空白、改行、説明文、記号、単語など、あらゆるトークンが競合します。

この場合、数字9行だけを直しても、モデルが最初に数字を出すとは限りません。 「Let's think」や説明文を出そうとする可能性もあります。

つまり、9-row repairは出力ヘッドの数字ベクトルを直しますが、モデルの生成経路全体を「直接数字を出すベクトル」へ変えるわけではありません。

ここが重要です。

出力ヘッドのズレは原因の一部だが、実用的な生成では上流のルーティングも問題になる。


8. Diagnostic Probe Steering:内部の答えを外から押し出す

論文では、もう1つの診断方法としてDPSを使います。 DPSは、Diagnostic Probe Steeringの略です。

これは、線形プローブが読んだカウント値を使って、該当する数字トークンのlogitを直接押し上げる方法です。

たとえば、プローブが「答えは7に近い」と読んだら、7のlogitを外部から強くします。

式としては、次のような形です。

logit(k) <- logit(k) + alpha * exp(-((c_hat - k)^2) / (2 * sigma^2))

ここで、

  • c_hat:プローブが読んだカウント
  • k:候補数字
  • alpha:押し上げる強さ
  • sigma:近い数字にもどれくらい影響させるか

です。

DPSは、通常の実用手法というより、仮説検証のための道具です。

もし内部に正しい数があるなら、それを外部からlogitへ注入すれば、正しい数字が出るはずです。 逆に、内部に数がないなら、DPSでもうまくいきません。

論文では、DPSは出力ヘッドを迂回してカウントベクトルを使う診断手法として位置づけられており、9-row repairと合わせて幾何学的ボトルネック仮説を検証する役割を持っています。


9. 実用的に直すには、LoRA Q/Vが効く

通常生成まで直すために、論文が使うのがLoRA Q/Vです。

これは、attentionのQueryとValueにLoRAを入れる方法です。

9-row repairは、最後の出力ヘッドを直す方法でした。 一方、LoRA Q/Vは、もっと上流のattention経路を直します。

なぜこれが必要なのでしょうか。

通常生成では、各ステップで新しい隠れ状態が作られ、それが出力ヘッドへ渡されます。 もしその隠れ状態の中で、カウント情報が出力に読まれやすい形に増幅されていなければ、出力ヘッドだけ直しても十分ではありません。

LoRA Q/Vは、attentionのルーティングを変えることで、カウント情報が最終層・最終位置・出力ヘッドへ届きやすくなるようにします。

論文では、LoRA Q/V rank-16が真のgreedy自己回帰生成で83.1% ± 7.2%を達成すると報告されています。これは、9-row repairが生成では0.0%だったことと対照的です。

GitHubのREADMEにも、主要結果として、LoRA Q/V generationが5 seedsで83.1% ± 7.2%、CoT baselineが20.2% ± 1.9%と整理されています。

ここから見えるのは、次のことです。

出力ヘッドだけ直すと、制約付き評価では改善する。しかし通常生成まで直すには、attentionの上流ルーティングを直す必要がある。


10. 結果を整理する

この論文の主要結果を簡単に並べると、次のようになります。

方法何をしているか何が分かるか
線形プローブ内部状態から数を読むモデル内部には数がある
logit-lens中間層をlm_headで読む出力ヘッドでは数を読みにくい
コサイン類似度カウントベクトルと数字行の向きを測るほぼ直交している
9-row repair数字9行だけ修正制約付き評価では改善する
DPSプローブで読んだ数をlogitへ注入内部の数を使えば正解できる
LoRA Q/VattentionのQ/V経路を修正通常生成でも大きく改善する

まず内部表現を確認する。 次に出力ヘッドとのズレを確認する。 さらに最小限の介入で原因を局在化する。 最後に、通常生成で効く修正対象を見つける。

この順番で、モデルの失敗を分解しています。


11. Chain-of-Thoughtでは十分ではないのか

カウントが苦手なら、Chain-of-Thoughtで「一つずつ数えて」と促せばいいのではないか、と思うかもしれません。

もちろん、実用上はCoTで改善することもあります。 しかしこの論文の実験では、CoTだけではLoRA Q/Vほど改善しません。

GitHubのREADMEでは、few-shot CoT baselineは20.2% ± 1.9%と整理されており、LoRA Q/V generationの83.1% ± 7.2%とはかなり差があります。

これは示唆的です。

CoTは、モデルに生成途中の経過を書かせる方法です。 しかし、問題が「内部表現と数字トークン出力のズレ」にあるなら、プロンプトだけで完全に直るとは限りません。

つまり、この論文は「考えさせれば直る」とは別の視点を出しています。


12. この論文が直しているもの、直していないもの

この論文は、LLMのすべての推論失敗を直したわけではありません。

対象としては、主に低語彙の集約タスクです。

たとえば、

  • 文章中の対象数を数える
  • 文字数を数える
  • リストの長さを答える
  • 足し算の答えを短い数字で出す
  • 多数決でラベルを選ぶ
  • 最大値を選ぶ

のようなタスクです。

共通点は、答えが短いトークン集合になることです。 内部に答えがすでにあり、それを数字やラベルとして出せばよいタイプの問題です。

一方で、GSM8Kのような多段階数学推論や、MMLUのような広範な知識選択では、同じボトルネックがそのまま当てはまるわけではありません。論文でも、MMLUやGSM8Kではこのreadout bottleneck効果は確認されず、DROPでは部分的な効果にとどまると説明されています。

つまり、この論文の主張は次のように限定して読むべきと思いました。

内部に答えがあるのに出力できないタイプの失敗には、この説明が効く。しかし、そもそも答えを多段階に作る必要がある問題には、そのままでは効かない。

「LLMは本当は全部分かっているが、出せないだけ」という話ではありません。 あくまで、カウントや低語彙集約タスクにおける、特定の失敗様式を説明している論文です。


まとめ

この論文は、Transformerのカウント失敗を、単なる「数えられない問題」として扱いません。

モデル内部を調べると、カウント情報はかなり正確に存在しています。 しかし、その情報を表すベクトルが、数字トークンを出すlm_headのベクトルとほぼ直交しています。

そのため、モデルは内部に正しい答えを持っていても、それを数字としてうまく出せません。

9-row repairは、このズレが出力ヘッドの数字行に局在していることを示します。 しかし通常生成では、出力ヘッドだけでは不十分です。 LoRA Q/Vでattentionの上流ルーティングを直すことで、実際のgreedy生成でも大きく改善します。

Transformerがカウント問題で失敗する原因は、常に「答えを持っていない」こととは限りません。答えが内部にあるにもかかわらず、その答えを出力トークンへ運ぶベクトルがズレ、うまく出力できなくなっていることがあります。