Med-U1: Incentivizing Unified Medical Reasoning in LLMs via Large-scale Reinforcement Learning

TL;DR

Med-U1 introduces a unified reinforcement learning framework that enhances large language models' medical reasoning capabilities across diverse question-answering tasks, outperforming specialized models and generalizing well to new, unseen tasks.

Contribution

The paper presents Med-U1, a novel reinforcement learning approach that unifies medical QA tasks with diverse outputs, improving reasoning and performance of LLMs in medical domains.

Findings

Significant performance improvements on multiple Med-QA benchmarks.

Outperforms larger, specialized, and proprietary models.

Demonstrates strong generalization to out-of-distribution tasks.

Abstract

Medical Question-Answering (QA) encompasses a broad spectrum of tasks, including multiple choice questions (MCQ), open-ended text generation, and complex computational reasoning. Despite this variety, a unified framework for delivering high-quality medical QA has yet to emerge. Although recent progress in reasoning-augmented large language models (LLMs) has shown promise, their ability to achieve comprehensive medical understanding is still largely unexplored. In this paper, we present Med-U1, a unified framework for robust reasoning across medical QA tasks with diverse output formats, ranging from MCQs to complex generation and computation tasks. Med-U1 employs pure large-scale reinforcement learning with mixed rule-based binary reward functions, incorporating a length penalty to manage output verbosity. With multi-objective reward optimization, Med-U1 directs LLMs to produce concise and verifiable reasoning chains. Empirical results reveal that Med-U1 significantly improves performance across multiple challenging Med-QA benchmarks, surpassing even larger specialized and proprietary models. Furthermore, Med-U1 demonstrates robust generalization to out-of-distribution (OOD) tasks. Extensive analysis presents insights into training strategies, reasoning chain length control, and reward design for medical LLMs. Our code is available here.