技術概要：多様体誘導アテンション・ステアリング（MAGS）

問題定義

大規模言語モデル（LLM）は、複雑な推論タスクに必要な基礎知識を有していることが多いにもかかわらず、生成中にその能力を信頼性高く発揮できないことが頻繁に起こる。同一のモデルとプロンプトに対する反復サンプリングでは、正解と誤った解の両方が得られるが、プロセスレベルのアノテーションによれば、エラーは通常、終端的な知識不足ではなく、中間的な推論段階で生じることが示されている。

既存の推論時活性化制御（activation steering）手法は、残差ストリームまたはアテンション出力に対して固定された事前計算された修正ベクトルを適用することで、これらの振る舞いを修正しようとする。しかし、これらのアプローチは推論タスクとの構造的な不一致という重大な欠陥を抱えている：

1. 無差別な摂動：モデルの現在の状態に関わらず、修正を均一に適用する。これは、すでに正しい段階を汚染するリスクがあり、かつ未発現のエラーを未然に防ぐことに失敗する。
2. 静的性質：推論軌跡がいつエラー状態へ逸脱するかを検知するメカニズムが欠如しており、すべての生成段階を等しく修正を必要とするものとして扱う。

本研究の核心的な仮説は、推論エラーが、個々のアテンションヘッドの出力空間における低次元部分空間内での構造化された逸脱として現れるというものである。正解と誤った軌跡は幾何学的に分離可能な領域を占有し、正解から誤った振る舞いへの遷移は、特定の低ランク方向に従う。

手法：多様体誘導アテンション・ステアリング（MAGS）

MAGS は、適応的かつ軌跡を認識する推論時介入であり、「正しさ多様体（correctness manifold）」からの逸脱が検出された場合のみ、アテンションヘッドの出力を誘導する。この手法は 3 つのフェーズで動作する：

1. 対照的エラー多様体の構築

訓練フェーズにおいて、システムは対照的な推論トレースのペア（同一の問題、一つは正解、一つは誤答）を分析し、エラーの幾何学的構造を学習する。

• ヘッドごとの差分ベクトル：各アテンションヘッド $(l, h)$ について、システムは正解トレースの平均活性化 $\mu_c$ と誤答トレースの平均活性化 $\mu_e$ を計算する。差分ベクトル $\delta = \mu_e - \mu_c$ は、問題固有の特徴を相殺し、エラーに起因する方向的なシフトを分離する。
• 部分空間学習：これらの差分ベクトルを複数の問題にわたって積み重ね、主成分分析（PCA）を行うことで、MAGS は低次元基底 $B$（トップ-$k$ 特異ベクトル）を抽出する。この基底は「エラー部分空間」、すなわち正解から誤答への逸脱が最大の分散を示すヘッド出力空間内の方向を張る。
• グローバル参照点：すべての正解トレースからグローバルな重心 $\mu_c$ を計算し、推論時の参照点として用いる。

2. 近傍に基づくエラー検知

推論中、MAGS は選択されたアテンションヘッドの出力をリアルタイムで監視する。

• 近傍スコア：各トークンステップ $t$ において、システムは中心化されたヘッド出力を学習されたエラー部分空間へ射影したもののノルム二乗を計算する：
  $$d_t = \| B(a_t - \mu_c) \|^2$$
  高いスコアは、ヘッドの活性化がエラー多様体方向に著しく逸脱したことを示す。
• 動的トリガー：ヘッドごとの閾値 $\tau$ は、正解トレース上で較正される（例えば、正解軌跡からのスコアの上位パーセンタイルに設定される）。修正が適用されるのは $d_t > \tau$ の場合のみである。これにより、すでに正解軌道上にある段階は攪乱されないことが保証される。

3. 標的化された修正

トリガーが作動すると、MAGS はアテンションヘッドの出力がレイヤーの出力投影へ渡される前に、その出力に対してその場（in-place）で修正を適用する：
$$\tilde{a}_t = a_t - \alpha B^\top B (a_t - \mu_c)$$
ここで、$\alpha \in (0, 1]$ はステアリング強度のハイパーパラメータである。

• 幾何学的解釈：この操作は、逸脱 $(a_t - \mu_c)$ をエラー部分空間へ射影し、それを減算する。これにより、エラー成分を実質的に破棄しつつ、直交補空間（エラー多様体と無関係な $(d_h - k)$ 次元）のすべての情報を保持する。
• 情報の保持：この手法は、エラー多様体と無関係な情報が保持されることを保証し、静的な手法で見られる「過剰修正」を回避する。

多目的拡張（MAGSu）

対立する目的（例えば、構文の妥当性対化学結合親和性）を処理するために、著者は MAGSu を提案する。これは各目的に対して独立した多様体を学習し、互いに重ならないヘッドの集合を通じて修正を適用することで、ステアリング目標間の干渉を防ぐ。

主要な貢献

1. 仮説の検証：本論文は、推論エラーが特定のアテンションヘッドの低次元部分空間内での構造化された逸脱として現れるという実証的証拠を提供する。診断実験により、正解と誤った軌跡はこれらの部分空間によって高度に分離可能であり、エラー信号はスパースで特定のアテンションヘッドに集中していることが示された。
2. 適応的介入メカニズム：MAGS は、ヘッドがエラー部分空間へ逸脱した場合にのみ修正を適用する動的な近傍ベースのトリガーを導入する。これは、現在の活性化状態に関わらず固定ベクトルを適用する静的な手法とは対照的である。
3. 実証的な優位性：この手法は、多様なベンチマークにおいて、ステアリングを行わないベースラインおよび既存の静的ステアリング手法（ITI、Angular Steering）を一貫して上回る性能を示した。

実験結果

著者は、3 つのモデルファミリー（Llama-3.1-8B-Instruct、Gemma-4-E4b-it、GPT-OSS-20B）を 3 つのドメインで評価した：

• 数学的推論（MATH-500、GSM8K）：MAGS は MATH-500 で最大の改善を達成し、ステアリングを行わない Llama ベースラインに対して精度を +5.2%、ステアリングを行わない Gemma ベースラインに対して +3.4% 向上させた。この改善は、エラーの蓄積が起きやすい多段階導出を要するタスクで最も顕著であった。
• コード生成（HumanEval、MBPP）：MAGS は、ステアリングを行わないベースラインに対して HumanEval の pass@1 を約 4 ポイント向上させた。注目すべきは、静的な手法（ITI、Angular Steering）が構文的一貫性を乱す可能性によりコード生成の性能を低下させたのに対し、MAGS は正しい段階での修正抑制によってこれを回避した点である。
• 分子生成（SMILES）：構文の妥当性と高い結合親和性の両方を要求する多目的タスクにおいて、MAGSu は両指標を同時に改善することに成功した。対照的に、ITI は妥当性を改善したが結合親和性を著しく低下させ、静的ステアリングにおける絡み合った目的による干渉を実証した。
• 流暢性とオーバーヘッド：MAGS は、ステアリングを行わないベースラインと同等のパープレキシティレベルを維持し、生成分布を歪めていないことを示した。計算オーバーヘッドは無視できるレベルであり、各ステップあたり $O(K \cdot k \cdot d_h)$ の演算に支配される（ここで $K$ は監視対象のヘッド数である）。

意義と主張

本論文は、MAGS が活性化制御を定常的な修正タスクではなく、逸脱検知問題として再定義すると主張している。エラーの発生を予測する低ランク基底を特定することにより、MAGS はスカラー強度の手法が欠く解釈可能な幾何学的構造を提供する。

著者は以下を主張している：

1. 正しさ多様体の普遍性：正解と誤った振る舞いを分離する低次元部分空間の存在は、推論、コーディング、科学的生成にわたって適用可能な、LLM アテンション幾何学の一般的な特徴であるように見える。
2. 軌跡認識の重要性：無差別な修正は有効な推論段階を攪乱する。逸脱が検出された場合のみ介入する能力こそが、MAGS の優れた性能とハイパーパラメータ感度に対する頑健性の主要な駆動力である。
3. 幾何学的制約による過剰修正の防止：活性化をエラー部分空間の補空間へ射影することにより、MAGS はモデルの内在的な能力を保持しつつ、エラー固有の成分のみを除去する。

本研究は結論として、MAGS は多様体適合のための対照的トレースのセットを必要とするが、LLM における推論の一貫性を向上させるための堅牢で低オーバーヘッドの解決策を提供し、生成軌跡の動的性質を考慮しない静的介入を上回るとしている。