技術要約：事後的な温度スケーリングを超えて：LLMのキャリブレーションのための二段階最適化

問題提起

大規模言語モデル（LLM）、特に人間からのフィードバックによる強化学習（RLHF）や直接選好最適化（DPO）を通じてアライメントされたモデルは、しばしば不適切なキャリブレーション（較正）を示し、過信（オーバーコンフィデンス）として現れることが頻繁にあります。信頼性はしばしば信頼性の代用指標として用いられますが、高リスクな領域（医療、法律など）における不適切なキャリブレーションは、誤った信頼や意思決定を招く可能性があります。

既存の解決策は、主に事後的な温度スケーリング（Temperature Scaling: TS）に依存しています。これは、学習後の検証セットを用いて単一のスカラー温度を適合させ、ロジットを再スケーリングする手法です。しかし、本論文はTSにはデータセット依存性という決定的な限界があることを指摘しています。あるドメインに対して最適化された温度は、他のドメインには通用しないことが多く、多様なタスクにわたって良好にキャリブレーションできる「普遍的な」温度を見つけることは困難です。さらに、事後的な手法はモデルパラメータを固定したままにするため、よりキャリブレーションに適したモデルを本質的に学習する機会を逃しています。

手法：CALM（二段階最適化による大規模モデルのキャリブレーション）

著者らは、キャリブレーションを二段階最適化問題として定式化する、学習時のフレームワークであるCALMを提案しています。事後的なパラメータを調整する代わりに、CALMは予測エントロピーを最大化しつつタスク性能を維持するように、モデルパラメータと「温度のような」損失ハイパーパラメータを共同で最適化します。

1. 二段階の定式化

問題は以下のように構成されています：

• 下位レベル（タスク最適化）： ハイパーパラメータ $\alpha$ が与えられたとき、訓練セット $S_T$ 上の標準的なパラメトリック損失（クロスエントロピー）を最小化するようにモデルパラメータ $\theta$ を学習します。
  $$\theta(\alpha) \in \arg\min_{\theta'} \mathcal{L}_{tr}(\theta'; \alpha)$$
• 上位レベル（キャリブレーション最適化）： $\theta$ が下位レベルのタスクに対して最適であり続けるという制約の下で、保持されたキャリブレーションセット $S_V$ 上の予測分布のエントロピーを最大化するように、ハイパーパラメータ $\alpha$ を最適化します。
  $$\min_{\alpha} \mathcal{L}_{ent}(\alpha, \theta(\alpha))$$

ハイパーパラメータ $\alpha$ は温度スケーリングの概念を一般化したものです。単一のスカラー $\tau$ の代わりに、CALMは高次元のパラメータ $\alpha = (l, \tau)$ を使用します。ここで $l \in \mathbb{R}^K$ は加法的なロジットシフトを導入し、$\tau \in \mathbb{R}^K_{>0}$ はクラス依存の温度を提供します。これにより、グローバルな再スケーリングではなく、語彙ごとの調整が可能になります。

2. 目的関数

• 上位レベルの目的関数： 過度に集中した（過信的な）予測を抑制するために、シャノン・エントロピー $H(f_{\alpha, \theta}(\cdot | x))$ を最大化します。これは、モデルを過信から過小信へと移動させる制御可能な方向を提供し、信頼度が精度と一致する点を見つけ出す可能性があります。
• 下位レベルの制約： 訓練データにおけるクロスエントロピーを最小化することで、識別能力を維持することを保証します。

3. 効率的な一次近似

厳密な二段階最適化には、ヘッセ行列の逆行列を含むハイパーグラディエントの計算が必要であり、数十億のパラメータを持つLLMにとっては計算コストが極めて高くなります。これに対処するため、CALMはBOME (Bilevel Optimization Made Easy) に着想を得た一次近似を採用しています。

• 問題を、結合変数 $u = (\alpha, \theta)$ に関する制約付き最適化として再定式化します。
• 現在の $\theta$ から $T$ 回の内部勾配ステップを実行することで、下位レベルの価値関数を近似的な代理 $\tilde{\theta}$ を用いて近似します。
• 「プラグイン残差」制約を構築し、明示的なヘッセ行列の逆行列計算を回避して、一次微分のみを用いて結合更新方向を計算します。

主な貢献

1. 学習時の視点： 本論文は、キャリブレーションのパラダイムを、データセット固有の事後的な調整から、学習時の最適化へと転換し、様々なデータセットに対してキャリブレーションの調整を許容するモデルを学習することを目指しています。
2. 二段階の損失設計： キャリブレーション指向の損失設計を、上位レベルでエントロピーベースの目的関数を用いて温度のようなハイパーパラメータを最適化し、下位レベルでモデルを訓練する二段階問題として定式化しました。
3. エントロピー最大化： 追加のハードラベルや補助ヘッドを必要とせずに、アライメントされたLLMの過信を軽減するために、エントロピー最大化を上位レベルの目的関数として提案しています。
4. スケーラブルなフレームワーク： 高価な二次計算を回避する効率的な一次近似による二段階最適化スキーム（CALM）を導入し、大規模なLLMのファインチューニングにおける実用性を実現しました。

実験結果

本手法は、4つのアライメント済みLLM（Llama-3.1-Tulu-8B, Vicuna-7B, OLMo-2-7B, Mistral-7B）を用い、多肢選択式QA（MCQA）およびオープンエンド生成QAにおいて評価されました。

ドメイン外（OOD）汎化

• 優れたキャリブレーション： CALMは、OOD設定において4つのモデルのうち3つで最高の信頼度期待キャリブレーション誤差（conf-ECE）を達成し、ベースラインと比較して誤差を大幅に減少させました。例えば、Llama-3.1において、conf-ECEを0.1784から0.1050に低減しました。
• 精度の維持： 一部の切除実験（例：「Iterate」や、二段階制約のない「Regularization」）では壊滅的な精度の崩壊が見られましたが、CALMは高いタスク精度を維持しながらキャリブレーションを向上させました。
• 堅牢性： オープンエンド生成QA（PopQAからTriviaQAへのクロスドメイン転移）において、CALMは最も低い平均セマンティックECE（0.0671）を達成し、未校正のベースラインや強力な学習時ベースラインであるCFT（Calibration-Aware Fine-Tuning）を上回りました。

ドメイン内（ID）性能

• 訓練分布と評価分布が一致するID設定では、事後的な手法として予想通り、温度スケーリングがECEにおいて最良の結果を示しました。しかし、CALMは、4つのモデルすべてにおいて、ベースラインと比較して精度とキャリブレーションの両方を一貫して向上させた唯一の学習ベースの手法でした。

アブレーション研究

• 二段階構造： 二段階の制約（Regularization）または結合項（Iterate）を取り除くと、キャリブレーションが悪化するか、大幅な精度低下を招きました。これは、タスク性能とキャリブレーションの間の結合が不可欠であることを示しています。
• パラメータ化： 語彙ごとのベクトルパラメータ（$d_y, l_y$）を使用することは、収束したスカラーパラメータ（学習されたTS）よりもはるかに優れており、回答の選択肢間の非対称な誤校正を扱うために微細な調整が必要であることを浮き彫りにしました。

意義と主張

本論文は、CALMがアライメントされたLLMの過信問題に対する堅牢な学習時ソリューションを提供すると主張しています。その主な意義は以下の通りです：

• 汎用性： 事後的なTSとは異なり、CALMはドメインを越えて汎用するキャリブレーションメカニズムを学習し、現在の手法における「データセット依存性」に対処します。
• トレードオフ管理： 単純な正則化によって生じるユーティリティの崩壊を回避し、優れた精度とキャリブレーションのトレードオフを実現します。
• 実用性： 一次近似を利用することで、二段階最適化を現代のLLMにおいて実行可能にし、「事後処理によって校正される」モデルではなく、「設計段階から校正された」モデルへの道を開きます。

著者らは、本手法がドメイン外の設定において特に有利であり、代替手法が不安定になったり大幅なユーティリティ損失を招いたりする中で、多様で現実世界のシナリオに信頼できるLLMをデプロイするための実行可能なアプローチであると結論付けています。