Functional-level Uncertainty Quantification for Calibrated Fine-tuning on LLMs

技術的概要：LLM における較正済みファインチューニングのための機能レベル不確実性定量化

問題定義
不確実性の定量化（UQ）の精度は、医療や科学的推論といった安全性が重要な分野における大規模言語モデル（LLM）の信頼性の高い展開にとって不可欠である。しかし、ファインチューニングされた LLM は、限られたデータでタスクに適応させられる際、頻繁に過剰な自信を示し、出力の真の信頼性を反映できずにいる。パラメータ効率型ファインチューニング（PEFT）に対する既存の UQ 手法、例えば低ランク適応（LoRA）に基づく手法の多くは、事後的（post hoc）である。これらは、モデルが訓練された後にパラメータ空間を分析する（アンサンブルやベイズ近似を介するなど）か、プロンプトを攪拌することによって不確実性を推定する。これらのアプローチは、過剰な自信の根本原因に対処することに失敗していることが多い。すなわち、適応プロセスそのものの間に、ファインチューニングによって引き起こされる不確実性や汎化ギャップをモデルが捉えられないという点である。さらに、複数のフォワードパスや広範なサンプリングに依存する手法は、計算オーバーヘッドを著しく増大させ、スケーラビリティを制限する。

手法：UQ4CT
著者は、パラメータ空間から機能空間への不確実性への焦点の移行を行うフレームワークである、較正済みファインチューニングのための機能レベル不確実性定量化（UQ4CT）を提案する。学習済みパラメータを訓練後に分析するのではなく、UQ4CT は、LoRA 専門家のプロンプト依存混合によって生み出される機能空間上の不確実性を、ファインチューニングプロセスの最中に較正する。

この手法は、以下の 3 つの中核的構成要素からなる：

1. MixLoRA による階層的機能分解：
   層全体にわたるすべての可能な専門家組み合わせをモデル化することによる組み合わせ爆発を回避するため、UQ4CT は階層的分解を採用する。これは、各層において、スパースなルーターが入力隠れ状態に基づいて LoRA 専門家のサブセットを動的に選択する、専門家混合（MoE）アーキテクチャを利用する。これにより、モデルの出力が入力に条件付けられた多様な基底関数（LoRA 専門家）の重み付き和となる、柔軟な機能空間が構築される。この機能空間は、層ごとの混合の再帰的構成として表現され、モデルが計算パスを適応的に構成することを可能にする。

2. 機能レベルの信頼度（FLC）推定：
   モデルは、専門家選択の集中度を分析することで不確実性を定量化する。具体的には、機能レベルの信頼度（FLC）は、各層においてルーターによって選択されたトップ-$k$ 専門家のルーター重みから導き出される。ルーター重みにおける高い集中度（すなわち、ルーターが特定の専門家を選択することに自信を持っていること）は、高い機能レベルの信頼度に対応し、分散した重みはより高い不確実性を示す。これにより、複数のサンプリングパスを必要とせず、モデルの確信度に関する原理的な内部シグナルが得られる。

3. 較正損失（$L_{cal}$）：
   中核的な革新は、訓練中に機能レベルの信頼度を予測の正しさと整合させる新しい較正損失である。この損失は以下のように定義される：
   $$L_{cal} = (1\{\text{MixLoRA}(x) = y^*\} - \text{FLC}(x))^2$$
   ここで、$1\{\cdot\}$ は正しい予測に対する指示関数である。この目的関数は、FLC がモデルの予測が正しい確率に近似することを強制する。標準的な交差エントロピー損失や負荷分散損失と併せてこの損失を最小化することで、モデルは、正しい可能性が高い場合のみ高い信頼度を割り当て、そうでない場合はより高い不確実性を割り当てることを学習する。これにより、ルーターは不確実性が高い場合に代替的な専門家混合を探求するよう促され、汎化性が向上する。

主要な貢献

• 機能レベルの視点： 本論文は、パラメータ空間から機能空間への不確実性定量化への転換を導入し、ファインチューニング中にモデルが実現し得る入力 - 出力マッピングに対するモデルの信頼度を明示的に較正する。
• 較正意識型ファインチューニング： プロンプト依存の LoRA 混合重みと予測の正しさを動的に整合させる新しい較正損失を提案し、精度を低下させることなく過剰な自信を軽減する。
• 効率的なアーキテクチャ： この手法は、最小限の計算オーバーヘッドで多様な機能関係を捉える階層的 MixLoRA アーキテクチャを活用し、推論時には単一のフォワードパスのみを必要とする。
• 理論的保証： 著者は、データ分布上の較正リスクを最適化することが、予測の正しさの真の確率に対応する最適な FLC を生み出すことを証明する。

実験結果
著者は、Llama-3.1-8B および Mistral-7B を用いて、4 つの多肢選択ベンチマーク（ARC-E, ARC-C, OBQA, ClimateQA）と 2 つの自由回答生成 QA タスク（TruthfulQA, TriviaQA）において UQ4CT を評価した。

• 分布内性能： UQ4CT は、強力なベースライン（LoRA、MC ドロップアウト、ディープアンサンブル、ベイズ LoRA など）と比較して、ベンチマーク全体で期待較正誤差（ECE）を**25%**以上削減し、かつ競争力のある、あるいは優れた精度を維持した。注目すべきは、UQ4CT が単一のフォワードパスでこれらの結果を達成したのに対し、BLoB（N=10）のような手法は 10 回のパスを必要としたにもかかわらず、より高い ECE を生み出したことである。
• 分布シフト： 分布シフト下（CS、工学、法、医療の MMLU サブタスクでテスト）において、UQ4CT は優れた較正性と競争力のある精度を維持し、ベースラインと比較して分布外入力に対する堅牢性の向上を示した。
• 生成タスク： 自由回答生成において、UQ4CT は TruthfulQA および TriviaQA で最低の ECE を達成し、機能空間の較正という視点が分類スタイルの設定を超えて拡張可能であることを示唆した。
• 効率性： メモリ分析によると、UQ4CT のピーク GPU メモリ使用量は標準的な LoRA よりも約 7.5% 高いのみであり、大規模展開にとって計算上実行可能であることを示している。

意義と主張
本論文は、UQ4CT が現在のファインチューニング慣行における根本的な限界、すなわち適応プロセス中の不確実性への意識の欠如に対処していると主張する。機能レベルのシグナルを介して不確実性較正をファインチューニング目的関数に直接統合することで、この手法はモデルの信頼性が実際の予測信頼性をよりよく反映することを保証する。著者は、この研究を、特に過剰な自信が重大なリスクをもたらすデータ不足環境において、信頼性の高いファインチューニング済み LLM に向けた基礎的な一歩として位置づけている。また、F1 ベースの代理指標を用いて自由回答生成に拡張可能である一方で、長文の推論や多段階生成へのスケーリングは今後の課題であると控えめに指摘している。