[논문 리뷰] Two layer Ensemble of Deep Learning Models for Medical Image Segmentation

Truong Dang, Tien Thanh Nguyen외 3명|arXiv (Cornell University)|2021. 04. 10.

Advanced Neural Network Applications참고 문헌 30인용 수 11

한 줄 요약

이 논문은 의료 영상 분할을 위한 이중층 앙상블 딥 러닝 프레임워크를 제안한다. 여기서 첫 번째 계층의 모델 예측 결과를 두 번째 계층 모델의 증강 입력으로 사용하며, 최종 출력은 선형 회귀를 통해 학습된 가중치를 활용해 조합된다. 이로 인해 CAMUS 및 Kvasir-SEG 데이터셋에서 최신 기술 수준(SOTA) 성능을 달성하였으며, Dice 점수는 최대 0.96에 이르고 기준 모델 대비 Hausdorff 거리가 최대 0.3 향상되었다.

ABSTRACT

In recent years, deep learning has rapidly become a method of choice for the segmentation of medical images. Deep Neural Network (DNN) architectures such as UNet have achieved state-of-the-art results on many medical datasets. To further improve the performance in the segmentation task, we develop an ensemble system which combines various deep learning architectures. We propose a two-layer ensemble of deep learning models for the segmentation of medical images. The prediction for each training image pixel made by each model in the first layer is used as the augmented data of the training image for the second layer of the ensemble. The prediction of the second layer is then combined by using a weights-based scheme in which each model contributes differently to the combined result. The weights are found by solving linear regression problems. Experiments conducted on two popular medical datasets namely CAMUS and Kvasir-SEG show that the proposed method achieves better results concerning two performance metrics (Dice Coefficient and Hausdorff distance) compared to some well-known benchmark algorithms.

연구 동기 및 목표

• 작은 의료 영상 데이터셋에서 학습된 딥 러닝 모델의 과적합 및 일반화 능력 부족 문제를 해결하기 위해.
• 앙상블 전략을 통해 다수의 딥 러닝 모델을 효과적으로 융합하여 분할 성능을 향상시키기 위해.
• 학습된 가중치를 활용해 다수의 분할 모델 예측 결과를 계산적으로 효율적으로 융합하는 방법을 개발하기 위해.
• 계층적 앙상블 아키텍처를 통해 개별 모델의 약점을 보완하고 상호 보완적인 강점을 활용하기 위해.
• 표준 분할 성능 지표를 사용하여 벤치마크 의료 영상 데이터셋에서 방법을 검증하기 위해.

제안 방법

• 첫 번째 계층은 다양한 백본(예: ResNet34, VGG16)을 가진 여러 딥 러닝 모델(예: UNet, LinkNet, FPN)로 구성되며, 원본 영상 데이터에서 학습된다.
• 첫 번째 계층 모델의 픽셀 단위 예측 결과는 원본 입력 영상과 결합되어 두 번째 계층 모델의 증강된 훈련 데이터를 형성한다.
• 두 번째 계층 모델은 이러한 증강된 데이터에서 첫 번째 계층 모델의 정보를 통합함으로써 분할 예측을 정밀하게 개선하도록 훈련된다.
• 최종 예측 결과는 선형 회귀를 통해 학습된 가중치를 활용한 가중 평균 방식으로 조합된다.
• 개별 모델의 기여도를 균형 있게 조정하기 위해 최소 제곱 회귀 기반 접근 방식을 사용해 각 클래스별 가중치를 최적화한다.
• 프레임워크는 두 개의 공개 데이터셋인 CAMUS 및 Kvasir-SEG에서 Dice 계수와 Hausdorff 거리 지표를 사용해 평가된다.

실험 결과

연구 질문

• RQ1 작은 의료 영상 데이터셋에서 개별 모델 대비 이중층 앙상블 딥 러닝 모델이 분할 정확도를 향상시킬 수 있는가?
• RQ2 첫 번째 계층 모델의 예측 결과를 두 번째 계층의 증강 입력으로 사용할 경우, 특징 표현 및 모델 일반화 능력이 향상되는가?
• RQ3 학습 가능한 가중치 기반 융합 전략이 단순 평균화나 투표 방식보다 성능이 뛰어나게 작용하는가?
• RQ4 초음파 영상에서 복잡한 해부학적 구조인 심장의 좌심방 및 심내막을 분할하는 데 있어 제안된 앙상블의 성능은 어떠한가?
• RQ5 Hausdorff 거리로 측정했을 때, 앙상블이 경계 검출 오류를 얼마나 줄이는가?

주요 결과

• 제안된 앙상블은 CAMUS 데이터셋의 심내막 클래스에서 0.96의 Dice 계수를 달성하여 두 번째로 우수한 성능을 보인 OLE(0.959)를 초월하였다.
• 좌심방 클래스에서는 0.907의 Dice 점수를 기록하여 OLE(0.900) 및 UNet-ResNet34(0.900)와 같은 다른 모델들을 모두 앞섰다.
• 수축기 말기 단계에서 좌심실에 대해 Hausdorff 거리는 4.4를 기록하여 두 번째로 우수한 모델인 LinkNet-ResNet34(4.7)보다 우수했다.
• 수확기 말기 단계에서 좌심방에 대해 Hausdorff 거리는 4.8을 기록하여 OLE(5.3) 및 기타 모델들을 능가하였다.
• 특히 좌심방의 오른쪽 아래부분과 같이 복잡한 영역을 분할하는 데서 구성 모델들이 실패한 상황에서도 앙상블이 개선된 성능을 보였다.
• VGG16 백본을 사용한 모델들은 성능이 열악하여 Dice 점수가 0.3 이하에 머물러, 백본 선택과 앙상블 융합의 중요성을 강조하였다.