Multi-Dataset Cross-Domain Knowledge Distillation for Unified Medical Image Segmentation, Classification, and Detection

기술 요약: 통합 의료 영상 분할, 분류 및 검출을 위한 다중 데이터셋 교차 도메인 지식 증류

문제 제기
딥러닝은 의료 영상 분석을 혁신하였으나, 단일 데이터셋으로 훈련된 모델은 촬영 프로토콜, 스캐너 하드웨어, 환자 군집의 차이로 인한 도메인 이동 (domain shifts) 으로 인해 새로운 임상 환경에 일반화되지 못하는 경우가 많습니다. 전이 학습과 데이터 증강은 부분적인 해결책을 제공하지만, 종종 단일 소스 사전 훈련이나 미세한 해부학적 세부 사항과 데이터셋 간 변이를 포착하지 못하는 합성 데이터에 의존합니다. furthermore, 기존 지식 증류 (KD) 접근법은 일반적으로 단일 작업 (예: 분할) 에 대해 동질적인 데이터로 단일 교사 모델을 활용하며, MRI, CT, X-선과 같은 이질적인 모달리티에 걸친 분할, 분류, 객체 검출을 포괄하는 임상 워크플로우의 다양하고 다중 작업적인 특성을 간과합니다.

방법론
저자들은 이질적인 의료 영상 데이터셋으로부터 지식을 집계하도록 설계된 통합적이고 작업 중립적인 교사 - 학생 프레임워크를 제안합니다. 이 프레임워크는 세 가지 연속적인 단계로 운영됩니다:

1. 도메인 적대적 교사 훈련:

   • 타겟 교사는 작업별 손실 (예: 분할용 Dice 손실) 을 사용하여 타겟 데이터셋 exclusively 로 훈련됩니다.
   • 소스 교사는 각각의 소스 데이터셋에서 훈련되지만 도메인 판별기가 추가됩니다. 도메인 적대적 목적 함수를 사용하여 이러한 인코더는 판별기를 혼란스럽게 하는 도메인 불변 특징을 생성하도록 강제되며, 이를 통해 소스 및 타겟 특징 분포를 정렬합니다. 이 단계는 지도 (레이블된 타겟 데이터 사용) 및 비지도 적응을 모두 지원합니다.

2. 교차 주의 융합을 통한 통합 교사 구성:

   • 통합 교사($\mathcal{T}^*$) 는 타겟 교사와 모든 소스 교사로부터의 다중 수준 인코더 및 병목 (bottleneck) 특징을 융합하여 구성됩니다.
   • 이 융합은 서로 다른 도메인의 보완적 표현을 집계하기 위해 교차 주의 메커니즘을 활용합니다. 이 과정에서 모든 구성 교사의 인코더 및 병목 레이어는 학습된 지식을 보존하기 위해 고정됩니다.
   • 새로운 작업별 헤드 (디코더, 분류기 또는 검출 헤드) 가 융합된 백본에 연결되어 타겟 데이터셋에서 처음부터 훈련됩니다. 이를 통해 통합 교사는 도메인 불변 특징을 유지하면서 타겟 작업에 특화될 수 있습니다.

3. 커리큘럼 기반 다중 수준 증류:

   • 컴팩트한 학생 모델이 타겟 데이터셋에서 통합 교사로부터 지식을 계승하도록 훈련됩니다.
   • 훈련 목적 함수는 주요 작업 손실과 여러 증류 항을 결합합니다: 특징 정렬 (MSE), 코사인 유사도, 대조 손실 (유사도 분포에 대한 KL 발산), 그리고 작업별 확장 (예: 주의 전이, 로그트 증류, 검출을 위한 ROI 특징 정렬).
   • 커리큘럼 인자($c_f$) 가 증류 손실의 가중치를 동적으로 스케줄링합니다. 훈련은 지도 작업 손실에 초점을 맞추어 시작되다가, 초기 학습을 불안정하게 하지 않으면서 학생의 내부 표현을 정제하기 위해 증류 항의 영향을 점차적으로 증가시킵니다.

이 프레임워크는 유연하도록 설계되어 분할, 분류 및 검출 작업 전반에 걸쳐 UNet, TResUNet, MedViT, EfficientNet, Faster R-CNN, RF-DETR 아키텍처를 지원합니다.

주요 기여

1. 통합 프레임워크: 이질적인 의료 데이터셋 전반에 걸쳐 분할, 이미지 수준 분류, 객체 검출을 처리할 수 있는 단일 확장 가능 아키텍처.
2. 통합 교사 아키텍처: 교차 주의를 통해 여러 소스 교사로부터의 다중 수준 특징을 융합하는 새로운 방법으로, 단일 모놀리식 사전 훈련 모델이 필요 없이 강력한 교차 도메인 표현 전이를 가능하게 함.
3. 커리큘럼 기반 증류: 커리큘럼 스케줄에 의해 지배되어 안정적인 최적화를 보장하는 분할, 대조, 정렬 및 유사도 손실을 결합한 다중 수준 증류 전략.
4. 종합 평가: 세 가지 작업, 세 가지 모달리티 (MRI, CT, X-선), 그리고 수많은 데이터셋 (BrainMetShare, ISLES, BraTS, LiTS, KiTS, COVIDx-CXR, OASIS, ADNI, DeepLesion 포함) 을 통한 광범위한 실험.

실험 결과
제안된 방법은 두 가지 기준선, 즉 데이터셋별 모델(각 데이터셋에서 독립적으로 훈련) 과 다중 헤드 다중 데이터셋 모델(데이터셋별 헤드가 있는 공유 백본) 에 대해 평가되었습니다.

• 분할: MRI 및 CT 데이터셋에서 학생 모델은 두 가지 기준선보다 일관되게 우수한 성능을 보였습니다. CT 기반 작업의 경우, 학생 모델은 데이터셋별 기준선 대비 평균 IoU 를 약 2.5% 향상시켰으며, 다중 헤드 모델 대비 약 4.8% 향상시켰고, Hausdorff Distance (HD95) 가 유의하게 감소하여 더 날카로운 경계 예측을 나타냈습니다.
• 분류: 폐질환 및 치매 분류 작업에서 증류된 학생 모델은 데이터셋별 및 다중 헤드 기준선보다 높은 정확도를 달성했습니다. Late-$k$ 증류 (더 깊은 레이어로부터 지식 전달) 가 일반적으로 더 우수한 결과를 산출했습니다.
• 객체 검출: 폐 CT 데이터셋에서 이 프레임워크는 평균 정밀도 (mAP) 를 향상시켰습니다. 통합 교사는 종종 최고 성능을 달성했으며, 증류된 학생 모델은 이러한 이점의 대부분을 유지하여 배포에 적합한 더 컴팩트한 모델을 제공했습니다.
• Ablation 연구: 결과들은 (i) 지도 타겟 적응이 경계 정확도를 향상시킨다는 점, (ii) 성능이 소스 데이터셋의 수에 따라 확장된다는 점, 그리고 (iii) 단일 수준 추출보다 인코더 + 병목의 다중 수준 특징 추출이 더 나은 결과를 산출한다는 점을 확인했습니다.
• 시각화: t-SNE 시각화는 증류된 학생 모델이 세 가지 작업 전반에 걸쳐 데이터셋별 기준선보다 더 컴팩트하고 더 잘 분리된 특징 클러스터를 학습함을 보여주었습니다.

의의 및 주장
본 논문은 다중 데이터셋 지식 증류가 이질적인 의료 영상 도메인 전반의 성능을 향상시키는 확장 가능하고 작업 중립적인 접근법이라고 주장합니다. 다양한 소스로부터 도메인 불변 표현을 집계하고 이를 컴팩트한 학생 모델로 증류함으로써, 이 프레임워크는 데이터 부족, 주석 비용, 도메인 이동의 문제를 해결합니다. 저자들은 그들의 접근 방식이 강력한 기준선 대비 일관된 개선을 이루어 분포 이동에 대한 향상된 견고성과 우수한 일반화를 입증한다고 강조합니다. 이 연구는 이전의 분할 중심 연구를 실질적으로 확장한 것으로, 이제 의료 영상 분석 작업의 전체 스펙트럼에서 검증되었습니다.