모드 너머로: 언어 모델의 분포적 추론을 위한 강화 학습

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

질문이 주어졌을 때, 언어 모델(LM)은 가능한 답변들에 대한 분포를 암묵적으로 인코딩합니다. 실제로는, LM의 후처리 과정에서 이러한 분포가 종종 단일한 우세한 모드로 붕괴됩니다. 이는 일반적으로 하나의 정답을 가정하는 벤치마크 평가에서는 문제가 되지 않지만, 많은 실제 작업은 본질적으로 여러 개의 유효한 답변을 포함하거나 해결할 수 없는 불확실성을 내포합니다. 이러한 예로는 의학적 진단, 애매한 질문 답변, 그리고 불완전한 정보가 있는 환경 등이 있습니다. 이러한 경우, 우리는 LM이 여러 개의 그럴듯한 가설을 생성하기를 원하며, 이상적으로는 각 가설에 대한 신뢰도 추정치를 함께 제공하고, 계산 집약적인 반복 샘플링 없이 비-모달 답변을 생성해야 합니다. 본 논문에서는 추론 과정에서 LM이 여러 답변에 대한 분포적 추론을 수행하도록 학습하기 위한 다중 답변 강화 학습(RL) 접근 방식을 설명합니다. 우리는 RL 목적 함수를 수정하여 모델이 단일 순전파 과정에서 여러 개의 후보 답변을 명시적으로 생성할 수 있도록 하고, 추론 시간의 검색 과정을 모델의 생성 과정에 통합합니다. 질문 답변, 의학적 진단, 그리고 코딩 벤치마크에서, 우리는 단일 답변으로 학습된 기준 모델과 비교하여 향상된 다양성, 커버리지 및 집합 수준의 교정 점수를 관찰했습니다. 우리의 접근 방식으로 학습된 모델은 경쟁적인 접근 방식보다 여러 답변을 생성하는 데 더 적은 토큰을 필요로 합니다. 코딩 작업에서, 이 모델은 또한 상당히 더 높은 정확도를 보입니다. 이러한 결과는 다중 답변 RL을 추론 시간 스케일링 절차(예: best-of-k)의 원칙적이고 계산 효율적인 대안으로 제시합니다. 코드 및 자세한 내용은 https://multi-answer-rl.github.io/ 에서 확인할 수 있습니다.

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