[논문 리뷰] Deep Expectation-Maximization for Semi-Supervised Lung Cancer Screening

Sumeet Menon, David Chapman외 4명|arXiv (Cornell University)|2020. 10. 02.

COVID-19 diagnosis using AI참고 문헌 16인용 수 9

한 줄 요약

이 논문은 뇌종양 스크리닝을 위한 3D 컨volution 신경망(CNN)을 훈련하기 위해 기대값 최대화(EM) 알고리즘을 사용하는 준지도 학습 딥 러닝 프레임워크를 제안한다. 이는 라벨이 부여된 데이터와 라벨이 없는 CT 스캔 데이터를 모두 활용한다. 라벨이 없는 라벨를 잠재 변수로 간주함으로써, 더 작은 데이터셋에서만 라벨이 있는 데이터를 사용하는 완전히 지도 학습 기반 모델을 초월하여 분류 정확도를 향상시키며, NLST 데이터에서 최대 81.1%의 정확도를 달성하여 지도 학습 기반 모델 대비 20% 향상된 성능을 기록한다.

ABSTRACT

We present a semi-supervised algorithm for lung cancer screening in which a 3D Convolutional Neural Network (CNN) is trained using the Expectation-Maximization (EM) meta-algorithm. Semi-supervised learning allows a smaller labelled data-set to be combined with an unlabeled data-set in order to provide a larger and more diverse training sample. EM allows the algorithm to simultaneously calculate a maximum likelihood estimate of the CNN training coefficients along with the labels for the unlabeled training set which are defined as a latent variable space. We evaluate the model performance of the Semi-Supervised EM algorithm for CNNs through cross-domain training of the Kaggle Data Science Bowl 2017 (Kaggle17) data-set with the National Lung Screening Trial (NLST) data-set. Our results show that the Semi-Supervised EM algorithm greatly improves the classification accuracy of the cross-domain lung cancer screening, although results are lower than a fully supervised approach with the advantage of additional labelled data from the unsupervised sample. As such, we demonstrate that Semi-Supervised EM is a valuable technique to improve the accuracy of lung cancer screening models using 3D CNNs.

연구 동기 및 목표

• 폐암 스크리닝에서 라벨이 부여된 의료 영상 데이터의 제한성 문제를 해결하기 위해.
• 크로스 도메인 폐암 분류에서 3D CNN의 일반화 능력과 정확도를 향상시키기 위해.
• EM 알고리즘을 사용하여 라벨이 없는 데이터에서 잠재 라벨을 추론하고 CNN 가중치를 동시에 최적화하는 것이 가능한지 탐색하기 위해.
• 준지도 학습을 통해 대량의 라벨이 없는 CT 스캔을 효과적으로 활용하여 모델 성능을 향상시킬 수 있는지 평가하기 위해.
• 준지도 학습 기반 EM 모델의 성능을 완전히 지도 학습 기반 모델과 전체 라벨이 부여된 훈련의 상한선 성능과 비교하기 위해.

제안 방법

• 이 방법은 관측된(라벨이 있는) 데이터와 관측되지 않은(잠재적) 라벨을 포함한 데이터가 존재하는 상황에서, 기대값 최대화(EM) 메타 알고리즘을 사용하여 3D CNN을 훈련한다.
• E단계에서는 현재 CNN 가중치를 사용하여 라벨이 없는 CT 스캔에 대한 클래스 확률을 예측한다.
• M단계에서는 E단계에서 생성된 가짜 라벨 예측값과 함께 라벨이 있는 데이터를 사용하여 CNN을 재학습함으로써 가능도를 최대화한다.
• 이 프레임워크는 서로 다른 데이터 양과 영상 촬영 프로토콜을 가진 두 개의 별도된 폐암 스크리닝 데이터셋인 Kaggle 데이터 사이언스 볼 2017(Kaggle17)과 국립 폐스크리닝 시험(NLST)을 결합한다.
• 모델은 2층 3D CNN 및 AlexNet 아키텍처를 사용하며, 데이터는 50×50×50 박자로 구성된 20층 조각으로 사전 처리된다.
• 크로스 도메인 평가는 한 데이터셋의 라벨이 있는 데이터로 훈련하고 다른 데이터셋의 라벨이 없는 데이터를 포함시켜 실제 세계의 데이터 부족 상황을 시뮬레이션한다.

실험 결과

연구 질문

• RQ1 EM 알고리즘이 라벨이 없는 CT 스캔을 활용한 준지도 학습에 적용되었을 때, 폐암 스크리닝에서 3D CNN 성능을 효과적으로 향상시킬 수 있는가?
• RQ2 라벨이 있는 데이터만으로 훈련된 완전히 지도 학습 모델과 비교했을 때, 준지도 학습 기반 EM 모델의 성능은 어떠한가?
• RQ3 더 큰 외부 도메인 데이터셋(NLST)의 라벨이 없는 데이터를 포함시켰을 때, 더 작은 내부 도메인 데이터셋(Kaggle17)에서의 분류 정확도는 어느 정도 향상되는가?
• RQ4 라벨이 있는 데이터에 비해 라벨이 없는 데이터의 크기가 증가함에 따라, 준지도 학습 기반 EM의 성능 향상 정도는 어떻게 변화하는가?
• RQ5 라벨이 있는 데이터가 더 작은 데이터셋에서 오고, 라벨이 없는 데이터가 더 큰 데이터셋에서 올 경우, EM 기반 접근 방식이 더 효과적인가?

주요 결과

• Kaggle17의 라벨이 있는 데이터로 훈련하고 NLST의 라벨이 없는 데이터를 포함시켰을 때, 2층 3D CNN의 정확도는 75.95%에서 77.5%로 향상되어 1.55%p의 절대적 향상이 이루어졌다.
• AlexNet 아키텍처의 경우 동일한 크로스 도메인 설정에서 정확도가 79.36%에서 81.1%로 상승하여 1.74%p 향상되었다.
• NLST의 라벨이 있는 데이터로 훈련하고 Kaggle17에서 테스트했을 때, 2층 3D CNN의 정확도는 76.75%에서 78.1%로 1.35%p 향상되었으며, AlexNet은 72.9%에서 74.4%로 상승했다.
• 준지도 학습 기반 EM의 성능 향상은 전체 라벨이 있는 데이터를 모두 사용한 완전히 지도 학습 기반 상한선 성능 향상의 약 절반 수준으로 일관되게 유지되었다.
• 라벨이 있는 데이터가 더 작은 Kaggle17 데이터셋에서 오고, 라벨이 없는 데이터가 더 큰 NLST 데이터셋에서 올 경우, 상대적 향상도가 더 크게 나타나, 라벨이 부족한 환경에서의 데이터 효율성이 높다는 것을 시사한다.
• 결과적으로, EM 기반 준지도 학습이 라벨이 부족한 상황에서 대량의 라벨이 없는 스캔 데이터를 활용할 수 있을 때, 의료 영상에서 3D CNN에 대해 효과적임을 확인하였다.