Test-Time Safety Alignment

Baturay Saglam と Dionysis Kalogerias による論文「Test-Time Safety Alignment」の詳細な技術的サマリーを以下に示します。

1. 問題定義

大規模言語モデル（LLM）は、有害なリクエストを拒否するために安全性トレーニング（例：RLHF、DPO）を受けます。しかし、敵対的プロンプト（ジャイルブレイク）は、モデルのトレーニング分布と無制限の入力空間との間の隙間を悪用することで、これらの防御を回避できます。これらのプロンプトは、有害なコンテンツを引き出すために、役割演技や隠蔽などの構造的な枠組みを頻繁に利用します。

既存のテスト時防御策は構造的な限界に直面しています：

• プロンプトレベルの戦略は離散トークンを操作しますが、意味論的攻撃に対して脆弱です。
• 活性化操作は隠れ状態へのアクセスを必要とし、しばしば汎化しない可能性のある固定された較正データセットに依存します。
• ガードモデルは個別の推論を必要とし、生成を誘導するのではなく、「検出後にフィルタリング」というパラダイムに従うことが一般的です。

核心的な問い： 入力単語埋め込みを**語彙未満（sub-lexical）**の方式で最適化（可視トークンを変更せずに連続ベクトル表現を修正する）することで、モデルが「拒否か従うか」の二峰性分布を生み出す場合でも、整列済みモデルを安全な出力へと誘導することは可能でしょうか？

2. 手法：テスト時安全性整列（TSA）

提案された手法TSAは、モデルの重みを変更することなく、追加のトレーニングを必要とせず、LLM をブラックボックスとして扱うことで、完全にテスト時に動作します。これはプロンプトの埋め込み行列を直接最適化します。

A. 目的関数

目標は、$y = f(X)$ で生成されるテキストの有害性を最小化することです。ここで、$f$ は LLM、$X$ はプロンプトの埋め込み行列です。

• 有害性オラクル（$h$）： 本手法はブラックボックスオラクルとしてOpenAI Moderation APIを使用します。これは、暴力、自傷行為、ヘイトなど 13 の有害カテゴリーにわたって出力をスコアリングします。
• 損失関数： $\Phi(X) = \max_{c} s_c(f(X))$ です。ここで、$s_c$ はカテゴリー $c$ のスコアです。目的は、この最大スコアを最小化し、すべての次元において応答が安全であることを保証することです。

B. 零次勾配推定

LLM（$f$）と Moderation API（$h$）はブラックボックスであるため、逆伝播は不可能です。TSA は零次勾配推定を使用します：

1. 摂動： 各反復 $k$ において、現在の埋め込み行列 $X_k$ に、$\mu$ でスケーリングされたガウスノイズ $U$ を加えて摂動させます：$X_k + \mu U$。
2. 有限差分： 勾配は、摂動入力とベース入力の損失を比較することで推定されます：
   $$g_k \approx \frac{1}{N} \sum_{i=1}^{N} \frac{\Phi(X_k + \mu U_i) - \Phi(X_k)}{\mu} U_i$$
   ここで、$N$ はモンテカルロサンプリングの数です。
3. 更新： 埋め込みは勾配降下法を通じて更新されます：$X_{k+1} \leftarrow X_k - \eta g_k$。

C. 正則化

最適化された埋め込みが有効であり、潜在空間の未トレーニング領域に逸脱しないことを保証するために：

• 勾配正規化： 更新ステップでは、ステップの大きさを制御するために正規化された勾配方向（$g_k / \|g_k\|_2$）を使用します。
• コサイン類似度制約： 類似度が閾値 $\kappa$ 未満に低下した場合、更新された埋め込みは元の埋め込み $X_0$ 周りのコサインボール上に投影されます。これにより、摂動が**語彙未満（sub-lexical）**であることが保証されます。つまり、最適化されたベクトルは元のトークンに復号されますが、そのトークンのボロノイ細胞内の連続表現は変更されます。

3. 主要な貢献

1. 整列済みモデルへの拡張： 以前はベースモデルの解毒に示されていた埋め込みレベルの制御が、本質的に異なる（二峰性の）出力分布を持つ安全性整列済みの指示調整モデルにおいても効果的に機能することを示しました。
2. 語彙未満の制御： トークン自体を変更せずに、そのトークンの最近傍セル内にある連続埋め込みベクトルを最適化することが、モデルの挙動を有害から安全へと転換するのに十分であることを証明しました。これは、埋め込みを単なる離散的なルックアップと見なす見方に挑戦するものです。
3. ブラックボックス最適化： この手法は、モデル内部（勾配、隠れ状態、またはログオッズ）へのアクセスを一切必要としません。入力埋め込みとブラックボックスのモデレーション API へのアクセスのみが必要です。
4. 無条件適用： この手法は、意図に関わらずすべてのプロンプトに適用されます。拒否を強制するのではなく、可能であれば安全で実質的な応答を生成し、必要であれば安全な拒否を生成するようにモデルを誘導します。

4. 実験結果

この手法は、2 つのベンチマーク（WildJailbreak と HarmBench）において、5 つの指示調整モデル（1B から 20B パラメータ：Gemma 3、Phi-3.5、Llama 3.1、Qwen3、GPT-OSS）で評価されました。

• WildJailbreak： 敵対的プロンプト（有害および良性）。
• HarmBench： 直接的な有害クエリ。

主要な発見：

• 安全性の排除： TSA は、標準ベンチマークにおいて、ほとんどのモデルでフラグ付き（有害）の応答数をゼロに削減しました。例えば、Llama 3.1-8B では、フラグ付きの有害プロンプトを 218 から 0 に削減しました。
• 収束： 最適化は通常、1〜2 回の勾配ステップ内で収束します。
• 良性の維持： 敵対的良性プロンプトにおいて、TSA はモデレーションスコアを増加させませんでした。これは、包括的な拒否バイアスを誘発したり、有用性を低下させたりしないことを確認しました。
• ベースラインとの比較： TSA は、SmoothLLM（ランダムトークン摂動）、AdaSteer（活性化操作）、RESTA（埋め込みノイズ）などのベースラインを大幅に上回りました。ベースラインは、意味論的ジャイルブレイクで失敗したり、直接的な有害クエリでフラグ付き率を増加させたりすることがよくありました。
• 効率性： 主な遅延コストは、モデル推論ではなく、モデレーションオラクルへの API 呼び出しです。

5. 意義と示唆

• 新たな制御面： この論文は、埋め込み層が以前考えられていたよりも豊かな制御面を提供することを明らかにしています。小さな語彙未満の摂動は、整列済みモデルの複雑な安全性境界をナビゲートできます。
• 展開の実用性： 「ドロップイン」型の安全性レイヤーとして、TSA はモデルの再トレーニングや微調整を必要としません。入力埋め込みへのアクセスが可能であれば（例：オープンウェイトモデル）、パブリックなモデレーション API と組み合わせて、任意のモデルに展開できます。
• 限界：
  • 遅延： 反復最適化はプロンプトごとの遅延（秒単位）を追加しますが、これは主に API の往復に支配されています。
  • モデル依存性： 効果はモデルの既存の安全性トレーニングの品質に比例します。薄い整列（例：Gemma 3-1B）を持つモデルでは、オプティマイザーが近くの安全な方向を見つけられなかったため、成功率が低下しました。
  • API アクセス： 入力埋め込みへのアクセス（純粋に API によって閉じられたモデルでは利用不可）とモデレーション API へのアクセスが必要です。

結論として、Test-Time Safety Alignmentは、連続入力表現を最適化することで敵対的攻撃を無力化する、堅牢でモデルに依存しないメカニズムを提供し、再トレーニングなしで理論的な安全性と実用的な展開の間のギャップを効果的に埋めます。