PaperlayerView exact revision on arXiv

Source-linked AI summary

Rewarding Reasoning, Not Answers: Fixing and Bounding Test-Time Reinforcement Learning on Medical QA

Kailong Fan, Anqi Pu, Yichen Wu, Wanhua Li, Yicong Li, Hanspeter Pfister, Huafeng Liu, Xiang Li, Quanzheng Li, Ning Guo

arXiv:2609.16660v1cs.CL

TL;DR

Majority-vote test-time reinforcement learning can collapse on medical multiple-choice QA because small answer spaces reinforce coincident wrong pseudo-labels. The paper introduces PROSE, which rewards reasoning steps with a medical process reward model and aggregates them by the minimum score. PROSE improves a general Llama model, internalizes the reward for reward-model-free inference, and transfers gains beyond adapted datasets.

  • Problem

    The paper investigates why majority-vote self-training succeeds in mathematics but stagnates and reduces diversity on medical multiple-choice QA.

  • Method

    PROSE performs label-free test-time adaptation using step-level medical process rewards, minimum trajectory aggregation, and answer-format constraints.

  • Results

    Without labels, PROSE gives Llama-3.1-8B the best average accuracy among evaluated 8–9B models and matches much larger systems.

  • Takeaways & Limitations

    The process signal is internalized into the policy, so adapted models need no reward model at inference and retain gains on unseen datasets.

  • Takeaways & Limitations

    PROSE requires a domain-relevant process reward model, and inference-time selection remains preferable below roughly 4B parameters in the reported experiments.

Abstract

from arXiv · show

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.

1 Introduction

Test-time reinforcement learning is attractive for unlabeled medical QA, but majority-vote self-training fails in structured multiple-choice settings. The paper attributes this failure to answer-space collisions and introduces PROSE, which rewards reasoning quality instead.

  • Motivation: Medical QA is a valuable target for label-free adaptation because expert annotation is costly, clinical distributions vary, and privacy constraints limit external labeling.
  • Observed failure: Standard TTRL produces only modest avg@16 gains, no maj@16 gains, and sharply reduced output diversity on medical QA, whereas PROSE improves both accuracy metrics.The contrast indicates that greater agreement does not necessarily improve correctness.
  • Diagnosis: A controlled experiment varying only the answer space supports a structural explanation: wrong rollouts collide on shared options in multiple choice but disperse in free-form answers.The protection gap between pseudo-label reward accuracy and true answer accuracy nearly vanishes under multiple choice.
  • PROSE: PROSE replaces majority-vote pseudo-labels with step-level medical process rewards, uses the minimum step score as trajectory reward, and enforces answer-format constraints.Minimum aggregation makes the reward sensitive to the weakest reasoning step.
  • Deployment: The adapted policy internalizes the process signal, needs no reward model at inference, outperforms PRM-based selection at the same budget, and transfers gains to unseen datasets.The PRM is used during adaptation rather than on every future query.

2 Related Work

Prior work improves reasoning through inference-time scaling, consensus, selection, or test-time adaptation. Process rewards address weaknesses of outcome-only evaluation by judging intermediate reasoning steps.

  • Test-Time Scaling and Self-Evolution: Inference-time scaling samples multiple reasoning trajectories and aggregates them through consensus or selection, while TTRL adapts model parameters using rewards derived from repeated sampling.
  • Test-Time Scaling and Self-Evolution: Medical test-time adaptation is especially appealing because real-time gold standards are rarely available, making unlabeled approaches practically relevant.
  • Inference-Time Verification: Outcome-based reward models can misjudge trajectories containing spurious correct reasoning or later derailments, motivating process reward models that assess reasoning steps.

3 Motivation: TTRL Fails on Medical MCQ

Majority-vote TTRL stagnates on medical multiple-choice QA because agreement over a small answer space is not a reliable correctness signal. Controlled comparisons link the collapse to answer-space structure and motivate process-based reward design.

  • Observed Collapse: TTRL yields limited sample-level accuracy gains, unchanged majority-vote accuracy, and sharply lower output diversity, making the model more self-consistent without becoming more correct.
  • Answer Space Matters: Holding questions, model, optimizer, and training procedure fixed while changing only the answer space isolates answer-space structure rather than domain difficulty.
  • Answer Space Matters: Avg@16 improves for free-form answers but collapses for four-option MCQs, while the protection gap remains positive for free-form answers and vanishes for MCQs.
  • The Lucky Hit Mechanism: In large free-form spaces, wrong rollouts disperse across distinct strings and usually avoid reinforcement; in small answer spaces, mistaken rationales converge on the same wrong option.
  • From Diagnosis to Reward Design: PROSE follows by replacing final-answer agreement with medical process-level rewards and assigning each trajectory the minimum score across reasoning steps.

4 Method

PROSE replaces answer-agreement rewards with process-level rewards during label-free test-time reinforcement learning, using minimum step quality and format constraints to guide GRPO updates.

  • Process-Level Reward: PROSE samples unlabeled-question rollouts, scores each reasoning step with a medical process reward model, takes the minimum step score, guards answer format, and updates the policy with GRPO.The adapted policy is trained entirely on unlabeled test questions.
  • Process-Level Reward: Unlike standard TTRL, PROSE rewards reasoning quality rather than agreement with other sampled final answers, avoiding direct dependence on answer-space collisions.The process score is independent of whether other rollouts select the same option.
  • Process-Level Reward: Minimum aggregation keeps trajectory reward controlled by the weakest reasoning step, limiting reward hacking that mean or sum aggregation can enable.Mean aggregation can dilute a flawed step with high-scoring steps, whereas the minimum improves only when the weakest step improves.
  • Format Guard: The format guard assigns zero reward to rollouts without a parseable answer, so high reward requires both valid output formatting and consistently reliable reasoning.This prevents optimization from favoring reasoning traces that never produce an extractable final answer.
  • Inference Cost: After adaptation, the policy is evaluated directly with ordinary sampling or majority voting, requiring no process reward model at inference.The PRM serves as a temporary training signal rather than a per-query selector.

5 Experiments Setup

The experiments evaluate PROSE across clinical multiple-choice benchmarks, model families, baselines, and matched sampling protocols under transductive unlabeled adaptation.

  • Policies: The study evaluates Llama-3.1-8B-Instruct and Qwen3 models at 1.7B, 4B, and 8B, using the 8B models as the main policies.The smaller Qwen3 models probe the effect of model scale.
  • Process reward model: Med-PRM is an 8B medical process reward model that assigns each reasoning step a correctness probability and is used only during test-time training.The adapted policy answers without Med-PRM at inference.
  • Datasets and protocol: Experiments use four clinical multiple-choice benchmarks, with each policy adapted separately on each unlabeled test set without ground-truth labels.The datasets are MedQA-5op, MedQA-4op, MedMCQA, and DDXPlus.
  • Baselines: Baselines include no adaptation, majority-vote TTRL, inference-time Med-PRM selection, and large general-purpose, open-source, and medical models.Inference-time selection includes best-of-N reranking and reward-weighted self-consistency.
  • Metrics: Evaluation uses avg@16 as the primary metric, while matched-budget comparisons with single-answer selection methods use maj@16; all methods receive 16 samples per question.Evaluation samples use temperature 0.6 and top-p 0.95.
  • Results protocol: Table 1 reports avg@16 across four medical QA datasets and marks the best and second-best models below 32B parameters.The table caption states that PROSE consistently improves standard TTRL across all datasets.

6 Results and Analysis

PROSE improves medical QA adaptation by internalizing process rewards, outperforming answer-level TTRL and inference-time selection while generalizing beyond the adaptation data. Its minimum step-score aggregation avoids reward hacking, though benefits depend on sufficient policy capacity.

  • 6.1 Main Results: 0.740 avg@16 makes PROSE the best 8–9B model, exceeding Llama-3.1-8B’s base score of 0.635, TTRL’s 0.671, and the strongest purpose-built medical model’s approximately 0.665.PROSE also nearly matches QwQ at 0.741 and GPT-4o-mini at 0.744, while trailing Gemini Flash 2.0 at 0.801.
  • 6.1 Main Results: At matched sampling budgets, PROSE improves average performance over TTRL by 6.6 points on Llama-3.1-8B and 4.3 points on Qwen3-8B, outperforming TTRL on every dataset.The comparison isolates the reward signal because both methods otherwise use the same test-time RL framework.
  • 6.2 Ablation: What matters in the reward design: PRM-derived labels raise avg@16 from 0.672 to 0.709, while continuous graded rewards reach 0.724 because they preserve ordering among partially correct trajectories.Collapsing step scores into a single pseudo-label makes incorrect trajectories equally wrong and discards their gradation.
  • 6.1 Main Results: After adaptation on MedQA-5op, frozen PROSE policies outperform the base model and standard TTRL on three unseen datasets without further adaptation or PRM scoring.This transfer indicates that the process reward is internalized in the policy rather than used only as an inference-time selector.
  • 6.2 Ablation: What matters in the reward design: Minimum aggregation improves avg@16 by 8.5 points over the untrained base, whereas mean aggregation eventually falls 0.9 points below it as proxy reward approaches 1.0.Mean aggregation dilutes weak steps and enables reward hacking; minimum aggregation forces optimization to improve the weakest reasoning step.
  • 6.3 Ablation: Effect of model scale: At 4B, PROSE leads all four datasets with a 0.734 average, but at 1.7B best-of-N leads on average at 0.633 versus PROSE’s 0.614, with PROSE winning only on DDXPlus.The reversal indicates that effective process-reward training requires sufficient policy capacity.

7 Conclusion

PROSE addresses the structural failure of test-time reinforcement learning on medical multiple-choice QA and achieves strong adaptation results while internalizing process rewards into the policy. Its scope is constrained by the need for a domain-relevant process reward model and sufficient policy capacity.

  • 7 Conclusion: PROSE gives Llama-3.1-8B the best average accuracy among evaluated 8–9B models and matches much larger systems without labels.Its gains also remain on unseen adaptation datasets.
  • 7 Conclusion: Small answer spaces cause incorrect rollouts to collide on wrong pseudo-labels, explaining why test-time RL collapses on medical multiple-choice QA.The protection gap is +0.246 for free-form answering versus +0.001 for multiple-choice answering.
  • 7 Conclusion: PROSE rewards reasoning through minimum step-level aggregation and an answer-format guard instead of relying on final-answer agreement.Minimum aggregation makes the trajectory reward depend on the weakest reasoning step.
  • 7 Conclusion: PROSE requires a domain-relevant process reward model, and inference-time selection remains preferable below roughly 4B parameters in these experiments.Future work concerns reducing reward-model dependence and characterizing how policy capacity affects reward internalization.
Loading 2609.16660v1…