데이터 라벨 부족을 극복한 3D MRI 뇌종양 분할을 위한 반지도 학습 교사와 학생 프레임워크

초록

본 논문은 불확실성을 활용한 교사 모델이 생성한 확률적 가짜 라벨을 기반으로, 이미지‑레벨 신뢰도에 따라 단계적으로 비라벨 데이터를 학생 모델에 투입하는 진행형 커리큘럼을 제안한다. 학생은 고신뢰 영역은 학습하고 저신뢰 영역은 “언러닝”하는 이중 손실을 사용하며, 교사‑학생 예측 일치를 통해 라벨을 지속적으로 정제한다. BraTS 2021 데이터셋에서 라벨 10 %만 사용했을 때 DSC 0.393이 100 % 사용 시 0.872로 크게 향상되었으며, 학생 모델은 교사보다 특정 종양 서브클래스에서 우수한 성능을 보였다.

상세 분석

이 연구는 세 가지 핵심 기술적 기여를 통해 뇌종양 3D MRI 분할에서 라벨 부족 문제를 해결한다. 첫째, 교사 네트워크는 기존 UNet에 Transformer‑ 기반 전역 어텐션과 Atrous Spatial Pyramid Pooling(ASPP)을 결합한 TransASPP‑UNet 구조를 채택한다. 이 설계는 로컬 컨볼루션의 세밀한 경계 검출 능력과 Transformer의 장거리 의존성 포착을 동시에 제공하여, 다양한 스캐너와 프로토콜에서 발생하는 이미지 이질성을 완화한다. 또한, 교사는 두 개의 출력 헤드를 갖추어 픽셀‑레벨 확률 마스크와 불확실성 맵(예측 분산)을 동시에 추정한다. 불확실성은 Monte‑Carlo dropout 혹은 데이터 증강을 통한 K번 전방패스 후 로그 분산 형태로 학습되며, 이는 라벨이 없는 데이터에 대한 신뢰도 평가의 근거가 된다.

둘째, 학생 모델은 교사가 생성한 가짜 라벨을 활용하되, 이미지‑레벨 평균 신뢰도 C(x)를 기준으로 비라벨 데이터를 단계적으로 선택한다. 초기 단계에서는 상위 m %의 고신뢰 샘플만 사용해 안정적인 학습 기반을 마련하고, 이후 K_t 비율을 점진적으로 확대한다. 이 진행형 커리큘럼은 “easy‑to‑hard” 학습 전략과 유사하지만, 여기서는 신뢰도 기반 정량적 기준을 사용해 자동화한다는 점이 차별점이다.

셋째, 학생 학습에서는 고신뢰 영역에 대해 전통적인 Dice + Cross‑Entropy 손실을 최소화하고, 저신뢰 영역에 대해서는 손실을 최대화하는 “언러닝” 손실을 도입한다. 구체적으로, 픽셀 집합을 고신뢰 L_h와 저신뢰 L_l로 분리하고, 총 손실 L_student = L_sup(L_h) − α·L_sup(L_l) 형태로 정의한다. 이는 모델이 불확실한 예측을 강화하기보다 억제하도록 유도해, 가짜 라벨의 노이즈 전파를 효과적으로 차단한다.

추가적으로, 교사와 학생의 예측이 일치하는 영역은 그대로 유지하고, 불일치 영역은 교사의 불확실성 맵을 활용해 재평가·정제한다. 이러한 피드백 루프는 가짜 라벨의 품질을 반복적으로 개선하고, 학생이 점진적으로 교사의 성능을 넘어설 수 있게 만든다. 실험 결과, 교사 모델 자체는 전체 DSC 0.922를 기록했으며, 학생 모델은 전체 0.872에 그쳤지만, 종양 서브클래스인 NCR/NET(0.797)와 Edema(0.980)에서 교사를 능가했고, 특히 Enhancing Tumor 영역에서 0.620이라는 의미 있는 DSC를 달성해 교사의 실패를 보완했다. 라벨 10 %만 사용했을 때의 저성능(DSC 0.393) 대비 100 % 사용 시 0.872로의 급격한 상승은 제안된 커리큘럼이 데이터 효율성을 크게 향상시킴을 입증한다.

전반적으로 이 논문은 (1) 불확실성 기반 가짜 라벨 생성, (2) 신뢰도 기반 진행형 데이터 선택, (3) 고‑저신뢰 영역을 구분하는 이중 손실, (4) 교사‑학생 예측 일치에 의한 라벨 정제라는 네 가지 메커니즘을 통합함으로써, 라벨이 제한된 의료 영상 분야에서 반지도 학습의 실용성을 크게 확대한다는 점에서 의의가 크다. 다만, 불확실성 추정에 의존하는 만큼 dropout·데이터 증강 파라미터 설정에 민감할 수 있으며, 커리큘럼 스케줄(K_t)의 하이퍼파라미터 튜닝이 필요함을 감안해야 한다.