Domain-Shift-Aware Conformal Prediction for Large Language Models

技術要約：大規模言語モデルのためのドメインシフトを考慮した共形予測（DS-CP）

問題提起
大規模言語モデル（LLM）は、しばしば事実とは異なるが自信に満ちた出力（ハルシネーション）を生成することがあり、これはヘルスケアや金融といった高リスクな領域において重大なリスクをもたらします。共形予測（Conformal Prediction, CP）は、有限標本における分布フリーの被覆保証を備えた不確実性定量化（UQ）の理論的枠組みを提供しますが、標準的なCPは交換可能性（exchangeability）の仮定に依存しています。この仮定は、較正（キャリブレーション）データとテストデータが同一の分布から抽出されていることを要求します。しかし、現実世界のデプロイメントにおいては、ドメインシフトが遍在しています。ユーザー行動の変化や入力分布の変化により、テスト時のプロンプトは歴史的な較正データとは大きく異なることが頻繁にあります。このようなシフトが発生した場合、標準的なCPは頻繁に**被覆不足（under-coverage）**に陥り、真の回答を含むには狭すぎる予測セットを生成してしまい、信頼できる不確実性の保証を提供できなくなります。

手法：DS-CPフレームワーク
著者らは、ラベル付きのターゲットデータ（target data）を必要とせずに、ドメインシフト下でのLLMへのCP適用を可能にするフレームワークである**ドメインシフトを考慮した共形予測（Domain-Shift-Aware Conformal Prediction, DS-CP）**を提案しています。この手法は、高次元で非構造的なプロンプト空間において密度比を推定することの困難さを、以下のステップを通じて解決します。

1. 意味的埋め込み（Semantic Embedding）: 生のトークン空間で操作する代わりに、本フレームワークは事前学習済みの埋め込みモデル $h: \mathcal{X} \to \mathcal{Z}$ を使用して、プロンプトを低次元の意味空間へと投影します。これにより、ドメイン間の類似性を捉え、ドメインシフトを埋め込み分布のシフトとして特徴付けることが可能になります。
2. 密度比推定（Density Ratio Estimation）: 埋め込み空間において、ソース（較正）プロンプトとターゲット（テスト）プロンプトを識別するためのドメイン分類器を学習させます。この分類器は、較正サンプルがターゲットドメインにどの程度近いかに基づく重みとして機能する密度比 $\tilde{r}(z) = \frac{dP'_Z}{dP_Z}(z)$ を推定します。
3. 正則化された重み付け（Regularized Weighting）: ドメイン比が極端に不均衡な場合（大きなシフトやサンプルサイズの小ささに起因する場合）、予測セットが全出力空間へと退化してしまうリスク（予測セットが広がりすぎるリスク）を防ぐため、著者らは正則化ステップを導入しています。テストポイントに割り当てられる重みは、較正データのみに依存する正則化パラメータ $\lambda$ に置き換えられます。
   • 経験的なスコア分布は以下のように構成されます：
     $$\hat{\mu}_n = \sum_{i=1}^n \frac{\hat{w}_i}{\sum \hat{w}_j + \lambda} \delta_{S_i} + \frac{\lambda}{\sum \hat{w}_j + \lambda} \delta_{\infty}$$
     ここで、$\hat{w}_i$ は推定された較正サンプルの密度比であり、$\lambda$ は妥当性と適応性のバランスを取るために選択されます（通常は $\max_i \hat{w}_i$ とされます）。
4. 予測セットの構築: 新しいプロンプトに対する最終的な予測セットは、この重み付き経験分布の $(1-\alpha)$ 分位点によって定義されます。

理論的保証
本論文は、DS-CPに関する有限標本被覆境界を確立しています。

• 下界（Lower Bound）: 被覆確率は、ギャップ項を差し引いた $1 - \alpha$ 以上であることが保証されます。このギャップは、正規化された重みと、埋め込み空間におけるスコア分布のソースとテスト間の全変動距離に依存します。
• 滑らかさの条件（Smoothness Condition）: 非適合度スコア（nonconformity score）が意味的埋め込みに対して滑らかに変化する場合、生のプロンプト分布が異なっていても、被覆ギャップは小さく保たれることを著者らは示しています。
• 保守性（Conservativeness）: 上界も提供されており、関連する較正データの有効サンプルサイズが少ない場合、正則化項 $\lambda$ が制御された程度の保守性（過剰被覆）をもたらすことが示されています。
• 標準CPへの回帰: 交換可能性の条件（ドメインシフトがない状態）の下では、本手法は標準的なCPと正確に一致し、理想的な設定における妥当性の損失がないことを保証します。

実験結果
本フレームワークは、ドメインシフトをシミュレートするために異なる科目カテゴリを異なるドメインとして扱う MMLUベンチマークを用いて評価されました。実験は、9つのモデルファミリーにわたる16のオープンソースLLMを用いて実施されました。

• 被覆率（Coverage）: 標準的なCPは、目標とする90%の被覆率を達成できないことが頻繁にあり、較正ドメインとテストドメインが乖離すると、このレベルを大幅に下回ることがありました。DS-CPは一貫して被覆率を改善し、すべてのモデルおよびドメインペアにおいて、目標レベルに近い被覆率を維持しました。
• 集合サイズ（Set Size）: DS-CPは、シフトを補償するために標準的なCPよりもわずかに大きな予測セットを生成しましたが、その増加は緩やかでした。これは、妥当性の回復と効率性の維持との間のトレードオフが成功していることを示しています。
• 適応性（Adaptivity）: 分析により、DS-CPは標準的なCPが失敗する困難なケースにおいて被覆不足を選択的に修正する一方で、標準的なCPが既に有効である場合には最小限の調整しか行わないことが示されました。これは、単なる一律な集合の拡大ではないことを証明しています。

意義と主張
本論文は、LLMの現実世界のデプロイメントにおける信頼できる不確実性定量化への実用的な一歩を提供すると主張しています。その主な意義は以下の通りです：

1. 交換可能性を超えたCPの拡張: 交換可能性の仮定が崩れるシナリオ、すなわちLLMにおけるドメインシフトに対処するための、体系的かつ理論に基づいた手法を提供します。
2. 意味的適応: テキストそのものではなく意味的埋め込みを活用することで、高次元のLLM入力に対する密度比推定を可能にします。
3. バランスの取れた信頼性: 統計的妥当性（被覆保証の維持）と適応性（特定のターゲットドメインへの予測セットの適合）のバランスを実現し、安全性が重要なアプリケーションにおける信頼できない予測という極めて重要な課題に対処します。

著者らは、埋め込みおよび密度比推定の品質への依存、ラベルなしのターゲットプロンプトの必要性、および現在の焦点が自由記述生成ではなく多肢選択タスクに限定されていることなどの限界についても言及しています。しかし、DS-CPは、ドメインシフト下におけるLLMプロンプトの高次元かつ非構造的な性質に特化した、原理に基づいた共形予測の適応であると述べています。