CUPID-Med: Conformal Uncertainty-aware Prototype Discovery for Clinically Interpretable Deep Medical Clustering

CUPID-Med 기술 요약: 임상적으로 해석 가능한 딥 메디컬 클러스터링을 위한 공형 불확실성 인지 프로토타입 발견 (Conformal Uncertainty-aware Prototype Discovery for Clinically Interpretable Deep Medical Clustering)

문제 정의
딥 메디컬 클러스터링은 라벨이 없는 의료 데이터로부터 잠재적인 질병 하위 그룹을 발견하는 것을 목표로 하며, 이는 전문가 주석의 높은 비용과 가변성을 고려할 때 매우 중요한 과제입니다. 그러나 기존의 딥 클러스터링 방식은 다음과 같은 세 가지 주요 한계점을 가집니다:

1. 결정론적 할당 (Deterministic Assignments): 대부분의 방법은 샘플을 단일한 하드 클러스터 레이블로 강제하며, 이는 임상적 표현형의 연속적이고 중첩되며 이질적인 특성을 반영하지 못합니다.
2. 보정된 불확실성 결여 (Lack of Calibrated Uncertainty): 표준 파이프라인은 보정된 불확실성 추정치를 제공하거나, 전문가의 검토가 필요한 모호한 사례를 식별하기 위한 공식적인 메커니즘을 갖추는 경우가 드뭅니다.
3. 제한된 해석 가능성 (Limited Interpretability): 클러스터링 결과는 대표적인 프로토타입이나 시각적 설명과 같은 임상적으로 의미 있는 요약을 결여하는 경우가 많아, 임상의가 이를 검사하고 신뢰하기 어렵게 만듭니다.

또한, 의료 파운데이션 모델은 강력한 표현력을 제공하지만, 고정된 임베딩은 비임상적 아티팩트(예: 획득 편향)를 인코딩할 수 있으며, 본질적으로 불확실성을 인지하는 하위 그룹 발견을 지원하지는 않습니다.

방법론: CUPID-Med
저자들은 고정된 의료 임베딩을 보정되고 해석 가능하며 불확실성을 인지하는 하위 그룹으로 변환하기 위해 설계된 프레임워크인 CUPID-Med를 제안합니다. 이 아키텍처는 다섯 가지 핵심 단계로 구성됩니다:

1. 의료 파운데이션 인코딩 (Medical Foundation Encoding): 프레임워크는 의료 파운데이션 인코더(또는 DenseNet121과 같은 사전 학습된 백본)를 사용하여 원시 입력(이미지 또는 멀티모달 데이터)을 통일된 표현으로 매핑합니다. 이미지 전용 설정(CheXpert)의 경우, 이는 패치 임베딩과 트랜스포머 기반 처리 또는 CNN 특징 추출을 포함합니다.
2. 분포적 잠재 표현 학습 (Distributional Latent Representation Learning): 샘플을 결정론적인 점으로 매핑하는 대신, CUPID-Med는 각 샘플을 가우시안 분포 $N(\mu_i, \text{diag}(\sigma_i^2))$로 표현합니다. 이는 평균 벡터와 로그 분산을 출력하는 프로젝션 헤드를 통해 달성됩니다. 모델은 견고한 잠재 분포를 학습하기 위해 자기 지도 일관성 손실(증강된 뷰 사이의 KL 발산)과 대조 학습 목적 함수를 사용합니다. 스칼라 불확실성 점수 $u_i$는 잠재 분산으로부터 도출됩니다.
3. 프로토타입 기반 딥 클러스터링 (Prototype-Based Deep Clustering): 프레임워크는 잠재 질병 하위 그룹을 나타내는 학습 가능한 메모리 뱅크인 $K$개의 프로토타입을 유지합니다. 모델은 다음을 결합한 클러스터링 목적 함수를 최적화합니다:
   • 타겟 분포 정교화 (소프트 할당을 정교화된 타겟에 정렬)
   • 프로토타입 압축성 (샘플과 할당된 프로토타입 사이의 거리 최소화)
   • 클러스터 균형 (균등한 프로토타입 활용 장려)
   • 프로토타입 분리 (서로 다른 프로토타입 간의 거리 최대화)
4. 공형 불확실성 보정 (Conformal Uncertainty Calibration): 신뢰할 수 있는 결정 경계를 구축하기 위해, CUPID-Med는 공형 예측(conformal prediction)을 비지도 학습 환경에 적응시킵니다. 이는 소프트 클러스터 할당을 **집합 값 예측(set-valued predictions)**으로 변환합니다.
   • 비정상성 점수(nonconformity score)는 낮은 할당 확신도와 잠재적 불확실성을 기반으로 계산됩니다.
   • 분포 불변 커버리지를 보장하기 위해 캘리브레이션 세트에서 임계값이 추정됩니다.
   • 테스트 샘플의 경우, 출력은 후보 클러스터 집합 $\Gamma_\alpha(x)$가 됩니다. 만약 집합에 여러 클러스터가 포함되어 있거나 불확실성이 높은 경우, 해당 샘플은 단일 레이블로 강제되는 대신 전문가 검토 대상으로 플래그가 지정됩니다.
5. 임상적으로 해석 가능한 출력 (Clinically Interpretable Outputs): 시스템은 대표 이미지(프로토타입에 가장 가까운 샘플), 임상적 풍부화 프로파일(사후에 레이블이 있는 경우), 불확실성 분포를 포함하는 "클러스터 카드"를 생성합니다. 시각적 설명(예: Grad-CAM)은 향후 확장 과제로 언급되었으며, 현재의 초점은 프로토타입 검색에 있습니다.

주요 기여

• 프레임워크 통합: 의료 파운데이션 인코딩, 분포적 잠재 학습, 프로토타입 발견 및 공형 보정을 결합한 통합 프레임워크 제안.
• 불확실성 인지 할당: 결정론적 레이블을 제공하는 대신, 집합 값 예측을 생성하여 인간 중심의 검토(human-in-the-loop)를 명시적으로 식별하는 공형 전략 도입.
• 프로토타입 기반 해석 가능성: 각 클러스터에 대한 대표 샘플을 검색하는 메커니즘을 통해 사례 기반 추론 및 시각적 검사를 가능하게 함.
• 무레이블 학습 프로토콜: CheXpert 데이터셋에서 질병 레이블을 학습에서 엄격히 제외하고 사후 임상적 풍부화 및 해석 가능성 분석에만 사용하는 엄격한 평가 프로토콜 정의.

실험 결과
프레임워크는 무레이블 학습 프로토콜을 사용하여 CheXpert 데이터셋의 5,000개 흉부 정면 이미지 서브셋에 대해 평가되었습니다.

• 베이스라인 성능: 고정된 DenseNet121 임베딩을 사용한 표준 방법(K-means, GMM, Agglomerative)의 예비 클러스터링은 약한 분리도를 보였습니다. 최적의 베이스라인(K-means, $K=5$)은 실루엣 점수 0.0510과 데이비스-볼딘 지수(Davies–Bouldin index) 3.6950을 기록했습니다.
• CUPID-Med 성능: 주요 CUPID-Med 실행($K=14$)은 실루엣 점수 0.5175, 데이비스-볼딘 지수 2.7809, 칼린스키-하라바츠(Calinski–Harabasz) 점수 29,733.16을 달성했습니다.
• 2차 실행: 2차 50-에폭 실행은 더 높은 내부 지표(실루엣: 0.6216)를 보였으나, 빈 클러스터 경고가 발생하여 저자들은 80-에폭 실행을 공식 재현 가능한 결과로 취급했습니다.
• 관찰 사항: CUPID-Med는 결정론적 베이스라인에 비해 잠재적 분리도를 크게 개선했지만, 주요 실행에서 클러스터 크기 불균형(클러스터당 2개에서 1,132개 샘플까지 다양)이 관찰되었습니다. 이는 프로토타입 균형 정규화의 추가적인 개선이 필요함을 시사합니다.

의의 및 주장
본 논문은 분포적 프로토타입 학습이 고정된 딥 임베딩의 결정론적 클러스터링보다 더 구조화되고 해석 가능한 의료 영상 하위 그룹을 생성할 수 있다고 주장합니다.

• 패러다임의 전환: 본 연구는 "불투명한 하드 클러스터"에서 "보정되고 해석 가능하며 불확실성을 인지하는 하위 그룹 발견"으로의 이동을 주장합니다.
• 임상적 유용성: 집합 값 예측을 생성하고 불확실한 사례를 식별함으로써, 중첩되는 표현형과 모호한 소견이 흔한 임상 현실에 더 잘 부합합니다.
• 배포에 대한 신중함: 저자들은 본 방법이 우수한 내부 분리도와 해석 가능성을 입증했음에도 불구하고, 관찰된 클러스터 불균형과 더 강력한 균형 정규화의 필요성 때문에 이 방법이 아직 안정적인 임상 배포를 위한 페노타입 발견 도구로서 준비되지 않았음을 명시적으로 밝힙니다. 결과는 더 신뢰할 수 있고 인간 중심적인 클러스터링 워크플로우를 위한 개념 증명(proof-of-concept)으로 제시됩니다.