cs.LGMar 25, 2026

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

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

Organizations: Massachusetts Institute of Technology

Abstract

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/.

Figures & tables

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Explore similar work

CardsList
  1. Escaping the Mode Lottery: Multi-Response Training Improves Language Model Generalization

    May 30, 2026Hasan Amin, Kian Ahrabian, Ming Yin +1LLM Fine-TuningOOD Generalization

  2. Rationale Distribution Distillation via Flow Matching for Distributional Continuous Latent Reasoning

    Oct 14, 2026Mingyeong Kim, Gyusang Cho, Juho Lee +1LLM ReasoningContinuous Latent Reasoning

  3. MaD-RL: Matching Distributions for Calibrating LLMs with Reinforcement Learning

    Sep 12, 2026Sourabh Kulkarni, Ksheeraj Sai Vepuri, Basar Demir +8Distribution MatchingRL for Language Models