Reaching Beyond the Mode: RL for Distributional Reasoning in Language Models

1. Problem Statement

Current Large Language Models (LLMs), particularly those trained with Reinforcement Learning with Verifiable Rewards (RLVR), are optimized to produce a single dominant answer (the "mode") for a given query. While effective for benchmarks with unique ground-truth answers, this approach fails in real-world scenarios characterized by:

• Inherent Ambiguity: Tasks like medical diagnosis or ambiguous question answering often have multiple valid answers simultaneously.
• Epistemic Uncertainty: Incomplete information may make it impossible to determine a single "correct" answer with certainty.
• Entropy Collapse: Standard RL training tends to collapse the model's output distribution, suppressing plausible alternative hypotheses and leading to overconfidence.
• Inefficiency of Inference-Time Scaling: Current workarounds (e.g., sampling $K$ times independently to get a set of answers) are computationally expensive and often result in redundant reasoning chains (high token overlap) and duplicate answers.

The paper argues that models should be trained to explicitly represent and generate distributions over plausible answers (sets of candidates with confidence estimates) rather than just a single mode.

2. Methodology: Multi-Answer Reinforcement Learning

The authors propose Multi-Answer RL, a framework that modifies the standard RL objective to train models to generate a structured set of $K$ distinct candidate answers and associated confidence scores in a single forward pass.

2.1 Multi-Answer RLVR (Reinforcement Learning with Verifiable Rewards)

• Objective: Instead of optimizing for a single correct answer, the model is trained to output a set $A = \{a_1, ..., a_K\}$.
• Reward Function: The reward is defined at the set level:
  $$R_{multi}^{RLVR}(A, Y^*) = \sum_{i=1}^{K} \mathbb{I}[a_i \in Y^*]$$
  where $Y^*$ is the set of ground-truth valid answers. This encourages the model to maximize the coverage of correct answers within the generated set.
• Constraints: A format reward enforces that all $K$ answers are distinct (uniqueness constraint).

2.2 Multi-Answer RLCR (Reinforcement Learning with Calibration Rewards)

• Objective: Extends RLVR to include calibrated confidence estimates. The model outputs a set of answers $A$ and a set of confidence scores $Q = \{q_1, ..., q_K\}$.
• Reward Function: Combines the correctness reward with a Multi-Brier Score (a proper scoring rule) to penalize miscalibration:
  $$R_{multi}^{RLCR} = R_{multi}^{RLVR} - \frac{1}{K} \sum_{i=1}^{K} (q_i - \mathbb{I}[a_i \in Y^*])^2$$
• Interpretation: This trains the model to output a distribution where $q_i$ approximates the probability that answer $a_i$ is correct. In single-answer settings ($N=1$), the confidences are constrained to sum to $\le 1$; in multi-answer settings ($N>1$), they represent independent probabilities.

2.3 Training Setup

• Base Model: Qwen3-8B.
• Algorithm: GRPO (Group Relative Policy Optimization) with modifications for set-level rewards.
• Datasets:
  • DDXPlus: Medical diagnosis (multiple correct diagnoses per case).
  • HotPotQA-Modified: Ambiguous QA (single gold answer but incomplete context).
  • MBPP: Coding (multiple distinct algorithmic implementations for the same task).

3. Key Contributions

1. Novel Training Objective: Introduces a principled RL framework that shifts the training goal from "finding the best answer" to "recovering the distribution of plausible answers."
2. Inference-Time Efficiency: Demonstrates that generating a set of answers in a single pass is significantly more token-efficient than repeated sampling (inference-time scaling), as it avoids redundant reasoning chains.
3. Calibration via Proper Scoring Rules: Extends the use of proper scoring rules (Brier score) to multi-answer settings, enabling models to output well-calibrated confidence scores for each hypothesis.
4. Empirical Validation: Provides comprehensive benchmarks showing that Multi-Answer RL outperforms single-answer baselines in coverage, diversity, and calibration without sacrificing top-1 accuracy.

4. Key Results

The paper evaluates the approach across medical, QA, and coding tasks with $K=3$ candidates:

• Coverage & Diversity:

  • DDXPlus (Medical): Multi-Answer RLVR achieved 0.79 average correct answers per set vs. 0.62 for single-answer baselines. Diversity (unique answers) reached 1.00 (perfect) compared to 0.62.
  • MBPP (Coding): Top-1 accuracy improved by >50% (0.49 vs 0.29 for RLVR Single), and the model generated significantly more unique algorithmic solutions.
  • HotPotQA: Improved Pass@K performance, showing the model could recover the correct answer even when information was incomplete, whereas single-answer models collapsed to a single (often wrong) mode.

• Token Efficiency:

  • Multi-Answer models required ~44% fewer tokens to generate 3 answers compared to sampling a single-answer model 3 times.
  • Analysis showed that single-answer sampling often repeats the same reasoning scaffolding (high n-gram overlap), whereas Multi-Answer RL diversifies the reasoning process within a single generation.

• Calibration:

  • RLCR Multi significantly outperformed RLVR Multi in calibration metrics (Brier Score and ECE).
  • On DDXPlus, RLCR Multi achieved a Set ECE of 0.02 compared to 0.13 for RLVR Multi, indicating much better alignment between predicted confidence and actual correctness.
  • The model successfully learned to assign lower confidence when uncertain (e.g., in ambiguous HotPotQA settings), though it occasionally struggled with the "sum-to-one" prior in extremely difficult single-label settings.

5. Significance and Conclusion

This work represents a paradigm shift in how LMs are trained for reasoning tasks. By moving beyond the "single best answer" objective, Multi-Answer RL:

• Aligns with Real-World Needs: Better suits high-stakes domains (medicine, law, coding) where exploring multiple hypotheses is critical.
• Solves the Efficiency Bottleneck: Offers a compute-efficient alternative to "best-of-$K$" inference scaling, reducing token usage by nearly half while improving accuracy.
• Improves Trustworthiness: Produces calibrated uncertainty estimates, allowing users to understand the model's confidence in each hypothesis.

The authors conclude that explicitly optimizing for distributional reasoning allows models to "reach beyond the mode," recovering valid alternatives that standard RL suppresses, while simultaneously improving the efficiency and reliability of the generation process.