Выход за пределы моды: обучение с подкреплением для распределительных рассуждений в языковых моделях

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

March 25, 2026

Авторы: Isha Puri, Mehul Damani, Idan Shenfeld, Marzyeh Ghassemi, Jacob Andreas, Yoon Kim

cs.AI

Аннотация

При заданном вопросе языковая модель (ЯМ) неявно кодирует распределение возможных ответов. На практике посттренировочные процедуры для ЯМ часто схлопывают это распределение до единственной доминирующей моды. Хотя это обычно не является проблемой для бенчмарк-оценок, предполагающих один правильный ответ, многие реальные задачи по своей природе связаны с множеством допустимых ответов или нередуцируемой неопределенностью. Примерами служат медицинская диагностика, ответы на неоднозначные вопросы и сценарии с неполной информацией. В таких случаях желательно, чтобы ЯМ генерировали несколько правдоподобных гипотез, в идеале с оценкой уверенности для каждой, без вычислительно затратного повторного сэмплирования для получения немодальных ответов. В данной статье описывается подход обучения с подкреплением для множественных ответов, который тренирует ЯМ выполнять дистрибутивные рассуждения над несколькими ответами во время вывода. Мы модифицируем цель RL, чтобы модели могли явно генерировать несколько кандидатных ответов за один прямой проход, интериоризируя аспекты поиска на этапе вывода в генеративный процесс модели. На бенчмарках вопросно-ответных систем, медицинской диагностики и программирования мы наблюдаем улучшенные показатели разнообразия, покрытия и калибровки на уровне множеств по сравнению с базовыми моделями, обученными на одиночные ответы. Модели, обученные нашим методом, требуют меньше токенов для генерации нескольких ответов, чем конкурирующие подходы. На задачах программирования они также существенно точнее. Эти результаты позиционируют multi-answer RL как принципиальную и вычислительно эффективную альтернативу процедурам масштабирования на этапе вывода, таким как best-of-k. Код и дополнительная информация доступны по адресу https://multi-answer-rl.github.io/.

English

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