MedJudgeRAG: Option-Wise Evidence Judgment with Dynamic Knowledge Graphs for Medical MCQA

기술 요약: MedJudgeRAG

문제 정의

의학 다지선다형 질의응답(MCQA)에서 검색 증강 생성(RAG)은 외부 도메인 정보를 활용하여 대규모 언어 모델(LLM)의 파라미터 지식을 보완하는 것을 목적으로 합니다. 그러나 단순히 검색된 문서를 입력값 앞에 붙이는 "바닐라(vanilla)" RAG 방식은 종종 모델의 성능을 저하시킵니다. 이러한 성능 저하는 다음과 같은 이유로 발생합니다:

1. 파편화된 증거: 정답은 여러 문서에 흩어져 있는 증거에 의존하는 경우가 많으나, 바닐라 RAG는 이를 구조적으로 연결하지 못합니다.
2. 검증의 부재: 모델이 검색된 정보가 특정 선택지에 대해 충분한 증거가 되는지를 체계적으로 판단할 수 있는 메커니즘이 부족합니다.
3. 의사결정 과정의 결여: 검색된 증거와 최종 답안 선택 사이를 잇는 명시적인 중간 의사결정 과정이 없습니다. 이로 인해 추론 정확도가 파라미터 지식에만 의존하는 모델보다 낮아질 수 있습니다.

방법론: MedJudgeRAG

저자들은 검색된 문서들을 **동적 지식 그래프(Dynamic Knowledge Graph, KG)**로 구조화하고 **선택지별 증거 판단(option-wise evidence judgment)**을 수행하도록 설계된 프레임워크인 MedJudgeRAG를 제안합니다. 이 시스템은 교사 모델(teacher model)과 학생 모델(student model)을 포함하는 지도 미세 조정(SFT) 파이프라인을 통해 작동합니다.

1. 동적 지식 그래프 구축

모델은 가공되지 않은 텍스트를 그대로 입력하는 대신, 질문($q$)과 선택지($O$)를 바탕으로 검색된 문서($D$)로부터 구조화된 KG를 추출합니다.

• 스키마: KG는 엔티티 튜플(이름, 유형, 설명, 증거)과 관계 튜플(출처, 관계 유형, 대상, 설명, 증거)로 구성됩니다.
• 제약 조건: 엔티티 유형은 15개의 통합 의료 용어 시스템(UMLS) 시맨틱 그룹과 일치합니다. 관계는 UMLS에서 채택되었습니다.
• 생성 원칙:
  • 변별적(Discriminative): 선택지를 구분할 수 있는 정보만을 추출합니다.
  • 문서 근거적(Document-Grounded): 환각을 방지하기 위해 모든 엔티티와 관계는 반드시 검색된 문서에 의해 뒷받침되어야 합니다.
  • 증거 추적(Evidence-Traced): 모든 요소는 출처 문서의 인덱스를 인용합니다.
  • 빈 그래프 허용(Empty Graph Allowed): 문서가 무관한 경우, 노이즈를 피하기 위해 모델은 빈 KG를 출력할 수 있습니다.

2. 선택지별 증거 판단

KG 생성 후, 모델은 세 단계로 구성된 구조화된 추론 시퀀스($r$)를 수행합니다.

• 판단 생성(Verdict Generation): 각 선택지에 대해 모델은 지지됨(SUPPORTED)(직접적인 증거 존재), 반박됨(CONTRADICTED)(오답임을 입증하는 증거 존재), 또는 불충분함(INSUFFICIENT)(증거 부족) 중 하나의 판정을 내립니다. 이는 "증거 우선" 방식으로 수행되며, 모델은 결론을 내리기 전에 반드시 특정 문서나 KG 요소를 인용해야 합니다.
• 전략 선택(Strategy Selection): 판단의 조합에 따라 모델은 다음의 추론 전략을 선택합니다.
  • 근거 기반(Grounded): 특정 옵션이 지지됨 상태라면, 이를 정답으로 선택합니다.
  • 소거법(Elimination): 지지되는 옵션은 없으나 일부가 반박됨 상태인 경우, 틀린 옵션들을 제거하고 남은 옵션들에 대해 파라미터 지식을 사용하여 추론합니다.
  • 파라미터 기반(Parametric): 모든 옵션이 불충분함 상태인 경우, 모델은 내부의 파라미터 지식으로 회귀합니다.
• 최종 답안: 선택된 전략에 따라 답안이 생성됩니다.

3. 학습 전략

• 교사-학생 설정(Teacher-Student Setup): 교사 모델(GPT-5.1)은 $(q, O, D)$로부터 정답(gold answer)에 접근하지 않고도 구조화된 추론 흔적(KG + 판단 + 결정 + 답안)을 생성합니다. 이 흔적들은 학생 모델(Mistral-7B 또는 Llama-3-8B)의 SFT 타겟이 됩니다.
• 가중치 손실 함수(Weighted Loss Function): 학습에는 가중치가 적용된 교차 엔트로피 손실이 사용됩니다. KG 세그먼트의 손실 가중치($\lambda_g$)는 조절 가능하며, 추론 세그먼트의 가중치는 1로 고정됩니다. 이를 통해 모델이 완벽한 표면 형태의 KG 재구축보다는 추론 전략을 배우는 데 우선순위를 두도록 합니다.
• 디코딩 모드:
  • 명시적(Explicit): 모델이 KG를 먼저 생성한 후 추론 흔적을 생성합니다.
  • 암묵적(Implicit): 모델이 명시적인 KG 생성을 건너뛰고, KG의 내부 표현을 조건으로 하여 직접 추론 흔적을 출력합니다.

주요 결과

실험은 Mistral-7B와 Llama-3-8B 백본을 사용하여 MedQA 및 MedMCQA 벤치마크에서 수행되었으며, 동일한 검색 설정(Contriever를 통한 PubMed/교과서 상위 5개 문서)을 사용했습니다.

• Vanilla RAG와의 성능 비교: 바닐라 RAG는 파라미터 베이스라인에 비해 지속적으로 성능이 저하되었습니다(예: Llama-3는 MedQA에서 약 10.8pp 하락). MedJudgeRAG는 이 추세를 역전시켜, Llama 기준 MedQA에서 최대 +17.1pp, Med$\text{MCQA}$에서 +11.7pp의 성능 향상을 보였습니다.
• 파라미터 모델과의 성능 비교: MedJudgeRAG는 파라미터 베이스라인을 일관되게 상회하며, 외부 지식의 성공적인 통합을 입증했습니다.
• 명시적 vs 암묵적 디코딩:
  • 암묵적 디코딩(KG 출력을 생략하는 방식)이 명시적 디코딩에 비해 일반적으로 더 높은 정확도를 보였으며, 모델이 답안을 선택하지 못하는 "무효(Invalid)" 출력 횟수가 현저히 적었습니다.
  • Llama-3의 경우, 두 벤치마크 모두에서 암묵적 디코딩이 더 우수했습니다.
• KG 손실($\lambda_g$)에 대한 절제 연구(Ablation):
  • $\lambda_g = 0.0$(KG 토큰에 대한 직접적인 손실 없음)으로 설정했을 때 정답 정확도가 가장 높았습니다.
  • $\lambda_g$를 높여 모델이 KG를 강제로 재구축하게 하는 것은 목적 함수 불일치 및 토큰 예산 제한으로 인해 정답 정확도를 선형적으로 개선하지 못했으며, 오히려 성능을 저하시키기도 했습니다.
  • 이는 KG가 추론 시 필수적인 명시적 출력물이라기보다, **그래프 조건부 감독(graph-conditioned supervision)**으로서 더 효과적임을 시사합니다.

의의 및 주장

본 논문은 MedJudgeRAG가 구조화된 추론 과정을 도입함으로써 의료 MCQA에서 발생하는 바닐라 RAG의 핵심 한계를 해결한다고 주장합니다. 주요 기여는 다음과 같습니다:

1. 프레임워크 혁신: 흩어진 증거를 동적 KG로 구조화하고 옵션별 증거 판단을 수행하는 방법을 제안하여, 표준 RAG에서 나타나는 성능 저하 문제를 극복했습니다.
2. 학습 통찰: 절제 연구를 통해, 동적 KG가 명시적인 추론 시 출력물이라기보다 구조적 학습 시간 감독 신호로서 가장 효과적임을 밝혀냈습니다. 모델은 그래프 구조에 조건화되어 추론하는 법을 배우며, 이를 통해 구조적 가이드와 추론 효율성 사이의 균형을 맞춥니다.
3. 강건성: 적절한 추론 전략(판단 및 소거)을 갖춘다면, 검색된 문서가 복잡한 의료 도메인에서 LLM의 성능을 저해하는 것이 아니라 오히려 향상시킬 수 있음을 입증했습니다.

저자들은 이러한 개선이 새로운 도메인 지식을 학습했기 때문이 아니라, 동적 지식 구축과 다지선다형 질문에 특화된 추론 전략을 학습했기 때문임을 강조합니다.