論文「Reaching Beyond the Mode: RL for Distributional Reasoning in Language Models」の技術的サマリー

本論文は、大規模言語モデル（LM）が従来の「単一の最尤回答」を出力する枠組みを超え、**確率的推論（Distributional Reasoning）**を行い、複数の妥当な回答候補とそれぞれの信頼度を単一のフォワードパスで生成することを可能にする新しい強化学習手法「Multi-Answer RL」を提案するものです。

以下に、問題定義、手法、主要な貢献、実験結果、および意義について詳述します。

---

1. 背景と問題定義

従来の限界

現在の言語モデルは、強化学習（特に RLVR: Reinforcement Learning with Verifiable Rewards）を用いてトレーニングされることが多く、これは「単一の正解」を最大化するように最適化されます。

• モードへの収束（Mode Collapse）: 強化学習により、モデルは分布の最も確率の高いモード（単一の回答）に収束しやすくなります。
• 現実世界の不確実性: 医療診断、曖昧な質問応答、不完全な情報に基づく推論など、現実の多くのタスクでは「単一の正解」が存在しないか、複数の妥当な仮説が存在します。
• 推論時の非効率性: 従来のモデルが複数の回答を生成するには、推論時に複数回サンプリング（Best-of-K など）を行う必要があります。しかし、これらは計算コストが高く、モデルがトレーニング時に「単一の回答にコミットする」ように学習されているため、推論時に多様な仮説を探索する挙動とミスマッチが生じます。

目指すゴール

• 単一のフォワードパスで、複数の妥当な仮説（回答セット）を生成する。
• 各回答に対する信頼度（確率）を明示的に出力し、分布を表現する。
• 推論時のサンプリングに依存せず、モデル内部で分布推論を内包化する。

---

2. 提案手法：Multi-Answer RL

著者らは、モデルが回答の分布 $P$ を直接生成・最適化することを目的とした強化学習フレームワークを提案しました。

2.1 基本構造

モデルは、プロンプト $x$ に対して、構造化された出力として $K$ 個の異なる候補回答の集合 $A = \{a_1, \dots, a_K\}$ と、それに対応する信頼度スコアの集合 $Q = \{q_1, \dots, q_K\}$ を単一の思考連鎖（Chain-of-Thought）内で生成するようにトレーニングされます。

2.2 報酬関数の設計

提案手法は、従来の RLVR と RLCR（Calibration Rewards）を多回答設定に拡張したものです。

1. Multi-Answer RLVR (集合レベルの正解性):
   生成された回答集合 $A$ と正解集合 $Y^*$ の被覆度を評価します。
   $$R_{RLVR}^{multi}(A, Y^*) = \sum_{i=1}^{K} \mathbb{I}[a_i \in Y^*]$$
   これにより、モデルは単一の正解だけでなく、複数の正解候補を網羅することを学習します。

2. Multi-Answer RLCR (較正された信頼度):
   正解性に加えて、各回答の信頼度 $q_i$ が実際の正誤と一致するよう、適切なスコアリング則（Proper Scoring Rules）に基づくペナルティを導入します。
   $$R_{Brier}^{multi}(A, Q, Y^*) = \frac{1}{K} \sum_{i=1}^{K} (q_i - \mathbb{I}[a_i \in Y^*])^2$$
   最終的な報酬は、正解性を最大化し、Brier スコア（較正誤差）を最小化する方向に設計されます：
   $$R_{RLCR}^{multi} = R_{RLVR}^{multi} - R_{Brier}^{multi}$$

2.3 出力形式の制約

• 一意性: 生成された $K$ 個の回答は互いに異なる必要があります（重複排除）。
• 確率の和: 単一正解タスク（$N=1$）では、信頼度の和が 1 以下になるよう制約を設け、離散確率分布として解釈できるようにしています。

---

3. 主要な貢献

1. 分布推論のための RL 手法の提案: 推論時のサンプリングに依存せず、トレーニング段階でモデルが「回答の分布」を生成・較正する能力を習得させる新しい RL 枠組みを確立しました。
2. 多様性とカバレッジの向上: 単一回答トレーニングのモデルが陥る「モード収束」を回避し、複数の妥当な仮説を網羅的に生成する能力を大幅に向上させました。
3. 計算効率の改善: 推論時に複数回サンプリングして複数の回答を得る従来の手法（Best-of-K）と比較し、単一のフォワードパスで同等以上の多様な回答セットを生成できるため、トークン使用量を大幅に削減しました。
4. 較正された不確実性の出力: 各回答に対して信頼度スコアを付与し、特に医療診断や曖昧なタスクにおいて、モデルが「何を知らないか」を適切に表現できるようにしました。

---

4. 実験結果

医療診断（DDXPlus）、曖昧な QA（HotPotQA-Modified）、コーディング（MBPP）の 3 つのベンチマークで評価を行いました。

4.1 正解性とカバレッジ

• カバレッジ（Coverage）: Multi-Answer RL は、単一回答ベースラインと比較して、正解の回答セットをより多く網羅しました（例：医療診断タスクでカバレッジが 0.62 → 0.79 に向上）。
• 多様性（Diversity）: 単一回答モデルを複数回サンプリングしても得られないような、構造的に異なる正解（コーディングタスクなど）を、Multi-Answer RL は単一生成で発見しました。
• Top-1 精度: 集合内の最上位回答の精度も維持、あるいは向上しました（コーディングタスクでは 50% 以上向上）。

4.2 計算効率

• トークン効率: 単一回答モデルで $K$ 個の回答を得るために必要なトークン数に対し、Multi-Answer RL は約半分以下のトークン数で同等の回答セットを生成しました。これは、推論時の重複した推論プロセス（思考連鎖のオーバーラップ）を排除したためです。

4.3 較正性（Calibration）

• Brier スコアと ECE: RLCR Multi（較正報酬付き）は、RLVR Multi と比較して、回答セット全体および個々の回答における信頼度スコアの較正性が大幅に向上しました。
• 過信の抑制: 従来の RLVR モデルは高確信で誤答する傾向（過信）が見られましたが、Multi-Answer RLCR はこれを抑制し、不確実な状況では適切な低い信頼度を出力しました。

4.4 具体例

• 医療診断: 患者の症状から「結核」「肺炎」「気管支炎」など複数の可能性を同時に提示し、それぞれに確率（例：結核 40%、肺炎 30%...）を付与しました。
• コーディング: 同じ機能を持つがアルゴリズムが異なる複数の実装（例：文字列変換 vs 数値計算）を一度に生成し、単一回答モデルが繰り返す「同じコードのバリエーション」を回避しました。

---

5. 意義と結論

本論文は、言語モデルのトレーニング目標を「単一の正解」から「妥当な回答の分布」へと転換する重要なステップを示しています。

• 実用性: 医療診断やリスク管理など、単一の正解が存在しない、あるいは複数の選択肢を検討する必要がある高リスク分野において、モデルの判断支援能力を飛躍的に向上させます。
• 効率性: 推論時の計算コストを削減しつつ、より豊かな推論能力を実現するため、大規模モデルの実運用におけるスケーラビリティが向上します。
• 将来的な展望: 本手法は、モデルが内部の確率分布を明示的に表現・操作する能力を身につけるための基盤となり、より高度な推論タスクや意思決定支援システムへの応用が期待されます。

結論として、Multi-Answer RL は、推論時のスケーリング（Best-of-K など）に依存するのではなく、モデル自体に分布推論の能力を内包させる「原理的かつ計算効率的な代替手段」として位置づけられます。