[논문 리뷰] Dual Multi-scale Mean Teacher Network for Semi-supervised Infection Segmentation in Chest CT Volume for COVID-19

Liansheng Wang, Jiacheng Wang외 7명|arXiv (Cornell University)|2022. 11. 10.

COVID-19 diagnosis using AI인용 수 5

한 줄 요약

이 논문은 흉부 CT 영상에서 반감성 3D 코로나19 폐 감염 세그먼테이션을 위한 이중 다중 척도 평균 티처 네트워크(DM²T-Net)를 제안한다. 이는 계층적 다중 척도 특징을 포착하기 위해 다차원 주의 메커니즘을 갖춘 다중 차원 주의 컨볼루션 네트워크(MDA-CNN)를 활용하며, 다양한 예측 헤드를 통해 학생 및 교사 네트워크 간 일관성을 강제한다. 제안된 방법은 COVID-19-P20 데이터셋에서 72.59%의 딱스 점수를 기록하고 MosMedData에서 60.19%를 기록하여, 데이터가 적은 환경에서도 뛰어난 정확도와 일반화 능력을 입증한다.

ABSTRACT

Automated detecting lung infections from computed tomography (CT) data plays an important role for combating COVID-19. However, there are still some challenges for developing AI system. 1) Most current COVID-19 infection segmentation methods mainly relied on 2D CT images, which lack 3D sequential constraint. 2) Existing 3D CT segmentation methods focus on single-scale representations, which do not achieve the multiple level receptive field sizes on 3D volume. 3) The emergent breaking out of COVID-19 makes it hard to annotate sufficient CT volumes for training deep model. To address these issues, we first build a multiple dimensional-attention convolutional neural network (MDA-CNN) to aggregate multi-scale information along different dimension of input feature maps and impose supervision on multiple predictions from different CNN layers. Second, we assign this MDA-CNN as a basic network into a novel dual multi-scale mean teacher network (DM${^2}$T-Net) for semi-supervised COVID-19 lung infection segmentation on CT volumes by leveraging unlabeled data and exploring the multi-scale information. Our DM${^2}$T-Net encourages multiple predictions at different CNN layers from the student and teacher networks to be consistent for computing a multi-scale consistency loss on unlabeled data, which is then added to the supervised loss on the labeled data from multiple predictions of MDA-CNN. Third, we collect two COVID-19 segmentation datasets to evaluate our method. The experimental results show that our network consistently outperforms the compared state-of-the-art methods.

연구 동기 및 목표

• 흉부 CT 영상에서 정확한 감염 세그먼테이션을 위해 2D 기반 방법이 3D 공간적 맥락과 상하좌우 슬라이스 간 관계를 포착하는 데 한계를 가진다는 점을 해결하기 위해.
• 기존 3D 세그먼테이션 모델이 다중 척도 수신장치를 갖추지 못한 점을 보완하기 위해 계층적 특징 학습을 위한 다차원 주의 메커니즘을 설계하기 위해.
• 라벨이 부족한 3D CT 데이터에 대한 과제를 해결하기 위해, 비라벨 데이터를 효과적으로 활용하는 반감성 학습 프레임워크를 개발하기 위해.
• 다양한 데이터셋 간 일반화 능력을 향상시키기 위해 다중 척도 일관성 손실과 다중 예측 감독을 통합하기 위해.

제안 방법

• 3D CT 영상의 다양한 공간 차원을 따라 다중 척도 특징을 집계하기 위해 다중 차원 주의 컨볼루션 네트워크(MDA-CNN)를 설계하여 계층적 표현 학습을 향상시킨다.
• MDA-CNN는 이중 다중 척도 평균 티처 네트워크(DM²T-Net)에 통합되며, 학생 및 교사 네트워크가 서로 다른 네트워크 레이어에서 다중 척도 예측을 생성한다.
• 비라벨 데이터에서 학생 및 교사 예측 간 일관성을 확보하기 위해 다중 척도 일관성 손실을 도입한다. 이는 다양한 척도와 네트워크 깊이에서 적용된다.
• 라벨이 있는 데이터에 대한 지도 학습 손실과 비라벨 데이터에 대한 다중 척도 일관성 손실을 결합하여 반감성 학습을 가능하게 하며, 일반화 능력 향상을 이룬다.
• 교사 네트워크의 가중치는 학생의 가중치를 지수 이동 평균으로 업데이트되어 학습을 안정화시키고 예측 일관성을 향상시킨다.
• 제안된 방법의 종합적 평가를 위해 새로운 3D 흉부 CT 데이터셋 두 개—COVID-19-P20와 MosMedData—를 수집하고 사용한다.

실험 결과

연구 질문

• RQ1 차원별 주의 메커니즘을 갖춘 다중 척도 3D CNN이 단일 척도 또는 2D 접근 방식보다 감염 세그먼테이션 정확도를 향상시키는가?
• RQ2 제안된 다중 척도 일관성 손실이 반감성 3D 세그먼테이션에서 비라벨 데이터를 효과적으로 활용하는 데 얼마나 유용한가?
• RQ3 기존 최신 기술(SOTA) 방법과 비교해 이중 다중 척도 평균 티처 프레임워크가 다양한 CT 영상 프로토콜과 감염 패턴에 대해 더 잘 일반화되는가?
• RQ4 다양한 네트워크 레이어에서의 다중 예측 감독이 모델의 강건성과 세그먼테이션 성능을 얼마나 향상시키는가?

주요 결과

• 제안된 DM²T-Net은 COVID-19-P20 데이터셋에서 72.59%의 딱스 점수를 기록하여 기존 최신 기술(SOTA) 방법을 뛰어넘는다.
• MosMedData 데이터셋에서는 60.19%의 딱스 점수를 기록하여 더 도전적인, 다양한 특성을 지닌 데이터셋에서도 뛰어난 성능을 보인다.
• 다중 척도 일관성 손실은 60 에포크 후 안정적으로 수렴하며, 초기 학습 단계에서 점진적인 증가를 보여 비라벨 데이터의 효과적인 활용이 가능함을 시사한다.
• 시각적 비교 결과, 교사 네트워크의 예측이 학생의 예측보다 항상 더 정확한 것으로 나타나, 다중 척도 감독을 통한 지식 전이의 효과를 확인한다.
• 일반화 분석 결과, 한 데이터셋에서 학습하고 다른 데이터셋에서 테스트했을 때 DM²T-Net은 뛰어난 성능을 유지한다(딱스: 48.81±23.51, HD95: 35.34±30.57), 2D U-Net++ 및 3D UA-MT보다도 뛰어나다.
• 모델은 영상 변동성과 감염 패턴의 다양성에 대해 더 강건한 성능을 보이며, 교차 데이터셋 평가에서 기준 모델보다 낮은 HD95 값과 더 높은 딱스 점수를 기록한다.