Halfway to 3D: Ensembling 2.5D and 3D Models for Robust COVID-19 CT Diagnosis

1. 문제 정의 (Problem Definition)

• 배경: COVID-19 및 폐 질환 진단에서 흉부 CT 스캔은 중요한 역할을 하지만, 다양한 기관, 스캐너, 촬영 프로토콜, 환자 군집에서 수집된 데이터로 인해 **도메인 시프트 (Domain Shift)**가 발생합니다. 이는 모델의 일반화 성능을 저하시킵니다.
• 도전 과제:
  1. 강건성 (Robustness): 이질적인 데이터 소스 (Multi-source) 에서 일관된 성능을 내는 모델 개발.
  2. 공정성 (Fairness): 성별 등 인구통계학적 그룹 간 성능 편차 (Performance Disparity) 를 줄이는 것.
  3. 정보의 한계: 3D 볼륨 모델은 전역적 맥락을 포착하지만 계산 비용이 크고 미세한 슬라이스 수준의 패턴을 놓칠 수 있으며, 2D/2.5D 모델은 국소적 특징은 잘 포착하지만 3D 구조적 맥락이 부족합니다.

2. 제안 방법론 (Methodology)

저자들은 2.5D 슬라이스 기반 표현과 3D 볼륨 기반 표현을 통합한 하이브리드 앙상블 프레임워크를 제안합니다.

A. 데이터 전처리

• 원시 axial CT 슬라이드를 스택하여 3D 볼륨 (128×128×128) 을 재구성합니다.
• 3D 가우시안 노이즈 제거, 마스크 샤프닝, 그리고 픽셀 값 정규화 ([0, 255]) 를 수행합니다.

B. 3D 표현 학습 (3D Branch)

• 아키텍처: ResNet-18 (비디오 처리용 3D 컨볼루션 기반).
• 학습 전략:
  1. 도메인 일반화 프리트레이닝: **VREx (Variance Risk Extrapolation)**를 사용하여 도메인 간 분산을 최소화하고 도메인 불변 표현을 학습합니다.
  2. 태스크 특화 파인튜닝: Supervised Contrastive Learning과 Cross-Entropy 손실을 결합하여 특징의 분리도를 높입니다. 또한 MixUp 증강을 적용하여 결정 경계를 부드럽게 만듭니다.

C. 2.5D 다중 뷰 표현 학습 (2.5D Branch)

• 접근 방식: 재구성된 3D 볼륨에서 Axial, Coronal, Sagittal 세 가지 직교 평면의 슬라이스를 추출합니다.
• 아키텍처: DINOv3 (Self-supervised Vision Transformer) 를 백본 (Backbone) 으로 사용합니다. 사전 학습된 강력한 시각적 표현을 의료 영상에 전이 (Transfer) 합니다.
• 학습 전략: 각 뷰의 슬라이스를 독립적으로 처리한 후, 특징 수준 (Feature-level) 에서 융합하여 다중 뷰 정보를 통합합니다.

D. 앙상블 통합 (Ensemble Integration)

• 3D 모델과 2.5D 모델의 예측 로짓 (Logits) 을 결합하여 최종 분류를 수행합니다.
• 이 방식은 볼륨적 맥락 (3D) 과 고해상도 슬라이스 수준의 특징 (2.5D) 을 모두 활용하여 강건성을 극대화합니다.

3. 주요 기여 (Key Contributions)

1. 하이브리드 2.5D-3D 프레임워크: CT 기반 진단을 위해 다중 뷰 슬라이스 표현과 볼륨 모델링을 결합한 새로운 구조 제안.
2. 강화된 3D 학습 전략: VREx 와 지도식 대비 학습 (Supervised Contrastive Learning) 을 도입하여 도메인 일반화 및 특징 분리 능력을 향상시킴.
3. DINOv3 기반 2.5D 파이프라인: 사전 학습된 대규모 비전 트랜스포머를 의료 영상 슬라이스 분석에 효과적으로 적용.
4. 성능 및 강건성 입증: PHAROS-AIF-MIH 벤치마크에서 앙상블 전략이 다중 소스 환경에서의 강건성과 공정성을 크게 개선함을 실험을 통해 증명.

4. 실험 결과 (Results)

벤치마크: PHAROS-AIF-MIH (다중 소스 COVID-19 탐지 및 공정한 질병 진단 챌린지)

Task 1: 이진 COVID-19 탐지 (Binary Detection)

• 성능: 앙상블 모델이 Macro F1-score 0.751을 기록하여 챌린지 2 위를 차지했습니다.
• 비교: 2.5D 모델 (DINOv3) 이 3D 모델 (ResNet-18) 보다 단일 모델로서 더 높은 성능을 보였으나, 두 모델을 앙상블하면 다양한 병원 (Source) 에서의 일반화 성능이 추가적으로 향상되었습니다. 특히 Source 2 와 같은 도메인 시프트가 큰 환경에서 강건성이 입증되었습니다.

Task 2: 공정한 질병 진단 (Multi-class Classification & Fairness)

• 성능: 4 가지 클래스 (Healthy, Adenocarcinoma, Squamous Cell Carcinoma, COVID-19) 분류에서 Macro F1-score 0.633을 기록하여 3 위를 달성했습니다.
• 공정성: 성별 간 성능 격차가 존재했으나, 가중치를 조절된 앙상블 (0.7-weighted) 을 통해 남성 환자군의 성능을 크게 개선하여 전체적인 균형을 맞추었습니다.
• 한계: 'Squamous Cell Carcinoma (Class G)'와 같은 특정 질병 분류에서는 2.5D 모델이 단독으로 가장 좋은 성능을 보였으며, 3D 볼륨 정보가 추가적인 이점을 제공하지는 못했습니다.

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

• 상호 보완적 강점: 3D 모델은 전역적 맥락 (Global Context) 을, 2.5D 모델은 고해상도 국소적 특징 (Local Features) 을 잘 포착하여, 이 둘을 결합함으로써 단일 접근법의 한계를 극복했습니다.
• 임상적 적용 가능성: 실제 병원 환경에서 발생하는 도메인 시프트 (다양한 장비, 프로토콜) 에 견딜 수 있는 강건한 진단 도구 개발의 실용적인 경로를 제시했습니다.
• 공정성 고려: 앙상블 가중치 조절을 통해 인구통계학적 그룹 간의 성능 편차를 줄일 수 있음을 보여주었습니다.
• 결론: 본 연구는 "Halfway to 3D"라는 개념을 통해, 완전한 3D 모델링과 2D 슬라이스 분석 사이의 최적 균형을 찾음으로써 의료 AI 의 신뢰성과 확장성을 높이는 데 기여했습니다.