[논문 리뷰] Hybrid Dual Mean-Teacher Network With Double-Uncertainty Guidance for Semi-Supervised Segmentation of MRI Scans

Jiayi Zhu, Bart Bolsterlee외 3명|arXiv (Cornell University)|2023. 03. 09.

Advanced Neural Network Applications인용 수 5

한 줄 요약

이 논문은 이중 2D 및 3D 평균-티처 네트워크를 융합하여 이종 특징과 불확실성 유도 병합을 활용함으로써, 최소한의 레이블 데이터로도 반감성 MRI 분할 성능을 향상시키는 하이브리드 듀얼 메인티처(HD-Teacher) 네트워크를 제안한다. 모델 및 작업 불확실성에 기반한 이중 불확실성 추정을 통해 양 모odal리티의 예측을 동적으로 가중함으로써, 이중 및 다중 클래스 MRI 분할에서 최신 기준 성능을 달성한다.

ABSTRACT

Semi-supervised learning has made significant progress in medical image segmentation. However, existing methods primarily utilize information acquired from a single dimensionality (2D/3D), resulting in sub-optimal performance on challenging data, such as magnetic resonance imaging (MRI) scans with multiple objects and highly anisotropic resolution. To address this issue, we present a Hybrid Dual Mean-Teacher (HD-Teacher) model with hybrid, semi-supervised, and multi-task learning to achieve highly effective semi-supervised segmentation. HD-Teacher employs a 2D and a 3D mean-teacher network to produce segmentation labels and signed distance fields from the hybrid information captured in both dimensionalities. This hybrid learning mechanism allows HD-Teacher to combine the `best of both worlds', utilizing features extracted from either 2D, 3D, or both dimensions to produce outputs as it sees fit. Outputs from 2D and 3D teacher models are also dynamically combined, based on their individual uncertainty scores, into a single hybrid prediction, where the hybrid uncertainty is estimated. We then propose a hybrid regularization module to encourage both student models to produce results close to the uncertainty-weighted hybrid prediction. The hybrid uncertainty suppresses unreliable knowledge in the hybrid prediction, leaving only useful information to improve network performance further. Extensive experiments of binary and multi-class segmentation conducted on three MRI datasets demonstrate the effectiveness of the proposed framework. Code is available at https://github.com/ThisGame42/Hybrid-Teacher.

연구 동기 및 목표

• 높은 이방성(MRI 스캔에서 복잡한 해부학적 구조를 가진 경우에 성능이 떨어지는) 2D 또는 3D 단일 차원 특징에 의존하는 기존 반감성 방법의 한계를 해결한다.
• 2D 및 3D 네트워크에서 유사하지만 상보적인 특징을 하이브리드 학습을 통해 융합하여 분할 정확도와 내성 향상을 도모한다.
• 모델 수준 및 작업 수준 불확실성의 이중 불확실성 유도 메커니즘을 도입하여 예측을 동적으로 가중하고 교사 예측을 정밀하게 보정한다.
• 학생 모델이 불확실성 가중 하이브리드 예측에 정렬되도록 하는 하이브리드 정규화 모듈을 개발하여 일반화 능력을 향상시키고 신뢰할 수 없는 교사 출력에서 발생하는 노이즈를 줄인다.
• 작은 수의 레이블된 MRI 스캔만으로도 완전 감독 모델과 인간 전문가 평가자 수준의 성능을 달성하여 임상 적용에 적합한 프레임워크를 확보한다.

제안 방법

• 2D 및 3D 각각의 입력 공간에서 세그멘테이션 마스크와 서명 거리 필드(SDF)를 생성하는 별도의 평균-티처 네트워크를 훈련한다.
• 각 교사 네트워크(2D 및 3D)의 예측 신뢰도를 평가하기 위해 모델 수준 불확실성 추정을 수행한다.
• 일致성 정규화 과정을 유도하기 위해 SDF 예측 작업에 대한 작업 수준 불확실성 추정을 수행한다.
• 개별적인 불확실성 점수에 기반한 소프트맥스 가중 방식을 사용하여 2D 및 3D 교사 출력을 동적으로 융합하여 단일 하이브리드 예측을 생성한다.
• 융합된 예측에서 하이브리드 불확실성 추정을 수행하여 신뢰할 수 없는 지식을 재교정하고 억제한 후, 이를 두 학생 모델의 정규화에 활용한다.
• 두 학생 네트워크가 모두 불확실성 가중 하이브리드 예측에 정렬되도록 하이브리드 정규화 손실을 적용함으로써 일반화 능력과 특징 학습을 향상시킨다.

실험 결과

연구 질문

• RQ1 2D 및 3D 평균-티처 네트워크를 융합하는 것이 단일 모odal리티 접근 방식에 비해 반감성 MRI 분할 성능을 향상시킬 수 있는가?
• RQ2 2D 및 3D 교사 예측의 불확실성 유도 병합이 낮은 데이터 환경에서 모델의 내성과 노이즈 감소에 어떻게 기여하는가?
• RQ3 제안된 이중 불확실성 유도(모델 수준 및 작업 수준)가 단일 불확실성 방법보다 더 신뢰할 수 있는 지식 전이를 이끌어내는가?
• RQ4 하이브리드 특징 학습은 복잡한 해부학적 구조를 가진 고도로 이방성 MRI 스캔에서 분할 정확도를 얼마나 향상시키는가?
• RQ5 제안된 프레임워크는 최소한의 레이블 데이터로 인간 수준의 성능을 달성할 수 있는가? 이는 임상 적용에 적합한가?

주요 결과

• HD-Teacher는 세 가지 다양한 MRI 데이터셋에서 최신 기준 성능을 달성하였으며, 이중 및 다중 클래스 분할 작업 모두에서 기존 반감성 및 완전 감독 기반 기준 모델을 초월하였다.
• MUGgLE 데이터셋에서 HD-Teacher는 레이블 데이터의 10%만을 사용했음에도 불구하고, 완전 감독 모델과 인간 전문가 평가자 수준의 성능을 재현하였다.
• 제거 분석 결과, 하이브리드 학습 및 불확실성 유도 정규화 구성 요소가 모두 분할 정확도 향상에 기여했으며, 특히 고도로 이방성 케이스에서 하이브리드 불확실성 모듈이 오차를 최대 18%까지 감소시켰다.
• 서명 거리 필드(SDF)를 보조 작업으로 사용함으로써 경계 정밀도가 향상되었으며, 특히 작은 또는 비정상적인 형태의 구조에서 두드러졌다.
• 다양한 MRI 모odal리티 및 병렬 해상도에서도 일관된 성능을 보였으며, 특히 0.4×0.4×5 mm와 같은 고도로 이방성 스캔에서 기존 3D 전용 모델이 성능을 낮추는 상황에서도 유의미한 성능 유지를 보였다.
• 하이브리드 불확실성 메커니즘이 신뢰도가 낮은 영역에서 예측을 효과적으로 걸러내어 더 안정적이고 일반화 능력이 뛰어난 학생 모델 훈련을 가능케 하였다.