[논문 리뷰] Multi-Task, Multi-Domain Deep Segmentation with Shared Representations and Contrastive Regularization for Sparse Pediatric Datasets

Arnaud Boutillon, Pierre-Henri Conze외 3명|arXiv (Cornell University)|2021. 05. 21.

COVID-19 diagnosis using AI참고 문헌 25인용 수 7

한 줄 요약

이 논문은 공유된 컨볼루션 특징, 도메인별 배치 정규화, 도메인별 세그멘테이션 헤드를 활용하는 다중 작업, 다중 도메인 딥 러닝 프레임워크를 제안한다. 이는 희소한 소아 근골격계 MRI 데이터셋에서 세그멘테이션 정확도를 크게 향상시키기 위해 감독형 대비 정규화를 통합한다. 이로 인해 최신 기술 대비 최대 2.0% 향상된 Dice 점수를 기록한다.

ABSTRACT

Automatic segmentation of magnetic resonance (MR) images is crucial for morphological evaluation of the pediatric musculoskeletal system in clinical practice. However, the accuracy and generalization performance of individual segmentation models are limited due to the restricted amount of annotated pediatric data. Hence, we propose to train a segmentation model on multiple datasets, arising from different parts of the anatomy, in a multi-task and multi-domain learning framework. This approach allows to overcome the inherent scarcity of pediatric data while benefiting from a more robust shared representation. The proposed segmentation network comprises shared convolutional filters, domain-specific batch normalization parameters that compute the respective dataset statistics and a domain-specific segmentation layer. Furthermore, a supervised contrastive regularization is integrated to further improve generalization capabilities, by promoting intra-domain similarity and impose inter-domain margins in embedded space. We evaluate our contributions on two pediatric imaging datasets of the ankle and shoulder joints for bone segmentation. Results demonstrate that the proposed model outperforms state-of-the-art approaches.

연구 동기 및 목표

• 근골격계 세그멘테이션을 위한 소아 MRI 영상에 대한 제한된 애너테이션 데이터 문제를 해결한다.
• 다양한 이질적인 데이터셋을 활용하여 소아 영상에서 딥 러닝 모델의 일반화 능력과 성능을 향상시킨다.
• 해부학적 도메인 간 공유 표현을 갖는 통합 세그멘테이션 프레임워크를 개발한다.
• 감독형 대비 정규화를 통해 특징 공간의 강건성을 향상시켜 도메인 내 유사성과 도메인 간 분리성을 증진시킨다.
• 발목 및 어깨 관절의 두 가지 실제 소아 MRI 데이터셋에서 본 방법의 유효성을 입증한다.

제안 방법

• 모든 도메인에서 공유된 컨볼루션 필터를 사용하는 U-Net 기반 아키텍처를 사용하여 일반적인 해부학적 특징을 학습한다.
• 각 데이터셋의 분포에 맞게 정규화 통계를 적응시키기 위해 도메인별 배치 정규화 레이어를 도입한다.
• 각 해부학적 영역에 맞게 출력 예측을 조정하기 위해 도메인별 세그멘테이션 헤드를 사용한다.
• 동일한 도메인의 특징 임베딩이 가까이 군집되도록 유도하고, 다른 도메인의 임베딩은 서로 멀어지도록 하는 감독형 대비 손실을 적용한다.
• 학습된 공유 표현을 시각화하고 대비 정규화 효과를 평가하기 위해 차원 축소 두 단계(PCA 후 t-SNE)를 활용한다.
• 발목 및 어깨의 서로 다른 해부학적 영역에서의 비짝업된 소아 MRI 데이터셋의 결합된 데이터로 모델을 엔드 투 엔드로 훈련한다.

실험 결과

연구 질문

• RQ1 다양한 해부학적 영역에서의 다수의 소아 MRI 데이터셋으로 훈련된 단일 딥 세그멘테이션 모델이 개별 모델 대비 성능 향상을 이룰 수 있는가?
• RQ2 도메인별 배치 정규화 및 세그멘테이션 레이어 통합이 희소한 소아 데이터에서 모델의 일반화에 어떤 영향을 미치는가?
• RQ3 감독형 대비 정규화가 도메인 내 공유 특징 군집화와 도메인 간 분리에 얼마나 기여하는가?
• RQ4 제안된 다중 작업, 다중 도메인 프레임워크가 소아 뼈 세그멘테이션 과제에서 최신 기술 대비 우월한 성능을 보일 수 있는가?
• RQ5 특히 운동 아티팩트가 더 많은 데이터셋에서, 모델의 이미지 품질 및 노이즈 변화에 대해 얼마나 강건한가?

주요 결과

• 제안된 모델은 발목 및 어깨 데이터셋의 2D 슬라이스에서 Dice 점수 88.9±17.0을 기록하여 기준 U-Net(86.9±18.7) 및 Att-UNet(87.8±18.5) 모델보다 유의미하게 높은 성능을 보였다.
• 대비 정규화를 추가함으로써 공유 표현 공간에서 특징 임베딩의 군집화가 더욱 명확해졌으며, t-SNE 시각화를 통해 도메인 내 클러스터와 도메인 간 간격이 뚜렷하게 관찰되었다.
• 통계 분석 결과, 제안된 방법이 기준 모델 및 공동 훈련 기반 접근법보다 유의미한 향상을 이룬 것으로 확인되었으며, 모든 비교에서 Dice 점수에 대해 p값 < 0.05였다.
• DSL(도메인별 레이어) 기반 접근법이 기준 모델 및 공동 훈련 방식을 모두 초월하여, 레이어 전문화의 유용성을 입증하였다.
• 특히 운동 아티팩트가 더 많은 어깨 데이터셋에서 노이즈 및 변동성에 대해 향상된 강건성을 보였다.
• 시각적 비교를 통해 제안된 모델이 두 데이터셋 모두에서 더 정확하고 일관된 뼈 세그멘테이션 결과를 도출하였으며, 경계 오류가 감소하고 해부학적 구조가 더 잘 유지됨을 확인하였다.