InfiMed-ORBIT: Aligning LLMs on Open-Ended Complex Tasks via Rubric-Based Incremental Training

Motivation

• Known: RL with verifiable rewards(RLVR)은 정답이 명확한 코드 생성 등에서 효과적이며, medical domain에서는 HealthBench와 같은 전문가 작성 rubric 기반 평가 체계가 clinical reasoning 평가를 발전시켜왔다.
• Gap: Open-ended medical dialogue는 feedback이 모호하고 context-dependent하여 단일 scalar reward로 요약하기 어렵고, 기존 reward model 기반 접근은 구축 비용이 크고 domain-dependent하며 reward hacking에 취약하다는 문제가 있다.
• Why: 단일 scalar reward 대신 해석 가능한 rubric 기반 평가로 전환함으로써, 별도의 task-specific judge fine-tuning이나 고비용 reward model 없이도 고위험(high-stakes) 의료 대화에서 안전하고 투명한 RL 정렬이 가능함을 보여준다.
• Approach: ORBIT은 소수의 전문가 작성 rubric seed로부터 retrieval-augmented in-context prompting을 통해 case-conditioned rubric을 동적으로 생성하고, 이를 general-purpose instruction-following LLM judge와 결합하여 GRPO 기반 incremental RL의 보상으로 사용한다.

Achievement

Figure 3. Performance comparison of ORBIT models across multiple clinical dimensions. Results are organized by clinical

1. HealthBench-Hard 대폭 향상: 단 2k개의 학습 샘플만으로 Qwen3-4B-Instruct의 HealthBench-Hard 점수를 7.0에서 27.5로 끌어올려 동급 크기 오픈소스 모델 중 state-of-the-art를 달성했다.
2. Rubric coverage 확장에 따른 추가 성능 향상: 더 큰 rubric set을 사용할 경우 37.3까지 성능이 상승하여, 자신보다 최대 8배 큰 오픈소스 모델을 능가했다.
3. Judge-sensitivity 및 rubric-similarity 분석을 통한 검증: ORBIT이 evaluator-specific preference나 재사용 가능한 rubric pattern에 의존하는 것이 아니라 실제로 case-conditioned clinical criteria를 학습함을 실증적으로 확인했다.

How

Figure 2. Overview of ORBIT, a rubric-guided reinforcement learning framework for aligning LLMs on open-ended medical di

• 실제 임상 대화 데이터셋으로부터 multi-turn medical dialogue query를 구성하고 length random split을 통해 dialogue simulation을 수행한다.
• Dialogue/rubric embedding을 기반으로 (dialogue, rubrics) pair를 indexing하여 rubric pool과 dialogue pool을 구축하고, 새로운 dialogue query에 대해 case retrieval을 수행한다.
• Rubric Generation Model이 system/task prompt와 in-context learning을 통해 후보 rubric을 생성하고, LLM rewrite 및 reranker, positive/negative rubric 필터링을 거쳐 최종 Q rubrics를 선별한다.
• Multi-stage filtering(Rubrics Filter, Sample Filter)과 pass@k 기반 dynamic sampling을 적용해 계산 효율을 높이면서 안정적인 incremental RL을 구현한다.
• Judge Model이 case-specific rubrics를 기준으로 policy(reference model)의 응답을 평가하여 reward를 산출하고, group computation과 KL divergence loss를 포함한 policy gradient update(GRPO)로 InfiMed-ORBIT-4B를 학습한다.

Originality

• 단일 scalar reward나 별도의 fine-tuned reward model 없이, 소수의 전문가 rubric seed만으로 case-conditioned rubric을 자동 생성하는 retrieval-augmented 파이프라인을 제안했다.
• General-purpose instruction-following LLM을 judge로 그대로 활용하여 task-specific judge fine-tuning을 회피한 점이 기존 domain-adaptive reward 접근과 차별화된다.
• Rubric coverage를 incremental하게 확장하며 RL을 진행하는 incremental training 개념을 도입해, 데이터 효율성과 성능을 동시에 확보했다.

Limitation & Further Study

• 의료 도메인, 특히 HealthBench라는 특정 벤치마크에 집중되어 있어 다른 open-ended 도메인(심리상담, 사회적 상호작용 등)으로의 일반화 가능성에 대한 실증이 부족하다.
• General-purpose LLM을 judge로 사용함에 따라 judge 모델 자체의 편향이나 한계가 rubric 평가 품질에 영향을 줄 수 있으며, judge-sensitivity 분석이 다양한 judge 모델 규모/종류에 대해 충분히 폭넓게 이루어졌는지 추가 검증이 필요하다.
• Rubric 생성 파이프라인의 retrieval 품질(embedding, reranker)에 성능이 크게 의존할 수 있어, retrieval 실패 시 rubric coverage나 정확도가 저하될 위험에 대한 분석이 더 필요하다.
• 향후 연구로 다른 고위험 open-ended 도메인에 대한 확장 적용, 더 다양한 judge/reward model과의 결합, 실제 임상 환경에서의 안전성 검증 등이 필요하다.

Evaluation

총평: Rubric 기반 incremental RL이라는 아이디어를 통해 opaque scalar reward의 한계를 실질적으로 해결하면서 소량 데이터로 큰 성능 향상을 보여준 인상적인 연구이며, 의료를 넘어 다른 open-ended 고위험 도메인으로의 확장 가능성이 기대된다.