[논문 리뷰] Parallel Multi-Dimensional LSTM, With Application to Fast Biomedical Volumetric Image Segmentation

Marijn Stollenga, Wonmin Byeon외 2명|arXiv (Cornell University)|2015. 06. 24.

Advanced Neural Network Applications참고 문헌 31인용 수 153

한 줄 요약

이 논문은 GPU에서 효율적인 병렬 처리가 가능한 새로운 병렬 처리 가능한 다차원 Long Short-Term Memory 네트워크인 PyraMiD-LSTM을 소개한다. 이는 LSTM 계산을 피라미드형, 격자 기반의 구조로 재조정하여 GPU 병렬 처리를 가능하게 한다. 기존의 MD-LSTM는 순차적 의존성과 낮은 GPU 활용도로 인해 문제가 되지만, PyraMiD-LSTM은 3D 생물의학 영상 분할 분야에서 최신 기준 성능(SOTA)을 달성하여 MRBrainS13 데이터셋에서 새로운 최고 기록을 수립했고, EM-ISBI12에서도 경쟁력 있는 성능을 보였다.

ABSTRACT

Convolutional Neural Networks (CNNs) can be shifted across 2D images or 3D videos to segment them. They have a fixed input size and typically perceive only small local contexts of the pixels to be classified as foreground or background. In contrast, Multi-Dimensional Recurrent NNs (MD-RNNs) can perceive the entire spatio-temporal context of each pixel in a few sweeps through all pixels, especially when the RNN is a Long Short-Term Memory (LSTM). Despite these theoretical advantages, however, unlike CNNs, previous MD-LSTM variants were hard to parallelize on GPUs. Here we re-arrange the traditional cuboid order of computations in MD-LSTM in pyramidal fashion. The resulting PyraMiD-LSTM is easy to parallelize, especially for 3D data such as stacks of brain slice images. PyraMiD-LSTM achieved best known pixel-wise brain image segmentation results on MRBrainS13 (and competitive results on EM-ISBI12).

연구 동기 및 목표

• 기존의 다차원 LSTM(MD-LSTM) 네트워크가 GPU 병렬 처리에 취약하여 대규모 3D 생물의학 영상 분할에 활용하기 어려운 문제를 해결하기 위해.
• MD-LSTM의 계산 구조를 재고함으로써 3D 부피 데이터(예: MRI 및 전자현미경 스택 등)에 대해 효율적이고 확장 가능한 추론을 가능하게 하기 위해.
• 완전한 시공간적 맥락을 포착하는 순환 아키텍처를 사용하여 MRBrainS13 및 EM-ISBI12와 같은 도전적인 벤치마크에서 최신 기준 성능을 달성하기 위해.
• 현대 GPU 하드웨어에 적절히 병렬화된 경우 MD-LSTM가 컨볼루션 네트워크를 능가할 수 있음을 입증하기 위해.

제안 방법

• 표준 MD-LSTM의 계산 순서를 피라미드형 구조로 재조정하여, 각 층이 3D 볼륨의 2D 슬라이스를 처리하도록 하되, 한 행의 모든 픽셀에 대해 완전한 병렬 처리가 가능하도록 한다.
• 표준 45도 스캔 구조를 회전시킬 경우 발생하는 맥락 갭을 제거하기 위해, 추가된 대각선 연결을 사용한 수정된 연결 방식을 도입한다.
• 피라미드 스캔을 확장된 또는 겹치는 컨볼루션의 형태로 간주함으로써 GPU에서 효율적인 CUDA 기반 컨볼루션 연산을 활용하여 고처리량 계산을 가능하게 한다.
• 피라미드 스캔을 3D 볼륨의 모든 볼륨 픽셀에 확장하여 3D 데이터에 동일한 아키텍처를 적용하고, 8개의 방향 스캔을 사용해 전체 부피 맥락을 포착한다.
• 입력, 포현, 출력 게이트 및 셀 상태를 포함한 표준 LSTM 방정식을 사용하지만, 계산 순서를 재정렬하여 전체 행 또는 열을 독립적으로 처리할 수 있도록 한다.
• 피라미드 구조로 인해 더 큰 유효 수신장(3×3보다 큼)을 가진 겹치는, 중복되는 맥락 창을 활용함으로써 정확도와 강건성을 향상시킨다.

실험 결과

연구 질문

• RQ1 다시 설계된 MD-LSTM 아키텍처가 순환 네트워크의 완전한 맥락 모델링 능력을 유지하면서도 GPU 병렬 처리를 효율적으로 가능하게 할 수 있는가?
• RQ2 제안된 PyraMiD-LSTM이 기존의 CNN 및 MD-LSTM 기반 모델 대비 3D 생물의학 영상 데이터셋에서 더 높은 분할 정확도를 달성할 수 있는가?
• RQ3 피라미드형 계산 순서가 3D 부피 분할 작업의 계산 병목 현상을 줄이고 추론 속도를 향상시킬 수 있는가?
• RQ4 PyraMiD-LSTM의 겹치는 다방향 맥락이 분할의 강건성과 일반화 능력을 얼마나 향상시키는가?
• RQ5 이러한 아키텍처는 실제 고해상도 3D 의료 영상 응용 분야에 대해 확장 가능하고 효과적인가?

주요 결과

• PyraMiD-LSTM는 MRBrainS13 벤치마크에서 최고의 성능 기록을 달성하여 DICE, 수정된 하우스도르프 거리, 볼륨 차이 등 모든 지표에서 새로운 최고 기록을 수립했다.
• EM-ISBI12 데이터셋에서 PyraMiD-LSTM는 후처리를 하지 않은 방법 중에서 가장 낮은 랜덤 오차를 기록했으며, 다른 주요 접근 방식을 능가했다.
• MRBrainS13 데이터셋에서 PyraMiD-LSTM는 회색질에 84.82%, 백질에 88.33%, 뇌척수액에 83.72%의 DICE 스코어를 기록하여 최종 순위에서 모든 팀을 압도했다.
• 드롭아웃 정규화는 성능 향상에 기여하지 않았고, 오히려 학습 시간만 증가시켰다. 이는 아키텍처 자체가 이미 강건하다는 것을 시사한다.
• 피라미드형 아키텍처는 효율적인 GPU 병렬 처리를 가능하게 하여 계산 병목 현상을 줄였고, 표준 MD-LSTM보다 3D 부피 데이터에 대한 추론 속도를 높였다.
• 이 방법은 현대 하드웨어에 적절히 병렬화된 경우 순환 네트워크가 전체 맥락 인식 능력을 갖추고 있음에도 불구하고, 분할 작업에서 CNN을 능가할 수 있음을 입증했다.