Test-time reinforcement learning improved by rewarding reasoning steps
Rewarding Reasoning, Not Answers: Fixing and Bounding Test-Time Reinforcement Learning on Medical QA
Computation and Language
Summary
Adapting AI models to answer medical multiple-choice questions during testing can fail because wrong answers tend to reinforce themselves. The authors found that this happens due to the way answers are structured rather than the difficulty of the questions. They created PROSE, a method that rewards the quality of reasoning steps instead of just the final answers, which helps the model improve without needing labeled data. PROSE also keeps the model’s gains when tested on new medical questions and doesn't need the reward system during actual use.
What this means in practice
- •For medical ai developers: Enhance AI models for medical question answering by improving reasoning quality without labeled data during testing.
- •For healthcare technology teams: Improve clinical decision support tools by adapting pre-trained models on unlabeled patient questions to boost accuracy and reliability.
Authors
Kailong Fan, Anqi Pu, Yichen Wu, Wanhua Li, Yicong Li, Hanspeter Pfister, Huafeng Liu, Xiang Li, Quanzheng Li, Ning Guo
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.