[논문 리뷰] W-net: Bridged U-net for 2D Medical Image Segmentation.

Wanli Chen, Yue Zhang외 5명|arXiv (Cornell University)|2018. 07. 12.

Advanced Neural Network Applications인용 수 18

한 줄 요약

이 논문은 2D 의료 영상 분할을 향상시키기 위해 더 깊은 U-net 변종인 W-net을 제안한다. 브리지드 아키텍처를 통해 성능을 향상시키며, 진동을 줄이기 위한 새로운 손실 함수와 훈련 안정성 및 성능을 향상시키기 위한 하이브리드 활성화 함수(ELU + ReLU)를 도입하여 소규모 의료 데이터셋에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

In this paper, we focus on three problems in deep learning based medical image segmentation. Firstly, U-net, as a popular model for medical image segmentation, is difficult to train when convolutional layers increase even though a deeper network usually has a better generalization ability because of more learnable parameters. Secondly, the exponential ReLU (ELU), as an alternative of ReLU, is not much different from ReLU when the network of interest gets deep. Thirdly, the Dice loss, as one of the pervasive loss functions for medical image segmentation, is not effective when the prediction is close to ground truth and will cause oscillation during training. To address the aforementioned three problems, we propose and validate a deeper network that can fit medical image datasets that are usually small in the sample size. Meanwhile, we propose a new loss function to accelerate the learning process and a combination of different activation functions to improve the network performance. Our experimental results suggest that our network is comparable or superior to state-of-the-art methods.

연구 동기 및 목표

• 기울기 소실과 일반화 능력 부족로 인해 소규모 의료 영상 데이터셋에서 깊은 U-net 아키텍처를 훈련시키는 데 도전하는 문제를 해결한다.
• 깊은 네트워크에서 지수형 ReLU(ELU)의 효과가 제한적인 점을 극복하기 위해 활성화 함수를 조합하여 기울기 흐름을 향상시킨다.
• 지속적인 진동과 최적화 수렴의 문제를 해결하기 위해 진짜값 근처의 분할에 특화된 새로운 손실 함수를 제안한다.
• 의료 영상 분할을 위한 아키텍처 혁신과 손실 함수 설계를 통해 모델 성능과 훈련 안정성을 향상시킨다.

제안 방법

• 에코더와 디코더 블록을 잔차 유사 스킵 연결로 연결하는 브리지드 U-net 아키텍처를 제안하여 더 깊은 네트워크에서 기울기 흐름을 향상시킨다.
• 깊은 층에서 비선형성과 기울기 안정성을 균형 잡기 위해 ELU와 ReLU를 조합한 하이브리드 활성화 함수를 도입한다.
• 예측값이 진짜값에 가까워질 때 기울기 업데이트 행동을 수정함으로써 훈련 중 진동을 줄이는 새로운 손실 함수를 설계한다.
• 표준 교차 엔트로피 손실과 Dice 손실과 함께 새로운 손실 함수를 적용하여 수렴을 안정화하고 가속화한다.
• 일반화 능력을 향상시키기 위해 데이터 증강 및 전이 학습 기법을 사용하여 소규모 의료 영상 데이터셋에서 모델을 훈련시킨다.
• 학습 효율성과 성능을 향상시키기 위해 Adam 옵timizer와 학습률 스케줄링을 적용하여 네트워크를 최적화한다.

실험 결과

연구 질문

• RQ1 소규모 의료 영상 데이터셋에서 과적합이나 기울기 소실 없이 더 깊은 U-net 아키텍처를 효과적으로 훈련시킬 수 있는가?
• RQ2 ELU와 ReLU를 조합하면 단독으로 사용할 경우보다 깊은 U-net 변종에서 훈련 안정성과 성능을 향상시키는가?
• RQ3 수정된 손실 함수는 특히 예측값이 진짜값에 가까워질 때 진동을 줄이고 수렴 속도를 향상시키는가?
• RQ4 소규모 데이터셋에서 상태의 기술 수준의 방법과 비교할 때 제안된 W-net 아키텍처는 분할 정확도와 내구성 측면에서 어떻게 성능을 내는가?

주요 결과

• 제안된 W-net는 제한된 훈련 샘플에도 불구하고 벤치마크 의료 영상 데이터셋에서 최신 기술 수준의 방법과 비교해 뛰어난 또는 유사한 분할 성능을 달성한다.
• 하이브리드 활성화 함수(ELU + ReLU)는 표준 ReLU 또는 ELU 단독 사용 대비 훈련 안정성과 모델 수렴을 향상시킨다.
• 새로운 손실 함수는 특히 예측값이 진짜값에 가까워지는 최종 훈련 단계에서 훈련 진동을 크게 줄인다.
• 브리지드 U-net 아키텍처는 더 깊은 네트워크 배포를 가능하게 하며 기울기 흐름을 향상시켜 특징 학습과 분할 정확도를 향상시킨다.
• 모델는 소규모 데이터셋에서 강력한 일반화 능력을 보이며, 의료 영상에서 흔한 데이터 부족 문제에 대한 내구성을 보여준다.
• 정량적 결과는 향상된 Dice 점수와 손실 분산 감소를 확인하여 향상된 훈련 안정성과 성능을 뒷받침한다.