reward-lens: A Mechanistic Interpretability Library for Reward Models

以下は、論文「reward-lens: A Mechanistic Interpretability Library for Reward Models」の詳細な技術的サマリーです。

1. 問題提起

人間のフィードバックに基づく強化学習（RLHF）は、大規模言語モデル（LLM）の行動を形成するために報酬モデル（RMs）に依存しています。報酬モデルが誤った選好を符号化している場合、その結果得られる方策はそれらの誤りを継承し増幅します。報酬モデルの行動評価（RewardBench など）は成熟していますが、報酬モデルに対する機械的解釈性は事実上存在しません。

既存の解釈性ツール（TransformerLens、Logit Lens、Direct Logit Attribution など）は、生成型 LLM 向けに設計されています。これらは、トークン分布を分析するために中間モデル状態を語彙のアン埋め行列（$W_U$）に射影することに依存しています。しかし、報酬モデルは $W_U$ をスカラー回帰ヘッド（$w_r$）に置き換えています。この構造的な違いにより、標準的なプリミティブは適用不可能となり、安全上重要な報酬モデルの内部メカニズムがブラックボックスのままとなる「構造的なギャップ」が生じています。

2. 手法：$w_r$ 中心の視点

本論文は、報酬モデル向けに機械的解釈性のプリミティブを適応させたオープンソースライブラリreward-lensを導入します。核心的な理論的洞察は、報酬ヘッドの重みベクトル$w_r$が、すべての解釈性の問いに対する自然な軸となるという点です。

このライブラリは、最終的な報酬 $r$ が最終残差ストリーム $h_{final}$ の線形汎関数であるという観察に基づいています：
$$r = w_r^\top h_{final} + b_r$$
残差ストリームは層間で加法的に分解される（$h_{final} = h_{embed} + \sum (attn^{(\ell)} + mlp^{(\ell)})$）ため、報酬は以下のように正確に分解できます：
$$r = w_r^\top h_{embed} + \sum (w_r^\top attn^{(\ell)} + w_r^\top mlp^{(\ell)}) + b_r$$

これに基づき、ライブラリは $W_U$ を $w_r$ に置き換えることで以下のプリミティブを実装します：

• Reward Lens: 層 $\ell$ における中間残差ストリームを $w_r$ に射影し、層全体にわたる選好信号の「結晶化」を追跡します。
• Component Attribution: 各サブレイヤー（Attention/MLP）の最終スカラー報酬への寄与を、$w_r^\top \cdot \text{contribution}$ を通じて計算します。
• Contrastive Activation Patching: 選好された補完と選好されない補完の間で活性化を交換し、$w_r^\top h_{final}$ の変化を測定します。これは、選好データが本質的に持つ対照構造により、RM にとってネイティブな手法です。
• SAE Feature Attribution: スパースオートエンコーダー（SAE）を用いて残差ストリームを特徴に分解し、特徴の整合性を $w_r^\top d_i$（ここで $d_i$ は辞書ベクトル）として計算します。

アーキテクチャ非依存性: このライブラリは、10 種類のメソッドアダプタープロトコルを持つHuggingFace ネイティブのアプローチを採用しています。これにより、重みの手術なしに、標準的なモデル（Llama、Mistral、Gemma-2）や複雑な多目的モデル（MoE ゲーティングを備えた ArmoRM など）をラップできます。

理論に基づく拡張: 本ライブラリは、最近の整合性理論の結果をコードとして実用化します：

• Distortion Index: 評価プローブによってカバーされていない品質次元を予測し（報酬ハッキングのリスク）、どの次元が不足しているかを特定します。
• Divergence-Aware Patching: パッチングが分布外活性化を誘発するかどうかを検出します。
• Misalignment Cascade Detector: 異なる整合性不全の次元間で相関する失敗モードを特定します。
• Reward-Term Conflict Analyzer: 異なる報酬目的が整合しているか、直交しているか、あるいは矛盾しているかをチェックします。
• Concept-Vector Analysis: 線形概念方向を抽出し、それらが報酬に対する因果的な用量反応を測定します。

3. 主要な貢献

1. ライブラリ（reward-lens）: 報酬モデルの機械的解釈性のための、最初のインストール可能かつアーキテクチャ非依存のツールキットです。これは、観測的（帰属）視点と因果的（パッチング）視点を単一の代数的枠組みの下で統合します。
2. 統合的な枠組み: RM に対するあらゆる有用な解釈性プリミティブは、$w_r$ への射影または $w_r$ に沿った分解であるという事実を確立します。
3. 実証的検証: 2 つの生産モデルへの体系的な適用：Skywork-Reward-Llama-3.1-8B-v0.2（単一スカラーヘッド）とArmoRM-Llama3-8B-v0.1（多目的ヘッド）。
4. 否定的な実証的発見: 論文は明示的に、線形帰属が報酬モデルにおける因果的重要性を予測しないと報告しています。

4. 主要な結果と知見

著者らは、モデルあたり約 695 組の RewardBench 選好ペアに対してライブラリを評価しました。

• 結晶化の深さ:
  • Skywork: 選好は層の深さの約 0.90（分数）で均一に遅れて結晶化し、出力付近で統合されたメカニズムが組み立てられていることを示唆しています。
  • ArmoRM: 選好はより早期に、かつより高い分散で結晶化し、その多目的かつ MoE ゲーティングされたアーキテクチャを反映しています。
• 帰属対因果性（忠実性のギャップ）:
  • 観測的帰属: 最終 MLP レイヤー（特に mlp_L31）が、両モデルにおいて報酬信号を支配しています。
  • 因果的パッチング: 初期層のコンポーネント（例：mlp_L0、attn_L0）が、しばしば最終報酬にとって最も重要です。
  • 相関: 帰属の大きさとパッチング効果の間のスピアマン順位相関は負またはほぼゼロです（Skywork: $\rho = -0.256$; ArmoRM: $\rho = -0.027$）。
  • 解釈: これはバグではなく、冗長な残差ストリームの特性です。後期層は並列経路が補償するため除去可能ですが、初期層は「荷重を支える」役割を果たしています。
• 報酬ハッキングのプロファイル:
  • Skyworkは、長さ、自信、フォーマットバイアスを罰します（これらを報酬する一般的な RLHF モデルとは異なり）。
  • ArmoRMは自信とフォーマットを報酬しますが、追従（sycophancy）を罰します。
• 概念の整合性:
  • Skywork は「合意」などの概念と強い線形整合性を示します（傾斜 $\approx 1.0$）。
  • ArmoRM は線形概念との整合性が弱く（大きさ $\approx 1\%$）、より複雑で非線形な報酬幾何学を示唆しています。

5. 重要性

• パラダイムシフト: 本論文は、研究者が報酬モデルにおいて観測的帰属を因果的重要性の代理として扱うことをやめるべきだと主張しています。両者の間のギャップは構造的であり、無視するのではなく測定されなければなりません。
• 安全性ツール: 「歪み」（カバーされていない次元）や「整合性不全の連鎖」を検出するツールを提供することで、このライブラリはトレーニング中の監視を可能にします。これにより、報酬モデルの監査は事後活動から、RLHF トレーニングループ中のリアルタイムの安全装置へと移行します。
• 編集可能性: この枠組みは、報酬信号を駆動する特定の特徴や概念を特定し、下流の方策トレーニング前に解釈性に基づく報酬モデルの編集（例：報酬ヘッドから「追従」という概念を抑制する）への道を開きます。
• オープンサイエンス: reward-lens を公開することで、著者らは報酬モデルの内部メカニズムへのアクセスを民主化し、最先端の AI システムの安全性を確保するための重要な一歩を踏み出しました。

結論として、reward-lensは、報酬モデルのスカラー出力という性質に合わせて機械的解釈性を適応させることで、AI 安全性における重要な空白を埋めています。それは、それらの内部論理がしばしば直感に反する（後期帰属 $\neq$ 初期因果性）ことを明らかにし、整合性の失敗が方策モデルに伝播する前に診断し修正するために必要なツールを提供します。