1. 문제 정의 (Problem Statement)

의료 영상 분석 분야에서 딥러닝 모델, 특히 비전 트랜스포머 (Vision Transformers, ViT) 는 뛰어난 성능을 보이고 있지만, 임상 현장에 안전하게 통합되기 위해서는 다음과 같은 치명적인 한계가 존재합니다.

• 과신 (Overconfidence) 및 불확실성 부재: 기존 모델은 예측 확률의 보정 (Calibration) 이 잘 되어 있지 않아, 모호한 입력이나 분포 밖 (Out-of-Distribution) 데이터에 대해서도 지나치게 자신감 있게 잘못된 예측을 하는 경향이 있습니다.
• 블랙박스 문제: 의사결정 과정에 대한 투명성이 부족하여, 왜 특정 진단이 내려졌는지에 대한 해석 가능한 근거를 제공하지 못합니다.
• 임상 데이터의 복잡성: 의료 데이터는 노이즈가 많고, 클래스 불균형이 심하며, 조직의 모호한 경계나 환자 간 변이로 인해 본질적인 불확실성이 내재되어 있습니다.

따라서 높은 정확도뿐만 아니라 신뢰할 수 있는 불확실성 정량화 (Uncertainty Quantification) 와 해석 가능한 근거를 제공하는 모델 개발이 시급합니다.

2. 제안 방법론 (Methodology)

저자들은 기존 계층적 의료 트랜스포머인 MedFormer를 확장하여 MedFormer-UR를 제안했습니다. 이 프레임워크는 불확실성을 단순한 출력값이 아닌, 학습 과정과 특징 전파를 제어하는 능동적 제어 신호로 활용합니다.

핵심 구성 요소:

1. 토큰 단위 증거 기반 불확실성 추정 (Per-Token Evidential Uncertainty Estimation):

   • 각 공간 토큰 (spatial token) 에 대해 디리클레 분포 (Dirichlet distribution) 를 사용하여 증거 (evidence) 를 모델링합니다.
   • 이를 통해 예측의 알레토릭 불확실성 (데이터 내재적 노이즈) 과 에피스테믹 불확실성 (모델 지식 부족) 을 명시적으로 분해하고 정량화합니다.
   • 불확실성은 샘플링 없이 실시간으로 계산됩니다.

2. 불확실성 유도 라우팅 및 특징 정제 (Uncertainty-Guided Routing & Refinement):

   • Uncertainty-Gated Token Refinement (UGTR): 불확실성 추정치를 기반으로 특징 맵을 동적으로 수정합니다.
   • 동작 원리:
     • 기저 경로: 표준 DSSA(Dual Sparse Selection Attention) 를 통해 전역적 맥락을 포착합니다.
     • 정제 경로: 경량 합성곱 네트워크를 통해 국소적 세부 사항 (질병 경계, 미세 질감 등) 을 강조합니다.
     • 라우팅 마스크: 토큰별 불확실성 ($\sigma_i$) 과 조직 마스크 (해부학적 정보) 를 결합하여 정제 정도를 조절합니다. 불확실성이 높은 영역에서는 정제를 억제하여 모델이 불안정한 특징을 증폭하는 것을 방지하고, 불확실성이 낮은 영역에서는 정제를 강화하여 진단에 중요한 부분을 부각시킵니다.

3. 프로토타입 기반 표현 학습 (Prototype-Based Representation Learning):

   • 기존 선형 분류기 대신 학습 가능한 클래스별 프로토타입 (Reference Patterns) 을 도입합니다.
   • 입력 특징과 프로토타입 간의 코사인 유사도를 기반으로 예측을 수행하여, "어떤 시각적 패턴과 유사하기 때문에 이 진단이 내려졌다"는 유사도 기반 추론을 가능하게 합니다.
   • 프로토타입 클러스터 손실과 다양성 정규화를 통해 임베딩 공간이 구조화되도록 하여 해석 가능성을 높입니다.

4. 학습 목적 함수 (Training Objective):

   • 분류 손실 (Cross-Entropy), 라우팅 감독 손실 (Tissue Mask 기반), 프로토타입 정규화 손실 (클러스터링 및 다양성) 을 결합한 복합 목적 함수를 사용합니다.

3. 주요 기여 (Key Contributions)

• 새로운 아키텍처 확장: MedFormer 에 증거 기반 헤드를 추가하여 토큰 단위 불확실성 분해를 가능하게 함.
• 능동적 불확실성 제어: 불확실성을 라우팅 메커니즘에 통합하여 모호한 특징의 영향을 학습 중 동적으로 줄임.
• 해석 가능한 분류: 프로토타입 학습을 통해 결정 근거를 시각적 유사성으로 설명 가능하게 함.
• 광범위한 검증: 유방촬영 (Mammography), 초음파 (Ultrasound), MRI, 조직병리학 (Histopathology) 등 4 가지 다른 의료 영상 모달리티에서 성능을 입증.

4. 실험 결과 (Experimental Results)

총 4 가지 데이터셋 (CBIS-DDSM, BUSI, IDC, Brain MRI) 에서 Baseline MedFormer 와 비교 평가되었습니다.

• 분류 성능:

  • BUSI (초음파) 및 BreakHis (조직병리): 명확한 병변 패턴이 있는 경우 정확도와 AUROC 가 크게 향상됨 (BUSI 정확도 0.800 → 0.893).
  • Brain MRI: 이미 Baseline 성능이 매우 높았으나, 프로토타입 기반 접근법으로 미세한 성능 향상 달성.
  • CBIS-DDSM (유방촬영): 낮은 대비와 조직 중첩으로 인해 어려운 작업이었으나, 정확도 (0.629 → 0.673) 와 AUROC (0.727 → 0.755) 에서 유의미한 개선.

• 보정 (Calibration) 및 신뢰성 (가장 중요한 성과):

  • ECE (Expected Calibration Error) 감소: Baseline 대비 최대 35% 까지 ECE 가 감소 (예: 0.1836 → 0.0688). 이는 모델의 예측 확률이 실제 정답률과 훨씬 잘 일치함을 의미합니다.
  • 선택적 예측 (Selective Prediction): 불확실성이 높은 경우 예측을 거부 (Reject) 하는 능력 (AURC) 이 향상되어, 임상적으로 안전한 의사결정을 지원함.
  • 매개변수 $\beta$의 영향: 불확실성 라우팅 강도 ($\beta$) 를 조절했을 때, $\beta=0.8$ 부근에서 가장 우수한 보정 성능을 보였으며, $\beta=1.0$(강한 라우팅) 은 분류 정확도는 높일 수 있으나 보정 성능이 떨어지는 트레이드오프가 존재함을 확인.

• 시각화:

  • 생성된 불확실성 지도 (Uncertainty Maps) 는 병변 경계가 모호하거나 조직이 밀집된 영역에서 높은 불확실성을 보여주어, 실제 임상적 난이도와 일치함을 입증.

5. 의의 및 결론 (Significance & Conclusion)

• 임상 신뢰성 확보: 본 연구는 의료 AI 가 단순히 "정답"을 맞추는 것을 넘어, "언제 자신이 틀릴 수 있는지" 를 인지하고 이를 의사에게 전달할 수 있어야 함을 강조합니다.
• 블랙박스 해결: 프로토타입 기반 유사도 비교를 통해 모델의 결정 근거를 시각적으로 설명 가능하게 하여, 임상 의사들의 수용성을 높입니다.
• 안전한 배포: 불확실성 기반 라우팅 메커니즘은 모호한 데이터에서 모델이 과도하게 자신감을 갖는 것을 방지하여, 오진으로 인한 임상적 위험을 줄이는 데 기여합니다.

결론적으로 MedFormer-UR는 의료 영상 분류에서 정확도와 신뢰성 (보정, 해석 가능성) 을 동시에 달성하기 위한 강력한 프레임워크로, 특히 고위험 의료 환경에서의 AI 도입을 위한 핵심 기술로 평가됩니다.