Before the Last Token: Diagnosing Final-Token Safety Probe Failures

技術的サマリー：最終トークンの前：最終トークン安全プローブの失敗の診断

問題定義
SafeSwitch などの大規模言語モデル（LLM）向けの現在の安全監視システムは、しばしば「最終トークンプローブ」に依存している。これらのシステムは、生成開始前にプロンプトのプリフィル段階における最終トークンの隠れ状態に基づいて分類器を訓練し、安全でないコンテンツを検出する。本論文は、このアーキテクチャにおける致命的な失敗モードを特定している。すなわち、これらのプローブは「クリーン」な有害プロンプトにおいて高い再現率を達成する一方で、「ジェイルブレイク」プロンプトの検出には頻繁に失敗するということである。ジェイルブレイクは、敵対的な枠組み、ロールプレイ、または隠蔽の中に有害なリクエストを埋め込む。著者らは、これらのラッパーが、訓練中に学習された特定の「有害対無害」の対比方向から最終トークンの表現を逸脱させ、有害なコンテンツがプロンプト内に存在するにもかかわらず、プローブが安全でないリクエストを見逃す原因となると主張している。この失敗は単なる容量の問題ではなく、プローブのボトルネック幅を増加させても検出ギャップを確実に解消できないことから、容量の問題ではないことが示唆される。

手法
本研究では、Llama-3.1-8B-Instruct、Mistral-7B-Instruct-v0.1、および OLMo3-7B-Instruct の 3 つのインストラクションチューニング済みモデルを評価対象とした。

1. プローブの訓練: SafeSwitch 風のプローブを、元の SafeSwitch 訓練分割から得られた「クリーン」な有害および無害プロンプトのみに基づいて訓練した。これらのプローブは、最終プロンプトトークンの隠れ状態を読み取り、それをボトルネック（初期幅 $w=64$）を介して投影し、安全性を予測する。
2. 評価: プローブを以下のデータでテストした。
   • 保持されたクリーンな有害プロンプト（SorryBench）。
   • HarmBench からのジェイルブレイクテンプレートでラップされた同じ有害プロンプト。
   • 偽陽性を測定するための無害かつ安全性に隣接するプロンプト（XSTest）。
   • 一般的な分布シフトを排除するための直接的な有害ベンチマーク（AdvBench）。
3. 幾何学的分析: 著者らはプローブの表現部分空間を分析した。彼らは、さまざまな集合（例：見逃されたジェイルブレイク対無害プロンプト）の平均値差ベクトルを計算し、それらの「プローブ可視エネルギー」（SVD によるプローブの行空間によって捕捉される方向のノルムの割合）を測定した。
4. 軌跡分析: 最終トークンのみに依存するのではなく、著者らはプリフィル段階中のトークンごとの隠れ状態軌跡を分析した。彼らは、安全でない証拠がシーケンスのより早期に現れていたが、最終トークンまでに崩壊しているかどうかを調査した。
5. 診断モデリング: 軌跡情報が見逃された検出を回復できるかどうかをテストするため、著者らはクリーンな訓練データ上で単純な PCA-HMM（主成分分析 - 隠れマルコフモデル）を適合させた。このモデルは、クリーンな有害プロンプトと無害プロンプトの軌跡ダイナミクスを学習し、ジェイルブレイクデータでの再調整なしにジェイルブレイク軌跡に適用された。

主要な結果

• ジェイルブレイクに対するプローブの失敗: 最終トークンプローブは、クリーンな有害プロンプトにおいて高い再現率（約 95%）を維持したが、ジェイルブレイクでは大幅に低下した（モデルに応じて約 63% から約 87% の範囲）。逆に、無害で安全性に隣接するプロンプトに対しては、23% から 36% の非 trivial な偽陽性を生じた。
• 容量はボトルネックではない: プローブのボトルネック幅を 64 から 1024 に増加させても、ジェイルブレイクの検出は確実に向上しなかった。OLMo3 の場合、パフォーマンスは実際には低下した。これは、失敗が（信号がどこにあるかという）幾何学的な問題であり、表現容量の欠如ではないことを示唆している。
• 幾何学的な不一致: 見逃されたジェイルブレイクと無害プロンプトを分離する方向（$d_{miss}$）は、クリーンな訓練対比（$d_{harm}$）と比較して、最も低いプローブ可視エネルギー（約 2.4% から 3.5%）を持っていた。プローブの行空間は、見逃されたジェイルブレイクと無害入力を区別する特定の機能に対して事実上直交している。
• 位置による失敗: 見逃されたジェイルブレイクにおける安全でない証拠は、中間トークン位置（特に埋め込まれた有害なリクエストの近く）でしばしば可視であるが、最終トークンまでに消えたり「崩壊」したりする。
• 単純なプーリングの失敗: すべてのトークン位置にわたって最大スコアを取得するだけでは viable な解決策ではない。これは XSTest の無害プロンプトの 100% でトリガーされる（偽陽性）ためであり、高い中間スコア自体が意味的に信頼できるものではないことを示している。
• 軌跡による回復: クリーンなデータのみで訓練された PCA-HMM 診断モデルは、最終トークンで見逃されたものの大部分を成功裡に回復した（回復率 97% から 98%）一方で、XSTest における偽陽性率は単純なプーリングと比較して著しく低い（6% から 15%）を維持した。見逃されたジェイルブレイク方向は、最終トークンプローブの行空間よりも、軌跡 PCA 部分空間において著しく高いエネルギーを示した。

意義と主張
本論文は、展開された防御メカニズムではなく、診断研究として位置づけられている。その主な貢献は、最終トークンプローブが失敗する「理由」を特定することである。すなわち、それらはクリーンな訓練分布には敏感だが、ジェイルブレイクラッパーによって引き起こされる特定の幾何学的シフトには盲目である、「ショートカットのような」読み出しに依存しているということである。

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

1. 最終トークンプローブは、ラッパーによって引き起こされるシフトが表現をプローブの敏感な部分空間から外れるため、ラップされた有害なリクエストの検出には不十分である。
2. 安全に関連する情報は、最終トークンの前に隠れ状態軌跡に存在することが多いが、最終読み出しが行われる時点では失われている。
3. クリーンなデータのみで訓練された単純な軌跡認識モデル（PCA-HMM など）は、最終トークンプローブに対する効果的な診断的補完として機能しうる。それらは、プロンプト全体にわたる安全シグナルの進化を分析することで見逃された検出を回復し、単純なトークンプーリングの壊滅的な偽陽性率なしに、プロンプトの内部状態に対するより堅牢な視点を提供する。

本研究は、現在のプローブベースの安全システムの限界をよりよく理解し、緩和するために、安全診断を単一点の読み出しから軌跡認識分析へと移行することを促している。