UCCI: Calibrated Uncertainty for Cost-Optimal LLM Cascade Routing

技術概要：UCCI – コスト最適化された LLM キャスケードルーティングのための較正された不確実性

問題定義
大規模言語モデル（LLM）は構造化予測タスクにおいて高い性能を発揮しますが、推論コストはモデルサイズとともに急激に増大します。キャスケード推論（簡単なクエリを小規模モデルにルーティングし、困難なクエリを大規模モデルにエスカレートさせる方式）はコスト削減を約束します。しかし、現在展開されているルーターは、プロンプトの文言やワークロードのシフトに敏感な、較正されていない信頼度スコア（生トークンエントロピーやマージンなど）に依存することが多く、その結果、閾値の調整は往々にしてワークロードごとの手作業となり、一般化に失敗します。これにより、キャスケードの理論的な可能性と本番環境での性能との間にギャップが生じています。

手法
本論文は、キャスケード設計を調整作業ではなく、較正された意思決定問題として扱うルーティングフレームワーク「UCCI（Uncertainty-Calibrated Cascade Inference：不確実性較正キャスケードア推論）」を提案します。手法は以下の 3 段階で進行します。

1. トークンレベルの不確実性抽出: 生成されたシーケンスに対し、システムはすべてのトークンにわたるトップ 1 とトップ 2 の次のトークンの確率の差（$p_{t,1} - p_{t,2}$）に基づき、スカラー不確実性スコア $u(x)$ を計算します。このマージンを集約することで、小規模モデルの出力全体の不確実性を表します。
2. アイソトニック較正: 生不確実性スコア $u(x)$ は確率ではありません。UCCI は、ホールドアウトされた較正セット上でアイソトニック回帰を用いて、このスコアを較正されたクエリごとの誤り確率 $\hat{p}(x)$ にマッピングします。このステップにより、較正されていない信号を、小規模モデルの出力が誤っている確率（$P(e(x)=1 | u(x))$）の信頼性の高い推定値に変換します。
3. 制約付きコスト最小化: 検証セット上で閾値ポリシー $\pi_\theta$ が選択されます。閾値 $\theta^*$ は、目標精度制約（$\tau$）の下で期待推論コストを最小化するように選ばれます。このポリシーは、$\hat{p}(x) \leq \theta$ のクエリを小規模モデルにルーティングし、$\hat{p}(x) > \theta$ のクエリを大規模モデルにエスカレートさせます。

理論的貢献
本論文は、以下の 2 つの主要な理論的結果を提供します。

• 閾値ポリシーの最適性: 3 つの明示的な仮定（大規模モデルのコスト > 小規模モデルのコスト；エスカレートされるクエリに関わらず大規模モデルの精度は不変である；真の較正済み誤り確率へのアクセスが可能である）の下で、著者らは較正済み誤り確率に対する閾値ポリシーがコスト最適であることを証明します。
• サンプル複雑性: 本論文は、アイソトニック回帰が $n$ 個の較正例を用いて $O(n^{-1/3})$ の期待較正誤差（ECE）を達成することを示し、較正ステップのデータ効率性に対する理論的限界を確立します。

実験設定と結果
このフレームワークは、企業向け写真管理システムから得られた 75,000 件の手動ラベル付けクエリを含む本番環境の**命名語認識（NER）**ワークロードで評価されました。

• モデル: 40 億パラメータの小規模モデル（$f_s$）と 120 億パラメータの大規模モデル（$f_\ell$）。どちらもインストラクションチューニングされ、H100 GPU 上で提供されました。
• コスト測定: コストは H100 GPU 上で測定されたエンドツーエンドのレイテンシから導出され、コスト比 $c_\ell/c_s \approx 3.02$ を得ました。
• ベースライン: 本研究では、UCCI を単一モデル推論、較正されていないエントロピー閾値法、分割適合ルーティング、および FrugalGPT 風の学習済み閾値と比較しました。

主要な知見:

• コスト削減: マイクロ F1 0.91 を目標とした場合、UCCI は大規模モデルのみを使用する場合と比較して、推論コストを**31%**削減しました（95% 信頼区間：[27%, 35%]）。
• 較正の質: アイソトニック回帰は、期待較正誤差（ECE）を較正前（0.12）から0.03に削減しました。
• ベースラインとの性能比較: 同じ動作点（F1 = 0.91）において、UCCI はエントロピー閾値法（2.31）、適合予測（2.18）、FrugalGPT 風ルーティング（2.24）と比較して、最低のコスト（2.08 の正規化単位）を達成しました。
• 仮定の検証: 大規模モデルの精度がルーティングに対して不変であるという仮定は実証的にテストされました。エスカレートされたクエリにおける精度（0.928）は、グローバルな精度（0.932）とほぼ同一であり、理論的フレームワークを支持しました。

意義と主張
本論文は、キャスケードシステムにおける主要な工学的価値は、最終的な閾値を手動で調整することではなく、ルーティングスコアを較正することにあると主張します。ノイズの多いトークンマージン信号を、小さなアイソトニックフィッティングを通じて確率的な誤り推定値に変換することで、UCCI は複雑な学習による遅延アーキテクチャや好意データなしに、ほぼ最適なルーティング性能を達成します。

著者らは、結果がグローバルな精度や名义的な API 価格に基づくシミュレーションに内在する選択バイアスを避け、実際のモデル出力と測定されたハードウェアレイテンシを用いたエンドツーエンドのルーティングに基づいていることを強調しています。フレームワークはタスク非依存ですが、著者らは特定のコスト削減（31%）と較正指標はワークロード固有であると指摘しています。また、「小規模モデルにとって難しい」クエリが「大規模モデルにとって」不釣り合いに難しい領域における不変性仮定に関する限界を認め、分布シフトを伴うストリーミング展開では静的な較正が適応を必要としうることに言及しています。