Rewarding Reasoning, Not Answers: Fixing and Bounding Test-Time Reinforcement Learning on Medical QA
Kailong Fan, Yichen Wu and colleagues (Harvard Medical School/MGH) show why majority-vote test-time RL collapses on medical multiple-choice QA and propose PROSE, which rewards reasoning steps instead of answer agreement.
Ask this paper
Diagnosis: A controlled experiment changing only the answer space shows that in small answer spaces wrong rollouts agree on the same wrong pseudo-label and reinforce it.
PROSE: A medical process reward model scores each step, the trajectory reward is the minimum step score, and an answer-format guard is enforced.
Results: Without labels, Llama-3.1-8B gets the best average accuracy among evaluated models of its size class, matches much larger systems, needs no reward model at inference and keeps gains on unseen datasets.
Aggregation matters: Mean aggregation is exploited: proxy reward saturates while accuracy falls. Below about 4B parameters, inference-time PRM selection remains better.
Abstract
Test-time reinforcement learning adapts a model on its own unlabeled test set using majority-vote pseudo-labels and has shown strong results in mathematics. We show that this recipe collapses on medical multiple-choice QA: accuracy stagnates while output diversity rapidly declines. Through a controlled experiment that keeps the questions, model, and optimizer fixed while changing only the answer space, we trace this failure to answer-space structure rather than domain difficulty. In small answer spaces, incorrect rollouts often collide on the same wrong pseudo-label and reinforce it; in large answer spaces, they disperse and receive little reward. This diagnosis motivates PROSE, Process Reward Guided Self-Training, which rewards reasoning quality instead of answer agreement. PROSE scores each reasoning step with a medical process reward model, assigns the trajectory reward as the minimum score across steps, and enforces answer-format constraints. Without labels, PROSE substantially improves a general Llama model, surpassing purpose-built medical models and matching much larger systems. Because the process signal is internalized into the policy, the adapted model requires no reward model at inference and transfers its gains to unseen datasets. We further show that the minimum aggregation is essential: mean aggregation can be exploited, saturating the proxy reward while degrading accuracy.