Mutual Distillation of Dual-Foundation Models for Semi-Supervised PET/CT Segmentation

기술 요약: 반지도 학습 기반 PET/CT 분할을 위한 이중 파운데이션 모델의 상호 증류(Mutual Distillation)

문제 정의
전신 PET/CT에서 정확한 장기 분할은 정량적 대사 분석, 질병 병기 결정 및 방사선 치료 계획 수립에 필수적입니다. 그러나 딥러닝 모델 개발은 두 가지 주요 과제로 인해 어려움을 겪고 있습니다:

1. 해부학적-기능적 불일치(Anatomical-Functional Mismatch): PET 영상은 높은 대사 민감도를 제공하지만 경계가 흐릿하고 스필오버(spill-over) 아티팩트가 발생하는 반면, CT는 명확한 구조적 세부 정보를 제공합니다. 공간적 불일치를 해결하기 위해 이러한 두 양식(modality)을 융합하는 것은 어렵습니다.
2. 데이터 부족: 고품질의 픽셀 단위 3D 주석(annotation)을 생성하려면 방사선 전문의와 핵의학 전문가의 동기화된 전문 지식이 필요하며, 이로 인해 라벨링된 데이터를 확보하는 데 많은 비용이 듭니다.
   반지도 학습(SSL)과 최근의 파운데이션 모델(SAM 등)이 해결책을 제시하고 있으나, 기존 방식들은 종종 단일 양식 시나리오에 국한되거나 구조적(CT) 및 기능적(PET) 영상 사이의 깊은 의미론적 상호 의존성을 활용하는 데 실패합니다. 또한, 의사 라벨링(pseudo-labeling)에 의존하는 표준 SSL 방법은 확증 편향(confirmation bias)을 일으켜 잘못된 예측을 강화할 위험이 있습니다.

방법론: MuDuo 프레임워크
저자들은 수동 프롬프트 없이 경량화된 학생 네트워크로 지식을 증류하기 위해, 구조적 및 기능적 파운데이션 모델을 시너지 효과를 내며 활용하도록 설계된 상호 증류 프레임워크인 MuDuo를 제안합니다. 이 프레임워크는 이중 양식 입력($X_{CT}$ 및 $X_{PET}$)을 대상으로 작동하며, 세 가지 주요 구성 요소로 이루어집니다:

1. 이중 파운데이션 모델 교사(Dual Foundation Model Teachers):

   • SAM-Med3D: 해부학적 사전 지식을 활용하여 CT에 특화된 고정된(frozen) 교사 모델입니다.
   • SegAnyPET: 대사적 사전 지식을 활용하여 PET에 특화된 고정된 교사 모델입니다.
   • 학생 네트워크(Student Network): 이중 양식 입력을 처리하는 학습 가능한 3D U-Net ($S_\theta$)입니다.
   • 프롬프트 생성기(Prompt Generator): 학생 모델과 PET 교사 모델 사이를 연결하는 경량 UNet ($G_\phi$)입니다.

2. 이중 경로 의사 주석 생성(Dual-Pathway Pseudo Annotation Generation):
   라벨이 없는 데이터에 대해, 프레임워크는 양식의 특성에 맞게 조정된 두 가지 별도의 경로를 통해 의사 라벨을 생성합니다:

   • CT 경로: 학생의 예측값(프롬프트 생성기를 통해)을 다운샘플링하여 SAM-Med3D를 위한 **마스크 프롬프트(mask prompts)**로 사용합니다. 교사는 CT의 해부학적 사전 지식을 바탕으로 이 예측을 정교화합니다.
   • PET 경로: 프롬프트 생성기는 조粗한 로짓(logits)을 생성하며, 이는 단일 중심점이 아닌 다지점 프롬프트(multi-point prompts)(상위 K개의 고확률 위치 샘플링)로 변환됩니다. 이는 PET 영상 특유의 이질적인 섭취(uptake) 현상을 해결하기 위함입니다. 이 지점들은 SegAnyPET이 정교한 마스크를 생성하도록 유도합니다.

3. 상호 증류 및 일관성(Mutual Distillation and Consistency):

   • IoU 기반 합의 필터링(IoU-Based Consensus Filtering): 신뢰성을 보장하기 위해, 프레임워크는 CT 브랜치와 PET 브랜치 예측 간의 IoU(Intersection-over-Union)를 계산합니다. 두 양식이 일치하는(높은 합의를 보이는) 상위 50%의 샘플만을 학습을 위해 유지합니다. 유효하지 않은 프롬프트를 가진 샘플은 제외됩니다.
   • 동적 의사 라벨 융합(Dynamic Pseudo-Label Fusion): 무작위 가중치 인자($\alpha$)를 사용하여 CT와 PET의 전문가 예측을 동적으로 혼합함으로써 고품질의 의사 라벨을 생성합니다. 이는 모델이 자신의 예측을 그대로 암기하는 것을 방지하기 위한 섭동(perturbation)을 도입합니다.
   • 일관성 손실(Consistency Loss): 학생 네트워크는 가우시안 램프업(Gaussian ramp-up) 가중치 스케줄을 사용하여 초기 학습 불안정성을 방지하면서, 이러한 융합된 고품질 의사 라벨과 일치하도록 MSE(Mean Squared Error)를 통해 학습됩니다.

주요 기여

• 이중 양식 파운데이션 통합: 단일 양식 SSL 방법과 달리, MuDuo는 독창적으로 SAM-Med3D(구조적)와 SegAnyPET(기능적)을 통합하여 상호 지식을 증류함으로써 PET/CT 고유의 해부학적-기능적 불일치를 해결합니다.
• 자동 프롬프트 메커니즘: 외부 프롬프트 생성기나 수동 프롬프트 없이, 학생 모델의 자체 예측을 사용하여 파운데이션 모델을 가이드함으로써 라벨 없는 데이터의 활용도를 극대화합니다.
• 합의 필터링: IoU 기반 필터링 메커니즘의 도입은 신뢰할 수 있고 합의된 의사 라벨만이 학생을 감독하도록 보장하여, 전통적인 의사 라벨링에서 흔히 발생하는 확증 편향을 완화합니다.
• PET를 위한 다지점 프롬프팅: PET 입력을 위한 다지점 샘플링 제안은 단일 지점 프롬프팅에 비해 이질적인 대사 섭취를 보이는 장기의 공간적 범위를 더 잘 포착합니다.

실험 결과
프레임워크는 AutoPET 데이터셋(1,038개 연구)을 대상으로 평가되었으며, 전문가 주석이 포함된 100개의 케이스 서브셋을 사용했습니다. 실험은 5개, 10개, 20개의 라벨된 케이스가 있는 극한의 데이터 희소성 설정 하에서 수행되었습니다.

• 성능: MuDuo는 모든 설정에서 최첨단(SOTA) 성능을 달랐습니다. 5개 라벨 케이스 설정에서 **46.93%**의 Dice 점수를 달성하여, 두 번째로 우수한 방법(CPS) 및 단일 양식 파운데이션 접근법(SemiSAM+)을 크게 앞질렀습니다.
• 경계 정확도: 이 방법은 탁월한 경계 정밀도를 보여주었으며, 5개 라벨 설정에서 95백분위수 Hausdorff 거리(HD95)를 8.50 복셀로 감소시켰습니다. 이는 SemiSAM+ 대비 41% 감소한 수치입니다.
• 전체 지도 학습과의 비교: 10개의 라벨 케이스를 사용할 때, 반지도 학습 방식인 MuDuo는 이중 양식 전체 지도 학습 베이스라인보다 Dice 점수에서 약 9% 더 높은 성능을 보였으며, 이는 교차 양식 합의를 통한 라벨 없는 데이터 마이닝의 효용성을 입증합니다.
• 절제 연구(Ablation Studies): 실험을 통해 이중 교사를 사용하는 것이 단일 양식 변형보다 우수하며, 적응형 상위 50% IoU 필터링이 고정 임계값보다 더 나은 결과를 낳고, 다지점 프롬프팅(K=3)이 PET 분할에 있어 단일 지점 프롬프팅보다 우수함을 확인했습니다.

의의 및 주장
본 논문은 MuDuo가 학생 모델의 작업 특정적 정밀도와 파운데이션 모델의 일반적 분할 사전 지식 사이의 간극을 성공적으로 메운다고 주장합니다. 구조적 및 기능적 사전 지식을 시너지 효과를 내며 활용함으로써, 프레임워크는 값비싼 이중 양식 주석의 부담 없이 자동화된 고정밀 분할을 달성합니다. 저자들은 이 접근 방식이 제한된 주석의 성능 병목 현상을 효과적으로 돌파하여, 라벨링된 데이터가 부족한 종양학 분야의 정량적 분석 및 방사선 치료 계획 수립에 실질적인 솔루션을 제공한다고 단언합니다. 재현성을 높이기 위해 소스 코드는 공개되었습니다.