技術要約: IG-Lens

問題提起

デコーダーのみのTransformerは、残差ストリームが次のトークンの実行中の推定値を洗練させていく反復的な推論プロセスとして概念化されることが多い。ここで、重要な解釈可能性に関する問いが生じる：具体的にどのレイヤーにおいて、予測されたトークンの確率が実際に決定されるのか？

既存のレイヤーごとの読み出し（readout）ツールは、確率空間においてこれを正確に回答できていない：

• Logit Lens & Tuned Lens: これらは各レイヤーにおける確率レベルを報告するが、加法的な分解は提供しない。それらの推定値はバイアスがあり、深さに対して単調ではなく、各レイヤーの確率の総和には数学的な意味を持たない。
• Direct Logit Attribution (DLA) & Embedding Methods: これらは加法的ではあるが、ロジット空間で動作する。Softmax関数は非常に非線形であるため、ロジットの加法的な分解は、実用的な関心事である確率の加法的な分解には直結しない。
• Layer Conductance & Layer Integrated Gradients (IG): これらの手法は勾配をレイヤーごとに積分するが、各レイヤーを独立したベースラインに対して属性付けする。その結果、各レイヤーの項は、ターゲットの予測確率の総変化量へと「テレスコーピング（望遠鏡のように積み重なる）」することはない。

手法: IG-Lens

著者らは、Transformerのレイヤー間にわたって、確率の正確な加法的分解を提供する手法である IG-Lens を導入する。

コアコンセプト

IG-Lensは、ベースラインから最終レイヤーまでの隠れ状態を通る単一のパスに沿って、Integrated Gradients (IG) をテレスコーピング的に適用する。

1. パスの定義: 手法は、レイヤーのシーケンス $0 = L_0 < L_1 < \dots < L_k = n$ を選択し、隠れ状態のパス $h_{L_0} \to h_{L_1} \to \dots \to h_{final}$ を定義する。
2. Head-Once Readout: 決定的な点として、読み出し関数 $f(h)$ は、単一の隠れベクトルに対して一度だけ適用される最終的なアンエンベディング（unembedding）として定義される：$f(h) = \text{softmax}(\text{head}(\text{norm}(h)))_{y_t}$。この関数にはAttentionもMLPも含まれない。すべてのトークン間の混合は、分離された隠れ状態を生成したフォワードパスですべて完了していると仮定される。
3. Telescoping Attribution: レイヤー $L_j$ の寄与は、$h_{L_{j-1}}$ から $h_{L_j}$ への直線に沿ったセグメントIGとして定義される。勾配の定理により、セグメントに沿った勾配の積分は、エンドポイントの値の差に等しい。これらのセグメントを合計することで、テレスコーピングな和が得られる：
   $$\sum_{j=1}^k IGL_j = f(h_{final}) - f(h_{base}) = \Delta p$$
   これにより、softmaxとLayerNormを線形化したり凍結したりするのではなく、積分経路の中に組み込むことで、属性の総和が正確にターゲット確率の変化量 $\Delta p$ となることが保証される。

予測認識型エスティメータ (Prediction-Aware Estimator)

積分は正確であるが、有限ステップのエスティメータは通常、離散化誤差を導入する。IG-Lensは、デフォルトとして予測認識型エスティメータ（IDGIに着想を得たもの）を採用する：

• 生の勾配に基づくのではなく、観測されたターゲット確率の変化（$f(h^{(s)}) - f(h^{(s-1)})$）によって各ステップにクレジットを割り当てる。
• スカラーの確率読み出しの場合、これは各セグメントをエンドポイント値のテレスコーピングな和へと収束させる。
• 結果: 完備性（Completeness）は、ステップ数 $m \ge 1$ のいかなる場合においても、浮動小数点精度で保持される。モデルが勾配の感度を示しながらも出力の変化を示さないステップは自然にフィルタリングされ、リーマンの離散化誤差を排除する。

実装

本手法は、バックプロパゲーション（autograd）を必要としない、シングルパスのバッチアルゴリズムを使用する。

• すべての補間点を単一のテンソルにスタックして評価することで、全トークンのレイヤーマップを一度のフォワードパスで計算し、勾配計算のオーバーヘッドを回避する。
• この効率性は、特に「head-once」設計によって可能となり、これによりトークン間の独立性が確保されている。

主な貢献

1. 正確な加法的分解: IG-Lensは、レイヤー間にわたって確率（ロジットではなく）の正確な加法的分解を提供する初めての手法であり、属性の総和が $\Delta p$ に一致する。
2. Telescoping Identity: 本論文は、セグメント積分が予測の総変化量へとテレスコープすることを、softmaxを積分経路内に保持したまま証明（Proposition 1）している。
3. 予測認識型エスティメーション: デフォルトのエスティメータは、ステップ数に関わらず機械精度での完備性を保証し、「応答のない感度」を持つステップを排除する。
4. 効率的な実装: バックワードコールなしで全属性マップを計算するシングルパスのフォワード評価を用いることで、計算効率が高い。
5. ポジショニング: 著者らは、なぜ既存の手法が確率空間で失敗するのか（ロジット空間の加法性、バイアス、またはテレスコーピングの欠如）を明確にし、IG-Lensをロジット空間の手法に対する確率空間の対応物として位置づけている。

結果と評価計画

• 数値的検証: 制御されたテストおよび実際のモデル（Llama-3.2-1B-Instruct）において、デフォルトのエスティメータは、各セグメントのエンドポイントの差を $\sim 10^{-7}$（float32精度）で再現する。テレスコーピングな和は、ステップ数 $m$ に依存せず、浮動小数点精度で正確に成立する。
• 例示的な出力: プロンプト "What is the capital of Vietnam?" における実験では、本手法は、Logit Lensの値が低から高へと遷移する特定のレイヤーセグメントに属性質量を集中させることに成功している。各トークンの属性の総和は、浮動小数点精度で $\Delta p$ と一致する。
• 評価戦略: 著者らは「発生（onset）」には固有の正解（ground truth）が存在しないことを認めている。したがって、彼らは**因果的アブレーション（causal ablation）**を通じて手法を検証することを提案している。すなわち、IG-Lensがレイヤー $L$ を発生として特定した場合、レイヤー $L$ をアブレーションすることは、より低い発生スコアを持つレイヤーをアブレーションすることよりも、ターゲット確率を抑制するはずである。このテストは、実行コストが低く、主張を反証可能なものとして設計されている。

意義と主張

本論文は、IG-Lensを反復的推論ビューの定量的洗練として位置づけている。Logit Lensが「このレイヤーにおける実行中の予測は何か？」と問うのに対し、IG-Lensは「各レイヤーは最終的な確率に対してどれだけの責任があるか？」と問い、それらの責任が全体に合致することを保証する。

控えめな主張と限界:

• 条件付き限界 (Conditional Marginal): 著者らは、IG-Lensは条件付き限界を測定していることを明示している。すなわち、$\text{norm}+\text{head}$ が $h_{L-1}$ に対して $h_L$ から読み取れる追加的な確率を測定している。これは、上層のAttentionやMLPブロックを経由した $h_L$ の効果を破棄している。
• トレードオフ: 本手法は、「総体的な因果効果（full forward passによるもの）」よりも「$\text{sum} = \Delta p$」を選択している。総体的な因果効果については、コンポーネント分解（DLAなど）が適切なツールであり続ける。
• 非単調性: 本手法は、符号付きの属性を用いることで、確率の軌跡が非単調（確率が構築された後に部分的に取り消される場合）であっても、絶対値ベースの手法よりも忠実なレイヤー寄与のビューを提供し、非単調な軌跡を扱う。

中心となる未解決の問いは、IG-Lensによる発生（onset）が、特に非単調な軌跡を持つトークンにおいて、因果的アブレーションが予測を最も抑制するレイヤーを予測するかどうかという、経験的な側面にある。