기술 요약: 방사선 보고서 생성을 위한 집합-거리 보상 (SDR)

문제 정의

시각-언어 모델(VLM)을 이용한 흉부 X-선 보고서 생성은 강화 학습(RL) 사후 학습(post-training) 시, 특히 보상 함수 설계 측면에서 독특한 과제를 안겨줍니다. 인과적 사슬의 중간 단계를 검증하는 프로세스 보상 모델(process reward models)을 사용하는 일반적인 RL 접근 방식은 추론 작업에 적합합니다. 그러나 방사선 보고서는 순차적이거나 인과적인 추론 구조를 따르지 않고, 대신 순서가 정해지지 않은 독립적인 소견들(unordered and orthogonal findings)로 구성됩니다 (예: "심장 크기가 커졌을 수 있음"과 "우측 폐는 깨끗함"은 서로 독립적인 관찰임).

결과적으로, 정확히 일치하는 정확도에 기반한 표준 이산(discrete) 보상은 이 도메인과 호환되지 않습니다. 이러한 방식은 보고서가 올바른 소견을 포함하고 있더라도 순서가 다르거나 약간의 구절 변형이 있는 경우 이를 패널티로 부여하여 노이즈를 유발하며, 이는 소견의 순서보다 집합(set) 자체가 중요하다는 임상적 실재를 포착하지 못하는 결과를 초래합니다. 또한, 프로세스 보상에 필요한 단계별 주석(step-level annotations)은 대개 가용하지 않으며, 외부 소스를 통해 각 단계를 검증하는 것은 계산 비용이 많이 듭니다.

방법론

저자들은 방사선 보고서를 순차적인 텍스트가 아닌, 문장 임베딩의 순서가 정해지지 않은 집합으로 취급하는 집합-거리 보상(Set-Distance Rewards, SDR) 프레임워크를 제안합니다. 이 방법론은 사후 학습(GRPO)과 추론 시간 스케일링(inference-time scaling)의 두 단계로 작동합니다.

1. 문장 수준 표현

흉부 X-선 보고서 $y$는 두 섹션인 소견(Findings)($y_F$)과 결론(Impression)($y_I$)으로 나뉩니다. 각 섹션은 개별 문장으로 분할됩니다. 동결된 사전 학습된 문장 트랜스포머(구체적으로 all-mpnet-base-v2)는 각 문장 $s$를 고정된 차원의 벡터 $e \in \mathbb{R}^d$로 임베딩합니다. 따라서 보고서는 다음과 같이 두 개의 순서가 정해지지 않은 임베딩 집합으로 표현됩니다:
$$E_F(y) = \{E_\phi(s^F_i) \mid 1 \le i \le n_F\}, \quad E_I(y) = \{E_\phi(s^I_j) \mid 1 \le j \le n_I\}$$
이 표현은 치환 불변성(permutation-invariant)을 가지며, 이는 소견의 순서가 진단적 의미를 갖지 않는다는 임상적 직관과 일치합니다.

2. 집합 간 거리 보상 (사후 학습)

그룹 상대 정책 최적화(GRPO) 과정 동안, 모델은 보고서 $\hat{y}$를 생성합니다. 의미론적 보상은 생성된 임베딩 집합과 참조(정답) 임베딩 집합 사이의 거리를 계산하여 산출됩니다.

• 거리 메트릭: 저자들은 단위 노름(unit-norm) 임베딩 간의 코사인 거리에 기반하여, $[0, 1]$로 정규화된 대칭적 집합 간 거리인 **샴 거리(Chamfer distance)**와 **하우스도르프 거리(Hausdorff distance)**를 활용합니다.
  • Chamfer: 최근접 이웃 비용을 평균화하며, "부드러운 커버리지(soft coverage)" 보상 역할을 합니다.
  • Hausdorff: 최대 최근접 이웃 비용을 취하며, 중요한 소견을 놓치는 것에 대한 "최악의 경우(worst-case)" 패널티 역할을 합니다.
• 보상 공식: 총 보상 $R(\hat{y}, y)$는 바이너리 형식 보상(출력이 구조화된 템플릿을 따르는지 확인)과 집합 거리로부터 유도된 연속적 의미론적 보상의 가중 합으로 구성됩니다:
  $$R(\hat{y}, y) = \lambda_{fmt} R_{fmt}(\hat{y}) + \lambda_{sem} (1 - D(E_S(\hat{y}), E_S(y)))$$
  여기서 $D$는 집합 거리입니다. 이는 부분적인 정확성에 보상을 주고 순서가 없는 소견들을 처리할 수 있는 연속적이고 미분 가능한 신호를 제공합니다.

3. 추론 시간 스케일링 및 프루닝(Pruning)

동일한 집합 거리 메트릭이 그래디언트 업데이트 없이 추론 시에도 적용됩니다:

• Best-of-N 선택: 주어진 이미지에 대해 모델은 $K$개의 후보 보고서를 생성합니다. 각 후보의 임베딩 집합은 훈련 코퍼스의 사전 계산된 임베딩 집합과 비교됩니다. 훈련 분포와의 거리가 최소인 후보가 선택됩니다. 이는 임상적으로 타당한 출력을 위한 그래디언트 프리(gradient-free) 사전 확률(prior) 역할을 합니다.
• 스트리밍 프루닝(Streaming Pruning): 계산 비용을 줄이기 위해, 시스템은 디코딩 중 생성되는 부분적 생성물의 실행 중인 집합 거리를 계산합니다. 부분 거리가 임계값을 초과하는 후보는 조기에 제거(pruning)됩니다. 이 "프루닝 후 디코딩(prune-then-decode)" 루프는 전체 Best-of-N 선택과 유사한 품질을 유지하면서 토큰 생성을 50% 이상 줄여줍니다.

주요 기여

1. GRPO를 위한 집합-거리 보상 함수: 본 논문은 방사선 보고서를 순서가 없는 집합으로 취급함으로써 프로세스 보상 모델링의 불가능 문제를 해결합니다. 샴 거리와 하우스도르프 거리를 사용하여 연속적인 보상을 제공함으로써, 세 가지 VLM 백본(Qwen3-VL-2B/4B, Gemma3-4B)과 두 개의 데이터셋(MIMIC-CXR, ReXGradient)에서 일관되게 우수한 성능을 보였습니다.
2. 집합-거리 유도 테스트 타임 선택: 저자들은 동일한 집합 거리 메트릭이 미세 조정된 모델뿐만 아니라 폐쇄형 소스 범용 LLM(Mistral-Small, Gemini-2.5 Flash-Lite, GPT-4o-mini, GPT-5 mini)에 대해서도 선택 기준으로 사용될 수 있음을 입증했습니다. 이 접근 방식은 모델 그래디언트나 도메인 특화 미세 조정 없이도 무작위 선택보다 향상된 성능을 보여줍니다.
3. 거리 기반 온더플라이(On-the-fly) 프루닝: 생성 도중에 낮은 점수의 후보를 제거하는 방법을 도입했습니다. 이는 생성되는 총 토큰 수를 50% 이상 줄이면서도 최종 선택의 품질을 유지하며, 더 효율적인 테스트 타임 스케일링을 제공합니다.

실험 결과

평가는 BERTScore, RadGraph F1, CheXbert F1을 포함한 핵심 지표를 다룹니다.

• 사후 학습 성능:
  • ReXGradient 데이터셋에서 샴 기반 보상이 모든 모델과 지표에서 가장 높은 성능을 달-성했습니다.
  • MIMIC-CXR에서 샴 및 하우스도르프 보상 모두 가장 높은 평균 성능에 도달했습니다.
  • SFT 및 이산 일치 GRPO 대비 상대적 개선 폭은 상당했습니다: 평균적으로 6.80%(BERTScore), 7.82%(RadGraph F1), 4.45%(CheXbert F1)의 개선을 보였습니다.
• 추론 시간 선택:
  • 폐쇄형 LLM에 집합-거리 선택을 적용했을 때 무작위 선택 대비 상당한 이득을 얻었습니다. 예를 들어, BERTScore의 경우 모델과 프롬프트 템플릿에 따라 +11.2%에서 +18.7% 사이의 개선을 보였습니다.
  • 이 방법은 특정 의료 도메인에 대한 미세 조정이 이루어지지 않은 범용 모델(예: GPT-4o mini, Gemini)의 성능도 성공적으로 향상시켰습니다.
• 프루닝 효율성:
  • 프루닝 전략은 생성 토큰을 **42.1%에서 60.1%**까지 절감했습니다.
  • 프루닝된 생성물은 전체 Best-of-N 선택과 대등한 품질을 유지했으며, 무작위 선택 대비 평균 상대적 개선율은 +12.7%(BS-F1), +17.1%(RG-F1), +6.2%(CXB-F1)였습니다.

의의 및 주장

본 논문은 집합-거리 보상을 흉부 X-선 보고서 생성을 위한 사후 학습 및 테스트 타임 스케일링을 위한 통합된 신호로 확립한다고 주장합니다.

• 통합된 신호: 동일한 수학적 공식(집합 간 거리)이 RL 사후 학습 과정을 효과적으로 유도할 뿐만 아니라, 추론 시점에 강력한 그래디언트 프리 선택 기준으로서 기능합니다.
• 도메인 적응: 방사선 소견의 순서가 정해지지 않은 특성을 인정함으로써, 순차적 제약으로 인해 발생하는 노이즈를 피하며, 이는 소견들이 흔히 독립적인 성격을 띠는 의료 보고서 생성에 특히 적합합니다.
• 효율성: 고품질의 보고서 생성이 반드시 모든 후보에 대해 모든 토큰을 생성할 필요는 없음을 보여줍니다. 집합 거리에 기반한 조기 프루닝은 추론 비용을 줄이는 실행 가능한 경로을 제시합니다.
• 일반화 가능성: 이 선택 규칙을 폐쇄형 LLM에 적용할 수 있다는 점은 값비싼 도메인 특화 미세 조정이나 모델 가중치에 대한 접근 없이도 임상 AI의 출력을 개선할 수 있는 경로를 시사합니다.

저자들은 자신들의 접근 방식이 방사선 보고서의 임상적 구조와 일치하는 실용적이고 연속적인 보상 신호를 제공하며, 기존의 이산적 또는 순차적 보상 메커니즘보다 우수하다고 결론짓습니다.