Reinforcement Learning Improves LLM Accuracy and Reasoning in Disease Classification from Radiology Reports

논문 요약: 방사선 보고서 기반 질병 분류에서 강화 학습을 통한 LLM 정확도 및 추론 능력 향상

1. 연구 배경 및 문제 제기 (Problem)

• 배경: 방사선 보고서에서 질병을 정확하게 분류하는 것은 임상 의사결정에 필수적입니다. 최근 경량화 된 대형 언어 모델 (Lightweight LLMs) 을 방사선 보고서 분석에 적용하려는 시도가 늘고 있으며, 이는 온프레미스 배포와 데이터 프라이버시 보호 측면에서 유리합니다.
• 문제점:
  • 기존 **지도 학습 미세 조정 (Supervised Fine-Tuning, SFT)**은 모델의 분류 정확도를 높일 수 있지만, 종종 모델이 **명확하고 일관된 추론 과정 (Reasoning Path)**을 생성하는 능력을 저하시킵니다.
  • 의료 분야에서는 투명성과 설명 가능성 (Explainability) 이 매우 중요하므로, 추론 능력이 결여된 모델은 임상 현장에서의 신뢰를 얻기 어렵습니다.
  • 기존에는 고품질의 추론 데이터를 대량으로 확보하기 어렵거나, 추론 라벨링에 막대한 비용이 들어가는 문제가 있었습니다.

2. 제안 방법론 (Methodology)

저자들은 **두 단계 (Two-stage)**의 학습 및 추론 프레임워크를 제안하여 SFT 의 한계를 극복하고 추론 능력을 복원/향상시켰습니다.

• 1 단계: 질병 라벨 기반 지도 학습 미세 조정 (SFT)

  • 입력: 방사선 보고서 텍스트.
  • 목표: 질병 라벨 (Ground Truth) 만을 사용하여 모델을 방사선 도메인에 적응시킵니다.
  • 특징: 이 단계에서는 명시적인 추론 (Reasoning) 라벨이 없으므로, 모델은 단순히 질병 목록을 출력하도록 학습됩니다.

• 2 단계: 그룹 상대적 정책 최적화 (Group Relative Policy Optimization, GRPO) 적용

  • 목적: 별도의 추론 라벨링 없이 모델의 분류 정확도와 추론 생성 능력을 동시에 최적화합니다.
  • 보상 함수 (Reward Function):
    1. 분류 정확도 (Accuracy): 예측된 질병과 정답 라벨 간의 정밀도 (Precision) 와 재현율 (Recall) 을 기반으로 계산.
    2. 포맷 및 추론 형식 (Format & Reasoning): 모델이 <reasoning> 태그 안에 근거를 명시적으로 포함하고 있는지 확인. 추론이 없거나 단순히 질병 목록만 반복하면 보상을 0 으로 부여하여 의미 있는 추론 생성을 유도합니다.
  • 학습 데이터: 1 단계에서 사용한 동일한 2,000 개의 보고서 데이터셋을 재사용합니다.

• 추론 단계 (Inference)

  • 다중 추론 및 다수결 투표 (Ensemble & Majority Voting): 각 보고서에 대해 5 번의 독립적인 추론을 수행하고, 최종 질병 분류는 다수결 투표로 결정하여 안정성을 높입니다.
  • 추론 요약 (Reasoning Summarization): 생성된 5 개의 개별 추론과 최종 질병 목록을 입력으로 받아, 원본 경량 LLM 이 이를 하나의 일관된 임상적 근거 보고서로 요약하도록 프롬프팅합니다.

3. 주요 실험 및 결과 (Results)

• 데이터셋: MIMIC-CXR, NIH-CXR, MIDRC 등 3 개의 전문가 (방사선과 전문의) 가 라벨링한 데이터셋을 사용했습니다.
• 평가 모델: LLaMA 3.1-8B, Qwen 2.5-3B, Phi-3 Min-128K 등 3 가지 경량 LLM 을 평가 대상으로 선정했습니다.
• 주요 성과:
  • 분류 성능 향상: SFT 만 적용한 모델보다 SFT+GRPO 모델이 9 개 코호트 중 8 개에서 미세 평균 F1 점수 (Micro-F1) 를 추가로 향상시켰습니다. 특히 Qwen 모델에서 MIMIC-CXR 데이터셋 기준 13.2% 의 큰 향상을 보였습니다.
  • 추론 능력 회복: SFT 단계에서 LLaMA 모델의 추론 능력이 완전히 소실되었으나 (Catastrophic forgetting), GRPO 적용 후 **추론 재현율 (Reasoning Recall)**과 **포괄성 (Comprehensiveness)**이 크게 개선되었습니다.
  • 전문가 평가: 두 명의 방사선과 전문의가 무작위 추출된 50 건의 보고서에 대해 추론의 정확성을 평가한 결과, GRPO 를 적용한 모델이 SFT 만 적용된 모델보다 더 일관되고 임상적으로 타당한 추론을 생성하는 것으로 나타났습니다.
  • 비교 분석:
    • 기존 텍스트 분류기 (BioClinicalBERT) 와 비교 시, 경량 LLM 이 추론 생성 능력에서 우위를 점했습니다.
    • 소량의 추론 데이터로 추가 SFT 를 수행한 경우 (SFT disease + SFT reasoning) 보다, GRPO 를 적용한 경우 (SFT + GRPO) 가 분류 정확도 면에서 더 우수하거나 유사한 성능을 보였습니다. 이는 명시적인 추론 라벨 없이도 강화 학습을 통해 추론 능력을 향상시킬 수 있음을 시사합니다.

4. 핵심 기여 (Key Contributions)

1. 라벨 없는 추론 강화: 명시적인 추론 라벨링 없이도 GRPO 와 규칙 기반 보상 함수를 활용하여 LLM 의 추론 능력을 복원하고 향상시키는 새로운 파이프라인을 제시했습니다.
2. 정확도와 설명 가능성의 동시 달성: 분류 정확도를 높이는 동시에 모델이 근거를 제시하는 능력을 유지/향상시켜, 의료 AI 의 임상 적용 장벽을 낮췄습니다.
3. 효율적인 학습 전략: 대규모 추론 데이터셋 구축 없이, 기존 질병 라벨 데이터만으로도 GRPO 를 통해 추론 능력을 개선할 수 있음을 입증했습니다.

5. 의의 및 한계 (Significance & Limitations)

• 의의:
  • 의료 분야에서 "블랙박스" 모델의 신뢰성 문제를 해결하는 실용적인 솔루션을 제공합니다.
  • 경량 LLM 을 저사양 환경에서도 고품질의 진단 보조 도구로 활용할 수 있는 가능성을 열었습니다.
  • 추론 능력을 강화함으로써 모델이 잘못된 예측을 할 경우 그 이유를 설명할 수 있어, 임상 의사들의 수용성을 높입니다.
• 한계:
  • 보상 함수가 분류 정확도와 형식 위주로 설계되어 추론의 길이나 논리적 일관성 등 더 세밀한 속성은 직접적으로 제어하지 못합니다.
  • 추론 품질 평가에 LLM (GPT-4o, Gemini) 을 '판심자 (Judge)'로 사용했으므로, 평가 자체의 편향 (Bias) 이 존재할 수 있습니다.
  • 모델의 예측에 대한 불확실성 (Uncertainty) 정량화나 확률 보정 (Calibration) 은 향후 과제로 남았습니다.

결론적으로, 이 연구는 강화 학습 (GRPO) 을 활용하여 경량 LLM 의 질병 분류 정확도와 추론 능력을 동시에 향상시키는 효과적인 방법을 제시함으로써, 방사선 보고서 분석을 위한 AI 도구의 실용성과 신뢰성을 크게 높였습니다.