[논문 리뷰] Deeper Image Quality Transfer: Training Low-Memory Neural Networks for 3D Images

Stefano B. Blumberg, Ryutaro Tanno외 2명|arXiv (Cornell University)|2018. 08. 16.

Advanced Image Processing Techniques참고 문헌 11인용 수 7

한 줄 요약

이 논문은 체크포인팅과 가역 네트워크(RevNets)를 조합하여 의료 영상에서 깊이 있는 3D 합성곱 신경망을 훈련하는 메모리 효율적인 프레임워크를 제안한다. 이로 인해 네트워크 깊이를 크게 증가시킬 수 있으며, 메모리 사용량은 오직 4.0% 증가에 그친다. 이 방법은 3D 확산 MRI 초해상도 분석에서 루트 평균 제곱 오차(RMSE)를 13% 감소시켜 이전 최고 성능 모델을 능가하면서도 표준 GPU에서 낮은 메모리 사용량을 유지한다.

ABSTRACT

In this paper we address the memory demands that come with the processing of 3-dimensional, high-resolution, multi-channeled medical images in deep learning. We exploit memory-efficient backpropagation techniques, to reduce the memory complexity of network training from being linear in the network's depth, to being roughly constant $ - $ permitting us to elongate deep architectures with negligible memory increase. We evaluate our methodology in the paradigm of Image Quality Transfer, whilst noting its potential application to various tasks that use deep learning. We study the impact of depth on accuracy and show that deeper models have more predictive power, which may exploit larger training sets. We obtain substantially better results than the previous state-of-the-art model with a slight memory increase, reducing the root-mean-squared-error by $ 13\% $. Our code is publicly available.

연구 동기 및 목표

• 고해상도 의료 영상에서 깊이 있는 3D 합성곱 신경망을 훈련할 때 발생하는 높은 메모리 요구량 문제를 해결하여, 표준 GPU에서 모델의 깊이와 성능을 제한하지 않도록 한다.
• 의료 영상 딥러닝에서 입력 패치 크기 축소나 배치 크기 감소 등의 타협을 유도하는 GPU 메모리 제약을 극복한다.
• 네트워크 깊이에 비례하여 선형에서 근접 상수 수준으로 메모리 복잡도를 감소시켜 더 깊은 네트워크의 훈련을 가능하게 한다.
• 낮은 해상도의 확산 MRI 스캔을 향상시키는 이미지 품질 전이(IQT) 파рад라임에서 제안된 방법의 효과성을 입증한다.
• 메모리 효율성 덕분에 증가된 모델 깊이가 예측 능력과 일반화 능력을 향상시키며, 특히 더 큰 훈련 세트에서 유의미한 성능 향상을 이룬다.

제안 방법

• 활성화 함수 재구성 기능을 통해 출력에서 입력을 복원할 수 있도록 가역 잔여 연결 네트워크(RevNets)를 3D 합성곱 아키텍처에 통합하여, 깊이와 관계없이 메모리 사용량을 근접 상수 수준으로 줄인다.
• 중간 활성화 값을 선택적으로 저장함으로써 메모리 소비를 줄이고, 역전파 과정에서의 메모리 사용량을 감소시키기 위해 체크포인팅을 적용한다.
• RevNets와 체크포인팅을 결합하여 전방 및 역방향 전파를 수동으로 제어함으로써 깊은 네트워크의 메모리 효율적인 훈련을 가능하게 한다.
• 계산 그래프 추상화를 사용하여 전방 및 역방향 전파를 모델링하고, 활성화 값을 저장하는 대신 필요 시에만 재계산한다.
• 3D 이미지 초해상도를 위한 ESPCN 기반 네트워크에 RevNet 블록을 층 간에 삽입하여 깊이를 증가시킨다.
• 표준 최적화 방법을 사용하여 단일 NVIDIA Pascal GPU에서 모델을 훈련하고, 검증 성능 기반으로 모델 선택 및 테스트 세트에서 평가한다.

실험 결과

연구 질문

• RQ1 메모리 효율적인 역전파 기법을 통해 의료 영상에서 깊이 있는 3D 합성곱 신경망을 훈련시킬 수 있는가? 이때 메모리 오버헤드는 최소한으로 유지되는가?
• RQ2 메모리 감소로 가능해진 네트워크 깊이 증가가 이미지 품질 전이 성능 향상에 실제로 기여하는가?
• RQ3 체크포인팅과 RevNets의 조합이 3D 의료 영상 네트워크에서 메모리 사용량을 얼마나 줄일 수 있으며, 훈련 효율성은 유지되는가?
• RQ4 모델 깊이가 3D 확산 MRI 초해상도 분석에서 성능에 어떤 영향을 미치며, 더 큰 훈련 세트에서 유의미한 이점을 얻을 수 있는가?
• RQ5 제안된 방법은 분할, 합성, 정렬과 같은 다른 3D 의료 영상 작업으로 일반화 가능한가?

주요 결과

• 제안된 방법은 단순 역전파 대비 메모리 사용량을 50% 이상 감소시켰으며, 기본 ESPCN 모델 대비 오직 4.0% 증가에 그친다.
• 최고 성능을 보인 모델인 ESPCN-RN4는 32명의 환자에서 이전 최고 성능 모델(9.76 ± 0.51) 대비 루트 평균 제곱 오차(RMSE)를 13% 상대적으로 감소시켰다(8.46 ± 0.67).
• 윌코크슨 부호 순위 검정을 통해 ESPCN-RN4와 이전 최고 성능 모델 간의 향상이 통계적으로 유의미하다는 것을 확인했다(p = 0.0117, W = 0).
• 더 깊은 모델들(예: ESPCN-RN4)은 뇌의 내부 및 외부 영역 모두에서 뛰어난 성능을 보이며, 고주파 수준의 세부 정보와 조직 경계를 잘 유지한다.
• 메모리 증가율이 극히 낮아 더 깊은 네트워크의 훈련이 가능해져 더 큰 훈련 세트의 활용과 일반화 능력 향상이 가능하다.
• 계산 비용 증가에도 불구하고, 모든 모델이 단일 GPU에서 24시간 이내에 훈련이 완료되어 실용적인 타당성을 입증했다.