[논문 리뷰] Neural Ordinary Differential Equations for Semantic Segmentation of Individual Colon Glands

Hans Pinckaers, Geert Litjens|arXiv (Cornell University)|2019. 10. 23.

Radiomics and Machine Learning in Medical ImagingMedicine인용 수 20

한 줄 요약

이 논문은 조직병리학적 영상에서 대장 샘을 위한 세분화를 위해 신경 미분방정식(NODE)을 통합한 U-Net인 U-Node를 제안한다. 표준 잔차 블록 대신 연속 깊이 NODE를 사용함으로써, 2M 대비 30M(93% 감소)의 파rameter 수를 줄이고 메모리 사용량도 감소시키면서도 정확도를 향상시켰으며, GlaS 데이터셋에서 추론 및 훈련 중에 함수 평가 횟수를 변동시켜 적응형 수신 영역을 구현하였다.

ABSTRACT

Automated medical image segmentation plays a key role in quantitative research and diagnostics. Convolutional neural networks based on the U-Net architecture are the state-of-the-art. A key disadvantage is the hard-coding of the receptive field size, which requires architecture optimization for each segmentation task. Furthermore, increasing the receptive field results in an increasing number of weights. Recently, Neural Ordinary Differential Equations (NODE) have been proposed, a new type of continuous depth deep neural network. This framework allows for a dynamic receptive field at a fixed memory cost and a smaller amount of parameters. We show on a colon gland segmentation dataset (GlaS) that these NODEs can be used within the U-Net framework to improve segmentation results while reducing memory load and parameter counts.

연구 동기 및 목표

• 표준 U-Net 아키텍처에서 고정된 수신 영역의 한계를 해결하기 위해, 수동적인 아키텍처 튜닝이 필요하고 깊이가 증가할수록 파rameter 수가 증가한다는 점을 해결하고자 한다.
• 신경 미분방정식(NODE)이 의료 영상 세분화에서 메모리 효율적이고 적응형 수신 영역을 가능하게 할 수 있는지 탐색하고자 한다.
• 기본 U-Net 및 잔차 기반 베이스라인과 비교하여, NODE 기반 네트워크의 GlaS 대장 샘 세분화 벤치마크에서의 성능을 평가하고자 한다.
• NODE를 통해 동적이고 맥락 인식 특징 학습을 통해 모델 크기와 메모리 사용량을 줄일 수 있으며, 이로 인해 세분화 정확도가 향상됨을 입증하고자 한다.

제안 방법

• 표준 잔차 블록 대신 ODE 시스템으로 정의된 연속 깊이 블록을 사용하여 U-Net의 인코더 및 디코더 경로에 신경 ODE를 통합한다.
• ODE 솔버를 통한 효율적 역전파를 위해 연역 방법을 사용하여 중간 활성화 값을 저장하지 않고 메모리 소비를 줄인다.
• 정확도와 계산 비용의 균형을 맞추기 위해 정확도 허용 오차 1e-3로 함께 사용된 오일러-프린스 5차 솔버(dopri5)를 사용한다.
• 이미지 및 레이어별로 함수 평가 횟수(NFE)를 변동시켜, 입력 복잡도에 따라 적응형 수신 영역을 구현한다.
• 데이터 증강(이동, 반전, 회전, 탄성 왜곡, 색상 왜곡)과 교차 엔트로피 손실을 사용하여 Adam 옵timizer로 모델을 훈련시킨다.
• 테스트 시간 증강을 통해 반전 및 평균화를 적용하여 정확도와 일반화 능력을 향상시킨다.

실험 결과

연구 질문

• RQ1 신경 ODE는 모델 크기와 메모리 사용량을 줄이며 의료 영상 세분화 성능을 향상시킬 수 있는가?
• RQ2 NODE의 적응형 수신 영역이 복잡한 조직병리학적 영상에서 더 나은 일반화와 세분화 정확도를 이끌어낼 수 있는가?
• RQ3 NODE 기반 U-Net의 성능은 표준 U-Net 및 잔차 U-Net 변종과 비교해 볼 때 대장 샘 세분화에서 어떻게 나타나는가?
• RQ4 NODE의 연속 깊이 아키텍처는 표준 U-Net 아키텍처보다 훨씬 적은 파rameter로도 높은 정확도를 유지할 수 있는가?

주요 결과

• U-Node 모델은 전체 GlaS 테스트 세트에서 객체 Dice 스코어 0.881을 기록하여 표준 U-Net(0.868)과 U-ResNet(0.757)을 모두 앞섰다.
• U-Node는 Hausdorff 거리 59.5를 기록하여 U-Net(69.6)과 U-ResNet(122)보다 유의미하게 낮아, 더 나은 형태 유사도를 나타냈다.
• U-Net의 3000만 개 파rameter 대비 단 200만 개 파rameter만을 사용했음에도 불구하고, U-Node는 모든 지표에서 원래 U-Net의 성능을 따라하거나 초월했다.
• 정성적 결과에서 볼 때, 노이즈가 적고 경계가 더 선명하게 드러나는 더 나은 세분화 품질을 보였으며, 특히 큰 샘에서 두드러졌다.
• 함수 평가 횟수를 변동시켜 실현된 적응형 수신 영역 덕분에, 네트워크가 맥락 사용을 동적으로 조정할 수 있었으며, 특히 어려운 패치에서는 높은 NFE가 유리하게 작용했다.
• 반복적인 ODE 솔버 호출로 인해 훈련 속도가 표준 U-Net보다 느렸지만, NFE 제어를 통해 추론 시 속도와 정확도 사이의 트레이드오���을 가능하게 했다.