[논문 리뷰] Invertible Residual Network with Regularization for Effective Medical Image Segmentation

Kashu Yamazaki, Vidhiwar Singh Rathour외 1명|arXiv (Cornell University)|2021. 03. 16.

Advanced Neural Network Applications참고 문헌 18인용 수 4

한 줄 요약

이 논문은 메모리 효율적인 3D 의료 영상 분할을 위한 두 가지 역가능한 잔차 신경망 변형—부분 역가능 잔차(Partially-InvRes)와 완전 역가능 잔차(Fully-InvRes)—을 제안한다. 덧셈 결합과 학습 가능한 압축 연산을 활용하여 역가능한 다운샘플링/업샘플링을 구현함으로써, 훈련 중 메모리 사용을 줄이고 iSeg 2019 및 BraTS 2020에서 경쟁적인 분할 성능를 유지한다. 특히 Fully-InvRes는 기준 3D Unet 대비 4.18GB 대비 2.12GB의 메모리 사용량을 기록한다.

ABSTRACT

Deep Convolutional Neural Networks (CNNs) i.e. Residual Networks (ResNets) have been used successfully for many computer vision tasks, but are difficult to scale to 3D volumetric medical data. Memory is increasingly often the bottleneck when training 3D Convolutional Neural Networks (CNNs). Recently, invertible neural networks have been applied to significantly reduce activation memory footprint when training neural networks with backpropagation thanks to the invertible functions that allow retrieving its input from its output without storing intermediate activations in memory to perform the backpropagation. Among many successful network architectures, 3D Unet has been established as a standard architecture for volumetric medical segmentation. Thus, we choose 3D Unet as a baseline for a non-invertible network and we then extend it with the invertible residual network. In this paper, we proposed two versions of the invertible Residual Network, namely Partially Invertible Residual Network (Partially-InvRes) and Fully Invertible Residual Network (Fully-InvRes). In Partially-InvRes, the invertible residual layer is defined by a technique called additive coupling whereas in Fully-InvRes, both invertible upsampling and downsampling operations are learned based on squeezing (known as pixel shuffle). Furthermore, to avoid the overfitting problem because of less training data, a variational auto-encoder (VAE) branch is added to reconstruct the input volumetric data itself. Our results indicate that by using partially/fully invertible networks as the central workhorse in volumetric segmentation, we not only reduce memory overhead but also achieve compatible segmentation performance compared against the non-invertible 3D Unet. We have demonstrated the proposed networks on various volumetric datasets such as iSeg 2019 and BraTS 2020.

연구 동기 및 목표

• 체적 의료 영상 데이터에서 훈련할 때 3D 컬러션 신경망(CNNs)의 높은 메모리 사용량 문제를 해결한다.
• 라벨이 제한된 의료 영상 데이터에서 과적합을 줄이기 위해 정규화를 통합함으로써 도전 과제를 극복한다.
• 역가능한 잔차 블록을 활용하여 제한된 GPU 메모리 조건에서도 더 깊고 큰 3D 네트워크를 훈련할 수 있도록 한다.
• 일반적인 3D Unet 기준 모델을 향상시키기 위해 역가능성과 변동형 오토에코더(Variational Autoencoder, VAE) 정규화를 통합하여 일반화 능력과 메모리 효율성을 향상시킨다.

제안 방법

• 입력을 두 부분으로 분할하고 임의의 함수를 통해 변환하며 잔차 연결을 적용함으로써, 덧셈 결합을 사용하여 역가능한 잔차 블록을 구현한 Partially-InvRes를 구현한다.
• 학습 가능한 픽셀 셔플링(압축)을 사용하여 역가능한 다운샘플링 및 업샘플링을 설계함으로써, 전체 네트워크의 완전한 역가능성을 보장하는 Fully-InvRes를 설계한다.
• 입력 체적 데이터를 재구성하기 위한 VAE 정규화 브랜치를 통합함으로써, 소규모 의료 데이터셋에서 일반화 능력을 향상시키고 과적합을 줄인다.
• 3D Unet 기반의 네트워크 아키텍처를 구성하며, 표준 잔차 블록 대신 역가능한 대체 블록을 사용하고, 자기지도 학습을 위한 VAE 브랜치를 추가한다.
• 기본 학습률이 2e-4인 ADAM 옵timizer를 사용하여 모델을 훈련하며, 모든 레이어에서 인스턴스 정규화와 Leaky ReLU를 적용한다.
• 중간 활성화 상태를 저장할 필요 없이 역가능한 함수 역전을 통해 메모리 효율적인 역전파를 구현하기 위해 PyTorch의 MemCNN 라이브러리를 활용한다.

실험 결과

연구 질문

• RQ1 역가능한 잔차 신경망은 분할 정확도를 희생시키지 않고 3D 의료 영상 분할에서 메모리 사용량을 줄일 수 있는가?
• RQ2 부분 역가능성과 완전 역가능성을 가진 잔차 신경망은 체적 데이터셋에서 메모리 효율성과 성능 측면에서 어떻게 비교되는가?
• RQ3 VAE 정규화 브랜치를 추가함으로써, 데이터가 적은 의료 영상 환경에서 일반화 능력 향상과 과적합 감소에 어느 정도 기여하는가?
• RQ4 역가능한 네트워크는 제한된 GPU 메모리 조건에서 3D Unet 아키텍처에 효과적으로 통합되어 더 깊은 모델의 훈련을 가능하게 하는가?
• RQ5 제안된 역가능한 네트워크는 iSeg 2019 및 BraTS 2020와 같은 표준 벤치마크에서 비역가능한 3D Unet 기준 모델 대비 어떻게 성능을 내는가?

주요 결과

• iSeg 2019 데이터셋에서 Fully-InvRes는 훈련 시 2.12GB의 최저 메모리 사용량을 기록하여 비역가능한 3D Unet 기준 모델의 4.18GB 대비 크게 감소시켰다.
• iSeg 2019 데이터셋에서 Fully-InvRes는 CSF에 대해 95.09%, GM에 대해 91.32%, WM에 대해 90.31%의 Dice 스코어를 기록하여 기준 모델보다 모든 구조에서 약간 높은 성능를 보였다.
• BraTS 2020 데이터셋에서 Fully-InvRes는 WT에 대해 88.20%, ET에 대해 72.01%, TC에 대해 77.74%의 Dice 스코어를 기록하여 모든 지표에서 기준 3D Unet을 초월했다.
• 두 데이터셋 모두에서 하우소프 거리가 감소했다. 예를 들어 BraTS 2020에서 WT의 경우 Fully-InvRes는 10.46을 기록하여 기준 모델의 11.45보다 개선된 경계 정렬을 보였다.
• 역가능한 블록과 VAE 정규화의 조합은 Partially-InvRes 및 Fully-InvRes가 제한된 훈련 데이터 조건에서도 과적합을 피하고 높은 정확도를 달성하도록 했다.
• 표 4는 두 제안된 모델이 모두 메모리 효율적이며 과적합을 방지하고 높은 정확도를 달성하고 있음을 확인한다. 반면 기준 3D Unet은 메모리 효율성과 과적합 정규화 기능을 모두 갖추지 못하고 있다.