Technical Summary: InfiMed-ORBIT

Problem Statement

Reinforcement Learning (RL) has driven significant breakthroughs in Large Language Models (LLMs) for tasks with verifiable, automatically computable rewards, such as code generation. However, applying RL to open-ended medical dialogue remains challenging. In this domain, feedback is often ambiguous, context-dependent, and difficult to summarize into a single scalar signal. Traditional approaches relying on learned Reward Models (RMs) trained on human preferences are costly to construct, highly domain-dependent, and fragile under distribution shifts. Furthermore, oversimplified binary rewards are prone to "reward hacking," posing serious risks in high-stakes clinical scenarios where safety, empathy, and appropriateness vary significantly case-by-case.

Existing medical LLMs struggle to bridge the gap between Question-Answering (QA) style optimization and realistic, context-dependent clinical consultation, often failing to meet the rigorous standards of benchmarks like HealthBench-Hard.

Methodology: The ORBIT Framework

The authors propose ORBIT (Open-ended Rubric-based Incremental Training), a framework designed to align LLMs on critical medical dialogues without relying on opaque scalar rewards or costly, task-specific reward models. Instead, ORBIT utilizes dynamically generated, case-conditioned rubrics as adaptive guides for incremental RL.

The framework operates through three primary stages:

1. Medical Dialogue Construction

ORBIT synthesizes realistic multi-turn medical dialogue data. It leverages existing datasets (e.g., DoctorAgent-RL, IMCS21, MedDG) and employs agentic synthesis workflows to generate multi-turn consultation histories. These dialogues serve as the query inputs for the subsequent rubric generation and RL stages.

2. Context-Aware Rubric Generation

Rather than using fixed templates, ORBIT employs a Retrieval-Augmented Generation (RAG) pipeline to create tailored rubrics for each specific medical case:

• Diagnostic Database: A vector database is constructed containing seed rubrics from HealthBench and their associated dialogue histories.
• Retrieval: For a new consultation query, the system retrieves semantically similar cases and candidate rubrics.
• Generation: A generative model (e.g., DeepSeek-R1) uses these retrieved examples as in-context prompts to synthesize novel, case-specific rubrics. These rubrics include both positive criteria (rewarding desirable behaviors like safety checks and empathy) and negative criteria (penalizing unsafe or misleading actions).
• Filtering: To ensure training efficiency, a two-stage Pass@k filtering mechanism is applied:
  • Sample-Level: Retains queries of intermediate difficulty (neither too easy nor unsolvable) to provide meaningful optimization signals.
  • Rubric-Level: Removes rubrics that are already consistently satisfied by the current policy (saturated constraints) to focus on discriminative feedback.

3. Rubric-Guided Reinforcement Learning

The policy is optimized using Group Relative Policy Optimization (GRPO).

• Reward Modeling: Instead of a single scalar, the reward is a dense, signed sum of criterion-level evaluations. A judge model evaluates the model's response against the specific rubric criteria, assigning weights based on the criticality of the criterion (e.g., -10 for life-threatening errors, +10 for critical safety actions).
• Stability Mechanisms: To prevent training instability caused by reward saturation or low variance, ORBIT implements:
  • Variance-Aware Filtering: Discards rollout groups where the reward variance is near zero, ensuring updates are driven by discriminative signals.
  • Staged Entropy Injection: Periodically re-initializes policy parameters from the best checkpoint and increases the sampling temperature to restore exploration capabilities.

Key Contributions

1. ORBIT Framework: A scalable, seed-based automated post-training framework that replaces opaque scalar rewards with interpretable, rubric-driven feedback tailored for high-stakes medical dialogue.
2. Data-Efficient Pipeline: A context-aware rubric generation and training pipeline that integrates retrieval-augmented prompting, multi-stage filtering, and dynamic sampling. It eliminates the need for new task-specific annotations or separate medical reward models, utilizing existing expert-crafted rubrics as seeds.
3. Empirical Validation: Extensive experiments on HealthBench demonstrate that ORBIT substantially improves medical consultation quality. The authors show that the model learns case-conditioned clinical criteria rather than relying on evaluator-specific preferences or reusable rubric patterns.

Experimental Results

The authors evaluated ORBIT using the Qwen3-4B-Instruct model as the base policy.

• Performance Gains: Using only 2,000 training samples, ORBIT raised the HealthBench-Hard score from 7.0 to 27.5, a relative improvement of 293%.
• Scaling: Increasing the training data to 8k and 28k samples further improved scores to 33.6 and 37.3, respectively.
• State-of-the-Art: The 4B-parameter InfiMed-ORBIT model achieved state-of-the-art performance among comparable-size open-source models, outperforming significantly larger models (e.g., Qwen3-30B-Instruct, Baichuan-M2-32B) and approaching the performance of proprietary models like GPT-4.1.
• Distributional Shift: Analysis of inference scaling versus RL training revealed that while increasing inference rollouts (K=8 to K=40) for the baseline yielded marginal gains, ORBIT induced a structural shift in the policy distribution, clustering responses in high-score regions and significantly improving rubric pass rates.

Significance and Claims

The paper claims that ORBIT provides a data-efficient mechanism for aligning compact medical LLMs in high-stakes clinical applications. By decomposing alignment objectives into verifiable atomic rubrics, the framework offers fine-grained control without the heavy supervision costs associated with traditional reward modeling.

The authors emphasize that their approach addresses the specific limitations of medical dialogue—where correctness is insufficient and safety, empathy, and context-awareness are paramount. They assert that ORBIT narrows the performance gap between small open-source models and much larger proprietary or specialized medical models, demonstrating that carefully structured rubric-based rewards can effectively substitute for extensive parameter scaling in alignment-intensive tasks.

The paper explicitly states that ORBIT is a post-training research framework designed for assistive medical dialogue, not an autonomous clinical deployment system, and emphasizes the necessity of expert oversight in any practical application.