技術概要：確率的較正は言語モデルにおいて学習可能な能力である

問題定義

言語モデル（LM）は、多様な出力の生成、選択確率のバランス調整、または特定の数学的分布からのサンプリングなど、ユーザーが指定するランダム性を必要とするシナリオにおいて、ますます展開されています。しかし、標準的な自己回帰型 LM は、生来の確率的生成器としては不十分です。それらが誘発する出力分布は、しばしば狭い部分集合に収束し、サポートの制約に違反し、要求された目標分布と整合しません。既存のプロンプトによる回避策（推論やシードの注入など）はこれらの問題を部分的に緩和できますが、推論時のコスト、プロンプトへの感度、計算オーバーヘッドを導入するだけであり、モデルのサンプリング行動を根本的に変えるものではありません。

ここで扱われる中心的な問いは、確率的較正（既知の目標分布を反映するサンプルを生成する LM の能力）が、推論時のヒューリスティックに依存するのではなく、微調整を通じて直接学習可能かどうかです。

手法：較正微調整

著者らは、数学的分布からのサンプリングを明示的に要求する合成プロンプト（例：「[分布] 分布のパラメータ [params] から、ちょうど 1 つのランダムな数を生成せよ」）を使用して、較正された確率的サンプリャとして LM を訓練するフレームワークを提案します。

2 つの異なる微調整戦略が、凍結されたベースモデル上の低ランク適応（LoRA）を利用することで比較されます。

1. ソフトターゲット微調整：

   • メカニズム：目標分布は、有限の標準的な出力空間に離散化されます。すべての有効な数値出力からプレフィックストライが構築され、各ノードはトークンプレフィックスを表します。
   • 監督：任意の与えられたプレフィックスに対して、トライ内のすべての有効な継続の残りの確率質量を正規化することで、ターゲット次トークン分布が計算されます。
   • 目的：モデルは、訪問されたすべてのプレフィックスにおいて、予測された次トークン分布とトライから導出されたターゲット分布との間のカルバック・ライブラー（KL）発散を最小化します。これにより、プレフィックスレベルの密な監督が提供されます。

2. ハードターゲット微調整：

   • メカニズム：モデルは、目標分布から直接サンプリングされた完了文で訓練されます。
   • 監督：各プロンプトに対して、ターゲット分布から数値回答がサンプリングされ、トークン化され、真の完了文として使用されます。
   • 目的：モデルは、サンプリングされた完了文における標準的なマスク付き自己回帰クロスエントロピーを最小化します。これにより、スパースなパスレベルの監督が提供され、同等のカバレッジを達成するためにより多くの訓練ステップ（プロンプトあたりのより多くのサンプル）が必要です。

実験設定：

• モデル：0.6B から 27B パラメータの範囲にある、4 つのファミリー（Qwen3、Gemma-3、Llama-3.2、GPT-OSS）にわたる 12 のモデル。
• データ：30 の分布ファミリー（24 は訓練用、6 は分布外（OOD）テスト用に保持）。
• 評価：
  • 構造化サンプリング：目標分布に対する Wasserstein-1 距離（W1）および logit レベルの KL 発散を介して測定。
  • 転送ベンチマーク：オープンエンドのランダム生成（多様性サポート）、MCQ 回答位置のバランス調整、および NoveltyBench（意味的多様性）。
  • 保持：TinyBenchmarks（MMLU、GSM8K など）における能力の保持、および PALOMA における言語モデル適合度。

主要な貢献

1. 較正の学習可能性：本論文は、確率的較正が学習可能な能力であることを実証しています。ソフトターゲットとハードターゲットの両方の微調整は、保持された分布ファミリーと未見のパラメータ設定における構造化サンプリング忠実度を大幅に向上させます。
2. 戦略の比較分析：
   • ハードターゲット：一般的に、特に未見のパラメータにおいて、ドメイン内構造化数値サンプリングベンチマークで優れた性能を達成します。正確な数値忠実度において優れています。
   • ソフトターゲット：数値タスクでは競争力がありますが、より広範な確率的生成設定への転送において優れています。オープンエンド生成において確率的サポートを一貫して拡大し、MCQ における回答位置のバランスを改善し、NoveltyBench においてより高い有用性を生み出します。
3. トレードオフとコスト：本研究は、較正の向上がしばしばトレードオフを伴うことを浮き彫りにしています。具体的には、算術推論（GSM8K）は低下する傾向にあり、そのコストはモデルのサイズと微調整バリアントによって異なります。しかし、言語モデル適合度（PALOMA パープレキシティ）はしばしば保持または向上しており、モデルが単にトークン確率を平坦化するのではなく、特定の確率的行動を学習していることを示唆しています。

結果

• 構造化サンプリング：両手法とも、ベースチェックポイントと比較して正規化された W1 および logit KL を大幅に削減します。ハードターゲット微調整は、多くの場合、数値ベンチマークで最も強力なパフォーマンスを示し、エラーを 1 桁から 2 桁削減します。
• オープンエンド生成：ソフトターゲット微調整は、すべてのモデルにおいてトップ 90% サポートサイズ（確率質量の 90% をカバーするために必要なトークンの数）を一貫して増加させ、しばしば 1 から 2 桁増加させます。これにより、都市、動物、名前などのカテゴリにおいて、より多様で集中度の低い出力が生まれます。
• MCQ と Novelty：ソフトターゲット構成は、一般的に MCQ 生成における回答位置のバランスを改善し、NoveltyBench においてより高い有用性スコアをもたらします。これは、学習された確率的行動が暗黙の自然言語制約にも転移することを示しています。
• 保持：一般的能力への影響は混合的です。一部の中型から大型モデルは TinyBenchmarks 上の集計スコアを維持または改善しますが、小型モデルや特定のタスク（特に GSM8K）は劣化を示します。これは、訓練目的（推論痕跡なしの短く直接的な完了）が推論集約型タスクと衝突する可能性があることを示唆しています。

意義と主張

本論文は、確率的較正が言語モデルの固有の限界ではなく、明示的に訓練可能な能力であると主張しています。この問題を、有効な出力空間と、その上の確率質量の配分の両方を学習する問題として捉えることで、著者らは LM を制御された確率的サンプリャへと変換できることを示しています。

その意義は、推論時のプロンプティングを超え、モデルの分布行動を根本的に変える訓練時の解決策へと移行することにあります。著者らは、向上が選択的であり、すべての下流タスクを普遍的に改善するわけではないこと（算術推論が顕著な例外である）を指摘しつつ、特定の確率的行動のためにモデルを調整する能力が、多様性、公平性、または統計的制約への厳密な準拠を必要とする応用に対する新たな道を開くと述べています。この研究は、将来の研究がこれらの目的によって学習されたメカニズムに焦点を当て、能力の劣化を緩和するための保持を考慮したバリアントを開発すべきであることを示唆しています。