コンフォーマル・シンキング：計算予算における推論のリスク制御
推論大規模言語モデル（LLM）はテスト時のスケーリングにより、トークン予算が増えるにつれてデータセットレベルの精度が向上するため、適応的推論（信頼性を向上させる場合のみトークンを使い、追加計算が役に立たない場合は早期停止）が動機づけられる。しかし、トークン予算や適応的推論の閾値の設定は、根本的なリスクと精度のトレードオフを伴う実用的な課題である。本論文では、予算設定問題をリスク制御として再定義し、誤り率を制限しながら計算量を最小化する。このフレームワークは、モデルが確信を持ったときに停止する上限閾値（誤った出力のリスク）と、解けないインスタンスを事前に停止する新しいパラメトリック下限閾値（早期停止のリスク）を導入する。目標リスクと検証セットが与えられれば、分布フリーのリスク制御を用いてこれらの停止メカニズムを最適に指定する。多様な推論タスクとモデルにおける実験結果は、ユーザー指定のリスク目標を満たしつつ、計算効率の向上を示している。
推論大規模言語モデル（LLM）は、テスト時にトークン予算を増やすことで精度が向上することが知られており、これにより適応的推論の研究が進められている。適応的推論では、推論が信頼性を向上させる場合にのみトークンを消費し、そうでない場合は早期に停止する。しかし、適応的推論のトークン予算や停止閾値の設定は、リスクと精度のトレードオフを伴う難しい問題である。ジョンズ・ホプキンス大学とAppleの研究チームは、この問題をリスク制御として捉え直す「コンフォーマル・シンキング」フレームワークを提案した。このフレームワークは、ユーザーが指定した誤り率の範囲内で計算量を最小化することを目的とする。
具体的には、二つの停止メカニズムを導入する。上限閾値は、モデルが現在の推論結果に高い確信を持った時点で推論を停止し、過剰な計算による誤った出力のリスクを回避する。下限閾値はパラメトリックに制御される新しいメカニズムで、モデルが問題を解けないと判断した場合に早期に停止し、無駄な計算を省く。これらの閾値は、分布フリーのリスク制御手法を用いて、目標リスクと検証セットから最適に設定される。分布フリーのため、データ分布に関する仮定を必要とせず、幅広いタスクに適用可能である。
研究チームは、数学的推論や常識推論など複数のタスクとモデルで実験を行った。結果は、コンフォーマル・シンキングのフレームワークが誤り率を効果的に制御しつつ、計算効率を大幅に向上させることを示した。特に下限閾値とアンサンブル停止機構が効率向上に寄与している。この研究は、LLM推論における計算予算管理に原理的なアプローチを提供し、リソースが限られた実環境での応用が期待される。コードはGitHubで公開されている。