技術要約：ロジット貢献度スコアリング（LOCOS）による非リテラルな検索ヘッドの特定

1. 問題提起

大規模言語モデル（LLM）は、関連するコンテキストのスパンから単にトークンを文字通りコピーするのではなく、その意味内容から回答を合成することが頻繁にある。既存のメカニスティック・インタープリタビリティ（機械論的解釈可能性）の研究は、コンテキストへのアクセスを担う「検索ヘッド（retrieval heads）」を特定してきたが、現在の検出手法は根本的にリテラルな（逐語的な）検索に限定されている。

既存の検出器（Wu et al., 2025; Fu et al., 2025など）は、アテンション重みの観測量に依存している。これらは、アテンションを向けたソース・トークンが生成された回答トークンと一致する場合に、高いスコアを与える。この「トークン一致」の基準は、コピー＆ペーストを行うヘッド（リテラルな検索）を特定することには成功しているが、非リテラルな（非逐語的な）検索を行うヘッドを検出するには失敗している。非リテラルなシナリオでは、モデルはソースのフレーズ（例：「エッフェル塔」）にアテンションを向けるが、その出力値（OV）回路を通じて異なる回答トークン（例：「Yuki」）を合成する場合がある。アテンションベースのスコアリングは、ヘッドが「どこを読んだか」は捉えるが、「何を書いたか」を無視しており、非リテラルな合成が行われる重要なメカニズムを見落としている。

2. 手法：ロジット貢献度スコアリング（LOCOS）

著者らは、アテンション分布ではなく、最終的な回答のロジットへの貢献度に基づいてアテンション・ヘッドを評価する「書き込み認識型（write-aware）」の検出器である**ロジット貢献度スコアリング（LOCOS）**を導入する。

コアメカニズム

LOCOSは、アテンション・ヘッドの動作を以下の2つの回路に分解する：

1. QK回路： ヘッドが「どこを」読むかを決定する（アテンション重み $\alpha$）。
2. OV回路： ヘッドが残差ストリームに「何を」書き込むかを決定する（出力射影 $W_O v$）。

本手法は、ヘッド $(l, h)$ の出力を、正解の回答トークンのアンエンベディング方向 ($u_{y_t}$) へと射影することでスコア化する。

3ステップの手順

1. 位置ごとのロジット貢献度：
   トークン $y_t$ を生成するデコーディング・ステップ $t$ において、ヘッド $(l, h)$ を通じたソース位置 $j$ の貢献度は次のように計算される：
   $$\phi^{(l,h)}_{t,j} = u_{y_t}^\top o^{(l,h)}_{t,j} = \alpha^{(l,h)}_{t,j} \cdot u_{y_t}^\top W^{(l,h)}_O v^{(l,h)}_{t,j}$$
   このスカラー $\phi$ は完全なパイプラインを捉える：ヘッドがどれほど高いアテンション（$\alpha$）を持っていても、そのOV出力が回答方向と直交していれば、貢献度はほぼゼロとなる。逆に、中程度の注意であっても、OV回路がアテンションを向けた表現を回答に整合したベクトルへと変換すれば、高い貢献度を持つ。

2. 空間的コントラスト：
   検索特有の信号を分離するため、LOCOSは**ニードル（needle）**位置（関連するコンテキスト・スパン $[s_\tau, e_\tau)$）の貢献度を、**オフ・ニードル（off-needle）**位置と比較する。
   $$\Phi^{(l,h),+}_t = \sum_{j \in \text{needle}} \phi^{(l,h)}_{t,j}$$
   $$\Phi^{(l,h),-}_t = \frac{e_\tau - s_\tau}{N_t - (e_\tau - s_\tau)} \sum_{j \notin \text{needle}} \phi^{(l,h)}_{t,j}$$
   この再スケーリング係数は、オフ・ニードルの合計が長さに対してニードルの合計と比較可能であることを保証する。単一ステップにおけるスコアは、この空間的コントラスト（単一ステップ内でのニードル vs オフ・ニードル）である $\Phi^{(l,h),+}_t - \Phi^{(l,h),-}_t$ である。これは、先行研究で使用されている時間的コントラスト（回答ステップ vs 非回答ステップ）に代わるものである。

3. 集計：
   最終的なスコア $S_{l,h}$ は、モデルが正解を生成するすべてのデコーディング・ステップにおけるこれらの空間的コントラストの平均であり、すべての成功した試行全体で集計される。負のスコアは、回答のロジットへの貢献が主にオフ・ニードルの位置から来ているヘッドを示す。

3. 主な貢献

• 書き込み認識型の検出： 非リテラルな検索を特定するには、アテンション重みだけでなく、OV回路の出力を測定する必要があることを確立した。
• LOCOSアルゴリズム： 回答のロジットへの射影された貢献度によってヘッドをスコアリングし、ニードルとオフ・ニードルのソースを対比させる、シングル・フォワードパスの手法。
• 因果的検証： 平均アブレーション実験を通じて、LOCOSが非リテラルな検索に因果的に重要なヘッドを特定する一方で、アテンションベースのベースラインはそれに失敗することを実証した。

4. 実験結果

本手法は、3つのモデルファミリー（Qwen3 (8B, 14B, 32B), Gemma-3 (12B, 27B), OLMo-3.1 (32B)）にわたる NoLiMa ベンチマーク（非リテラルな「針を干し草の中に探す」タスク）で評価された。

アブレーション性能 (NoLiMa)

• 優れた劣化： 上位 $k$ 個のLOCOSヘッドを平均アブレーション（mean-ablating）すると、すべての注意ベースのベースライン（Wu/NIAHスコア、Wu/NoLiMaスコア、およびランダムなヘッド）と比較して、ROUGE-Lがより急激に低下した。
• Qwen3-8Bのケーススタディ： 上位50個のLOCOSヘッドをアブレーションすると、ROUGE-Lは 0.401 から 0.000 へと崩壊した。対照的に、最も強力な注意ベースのベースライン（Wu/NIAH）は、同じ深さにおいて 0.292 のROUGE-Lを維持している。
• 特異性： LOCOSヘッドのアブレーションは、文脈の検索を激しく劣化させる一方で、パラメトリックな想起（例：都市と国の関連付け）や算術推論は維持した。これにより、特定されたヘッドが、一般的な出力ではなく、検索に特化したものであることが確認された。

リテラル vs 非リテラルな特異性

• 上位 $k$ 個のLOCOSヘッドのアブレーションは、非リテラル（NoLiMa）とリテラル（NIAH）の両方の検索を劣化させる。
• しかし、その低下はNoLiMaにおいて著しく大きい。これは、LOCOSが両方のレジームに関連するヘッドを捉える一方で、トークン一致検出器が見逃してしまう非リテラルな合成を担当するサブセットのヘッドを独自に特定していることを示している。

ダウンストリーム転移

• 下流のロングコンテキスト・ベンチマーク（MuSiQue および BABILong）において、上位50個のLOCOSヘッドのアブレーションは、特にQwen3ファミリーにおいて性能を著しく低下させた（例：Qwen3-8BにおけるMuSiQueの精度は0.55から0.08に低下）。
• LOCOSとアテンションベースのベースラインのランキングはベンチマークによって異なり、マルチホップQAタスクがリテラルおよび非リテラルな検索回路の両方を動員していることが示唆される。

レイヤー分布

• Qwen3ファミリーおよびGemma-3-27Bにおいて、LOCOSスコアは後半のレイヤーに強く集中している。
• 著者らは、この集中が「ダイレクトパス仮定（出力に近いヘッドほどロジットへの直接的な影響が大きい）」の影響を受けている可能性があると述べているが、因果的パッチング実験は、この後半レイヤーのパターンが単なるアーティファクトではないことを示唆している。

5. 意義と主張

本論文は、LOCOSが非リテラルな検索ヘッドの最初のメカニスティックな特定を提供すると主張している。

• 先行研究の限界： 既存の検出器は、アテンションを向けたトークンが生成されたトークンと一致する場合に報酬を与えるため、非リテラルな検索に対して「盲目」である。これは、ソースに注目しながら異なる回答を合成するヘッドを見落とす。
• メカニスティックな洞察： OV回路の射影を測定することで、LOCOSは合成に不可欠な「書き込み」操作を捉える。
• 実用的な有用性： 特定されたヘッドは、長いコンテキストの推論において因果的に重要である。この知見は以下に影響を与える：
  • 解釈可能性： LLMがどのようにコンテキストにアクセスし、それを合成するかについての、より正確なマップを提供する。
  • 効率性： 一様な、あるいはアテンションベースのヒューリスティックではなく、検索を担当する特定のヘッドのみを保持することで、より標的を絞ったKVキャッシュ圧縮を可能にする。
  • ハルシネーションの軽減： 関連するコンテキストを正しく検索できずに失敗するヘッドを特定し、介入するためのメカニズムを提供する。

著者らは、自らの知見がテストされたアーキテクチャ（標準的なデコーダー専用トランスフォーマー）およびNoLiMaベンチマークに特有のものであることを認め、一般化については慎重な姿勢をとっている。彼らは、LOCOSがすべての検索メカニズムを特定すると主張しているのではなく、トークン一致基準によって見逃されるサブセットを、LOCOSが正常に分離できることを主張している。