技術要約: Best-of-Better-N (BoBN)

問題提起

Best-of-N (BoN) に代表される推論時アライメント（Inference-time alignment; ITA）手法は、モデルの重みを変更することなく、大規模言語モデル（LLM）を人間の好みや特定のタスク要件に適合させることを目的としている。これらの手法は、「サンプリングと評価」のフレームワークに基づいて動作する。すなわち、参照方策 $\pi_{ref}$ から $N$ 個の応答をサンプリングし、報酬モデル $r$ に基づいて最も高い報酬を持つものを選択する。

ITAの有効性は、根本的に以下の2つの仮定によって制約を受ける：

1. カバレッジ (A1): 参照方策 $\pi_{ref}$ が、少なくとも一つの高品質（高報酬）な応答をサンプリングできること。
2. 選択 (A2): 選択基準が、サンプルの中から高品質な応答を正しく識別できること。

現在のITA研究は、主に報酬ハッキングを軽減するために選択基準を洗練させることで (A2) を改善することに焦лоんでいる。しかし、決定的なボトルネックとして (A1) が残っている。もし $\pi_{ref}$ がダウンストリームタスクに適合していない場合（例：安全性訓練を受けていないモデルが有害なプロンプトに対して拒絶を行おうとする場合）、高報酬な応答に割り当てられる確率は極めて低くなる。このような場合、いかなる選択戦略も成功し得ない。なぜなら、出力空間における高報酬領域が参照モデルにとって事実上アクセス不能だからである。既存の解決策は、多くの場合、ファインチューニング（コストがかかり、重みへのアクセスが必要）を必要とするか、あるいはタスク固有の関連性に欠ける可能性のある静的なインコンテキスト例を使用している。

手法: Best-of-Better-N (BoBN)

著者らは、重みのファインチューニングや重みへのアクセスなしに、$\pi_{ref}$ のサンプリング分布を高報酬領域へとシフトさせるための推論時フレームワークとして Best-of-Better-N (BoBN) を提案する。BoBNは、インコンテキスト学習（ICL）を、検索およびリスタイリング（再構成）のパイプラインと組み合わせたものである。

BoBNのワークフローは、以下の3つのステージで構成される：

1. 検索 (Retrieval): 入力プロンプト $x$ に対して、埋め込みモデルを用いてデータセット $D$ から最も類似したプロンプトと応答のペア $(x_i, y_i)$ を上位 $K$ 個検索する。これらの例は、特定のクエリに関連する高報酬行動のデモンストレーションとして機能する。
2. リスタイリング (Restyling): 検索された例がターゲットタスクの特定の制約、フォーマット、およびペルソナに適合することを保証するため、参照モデル $\pi_{ref}$ を用いて、検索された応答 $y_i$ を「リスタイル」する。具体的には、$\pi_{ref}$ は「$y_i$ をスタイル $T$ に合わせて書き換えてください」といったプロンプトに基づき、$y_i$ を $\hat{y}_i$ へと書き換える。このステップにより、検索されたデータの整合性を保ちつつ、コンテキストと参照モデル間のトークンレベルでの分布的一貫性を確保する。
3. 生成と選択 (Generation and Selection): リスタイルされた例 $C_K = \{(x_i, \hat{y}_i)\}$ を、インコンテキスト例として入力プロンプト $x$ の前方に付加する。その後、参照モデルは条件付き分布 $\pi_{ref}(\cdot | C_K, x)$ から $N$ 個の応答をサンプリングする。最後に、選択基準（例：BoN）に基づき、最良の応答を選択する。

$C_K$ によって条件付けることで、BoBNは実質的に $\pi_{ref}$ の確率質量を高報酬領域へとシフトさせ、高品質な候補が最初に生成される確率を高める。

主な貢献

1. 新規フレームワーク: 検索およびリスタイルされた例を活用することで、$\pi_{ref}$ の高品質な応答に対するカバレッジを向上させる、ICLベースのアプローチであるBoBNの提案。
2. 理論的分析: 検索の質がいかにカバレッジとダウンストリームの性能を制御するかについての理論的特性付け。著者らは、静的なICLやICLなしの場合と比較して、タスクに関連する検索がカバレッジ（ゴールデン方策と参照方策の間のカイ二乗ダイバージェンスとして定義）を確実に減少させることを示している。また、このカバレッジの減少がBoNによる選択と組み合わさることで、タスク性能を向上させることを示している。
3. 実証的検証: 安全性アライメント（StrongReject, SORRY-Bench）および数学的推論（GSM8k, MATH500）のベンチマークにおける広範な実験。結果は、複数の参照LLM（zephyr-7b, tulu3-8b）において、$N=1$（個別の応答品質）の場合と $N>1$（選択との組み合わせ）の場合の両方で、BoBNが性能を向上させることを示している。

実験結果

著者らは、標準的なBoN（$\pi_{ref}$ を直接使用）およびURIAL（静的な例を用いたバニラなICL）と比較してBoBNを評価した。

• 安全性アライメント: StrongRejectベンチマークにおいて、BoBNはベースラインを大幅に上回った。zephyr-7bモデルにおいて、$N=16$ のBoBNは0.033というスコア（ほぼ最適）を達成し、Claude Sonnet-3.7のような最先端モデルに匹敵した。一方、$\pi_{ref}$ 単体では0.425であった。また、BoBNはURIALや$\pi_{ref}$ と比較して、より優れたカバレッジ（より低い $\log C$）を示した。
• 数学的推論: GSM8kおよびMATH500において、BoBNは一貫した精度の向上をもたらした。GSM8kにおけるzephyr-7bについて、BoNを用いた$\pi_{ref}$ の精度が43.9%であったのに対し、BoBNは$N=16$ で58.7%の精度を達成した。特筆すべきは、$N=4$ のBoBNが、しばしば$N=16$ のベースラインと同等またはそれを上回る性能を示したことであり、これはサンプリング効率が高いことを示唆している。
• コストとレイテンシ: BoBNは生成前のリスタイリングステップ（生成前に $\pi_{ref}$ への $K$ 回の呼び出しを必要とする）により追加のレイテンシが発生するが、大幅なトークン効率を実現している。$N=4$ のBoBNは、同等またはそれ以上の性能を達成しながら、URIALの$N=16$ と比較してトークン使用量を最大50%削減した。
• アブレーション研究:
  • リスタイリング: リスタイリングステップを除去すると、安全性アライメントが弱まり、カバレッジが高くなった（性能が悪化した）。これは、検索された例をターゲットタスクのスタイルに適合させる必要性を裏付けている。
  • 検索数 ($K$): $K$ が1から8に増加するにつれて性能は向上したが、$K=16$ になると低下した。これは、あまりに多くの例を検索しすぎると、関連性の低いデータが混入し、インコンテキスト信号を希釈してしまうことを示唆している。

意義と主張

本論文は、BoBNが推論時アライメントの根本的な限界、すなわち「アライメントされていない参照モデルが高報酬な応答をサンプリングできないこと」に対処していると主張している。検索およびリスタイルされたインコンテキスト例を通じてサンプリング分布をシフトさせることで、BoBNは以下を可能にする：

• 固定予算での高品質化: 固定されたサンプル数 $N$ において、より高いタスク性能を実現する。
• 効率性: より少ない $N$ で目標性能を達成し、推論コストを削減する。
• トレーニングフリーのアライメント: モデルの重みや尤度へのアクセスを必要とせず、ダウンストリームタスクのために明示的にファインチューニングされていないモデルのアライメントを向上させる。

著者らは、線形トランスフォーマーに基づく理論的分析が、検索の質、カバレッジの減少、およびダウンストリームのタスク性能の間の、証明可能なリンクを提供しており、単なる選択のヒューリスティックを超えて、ITAを改善するための原理的なアプローチを提供していることを強調している。