Reproduce “Reaching Beyond the Mode: RL for Distributional Reasoning in Language Models” on OpenResearch

About this paper

Isha Puri, Mehul Damani, Idan Shenfeld, Marzyeh Ghassemi, Jacob Andreas, Yoon Kim

Multi-Answer Reinforcement Learning enables large language models to perform distributional reasoning by generating diverse sets of plausible answers with calibrated confidence estimates in a single pass. This approach enhances answer coverage and diversity across medical diagnosis, ambiguous QA, and coding tasks, while also improving computational efficiency compared to repeatedly sampling from single-answer models.

Key insights

• Training LMs with a multi-answer objective directly counteracts entropy collapse, allowing models to naturally explore and represent a broader distribution of plausible responses.
• Explicitly optimizing for set-level correctness and calibration during training is more effective and computationally efficient than relying on post-hoc inference-time sampling or prompting techniques.
• The proposed framework generalizes to various task complexities, including single-correct-answer tasks with inherent ambiguity and multi-correct-answer tasks, fostering diverse and relevant outputs.

Summary by alphaXiv.