Uncertainty Estimation in Pathology Foundation Models via Deep Mutual Learning

기술 요약: 심층 상호 학습을 통한 병리 파운데이션 모델의 불확실성 추정

문제 정의
병리 파운데이션 모델(Pathology Foundation Models, PFMs)은 생존 분석, 바이오마커 발견, 종양 아형 분류와 같은 과업을 위해 전체 슬라이드 이미지(WSI)로부터 표현형을 추출하는 데 강력한 일반화 능력을 입증해 왔다. 그러나 이들의 일상적인 임상 적용은 다음과 같은 두 가지 결정적인 한계로 인해 저해되고 있다:

1. 신뢰도 추정치의 부재: 기존의 PFM 파이프라인은 일반적으로 명시적인 신뢰도나 불확실성 측정치 없이 예측값만을 출력한다. 임상 환경에서 모델은 정확할 뿐만 아니라, 안전한 의사결정을 지원하기 위해 자신의 예측이 불확실할 때 이를 나타낼 수 있어야 한다.
2. 과업 의존적 성능: 모든 다운스트림 과업, 데이터셋 및 임상 맥락에 대해 보편적으로 우월한 단일 PFM은 존재하지 않는다. 최적의 모델을 선택하는 것은 도전적인 과제이며, 단일 모델에 의존하는 것은 성능이 저하되거나 실패 가능성이 높은 사례에서 과도한 확신을 가질 위험이 있다.

현재의 융합 방법들은 성능을 개선할 수는 있지만, 명시적이고 과업 불가지론적인(task-agnostic) 불확실성 추정치를 제공하지는 못한다. 반대로, 딥 앙상블(deep ensembles)이나 몬테카를로 드롭아웃(Monte Carlo dropout)과 같은 불확실성 정량화 방법들은 PFM 앙상블에 특화되어 설계되지 않았거나 특정 과업(예: 글리오블라스트종 진단)에 국한되어 있다.

방법론: DICE 프레임워크
저자들은 $K$개의 동결된 PFM을 앙상블하여 불확실성을 추정하고 다운스트림 성능을 향상시키는 플러그 앤 플레이(plug-and-play) 프레임워크인 DICE(Disagreement-Informed Coordination of Experts)를 제안한다. 핵심 직관은 이질적인 PFM 간의 불일치가 최적화 노이즈에 의해 지배되지 않는다면, 유의미한 불확실성 신호 역할을 한다는 것이다.

프레임워크는 다음과 같이 작동한다:

1. 구조: 시스템은 $K$개의 동결된 PFM(예: Virchow2, UNI2-h, H-optimus-1, CONCHv1.5, Hibou-L)을 활용한다. 각 PFM은 하나의 "전문가(expert)" $E^{(k)}$로 작용하며, 독립적인 다중 인스턴스 학습(Multiple Instance Learning, MIL) 집계기(구체적으로 ABMIL) 및 분류 헤드와 쌍을 이룬다.
2. 학습 목표: 전문가들이 노이즈가 아닌 의미 있는 불확실성을 반영하도록 하기 위해, 세 가지 목표를 통해 학습 과정에서 전문가들을 결합한다:
   • 지도 학습 손실 ($L_{sup}$): 정확한 슬라이드 수준 분류를 보장하기 위한 표준 교차 엔트로피 손실.
   • 심층 상호 학습 ($L_{DML}$): 예측 공간 정렬 목표로서, 각 전문가는 자신의 예측 분포와 나머지 $K-1$개 전문가의 예측 분포 사이의 쿨백-라이블러(Kullback–Leibler, KL) 발산을 최소화한다. 이는 정답과의 정렬을 유지하면서도 합의를 장려한다.
   • 그라미안 정규화 ($L_{Gram}$): 그라미안 측도(Gramian measure)에 기반한 표현 공간 정렬 목표이다. 이는 전문가들의 슬라이드 수준 임베딩이 낮은 볼륨의 부분 공간을 형성하도록 유도하여, 불확지성을 가장할 수 있는 기하학적 불일치를 방지한다.
   • 전체 손실: 학습 목표는 다음을 결합한다: $L = L_{sup} + L_{DML} + \lambda_{Gram}(t) L_{Gram}$, 여기서 $\lambda_{Gram}$은 웜업 스케줄을 따른다.
3. 추론 및 불확실성 추정:
   • 불확실성 대리 지표: 추론 시, 앙상블 예측은 $K$개 전문가 분포의 평균이다. 전체 예측 불확실성은 우연적(aleatoric, 데이터) 성분과 인식적(epistemic, 모델) 성분으로 분해된다. 저자들은 이론적으로 인식적 불확실성(잔여 불일치)이 집계된 DML 손실에 의해 상한이 결정됨을 증명한다. 따라서 학습 중 DML 손실을 최소화하는 것은 모델 불확실성의 상한을 축소시킨다.
   • 국소화(Localization): 패치 수준의 병변 국소화는 모든 전문가의 어텐션 가중치에 대한 합의를 계산함으로써 비지도 방식으로 이루어진다.

주요 기여

1. 최초의 과업 불가지론적 프레임워크: DICE는 병리 계산 분야에서 과업 불가지론적인 불확실성 추정을 위해 PFM 앙상블을 활용한 최초의 프레임워크이다.
2. 이론적 토대: 본 논문은 학습 시 DML 정렬 목표를 최소화하는 것이 앙상블의 인식적 불확실성에 대한 상한을 축소한다는 이론적 증명을 제공하여, 잔여 불일치를 신뢰할 수 있는 불확실성 대리 지표로 사용하는 것을 검증한다.
3. 실증적 성능: 광범한 실험을 통해 DICE가 분류, 캘리브레이션(calibration) 및 비지도 병변 국소화에서 단일 PFM 및 퓨전 베이스라인보다 정확하게 실패 가능성이 높은 샘플(내부 및 외부 분포 모두)을 식별할 뿐만 아니라, SOTA 단일 PFM과 대등하거나 더 나은 성능을 보임을 입증했다.

실험 결과
프레임워크는 세 가지 약지도 학습 WSI 벤치마크인 PANDA(전립선암 등급), CAMELYON16, CAMELYON17(유방 림프절 전이 탐지)에서 평가되었다.

• 분류 및 캘리브레이션: DICE는 대부분의 데이터셋에서 가장 높은 F1 점수와 가장 낮은 음의 로그 가능도(Negative Log-Likelihood, NLL)를 달 기록했다. PANDA 데이터셋에서 DICE는 최고의 단일 PFM보다 3.8%포인트, 최고의 퓨전 베심라인보다 2.8%포인트 높은 F1 점수를 개선했다.
• 불확실성 추정: 선택적 예측 실험에서 DICE는 불확실성이 높은 슬라이드를 효과적으로 유예함으로써 가장 낮은 잔여 오류율을 달성했다. DICE는 몬테카를로 드롭아웃 및 기타 퓨전 방법들과 비교하여 정답과 오답 예측 사이의 분리 능력(Cliff's $\Delta$로 측정)에서 우수한 성능을 보였다. 이러한 능력은 코호트 변화(예: CAMELYON16에서 학습하고 CAMELYON17에서 테스트) 상황에서도 유지되었다.
• 국소화: 어텐션 맵을 집계함으로써, DICE는 단일 모델보다 경쟁력 있거나 더 우수한 비지도 패치 수준 국소화를 제공했으며, 특히 주변 기질 조직이 아닌 주석 처리된 종양 영역 내에 어텐션을 집중시키는 데 탁월했다.
• 절제 연구(Ablation Studies): 프레임워크는 MIL 집계기(DSMIL, TransMIL)의 선택에 견고함을 보였으며, 특징 공간이 충분히 커버되면 전문가 수를 늘려도 수익 체감의 법칙이 나타남을 보여주었다. 이는 단순한 앙상블 크기보다 전문가 간의 상보성이 중요함을 강조한다.

의의 및 주장
본 논문은 DICE를 불확실성을 인지하는 의사결정 지원 시스템으로 전환하기 위한 중요한 단계로 자리매김한다. 이론적으로 근거가 있는 PFM 융합 접근 방식을 구축함으로써, DICE는 현재의 "블랙박스" 성격을 해결하고, 오류 발생 가능성이 높은 사례를 식별하여 임상의에게 안전하게 판단을 유예할 수 있는 메커니즘을 제공한다. 저자들은 DICE가 국소화나 신뢰도 캘리브레이션을 위한 명시적인 지도 학습 없이도 신뢰할 수 있는 불확실성 추정치를 제공함으로써, 임상 환경에서 PFM을 더욱 신뢰할 수 있게 만든다고 주장한다.

저자들은 $K$개의 동결된 모델로부터 특징을 추출하는 데 드는 계산 비용(비록 학습 비용은 엔드 투 엔드 미세 조정보다 낮지만)과, 불확실성 추정에 필요한 불일치 자체를 억제하지 않도록 정렬과 다양성 사이의 균형을 신중하게 맞출 필요성 등의 한계를 인정한다. 그들은 전문가의 수가 현재 고정되어 있지만, 이 프레임워크가 임상 워크플로우에 통합될 수 있는 더 신뢰할 수 있고 불확실성을 인지하는 병리 모델을 위한 길을 열어준다고 결론짓는다.