MediX-R1: 개방형 의료 강화 학습 프레임워크

기존 의료 AI는 객관식 문제 풀이에 치중되어 실제 임상 현장의 자유로운 질의응답에 한계가 있었다. 이 논문은 복합 보상 체계를 갖춘 강화 학습을 통해 AI가 의학적 근거를 바탕으로 사고하고, 전문 용어의 변형을 이해하며, 정확한 진단 보고서를 작성할 수 있게 한다.

기존 의료 AI 학습은 정답과 모델 출력이 글자 하나까지 일치해야 점수를 주는 방식(Exact Match)이나 BLEU/ROUGE 같은 통계적 유사도에 의존했다. 하지만 의료 현장에서는 '폐렴 가능성'과 '폐의 염증 소견'이 같은 의미임에도 불구하고, 이러한 방식으로는 모델의 정답을 오답으로 처리하거나 복잡한 의학적 추론 과정을 평가하기 어려웠다.

MediX-R1은 이를 해결하기 위해 강화 학습의 보상 함수를 다각화한다. LLM을 판사로 세워 의미적 정답 여부를 가리고, 의료 전문 임베딩 모델을 통해 용어의 변형을 포착하며, 태그를 강제해 사고 과정을 투명하게 만든다. 이는 마치 의대생이 단순히 답만 맞히는 게 아니라, 진단 근거를 논리적으로 설명하고 전문 용어를 적재적소에 사용하도록 훈련받는 과정과 유사하다.

결과적으로 모델은 단순히 확률적으로 높은 단어를 나열하는 것이 아니라, 입력된 의료 영상(X-ray, MRI 등)의 모달리티를 정확히 인식하고 그에 맞는 의학적 지식을 인출하여 답변하게 된다. 30B 규모의 모델에서 평균 73.6%의 정확도를 기록하며, 훨씬 더 많은 데이터를 사용한 기존 모델들보다 뛰어난 임상 보고서 생성 능력을 보여준다.

MediX-R1은 GRPO(Group Relative Policy Optimization)를 기반으로 한 단일 단계 강화 학습을 수행한다. 입력값 v(이미지 I + 질문 q)에 대해 G개의 답변 후보를 생성하고, 각 답변의 보상 ri를 계산한 뒤 그룹 내 상대적 이득(Advantage) Ai를 산출한다. Ai = (ri - mean(rj)) / std(rj) [개별 보상 ri와 그룹 평균/표준편차 입력 → 정규화 연산 → 이득 Ai 출력 → 그룹 내 상대적 우수성 의미] 과정을 통해 별도의 가치 함수 없이도 안정적인 학습이 가능하다.

보상 함수는 네 가지 요소의 가중 합으로 구성된다. Rllm은 Qwen3-4B 모델을 사용하여 정답과 모델 출력이 의미적으로 일치하는지 YES/NO로 판정하며, Rembed는 MedEmbed-large 모델로 코사인 유사도를 계산해 0.8 이상의 값을 가질 때 보상을 부여한다. [예측 답변과 참조 답변 임베딩 입력 → 코사인 유사도 연산 → 임계값 비교 → 이진 보상 출력] 순으로 작동하여 전문 용어의 유의어를 수용한다.

구조적 제약을 위해 Rformat은 정규표현식을 사용하여 와 태그의 존재 여부를 확인하고, Rmodality는 답변 시작 전 올바른 영상 종류(예: <X_RAY>)를 명시했는지 검사한다. 이러한 복합 보상 체계는 모델이 보상 해킹에 빠지지 않고 임상적으로 유효한 답변을 생성하도록 유도한다.

MediX-R1 30B 모델은 텍스트 전용(LLM) 및 이미지+텍스트(VLM) 벤치마크 통합 평가에서 평균 73.6%의 정확도를 기록하며 SOTA를 달성했다. 특히 MedGemma 27B(68.4%)와 비교했을 때 훨씬 적은 학습 데이터를 사용하고도 더 높은 성능을 보였다.

개방형 임상 과제인 MIMIC-CXR 요약 및 보고서 생성에서 각각 0.786, 0.350의 점수를 기록하며 기존 모델들을 압도했다. 이는 모델이 단순 지식 암기를 넘어 복잡한 의료 영상을 해석하고 서술하는 능력이 뛰어남을 시사한다.

Ablation Study 결과, LLM 판정 보상과 임베딩 보상을 결합했을 때 학습 곡선이 가장 안정적이었으며 보상 해킹 현상이 현저히 줄어들었다. 또한 GRPO 외에도 DAPO, GSPO 등 다양한 그룹 기반 RL 알고리즘에서도 일관된 성능 향상이 확인됐다.

아키텍처는 비전 인코더와 언어 모델 백본이 융합된 멀티모달 구조를 사용하며, Qwen2.5-VL, Qwen3-VL 등을 백본으로 실험했다. 학습은 8개의 A100(80GB) GPU에서 약 25시간 동안 진행되었으며, FSDP(Fully Sharded Data Parallel) 전략을 적용했다.

핵심 알고리즘인 GRPO는 PPO의 클리핑 메커니즘을 유지하면서도 그룹 내 통계를 활용해 기준점(Baseline)을 설정한다. 손실 함수 J_GRPO(theta)는 중요도 샘플링 비율 rho_i(theta)와 이득 Ai의 곱을 최소화하며, KL Divergence를 통해 참조 정책 pi_ref로부터 과도하게 벗어나는 것을 방지한다.

보상 설계 시 w_fmt=0.10, w_mod=0.045, w_llm=0.5175, w_emb=0.3375의 가중치를 사용하여 의미적 정확성에 가장 높은 비중을 두었다. 임베딩 보상에는 gating 메커니즘을 도입하여 너무 짧거나 문장 부호만 있는 답변에 보상이 주어지는 것을 차단했다.

모델이 존재하지 않는 소견을 만들어내는 환각(Hallucination) 가능성이 여전히 존재하며, 참조 기반 LLM 판정기 자체가 편향되거나 오판할 위험이 있다. 또한 공개된 데이터셋 위주로 학습되어 특정 희귀 질환에 대한 일반화 성능은 검증이 더 필요하다.