[논문 리뷰] Unpaired Multi-modal Segmentation via Knowledge Distillation

Qi Dou, Quande Liu외 2명|arXiv (Cornell University)|2020. 01. 06.

Domain Adaptation and Few-Shot Learning참고 문헌 34인용 수 10

한 줄 요약

이 논문은 CT와 MRI 간에 모든 컨볼루션 커널을 공유하면서 모odal별 배치 정규화 레이어를 사용하는 파라미터 효율적인, 쌍이 없는 다중 모odal 세그멘테이션 프레임워크를 제안한다. 예측 분포의 KL-분산을 기반으로 한 새로운 지식 정련 손실을 도입하여 모달 간 의미적 특징을 정렬함으로써, 2D 및 3D 네트워크를 사용하여 심장 및 Abdominal 다기관 세그멘테이션 작업에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Multi-modal learning is typically performed with network architectures containing modality-specific layers and shared layers, utilizing co-registered images of different modalities. We propose a novel learning scheme for unpaired cross-modality image segmentation, with a highly compact architecture achieving superior segmentation accuracy. In our method, we heavily reuse network parameters, by sharing all convolutional kernels across CT and MRI, and only employ modality-specific internal normalization layers which compute respective statistics. To effectively train such a highly compact model, we introduce a novel loss term inspired by knowledge distillation, by explicitly constraining the KL-divergence of our derived prediction distributions between modalities. We have extensively validated our approach on two multi-class segmentation problems: i) cardiac structure segmentation, and ii) abdominal organ segmentation. Different network settings, i.e., 2D dilated network and 3D U-net, are utilized to investigate our method's general efficacy. Experimental results on both tasks demonstrate that our novel multi-modal learning scheme consistently outperforms single-modal training and previous multi-modal approaches.

연구 동기 및 목표

• 쌍이 없는 다중 모달 의료 영상 세그멘테이션 문제, 특히 이미지 레지지스트레이션과 쌍이 맞는 데이터가 없는 CT와 MRI 간의 도전 과제를 해결하기 위해.
• 쌍이 맞는 훈련 데이터나 복잡한 네트워크 아키텍처 없이도 교차 모달 지식을 활용하여 세그멘테이션 성능을 향상시키기 위해.
• 모든 컨볼루션 커널을 모달 간에 공유하면서도 모달별 정규화와 새로운 정련 손실을 통해 분포 이탈을 최소화함으로써 매우 파라미터 효율적인 모델을 개발하기 위해.
• 다양한 네트워크 아키텍처와 데이터 품질 조건에서 실제 임상 세그멘테이션 작업에 대해 방법을 검증하기 위해.

제안 방법

• 모든 컨볼루션 커널을 CT와 MRI 간에 공유함으로써 모델 파라미터를 줄이고 파라미터 효율성을 달성한다.
• CT와 MRI의 서로 다른 통계적 분포에 적응하기 위해 모달별 내부 정규화 레이어(예: 배치 정규화)를 사용한다.
• 동일한 입력에 대해 서로 다른 모달에서 유도된 예측 분포 간의 KL-분산을 최소화하여 새로운 지식 정련 손실을 도입한다.
• 예측 분포의 정렬을 위해 소프트맥스 이전 특징에서 손실을 계산함으로써 의미적 표현을 정렬하고, 외관 차이가 있음에도 불구하고 공유된 강력한 특징을 학습하도록 유도한다.
• 이 프레임워크는 2D 확장 컨볼루션 신경망과 3D U-Net 모두에 적용되어 아키텍처와 작업 간의 일반화 능력을 입증한다.
• 표준 교차 엔트로피 손실과 제안된 KD-손실을 함께 사용하여 종단 간 훈련을 수행함으로써 모달별 및 공유 표현의 공동 최적화를 가능하게 한다.

실험 결과

연구 질문

• RQ1 같은 입력에 대해 쌍이 없는 CT와 MRI 데이터에서 효과적으로 학습할 수 있는 단일이고 콤팩트한 신경망은 가능한가?
• RQ2 두 모달 간에 심하게 다른 강도 분포와 외관을 가질 경우, 교차 모달 지식을 효과적으로 정련하는 방법은 무엇인가?
• RQ3 모든 컨볼루션 가중치를 공유하면서도 모달별 정규화 레이어를 사용하는 것이 쌍이 없는 다중 모달 학습에서 별도의 인코더보다 더 나은 성능을 내는가?
• RQ4 예측 분포 정렬 기반의 지식 정련 손실은 저품질 또는 이방성 영상 환경에서 정확도와 강건성을 향상시키는가?
• RQ5 제안된 방법은 다양한 네트워크 아키텍처(2D 대 3D)와 다중 클래스 세그멘테이션 작업 간에 일반화 가능한가?

주요 결과

• 이중적인 이상치 케이스(이미지 아티팩트 존재)에서 복부 다기관 세그멘테이션 작업에서 제안된 방법은 81.1의 Dice 스코어를 기록하여 모든 베이스라인을 초월했다.
• 동일한 이상치 케이스에서 Hausdorff 거리는 3.78을 기록하여 다음으로 좋은 방법(4.89)보다 유의미하게 낮아 경계 정렬 성능이 뛰어났다.
• 일반적인 케이스에서는 싱글 모달 훈련과 최신 기술 수준의 다중 모달 접근 방식(예: 'Y'-형 및 'X'-형 네트워크 포함)을 모두 초월하여 심장 및 복부 세그멘테이션 작업에서 뛰어난 성능을 보였다.
• 2D 확장 컨볼루션 신경망과 3D U-Net을 포함한 다양한 네트워크 아키텍처에서 일관된 성능 향상을 보이며 일반화 능력을 확인했다.
• 절단 실험을 통해 제안된 지식 정련 손실이 특히 도전적인 영상 조건에서 성능에 핵심적인 역할을 한다는 점을 입증했다.
• 이중적인 아티팩트와 저품질 데이터에 대해 강건성을 보이며, 다른 방법이 심각하게 성능 저하를 보이는 상황에서도 높은 성능 유지를 유지했다.