🤖 AI

Inference Time Causal Probing in LLMs

本論文は、既存の分類器依存型アプローチよりも信頼性の高い因果的介入を実現するためにモデルのネイティブ出力を用いてLLMの隠れ状態を直接誘導するプローブ不要の勾配ベース手法である隠れ状態駆動マージン介入（HDMI）を提案し、あわせてテキスト編集のための先見型変種を導入する。

原著者： Sadegh Khorasani, Saber Salehkaleybar, Negar Kiyavash, Matthias Grossglauser

公開日 2026-05-11

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Sadegh Khorasani, Saber Salehkaleybar, Negar Kiyavash, Matthias Grossglauser

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

この論文を、平易な言葉と日常的な比喩を用いて解説します。

全体像：「ブラックボックス」の修正

Large Language Model（LLM）を、厨房にいる超優秀だが少し不透明なシェフだと想像してください。このシェフは素晴らしい物語を書くことができますが、「このスープに塩を入れたのはなぜですか？」と尋ねると、明確な答えを出せないかもしれません。彼らは単に「味が合う」と知っているだけです。

科学者たちは、このシェフがどのように決定を下すのかを理解したいと考えています。具体的には、「シェフは動詞 'run' を選ぶ際に、実際に 'cats' のような『複数形』という概念を使っているのか、それとも単なる偶然なのか？」を知りたいのです。

これを検証するために、研究者たちは**Causal Probing（因果的プロービング）**という手法を使用します。これは「もしも」の実験のようなものです。「では、主語を単数形ではなく複数形だと仮定してみましょう。シェフは動詞を 'runs' から 'run' に変更するでしょうか？」と問うのです。

旧来の手法の問題点：「翻訳者」の問題

以前、この実験を行うためには、研究者たちは翻訳者（「プローブ」と呼ばれる）を雇う必要がありました。

まず、この翻訳者を訓練して、シェフの秘密のメモ（隠れ状態）を読み、主語が単数形か複数形かを推測させます。
次に、翻訳者のフィードバックを用いて、シェフのメモを微調整し、意味を変更させます。

欠陥： この翻訳者は、シェフの言語を完璧に話しているとは限りません。翻訳者自身が「複数形」のあり方について独自のルールを持っており、それがシェフの実際の思考プロセスと一致しない可能性があります。これは、異なるブランドの車のマニュアルを使ってエンジン修理を試みるようなものです。正しいボタンを押すかもしれませんが、エンジンの真の構造を理解していないため、他の何かを壊してしまう可能性があります。

新しい解決策：HDMI（「直接の微調整」）

著者たちは、**HDMI（Hidden-state Driven Margin Intervention：隠れ状態駆動マージン介入）**と呼ばれる新しい手法を提案しています。

比喩： 翻訳者を雇う代わりに、HDMI はシェフの頭と直接話します。

目標： シェフが「runs（単数形）」と言おうとしているところを、「run（複数形）」と言えるようにします。
旧来の方法： 翻訳者に「複数形」のボタンを探させて押させる。
HDMI の方法： HDMI はシェフの最終的な意思決定プロセス（出力）を眺めます。単語「run」の確率をわずかに上げ、「runs」の確率をわずかに下げるために必要な、正確な数学的な「微調整」を計算します。これは、2 つの単語の間の差（マージン）を見ることによって行われます。

なぜ優れているのか：

翻訳者不要： 概念を理解するために別の AI を訓練する必要はありません。モデル自身の脳を使って、出力を変更する方法を特定します。
精度： モデル自身の「幾何学構造」（単語をどのように自然に配置しているか）を利用するため、モデルが実際にどのように思考しているかと完全に整合します。
効率性： 長い迂回路ではなく、ステアリングホイールを素早く軽く叩くような、シンプルで高速な計算です。

「先読み」アップグレード：LA-HDMI

この論文では、テキスト編集のために、**LA-HDMI（Lookahead HDMI：先読み HDMI）**と呼ばれる洗練されたバージョンも紹介しています。

シナリオ： あなたが物語を書いていると想像してください。「The cat was sleeping.（猫は眠っていた）」と書きます。これを「The cats were sleeping.（猫たちは眠っていた）」に変更したいとします。

問題点： 単に「was」を「were」に変えるだけでは、文が壊れる可能性があります。文法がスムーズに流れるように、前の単語（「The」）や後の単語も変更する必要があるかもしれないからです。
LA-HDMI の解決策： この手法は現在の単語だけでなく、先読みを行います。変更を加えながら、文の次の数ステップをシミュレーションします。
- 「was」を「were」に変える場合、「The」は「The」のまま（変更なし）ですが、名詞が変わる場合はその前の冠詞を調整する必要があることを認識します。
- 単語が書かれる前に、モデルの隠れた思考を優しく誘導し、変更した単語だけでなく、文全体が滑らかで流暢であることを保証します。

これは、今すぐ左折するよう指示するだけでなく、後で行き止まりに陥らないように、先のルート全体を計算する GPS のようなものです。

効果はあったか？（結果）

著者たちは、文法と論理をテストするために設計された 2 つの主要な「トレーニングジム（データセット）」でこれをテストしました。

LGD Agreement（主語と動詞の一致）： 「he is」対「they are」のように、主語と動詞が一致しているか確認します。
CausalGym： 複雑な文法パズルのセットです。

スコアボード：
彼らは 2 つのことを測定しました。

完全性（Completeness）： 意味は正常に変更されましたか？（例：「runs」が「run」になりましたか？）
選択性（Selectivity）： 誤って他のものを壊しましたか？（例：時制や文全体の意味を誤って変更してしまいましたか？）

結論：
HDMI は、旧来の手法よりも一貫して高いスコアを記録しました。文の残りを混乱させることなく、変更すべきものを正確に変更することに優れていました。これは、Llama や Pythia など、さまざまな種類の AI モデルでも機能することを証明しており、堅牢なツールであることを示しています。

まとめ

旧来の方法： AI の考えを変える方法を推測するために、別のツールを使用する（しばしば不正確）。
HDMI： AI の出力自体を使って、考えを変えるための完璧な微調整を計算する（正確かつ効率的）。
LA-HDMI： 「水晶玉」を使って HDMI を活用し、編集された単語だけでなく、文全体が自然に流れるように、テキストを滑らかに編集する。

この論文は、この「直接の微調整」アプローチが、AI モデルの思考と執筆のあり方を理解し制御する、より信頼性の高い方法であると結論付けています。

技術概要：LLM における推論時因果的プロービング（HDMI）

1. 問題定義

本論文は、大規模言語モデル（LLM）が予測のために内部表現をどのように利用するかを理解するために用いられる、現在の因果的プロービング手法の限界に焦点を当てています。

相関的プロービングの限界: 従来の手法は、隠れ状態に補助分類器（「プローブ」）を訓練して言語的性質（例：主語の単複）を検出します。これらは何が符号化されているかを示しますが、モデルがその情報を次のトークン生成にどのように利用するかを証明するものではありません。
因果的プロービングの限界: 既存の因果的アプローチは、隠れ状態に介入して性質を変更（例：単数を複数形に変更）し、出力の変化を観察しようとします。しかし、これらの手法は通常、介入の方向性を決定するために（プローブの勾配を使用してラベルを反転させるなど）、別途プローブを訓練することに依存しています。
- 欠点: これにより追加の訓練コストが生じ、性質固有の教師信号が必要となり、ミスマッチのリスクが高まります。プローブの分類境界はモデルの内部「予測幾何学」と一致しない可能性があり、モデルの生成を効果的に誘導しない介入や、無関係な性質を意図せず変更してしまう介入につながる恐れがあります。

著者らは、モデルの再訓練や補助プローブの訓練なしに、完全に推論時に介入を適用する推論時因果的プロービングという枠組みを提案します。

2. 手法：隠れ状態駆動マージン介入（HDMI）

中核的な貢献は、HDMIです。これはプローブ不要の勾配ベース手法であり、生成モデルのネイティブな出力ヘッドを直接活用して隠れ状態を誘導します。

2.1 中核メカニズム

HDMI は、「ソース」トークン（モデルの元の優先される継続）と「ターゲット」トークン（望む反事実的継続）との間でマージン目的関数を定義することで動作します。

目的関数: $\phi(x)_\tau - \phi(x)_\sigma$ というログイットマージン $L(x)$ を最大化します。ここで、 $\phi(x)$ はログイットベクトル、 $\tau$ はターゲットトークンインデックス、 $\sigma$ はソーストークンインデックスです。
介入: この手法は、特定の層 $\ell$ における隠れ状態 $h_\ell(x)$ に対して勾配上昇を行い、ターゲットトークンの確率を増加させつつソーストークンを減少させます。
計算効率: 最終層の場合、勾配は閉形式の行列 - ベクトル積 $\nabla_{h_L} L(x) = W_U^\top (e_\tau - e_\sigma)$ となります。これにより、反復的なプローブ訓練の必要性が回避されます。
マルチトークン拡張: この手法は、ターゲット集合（ $T^+$ ）とソース集合（ $T^-$ ）のログイットを合計して定義される $L_{set}(x) = \sum \phi(x)_{\tau_i} - \sum \phi(x)_{\sigma_j}$ として、トークンの集合（例：複数の有効な複数形動詞）をサポートします。

2.2 テキスト編集のための先読み HDMI（LA-HDMI）

ユーザーが修正されたシーケンスを提供するテキスト編集に対処するため、著者らはLA-HDMIを導入します。

課題: 標準的な自己回帰生成では、ステップ $t$ における隠れ状態が $t' > t$ の将来の編集を認識することができません。
解決策: LA-HDMI は、期待されるトークン埋め込みを通じて勾配を逆伝播させる「先読み」メカニズムを導入します。
- 前方パス（流暢性を維持するため低温 $\beta_f$ を使用）と勾配パス（期待埋め込みを計算するため高温 $\beta_g$ を使用）を分離します。
- 目的関数 $J_t(h_t)$ は、将来の編集位置（ $S_{max}$ ステップ先）に対するマージン目的関数を累積し、現在の隠れ状態 $h_t$ を調整して、将来の生成をユーザーの編集に向けて誘導しつつ、編集されていない領域での流暢性を維持します。

3. 評価枠組み

本論文は、Canby ら [2024] の枠組みに従い、2 つの主要な指標を用いて介入を評価しています：

完全性（Completeness）: 介入が隠れ状態において対象の性質（例：主語の単複の変更）を成功裡にシフトさせたかどうかを測定します。これは別途の検証プローブによって確認されます。
選択性（Selectivity）: 介入後に無関係な性質（ノイズ変数、例：主語以外の名詞の単複）が安定して維持されるかどうかを測定します。
信頼性（Reliability）: 完全性と選択性の調和平均であり、介入の質全体の指標として機能します。

4. 主要な結果

著者らは、LGD 合意コーパスとCausalGym ベンチマークを用いて、Meta-Llama-3-8B-InstructおよびPythia-70M上で HDMI を評価しました。

性能: HDMI は、ほとんどのタスクにおいて、プローブ駆動のベースライン（AlterRep、FGSM、PGD）よりも一貫して高い信頼性を達成します。
- 合意スイート（例：主語 - 動詞の一致）において、HDMI はしばしば完全性を完璧に、選択性を高く達成します。
- ベースラインとは異なり、HDMI の性能は介入用プローブの品質に依存しません（なぜならプローブを使用しないため）。AlterRep や PGD などのベースラインは、プローブ精度が低い場合、小規模モデル（Pythia-70M）で顕著な性能低下を示しました。
頑健性: HDMI は異なるモデルサイズやタスクにわたって一貫した信頼性を維持しますが、プローブベースの手法は特定の性質やモデルに敏感で、結果がばらつきました。
テキスト編集: LA-HDMI は、「girl」を「owl」へ、「was」を「were」へ変更するなど、ユーザー指定の編集に向けて生成を誘導する能力を成功裏に実証しました。同時に、文法的な流暢性を維持し、必要な文脈（例：「a」を「an」に変更）を調整しました。

5. 意義と主張

本論文は、HDMI がプローブ依存からモデルネイティブな因果的プロービングへの転換を表すと主張しています。

整合性: モデル自身の出力ヘッドを用いて介入方向を定義することで、HDMI は摂動がモデルの実際の予測幾何学と整合することを保証し、外部プローブのミスマッチリスクを回避します。
効率性: この手法は計算的に軽量であり、補助分類器の訓練というオーバーヘッドなしに、隠れ状態に対する勾配上昇のみを必要とします。
汎用性: このアプローチは自然にマルチトークン性質に拡張され、先読み勾配を介した流暢なテキスト編集のメカニズムを提供します。

著者らは、ハイパーパラメータ（ステップサイズ、温度、ホライズン）への感度や、長距離テキスト編集タスクにおける勾配の減衰といった限界を認めています。さらに完全性を犠牲にすることなく選択性を向上させるため、適応的な層選択とステップサイズスケジューリングに関する将来の作業を提案しています。