Source-linked AI summary
Reaching Beyond the Mode: RL for Distributional Reasoning in Language Models
Isha Puri, Mehul Damani, Idan Shenfeld, Marzyeh Ghassemi, Jacob Andreas, Yoon Kim
TL;DR
Single-answer RL can suppress valid alternatives and poorly represent uncertainty in tasks with multiple correct answers or incomplete information. The paper introduces Multi-Answer RL, which trains models to generate distinct answer sets and confidence estimates directly, yielding broader coverage and diversity, better set-level calibration, lower token use, and higher coding accuracy, with limitations in calibration and experimental scope.
Problem
Single-answer post-training can collapse answer distributions despite real-world tasks involving multiple valid answers or uncertainty from incomplete information.
Method
Multi-Answer RL trains models to generate multiple distinct candidate answers in one generation and uses calibration rewards to produce confidence estimates for each answer.
Results
Across QA, medical, and coding benchmarks, Multi-Answer RL improves coverage, diversity, token efficiency, and set-level calibration; coding top-1 accuracy rises by over 50%.
Takeaways & Limitations
The results support multi-answer training as a compute-efficient alternative to repeated sampling for generating plausible answer distributions.
Takeaways & Limitations
Single-answer objectives retain higher top-1 accuracy, calibration can degrade in challenging single-gold settings, and experiments are confined to QA while serial generation limits parallelism.
Abstract
from arXiv · showhide
Given a question, a language model (LM) implicitly encodes a distribution over possible answers. In practice, post-training procedures for LMs often collapse this distribution onto a single dominant mode. While this is generally not a problem for benchmark-style evaluations that assume one correct answer, many real-world tasks inherently involve multiple valid answers or irreducible uncertainty. Examples include medical diagnosis, ambiguous question answering, and settings with incomplete information. In these cases, we would like LMs to generate multiple plausible hypotheses, ideally with confidence estimates for each one, and without computationally intensive repeated sampling to generate non-modal answers. This paper describes a multi-answer reinforcement learning approach for training LMs to perform distributional reasoning over multiple answers during inference. We modify the RL objective to enable models to explicitly generate multiple candidate answers in a single forward pass, internalizing aspects of inference-time search into the model's generative process. Across question-answering, medical diagnostic, and coding benchmarks, we observe improved diversity, coverage, and set-level calibration scores compared to single answer trained baselines. Models trained with our approach require fewer tokens to generate multiple answers than competing approaches. On coding tasks, they are also substantially more accurate. These results position multi-answer RL as a principled and compute-efficient alternative to inference-time scaling procedures such as best-of-k. Code and more information can be found at https://multi-answer-rl.github.io/.
1 Introduction
Single-answer reinforcement learning can suppress valid alternative answers and uncertainty, motivating training objectives that generate calibrated answer sets directly. Multi-Answer RL explicitly optimizes this behavior and improves diversity, coverage, calibration, and coding efficiency across evaluated tasks.
- Motivation: Single-answer RL objectives can collapse models onto dominant answers even when multiple answers are valid or uncertainty is irreducible.The paper highlights incomplete information, ambiguous questions, and clinical diagnosis as motivating settings.
- Motivation: Inference-time sampling and uncertainty verbalization can produce answer sets or confidence estimates, but the paper instead optimizes these properties during training.
- Approach: Multi-Answer RL trains models to reason over multiple hypotheses and generate structured candidate sets in one generation.A proper-scoring-rule-inspired reward extends the approach to confidence scores for each answer.
- Results: Across question answering, medical diagnosis, and coding, the approach improves coverage, answer diversity, and token efficiency relative to baselines.
- Results: On coding tasks, top-1 accuracy increases by over 50% while token usage decreases by more than half.
2 Background
The background contrasts standard language-model reinforcement learning, which maximizes expected reward and can concentrate probability mass on single outputs, with calibration rewards that combine correctness and confidence calibration. Proper scoring rules provide the basis for the latter objective.
- Standard RL: Standard language-model RL maps prompts to output distributions and maximizes expected reward over prompt–answer pairs.
- Standard RL: Expected-reward training can concentrate a policy’s probability mass on single outputs.
- RLVR: RLVR uses a deterministic binary correctness reward that indicates whether an output matches the ground-truth answer.
- RLCR: RLCR adds confidence estimates and uses a proper scoring rule to jointly incentivize correctness and calibration.
3 Multi-Answer Reinforcement Learning
Multi-Answer RL generalizes reinforcement learning from single outputs to structured sets of distinct candidates, rewarding recovery of valid answers. Multi-Answer RLCR additionally trains answer-level confidence estimates and interprets outputs as distributions over plausible answers.
- Multi-Answer RLVR: The generalized setting associates each prompt with a variable-sized set of valid answers and aims to recover the full set or a high-coverage subset in one generation.
- Multi-Answer RLVR: The policy generates K distinct candidate answers within a single chain of thought and receives a set-level reward based on valid-answer membership.
- Objective connections: When N = 1 and K = 1, the objective reduces to standard RLVR; when N = 1 and K > 1, it is equivalent to pass@K.
- Objective connections: For N > 1, K ≤ N encourages maximal coverage of distinct valid answers, while K ≥ N permits recovery of the entire ground-truth set.
- Multi-Answer RLCR: Multi-Answer RLCR adds a confidence value q_i for each candidate answer and combines correctness rewards with a proper scoring rule.
- Distributional interpretation: The output (A, Q) represents a distribution over plausible answers, with confidence scores summing to at most one for N = 1 and forming multivariate correctness events when N > 1.
4 Experiments
Across three set-valued reasoning tasks, multi-answer reinforcement learning improves coverage, diversity, and calibration while reducing redundant generation costs. The experiments also show that explicit multi-answer training is needed to recover distinct hypotheses reliably.
- Experimental Setup: The evaluation spans DDXPLUS, HotPotQA-Modified, and MBPP, covering multiple diagnoses, incomplete-information ambiguity, and distinct correct program implementations.DDXPLUS allows multiple correct diagnoses, HotPotQA-Modified uses incomplete context around a single gold answer, and MBPP permits structurally distinct correct implementations.
- Correctness: Multi-Answer models substantially outperform Single-Answer models on set-level correctness metrics, including coverage.Simply prompting Single-Answer models to produce multiple answers performs substantially worse than trained Multi-Answer models.
- Correctness: Multi-Answer training recovers more alternatives in DDXPLUS, more distinct implementations in MBPP, and higher pass@k in ambiguous HotPotQA-Modified questions.The HotPotQA-Modified gains arise from reasoning over multiple hypotheses within a single generation, whereas Single-Answer models repeatedly regenerate the dominant answer.
- Calibration: RLCR consistently improves calibration over RLVR across all three datasets and both Single- and Multi-Answer settings.On DDXPLUS and MBPP, RLCR-Multi improves set-level calibration across metrics; on DDXPLUS, RLVR-Multi remains systematically overconfident.
- Calibration: Multi-Answer training matches Single-Answer RL on Top-1 ECE and Top-1 Brier, so set-valued optimization does not degrade calibration of the top answer.These are the only calibration metrics reported as directly comparable between Single- and Multi-Answer methods.
- Diversity and Efficiency: Multi-Answer RL produces nearly twice as many unique answers on average, 8 versus 4, while retaining higher average coverage.It also generates the same number of answers using only 56% of the token length required by independent Single-Answer sampling in the medical domain.
5 Related Work
Prior work studies calibration and answer-set generation through probing, sampling, verbalization, prompting, and related training objectives. These approaches motivate explicit methods for representing multiple plausible answers and their confidence.
- Recent work documents trade-offs between reinforcement-learning gains in pass@1 and contraction of the model’s reasoning support.Other studies report base-model advantages over aligned models on creative tasks.
- Prior calibration methods include intent probing, sampling-based uncertainty estimation, and post-hoc confidence verbalization.
- Calibration: Proper-scoring-rule rewards have been used to align predicted confidence with empirical correctness.Some methods optimize calibration alone, while others jointly optimize correctness and calibration.
- Generating Answer Sets: Prompting methods elicit distributions over fixed labels, verbalized distributions, or enumerated and iteratively sampled answers.
6 Conclusion
The paper concludes that Multi-Answer RL trains reasoning language models to generate plausible answer sets, improving coverage, diversity, and efficiency across evaluated settings. It also identifies higher single-answer top-1 accuracy, QA-only evaluation, and limited parallelism as remaining boundaries.
- Multi-Answer RL trains reasoning language models to generate sets of plausible answers rather than a single most probable response.
- Across medical and question-answering benchmarks, the objective improves coverage and diversity while recovering alternatives missed by standard reinforcement learning.
- Multi-Answer RL uses fewer total tokens to generate answer sets than the compared single-answer approach.
- Limitations: Single-answer objectives retain higher top-1 accuracy, while the experiments are confined to question answering.
- Limitations: Serial generation limits parallelism despite improved token efficiency.
A System Prompts
The system prompts specify how assistants should structure reasoning, answers, and confidence. Separate formats support single-answer and multi-answer responses with distinct candidates and confidence scores.
- Confidence-aware prompts add uncertainty analysis and a confidence level between 0 and 1 inside <confidence> tags.
- Single-answer prompts enclose reasoning and the final answer in <think> and <answer> tags.
- Multi-answer prompts require multiple candidate answers, each paired with a confidence score.
- The multi-answer format requires exactly {K} answer pairs and specifies that candidate confidences sum to one.
- A stricter multi-answer prompt requires exactly {K} distinct candidates and separate reasoning about why each could be correct or incorrect.
B Additional Training Details
Training uses fixed sampling and optimization settings, structured format rewards, and additional constraints for uniqueness and confidence validity. Format and calibration rewards are balanced equally.
- Training samples 32 responses per prompt at temperature 0.7 with an effective batch size of 1536.The base learning rate is 1 × 10−6 with a warmup ratio of 0.05.
- Optimization removes standard-deviation advantage normalization and uses BNPO with token losses normalized by active tokens.Completions are capped at 1536 tokens and training runs for one epoch.
- Format rewards enforce the required tags for single-answer and multi-answer outputs, while format and calibration rewards receive equal weights.
- Multi-answer training zeroes all rewards when extracted answers are not unique.For N = 1 datasets, the format reward also zeroes out when confidence sums exceed 1.
C Token Efficiency Comparison: Single Answer vs. Multi Answer RL
The comparison visualizes average token usage for single- and multi-answer RL variants on DDXPlus. The accompanying text reports that Multi-Answer training significantly improves compute efficiency.
- The figure presents token compute usage for sets of answers generated by the four training variants.
- Multi-Answer training significantly increases compute efficiency in the token-usage comparison.
- Average token usage is compared across RLVR Single, RLVR Multi, RLCR Single, and RLCR Multi on DDXPlus.
D Unique Answer Analysis: RLVR Single vs. RLVR Multi
An example medical-question analysis compares answer diversity from repeated RLVR Single runs with sets generated by RLVR Multi. RLVR Multi admits substantially more unique answers in fewer runs.
- RLVR Single is run 30 times, whereas RLVR Multi is run 10 times with k = 3 answers per set.
- RLVR Single collects only 3 unique answers across 30 runs, while RLVR Multi admits significantly more unique answers.
- The word cloud compares unique answers from RLVR Single and RLVR Multi on an example medical question.
E Increasing k in RLVR-Multi
Table 3 measures coverage as the number of correct answers generated in a set while k increases, using 5,000 evaluation questions.
- 5,000 evaluation questions are used to measure coverage as k increases.Coverage is defined as the average number of correct answers generated in a set.
F Full Example - Medical RLVR Single, RLVR Multi, RLCR Single, RLCR Multi
The medical example presents symptoms and antecedents for differential diagnosis, followed by multi-answer outputs with ranked diagnoses and confidence estimates. The generated set includes pulmonary embolism, pneumonia, and tuberculosis, while the gold list contains several plausible conditions.
- Medical case: The case includes hemoptysis, sudden chest pain, shortness of breath, fever, chills, and a localized peeling rash.The patient also has alcohol addiction, heart failure, COPD, and is underweight.
- Gold answers: The gold differential includes pneumonia, pulmonary neoplasm, bronchitis, tuberculosis, several cardiac conditions, and GERD.The listed gold answers contain nine conditions or diagnostic possibilities.
- RLCR Multi: The multi-answer output ranks pulmonary embolism, pneumonia, and tuberculosis as the three leading diagnoses.The diagnoses are presented as answer1, answer2, and answer3.
- RLCR Multi: The ranked diagnoses receive confidence estimates of 0.45, 0.35, and 0.20, respectively.These values correspond to pulmonary embolism, pneumonia, and tuberculosis in output order.